Allow ASE constraint to be used in geomopt

- Separate parameter from filter_func
- kwargs come from minimize_args, same as filter args
- Check if the constraint needs "atoms" and insert if so

janus_core/calculations/geom_opt.py

@@ -1,9 +1,10 @@
 """Prepare and run geometry optimization."""
 
+import inspect
 from typing import Any, Callable, Optional, Union
 import warnings
 
-from ase import Atoms, filters, units
+from ase import Atoms, constraints, filters, units
 from ase.filters import FrechetCellFilter
 from ase.io import read
 import ase.optimize
@@ -62,9 +63,14 @@ class GeomOpt(BaseCalculation):
     angle_tolerance : float
         Angle precision for spglib symmetry determination, in degrees. Default is -1.0,
         which means an internally optimized routine is used to judge symmetry.
+    constraint_func : Optional[Union[Callable, str]]
+        Constraint function, or name of function from ase.constraints. Default is None.
+    constraint_kwargs : Optional[dict[str, Any]]
+        Keyword arguments to pass to constraint_func. Default is {}.
     filter_func : Optional[Union[Callable, str]]
-        Filter function, or name of function from ase.filters to apply constraints to
-        atoms. Default is `FrechetCellFilter`.
+        Filter function, or name of function from ase.filters to apply filters
+        to atoms. (These are used for lattice-vector optimisation.) Default is
+        `FrechetCellFilter`.
     filter_kwargs : Optional[dict[str, Any]]
         Keyword arguments to pass to filter_func. Default is {}.
     optimizer : Union[Callable, str]
@@ -106,6 +112,8 @@ def __init__(
         steps: int = 1000,
         symmetry_tolerance: float = 0.001,
         angle_tolerance: float = -1.0,
+        constraint_func: Optional[Union[Callable, str]] = None,
+        constraint_kwargs: Optional[dict[str, Any]] = None,
         filter_func: Optional[Union[Callable, str]] = FrechetCellFilter,
         filter_kwargs: Optional[dict[str, Any]] = None,
         optimizer: Union[Callable, str] = LBFGS,
@@ -173,16 +181,18 @@ def __init__(
             Keyword arguments to pass to ase.io.write to save optimization trajectory.
             Must include "filename" keyword. Default is {}.
         """
-        (read_kwargs, filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs) = (
+        (read_kwargs, constraint_kwargs, filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs) = (
             none_to_dict(
-                (read_kwargs, filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs)
+                (read_kwargs, constraint_kwargs, filter_kwargs, opt_kwargs, write_kwargs, traj_kwargs)
             )
         )
 
         self.fmax = fmax
         self.steps = steps
         self.symmetry_tolerance = symmetry_tolerance
         self.angle_tolerance = angle_tolerance
+        self.constraint_func = constraint_func
+        self.constraint_kwargs = constraint_kwargs
         self.filter_func = filter_func
         self.filter_kwargs = filter_kwargs
         self.optimizer = optimizer
@@ -232,12 +242,30 @@ def __init__(
         # Configure optimizer dynamics
         self.set_optimizer()
 
+    def _get_constraint_args(self, constraint_class: object) -> list[Any]:
+        """Inspect constraint class for mandatory arguments
+
+        For now we are just looking for the "atoms" parameter of FixSymmetry
+        """
+        parameters = inspect.signature(constraint_class.__init__).parameters
+        if "atoms" in parameters:
+            return [self.struct]
+
+        return []
+
     def set_optimizer(self) -> None:
         """Set optimizer for geometry optimization."""
         self._set_functions()
         if self.logger:
             self.logger.info("Using optimizer: %s", self.optimizer.__name__)
 
+        if self.constraint_func is not None:
+            constraint_args = self._get_constraint_args(self.constraint_func)
+            self.struct.set_constraint(self.constraint_func(*constraint_args, **self.constraint_kwargs))
+
+            if self.logger:
+                self.logger.info("Using constraint: %s", self.constraint_func.__name__)
+
         if self.filter_func is not None:
             if "scalar_pressure" in self.filter_kwargs:
                 self.filter_kwargs["scalar_pressure"] *= units.GPa
@@ -263,13 +291,19 @@ def set_optimizer(self) -> None:
             self.dyn = self.optimizer(self.struct, **self.opt_kwargs)
 
     def _set_functions(self) -> None:
-        """Set optimizer and filter functions."""
+        """Set optimizer, constraint and filter functions."""
         if isinstance(self.optimizer, str):
             try:
                 self.optimizer = getattr(ase.optimize, self.optimizer)
             except AttributeError as e:
                 raise AttributeError(f"No such optimizer: {self.optimizer}") from e
 
+        if self.constraint_func is not None and isinstance(self.constraint_func, str):
+            try:
+                self.constraint_func = getattr(constraints, self.constraint_func)
+            except AttributeError as e:
+                raise AttributeError(f"No such constraint: {self.constraint_func}") from e
+
         if self.filter_func is not None and isinstance(self.filter_func, str):
             try:
                 self.filter_func = getattr(filters, self.filter_func)

janus_core/cli/geomopt.py

@@ -126,6 +126,12 @@ def geomopt(
             )
         ),
     ] = None,
+    constraint_func: Annotated[
+        str,
+        Option(
+            help="Name of ASE constraint function to use."
+        ),
+    ] = None,
     pressure: Annotated[
         float, Option(help="Scalar pressure when optimizing cell geometry, in GPa.")
     ] = 0.0,
@@ -177,9 +183,13 @@ def geomopt(
     opt_cell_fully : bool
         Whether to fully optimize the cell vectors, angles, and atomic positions.
         Default is False.
+    constraint_func : Optional[str]
+        Name of constraint function from ase.constraints, to apply constraints
+        to atoms. Parameters should be included as a "constraint_kwargs" dict
+        within "minimize_kwargs". Default is None
     filter_func : Optional[str]
-        Name of filter function from ase.filters or ase.constraints, to apply
-        constraints to atoms. If using --opt-cell-lengths or --opt-cell-fully, defaults
+        Name of filter function from ase.filters, to apply (unit cell) filter to atoms.
+    If using --opt-cell-lengths or --opt-cell-fully, defaults
         to `FrechetCellFilter` if available, otherwise `ExpCellFilter`.
     pressure : float
         Scalar pressure when optimizing cell geometry, in GPa. Passed to the filter
@@ -221,6 +231,7 @@ def geomopt(
     # Check optimized structure path not duplicated
     if "filename" in write_kwargs:
         raise ValueError("'filename' must be passed through the --out option")
+
     if out:
         write_kwargs["filename"] = out
 
@@ -255,6 +266,7 @@ def geomopt(
         "optimizer": optimizer,
         "fmax": fmax,
         "steps": steps,
+        "constraint_func": constraint_func,
         **opt_cell_fully_dict,
         **minimize_kwargs,
         "write_results": True,