Merge pull request #46 from jpd-de/main

Added Step Back and Restart feature

znvis/particle/particle.py

@@ -53,8 +53,9 @@ class Particle:
             A list of mesh objects, one for each time step.
     static : bool (default=False)
             If true, only render the mesh once at initialization. Be careful
-            as this changes the shape of the required position and director
-            to (n_particles, n_dims)
+            as this changes the shape of the required position and director 
+            to (n_particles, n_dims) 
+
     smoothing : bool (default=False)
             If true, apply smoothing to each mesh object as it is rendered.
             This will slow down the initial construction of the mesh objects
@@ -113,7 +114,7 @@ def construct_mesh_list(self):
         """
         self.mesh_list = []
         try:
-            if not self.static:
+            if not self.static: 
                 n_particles = int(self.position.shape[1])
                 n_time_steps = int(self.position.shape[0])
             else:

znvis/particle/vector_field.py

@@ -49,8 +49,8 @@ class VectorField:
             A list of mesh objects, one for each time step.
     static : bool (default=False)
             If true, only render the mesh once at initialization. Be careful
-            as this changes the shape of the required position and direction
-            to (n_particles, n_dims)
+            as this changes the shape of the required position and direction 
+            to (n_particles, n_dims) 
     smoothing : bool (default=False)
             If true, apply smoothing to each mesh object as it is rendered.
             This will slow down the initial construction of the mesh objects
@@ -108,7 +108,7 @@ def construct_mesh_list(self):
         """
         self.mesh_list = []
         try:
-            if not self.static:
+            if not self.static: 
                 n_particles = int(self.position.shape[1])
                 n_time_steps = int(self.position.shape[0])
             else:

znvis/visualizer/visualizer.py

@@ -105,7 +105,11 @@ def __init__(
             for particle in particles:
                 if not particle.static:
                     len_list.append(len(particle.position))
-            self.number_of_steps = min(len_list)
+
+            if len_list == []:
+                self.number_of_steps = 1
+            else:
+                self.number_of_steps = min(len_list)
 
         self.output_folder = pathlib.Path(output_folder).resolve()
         self.frame_folder = self.output_folder / "video_frames"
@@ -114,6 +118,8 @@ def __init__(
         self.renderer_spp = renderer_spp
         self.keep_frames = keep_frames
         self.renderer = renderer
+        self.play_speed = 1
+        self.do_rewind = False
 
         self.obj_folder = self.output_folder / "obj_files"
 
@@ -139,8 +145,14 @@ def _initialize_app(self):
         self.vis.reset_camera_to_default()
 
         # Add actions to the visualizer.
-        self.vis.add_action("Step", self._update_particles)
+        self.vis.add_action("<<<", self._toggle_play_speed_back)
+        self.vis.add_action("<<", self._update_particles_back)
         self.vis.add_action("Play", self._continuous_trajectory)
+        self.vis.add_action(">>", self._update_particles)
+        self.vis.add_action(">>>", self._toggle_play_speed)
+        self.vis.add_action("Toggle Direction", self._toogle_play_direction)
+        self.vis.add_action("Slow", self._toggle_slowmotion)
+        self.vis.add_action("Restart", self._restart_trajectory)
         self.vis.add_action("Export Scene", self._export_scene)
         self.vis.add_action("Screenshot", self._take_screenshot)
         self.vis.add_action("Export Video", self._export_video)
@@ -174,6 +186,7 @@ def _export_video(self, vis):
         -------
         Saves a video locally.
         """
+        self.do_rewind = False
         self.interrupt = 0  # stop live feed if running.
 
         # Create temporary directory
@@ -252,9 +265,20 @@ def _export_scene(self, vis):
         """
         old_state = self.interrupt  # get old state
         self.interrupt = 0  # stop live feed if running.
+        created_mesh = False
         for i, item in enumerate(self.particles):
-            if i == 0:
+            if item.static:
+                if i == 0 and not created_mesh:
+                    mesh = item.mesh_list[0]
+                    created_mesh = True
+                elif i == 0 and created_mesh:
+                    mesh += item.mesh_list[0]
+                else:
+                    continue
+
+            if i == 0 and not created_mesh:
                 mesh = item.mesh_list[self.counter]
+                created_mesh = True
             else:
                 mesh += item.mesh_list[self.counter]
 
@@ -423,6 +447,7 @@ def _continuous_trajectory(self, vis):
         vis : visualizer
                 Object passed during the callback.
         """
+        self.do_rewind = False
         if self.interrupt == 1:
             self._pause_run(vis)
         else:
@@ -576,7 +601,7 @@ def _run_trajectory(self):
         """
         self.interrupt = 1  # set global run state.
         while self.counter < self.number_of_steps:
-            time.sleep(1 / self.frame_rate)
+            time.sleep(1 / (self.frame_rate * self.play_speed))
             o3d.visualization.gui.Application.instance.post_to_main_thread(
                 self.vis, self._update_particles
             )
@@ -607,6 +632,63 @@ def _update_particles(self, visualizer=None, step: int = None):
             else:
                 self.counter += 1
             step = self.counter
+        if self.do_rewind == True:
+            if self.counter <= 1:
+                self.counter = self.number_of_steps - 2
+            else:
+                self.counter -= 2 
+
+        self._draw_particles(visualizer=visualizer)  # draw the particles.
+
+        # draw the vector field if it exists.
+        if self.vector_field is not None:
+            self._draw_vector_field(visualizer=visualizer)
+
+        visualizer.post_redraw()  # re-draw the window.
+
+    def _update_particles_back(self, visualizer=None, step: int = None):
+        """
+        Update the positions of the particles one step back (rewind-feature)
+
+        Parameters
+        ----------
+        step : int
+                Step to update to.
+
+        Returns
+        -------
+        Updates the positions of the particles in the box.
+        """
+        if visualizer is None:
+            visualizer = self.vis
+        if step is None:
+            if self.counter == 0:
+                self.counter = self.number_of_steps - 1
+            else:
+                self.counter -= 1
+            step = self.counter
+        self._draw_particles(visualizer=visualizer)  # draw the particles.
+
+        # draw the vector field if it exists.
+        if self.vector_field is not None:
+            self._draw_vector_field(visualizer=visualizer)
+
+        visualizer.post_redraw()  # re-draw the window.
+
+    def _toogle_play_direction(self, visualizer=None):
+        """
+        Reverts the direction of play.
+
+        Returns
+        -------
+        Rewinds the trajectory.
+        """
+        self.do_rewind = not self.do_rewind
+
+    def _restart_trajectory(self, visualizer=None):
+        if visualizer is None:
+            visualizer = self.vis
+        self.counter = 0
 
         self._draw_particles(visualizer=visualizer)  # draw the particles.
 
@@ -616,6 +698,55 @@ def _update_particles(self, visualizer=None, step: int = None):
 
         visualizer.post_redraw()  # re-draw the window.
 
+    def _toggle_play_speed(self, visualizer=None):
+        """
+        Toggle the play speed from 1 to 2 to 4 to 8 and back to 1.
+
+        """
+        if self.do_rewind == True:
+            self.do_rewind = False
+            self.play_speed = 1
+
+        if self.play_speed == 1:
+            self.play_speed = 2
+        elif self.play_speed == 2:
+            self.play_speed = 4
+        elif self.play_speed == 4:
+            self.play_speed = 8
+        else:
+            self.play_speed = 1
+
+    def _toggle_play_speed_back(self, visualizer=None):
+        """
+        Toggle the play speed from 1 to 2 to 4 to 8 and back to 1.
+        """
+        if self.do_rewind == False:
+            self.play_speed = 1
+
+        self.do_rewind = True
+
+        if self.play_speed == 1:
+            self.play_speed = 2
+        elif self.play_speed == 2:
+            self.play_speed = 4
+        elif self.play_speed == 4:
+            self.play_speed = 8
+        else:
+            self.play_speed = 1
+
+    def _toggle_slowmotion(self, visualizer=None):
+        """
+        Toggle the play speed from 1 to 1/2 to 1/4 to 1/8 and back to 1
+        """
+        if self.play_speed >= 1:
+            self.play_speed = 1 / 2
+        elif self.play_speed == 1 / 2:
+            self.play_speed = 1 / 4
+        elif self.play_speed == 1 / 4:
+            self.play_speed = 1 / 8
+        else:
+            self.play_speed = 1
+
     def run_visualization(self):
         """
         Run the visualization.