From 165f331a7e0807971cd7159daacc476f8c116512 Mon Sep 17 00:00:00 2001 From: Austin Klein Date: Thu, 25 Jul 2024 11:21:34 -0700 Subject: [PATCH] Changed WormView argument behavior and added plotting --- WormSim/Model/WormPlot.py | 57 ++++++++++++++++++++++++++++++++++++--- WormSim/Model/WormView.py | 42 ++++++++++++++++++----------- 2 files changed, 79 insertions(+), 20 deletions(-) diff --git a/WormSim/Model/WormPlot.py b/WormSim/Model/WormPlot.py index 932671a..5c5ac12 100644 --- a/WormSim/Model/WormPlot.py +++ b/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,7 +38,7 @@ 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 @@ -45,7 +46,55 @@ def main(): 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() diff --git a/WormSim/Model/WormView.py b/WormSim/Model/WormView.py index 4aa9d75..3c2dd36 100644 --- a/WormSim/Model/WormView.py +++ b/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__":