Automatically pass radius and length to the channels

jaxleyverse · Feb 3, 2024 · 4a74630 · 4a74630
1 parent 58b0d86
commit 4a74630
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/jaxley/modules/base.py b/jaxley/modules/base.py
@@ -564,6 +564,10 @@ def _step_channels_state(
             channel_params = {}
             for p in channel_param_names:
                 channel_params[p] = params[p][indices]
+            channel_params["radius"] = params["radius"][indices]
+            channel_params["length"] = params["length"][indices]
+            channel_params["axial_resistivity"] = params["axial_resistivity"][indices]
+
             channel_states = {}
             for s in channel_state_names:
                 channel_states[s] = states[s][indices]
@@ -608,6 +612,10 @@ def _channel_currents(
             channel_params = {}
             for p in channel_param_names:
                 channel_params[p] = params[p][indices]
+            channel_params["radius"] = params["radius"][indices]
+            channel_params["length"] = params["length"][indices]
+            channel_params["axial_resistivity"] = params["axial_resistivity"][indices]
+
             channel_states = {}
             for s in channel_state_names:
                 channel_states[s] = states[s][indices]