more accurate line width control for plotting

Cai-Lab-at-University-of-Michigan · Apr 30, 2021 · 6c1d5b3 · 6c1d5b3
1 parent 3944845
commit 6c1d5b3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/ngauge/Neuron.py b/ngauge/Neuron.py
@@ -23,7 +23,7 @@ def __init__(self):
         self.metadata = ""
         """An empty attribute to store metadata from the SWC file"""
 
-    def plot(self, fig=None, ax=None, axis="z", color="blue"):
+    def plot(self, fig=None, ax=None, axis="z", color="blue", linewidth=1):
         """Draws this Neuron as a figure
 
         :param ax: A matplotlib axis object to draw upon
@@ -51,17 +51,18 @@ def plot(self, fig=None, ax=None, axis="z", color="blue"):
         if not ax and not fig:
             fig = plt.figure()
             ax = fig.add_subplot(111)
+            ax.set_aspect("equal")
 
         for layer in self.soma_layers.values():
             if axis == "z":
-                ax.plot([i.x for i in layer], [i.y for i in layer], color=color)
+                ax.plot([i.x for i in layer], [i.y for i in layer], color=color, linewidth=linewidth)
             elif axis == "x":
-                ax.plot([i.y for i in layer], [i.z for i in layer], color=color)
+                ax.plot([i.y for i in layer], [i.z for i in layer], color=color, linewidth=linewidth)
             elif axis == "y":
-                ax.plot([i.x for i in layer], [i.z for i in layer], color=color)
+                ax.plot([i.x for i in layer], [i.z for i in layer], color=color, linewidth=linewidth)
 
         for branch in self.branches:
-            branch.plot(ax=ax, fig=fig, axis=axis, color=color)
+            branch.plot(ax=ax, fig=fig, axis=axis, color=color, linewidth=linewidth)
 
         return fig
 

diff --git a/ngauge/TracingPoint.py b/ngauge/TracingPoint.py
@@ -81,7 +81,7 @@ def __init__(self, x, y, z, r, t, children=None, fid=None, parent=None):
         else:
             raise TypeError("Children must be list of TracingPoint objects")
 
-    def plot(self, ax=None, fig=None, color="blue", axis="z"):
+    def plot(self, ax=None, fig=None, color="blue", axis="z", linewidth=1):
         """Draws this TracingPoint as a figure
 
         :param ax: A matplotlib axis object to draw upon
@@ -124,16 +124,19 @@ def plot(self, ax=None, fig=None, color="blue", axis="z"):
             lc = mc.LineCollection(
                 [[(ix, iy), (jx, jy)] for ix, iy, iz, jx, jy, jz in segments],
                 color=color,
+                linewidths=linewidth
             )
         elif axis == "x":
             lc = mc.LineCollection(
                 [[(iy, iz), (jy, jz)] for ix, iy, iz, jx, jy, jz in segments],
                 color=color,
+                linewidths=linewidth
             )
         elif axis == "y":
             lc = mc.LineCollection(
                 [[(ix, iz), (jx, jz)] for ix, iy, iz, jx, jy, jz in segments],
                 color=color,
+                linewidths=linewidth
             )
         ax.add_collection(lc)
         ax.autoscale()