update package access (#5)

dendritex/__init__.py

@@ -16,19 +16,25 @@
 
 __version__ = "0.0.1"
 
+from . import channels
+from . import ions
+# from .channels import *
+# from .channels import __all__ as _channels_all
+# from .ions import *
+# from .ions import __all__ as _ions_all
+# from .neurons import *
+# from .neurons import __all__ as _membranes_all
+from . import neurons
 from ._base import *
 from ._base import __all__ as _base_all
 from ._integrators import *
 from ._integrators import __all__ as _integrators_all
-from .channels import *
-from .channels import __all__ as _channels_all
-from .ions import *
-from .ions import __all__ as _ions_all
-from .neurons import *
-from .neurons import __all__ as _membranes_all
 
-__all__ = (_base_all +
-           _integrators_all +
-           _ions_all +
-           _channels_all +
-           _membranes_all)
+__all__ = (
+    ['neurons', 'ions', 'channels'] +
+    _base_all +
+    _integrators_all
+  # _ions_all +
+  # _channels_all +
+  # _membranes_all
+)

dendritex/_base.py

@@ -97,7 +97,7 @@ def __init__(
       raise ValueError('size must be int, or a tuple/list of int.'
                        f'But we got {type(size)}')
     self.size = size
-    assert len(size) >= 2, 'The size of the dendritic dynamics should be at least 2D: (n_neuron, n_compartment).'
+    assert len(size) >= 1, 'The size of the dendritic dynamics should be at least 1D: (..., n_neuron, n_compartment).'
     self.pop_size: Tuple[int, ...] = size[:-1]
     self.n_compartment: int = size[-1]
 

examples/hh_neuron.py

@@ -22,17 +22,17 @@
 bst.environ.set(dt=0.01 * bu.ms)
 
 
-class HH(dx.SingleCompartmentNeuron):
+class HH(dx.neurons.SingleCompartmentNeuron):
   def __init__(self, size):
     super().__init__(size)
 
-    self.na = dx.SodiumFixed(size, E=50. * bu.mV)
-    self.na.add_elem(dx.INa_HH1952(size))
+    self.na = dx.ions.SodiumFixed(size, E=50. * bu.mV)
+    self.na.add_elem(dx.channels.INa_HH1952(size))
 
-    self.k = dx.PotassiumFixed(size, E=-77. * bu.mV)
-    self.k.add_elem(dx.IK_HH1952(size))
+    self.k = dx.ions.PotassiumFixed(size, E=-77. * bu.mV)
+    self.k.add_elem(dx.channels.IK_HH1952(size))
 
-    self.IL = dx.IL(size, E=-54.387 * bu.mV, g_max=0.03 * (bu.mS / bu.cm ** 2))
+    self.IL = dx.channels.IL(size, E=-54.387 * bu.mV, g_max=0.03 * (bu.mS / bu.cm ** 2))
 
   def step_fun(self, t):
     # dx.euler_step(hh, t, 10 * bu.nA)

examples/thalamus_single_compartment_neurons.py

@@ -131,7 +131,8 @@ def __init__(self, size, V_initializer=bst.init.Constant(-70. * bu.mV)):
 def try_trn_neuron():
   bst.environ.set(dt=0.01 * bu.ms)
 
-  trn = RTC([1, 1])  # [n_neuron, n_compartment]
+  # trn = RTC([1, 1])  # [n_neuron, n_compartment]
+  trn = RTC(1)  # [n_neuron, n_compartment]
   # trn = TRN([1, 1])  # [n_neuron, n_compartment]
   trn.init_state()
 

requirements.txt

@@ -1,6 +1,3 @@
-numpy
-jax
-jaxlib
 brainunit
 brainstate