Updated plot_morphology ranges & aspect code (#175)

* Rewrote morphology plot range

* fixed offset issue

* extracted extra kwargs to always show full soma traces

bsb/morphologies.py

@@ -362,17 +362,6 @@ def get_compartment_positions(self, labels=None):
             return self.compartment_tree.get_arrays()[0]
         return [c.end for c in self.get_compartments(labels=labels)]
 
-    def get_plot_range(self, offset=[0.0, 0.0, 0.0]):
-        compartments = self.compartment_tree.get_arrays()[0]
-        n_dimensions = range(compartments.shape[1])
-        mins = np.array([np.min(compartments[:, i]) + offset[i] for i in n_dimensions])
-        max = np.max(
-            np.array(
-                [np.max(compartments[:, i]) - mins[i] + offset[i] for i in n_dimensions]
-            )
-        )
-        return list(zip(mins.tolist(), (mins + max).tolist()))
-
     def get_compartment_tree(self, labels=None):
         if labels is not None:
             return _compartment_tree(self.get_compartments(labels=labels))

bsb/plotting.py

@@ -110,6 +110,52 @@ def wrapper_function(*args, fig=None, cubic=True, swapaxes=True, **kwargs):
     return wrapper_function
 
 
+def _morpho_figure(f):
+    """
+    Decorator for functions that produce a Figure of a morphology. Applies ``@_figure``
+    and can set the offset, range & aspectratio and can swap the Y & Z axis labels.
+
+    Adds the `offset`, `set_range` and `swapaxes` keyword arguments.
+    """
+
+    @functools.wraps(f)
+    @_figure
+    def wrapper_function(
+        morphology,
+        *args,
+        offset=None,
+        set_range=True,
+        fig=None,
+        swapaxes=True,
+        soma_radius=None,
+        **kwargs
+    ):
+        if offset is None:
+            offset = [0.0, 0.0, 0.0]
+        r = f(
+            morphology,
+            *args,
+            fig=fig,
+            offset=offset,
+            set_range=set_range,
+            swapaxes=swapaxes,
+            soma_radius=soma_radius,
+            **kwargs
+        )
+        if set_range:
+            rng = get_morphology_range(morphology, offset=offset, soma_radius=soma_radius)
+            set_scene_range(fig.layout.scene, rng)
+            set_scene_aspect(fig.layout.scene, rng)
+        if swapaxes:
+            axis_labels = dict(xaxis_title="X", yaxis_title="Z", zaxis_title="Y")
+        else:
+            axis_labels = dict(xaxis_title="X", yaxis_title="Y", zaxis_title="Z")
+        fig.update_layout(scene=axis_labels)
+        return r
+
+    return wrapper_function
+
+
 def _input_highlight(f, required=False):
     """
     Decorator for functions that highlight an input region on a Figure.
@@ -355,12 +401,11 @@ def get_branch_traces(branches, traces):
     return fig
 
 
-@_network_figure
+@_morpho_figure
 def plot_morphology(
     morphology,
-    offset=[0.0, 0.0, 0.0],
+    offset=None,
     fig=None,
-    cubic=True,
     swapaxes=True,
     show=True,
     legend=True,
@@ -393,8 +438,6 @@ def plot_morphology(
     )
     for trace in traces:
         fig.add_trace(trace)
-    if set_range:
-        set_scene_range(fig.layout.scene, morphology.get_plot_range(offset=offset))
     return fig
 
 
@@ -507,14 +550,24 @@ def set_scene_range(scene, bounds):
     scene.zaxis.range = bounds[1]
 
 
+def set_scene_aspect(scene, bounds, mode="equal", swapaxes=True):
+    if mode == "equal":
+        ratios = np.array([d[1] - d[0] for d in bounds])
+        ratios = ratios / np.max(ratios)
+        items = zip(["x", "z", "y"] if swapaxes else ["x", "y", "z"], ratios)
+        scene.aspectratio = dict(items)
+    else:
+        scene.aspectmode = mode
+
+
 def set_morphology_scene_range(scene, offset_morphologies):
     """
     Set the range on a scene containing multiple morphologies.
 
     :param scene: A scene of the figure. If the figure itself is given, ``figure.layout.scene`` will be used.
     :param offset_morphologies: A list of tuples where the first element is offset and the 2nd is the :class:`Morphology`
     """
-    bounds = np.array(list(map(lambda m: m[1].get_plot_range(m[0]), offset_morphologies)))
+    bounds = np.array([get_morphology_range(m[1], m[0]) for m in offset_morphologies])
     combined_bounds = np.array(
         list(zip(np.min(bounds, axis=0)[:, 0], np.max(bounds, axis=0)[:, 1]))
     )
@@ -523,6 +576,15 @@ def set_morphology_scene_range(scene, offset_morphologies):
     set_scene_range(scene, combined_bounds)
 
 
+def get_morphology_range(morphology, offset=None, soma_radius=None):
+    if offset is None:
+        offset = [0.0, 0.0, 0.0]
+    r = soma_radius or 0.0
+    itr = enumerate(morphology.flatten(vectors=["x", "y", "z"]))
+    r = [[min(min(v), -r) + offset[i], max(max(v), r) + offset[i]] for i, v in itr]
+    return r
+
+
 def hdf5_plot_spike_raster(spike_recorders, input_region=None, show=True):
     """
     Create a spike raster plot from an HDF5 group of spike recorders.
@@ -846,5 +908,5 @@ def prepare_plot(self):
         if len(self._morphologies) == 0:
             raise MorphologyError("Cannot show empty MorphologyScene")
         for o, m, k in self._morphologies:
-            plot_morphology(m, offset=o, show=False, fig=self.fig, **k)
+            plot_morphology(m, offset=o, show=False, set_range=False, fig=self.fig, **k)
         set_morphology_scene_range(self.fig.layout.scene, self._morphologies)

tests/test_morphologies.py

@@ -353,7 +353,6 @@ def test_legacy_runs_without_errors(self):
         m.get_compartment_network()
         m.get_compartment_positions()
         m.get_compartment_positions(labels=["A"])
-        m.get_plot_range()
         m.get_compartment_tree()
         m.get_compartment_tree(labels=["B"])
         m.get_compartment_submask(["C"])