Simpler edge indexing

CHANGELOG.md

@@ -23,6 +23,8 @@ net.vis()
 
 - changelog added to CI (#537, #558,  @jnsbck)
 
+- Refactor of channel and synapse stepping internals and how the model is transferred to jax for more efficient and readable code (#487, @jnsbck).
+
 # 0.5.0
 
 ### API changes

jaxley/channels/channel.py

@@ -59,22 +59,22 @@ def change_name(self, new_name: str):
         new_prefix = new_name + "_"
 
         self._name = new_name
-        self.channel_params = {
+        self.params = {
             (
                 new_prefix + key[len(old_prefix) :]
                 if key.startswith(old_prefix)
                 else key
             ): value
-            for key, value in self.channel_params.items()
+            for key, value in self.params.items()
         }
 
-        self.channel_states = {
+        self.states = {
             (
                 new_prefix + key[len(old_prefix) :]
                 if key.startswith(old_prefix)
                 else key
             ): value
-            for key, value in self.channel_states.items()
+            for key, value in self.states.items()
         }
         return self
 

jaxley/channels/hh.py

@@ -17,15 +17,15 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gNa": 0.12,
             f"{prefix}_gK": 0.036,
             f"{prefix}_gLeak": 0.0003,
             f"{prefix}_eNa": 50.0,
             f"{prefix}_eK": -77.0,
             f"{prefix}_eLeak": -54.3,
         }
-        self.channel_states = {
+        self.states = {
             f"{prefix}_m": 0.2,
             f"{prefix}_h": 0.2,
             f"{prefix}_n": 0.2,

jaxley/channels/pospischil.py

@@ -40,11 +40,11 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gLeak": 1e-4,
             f"{prefix}_eLeak": -70.0,
         }
-        self.channel_states = {}
+        self.states = {}
         self.current_name = f"i_{prefix}"
 
     def update_states(
@@ -77,12 +77,12 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gNa": 50e-3,
             "eNa": 50.0,
             "vt": -60.0,  # Global parameter, not prefixed with `Na`.
         }
-        self.channel_states = {f"{prefix}_m": 0.2, f"{prefix}_h": 0.2}
+        self.states = {f"{prefix}_m": 0.2, f"{prefix}_h": 0.2}
         self.current_name = f"i_Na"
 
     def update_states(
@@ -148,12 +148,12 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gK": 5e-3,
             "eK": -90.0,
             "vt": -60.0,  # Global parameter, not prefixed with `Na`.
         }
-        self.channel_states = {f"{prefix}_n": 0.2}
+        self.states = {f"{prefix}_n": 0.2}
         self.current_name = f"i_K"
 
     def update_states(
@@ -204,12 +204,12 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gKm": 0.004e-3,
             f"{prefix}_taumax": 4000.0,
             f"eK": -90.0,
         }
-        self.channel_states = {f"{prefix}_p": 0.2}
+        self.states = {f"{prefix}_p": 0.2}
         self.current_name = f"i_K"
 
     def update_states(
@@ -261,11 +261,11 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gCaL": 0.1e-3,
             "eCa": 120.0,
         }
-        self.channel_states = {f"{prefix}_q": 0.2, f"{prefix}_r": 0.2}
+        self.states = {f"{prefix}_q": 0.2, f"{prefix}_r": 0.2}
         self.current_name = f"i_Ca"
 
     def update_states(
@@ -329,12 +329,12 @@ def __init__(self, name: Optional[str] = None):
 
         super().__init__(name)
         prefix = self._name
-        self.channel_params = {
+        self.params = {
             f"{prefix}_gCaT": 0.4e-4,
             f"{prefix}_vx": 2.0,
             "eCa": 120.0,  # Global parameter, not prefixed with `CaT`.
         }
-        self.channel_states = {f"{prefix}_u": 0.2}
+        self.states = {f"{prefix}_u": 0.2}
         self.current_name = f"i_Ca"
 
     def update_states(