Add support for constraints

pyiron_lammps/compatibility/constraints.py

@@ -0,0 +1,106 @@
+import numpy as np
+from ase.atoms import Atoms
+
+
+def _get_fixed_atom_boolean_vector(structure: Atoms):
+    fixed_atom_vector = np.array([[False, False, False]] * len(structure))
+    for c in structure.constraints:
+        c_dict = c.todict()
+        if c_dict["name"] == "FixAtoms":
+            fixed_atom_vector[c_dict["kwargs"]["indices"]] = [True, True, True]
+        elif c_dict["name"] == "FixedPlane":
+            if all(np.isin(c_dict["kwargs"]["direction"], [0, 1, 1 / np.sqrt(2)])):
+                if "indices" in c_dict["kwargs"]:
+                    fixed_atom_vector[c_dict["kwargs"]["indices"]] = np.array(
+                        c_dict["kwargs"]["direction"]
+                    ).astype(bool)
+                elif "a" in c_dict["kwargs"]:
+                    fixed_atom_vector[c_dict["kwargs"]["a"]] = np.array(
+                        c_dict["kwargs"]["direction"]
+                    ).astype(bool)
+            else:
+                raise ValueError(
+                    "Currently the directions are limited to [1, 0, 0], [1, 1, 0], [1, 1, 1] and its permutations."
+                )
+        else:
+            raise ValueError("Only FixAtoms and FixedPlane are currently supported. ")
+    return fixed_atom_vector
+
+
+def set_selective_dynamics(structure: Atoms, calc_md: bool = False):
+    control_dict = {}
+    if len(structure.constraints) > 0:
+        sel_dyn = _get_fixed_atom_boolean_vector(structure=structure)
+        # Enter loop only if constraints present
+        if len(np.argwhere(np.any(sel_dyn, axis=1)).flatten()) != 0:
+            all_indices = np.arange(len(structure), dtype=int)
+            constraint_xyz = np.argwhere(np.all(sel_dyn, axis=1)).flatten()
+            not_constrained_xyz = np.setdiff1d(all_indices, constraint_xyz)
+            # LAMMPS starts counting from 1
+            constraint_xyz += 1
+            ind_x = np.argwhere(sel_dyn[not_constrained_xyz, 0]).flatten()
+            ind_y = np.argwhere(sel_dyn[not_constrained_xyz, 1]).flatten()
+            ind_z = np.argwhere(sel_dyn[not_constrained_xyz, 2]).flatten()
+            constraint_xy = not_constrained_xyz[np.intersect1d(ind_x, ind_y)] + 1
+            constraint_yz = not_constrained_xyz[np.intersect1d(ind_y, ind_z)] + 1
+            constraint_zx = not_constrained_xyz[np.intersect1d(ind_z, ind_x)] + 1
+            constraint_x = (
+                not_constrained_xyz[np.setdiff1d(np.setdiff1d(ind_x, ind_y), ind_z)] + 1
+            )
+            constraint_y = (
+                not_constrained_xyz[np.setdiff1d(np.setdiff1d(ind_y, ind_z), ind_x)] + 1
+            )
+            constraint_z = (
+                not_constrained_xyz[np.setdiff1d(np.setdiff1d(ind_z, ind_x), ind_y)] + 1
+            )
+            control_dict = {}
+            if len(constraint_xyz) > 0:
+                control_dict["group constraintxyz"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_xyz]
+                )
+                control_dict["fix constraintxyz"] = "constraintxyz setforce 0.0 0.0 0.0"
+                if calc_md:
+                    control_dict["velocity constraintxyz"] = "set 0.0 0.0 0.0"
+            if len(constraint_xy) > 0:
+                control_dict["group constraintxy"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_xy]
+                )
+                control_dict["fix constraintxy"] = "constraintxy setforce 0.0 0.0 NULL"
+                if calc_md:
+                    control_dict["velocity constraintxy"] = "set 0.0 0.0 NULL"
+            if len(constraint_yz) > 0:
+                control_dict["group constraintyz"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_yz]
+                )
+                control_dict["fix constraintyz"] = "constraintyz setforce NULL 0.0 0.0"
+                if calc_md:
+                    control_dict["velocity constraintyz"] = "set NULL 0.0 0.0"
+            if len(constraint_zx) > 0:
+                control_dict["group constraintxz"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_zx]
+                )
+                control_dict["fix constraintxz"] = "constraintxz setforce 0.0 NULL 0.0"
+                if calc_md:
+                    control_dict["velocity constraintxz"] = "set 0.0 NULL 0.0"
+            if len(constraint_x) > 0:
+                control_dict["group constraintx"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_x]
+                )
+                control_dict["fix constraintx"] = "constraintx setforce 0.0 NULL NULL"
+                if calc_md:
+                    control_dict["velocity constraintx"] = "set 0.0 NULL NULL"
+            if len(constraint_y) > 0:
+                control_dict["group constrainty"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_y]
+                )
+                control_dict["fix constrainty"] = "constrainty setforce NULL 0.0 NULL"
+                if calc_md:
+                    control_dict["velocity constrainty"] = "set NULL 0.0 NULL"
+            if len(constraint_z) > 0:
+                control_dict["group constraintz"] = "id " + " ".join(
+                    [str(ind) for ind in constraint_z]
+                )
+                control_dict["fix constraintz"] = "constraintz setforce NULL NULL 0.0"
+                if calc_md:
+                    control_dict["velocity constraintz"] = "set NULL NULL 0.0"
+    return control_dict

pyiron_lammps/compatibility/file.py

@@ -9,6 +9,7 @@
     calc_minimize,
     calc_static,
 )
+from pyiron_lammps.compatibility.constraints import set_selective_dynamics
 from pyiron_lammps.output import parse_lammps_output
 from pyiron_lammps.potential import get_potential_by_name
 from pyiron_lammps.structure import write_lammps_datafile
@@ -97,8 +98,20 @@ def lammps_file_interface_function(
     ]
 
     if calc_mode == "static":
+        lmp_str_lst += [
+            k + " " + v
+            for k, v in set_selective_dynamics(
+                structure=structure, calc_md=False
+            ).items()
+        ]
         lmp_str_lst += calc_static()
     elif calc_mode == "md":
+        lmp_str_lst += [
+            k + " " + v
+            for k, v in set_selective_dynamics(
+                structure=structure, calc_md=True
+            ).items()
+        ]
         if "n_ionic_steps" in calc_kwargs.keys():
             n_ionic_steps = int(calc_kwargs.pop("n_ionic_steps"))
         else:
@@ -108,6 +121,12 @@ def lammps_file_interface_function(
         lmp_str_lst += ["run {} ".format(n_ionic_steps)]
     elif calc_mode == "minimize":
         calc_kwargs["units"] = units
+        lmp_str_lst += [
+            k + " " + v
+            for k, v in set_selective_dynamics(
+                structure=structure, calc_md=False
+            ).items()
+        ]
         lmp_str_tmp_lst, structure = calc_minimize(structure=structure, **calc_kwargs)
         lmp_str_lst += lmp_str_tmp_lst
     else:

tests/test_constraints.py

@@ -0,0 +1,171 @@
+import unittest
+
+from ase.build import bulk
+from ase.constraints import FixAtoms, FixCom, FixedPlane
+
+from pyiron_lammps.compatibility.constraints import (
+    set_selective_dynamics,
+)
+
+
+class TestConstraints(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        structure = bulk("Cu", cubic=True)
+        structure.symbols[2:] = "Al"
+        cls.structure = structure
+
+    def test_selective_dynamics_mixed_calcmd_x(self):
+        atoms = self.structure.copy()
+        c1 = FixAtoms(indices=[atom.index for atom in atoms if atom.symbol == "Cu"])
+        c2 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [1, 0, 0],
+        )
+        atoms.set_constraint([c1, c2])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=True)
+        self.assertEqual(len(control_dict), 6)
+        self.assertTrue(control_dict["group constraintxyz"], "id 1 2")
+        self.assertTrue(
+            control_dict["fix constraintxyz"], "constraintxyz setforce 0.0 0.0 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintxyz"], "set 0.0 0.0 0.0")
+        self.assertTrue(control_dict["group constraintx"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constraintx"], "constraintx setforce NULL NULL 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintx"], "set NULL NULL 0.0")
+
+    def test_selective_dynamics_mixed_calcmd_y(self):
+        atoms = self.structure.copy()
+        c1 = FixAtoms(indices=[atom.index for atom in atoms if atom.symbol == "Cu"])
+        c2 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [0, 1, 0],
+        )
+        atoms.set_constraint([c1, c2])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=True)
+        self.assertEqual(len(control_dict), 6)
+        self.assertTrue(control_dict["group constraintxyz"], "id 1 2")
+        self.assertTrue(
+            control_dict["fix constraintxyz"], "constraintxyz setforce 0.0 0.0 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintxyz"], "set 0.0 0.0 0.0")
+        self.assertTrue(control_dict["group constrainty"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constrainty"], "constrainty setforce NULL 0.0 NULL"
+        )
+        self.assertTrue(control_dict["velocity constrainty"], "set NULL 0.0 NULL")
+
+    def test_selective_dynamics_mixed_calcmd_z(self):
+        atoms = self.structure.copy()
+        c1 = FixAtoms(indices=[atom.index for atom in atoms if atom.symbol == "Cu"])
+        c2 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [0, 0, 1],
+        )
+        atoms.set_constraint([c1, c2])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=True)
+        self.assertEqual(len(control_dict), 6)
+        self.assertTrue(control_dict["group constraintxyz"], "id 1 2")
+        self.assertTrue(
+            control_dict["fix constraintxyz"], "constraintxyz setforce 0.0 0.0 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintxyz"], "set 0.0 0.0 0.0")
+        self.assertTrue(control_dict["group constraintz"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constraintz"], "constraintz setforce NULL NULL 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintz"], "set NULL NULL 0.0")
+
+    def test_selective_dynamics_calcmd_xy(self):
+        atoms = self.structure.copy()
+        c1 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [1, 1, 0],
+        )
+        atoms.set_constraint([c1])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=True)
+        self.assertEqual(len(control_dict), 3)
+        self.assertTrue(control_dict["group constraintxy"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constraintxy"], "constraintxy setforce NULL 0.0 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintxy"], "set NULL 0.0 0.0")
+
+    def test_selective_dynamics_calcmd_xz(self):
+        atoms = self.structure.copy()
+        c1 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [1, 0, 1],
+        )
+        atoms.set_constraint([c1])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=True)
+        self.assertEqual(len(control_dict), 3)
+        self.assertTrue(control_dict["group constraintxz"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constraintxz"], "constraintxz setforce NULL 0.0 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintxz"], "set NULL 0.0 0.0")
+
+    def test_selective_dynamics_calcmd_yz(self):
+        atoms = self.structure.copy()
+        c1 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [0, 1, 1],
+        )
+        atoms.set_constraint([c1])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=True)
+        self.assertEqual(len(control_dict), 3)
+        self.assertTrue(control_dict["group constraintyz"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constraintyz"], "constraintyz setforce NULL 0.0 0.0"
+        )
+        self.assertTrue(control_dict["velocity constraintyz"], "set NULL 0.0 0.0")
+
+    def test_selective_dynamics_mixed_x(self):
+        atoms = self.structure.copy()
+        c1 = FixAtoms(indices=[atom.index for atom in atoms if atom.symbol == "Cu"])
+        c2 = FixedPlane(
+            [atom.index for atom in atoms if atom.symbol == "Al"],
+            [1, 0, 0],
+        )
+        atoms.set_constraint([c1, c2])
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=False)
+        self.assertEqual(len(control_dict), 4)
+        self.assertTrue(control_dict["group constraintxyz"], "id 1 2")
+        self.assertTrue(
+            control_dict["fix constraintxyz"], "constraintxyz setforce 0.0 0.0 0.0"
+        )
+        self.assertTrue(control_dict["group constraintx"], "id 3 4")
+        self.assertTrue(
+            control_dict["fix constraintx"], "constraintx setforce 0.0 NULL NULL"
+        )
+
+    def test_selective_dynamics_single_fix(self):
+        atoms = self.structure.copy()
+        c1 = FixAtoms(indices=[atom.index for atom in atoms if atom.symbol == "Cu"])
+        atoms.set_constraint(c1)
+        control_dict = set_selective_dynamics(structure=atoms, calc_md=False)
+        self.assertEqual(len(control_dict), 2)
+        self.assertTrue(control_dict["group constraintxyz"], "id 1 2")
+        self.assertTrue(
+            control_dict["fix constraintxyz"], "constraintxyz setforce 0.0 0.0 0.0"
+        )
+
+    def test_selective_dynamics_errors(self):
+        atoms = self.structure.copy()
+        atoms.set_constraint(FixCom())
+        with self.assertRaises(ValueError):
+            set_selective_dynamics(structure=atoms, calc_md=False)
+
+    def test_selective_dynamics_wrong_plane(self):
+        atoms = self.structure.copy()
+        atoms.set_constraint(
+            FixedPlane(
+                [atom.index for atom in atoms if atom.symbol == "Al"],
+                [2, 1, 0],
+            )
+        )
+        with self.assertRaises(ValueError):
+            set_selective_dynamics(structure=atoms, calc_md=False)