Added animation option

05-geometrien/support.py

@@ -3,6 +3,8 @@ import plotly.graph_objects as go
 import atom_properties as atomp
 from itertools import accumulate
 import qc_parser as qcp
+import subprocess as sub
+import os
 
 class XYZ:
     sphere_mode = "ball"
@@ -30,6 +32,12 @@ class XYZ:
                          "y": float(line.split()[2]),
                          "z": float(line.split()[3])} for line in lines])
     
+    @classmethod
+    def from_list(cls, lines: list[list[str|float]]):
+        return cls([{"element": line[0],
+                     "x": float(line[1]),
+                     "y": float(line[2]),
+                     "z": float(line[3])} for line in lines])
     
     def get_bond_length(self, i: int, j: int) -> float:
         """
@@ -260,6 +268,84 @@ class Trajectory:
         parsed = qcp.read_qc_file(output_file)
         return cls([XYZ(frame) for frame in parsed["coordinates"]])
 
+    @classmethod
+    def from_trajectory_file(cls, trajectory_file: str):
+        """
+        Creates a Trajectory object from a trajectory file.
+
+        Args:
+            trajectory_file (str): The path to the trajectory file.
+
+        Returns:
+            Trajectory: A Trajectory object.
+        """
+        return cls([frame for frame in Trajectory.xyz_generator(trajectory_file)])
+    
+    @staticmethod
+    def xyz_generator(trajectory_file: str):
+        """
+        Generator that yields XYZ objects from a trajectory file.
+
+        Args:
+            trajectory_file (str): The path to the trajectory file.
+
+        Yields:
+            XYZ: An XYZ object.
+        """
+        with open(trajectory_file) as f:
+            lines = f.readlines()
+        block_size = int(lines[0].strip()) + 2
+        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:
+        """
+        Returns a plotly figure with the molecular structure of the trajectory. The figure comes with an animation.
+
+        Args:
+            resolution (int, optional): The number of points to be generated. Defaults to 64.
+            rel_cutoff (float, optional): Bond are only shown if the interatomic distance is smaller than the sum of the atomic vdw radii multiplied with this parameter.
+                Defaults to 0.5.
+
+        Returns:
+            go.Figure: A plotly figure with the molecular structure.
+        """
+        fig = go.Figure()
+
+        # Get and add all the meshes
+        mesh_elements = [0]
+        frames = []
+        for xyz in self.coordinates:
+            data = []
+            meshes = xyz.get_molecular_mesh(resolution=resolution, rel_cutoff=rel_cutoff)
+            mesh_elements.append(len(meshes))
+            for mesh in meshes:
+                data.append((mesh))
+            frames.append(go.Frame(data=data))
+        
+        fig = go.Figure(
+            data=data, # initial frame before animation
+            frames=frames
+        )
+        fig.update_layout(
+            updatemenus=[{
+                "type": "buttons",
+                "buttons": [{
+                    "label": "Play",
+                    "method": "animate",
+                    "args": [None, {
+                        "frame": {"duration": 200},
+                        "transition": {"easing": "linear"}
+                    }]
+                }],
+            }],
+            scene_aspectmode='data'
+        )
+
+
+
+        return fig
+
     def get_molecular_viewer(self, resolution: int = 64, rel_cutoff: float = 0.5) -> go.Figure:
         """
         Returns a plotly figure with the molecular structure of the trajectory. The figure comes with a slider that allows to switch between frames.
@@ -302,5 +388,6 @@ class Trajectory:
         fig.update_layout(
             sliders=[slider]
         )
-
+        
+        fig.update_layout(scene_aspectmode='data')
         return fig
\ No newline at end of file