From 96ef8c29f2166ebc31668e5b98d2e59bd9af315c Mon Sep 17 00:00:00 2001
From: mjpelah <mjpelah@gmail.com>
Date: Wed, 20 Jul 2022 14:01:21 -0400
Subject: [PATCH] cleaned, general update

---
 .../plot_record_extracellular_potentials.py   |  4 ++++
 hnn_core/network.py                           | 21 +++--------------
 hnn_core/viz.py                               | 23 ++++++-------------
 3 files changed, 14 insertions(+), 34 deletions(-)

diff --git a/examples/howto/plot_record_extracellular_potentials.py b/examples/howto/plot_record_extracellular_potentials.py
index 184eaace6e..79dca9d008 100644
--- a/examples/howto/plot_record_extracellular_potentials.py
+++ b/examples/howto/plot_record_extracellular_potentials.py
@@ -44,6 +44,10 @@
 # each electrode.
 net.plot_cells()
 
+###############################################################################
+# Plotting the cell morphologies of the network cells
+net.plot_cell_morphologies()
+
 ###############################################################################
 # The default network consists of 2 layers (L2 and L5), within which the cell
 # somas are arranged in a regular grid, and apical dendrites are aligned along
diff --git a/hnn_core/network.py b/hnn_core/network.py
index a2a4aa66cc..6aecea7902 100644
--- a/hnn_core/network.py
+++ b/hnn_core/network.py
@@ -17,7 +17,7 @@
 from .drives import _check_drive_parameter_values, _check_poisson_rates
 from .cells_default import pyramidal, basket
 from .params import _long_name, _short_name
-from .viz import plot_cell_morphologies, plot_cells, plot_cell_morphology
+from .viz import plot_cells, plot_cell_morphology
 from .externals.mne import _validate_type, _check_option
 from .extracellular import ExtracellularArray
 from .check import _check_gids, _gid_to_type, _string_input_to_list
@@ -1327,24 +1327,9 @@ def plot_cells(self, ax=None, show=True):
         """
         return plot_cells(net=self, ax=ax, show=show)
 
-    def plot_cell_morphologies(self, ax=None, show=True):
-        """Plot the morphology of the network cells
-
-        Parameters
-        ----------
-        ax : instance of matplotlib Axes3D | None
-            An axis object from matplotlib. If None,
-            a new figure is created.
-        show : bool
-            If True, show the figure.
-
-        Returns
-        -------
-        fig : instance of matplotlib Figure
-            The matplotlib figure handle.
-        """
 
-        return plot_cell_morphologies(self, ax=ax, show=show)
+    def plot_cell_morphologies(self, ax=None, show=True):
+        return plot_cell_morphology(self.cell_types, ax=ax, show=show)
 
 
 class _Connectivity(dict):
diff --git a/hnn_core/viz.py b/hnn_core/viz.py
index 551c3896c7..a966fa81f8 100644
--- a/hnn_core/viz.py
+++ b/hnn_core/viz.py
@@ -776,9 +776,8 @@ def plot_cell_morphology(cell, ax, show=True):
     import matplotlib.pyplot as plt
     from mpl_toolkits.mplot3d import Axes3D  # noqa
     cell_list = list()
-    clr_index=0
     colors = ['b', 'c', 'r', 'm']
-    multiple = 0
+    clr_index=0
 
     if ax is None:
         plt.figure()
@@ -787,20 +786,18 @@ def plot_cell_morphology(cell, ax, show=True):
     if type(cell) is dict:
         for ind_cell in cell:
             cell_list = list(cell.values())
-        print("is dict")
     else:
-        print("is not dict")
         cell_list[0]=cell
 
     # Cell is in XZ plane
     #ax.set_xlim((cell_list[0].pos[1] - 250, cell_list[0].pos[1] + 150))
     #ax.set_zlim((cell_list[0].pos[2] - 100, cell_list[0].pos[2] + 1200))
-    cell_radii = [0]
-    total_radius=0
+    cell_radii = list()
+    cell_radii.append(clr_index)
     for clr_index, cell in enumerate(cell_list):
-        radius = 0
 
         # Calculating the radius for cell offset
+        radius = 0
         for sec_name, section in cell.sections.items():
             end_pts = section.end_pts
             xs, ys, zs = list(), list(), list()
@@ -810,11 +807,11 @@ def plot_cell_morphology(cell, ax, show=True):
                 dz = cell.pos[2] - cell.sections['soma'].end_pts[0][2]
                 if radius < pt[0]:
                     radius = pt[0]
-            total_radius+=radius
+        cell_radii.append(radius)
 
         # Plotting the cell
         for sec_name, section in cell.sections.items():
-            ax.set_xlim((total_radius+200))
+            ax.set_xlim((sum(cell_radii, 100)))
             ax.set_zlim((cell.pos[2] - 100, cell.pos[2] + 1200))
             linewidth = _linewidth_from_data_units(ax, section.diam)
             end_pts = section.end_pts
@@ -823,19 +820,13 @@ def plot_cell_morphology(cell, ax, show=True):
                 dx = cell.pos[0] - cell.sections['soma'].end_pts[0][0]
                 dy = cell.pos[1] - cell.sections['soma'].end_pts[0][1]
                 dz = cell.pos[2] - cell.sections['soma'].end_pts[0][2]
-                xs.append(pt[0] + dx + ((radius + cell_radii[-1])*multiple)+100)
+                xs.append(pt[0] + dx + (radius + cell_radii[-1]+100))
                 ys.append(pt[1] + dz)
                 zs.append(pt[2] + dy)
             ax.plot(xs, ys, zs, color=colors[clr_index], linewidth=linewidth)
-        cell_radii.append(radius)
         ax.view_init(0, -90)
         ax.axis('on')
         ax.grid('off')
-        ax.set_yticks([])
-        ax.set_xticks([])
-        print("iteration: ")
-        print(multiple)
-        multiple = multiple + 1
     plt.tight_layout()
     plt_show(show)
     return ax