From 883b0162dc4c82792671ddedfc315dc1b00eb61c Mon Sep 17 00:00:00 2001 From: Paul Hohenberger <49005843+phohenberger@users.noreply.github.com> Date: Fri, 24 May 2024 15:54:45 +0200 Subject: [PATCH 1/6] Add vector field support --- znvis/__init__.py | 4 ++ znvis/mesh/arrow.py | 89 ++++++++++++++++++++++++++ znvis/particle/vector_field.py | 110 +++++++++++++++++++++++++++++++++ znvis/visualizer/visualizer.py | 43 +++++++++++++ 4 files changed, 246 insertions(+) create mode 100644 znvis/mesh/arrow.py create mode 100644 znvis/particle/vector_field.py diff --git a/znvis/__init__.py b/znvis/__init__.py index ba99d48..5306b99 100644 --- a/znvis/__init__.py +++ b/znvis/__init__.py @@ -28,12 +28,16 @@ from znvis.mesh.custom import CustomMesh from znvis.mesh.cylinder import Cylinder from znvis.mesh.sphere import Sphere +from znvis.mesh.arrow import Arrow from znvis.particle.particle import Particle +from znvis.particle.vector_field import VectorField from znvis.visualizer.visualizer import Visualizer __all__ = [ Particle.__name__, Sphere.__name__, + Arrow.__name__, + VectorField.__name__, Visualizer.__name__, Cylinder.__name__, CustomMesh.__name__, diff --git a/znvis/mesh/arrow.py b/znvis/mesh/arrow.py new file mode 100644 index 0000000..cdc34cf --- /dev/null +++ b/znvis/mesh/arrow.py @@ -0,0 +1,89 @@ +""" +ZnVis: A Zincwarecode package. +License +------- +This program and the accompanying materials are made available under the terms +of the Eclipse Public License v2.0 which accompanies this distribution, and is +available at https://www.eclipse.org/legal/epl-v20.html +SPDX-License-Identifier: EPL-2.0 +Copyright Contributors to the Zincwarecode Project. +Contact Information +------------------- +email: zincwarecode@gmail.com +github: https://github.com/zincware +web: https://zincwarecode.com/ +Citation +-------- +If you use this module please cite us with: + +Summary +------- +Create a sphere mesh. +""" + +from dataclasses import dataclass + +import numpy as np +import open3d as o3d + +from znvis.transformations.rotation_matrices import rotation_matrix + +from .mesh import Mesh + + +@dataclass +class Arrow(Mesh): + """ + A class to produce arrow meshes. + + Attributes + ---------- + scale : float + Scale of the arrow + """ + scale: float = 1.0 + resolution = 10 + + def create_mesh( + self, starting_position: np.ndarray, direction: np.ndarray = None + ) -> o3d.geometry.TriangleMesh: + """ + Create a mesh object defined by the dataclass. + + Parameters + ---------- + starting_position : np.ndarray shape=(3,) + Starting position of the mesh. + direction : np.ndarray shape=(3,) (default = None) + Direction of the mesh. + + Returns + ------- + mesh : o3d.geometry.TriangleMesh + """ + + # Calculate the length of the direction vector + direction_length = np.linalg.norm(direction) + + # Scale the arrow size with the direction length + cylinder_radius = 0.06 * direction_length * self.scale + cylinder_height = 0.85 * direction_length * self.scale + cone_radius = 0.15 * direction_length * self.scale + cone_height = 0.15 * direction_length * self.scale + + + arrow = o3d.geometry.TriangleMesh.create_arrow( + cylinder_radius=cylinder_radius, + cylinder_height=cylinder_height, + cone_radius=cone_radius, + cone_height=cone_height, + resolution=self.resolution + ) + + arrow.compute_vertex_normals() + matrix = rotation_matrix(np.array([0, 0, 1]), direction) + arrow.rotate(matrix) + + arrow.translate(starting_position.astype(float) + direction * 0.5 * self.scale) + + return arrow \ No newline at end of file diff --git a/znvis/particle/vector_field.py b/znvis/particle/vector_field.py new file mode 100644 index 0000000..bd3d00e --- /dev/null +++ b/znvis/particle/vector_field.py @@ -0,0 +1,110 @@ +""" +ZnVis: A Zincwarecode package. +License +------- +This program and the accompanying materials are made available under the terms +of the Eclipse Public License v2.0 which accompanies this distribution, and is +available at https://www.eclipse.org/legal/epl-v20.html +SPDX-License-Identifier: EPL-2.0 +Copyright Contributors to the Zincwarecode Project. +Contact Information +------------------- +email: zincwarecode@gmail.com +github: https://github.com/zincware +web: https://zincwarecode.com/ +Citation +-------- +If you use this module please cite us with: + +Summary +------- +Module for the particle parent class +""" + +import typing +from dataclasses import dataclass + +import numpy as np +from rich.progress import track + +from znvis.mesh import Mesh + + +@dataclass +class VectorField: + """ + Parent class for a ZnVis particle. + + Attributes + ---------- + name : str + Name of the particle + mesh : Mesh + Mesh to use e.g. sphere + position : np.ndarray + Position tensor of the shape (n_confs, n_particles, n_dims) + + mesh_list : list + A list of mesh objects, one for each time step. + 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 + but not the deployment. + """ + + name: str + mesh: Mesh = None + position: np.ndarray = None + direction: np.ndarray = None + mesh_list: typing.List[Mesh] = None + + smoothing: bool = False + + def _create_mesh(self, position, direction): + """ + Create a mesh object for the particle. + + Parameters + ---------- + position : np.ndarray + Position of the particle + director : np.ndarray + Director of the particle + + Returns + ------- + mesh : o3d.geometry.TriangleMesh + A mesh object + """ + + mesh = self.mesh.create_mesh(position, direction) + if self.smoothing: + return mesh.filter_smooth_taubin(100) + else: + return mesh + + def construct_mesh_list(self): + """ + Constructor the mesh list for the class. + + The mesh list is a list of mesh objects for each + time step in the parsed trajectory. + + Returns + ------- + Updates the class attributes mesh_list + """ + self.mesh_list = [] + try: + n_particles = int(self.position.shape[1]) + n_time_steps = int(self.position.shape[0]) + except ValueError: + raise ValueError("There is no data for these particles.") + + for i in track(range(n_time_steps), description=f"Building {self.name} Mesh"): + for j in range(n_particles): + if j == 0: + mesh = self._create_mesh(self.position[i][j], self.direction[i][j]) + else: + mesh += self._create_mesh(self.position[i][j], self.direction[i][j]) + self.mesh_list.append(mesh) \ No newline at end of file diff --git a/znvis/visualizer/visualizer.py b/znvis/visualizer/visualizer.py index bfa88f3..84d6995 100644 --- a/znvis/visualizer/visualizer.py +++ b/znvis/visualizer/visualizer.py @@ -61,6 +61,7 @@ class Visualizer: def __init__( self, particles: typing.List[znvis.Particle], + vector_field: typing.List[znvis.VectorField] = None, output_folder: typing.Union[str, pathlib.Path] = "./", frame_rate: int = 24, number_of_steps: int = None, @@ -89,6 +90,7 @@ def __init__( The format of the video to be generated. """ self.particles = particles + self.vector_field = vector_field self.frame_rate = frame_rate self.bounding_box = bounding_box() if bounding_box else None @@ -305,6 +307,11 @@ def _initialize_particles(self): self._draw_particles(initial=True) + def _initialize_vector_field(self): + for item in self.vector_field: + item.construct_mesh_list() + self._draw_vector_field(initial=True) + def _draw_particles(self, visualizer=None, initial: bool = False): """ Draw the particles on the visualizer. @@ -344,6 +351,36 @@ def _draw_particles(self, visualizer=None, initial: bool = False): item.name, item.mesh_list[self.counter], item.mesh.o3d_material ) + + def _draw_vector_field(self, visualizer=None, initial: bool = False): + """ + Draw the vector field on the visualizer. + + Parameters + ---------- + initial : bool (default = True) + If true, no particles are removed. + + Returns + ------- + updates the information in the visualizer. + ----- + """ + if visualizer is None: + visualizer = self.vis + + if initial: + for i, item in enumerate(self.vector_field): + visualizer.add_geometry( + item.name, item.mesh_list[self.counter], item.mesh.o3d_material + ) + else: + for i, item in enumerate(self.vector_field): + visualizer.remove_geometry(item.name) + visualizer.add_geometry( + item.name, item.mesh_list[self.counter], item.mesh.o3d_material + ) + def _continuous_trajectory(self, vis): """ Button command for running the simulation in the visualizer. @@ -509,6 +546,10 @@ def _update_particles(self, visualizer=None, step: int = None): step = self.counter self._draw_particles(visualizer=visualizer) # draw the particles. + + if self.vector_field is not None: + self._draw_vector_field(visualizer=visualizer) + visualizer.post_redraw() # re-draw the window. def run_visualization(self): @@ -521,6 +562,8 @@ def run_visualization(self): """ self._initialize_app() self._initialize_particles() + if self.vector_field is not None: + self._initialize_vector_field() self.vis.reset_camera_to_default() self.app.run() From 79ef8100d4b1deb235b07240213f41e9de77a921 Mon Sep 17 00:00:00 2001 From: Paul Hohenberger <49005843+phohenberger@users.noreply.github.com> Date: Fri, 24 May 2024 18:24:23 +0200 Subject: [PATCH 2/6] Fix docstrings, comments and typing --- znvis/mesh/arrow.py | 4 +--- znvis/particle/vector_field.py | 27 ++++++++++++++------------- znvis/visualizer/visualizer.py | 3 ++- 3 files changed, 17 insertions(+), 17 deletions(-) diff --git a/znvis/mesh/arrow.py b/znvis/mesh/arrow.py index cdc34cf..db2f0c7 100644 --- a/znvis/mesh/arrow.py +++ b/znvis/mesh/arrow.py @@ -62,16 +62,13 @@ def create_mesh( mesh : o3d.geometry.TriangleMesh """ - # Calculate the length of the direction vector direction_length = np.linalg.norm(direction) - # Scale the arrow size with the direction length cylinder_radius = 0.06 * direction_length * self.scale cylinder_height = 0.85 * direction_length * self.scale cone_radius = 0.15 * direction_length * self.scale cone_height = 0.15 * direction_length * self.scale - arrow = o3d.geometry.TriangleMesh.create_arrow( cylinder_radius=cylinder_radius, cylinder_height=cylinder_height, @@ -84,6 +81,7 @@ def create_mesh( matrix = rotation_matrix(np.array([0, 0, 1]), direction) arrow.rotate(matrix) + # Translate the arrow to the starting position and center the origin arrow.translate(starting_position.astype(float) + direction * 0.5 * self.scale) return arrow \ No newline at end of file diff --git a/znvis/particle/vector_field.py b/znvis/particle/vector_field.py index bd3d00e..36e43df 100644 --- a/znvis/particle/vector_field.py +++ b/znvis/particle/vector_field.py @@ -27,23 +27,24 @@ import numpy as np from rich.progress import track -from znvis.mesh import Mesh +from znvis.mesh.arrow import Arrow @dataclass class VectorField: """ - Parent class for a ZnVis particle. + A class to represent a vector field. Attributes ---------- name : str - Name of the particle + Name of the vector field mesh : Mesh - Mesh to use e.g. sphere + Mesh to use position : np.ndarray - Position tensor of the shape (n_confs, n_particles, n_dims) - + Position tensor of the shape (n_steps, n_vectors, n_dims) + direction : np.ndarray + Direction tensor of the shape (n_steps, n_vectors, n_dims) mesh_list : list A list of mesh objects, one for each time step. smoothing : bool (default=False) @@ -53,23 +54,23 @@ class VectorField: """ name: str - mesh: Mesh = None + mesh: Arrow = None # Should be an instance of the Arrow class position: np.ndarray = None direction: np.ndarray = None - mesh_list: typing.List[Mesh] = None + mesh_list: typing.List[Arrow] = None smoothing: bool = False def _create_mesh(self, position, direction): """ - Create a mesh object for the particle. + Create a mesh object for the vector field. Parameters ---------- position : np.ndarray - Position of the particle - director : np.ndarray - Director of the particle + Position of the arrow + direction : np.ndarray + Direction of the arrow Returns ------- @@ -99,7 +100,7 @@ def construct_mesh_list(self): n_particles = int(self.position.shape[1]) n_time_steps = int(self.position.shape[0]) except ValueError: - raise ValueError("There is no data for these particles.") + raise ValueError("There is no data for this vector field.") for i in track(range(n_time_steps), description=f"Building {self.name} Mesh"): for j in range(n_particles): diff --git a/znvis/visualizer/visualizer.py b/znvis/visualizer/visualizer.py index 84d6995..1614938 100644 --- a/znvis/visualizer/visualizer.py +++ b/znvis/visualizer/visualizer.py @@ -547,8 +547,9 @@ def _update_particles(self, visualizer=None, step: int = None): 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) + self._draw_vector_field(visualizer=visualizer) visualizer.post_redraw() # re-draw the window. From ee85d0906da493ffe46b2122b129a37fbef7cbaf Mon Sep 17 00:00:00 2001 From: Paul Hohenberger <49005843+phohenberger@users.noreply.github.com> Date: Fri, 24 May 2024 20:09:58 +0200 Subject: [PATCH 3/6] Fix arrow orientation --- znvis/mesh/arrow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/znvis/mesh/arrow.py b/znvis/mesh/arrow.py index db2f0c7..a2d0044 100644 --- a/znvis/mesh/arrow.py +++ b/znvis/mesh/arrow.py @@ -79,7 +79,7 @@ def create_mesh( arrow.compute_vertex_normals() matrix = rotation_matrix(np.array([0, 0, 1]), direction) - arrow.rotate(matrix) + arrow.rotate(matrix, center=(0, 0, 0)) # Translate the arrow to the starting position and center the origin arrow.translate(starting_position.astype(float) + direction * 0.5 * self.scale) From ced378ebaafacc411fb7691a5225d15275bf569f Mon Sep 17 00:00:00 2001 From: Paul Hohenberger <49005843+phohenberger@users.noreply.github.com> Date: Fri, 24 May 2024 20:14:51 +0200 Subject: [PATCH 4/6] Fix arrow position --- znvis/mesh/arrow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/znvis/mesh/arrow.py b/znvis/mesh/arrow.py index a2d0044..0b1247e 100644 --- a/znvis/mesh/arrow.py +++ b/znvis/mesh/arrow.py @@ -82,6 +82,6 @@ def create_mesh( arrow.rotate(matrix, center=(0, 0, 0)) # Translate the arrow to the starting position and center the origin - arrow.translate(starting_position.astype(float) + direction * 0.5 * self.scale) + arrow.translate(starting_position.astype(float)) return arrow \ No newline at end of file From 6e5393367f1dcd718bccd4ee7501c9a8de602d87 Mon Sep 17 00:00:00 2001 From: Samuel Tovey Date: Fri, 24 May 2024 22:28:32 +0100 Subject: [PATCH 5/6] Update arrow.py Add small changes to the module. --- znvis/mesh/arrow.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/znvis/mesh/arrow.py b/znvis/mesh/arrow.py index 0b1247e..5f9382c 100644 --- a/znvis/mesh/arrow.py +++ b/znvis/mesh/arrow.py @@ -28,7 +28,7 @@ from znvis.transformations.rotation_matrices import rotation_matrix -from .mesh import Mesh +from znvis.mesh import Mesh @dataclass @@ -40,9 +40,11 @@ class Arrow(Mesh): ---------- scale : float Scale of the arrow + resolution : int + Resolution of the mesh. """ scale: float = 1.0 - resolution = 10 + resolution: int = 10 def create_mesh( self, starting_position: np.ndarray, direction: np.ndarray = None @@ -84,4 +86,4 @@ def create_mesh( # Translate the arrow to the starting position and center the origin arrow.translate(starting_position.astype(float)) - return arrow \ No newline at end of file + return arrow From 28b8592e609bbdffc5e77aee1d39a857d1b80ab4 Mon Sep 17 00:00:00 2001 From: Samuel Tovey Date: Fri, 24 May 2024 22:29:33 +0100 Subject: [PATCH 6/6] Update vector_field.py Add type hints --- znvis/particle/vector_field.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/znvis/particle/vector_field.py b/znvis/particle/vector_field.py index 36e43df..871a3b6 100644 --- a/znvis/particle/vector_field.py +++ b/znvis/particle/vector_field.py @@ -61,7 +61,7 @@ class VectorField: smoothing: bool = False - def _create_mesh(self, position, direction): + def _create_mesh(self, position: np.ndarray, direction: np.ndarray): """ Create a mesh object for the vector field. @@ -108,4 +108,4 @@ def construct_mesh_list(self): mesh = self._create_mesh(self.position[i][j], self.direction[i][j]) else: mesh += self._create_mesh(self.position[i][j], self.direction[i][j]) - self.mesh_list.append(mesh) \ No newline at end of file + self.mesh_list.append(mesh)