Support MDAnalysis.Universe as the frame parameter

python/chemiscope/jupyter.py

@@ -246,7 +246,13 @@ def show(
             # also adds an index property to have something to show in the info panel
             properties["index"] = {
                 "target": "structure",
-                "values": list(range(len(frames))),
+                "values": (
+                    list(range(len(frames)))
+                    if hasattr(frames, "__len__")
+                    # MDAnalysis `frames` will not be a list, and does not have
+                    # a `__len__`
+                    else list(range(len(frames.trajectory)))
+                ),
             }
 
         widget_class = StructureWidget

python/chemiscope/structures/__init__.py

@@ -18,6 +18,12 @@
     ase_tensors_to_ellipsoids,
     ase_vectors_to_arrows,
 )
+from ._mda import (  # noqa: F401
+    _mda_to_json,
+    _mda_valid_structures,
+    _mda_list_atom_properties,
+    _mda_list_structure_properties,
+)
 from ._stk import (  # noqa: F401
     _stk_valid_structures,
     _stk_to_json,
@@ -42,6 +48,10 @@ def _guess_adapter(frames):
     if use_stk:
         return stk_frames, "stk"
 
+    mda_frames, use_mda = _mda_valid_structures(frames)
+    if use_mda:
+        return mda_frames, "mda"
+
     raise Exception(f"unknown frame type: '{frames[0].__class__.__name__}'")
 
 
@@ -60,6 +70,10 @@ def frames_to_json(frames):
         return [_ase_to_json(frame) for frame in frames]
     elif adapter == "stk":
         return [_stk_to_json(frame) for frame in frames]
+    elif adapter == "mda":
+        # Be careful of the lazy loading of `frames.atoms`, which is updated during the
+        # iteration of the trajectory
+        return [_mda_to_json(frames.atoms) for _ in frames.trajectory]
     else:
         raise Exception("reached unreachable code")
 
@@ -76,7 +90,8 @@ def _list_atom_properties(frames):
         return _ase_list_atom_properties(frames)
     elif adapter == "stk":
         return _stk_list_atom_properties(frames)
-
+    elif adapter == "mda":
+        return _mda_list_atom_properties(frames)
     else:
         raise Exception("reached unreachable code")
 
@@ -93,6 +108,8 @@ def _list_structure_properties(frames):
         return _ase_list_structure_properties(frames)
     elif adapter == "stk":
         return _stk_list_structure_properties(frames)
+    elif adapter == "mda":
+        return _mda_list_atom_properties(frames)
     else:
         raise Exception("reached unreachable code")
 
@@ -119,6 +136,12 @@ def extract_properties(frames, only=None, environments=None):
             "stk molecules do not contain properties, you must manually provide them"
         )
 
+    elif adapter == "mda":
+        raise RuntimeError(
+            "MDAnalysis molecules do not contain properties, you"
+            "must manually provide them"
+        )
+
     else:
         raise Exception("reached unreachable code")
 

python/chemiscope/structures/_mda.py

@@ -0,0 +1,44 @@
+import numpy as np
+
+try:
+    import MDAnalysis as mda
+
+    HAVE_MDA = True
+except ImportError:
+    HAVE_MDA = False
+
+
+def _mda_valid_structures(frames):
+    if HAVE_MDA and isinstance(frames, mda.Universe):
+        return frames, True
+    else:
+        return [], False
+
+
+def _mda_to_json(ag):
+    data = {}
+    data["size"] = len(ag)
+    data["names"] = [atom.type for atom in ag]
+    data["x"] = [float(value) for value in ag.positions[:, 0]]
+    data["y"] = [float(value) for value in ag.positions[:, 1]]
+    data["z"] = [float(value) for value in ag.positions[:, 2]]
+    if ag.dimensions is not None:
+        data["cell"] = list(
+            np.concatenate(
+                mda.lib.mdamath.triclinic_vectors(ag.dimensions),
+                dtype=np.float64,
+                # should be np.float64 otherwise not serializable
+            )
+        )
+
+    return data
+
+
+def _mda_list_atom_properties(frames) -> list:
+    # mda cannot have atom properties or structure properties, so skipping.
+    return []
+
+
+def _mda_list_structure_properties(frames) -> list:
+    # mda cannot have atom properties or structure properties, so skipping.
+    return []

python/chemiscope/structures/_stk.py

@@ -18,7 +18,7 @@ def _stk_valid_structures(
     if HAVE_STK and isinstance(frames, Molecule):
         # deal with the user passing a single frame
         return [frames], True
-    elif HAVE_STK and isinstance(frames[0], Molecule):
+    elif HAVE_STK and isinstance(frames, list) and isinstance(frames[0], Molecule):
         for frame in frames:
             assert isinstance(frame, Molecule)
         return frames, True

python/tests/mda_structures.py

@@ -0,0 +1,60 @@
+import unittest
+
+import MDAnalysis as mda
+import numpy as np
+
+import chemiscope
+
+BASE_FRAME = mda.Universe.empty(n_atoms=3, trajectory=True)
+BASE_FRAME.add_TopologyAttr("type", ["C", "O", "O"])
+BASE_FRAME.atoms.positions = np.array([[0, 0, 0], [1, 1, 1], [2, 2, 5]])
+
+
+class TestStructures(unittest.TestCase):
+    """Conversion of structure data to chemiscope JSON"""
+
+    def test_structures(self):
+        data = chemiscope.create_input(BASE_FRAME)
+        self.assertEqual(len(data["structures"]), 1)
+        self.assertEqual(data["structures"][0]["size"], 3)
+        self.assertEqual(data["structures"][0]["names"], ["C", "O", "O"])
+        self.assertEqual(data["structures"][0]["x"], [0, 1, 2])
+        self.assertEqual(data["structures"][0]["y"], [0, 1, 2])
+        self.assertEqual(data["structures"][0]["z"], [0, 1, 5])
+        self.assertEqual(data["structures"][0].get("cell"), None)
+
+        frame = BASE_FRAME.copy()
+        frame.dimensions = [23, 22, 11, 90, 90, 90]
+        data = chemiscope.create_input(frame)
+        self.assertEqual(len(data["structures"]), 1)
+        self.assertEqual(data["structures"][0]["cell"], [23, 0, 0, 0, 22, 0, 0, 0, 11])
+
+        frame = BASE_FRAME.copy()
+        frame.dimensions = [23, 22, 11, 120, 90, 70]
+        data = chemiscope.create_input(frame)
+        self.assertEqual(len(data["structures"]), 1)
+
+        cell = [
+            23.0,
+            0.0,
+            0.0,
+            7.5244431531647145,
+            20.673237657289985,
+            0.0,
+            0.0,
+            -5.852977748617515,
+            9.313573507209156,
+        ]
+        self.assertTrue(np.allclose(data["structures"][0]["cell"], cell))
+
+
+class TestExtractProperties(unittest.TestCase):
+    """Properties extraction"""
+
+    def test_exception(self):
+        with self.assertRaises(RuntimeError):
+            chemiscope.extract_properties(BASE_FRAME)
+
+
+if __name__ == "__main__":
+    unittest.main()

tox.ini

@@ -29,6 +29,7 @@ deps =
     ase==3.22.1
     rdkit==2024.3.4
     stk
+    MDAnalysis
 
 commands =
     pip install chemiscope[explore]