export trajectory

weas_widget/index.js

@@ -68,10 +68,16 @@ function render({ model, el }) {
     });
     // Listen for the custom 'atomsUpdated' event
     viewerElement.addEventListener('atomsUpdated', (event) => {
-        const updatedAtoms = event.detail.to_dict(); // event.detail contains the updated atoms
-        updatedAtoms.uuid = avr.atoms.uuid;
-        model.set("atoms", updatedAtoms);
+        // event detail is a trajectory: a array of atoms data
+        // loop all the atoms and export to a dict
+        const trajectory = [];
+        event.detail.forEach((atomsData) => {
+            trajectory.push(atomsData.to_dict());
+        });
+        trajectory.uuid = avr.uuid;
+        model.set("atoms", trajectory);
         model.save_changes();
+        // console.log("updatedAtoms: ", trajectory);
         console.log("Updated atoms from event.")
     });
     // Listen for the custom 'viewerUpdated' event

weas_widget/utils.py

@@ -50,7 +50,8 @@ def to_ase(cls, weas_atoms):
 
         # if atoms is a list of atoms, convert all atoms to a list of ase atoms
         if isinstance(weas_atoms, list):
-            return [cls.to_ase(atom) for atom in weas_atoms]
+            trajectory = [cls.to_ase(atom) for atom in weas_atoms]
+            return trajectory[0] if len(trajectory) == 1 else trajectory
         symbols = [weas_atoms["species"][s][0] for s in weas_atoms["speciesArray"]]
         positions = weas_atoms["positions"]
         cell = np.array(weas_atoms["cell"]).reshape(3, 3)