Bugfix for .initialize() on network

jaxleyverse · Nov 22, 2023 · 084f865 · 084f865
1 parent 1e70458
commit 084f865
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 37 deletions.
diff --git a/jaxley/modules/base.py b/jaxley/modules/base.py
@@ -195,6 +195,7 @@ def _append_to_channel_params_and_state(
     def _append_to_channel_nodes(self, index, channel):
         """Adds channel nodes from constituents to `self.channel_nodes`."""
         name = type(channel).__name__
+
         if name in self.channel_nodes:
             self.channel_nodes[name] = pd.concat(
                 [self.channel_nodes[name], index]

diff --git a/jaxley/modules/network.py b/jaxley/modules/network.py
@@ -52,6 +52,44 @@ def __init__(
         ]
 
         self.initialize()
+
+        # Indexing.
+        self.nodes = pd.DataFrame(
+            dict(
+                comp_index=np.arange(self.nseg * self.total_nbranches).tolist(),
+                branch_index=(
+                    np.arange(self.nseg * self.total_nbranches) // self.nseg
+                ).tolist(),
+                cell_index=list(
+                    itertools.chain(
+                        *[
+                            [i] * (self.nseg * b)
+                            for i, b in enumerate(self.nbranches_per_cell)
+                        ]
+                    )
+                ),
+            )
+        )
+
+        # Channel indexing.
+        for i, cell in enumerate(self.cells):
+            for channel in cell.channels:
+                name = type(channel).__name__
+                comp_inds = deepcopy(cell.channel_nodes[name]["comp_index"].to_numpy())
+                branch_inds = deepcopy(
+                    cell.channel_nodes[name]["branch_index"].to_numpy()
+                )
+                comp_inds += self.nseg * self.cumsum_nbranches[i]
+                branch_inds += self.cumsum_nbranches[i]
+                index = pd.DataFrame.from_dict(
+                    dict(
+                        comp_index=comp_inds,
+                        branch_index=branch_inds,
+                        cell_index=[i] * len(comp_inds),
+                    )
+                )
+                self._append_to_channel_nodes(index, channel)
+
         self.initialized_conds = False
 
     def _append_synapses_to_params_and_state(self, connectivities):
@@ -106,43 +144,6 @@ def init_morph(self):
             exclude_first=False,
         )
 
-        # Indexing.
-        self.nodes = pd.DataFrame(
-            dict(
-                comp_index=np.arange(self.nseg * self.total_nbranches).tolist(),
-                branch_index=(
-                    np.arange(self.nseg * self.total_nbranches) // self.nseg
-                ).tolist(),
-                cell_index=list(
-                    itertools.chain(
-                        *[
-                            [i] * (self.nseg * b)
-                            for i, b in enumerate(self.nbranches_per_cell)
-                        ]
-                    )
-                ),
-            )
-        )
-
-        # Channel indexing.
-        for i, cell in enumerate(self.cells):
-            for channel in cell.channels:
-                name = type(channel).__name__
-                comp_inds = deepcopy(cell.channel_nodes[name]["comp_index"].to_numpy())
-                branch_inds = deepcopy(
-                    cell.channel_nodes[name]["branch_index"].to_numpy()
-                )
-                comp_inds += self.nseg * self.cumsum_nbranches[i]
-                branch_inds += self.cumsum_nbranches[i]
-                index = pd.DataFrame.from_dict(
-                    dict(
-                        comp_index=comp_inds,
-                        branch_index=branch_inds,
-                        cell_index=[i] * len(comp_inds),
-                    )
-                )
-                self._append_to_channel_nodes(index, channel)
-
         self.initialized_morph = True
 
     def init_conds(self, params):