Merge pull request #160 from Exabyte-io/feature/SOF-7427

feature/SOF-7427 feat: ASE distance norm calculator

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

@@ -1,9 +1,9 @@
-from typing import Optional
+from typing import List, Optional, Union
 
 from ...material import Material
 from ..analyze import get_surface_area
 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 ..third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
 from .interaction_functions import sum_of_inverse_distances_squared
 

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

@@ -0,0 +1,133 @@
+from typing import List, Optional, Union
+
+import numpy as np
+from mat3ra.made.material import Material
+
+from ...convert import from_ase
+from ...convert.utils import INTERFACE_LABELS_MAP, InterfacePartsEnum
+from ...third_party import ASEAtoms, ASECalculator, ASEFixAtoms, ASEFixedPlane, ase_all_changes
+from ..calculators import InterfaceMaterialCalculator, MaterialCalculator
+from .constraints import RigidFilmXYInterfaceConstraint
+
+
+def get_interface_part_indices(atoms: ASEAtoms, part: InterfacePartsEnum) -> List[int]:
+    """
+    Get the indices of the atoms in the interface part.
+
+    Args:
+        atoms (ASEAtoms): The ASE Atoms object.
+        part (str): The interface part to get the indices of.
+
+    Returns:
+        List[int]: The indices of the atoms in the interface part.
+    """
+    return [i for i, tag in enumerate(atoms.get_tags()) if tag == INTERFACE_LABELS_MAP[part]]
+
+
+class FilmSubstrateDistanceASECalculator(ASECalculator):
+    """
+    ASE calculator that calculates the interaction energy between a film and substrate in an interface material.
+
+    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.
+        is_substrate_fixed (bool): Whether to fix the substrate atoms.
+        is_z_axis_fixed (bool): Whether to fix atoms movement in the z direction.
+        symprec (float): The symmetry precision for the ASE calculator. This parameter determines the tolerance
+                         for symmetry operations, affecting the identification of equivalent atoms and the overall
+                         symmetry of the system. For more details, refer to the ASE documentation:
+                         https://wiki.fysik.dtu.dk/ase/ase/constraints.html#ase.constraints.FixSymmetry
+
+    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,
+        is_substrate_fixed: bool = True,
+        is_z_axis_fixed: bool = True,
+        symprec: float = 0.01,
+        material_calculator: Union[MaterialCalculator, InterfaceMaterialCalculator] = InterfaceMaterialCalculator(),
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.shadowing_radius = shadowing_radius
+        self.fix_substrate = is_substrate_fixed
+        self.fix_z = is_z_axis_fixed
+        self.symprec = symprec
+        self.material_calculator = material_calculator
+
+    def _add_constraints(self, atoms: ASEAtoms) -> ASEAtoms:
+        constraints: List[Union[ASEFixAtoms, ASEFixedPlane, RigidFilmXYInterfaceConstraint]] = []
+        if self.fix_substrate:
+            substrate_indices = get_interface_part_indices(atoms, InterfacePartsEnum.SUBSTRATE)
+            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]))
+
+        film_indices = get_interface_part_indices(atoms, InterfacePartsEnum.FILM)
+        if film_indices:
+            constraints.append(RigidFilmXYInterfaceConstraint(film_indices=film_indices))
+
+        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 = self.material_calculator.get_energy(material_plus)
+
+                forces[i, j] = -(energy_plus - energy) / dx
+
+        return forces
+
+    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
+
+        super().calculate(atoms, properties, system_changes)
+        material = Material(from_ase(atoms))
+        energy = self.material_calculator.get_energy(material)
+
+        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/calculate/ase/constraints.py

@@ -0,0 +1,91 @@
+from typing import List, Optional
+
+import numpy as np
+from pydantic import BaseModel
+
+from ...third_party import ASEAtoms
+
+
+class InterfaceConstraint(BaseModel):
+    """
+    Based on  https://wiki.fysik.dtu.dk/ase/ase/constraints.html#making-your-own-constraint-class
+    """
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
+        """
+        Adjust the positions of the atoms in the system.
+
+        Args:
+            atoms (ASEAtoms): The ASE Atoms object.
+            new_positions (numpy.ndarray): The new positions to adjust in place.
+        """
+        raise NotImplementedError
+
+    def adjust_forces(self, forces: np.ndarray) -> None:
+        """
+        Adjust the forces on the atoms in the system.
+
+        Args:
+            forces (numpy.ndarray): The forces array to adjust in place.
+        """
+        raise NotImplementedError
+
+
+class RigidFilmXYInterfaceConstraint(InterfaceConstraint):
+    """
+    Custom constraint to allow only rigid translation in x and y for film atoms.
+    """
+
+    film_indices: List[int] = []
+    initial_positions: Optional[np.ndarray] = None
+    reference_center: Optional[np.ndarray] = None
+
+    def adjust_positions(self, atoms: ASEAtoms, new_positions: np.ndarray) -> None:
+        """
+        Adjust the positions of film atoms to maintain rigidity in x and y.
+
+        Args:
+            atoms (ase.Atoms): The ASE Atoms object.
+            new_positions (numpy.ndarray): The new positions to adjust in place.
+        """
+        if self.initial_positions is None:
+            self.initial_positions = atoms.positions[self.film_indices].copy()
+            self.reference_center = self.initial_positions.mean(axis=0)
+
+        current_center = new_positions[self.film_indices].mean(axis=0)
+        desired_translation = current_center - self.reference_center
+        desired_translation[2] = 0.0
+
+        # Update positions to maintain rigid body translation in x and y
+        for i, idx in enumerate(self.film_indices):
+            new_positions[idx, 0] = self.initial_positions[i, 0] + desired_translation[0]
+            new_positions[idx, 1] = self.initial_positions[i, 1] + desired_translation[1]
+
+    def adjust_forces(self, forces: np.ndarray) -> None:
+        """
+        Adjust the forces on film atoms to ensure rigidity.
+
+        Args:
+            forces (numpy.ndarray): The forces array to adjust in place.
+        """
+        if self.initial_positions is None:
+            return  # No adjustment needed if initial positions are not set
+
+        # Project forces onto x and y for collective translation
+        net_force_x = np.sum(forces[self.film_indices, 0])
+        net_force_y = np.sum(forces[self.film_indices, 1])
+
+        # Distribute the net force equally to all film atoms to maintain rigidity
+        num_film_atoms = len(self.film_indices)
+        if num_film_atoms == 0:
+            return  # Avoid division by zero
+
+        distributed_force_x = net_force_x / num_film_atoms
+        distributed_force_y = net_force_y / num_film_atoms
+
+        for idx in self.film_indices:
+            forces[idx, 0] = distributed_force_x
+            forces[idx, 1] = distributed_force_y

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",