Merge pull request #31 from AllenCellModeling/nma_fig_update

plot frequencies and put all timestep vtk files in a dir

Author: jcass11

mti_nma/steps/nma/nma.py

@@ -3,7 +3,6 @@
 from typing import Optional, List
 
 import matplotlib
-import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
 import vtk
@@ -53,7 +52,7 @@ def __init__(
     @log_run_params
     def run(
         self,
-        mode_list=list(range(6)),
+        mode_list=list(range(10)),
         avg_df=None,
         struct="Nuc",
         norm_vecs=True,
@@ -119,13 +118,17 @@ def run(
 
         verts, faces = get_vtk_verts_faces(polydata)
         w, v = run_nma(verts, faces)
-        draw_whist(w)
+        fpaths = draw_whist(w, int(verts.shape[0]), nma_data_dir, struct)
         vmags = get_eigvec_mags(v)
 
         # Working the visualization of eigenvectors on VTK
         n = polydata.GetNumberOfPoints()
         writer = vtk.vtkPolyDataWriter()
 
+        # Create directory to hold vtk files
+        vtk_dir = nma_data_dir / "vtk_timestep_files"
+        vtk_dir.mkdir(parents=True, exist_ok=True)
+
         for id_mode in range(9):
 
             # 1st Get eigenvector of interest as a Nx3 array
@@ -145,7 +148,7 @@ def run(
                 num_array=arr_eigenvec,
                 deep=True,
                 array_type=vtk.VTK_DOUBLE)
-            eigenvec.SetName('Eigenvector')
+            eigenvec.SetName("Eigenvector")
 
             # Assign eigenvectors as mesh points
             polydata.GetPointData().AddArray(eigenvec)
@@ -165,11 +168,9 @@ def run(
                 # Write mesh with new coordinates
                 writer.SetInputData(polydata)
                 writer.SetFileName(str(
-                    nma_data_dir / f"avgshape_{struct}_M{id_mode}_T{id_theta:03d}.vtk"))
+                    vtk_dir / f"avgshape_{struct}_M{id_mode}_T{id_theta:03d}.vtk"))
                 writer.Write()
 
-        fig_path = nma_data_dir / f"w_fig_{struct}.pdf"
-        plt.savefig(fig_path, format="pdf")
         w_path = nma_data_dir / f"eigvals_{struct}.npy"
         np.save(w_path, w)
         v_path = nma_data_dir / f"eigvecs_{struct}.npy"
@@ -183,7 +184,9 @@ def run(
             "w_FilePath": w_path,
             "v_FilePath": v_path,
             "vmag_FilePath": vmags_path,
-            "fig_FilePath": fig_path,
+            "eigval_hist_FilePath": fpaths[0],
+            "freq_hist_FilePath": fpaths[1],
+            "time_hist_FilePath": fpaths[2],
             "Structure": struct
         }, index=[0])
 

mti_nma/steps/nma/nma_utils.py

@@ -223,37 +223,78 @@ def get_faces(i, polydata):
     return mesh_verts, mesh_faces
 
 
-def draw_whist(w):
+def draw_whist(w, npts, nma_data_dir, struct, logscale=False):
     """
     Draws a histogram figure of the eigenvalues
 
     Parameters
     ----------
     w: Numpy 1D array
         An array of eigenvalue values
+    npts: int
+        Number of points in mesh
     Returns
     -------
     fig: Matplotlib Figure object
         Figure containg a histogram of eigenvalues
     """
-
-    plt.clf()
-    fig = plt.figure()
-
-    # set binning
-    minval = min(w) - 0.5
-    maxval = max(w) + 0.5
-    if len(w) < 20:
-        N = int(max(w) + 2)
-    else:
-        N = 30
-    bins = np.linspace(minval, maxval, N)
-
-    sb.distplot(w, kde=False, bins=bins)
-    plt.xlabel("Eigenvalues (w2*m/k)")
-    plt.ylabel("Counts")
-
-    return fig
+    xlabels = ["Eigenvalues (w2*m/k)", "Frequencies (Hz)", "Timescales (s)"]
+    fpaths = [nma_data_dir / f"eigval_hist_{struct}.pdf",
+              nma_data_dir / f"freq_hist_{struct}.pdf",
+              nma_data_dir / f"time_hist_{struct}.pdf"]
+
+    # Fix rounding errors to 0
+    idx = [i for i, n in enumerate(w) if n < 10**-10]
+    w[idx] = 0
+    print(w[0:7])
+
+    # set values to eigenvalues, frequencies, or timescales (1/f)
+    for i in range(3):
+        if i == 0:
+            vals = w
+        else:
+            m_total = 10**-13  # kg
+            m_vert = m_total / npts
+            k = 0.0001  # N/m
+            vals = np.sqrt(w * k / m_vert) / (2 * np.pi)
+            if i == 1:
+                tmp = [i for i in vals if i != 0]
+            if i == 2:
+                tmp = [1 / i for i in vals if i != 0]
+            vals = tmp
+
+        plt.clf()
+        fig = plt.figure()
+
+        # set binning
+        minval = min(vals)
+        maxval = max(vals)
+        if len(vals) < 20:
+            N = int(max(vals) + 2)
+        else:
+            N = 40
+        bins = np.linspace(minval, maxval, N)
+
+        if i == 0:
+            sb.distplot(vals, kde=False, bins=bins)
+        if i == 1:
+            if logscale:
+                logbins = np.logspace(np.log10(bins[0]), np.log10(bins[-1]), len(bins))
+                sb.distplot(vals, kde=False, bins=logbins)
+                plt.xscale("log")
+            else:
+                sb.distplot([i / 1000 for i in vals], kde=False, bins=bins / 1000)
+                xlabels[i] = "Frequencies (kHz)"
+        if i == 2:
+            plt.yscale("log")
+            sb.distplot(vals, kde=False, bins=bins)
+        plt.xlabel(xlabels[i])
+        plt.ylabel("Counts")
+
+        fig.savefig(fpaths[i], format="pdf")
+        plt.close(fig)
+
+    return fpaths
 
 
 def color_vertices_by_magnitude(

setup.py

@@ -63,7 +63,7 @@
     "python-dateutil<=2.8.0",  # need <=2.8.0 for quilt3 in step
     "scikit-image",
     "seaborn",
-    "vtk",
+    "vtk>=9.0",
     "prefect[viz]"
 ]