Merge pull request #50 from pyiron/structure_optimization

Implement structure optimization protocol

.ci_support/environment-lammps.yml

@@ -4,3 +4,4 @@ dependencies:
 - lammps =2023.08.02
 - pandas =2.1.3
 - pylammpsmpi =0.2.5
+- jinja2 =3.1.2

atomistics/calculators/ase.py

@@ -7,17 +7,42 @@
 if TYPE_CHECKING:
     from ase import Atoms
     from ase.calculators.calculator import Calculator as ASECalculator
+    from ase.optimize.optimize import Optimizer
     from atomistics.calculators.interface import TaskName
 
 
+def ase_optimize_structure(
+    structure, ase_calculator, ase_optimizer, ase_optimizer_kwargs
+):
+    structure_optimized = structure.copy()
+    structure_optimized.calc = ase_calculator
+    ase_optimizer_obj = ase_optimizer(structure_optimized)
+    ase_optimizer_obj.run(**ase_optimizer_kwargs)
+    return structure_optimized
+
+
 @as_task_dict_evaluator
 def evaluate_with_ase(
-    structure: Atoms, tasks: list[TaskName], ase_calculator: ASECalculator
+    structure: Atoms,
+    tasks: list[TaskName],
+    ase_calculator: ASECalculator,
+    ase_optimizer: Optimizer = None,
+    ase_optimizer_kwargs: dict = None,
 ):
-    structure.calc = ase_calculator
     results = {}
-    if "calc_energy" in tasks:
-        results["energy"] = structure.get_potential_energy()
-    if "calc_forces" in tasks:
-        results["forces"] = structure.get_forces()
+    if "optimize_positions" in tasks:
+        results["structure_with_optimized_positions"] = ase_optimize_structure(
+            structure=structure,
+            ase_calculator=ase_calculator,
+            ase_optimizer=ase_optimizer,
+            ase_optimizer_kwargs=ase_optimizer_kwargs,
+        )
+    elif "calc_energy" in tasks or "calc_forces" in tasks:
+        structure.calc = ase_calculator
+        if "calc_energy" in tasks:
+            results["energy"] = structure.get_potential_energy()
+        if "calc_forces" in tasks:
+            results["forces"] = structure.get_forces()
+    else:
+        raise ValueError("The ASE calculator does not implement:", tasks)
     return results

atomistics/calculators/interface.py

@@ -17,11 +17,15 @@ def __str__(self):
 class TaskEnum(StrEnum):
     calc_energy = "calc_energy"
     calc_forces = "calc_forces"
+    optimize_positions = "optimize_positions"
+    optimize_positions_and_volume = "optimize_positions_and_volume"
 
 
 class TaskOutputEnum(Enum):
     energy = "calc_energy"
     forces = "calc_forces"
+    structure_with_optimized_positions = "optimize_positions"
+    structure_with_optimized_positions_and_volume = "optimize_positions_and_volume"
 
 
 if TYPE_CHECKING:

atomistics/calculators/lammps.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING
 
+from jinja2 import Template
 from pylammpsmpi import LammpsASELibrary
 
 from atomistics.calculators.lammps_potentials import (
@@ -17,31 +18,92 @@
     from atomistics.calculators.interface import TaskName
 
 
-LAMMPS_INPUT_TEMPLATE = """\
+LAMMPS_STATIC_RUN_INPUT_TEMPLATE = """\
 thermo_style custom step temp pe etotal pxx pxy pxz pyy pyz pzz vol
 thermo_modify format float %20.15g
 thermo 100
 run 0"""
 
 
+LAMMPS_MINIMIZE_POSITIONS_INPUT_TEMPLATE = """\
+thermo_style custom step temp pe etotal pxx pxy pxz pyy pyz pzz vol
+thermo_modify format float %20.15g
+thermo 10
+min_style {{min_style}}
+minimize {{etol}} {{ftol}} {{maxiter}} {{maxeval}}"""
+
+
+LAMMPS_MINIMIZE_POSITIONS_AND_VOLUME_INPUT_TEMPLATE = """\
+fix ensemble all box/relax iso 0.0
+thermo_style custom step temp pe etotal pxx pxy pxz pyy pyz pzz vol
+thermo_modify format float %20.15g
+thermo 10
+min_style {{min_style}}
+minimize {{etol}} {{ftol}} {{maxiter}} {{maxeval}}"""
+
+
+def template_render(
+    template_str,
+    min_style="cg",
+    etol=0.0,
+    ftol=0.0001,
+    maxiter=100000,
+    maxeval=10000000,
+):
+    return Template(template_str).render(
+        min_style=min_style, etol=etol, ftol=ftol, maxiter=maxiter, maxeval=maxeval
+    )
+
+
 @as_task_dict_evaluator
 def evaluate_with_lammps_library(
     structure: Atoms,
     tasks: list[TaskName],
     potential_dataframe: DataFrame,
     lmp: LammpsASELibrary,
+    lmp_optimizer_kwargs: dict,
 ):
-    lmp = _run_simulation(
-        structure=structure,
-        potential_dataframe=potential_dataframe,
-        input_template=LAMMPS_INPUT_TEMPLATE,
-        lmp=lmp,
-    )
     results = {}
-    if "calc_energy" in tasks:
-        results["energy"] = lmp.interactive_energy_pot_getter()
-    if "calc_forces" in tasks:
-        results["forces"] = lmp.interactive_forces_getter()
+    if "optimize_positions_and_volume" in tasks:
+        lmp = _run_simulation(
+            structure=structure,
+            potential_dataframe=potential_dataframe,
+            input_template=template_render(
+                template_str=LAMMPS_MINIMIZE_POSITIONS_AND_VOLUME_INPUT_TEMPLATE,
+                **lmp_optimizer_kwargs,
+            ),
+            lmp=lmp,
+        )
+        structure_copy = structure.copy()
+        structure_copy.set_cell(lmp.interactive_cells_getter(), scale_atoms=True)
+        structure_copy.positions = lmp.interactive_positions_getter()
+        results["structure_with_optimized_positions_and_volume"] = structure_copy
+    elif "optimize_positions" in tasks:
+        lmp = _run_simulation(
+            structure=structure,
+            potential_dataframe=potential_dataframe,
+            input_template=template_render(
+                template_str=LAMMPS_MINIMIZE_POSITIONS_INPUT_TEMPLATE,
+                **lmp_optimizer_kwargs,
+            ),
+            lmp=lmp,
+        )
+        structure_copy = structure.copy()
+        structure_copy.positions = lmp.interactive_positions_getter()
+        results["structure_with_optimized_positions"] = structure_copy
+    elif "calc_energy" in tasks or "calc_forces" in tasks:
+        lmp = _run_simulation(
+            structure=structure,
+            potential_dataframe=potential_dataframe,
+            input_template=LAMMPS_STATIC_RUN_INPUT_TEMPLATE,
+            lmp=lmp,
+        )
+        if "calc_energy" in tasks:
+            results["energy"] = lmp.interactive_energy_pot_getter()
+        if "calc_forces" in tasks:
+            results["forces"] = lmp.interactive_forces_getter()
+    else:
+        raise ValueError("The LAMMPS calculator does not implement:", tasks)
     lmp.interactive_lib_command("clear")
     return results
 
@@ -56,6 +118,7 @@ def evaluate_with_lammps(
     log_file=None,
     library=None,
     diable_log_file=True,
+    lmp_optimizer_kwargs={},
 ):
     lmp = LammpsASELibrary(
         working_directory=working_directory,
@@ -67,7 +130,10 @@ def evaluate_with_lammps(
         diable_log_file=diable_log_file,
     )
     results_dict = evaluate_with_lammps_library(
-        task_dict=task_dict, potential_dataframe=potential_dataframe, lmp=lmp
+        task_dict=task_dict,
+        potential_dataframe=potential_dataframe,
+        lmp=lmp,
+        lmp_optimizer_kwargs=lmp_optimizer_kwargs,
     )
     lmp.close()
     return results_dict

atomistics/calculators/wrapper.py

@@ -33,6 +33,8 @@ def _convert_task_dict(old_task_dict: dict[TaskName, dict[str, Atoms]]) -> TaskD
     """
     task_dict = {}
     for method_name, subdict in old_task_dict.items():
+        if not isinstance(subdict, dict):
+            subdict = {"label_hidden": subdict}
         for label, structure in subdict.items():
             try:
                 task_dict[label][1].append(method_name)
@@ -72,10 +74,13 @@ def evaluate_with_calculator(
             output = calculate(structure, tasks, *calculate_args, **calculate_kwargs)
             for task_name in tasks:
                 result_name = TaskOutputEnum(task_name).name
-                try:
-                    results_dict[result_name][label] = output[result_name]
-                except KeyError:
-                    results_dict[result_name] = {label: output[result_name]}
+                if label != "label_hidden":
+                    try:
+                        results_dict[result_name][label] = output[result_name]
+                    except KeyError:
+                        results_dict[result_name] = {label: output[result_name]}
+                else:
+                    results_dict[result_name] = output[result_name]
         return results_dict
 
     return evaluate_with_calculator

atomistics/workflows/structure_optimization/workflow.py

@@ -0,0 +1,6 @@
+def optimize_positions_and_volume(structure):
+    return {"optimize_positions_and_volume": structure}
+
+
+def optimize_positions(structure):
+    return {"optimize_positions": structure}

setup.py

@@ -37,6 +37,7 @@
     extras_require={
         "phonopy": ['phonopy==2.20.0', 'seekpath==2.1.0', 'structuretoolkit==0.0.11'],
         "gpaw": ['gpaw==23.9.1'],
+        "lammps": ['pylammpsmpi==0.2.5', 'jinja2==3.1.2', 'pandas==2.1.3']
     },
     cmdclass=versioneer.get_cmdclass(),
 )

tests/test_elastic_lammps.py

@@ -5,6 +5,7 @@
 import unittest
 
 from atomistics.workflows.elastic.workflow import ElasticMatrixWorkflow
+from atomistics.workflows.structure_optimization.workflow import optimize_positions_and_volume
 
 try:
     from atomistics.calculators.lammps import (
@@ -23,14 +24,19 @@ class TestElastic(unittest.TestCase):
     def test_calc_elastic(self):
         potential = '1999--Mishin-Y--Al--LAMMPS--ipr1'
         resource_path = os.path.join(os.path.dirname(__file__), "static", "lammps")
-        structure = bulk("Al", a=4.05, cubic=True)
+        structure = bulk("Al", cubic=True)
         df_pot = get_potential_dataframe(
             structure=structure,
             resource_path=resource_path
         )
         df_pot_selected = df_pot[df_pot.Name == potential].iloc[0]
+        task_dict = optimize_positions_and_volume(structure=structure)
+        result_dict = evaluate_with_lammps(
+            task_dict=task_dict,
+            potential_dataframe=df_pot_selected,
+        )
         calculator = ElasticMatrixWorkflow(
-            structure=structure,
+            structure=result_dict["structure_with_optimized_positions_and_volume"],
             num_of_point=5,
             eps_range=0.005,
             sqrt_eta=True,
@@ -44,4 +50,4 @@ def test_calc_elastic(self):
         elastic_dict = calculator.analyse_structures(output_dict=result_dict)
         self.assertTrue(np.isclose(elastic_dict["C"][0, 0], 114.10393023))
         self.assertTrue(np.isclose(elastic_dict["C"][0, 1], 60.51098897))
-        self.assertTrue(np.isclose(elastic_dict["C"][3, 3], 51.23931149))
+        self.assertTrue(np.isclose(elastic_dict["C"][3, 3], 51.23853765))

tests/test_evcurve_lammps.py

@@ -5,6 +5,7 @@
 import unittest
 
 from atomistics.workflows.evcurve.workflow import EnergyVolumeCurveWorkflow
+from atomistics.workflows.structure_optimization.workflow import optimize_positions_and_volume
 
 
 try:
@@ -24,14 +25,19 @@ class TestEvCurve(unittest.TestCase):
     def test_calc_evcurve(self):
         potential = '1999--Mishin-Y--Al--LAMMPS--ipr1'
         resource_path = os.path.join(os.path.dirname(__file__), "static", "lammps")
-        structure = bulk("Al", a=4.05, cubic=True)
+        structure = bulk("Al", cubic=True)
         df_pot = get_potential_dataframe(
             structure=structure,
             resource_path=resource_path
         )
         df_pot_selected = df_pot[df_pot.Name == potential].iloc[0]
+        task_dict = optimize_positions_and_volume(structure=structure)
+        result_dict = evaluate_with_lammps(
+            task_dict=task_dict,
+            potential_dataframe=df_pot_selected,
+        )
         calculator = EnergyVolumeCurveWorkflow(
-            structure=structure,
+            structure=result_dict["structure_with_optimized_positions_and_volume"],
             num_points=11,
             fit_type='polynomial',
             fit_order=3,
@@ -47,4 +53,4 @@ def test_calc_evcurve(self):
         fit_dict = calculator.analyse_structures(output_dict=result_dict)
         self.assertTrue(np.isclose(fit_dict['volume_eq'], 66.43019853103964))
         self.assertTrue(np.isclose(fit_dict['bulkmodul_eq'], 77.7250135953191))
-        self.assertTrue(np.isclose(fit_dict['b_prime_eq'], 1.279502459079921))
+        self.assertTrue(np.isclose(fit_dict['b_prime_eq'], 1.2795467367276832))

tests/test_optimize_positions_emt.py

@@ -0,0 +1,30 @@
+import numpy as np
+from ase.build import bulk
+from ase.calculators.emt import EMT
+from ase.optimize import BFGS
+import unittest
+
+from atomistics.calculators.ase import evaluate_with_ase
+from atomistics.workflows.structure_optimization.workflow import optimize_positions
+
+
+class TestOptimizePositionsEMT(unittest.TestCase):
+    def test_optimize_positions(self):
+        structure = bulk("Al", a=4.0, cubic=True)
+        positions_before_displacement = structure.positions.copy()
+        structure.positions[0] += [0.01, 0.01, 0.01]
+        task_dict = optimize_positions(structure=structure)
+        result_dict = evaluate_with_ase(
+            task_dict=task_dict,
+            ase_calculator=EMT(),
+            ase_optimizer=BFGS,
+            ase_optimizer_kwargs={"fmax": 0.000001}
+        )
+        structure_optimized = result_dict["structure_with_optimized_positions"]
+        self.assertTrue(
+            all(np.isclose(
+                positions_before_displacement,
+                structure_optimized.positions-structure_optimized.positions[0],
+            ).flatten())
+        )
+

tests/test_optimize_positions_lammps.py

@@ -0,0 +1,47 @@
+import os
+
+from ase.build import bulk
+import numpy as np
+import unittest
+
+from atomistics.workflows.structure_optimization.workflow import optimize_positions
+
+try:
+    from atomistics.calculators.lammps import (
+        evaluate_with_lammps, get_potential_dataframe
+    )
+
+    skip_lammps_test = False
+except ImportError:
+    skip_lammps_test = True
+
+
+@unittest.skipIf(
+    skip_lammps_test, "LAMMPS is not installed, so the LAMMPS tests are skipped."
+)
+class TestOptimizePositionsLAMMPS(unittest.TestCase):
+    def test_optimize_positions(self):
+        potential = '1999--Mishin-Y--Al--LAMMPS--ipr1'
+        resource_path = os.path.join(os.path.dirname(__file__), "static", "lammps")
+        structure = bulk("Al", cubic=True)
+        positions_before_displacement = structure.positions.copy()
+        structure.positions[0] += [0.01, 0.01, 0.01]
+        df_pot = get_potential_dataframe(
+            structure=structure,
+            resource_path=resource_path
+        )
+        df_pot_selected = df_pot[df_pot.Name == potential].iloc[0]
+        task_dict = optimize_positions(structure=structure)
+        result_dict = evaluate_with_lammps(
+            task_dict=task_dict,
+            potential_dataframe=df_pot_selected,
+            lmp_optimizer_kwargs={"ftol": 0.000001},
+        )
+        structure_optimized = result_dict["structure_with_optimized_positions"]
+        self.assertTrue(
+            all(np.isclose(
+                positions_before_displacement,
+                structure_optimized.positions-structure_optimized.positions[0],
+            ).flatten())
+        )
+