Update quantum_espresso_workflow.py

Author: jan-janssen

quantum_espresso_workflow.py

@@ -7,14 +7,16 @@
 from adis_tools.parsers import parse_pw
 import matplotlib.pyplot as plt
 import numpy as np
+from optimade.adapters.structures.ase import from_ase_atoms, get_ase_atoms
+from optimade.models.structures import StructureResourceAttributes, StructureResource
 
 
 def write_input(input_dict, working_directory="."):
     filename = os.path.join(working_directory, "input.pwi")
     os.makedirs(working_directory, exist_ok=True)
     write(
         filename=filename,
-        images=Atoms(**input_dict["structure"]),
+        images=json_to_ase(atoms_json=input_dict["structure"]),
         Crystal=True,
         kpts=input_dict["kpts"],
         input_data={
@@ -31,7 +33,7 @@ def write_input(input_dict, working_directory="."):
 def collect_output(working_directory="."):
     output = parse_pw(os.path.join(working_directory, "pwscf.xml"))
     return {
-        "structure": atoms_to_json_dict(atoms=output["ase_structure"]),
+        "structure": ase_to_json(atoms=output["ase_structure"]),
         "energy": output["energy"],
         "volume": output["ase_structure"].get_volume(),
     }
@@ -53,12 +55,12 @@ def calculate_qe(working_directory, input_dict):
 def generate_structures(structure, strain_lst):
     structure_lst = []
     for strain in strain_lst:
-        structure_strain = Atoms(**structure)
+        structure_strain = json_to_ase(atoms_json=structure)
         structure_strain.set_cell(
             structure_strain.cell * strain ** (1 / 3), scale_atoms=True
         )
         structure_lst.append(structure_strain)
-    return {f"s_{i}": atoms_to_json_dict(atoms=s) for i, s in enumerate(structure_lst)}
+    return {f"s_{i}": ase_to_json(atoms=s) for i, s in enumerate(structure_lst)}
 
 
 def plot_energy_volume_curve(volume_lst, energy_lst):
@@ -74,40 +76,14 @@ def get_bulk_structure(element, a, cubic):
         a=a,
         cubic=cubic,
     )
-    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
+    return ase_to_json(atoms=ase_atoms)
+
+
+def ase_to_json(atoms):
+    struct_opt = from_ase_atoms(atoms=atoms)
+    return json.dumps(struct_opt.model_dump(mode="json"))
+
+
+def json_to_ase(atoms_json):
+    structure_restore = StructureResourceAttributes.model_validate_json(atoms_json))
+    return get_ase_atoms(optimade_structure=StructureResource(id="ase", attributes=structure_restore))