Clean up and test

jaxley/modules/base.py

@@ -16,6 +16,7 @@
     _compute_index_of_child,
     _compute_num_children,
     compute_levels,
+    loc_of_index,
 )
 from jaxley.utils.plot_utils import plot_morph
 
@@ -656,8 +657,10 @@ def _scatter(self, ax, col, dims, view, morph_plot_kwargs):
             np.isnan(self.xyzr[branch_ind][:, dims])
         ), "No coordinates available. Use `vis(detail='point')` or run `.compute_xyz()` before running `.vis()`."
 
-        comp_fraction = comp_ind / self.nseg
+        comp_fraction = loc_of_index(comp_ind, self.nseg)
         coords = self.xyzr[branch_ind]
+
+        # Perform a linear interpolation between coordinates to get the location.
         interpolated_loc_x = np.interp(
             comp_fraction, np.linspace(0, 1, len(coords)), coords[:, dims[0]]
         )

jaxley/modules/compartment.py

@@ -159,11 +159,10 @@ def vis(
         dims=(0, 1),
         morph_plot_kwargs: Dict = {},
     ):
-        nodes = self.set_global_index_and_index(self.view)
         return self.pointer._scatter(
             ax=ax,
             col=col,
             dims=dims,
-            view=nodes,
+            view=self.view,
             morph_plot_kwargs=morph_plot_kwargs,
         )

tests/test_plotting_api.py

@@ -21,6 +21,7 @@ def test_cell():
     _, ax = plt.subplots(1, 1, figsize=(3, 3))
     ax = cell.vis(ax=ax)
     ax = cell.branch([0, 1, 2]).vis(ax=ax, col="r")
+    ax = cell.branch(1).comp(0.9).vis(ax=ax, col="b")
 
     # Plot 2.
     cell.branch(0).add_to_group("soma")