diff --git a/znvis/particle/particle.py b/znvis/particle/particle.py index b0234c2..719073d 100644 --- a/znvis/particle/particle.py +++ b/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: diff --git a/znvis/particle/vector_field.py b/znvis/particle/vector_field.py index a5f746f..c47dba0 100644 --- a/znvis/particle/vector_field.py +++ b/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: diff --git a/znvis/visualizer/visualizer.py b/znvis/visualizer/visualizer.py index 41ab1f7..cf0785d 100644 --- a/znvis/visualizer/visualizer.py +++ b/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.