Changed WormView argument behavior and added plotting

WormSim/Model/WormPlot.py

@@ -1,9 +1,11 @@
 import numpy as np
 from matplotlib import pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap, Normalize
 import json
 import argparse
 import sys
 import os
+import math
 
 
 def validate_file(file_path):
@@ -21,8 +23,7 @@ def main():
 
     args = parser.parse_args()
 
-    fig, ax = plt.subplots()
-    plt.get_current_fig_manager().set_window_title("2D worm motion")
+    fig, axs = plt.subplots(1, 2, figsize=(14, 3))
 
     with open(args.wcon_file, 'r') as f:
         wcon = json.load(f)
@@ -37,15 +38,63 @@ def main():
     tmax = num_steps
     num = 60.
 
-    ax.set_title('2D worm motion', fontsize=title_font_size)
+    axs[0].set_title('2D worm motion', fontsize=title_font_size)
 
     for t in range(0, tmax, int(tmax/num)):
         f = float(t)/tmax
 
         color = "#%02x%02x00" % (int(0xFF*(f)),int(0xFF*(1-f)*0.8))
 
         for x, y in zip(x_arr[t], y_arr[t]):
-            ax.plot(x, y, '.', color=color, markersize=3 if t==0 else 0.4)
+            axs[0].plot(x, y, '.', color=color, markersize=3 if t==0 else 0.4)
+
+    axs[0].set_aspect('equal')
+
+    x_transposed = x_arr.T
+    y_transposed = y_arr.T
+
+    # TODO: Below is same utility function as in WormView.py. Move to a utility file.
+    r = 40e-3
+    n_bar = x_transposed.shape[0]
+    num_steps = x_transposed.shape[1]
+
+    n_seg = int(n_bar - 1)
+
+    # radii along the body of the worm
+    r_i = np.array([
+        r * abs(math.sin(math.acos(((i) - n_seg / 2.0) / (n_seg / 2.0 + 0.2))))
+        for i in range(n_bar)
+    ]).reshape(-1, 1)
+
+    diff_x = np.diff(x_transposed, axis=0)
+    diff_y = np.diff(y_transposed, axis=0)
+
+    arctan = np.arctan2(diff_x, -diff_y)
+    d_arr = np.zeros((n_bar, num_steps))
+
+    # d of worm endpoints is based off of two points, whereas d of non-endpoints is based off of 3 (x, y) points
+    d_arr[:-1, :] = arctan
+    d_arr[1:, :] = d_arr[1:, :] + arctan
+    d_arr[1:-1, :] = d_arr[1:-1, :] / 2
+
+    # TODO: Determine what "0" starting position should be. Currently, worm facing left while horizontal is "white" in curvature plot.
+    d_arr = d_arr - np.pi/2
+
+    # Blue corresponds to -pi, white corresponds to 0, and red corresponds to pi
+    colors = [
+        (0, 'blue'),
+        (0.5, 'white'),
+        (1, 'red')
+    ]
+    custom_cmap = LinearSegmentedColormap.from_list('custom_cmap', colors)
+    norm = Normalize(vmin=-np.pi, vmax=np.pi)
+
+    axs[1].set_title('Body curvature', fontsize=title_font_size)
+    axs[1].imshow(d_arr, aspect='auto', cmap=custom_cmap, norm=norm)
+
+    print(np.max(d_arr))
+    print(np.min(d_arr))
+    print(d_arr)
 
     plt.show()
 

WormSim/Model/WormView.py

@@ -55,11 +55,14 @@ def get_perimeter(x, y, r):
 
     return px, py
 
+
 def main():
 
+    # Default behavior is to use (px, py) if it exists, and if it doesn’t then automatically generate the perimeter from the midline. 
     parser = argparse.ArgumentParser(description="Process some arguments.")
     parser.add_argument('-f', '--wcon_file', type=validate_file, help='WCON file path')
-    parser.add_argument('-c', '--use_computed_perimeter', action='store_true', help='Use a perimeter that is computed from the midline of the worm')
+    parser.add_argument('-s', '--suppress_automatic_generation', action='store_true', help='Suppress the automatic generation of a perimeter which would be computed from the midline of the worm. If (px, py) is not specified in the WCON, a perimeter will not be shown.')
+    parser.add_argument('-i', '--ignore_wcon_perimeter', action='store_true', help='Ignore (px, py) values in the WCON. Instead, a perimeter is automatically generated based on the midline of the worm.')
     parser.add_argument('-r', '--minor_radius', type=float, default=40e-3, help='Minor radius of the worm in millimeters (default: 40e-3)', required=False)
 
     args = parser.parse_args()
@@ -92,19 +95,22 @@ def main():
 
     num_steps = t.size
 
-    if args.use_computed_perimeter:
-        print("Using computed perimeter")
-        minor_radius = args.minor_radius
-        px, py = get_perimeter(x, y, minor_radius)
-    else:
-        print("Using px, py from WCON file")
-
-        try:
+    if "px" in wcon["data"][0] and "py" in wcon["data"][0]:
+        if args.ignore_wcon_perimeter:
+            print("Ignoring (px, py) values in WCON file and computing perimeter from midline.")
+            px, py = get_perimeter(x, y, args.minor_radius)
+        else:
+            print("Using (px, py) from WCON file")
             px = np.array(wcon["data"][0]["px"]).T
             py = np.array(wcon["data"][0]["py"]).T
-        except KeyError:
-            print("Error: px, py not found in WCON file. Run with --use_computed_perimeter to compute the perimeter from the midline.")
-            return 1
+    else:
+        if not args.suppress_automatic_generation:
+            print("Computing perimeter from midline")
+            px, py = get_perimeter(x, y, args.minor_radius)
+        else:
+            print("Not computing perimeter from midline")
+            px = None
+            py = None
 
     def update(ti):
         global midline_plot, perimeter_plot
@@ -118,13 +124,17 @@ def update(ti):
         else:
             midline_plot.set_data(x[:, ti], y[:, ti])
 
-        if perimeter_plot is None:
-            (perimeter_plot,) = ax.plot(px[:, ti], py[:, ti], color="grey", linewidth=1)
-        else:
-            perimeter_plot.set_data(px[:, ti], py[:, ti])
+        if px is not None and py is not None:
+            if perimeter_plot is None:
+                (perimeter_plot,) = ax.plot(px[:, ti], py[:, ti], color="grey", linewidth=1)
+            else:
+                perimeter_plot.set_data(px[:, ti], py[:, ti])
 
     ani = Player(fig, update, maxi=num_steps - 1)
 
+    # TODO WormViewCSV and WormViewWCON - should WormViewCSV just be the original WormView? That's what it initially did. 
+    # TODO Could take out Player and WormViewWCON into separate repo - Taking out Player could be ugly. It is quite coupled with WormView due to the update function. 
+
     plt.show()
 
 if __name__ == "__main__":