feature/SOF-7427 feat: ASE distance norm calculator

src/py/mat3ra/made/tools/calculate.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional
+from typing import Callable, List, Optional, Union
 
 import numpy as np
 from mat3ra.made.tools.convert.utils import InterfacePartsEnum
@@ -7,10 +7,10 @@
 from ..material import Material
 from .analyze import get_surface_area, get_surface_atom_indices
 from .build.interface.utils import get_slab
-from .convert import decorator_convert_material_args_kwargs_to_atoms
+from .convert import decorator_convert_material_args_kwargs_to_atoms, from_ase
 from .enums import SurfaceTypes
 from .modify import get_interface_part
-from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
+from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT, ASEFixAtoms, ASEFixedPlane, ase_all_changes
 from .utils import decorator_handle_periodic_boundary_conditions, get_sum_of_inverse_distances_squared
 
 
@@ -173,3 +173,107 @@ def calculate_film_substrate_interaction_metric(
     substrate_coordinates_values = np.array(substrate_atoms_surface_coordinates.values)
 
     return interaction_function(film_coordinates_values, substrate_coordinates_values)
+
+
+class SurfaceDistanceCalculator(ASECalculator):
+    """
+    ASE calculator that computes the norm of distances between interfacial gap facing atoms
+    of the film and the substrate.
+
+    Args:
+        shadowing_radius (float): The radius for atom to shadow underlying from being considered surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the forces.
+        fix_substrate (bool): Whether to fix the substrate atoms.
+        fix_z (bool): Whether to fix atoms movement in the z direction.
+        symprec (float): The symmetry precision for the ASE calculator.
+
+    Example usage:
+    ```python
+    from ase.optimize import BFGS
+    atoms = to_ase(material)
+    calc = SurfaceDistanceCalculator(shadowing_radius=2.5)
+
+    atoms.calc = calc
+    opt = BFGS(atoms)
+    opt.run(fmax=0.05)
+    ```
+    Args:
+        shadowing_radius (float): Radius for atoms shadowing underlying from being treated as a surface, in Angstroms.
+        force_constant (float): The force constant for the finite difference approximation of the
+    Note:
+        Built following: https://wiki.fysik.dtu.dk/ase/development/calculators.html
+
+        The calculate method is responsible for computing the energy and forces (if requested).
+        Forces are estimated using a finite difference method, which is a simple approximation
+        and might not be the most accurate or efficient for all cases.
+    """
+
+    implemented_properties = ["energy", "forces"]
+
+    def __init__(
+        self,
+        shadowing_radius: float = 2.5,
+        force_constant: float = 1.0,
+        fix_substrate: bool = True,
+        fix_z: bool = True,
+        symprec: float = 0.01,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.shadowing_radius = shadowing_radius
+        self.force_constant = force_constant
+        self.fix_substrate = fix_substrate
+        self.fix_z = fix_z
+        self.symprec = symprec
+
+    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane]] = []
+        if self.fix_substrate:
+            substrate_indices = [i for i, tag in enumerate(atoms.get_tags()) if tag == 0]
+            constraints.append(ASEFixAtoms(indices=substrate_indices))
+        if self.fix_z:
+            all_indices = list(range(len(atoms)))
+            constraints.append(ASEFixedPlane(indices=all_indices, direction=[0, 0, 1]))
+
+        atoms.set_constraint(constraints)
+        return atoms
+
+    def _calculate_forces(self, atoms: ASEAtoms, energy: float) -> np.ndarray:
+        forces = np.zeros((len(atoms), 3))
+        dx = 0.01
+        for i in range(len(atoms)):
+            for j in range(3):
+                atoms_plus = atoms.copy()
+                atoms_plus.positions[i, j] += dx
+                material_plus = Material(from_ase(atoms_plus))
+                energy_plus = calculate_film_substrate_interaction_metric(material_plus, self.shadowing_radius)
+
+                forces[i, j] = -self.force_constant * (energy_plus - energy) / dx
+
+        return forces
+
+    @decorator_convert_material_args_kwargs_to_atoms
+    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=ase_all_changes):
+        if atoms is None:
+            atoms = self.atoms.copy()
+
+        atoms = self._add_constraints(atoms)
+        constraints = atoms.constraints
+
+        ASECalculator.calculate(self, atoms, properties, system_changes)
+        material = Material(from_ase(atoms))
+        energy = calculate_film_substrate_interaction_metric(material, self.shadowing_radius)
+
+        self.results = {"energy": energy}
+
+        if "forces" in properties:
+            forces = self._calculate_forces(atoms, energy)
+            for constraint in constraints:
+                constraint.adjust_forces(atoms, forces)
+            self.results["forces"] = forces
+
+    def get_potential_energy(self, atoms=None, force_consistent=False):
+        return self.get_property("energy", atoms)
+
+    def get_forces(self, atoms=None):
+        return self.get_property("forces", atoms)

src/py/mat3ra/made/tools/third_party.py

@@ -2,7 +2,10 @@
 from ase.build import add_vacuum as ase_add_vacuum
 from ase.build.supercells import make_supercell as ase_make_supercell
 from ase.calculators.calculator import Calculator as ASECalculator
+from ase.calculators.calculator import all_changes as ase_all_changes
 from ase.calculators.emt import EMT as ASECalculatorEMT
+from ase.constraints import FixAtoms as ASEFixAtoms
+from ase.constraints import FixedPlane as ASEFixedPlane
 from pymatgen.analysis.defects.core import Interstitial as PymatgenInterstitial
 from pymatgen.analysis.defects.core import Substitution as PymatgenSubstitution
 from pymatgen.analysis.defects.core import Vacancy as PymatgenVacancy
@@ -25,6 +28,9 @@
     "ASEAtoms",
     "ASECalculator",
     "ASECalculatorEMT",
+    "ASEFixAtoms",
+    "ASEFixedPlane",
+    "ase_all_changes",
     "PymatgenLattice",
     "PymatgenStructure",
     "PymatgenIStructure",