implementation of constrains in asegeoopt (#249)

* implementation of constrains in asegeoopt

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Fabian Zills <[email protected]>

Author: 3 people

Diff for: ipsuite/calculators/__init__.pyi

@@ -3,6 +3,7 @@ from .ase_geoopt import ASEGeoOpt
 from .ase_md import (
     ASEMD,
     BoxOscillatingRampModifier,
+    FixedLayerConstraint,
     FixedSphereConstraint,
     LangevinThermostat,
     NPTThermostat,
@@ -36,4 +37,5 @@ __all__ = [
     "NPTThermostat",
     "OrcaSinglePoint",
     "LammpsSimulator",
+    "FixedLayerConstraint",
 ]

Diff for: ipsuite/calculators/ase_geoopt.py

@@ -28,6 +28,7 @@ class ASEGeoOpt(base.ProcessSingleAtom):
     model_outs = zntrack.dvc.outs(zntrack.nwd / "model_outs")
     optimizer: str = zntrack.zn.params("FIRE")
     checker_list: list = zntrack.deps(None)
+    constraint_list: list = zntrack.deps(None)
 
     repeat: list = zntrack.zn.params([1, 1, 1])
     run_kwargs: dict = zntrack.zn.params({"fmax": 0.05})
@@ -39,6 +40,8 @@ class ASEGeoOpt(base.ProcessSingleAtom):
     def run(self):
         if self.checker_list is None:
             self.checker_list = []
+        if self.constraint_list is None:
+            self.constraint_list = []
 
         self.model_outs.mkdir(parents=True, exist_ok=True)
         (self.model_outs / "outs.txt").write_text("Lorem Ipsum")
@@ -48,6 +51,9 @@ def run(self):
         atoms = atoms.repeat(self.repeat)
         atoms.calc = calculator
 
+        for constraint in self.constraint_list:
+            atoms.set_constraint(constraint.get_constraint(atoms))
+
         atoms_cache = []
 
         db = znh5md.io.DataWriter(self.traj_file)

Diff for: ipsuite/nodes.py

@@ -76,6 +76,7 @@ class _Nodes:
     ThresholdCheck = "ipsuite.analysis.ThresholdCheck"
     TemperatureCheck = "ipsuite.analysis.TemperatureCheck"
     FixedSphereConstraint = "ipsuite.calculators.FixedSphereConstraint"
+    FixedLayerConstraint = "ipsuite.calculators.FixedLayerConstraint"
     AnalyseDensity = "ipsuite.analysis.AnalyseDensity"
 
     # calculators