diff --git a/CI/integration_tests/test_different_shapes.py b/CI/integration_tests/test_different_shapes.py index 7687c1f..0c91422 100644 --- a/CI/integration_tests/test_different_shapes.py +++ b/CI/integration_tests/test_different_shapes.py @@ -112,6 +112,15 @@ def run_process(): vis.Particle(name="Custom", mesh=mesh, position=trajectory, static=True) ) + trajectory = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + particle_list.append( + vis.Particle( + name="single", + mesh=vis.Sphere(radius=1.0), + position=trajectory, + ) + ) + # Create a bounding box bounding_box = vis.BoundingBox( center=np.array([0, 0, 0]), diff --git a/CI/integration_tests/test_simple_vector_fields.py b/CI/integration_tests/test_simple_vector_fields.py index 7a518e0..7914d4c 100644 --- a/CI/integration_tests/test_simple_vector_fields.py +++ b/CI/integration_tests/test_simple_vector_fields.py @@ -110,6 +110,25 @@ def run_process(): "Dynamic", mesh=dyn_mesh, position=dynamic_pos, direction=dynamic_directors ) + single_position = np.array([[[1.0, 2.0, 3.0]]]) + single_direction = np.array([[[0.0, 1.0, 0.0]]]) + single_dynamic_position = np.repeat(single_position, n_frames, axis=0) + single_dynamic_direction = np.repeat(single_direction, n_frames, axis=0) + single_dynamic_field = vis.VectorField( + "single_dynamic", + mesh=vis.Arrow(scale=1.0, material=Material()), + position=single_dynamic_position, + direction=single_dynamic_direction, + ) + + single_particle_field = vis.VectorField( + "single", + mesh=vis.Arrow(scale=1.0, material=Material()), + position=single_position + 1, + direction=single_direction, + static=True, + ) + # Create a bounding box bounding_box = vis.BoundingBox( center=np.array([0, 0, 0]), @@ -123,6 +142,8 @@ def run_process(): vector_field=[ static_field, dynamic_field, + single_dynamic_field, + single_particle_field, ], frame_rate=20, bounding_box=bounding_box, diff --git a/CI/unit_tests/particle/test_particle.py b/CI/unit_tests/particle/test_particle.py index 742ad8e..b6cf998 100644 --- a/CI/unit_tests/particle/test_particle.py +++ b/CI/unit_tests/particle/test_particle.py @@ -214,6 +214,21 @@ def test_construct_nan_mesh_list(self): str(context.exception), ) + def test_construct_squeezed_dynamic_mesh_list(self): + """ + Test that squeezed dynamic particle data is normalized. + """ + squeezed_particle = Particle( + name="squeezed_particle", + position=np.random.uniform(-5, 5, (10, 3)), + director=np.random.uniform(-5, 5, (10, 3)), + mesh=Sphere(), + ) + squeezed_particle.construct_mesh_list() + self.assertEqual(np.asarray(squeezed_particle.position).shape, (10, 1, 3)) + self.assertEqual(np.asarray(squeezed_particle.director).shape, (10, 1, 3)) + self.assertEqual(len(squeezed_particle.mesh_list), 10) + class TestVariableParticleCount(unittest.TestCase): """ diff --git a/CI/unit_tests/particle/test_vector_fields.py b/CI/unit_tests/particle/test_vector_fields.py index b48c13f..9641c26 100644 --- a/CI/unit_tests/particle/test_vector_fields.py +++ b/CI/unit_tests/particle/test_vector_fields.py @@ -199,9 +199,7 @@ def test_construct_empty_mesh_dict(self): with self.assertRaises(IndexError) as context: self.empty_vector_field.construct_mesh_list() # Check if error message is correct - self.assertEqual( - str(context.exception), "The provided data has an incompatible shape." - ) + self.assertIn("incompatible shape", str(context.exception)) def test_construct_nan_mesh_dict(self): """ @@ -248,6 +246,21 @@ def test_construct_none_dir_mesh_dict(self): # Check if error message is correct self.assertEqual(str(context.exception), "Director data cannot be None.") + def test_construct_squeezed_dynamic_mesh_dict(self): + """ + Test that squeezed dynamic vector field data is normalized. + """ + squeezed_field = VectorField( + name="squeezed_vector_field", + position=np.random.uniform(-5, 5, (10, 3)), + direction=np.random.uniform(-5, 5, (10, 3)), + mesh=Arrow(scale=10, material=self.material), + ) + squeezed_field.construct_mesh_list() + self.assertEqual(squeezed_field.position.shape, (10, 1, 3)) + self.assertEqual(squeezed_field.direction.shape, (10, 1, 3)) + self.assertEqual(len(squeezed_field.mesh_list), 10) + class TestVariableVectorFieldCount(unittest.TestCase): """ diff --git a/znvis/mesh/arrow.py b/znvis/mesh/arrow.py index 87b7f8c..4cdb6d1 100644 --- a/znvis/mesh/arrow.py +++ b/znvis/mesh/arrow.py @@ -60,7 +60,13 @@ def instantiate_mesh( mesh.rotate(matrix, center=(0, 0, 0)) # Translate the arrow to the starting position and center the origin - mesh.translate(starting_position.astype(float)) + starting_position = np.asarray(starting_position, dtype=float).reshape(-1) + if starting_position.size != 3: + raise ValueError( + "starting_position must describe exactly one 3D point, got " + f"shape {starting_position.shape}." + ) + mesh.translate(starting_position) return mesh diff --git a/znvis/mesh/custom.py b/znvis/mesh/custom.py index 7d847b3..cfd556d 100644 --- a/znvis/mesh/custom.py +++ b/znvis/mesh/custom.py @@ -57,7 +57,13 @@ def instantiate_mesh( mesh = o3d.io.read_triangle_mesh(self.file) mesh.compute_vertex_normals() mesh.scale(self.scale, center=mesh.get_center()) - mesh.translate(starting_position.astype(float)) + starting_position = np.asarray(starting_position, dtype=float).reshape(-1) + if starting_position.size != 3: + raise ValueError( + "starting_position must describe exactly one 3D point, got " + f"shape {starting_position.shape}." + ) + mesh.translate(starting_position) if starting_orientation is not None: matrix = rotation_matrix(self.base_direction, starting_orientation) mesh.rotate(matrix) diff --git a/znvis/mesh/mesh.py b/znvis/mesh/mesh.py index 0d76e80..bb88bb2 100644 --- a/znvis/mesh/mesh.py +++ b/znvis/mesh/mesh.py @@ -100,7 +100,13 @@ def instantiate_mesh( """ mesh = self.create_mesh() mesh.compute_vertex_normals() - mesh.translate(starting_position.astype(float)) + starting_position = np.asarray(starting_position, dtype=float).reshape(-1) + if starting_position.size != 3: + raise ValueError( + "starting_position must describe exactly one 3D point, got " + f"shape {starting_position.shape}." + ) + mesh.translate(starting_position) if starting_orientation is not None: matrix = rotation_matrix(self.base_direction, starting_orientation) mesh.rotate(matrix) diff --git a/znvis/particle/particle.py b/znvis/particle/particle.py index d00dd54..7690c01 100644 --- a/znvis/particle/particle.py +++ b/znvis/particle/particle.py @@ -134,8 +134,20 @@ def construct_mesh_list(self): "-------\nWARNING: The provided director array is empty." "Setting to None.\n-------", ) + if not self.static and self.position.ndim not in (2, 3): + raise IndexError("The provided data has an incompatible shape.") + if ( + not self.static + and self.director is not None + and self.director.ndim not in (2, 3) + ): + raise IndexError("The provided data has an incompatible shape.") # Static case differentiation if not self.static: + if self.position.ndim == 2: + self.position = self.position[:, np.newaxis, :] + if self.director is not None and self.director.ndim == 2: + self.director = self.director[:, np.newaxis, :] n_time_steps = self.position.shape[0] self.position = [self.position[i] for i in range(n_time_steps)] if self.director is not None: @@ -196,9 +208,11 @@ def construct_mesh_list(self): for frame_index in track( range(n_time_steps), description=f"Building {self.name} Mesh" ): - frame_pos = self.position[frame_index] + frame_pos = np.atleast_2d(self.position[frame_index]) frame_dir = ( - self.director[frame_index] if self.director is not None else None + np.atleast_2d(self.director[frame_index]) + if self.director is not None + else None ) n_particles = frame_pos.shape[0] meshes = [] @@ -249,6 +263,10 @@ def get_mesh_for_frame(self, frame_index: int): if frame_dir is not None: frame_dir = frame_dir[idx] + frame_pos = np.atleast_2d(frame_pos) + if frame_dir is not None: + frame_dir = np.atleast_2d(frame_dir) + n_particles = frame_pos.shape[0] meshes = [] time_index = 0 if self.static else frame_index diff --git a/znvis/particle/vector_field.py b/znvis/particle/vector_field.py index 0e38c7d..7a14f74 100644 --- a/znvis/particle/vector_field.py +++ b/znvis/particle/vector_field.py @@ -122,6 +122,15 @@ def construct_mesh_list(self): if variable_particle_count: n_time_steps = len(self.position) elif not self.static: + if self.position.ndim not in (2, 3) or self.direction.ndim not in ( + 2, + 3, + ): + raise IndexError("The provided data has an incompatible shape.") + if self.position.ndim == 2: + self.position = self.position[:, np.newaxis, :] + if self.direction.ndim == 2: + self.direction = self.direction[:, np.newaxis, :] n_particles = int(self.position.shape[1]) n_time_steps = int(self.position.shape[0]) else: