Improved display of times etc.

Author: pgleeson

SiberneticReplay.py

@@ -22,13 +22,16 @@
 all_points = []
 all_point_types = []
 
-
 plotter = None
 offset3d_ = (0, 0, 0)
 slider = None
 
 show_boundary = False
 
+max_time = None
+
+verbose = False
+
 
 def get_color_info_for_type(type_):
     """
@@ -37,15 +40,19 @@ def get_color_info_for_type(type_):
     """
 
     if type_ == 1.1:
-        return "#8BEBFC", "liquid 1", 5
+        return "cyan", "liquid 1", 2
     elif type_ == 1.2:
-        return "#3ACFF0", "liquid 2", 5
+        return "red", "liquid 2", 10
+
     elif type_ == 2.1:
-        return "yellow", "elastic 1", 5
+        return "pink", "elastic 1", 5
     elif type_ == 2.2:
         return "#FF0000", "elastic 2", 5
+    elif type_ >= 2.3 and type_ < 2.4:
+        return "lightyellow", "elastic variable", 5
     elif type_ > 2 and type_ < 3:
-        return "#00cc00", "elastic variable", 5
+        return "#00cc00", "elastic variable", 7
+
     elif type_ == 3:
         return "grey", "boundary 0", 3
     elif type_ == 3.1:
@@ -69,7 +76,8 @@ def add_sibernetic_model(
         plotter, \
         offset3d_, \
         slider, \
-        show_boundary
+        show_boundary, \
+        max_time
 
     offset3d_ = offset3d
     plotter = pl
@@ -83,6 +91,8 @@ def add_sibernetic_model(
     time_count = 0
     logStep = None
 
+    dt = None
+
     report_data = None
     count_point_types = {}
 
@@ -91,6 +101,14 @@ def add_sibernetic_model(
         report_data = json.load(open(report_file, "r"))
         print(report_data)
         position_file = os.path.join(sim_dir, "position_buffer.txt")
+        dt = float(report_data.get("dt").split(" ")[0])
+        duration = float(report_data.get("duration").split(" ")[0])
+        max_time = duration
+        time_points = np.linspace(0, duration, int(duration / dt) + 1)
+        print(
+            "Simulation dt: %s ms, duration: %s ms, times: %s"
+            % (dt, duration, time_points)
+        )
 
         if "worm" in report_data["configuration"]:
             muscle_activation_file = os.path.join(
@@ -104,8 +122,14 @@ def add_sibernetic_model(
 
             f, ax = plt.subplots(tight_layout=True)
             ax.imshow(musc_dat, interpolation="none", aspect="auto", cmap="YlOrRd")
+
+            num_ticks = 5
+            ax.set_xticks(np.linspace(0, musc_dat.shape[1], num_ticks))
+            ax.set_xticklabels(np.linspace(0, duration, num_ticks))
+            # quit()
+
             # ax.set_ylim([-1, 1])
-            ax.set_xlabel("Time (s)")
+            ax.set_xlabel("Time (ms)")
             _ = ax.set_ylabel("Muscle")
 
             h_chart = pv.ChartMPL(f, size=(0.35, 0.35), loc=(0.02, 0.06))
@@ -187,23 +211,38 @@ def add_sibernetic_model(
 
     create_mesh(0)
 
-    max_time = len(all_points) - 1
+    slider_text = "Time point"
+    if max_time is None:
+        max_time = len(all_points) - 1
+    else:
+        slider_text = "Time (ms)"
 
     slider = pl.add_slider_widget(
-        create_mesh, rng=[0, max_time], value=0, title="Time point", style="modern"
+        create_mesh, rng=[0, max_time], value=0, title=slider_text, style="modern"
     )
 
     pl.add_checkbox_button_widget(play_animation, value=False)
 
 
-def play_animation(play):
+def play_animation(play_button_active):
     global plotter, last_meshes, all_points, all_point_types, replaying, slider
-    print("Playing animation: %s" % play)
+    print(
+        f"Animation button pressed. Button active {play_button_active}; replaying: {replaying}, slider value: {slider.GetSliderRepresentation().GetValue()}"
+    )
 
-    if not play:
-        replaying = False
-        print("Animation stopped.")
-        return
+    if not play_button_active:
+        if not replaying:
+            print("Animation already stopped - restarting")
+            slider.GetSliderRepresentation().SetValue(0)
+            curr_time = slider.GetSliderRepresentation().GetValue()
+            replaying = True
+            plotter.update()
+            plotter.render()
+        else:
+            curr_time = slider.GetSliderRepresentation().GetValue()
+            replaying = False
+            print("Animation stopped at %s." % curr_time)
+            return
     else:
         replaying = True
         print("Animation started.")
@@ -218,8 +257,8 @@ def play_animation(play):
         curr_time = slider.GetSliderRepresentation().GetValue()
 
         print(
-            " --- Animating step %i (curr_time: %s) of %i, %s"
-            % (i, curr_time, len(all_points), play)
+            " --- Animating step %i/%i (curr_time: %s) of %i, %s"
+            % (i, len(all_points), curr_time, len(all_points), play_button_active)
         )
         next_time = curr_time + 1
         slider.GetSliderRepresentation().SetValue(next_time)
@@ -229,6 +268,8 @@ def play_animation(play):
         plotter.render()
         time.sleep(replay_speed)
 
+    replaying = False
+
 
 def create_mesh(step):
     step_count = step
@@ -255,10 +296,11 @@ def create_mesh(step):
         if show_boundary is False and is_boundary:
             continue
 
-        print(
-            "       - Plotting %i points of type '%s' (%s), color: %s, size: %i"
-            % (len(curr_points), type_, info, color, size)
-        )
+        if verbose:
+            print(
+                "       - Plotting %i points of type '%s' (%s), color: %s, size: %i"
+                % (len(curr_points), type_, info, color, size)
+            )
 
         if len(curr_points) == 0:
             continue
@@ -320,12 +362,22 @@ def create_mesh(step):
         swap_y_z=True,
         include_boundary=include_boundary,
     )
+    plotter.window_size = [1600, 800]
+
     plotter.set_background("white")
     plotter.add_axes()
     plotter.camera_position = "zx"
     plotter.camera.roll = 90
     plotter.camera.elevation = 45
     print(plotter.camera_position)
 
+    def on_close_callback(plotter):
+        global replaying
+        print(
+            f"Plotter window is closing. Performing actions now (replaying: {replaying})."
+        )
+        replaying = False
+
     if "-nogui" not in sys.argv:
-        plotter.show()
+        plotter.show(before_close_callback=on_close_callback, auto_close=True)
+        print("Done showing")