Fix squeezed single-particle trajectories in ZnVis

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]),

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,

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):
     """

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):
     """

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
 

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)

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)

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

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: