Working animation with mode vectors

05-geometrien/support.py

@@ -250,12 +250,64 @@ class XYZ:
 
         return fig
 
+def circle_coordinates(point: np.ndarray, v: np.ndarray, radius: float, num_points: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """
+    Generate cartesian coordinates of a circle in 3D space.
+
+    Args:
+        point (np.ndarray): A point on the circle.
+        v (np.ndarray): A vector perpendicular to the circle.
+        radius (float): The radius of the circle.
+        num_points (int): The number of points to generate on the circle.
+
+    Returns:
+        np.ndarray: A 2D array of shape (num_points, 3) containing the cartesian coordinates of the circle.
+    """
+
+    # Normalize the vector v
+    v = v / np.linalg.norm(v)
+
+    # Create a vector u that is perpendicular to v
+    u = np.cross(v, [1, 0, 0]) if np.linalg.norm(np.cross(v, [1, 0, 0])) > 1e-6 else np.cross(v, [0, 1, 0])
+    u = u / np.linalg.norm(u)
+
+    # Create a vector w that is perpendicular to both u and v
+    w = np.cross(u, v)
+
+    # Generate the cartesian coordinates of the circle
+    theta = np.linspace(0, 2 * np.pi, num_points)
+    circle_coords = point + radius * (u * np.cos(theta)[:, np.newaxis] + w * np.sin(theta)[:, np.newaxis])
+
+    x, y, z = circle_coords[:, 0], circle_coords[:, 1], circle_coords[:, 2]
+    return x, y, z
+
+def make_arrow_mesh(at: np.ndarray, dir: np.ndarray, resolution: int = 16, radius: float = 0.05) -> go.Mesh3d:
+    tip = at + dir
+    bottom = at - dir
+    v = bottom - tip
+    x1, y1, z1 = circle_coordinates(bottom, v, radius=radius, num_points=resolution//4)
+    x2, y2, z2 = circle_coordinates(tip, v, radius=radius, num_points=resolution//4)
+    x, y, z = np.concatenate((x1, x2)), np.concatenate((y1, y2)), np.concatenate((z1, z2))
+
+    arrow_mesh = go.Mesh3d(
+            x=x, 
+            y=y, 
+            z=z, 
+            color="#444444", 
+            opacity=1, 
+            alphahull=0,
+            # name=f'{self.elements[i]}-{self.elements[j]}',
+            hoverinfo='none',  # No hover info at all
+            )
+    return arrow_mesh
+
 class Trajectory:
     def __init__(self, coordinates: list[XYZ]):
         self.coordinates = coordinates
+        self.vibration_vectors = None
     
     @classmethod
-    def from_orca_output(cls, output_file: str):
+    def from_opt_output(cls, output_file: str):
         """
         Creates a Trajectory object from an ORCA output file.
 
@@ -281,6 +333,25 @@ class Trajectory:
         """
         return cls([frame for frame in Trajectory.xyz_generator(trajectory_file)])
     
+    @classmethod
+    def from_vibration_output(cls, output_file: str, mode: int):
+        """
+        Creates a Trajectory object from an ORCA output file
+        containing a frequency calculation at the given mode.
+
+        Args:
+            output_file (str): The path to the ORCA output file.
+
+        Returns:
+            Trajectory: A Trajectory object.
+        """
+        proc = sub.run(f"orca_pltvib {output_file} {mode}", shell=True, text=True, capture_output=True)
+        trajectory_file = f"{output_file}.v{mode:03d}.xyz"
+        xyz = Trajectory.from_trajectory_file(trajectory_file)
+        xyz.vibration_vectors = Trajectory.get_vibration_vectors(trajectory_file)
+        os.remove(trajectory_file)
+        return xyz
+
     @staticmethod
     def xyz_generator(trajectory_file: str):
         """
@@ -298,7 +369,38 @@ class Trajectory:
         for i in range(0, len(lines), block_size):
             yield XYZ.from_list([line.split() for line in lines[i+2:i+block_size]])
 
-    def get_molecular_viewer_animated(self, resolution: int = 64, rel_cutoff: float = 0.5) -> go.Figure:
+    @staticmethod
+    def get_vibration_vectors(trajectory_file: str):
+        """
+        Returns the vibration vectors from a trajectory file. Must be generated
+        by orca_pltvib, so it contains the vectors in the last three columns.
+
+        Args:
+            trajectory_file (str): The path to the trajectory file.
+
+        Returns:
+            list[list[float]]: A list of vibration vectors.
+        """
+        with open(trajectory_file) as f:
+            lines = f.readlines()
+        block_size = int(lines[0].strip())
+        lines = lines[2:]
+        vibration_vectors = [[float(element) for element in line.split()[-3:]] for line in lines[:block_size]]
+        print(vibration_vectors)
+        return vibration_vectors
+    
+    def get_vibration_vectors_mesh(self) -> list[go.Mesh3d]:
+        """
+        """
+        assert self.vibration_vectors is not None, "Vibration vectors must be generated first"
+        meshes = []
+        first_frame_coords = self.coordinates[0].xyz
+        for i,vector in enumerate(self.vibration_vectors):
+            meshes.append(make_arrow_mesh(at=first_frame_coords[i], dir=vector))
+        return meshes
+
+    
+    def get_molecular_viewer_animated(self, resolution: int = 64, rel_cutoff: float = 0.5, arrows: bool = False) -> go.Figure:
         """
         Returns a plotly figure with the molecular structure of the trajectory. The figure comes with an animation.
 
@@ -320,7 +422,10 @@ class Trajectory:
             meshes = xyz.get_molecular_mesh(resolution=resolution, rel_cutoff=rel_cutoff)
             mesh_elements.append(len(meshes))
             for mesh in meshes:
-                data.append((mesh))
+                data.append(mesh) # was in double parenthesis
+            if arrows and self.vibration_vectors is not None:
+                data += self.get_vibration_vectors_mesh()
+
             frames.append(go.Frame(data=data))
         
         fig = go.Figure(
@@ -334,16 +439,13 @@ class Trajectory:
                     "label": "Play",
                     "method": "animate",
                     "args": [None, {
-                        "frame": {"duration": 200},
-                        "transition": {"easing": "linear"}
+                        "frame": {"duration": 200}
                     }]
                 }],
             }],
             scene_aspectmode='data'
         )
 
-
-
         return fig
 
     def get_molecular_viewer(self, resolution: int = 64, rel_cutoff: float = 0.5) -> go.Figure: