Trajectory viewer (#18)

Author: BTSpeake

src/alc_aiidalab_widgets/viewers.py

@@ -14,6 +14,7 @@
     # AiiDA data type labels
     "data.core.structure.StructureData.": StructureViewWidget,
     "data.core.singlefile.SinglefileData.": SinglefileDataViewer,
+    "data.core.array.trajectory.TrajectoryData.": StructureViewWidget,
     # Custom redirections
     "xyz_table": XYZArrayDataTableWidget,
 }

src/alc_aiidalab_widgets/widgets/file_handling.py

@@ -65,7 +65,6 @@ def _on_file_upload(self, _):
                 self.file = self.get_aiida_file_object()
         else:
             self.file_handle.value = ""
-        print(self.file_dict)
         return
 
     def get_file_contents(self) -> BytesIO | None:

src/alc_aiidalab_widgets/widgets/structure.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 
-from aiida.orm import Node, SinglefileData, StructureData
+from aiida.orm import Node, SinglefileData, StructureData, TrajectoryData
 from ase import Atoms
 from ase import io as ase_io
 from ipywidgets import HTML, VBox
@@ -25,6 +25,8 @@ def __init__(self, node: Node | None = None, **kwargs):
             self.assign_structure_from_structuredata(node)
         elif isinstance(node, SinglefileData):
             self.assign_structure_from_file(node.filename, node.content)
+        elif isinstance(node, TrajectoryData):
+            self.assign_structure_from_trajectorydata(node)
         elif node:
             self.message.value = (
                 "<p>AiiDA Node type not supported by the structure viewer."
@@ -61,13 +63,13 @@ def assign_structure_from_file(self, fname: str, content: bytes) -> None:
                 ]
         return
 
-    def assign_structure_from_ase(self, structure: Atoms) -> None:
+    def assign_structure_from_ase(self, structure: Atoms | list[Atoms]) -> None:
         """Visualise the given ASE structure.
 
         Parameters
         ----------
-        structure: Atoms
-            The ASE atoms structure object.
+        structure: Atoms | list[Atoms]
+            The ASE atoms structure(s) object.
         """
         self.viewer = WeasWidget()
         self.viewer.from_ase(structure)
@@ -86,3 +88,22 @@ def assign_structure_from_structuredata(self, structure: StructureData) -> None:
         """
         self.assign_structure_from_ase(structure._get_object_ase())
         return
+
+    def assign_structure_from_trajectorydata(self, trajectory: TrajectoryData) -> None:
+        """
+        Visualise a series of structures contained in an AiiDA TrajectoryData node.
+
+        Parameters
+        ----------
+        trajectory: TrajectoryData
+            The AiiDA TrajectoryData node containing the structure series to visualise.
+        """
+        symbols = trajectory.symbols
+        positions = trajectory.get_positions()
+        nsteps = trajectory.numsteps
+        atoms = []
+        for i in range(nsteps):
+            step = Atoms(symbols=symbols, positions=positions[i])
+            atoms.append(step)
+        self.assign_structure_from_ase(atoms)
+        return

tests/test_structure.py

@@ -3,7 +3,7 @@
 from pathlib import Path
 from tempfile import NamedTemporaryFile
 
-from aiida.orm import SinglefileData, StructureData
+from aiida.orm import SinglefileData, StructureData, TrajectoryData
 from ipywidgets import HTML
 from weas_widget import WeasWidget
 
@@ -69,6 +69,26 @@ def test_structure_initialise_from_singlefiledata():
     assert structure_ase.get_atomic_numbers()[2] == 1
 
 
+def test_structure_initialise_from_trajectorydata():
+    """Create a WeasWidget viewer initialised with an AiiDA TrajectoryData node."""
+    structure_1 = StructureData()
+    with open(Path(__file__).resolve().parent / "data/water.xyz") as f:
+        structure_1._parse_xyz(f.read())
+    structure_2 = StructureData()
+    with open(Path(__file__).resolve().parent / "data/water.xyz") as f:
+        structure_2._parse_xyz(f.read())
+    trajectory = TrajectoryData()
+    trajectory.set_structurelist([structure_1, structure_2])
+    widget = StructureViewWidget(trajectory)
+    assert isinstance(widget.viewer, WeasWidget)
+    assert len(widget.children) == 1
+    structure_ase = widget.viewer.to_ase()
+    assert len(structure_ase) == 2
+    assert structure_ase[0].get_atomic_numbers()[0] == 8
+    assert structure_ase[0].get_atomic_numbers()[1] == 1
+    assert structure_ase[0].get_atomic_numbers()[2] == 1
+
+
 def test_invalid_file_type():
     """Load an invalid filetype into a StructureViewWidget."""
     widget = StructureViewWidget()