Add WCON viewers

.github/workflows/build.yml

@@ -42,15 +42,15 @@ jobs:
     - name: Install Python dependencies
       run: |
 
-        pip install numpy matplotlib  
+        pip install numpy matplotlib jsonschema
 
     - name: Build and run model
       run: |
         cd WormSim/Model
         ./test_all.sh
+        ls -alth
 
-
-
+
     - name: Print some output of the executed model
       run: |
         cd WormSim/Model

.gitignore

@@ -28,3 +28,5 @@
 .DS_Store
 .ipynb_checkpoints
 __pycache__
+/WormSim/Model/simdata.png
+/WormSim/Model/simdata.wcon

README.md

@@ -10,3 +10,4 @@ A model of <i>C. elegans</i> locomotion described in Boyle, Berri and Cohen, [Ga
 For information on installing and running this code, see [here](https://github.com/OpenSourceBrain/CelegansNeuromechanicalGaitModulation/tree/master/WormSim).
 
 
+

WormSim/Model/WormPlot.py

@@ -0,0 +1,110 @@
+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="Loads a WCON file and plots/saves an image of the path of the worm and body curvature")
+    parser.add_argument('-f', '--wcon_file', type=validate_file, help='WCON file path', required=True)
+    parser.add_argument('-nogui', action='store_true', help="If specified, only load file and generate the png, don't show GUI")
+
+    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)
+
+    filename = args.wcon_file.replace(".wcon", ".png")
+    plt.savefig(filename, bbox_inches="tight", dpi=300)
+    print("Saved plot image to: %s" % filename)
+
+
+    if not args.nogui:
+        plt.show()
+
+
+if __name__ == "__main__":
+    sys.exit(main())