Feedback

jaxley/modules/cell.py

@@ -179,7 +179,7 @@ def _init_morph_jaxley_spsolve(self):
         )
         self.root_inds = jnp.asarray([0])
 
-        # Generate mapping to dealing with the masking which allows using the custom
+        # Generate mapping to deal with the masking which allows using the custom
         # sparse solver to deal with different nseg per branch.
         self._remapped_node_indices = remap_index_to_masked(
             self._internal_node_inds,

jaxley/utils/solver_utils.py

@@ -14,7 +14,10 @@ def remap_index_to_masked(
     """Convert actual index of the compartment to the index in the masked system.
 
     E.g. if `nsegs = [2, 4]`, then the index `3` would be mapped to `5` because the
-    masked `nsegs` are `[4, 4]`.
+    masked `nsegs` are `[4, 4]`. I.e.:
+
+    original: [0, 1,           2, 3, 4, 5]
+    masked:   [0, 1, (2) ,(3) ,4, 5, 6, 7]
     """
     cumsum_nseg_per_branch = jnp.concatenate(
         [