add momenta to Frame (#295)

* add momenta * add forces - what is going on here? * Bugfix for additional info from ase --------- Co-authored-by: Paul Hohenberger <[email protected]>
zincware · Dec 21, 2023 · 4a2bb33 · 4a2bb33
1 parent 900fed1
commit 4a2bb33
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/zndraw/frame.py b/zndraw/frame.py
@@ -56,6 +56,8 @@ class Frame:
     numbers: np.ndarray = None
     colors: np.ndarray = None
     radii: np.ndarray = None
+    momenta: np.ndarray = None
+    forces: np.ndarray = None
     pbc: bool = False
     connectivity: nx.Graph() = nx.empty_graph()
     calc: dict = None
@@ -80,7 +82,9 @@ def from_atoms(cls, atoms: ase.Atoms):
         """
         Creates an instance of the frame class from an ase.Atoms object
         """
-        frame = cls(**atoms.arrays)
+        data = atoms.arrays
+
+        frame = cls(positions=data["positions"], numbers=data["numbers"])
 
         frame.cell = np.array(atoms.cell)
         frame.pbc = atoms.pbc