Merge pull request #5 from austinklein/aklein/generate_wcon

Aklein/generate wcon

WormSim/Model/WormPlot.py

@@ -0,0 +1,103 @@
+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):
+    if not os.path.exists(file_path):
+        raise argparse.ArgumentTypeError(f"The file {file_path} does not exist.")
+    if not os.path.isfile(file_path):
+        raise argparse.ArgumentTypeError(f"{file_path} is not a valid file.")
+    return file_path
+
+
+def main():
+
+    parser = argparse.ArgumentParser(description="Process some arguments.")
+    parser.add_argument('-f', '--wcon_file', type=validate_file, help='WCON file path')
+
+    args = parser.parse_args()
+
+    fig, axs = plt.subplots(1, 2, figsize=(14, 3))
+
+    with open(args.wcon_file, 'r') as f:
+        wcon = json.load(f)
+
+    title_font_size = 10
+
+    t_arr = np.array(wcon["data"][0]["t"])
+    x_arr = np.array(wcon["data"][0]["x"])
+    y_arr = np.array(wcon["data"][0]["y"])
+
+    num_steps = t_arr.size
+    tmax = num_steps
+    num = 60.
+
+    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]):
+            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()
+
+
+if __name__ == "__main__":
+    sys.exit(main())

WormSim/Model/WormView.py

@@ -1,122 +1,141 @@
 from matplotlib import pyplot as plt
-
-from numpy import genfromtxt
 import numpy as np
+import json
 import math
 import os
+import argparse
+import sys
+from Player import Player
 
-data = genfromtxt("simdata.csv", delimiter=",").T
-
-print("Loaded data: %s" % (str(data.shape)))
-
-t = data[0]
-
-x_offset = 1
-y_offset = 2
-d_offset = 3
-
-"""
-for i in [(j*3)+y_offset for j in range(49)]:
-    plt.plot(t,my_data[i],label=i)
-plt.legend()"""
-
-fig, ax = plt.subplots()
-plt.get_current_fig_manager().set_window_title("2D WormSim replay")
-ax.set_aspect("equal")
-
-usingObjects = os.path.isfile("objects.csv")
-if usingObjects:
-    Objects = genfromtxt("objects.csv", delimiter=",")
-    for o in Objects:
-        x = o[0]
-        y = o[1]
-        r = o[2]
-        # print("Circle at (%s, %s), radius %s"%(x,y,r))
-        circle1 = plt.Circle((x, y), r, color="b")
-        plt.gca().add_patch(circle1)
-else:
-    print("No objects found")
-
-num_t = 30
-timesteps = len(t)
-
-# Using same variable names as WormView.m
-Sz = len(data)
-Nt = len(data[0])
-Nbar = int((Sz - 1) / 3)
-NSEG = int(Nbar - 1)
-D = 80e-6
-
-R = [
-    D / 2.0 * abs(math.sin(math.acos(((i) - NSEG / 2.0) / (NSEG / 2.0 + 0.2))))
-    for i in range(Nbar)
-]
-
-CoM = np.zeros([Nt, Nbar, 3])
-CoMplot = np.zeros([Nt, 2])
-Dorsal = np.zeros([Nbar, 2])
-Ventral = np.zeros([Nbar, 2])
-
-print(f"Sz: {Sz}, Nt: {Nt}, Nbar: {Nbar}, NSEG: {NSEG}")
-# Dt = data(2,1) - data(1,1);
-
-ventral_plot = None
+# Global variables
 midline_plot = None
-dorsal_plot = None
+perimeter_plot = None
 
-from Player import Player
+def validate_file(file_path):
+    if not os.path.exists(file_path):
+        raise argparse.ArgumentTypeError(f"The file {file_path} does not exist.")
+    if not os.path.isfile(file_path):
+        raise argparse.ArgumentTypeError(f"{file_path} is not a valid file.")
+    return file_path
 
+def get_perimeter(x, y, r):
+    n_bar = x.shape[0]
+    num_steps = x.shape[1]
 
-def update(ti):
-    global dorsal_plot, ventral_plot, midline_plot
-    f = ti / timesteps
+    n_seg = int(n_bar - 1)
 
-    color = "#%02x%02x00" % (int(0xFF * (f)), int(0xFF * (1 - f) * 0.8))
-    print("Time step: %s, fract: %f, color: %s" % (ti, f, color))
-    ds = []
-    xs = []
-    ys = []
+    # 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)
 
-    for i in [(j * 3) + d_offset for j in range(Nbar)]:
-        ds.append(data[i][ti])
+    diff_x = np.diff(x, axis=0)
+    diff_y = np.diff(y, axis=0)
 
-    for i in [(j * 3) + x_offset for j in range(Nbar)]:
-        xs.append(data[i][ti])
+    arctan = np.arctan2(diff_x, -diff_y)
+    d_arr = np.zeros((n_bar, num_steps))
 
-    for i in [(j * 3) + y_offset for j in range(Nbar)]:
-        ys.append(data[i][ti])
+    # 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
 
-    for j in range(Nbar):
-        dX = R[j] * math.cos(ds[j])
-        dY = R[j] * math.sin(ds[j])
+    dx = np.cos(d_arr)*r_i
+    dy = np.sin(d_arr)*r_i
 
-        Dorsal[j, 0] = xs[j] + dX
-        Dorsal[j, 1] = ys[j] + dY
-        Ventral[j, 0] = xs[j] - dX
-        Ventral[j, 1] = ys[j] - dY
+    px = np.zeros((2*n_bar, x.shape[1]))
+    py = np.zeros((2*n_bar, x.shape[1]))
 
-    if dorsal_plot == None:
-        (dorsal_plot,) = ax.plot(Dorsal[:, 0], Dorsal[:, 1], color="grey", linewidth=1)
-    else:
-        dorsal_plot.set_data(Dorsal[:, 0], Dorsal[:, 1])
+    px[:n_bar, :] = x - dx
+    px[n_bar:, :] = np.flipud(x + dx) # Make perimeter counter-clockwise
 
-    if ventral_plot == None:
-        (ventral_plot,) = ax.plot(
-            Ventral[:, 0], Ventral[:, 1], color="grey", linewidth=1
-        )
-    else:
-        ventral_plot.set_data(Ventral[:, 0], Ventral[:, 1])
+    py[:n_bar, :] = y - dy
+    py[n_bar:, :] = np.flipud(y + dy) # Make perimeter counter-clockwise
 
-    if midline_plot == None:
-        (midline_plot,) = ax.plot(
-            xs, ys, color="g", label="t=%sms" % t[ti], linewidth=0.5
-        )
-    else:
-        midline_plot.set_data(xs, ys)
+    return px, py
 
 
-ax.plot()  # Causes an autoscale update.
+def main():
 
-ani = Player(fig, update, maxi=timesteps - 1)
+    # 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('-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)
 
-plt.show()
+    args = parser.parse_args()
+
+    fig, ax = plt.subplots()
+    plt.get_current_fig_manager().set_window_title("2D WormSim replay")
+    ax.set_aspect("equal")
+
+    with open(args.wcon_file, 'r') as f:
+        wcon = json.load(f)
+
+    if "@CelegansNeuromechanicalGaitModulation" in wcon:
+        center_x_arr = wcon["@CelegansNeuromechanicalGaitModulation"]["objects"]["circles"]["x"]
+        center_y_arr = wcon["@CelegansNeuromechanicalGaitModulation"]["objects"]["circles"]["y"]
+        radius_arr = wcon["@CelegansNeuromechanicalGaitModulation"]["objects"]["circles"]["r"]
+
+        for center_x, center_y, radius in zip(center_x_arr, center_y_arr, radius_arr):
+            circle = plt.Circle((center_x, center_y), radius, color="b")
+            ax.add_patch(circle)
+    else:
+        print("No objects found")
+
+        # Set the limits of the plot since we don't have any objects to help with autoscaling
+        ax.set_xlim([-1.5, 1.5])
+        ax.set_ylim([-1.5, 1.5])
+
+    t = np.array(wcon["data"][0]["t"])
+    x = np.array(wcon["data"][0]["x"]).T
+    y = np.array(wcon["data"][0]["y"]).T
+
+    num_steps = t.size
+
+    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
+    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
+        f = ti / num_steps
+
+        color = "#%02x%02x00" % (int(0xFF * (f)), int(0xFF * (1 - f) * 0.8))
+        print("Time step: %s, fract: %f, color: %s" % (ti, f, color))
+
+        if midline_plot is None:
+            (midline_plot,) = ax.plot(x[:, ti], y[:, ti], color="g", label="t=%sms" % t[ti], linewidth=0.5)
+        else:
+            midline_plot.set_data(x[:, ti], y[:, 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__":
+    sys.exit(main())