doc

jaxley/modules/base.py

@@ -88,13 +88,22 @@ def __str__(self):
         return f"jx.{type(self).__name__}"
 
     def _append_params_and_states(self, param_dict, state_dict):
+        """Insert the default params of the module (e.g. radius, length).
+        
+        This is run at `__init__()`. It does not deal with channels.
+        """
         for param_name, param_value in param_dict.items():
             self.nodes[param_name] = param_value
         for state_name, state_value in state_dict.items():
             self.nodes[state_name] = state_value
 
     def _gather_channels_from_constituents(self, constituents: List) -> None:
-        """Modifies `self.channels` and `self.nodes`."""
+        """Modify `self.channels` and `self.nodes` with channel info from constituents.
+        
+        This is run at `__init__()`. It takes all branches of constituents (e.g.
+        of all branches when the are assembled into a cell) and adds columns to
+        `.nodes` for the relevant channels.
+        """
         for module in constituents:
             for channel in module.channels:
                 if type(channel).__name__ not in [

jaxley/modules/branch.py

@@ -40,7 +40,6 @@ def __init__(
         self.cumsum_nbranches = jnp.asarray([0, 1])
 
         # Indexing.
-        # TODO: need to take care of setting the `HH` column to False, not NaN.
         self.nodes = pd.concat([c.nodes for c in compartment_list], ignore_index=True)
         self._append_params_and_states(self.branch_params, self.branch_states)
         self.nodes["comp_index"] = np.arange(self.nseg).tolist()