Changes to quantum_espresso_workflow.py

- Add `scale_structure` task
- Prepend return dict keys of `generate_structures` with `s_`
- Add function to serialize return atoms dict to pure python data types
- Rename `name` -> `element` for `get_bulk_structure`, as `name` clashes
  with WG

Author: GeigerJ2

quantum_espresso_workflow.py

@@ -1,5 +1,6 @@
 import os
 import subprocess
+import numpy as np
 
 from ase.atoms import Atoms
 from ase.build import bulk
@@ -57,7 +58,15 @@ def generate_structures(structure, strain_lst):
             structure_strain.cell * strain ** (1 / 3), scale_atoms=True
         )
         structure_lst.append(structure_strain)
-    return {str(i): s.todict() for i, s in enumerate(structure_lst)}
+    return {f"s_{i}": atoms_to_json_dict(atoms=s) for i, s in enumerate(structure_lst)}
+
+
+def scale_structure(structure, strain):
+    structure_strain = Atoms(**structure)
+    structure_strain.set_cell(
+        structure_strain.cell * strain ** (1 / 3), scale_atoms=True
+    )
+    return structure_strain
 
 
 def plot_energy_volume_curve(volume_lst, energy_lst):
@@ -67,9 +76,46 @@ def plot_energy_volume_curve(volume_lst, energy_lst):
     plt.savefig("evcurve.png")
 
 
-def get_bulk_structure(name, a, cubic):
-    return bulk(
-        name=name,
+def get_bulk_structure(element, a, cubic):
+    ase_atoms = bulk(
+        name=element,
         a=a,
         cubic=cubic,
-    ).todict()
+    )
+    return atoms_to_json_dict(atoms=ase_atoms)
+
+
+def atoms_to_json_dict(atoms):
+    """
+    Convert an ASE Atoms object to a fully JSON-serializable dictionary
+    that uses only Python base data types.
+
+    Parameters:
+    -----------
+    atoms : ase.Atoms
+        The Atoms object to convert
+
+    Returns:
+    --------
+    dict
+        A dictionary representation using only Python base types
+    """
+    # Get the dictionary representation from ASE
+    atoms_dict = atoms.todict()
+
+    # Create a new dictionary with JSON-serializable values
+    json_dict = {}
+
+    # Convert numpy arrays to lists
+    for key, value in atoms_dict.items():
+        if isinstance(value, np.ndarray):
+            # Convert numpy boolean values to Python booleans
+            if value.dtype == np.bool_ or value.dtype == bool:
+                json_dict[key] = value.tolist()
+            # Convert numpy arrays of numbers to Python lists
+            else:
+                json_dict[key] = value.tolist()
+        else:
+            json_dict[key] = value
+
+    return json_dict