diff --git a/src/py/mat3ra/made/tools/calculate/__init__.py b/src/py/mat3ra/made/tools/calculate/__init__.py index a4d1e1f0..88b728f5 100644 --- a/src/py/mat3ra/made/tools/calculate/__init__.py +++ b/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 diff --git a/src/py/mat3ra/made/tools/calculate/ase/__init__.py b/src/py/mat3ra/made/tools/calculate/ase/__init__.py new file mode 100644 index 00000000..4ba38758 --- /dev/null +++ b/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) diff --git a/src/py/mat3ra/made/tools/calculate/ase/constraints.py b/src/py/mat3ra/made/tools/calculate/ase/constraints.py new file mode 100644 index 00000000..bb39cf5f --- /dev/null +++ b/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 diff --git a/src/py/mat3ra/made/tools/third_party.py b/src/py/mat3ra/made/tools/third_party.py index 9013c7de..9858d917 100644 --- a/src/py/mat3ra/made/tools/third_party.py +++ b/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",