more options for plotting (#437)

* enh: add more options for plotting

* fix: add more doc to vis

* fix: address comments

* fix: fix copy paste error

* enh: update morph notebook

* fix: fix warning div 0 for 0 len comps

* enh: more tutorial changes and small fix in plot_utils

* fix: blacked

jaxley/modules/base.py

@@ -38,7 +38,7 @@
     concat_and_ignore_empty,
     cumsum_leading_zero,
 )
-from jaxley.utils.plot_utils import plot_comps, plot_morph
+from jaxley.utils.plot_utils import plot_comps, plot_graph, plot_morph
 from jaxley.utils.solver_utils import convert_to_csc
 from jaxley.utils.swc import build_radiuses_from_xyzr
 
@@ -141,7 +141,10 @@ def _update_nodes_with_xyz(self):
                 np.sum(np.diff(xyzr[:, :3], axis=0) ** 2, axis=1)
             ).cumsum()
             branch_len = np.hstack([np.array([0]), branch_len])
-            branch_len = branch_len / branch_len.max() + 2 * i  # add padding like above
+            max_len = branch_len.max()
+            branch_len = (
+                branch_len / (max_len if max_len > 0 else 1) + 2 * i
+            )  # add padding like above
             branch_len[np.isnan(branch_len)] = 0
             branch_lens.append(branch_len)
         branch_lens = np.hstack(branch_lens)
@@ -1357,11 +1360,27 @@ def vis(
     ) -> Axes:
         """Visualize the module.
 
+        Modules can be visualized on one of the cardinal planes (xy, xz, yz) or
+        even in 3D.
+
+        Several options are available:
+        - `line`: All points from the traced morphology (`xyzr`), are connected
+        with a line plot.
+        - `scatter`: All traced points, are plotted as scatter points.
+        - `comp`: Plots the compartmentalized morphology, including radius
+        and shape. (shows the true compartment lengths per default, but this can
+        be changed via the `morph_plot_kwargs`, for details see
+        `jaxley.utils.plot_utils.plot_comps`).
+        - `morph`: Reconstructs the 3D shape of the traced morphology. For details see
+        `jaxley.utils.plot_utils.plot_morph`. Warning: For 3D plots and morphologies
+        with many traced points this can be very slow.
+
         Args:
             ax: An axis into which to plot.
             col: The color for all branches.
             dims: Which dimensions to plot. 1=x, 2=y, 3=z coordinate. Must be a tuple of
                 two of them.
+            type: The type of plot. One of ["line", "scatter", "comp", "morph"].
             morph_plot_kwargs: Keyword arguments passed to the plotting function.
         """
         return self._vis(
@@ -1384,10 +1403,14 @@ def _vis(
     ) -> Axes:
         branches_inds = view["branch_index"].to_numpy()
 
-        if type == "volume":
+        if "comp" in type.lower():
             return plot_comps(
                 self, view, dims=dims, ax=ax, col=col, **morph_plot_kwargs
             )
+        if "morph" in type.lower():
+            return plot_morph(
+                self, view, dims=dims, ax=ax, col=col, **morph_plot_kwargs
+            )
 
         coords = []
         for branch_ind in branches_inds:
@@ -1396,7 +1419,7 @@ def _vis(
             ), "No coordinates available. Use `vis(detail='point')` or run `.compute_xyz()` before running `.vis()`."
             coords.append(self.xyzr[branch_ind])
 
-        ax = plot_morph(
+        ax = plot_graph(
             coords,
             dims=dims,
             col=col,
@@ -1420,7 +1443,7 @@ def _scatter(self, ax, col, dims, view, morph_plot_kwargs):
         coords = self.xyzr[branch_ind]
         interpolated_xyz = interpolate_xyz(comp_fraction, coords)
 
-        ax = plot_morph(
+        ax = plot_graph(
             np.asarray([[interpolated_xyz]]),
             dims=dims,
             col=col,
@@ -1852,11 +1875,25 @@ def vis(
     ) -> Axes:
         """Visualize the module.
 
+        Modules can be visualized on one of the cardinal planes (xy, xz, yz) or
+        even in 3D.
+
+        Several options are available:
+        - `line`: All points from the traced morphology (`xyzr`), are connected
+        with a line plot.
+        - `scatter`: All traced points, are plotted as scatter points.
+        - `comp`: Plots the compartmentalized morphology, including radius
+        and shape. (shows the true compartment lengths per default, but this can
+        be changed via the `morph_plot_kwargs`, for details see
+        `jaxley.utils.plot_utils.plot_comps`).
+        - `morph`: Reconstructs the 3D shape of the traced morphology. For details see
+        `jaxley.utils.plot_utils.plot_morph`. Warning: For 3D plots and morphologies
+        with many traced points this can be very slow.
+
         Args:
             ax: An axis into which to plot.
             col: The color for all branches.
-            type: Whether to plot as points ("scatter"), a line ("line") or the
-                actual volume of the compartment("volume").
+            type: The type of plot. One of ["line", "scatter", "comp", "morph"].
             dims: Which dimensions to plot. 1=x, 2=y, 3=z coordinate. Must be a tuple of
                 two of them.
             morph_plot_kwargs: Keyword arguments passed to the plotting function.