Restructering/modularizyation of the node library

Optimized import statements in nodes (pyiron_atomistics -> ase atom object) reduces import time from 5s to < 50 ms!
Implementation of universal force field M3GNet

pyiron_workflow/node_library/calculator/ase.py

@@ -0,0 +1,62 @@
+from pyiron_workflow.function import single_value_node, function_node
+
+
+@single_value_node()
+def static(atoms=None, engine=None, _internal=None):
+    if engine is None:
+        from ase.calculators.emt import EMT
+
+        engine = EMT()
+
+    atoms.calc = engine
+
+    out = {}
+    # out['structure'] = atoms # not needed since identical to input
+    out["forces"] = atoms.get_forces()
+    out["energy"] = atoms.get_potential_energy()
+    if _internal is not None:
+        out["iter_index"] = _internal[
+            "iter_index"
+        ]  # TODO: move _internal argument to decorator class
+    return out
+
+
+@function_node("structure", "out")
+def minimize(atoms=None, engine=None, fmax=0.005, log_file="tmp.log"):
+    from ase.optimize import BFGS
+    import numpy as np
+
+    if engine is None:
+        from ase.calculators.emt import EMT
+
+        engine = EMT()
+
+    atoms.calc = engine
+
+    if log_file is None:  # write to standard io
+        log_file = "-"
+
+    dyn = BFGS(atoms, logfile=log_file)
+    dyn.run(fmax=fmax)
+
+    # it appears that r0 is the structure of the second to last step (check)
+    atoms_relaxed = atoms.copy()
+    atoms_relaxed.calc = atoms.calc
+    if dyn.r0 is not None:
+        atoms_relaxed.positions = dyn.r0.reshape(-1, 3)
+
+    out = {}
+    out["relaxed_structure"] = atoms_relaxed
+    # out["forces"] = dyn.f0.reshape(-1, 3)
+    out["forces"] = atoms_relaxed.get_forces()
+    out["energy"] = atoms_relaxed.get_potential_energy()
+    out["energy_initial"] = atoms.get_potential_energy()
+    print("energy: ", out["energy"], "max_force: ", np.min(np.abs(out["forces"])))
+
+    return atoms_relaxed, out
+
+
+nodes = [
+    static,
+    minimize,
+]

pyiron_workflow/node_library/engine/ase.py

@@ -0,0 +1,22 @@
+from pyiron_workflow.function import single_value_node
+
+
+@single_value_node("engine")
+def EMT():
+    from ase.calculators.emt import EMT
+
+    return EMT()
+
+
+@single_value_node("engine")
+def M3GNet():
+    import matgl
+    from matgl.ext.ase import M3GNetCalculator
+
+    return M3GNetCalculator(matgl.load_model("M3GNet-MP-2021.2.8-PES"))
+
+
+nodes = [
+    EMT,
+    M3GNet,
+]

pyiron_workflow/node_library/phonopy.py

@@ -72,17 +72,18 @@ def generate_supercells(phonopy, parameters: InputPhonopyGenerateSupercells):
 @function_node()
 def create_phonopy(
     structure,
-    calculator=None,
+    engine=None,
     executor=None,
     parameters=InputPhonopyGenerateSupercells(),
 ):
     from phonopy import Phonopy
     from structuretoolkit.common import atoms_to_phonopy
+    from pyiron_workflow.node_library import calculator
 
     phonopy = Phonopy(unitcell=atoms_to_phonopy(structure))
 
     cells = generate_supercells(phonopy, parameters=parameters)  # .run()
-    gs = calc_static()
+    gs = calculator.ase.static(engine=engine)
     df = gs.iter(atoms=cells, executor=executor)
     phonopy.forces = df.forces
 
@@ -126,7 +127,6 @@ def check_consistency(phonopy, tolerance: float = 1e-10):
         has_imaginary_modes = True
     else:
         has_imaginary_modes = False
-        print("alles ok")
     return has_imaginary_modes
 
 
@@ -141,71 +141,11 @@ def get_total_dos(phonopy, mesh=3 * [10]):
     return total_dos
 
 
-@single_value_node()
-def calc_static(atoms=None, engine=None, _internal=None):
-    # move later to other package
-    # print("atoms: ", atoms)
-    if engine is None:
-        from ase.calculators.emt import EMT
-
-        engine = EMT()
-
-    atoms.calc = engine
-
-    out = {}
-    # out['structure'] = atoms # not needed since identical to input
-    out["forces"] = atoms.get_forces()
-    out["energy"] = atoms.get_potential_energy()
-    if _internal is not None:
-        out["iter_index"] = _internal[
-            "iter_index"
-        ]  # TODO: move _internal argument to decorator class
-    return out
-
-
-@function_node("structure", "out")
-def calc_minimize(atoms=None, engine=None, fmax=0.005, log_file="tmp.log"):
-    # move later to other package
-    from ase.optimize import BFGS
-    import numpy as np
-
-    if engine is None:
-        from ase.calculators.emt import EMT
-
-        engine = EMT()
-
-    atoms.calc = engine
-
-    if log_file is None:  # write to standard io
-        log_file = "-"
-
-    dyn = BFGS(atoms, logfile=log_file)
-    dyn.run(fmax=fmax)
-
-    # it appears that the is the structure of the second to last step (check)
-    atoms_relaxed = atoms.copy()
-    atoms_relaxed.calc = atoms.calc
-    if dyn.r0 is not None:
-        atoms_relaxed.positions = dyn.r0.reshape(-1, 3)
-
-    out = {}
-    out["relaxed_structure"] = atoms_relaxed
-    # out["forces"] = dyn.f0.reshape(-1, 3)
-    out["forces"] = atoms_relaxed.get_forces()
-    out["energy"] = atoms_relaxed.get_potential_energy()
-    out["energy_initial"] = atoms.get_potential_energy()
-    print("energy: ", out["energy"], "max_force: ", np.min(np.abs(out["forces"])))
-
-    return atoms_relaxed, out
-
-
 nodes = [
     #    generate_supercells,
     create_phonopy,
     get_dynamical_matrix,
     get_eigenvalues,
     check_consistency,
     get_total_dos,
-    calc_static,
-    calc_minimize,
 ]

pyiron_workflow/node_library/structure_build.py → pyiron_workflow/node_library/structure/build.py

@@ -31,7 +31,7 @@ def bulk(
 def cubic_bulk_cell(
     wf, element: str, cell_size: int = 1, vacancy_index: int | None = None
 ):
-    from pyiron_workflow.node_library.structure_transform import create_vacancy, repeat
+    from pyiron_workflow.node_library.structure.transform import create_vacancy, repeat
 
     wf.structure = bulk(name=element, cubic=True)
     wf.cell = repeat(structure=wf.structure, repeat_scalar=cell_size)

pyiron_workflow/node_library/structure_transform.py → pyiron_workflow/node_library/structure/transform.py

@@ -1,6 +1,9 @@
 from pyiron_workflow.function import single_value_node
 from typing import Optional, Union
-from pyiron_atomistics.atomistics.structure.atoms import Atoms
+
+# Huge savings when replacing pyiron_atomistics atoms class with ase one!! (> 5s vs 40 ms)
+# from pyiron_atomistics.atomistics.structure.atoms import Atoms
+from ase import Atoms
 
 
 @single_value_node("structure")