feature/SOF-7427 feat: ASE distance norm calculator

src/py/mat3ra/made/tools/analyze.py

@@ -1,10 +1,13 @@
-from typing import Callable, List, Literal, Optional
+from typing import Callable, List, Literal, Optional, Tuple
 
 import numpy as np
+from scipy.spatial import cKDTree
 
 from ..material import Material
+from .build.passivation.enums import SurfaceTypes
 from .convert import decorator_convert_material_args_kwargs_to_atoms, to_pymatgen
 from .third_party import ASEAtoms, PymatgenIStructure, PymatgenVoronoiNN
+from .utils import decorator_handle_periodic_boundary_conditions
 
 
 @decorator_convert_material_args_kwargs_to_atoms
@@ -265,13 +268,16 @@ def get_atom_indices_with_condition_on_coordinates(
 def get_nearest_neighbors_atom_indices(
     material: Material,
     coordinate: Optional[List[float]] = None,
+    tolerance: float = 0.1,
+    cutoff: float = 13.0,
 ) -> Optional[List[int]]:
     """
     Returns the indices of direct neighboring atoms to a specified position in the material using Voronoi tessellation.
 
     Args:
         material (Material): The material object to find neighbors in.
         coordinate (List[float]): The position to find neighbors for.
+        cutoff (float): The cutoff radius for identifying neighbors.
 
     Returns:
         List[int]: A list of indices of neighboring atoms, or an empty list if no neighbors are found.
@@ -280,17 +286,27 @@ def get_nearest_neighbors_atom_indices(
         coordinate = [0, 0, 0]
     structure = to_pymatgen(material)
     voronoi_nn = PymatgenVoronoiNN(
-        tol=0.5,
-        cutoff=15.0,
-        allow_pathological=False,
+        tol=tolerance,
+        cutoff=cutoff,
         weight="solid_angle",
-        extra_nn_info=True,
+        extra_nn_info=False,
         compute_adj_neighbors=True,
     )
-    structure.append("X", coordinate, validate_proximity=False)
-    neighbors = voronoi_nn.get_nn_info(structure, len(structure.sites) - 1)
+    coordinates = material.basis.coordinates
+    site_index = coordinates.get_element_id_by_value(coordinate)
+    remove_dummy_atom = False
+    if site_index is None:
+        structure.append("X", coordinate, validate_proximity=False)
+        site_index = len(structure.sites) - 1
+
+        remove_dummy_atom = True
+    try:
+        neighbors = voronoi_nn.get_nn_info(structure, site_index)
+    except ValueError:
+        return None
     neighboring_atoms_pymatgen_ids = [n["site_index"] for n in neighbors]
-    structure.remove_sites([-1])
+    if remove_dummy_atom:
+        structure.remove_sites([-1])
 
     all_coordinates = material.basis.coordinates
     all_coordinates.filter_by_indices(neighboring_atoms_pymatgen_ids)
@@ -324,3 +340,185 @@ def get_atomic_coordinates_extremum(
     coordinates = new_material.basis.coordinates.to_array_of_values_with_ids()
     values = [coord.value[{"x": 0, "y": 1, "z": 2}[axis]] for coord in coordinates]
     return getattr(np, extremum)(values)
+
+
+def get_local_extremum_atom_index(
+    material: Material,
+    coordinate: List[float],
+    extremum: Literal["max", "min"] = "max",
+    vicinity: float = 1.0,
+    use_cartesian_coordinates: bool = False,
+) -> int:
+    """
+    Return the id of the atom with the minimum or maximum z-coordinate
+    within a certain vicinity of a given (x, y) coordinate.
+
+    Args:
+        material (Material): Material object.
+        coordinate (List[float]): (x, y, z) coordinate to find the local extremum.
+        extremum (str): "min" or "max".
+        vicinity (float): Radius of the vicinity, in Angstroms.
+        use_cartesian_coordinates (bool): Whether to use Cartesian coordinates.
+
+    Returns:
+        int: id of the atom with the minimum or maximum z-coordinate.
+    """
+    new_material = material.clone()
+    new_material.to_cartesian()
+    if not use_cartesian_coordinates:
+        coordinate = new_material.basis.cell.convert_point_to_cartesian(coordinate)
+
+    coordinates = np.array(new_material.basis.coordinates.values)
+    ids = np.array(new_material.basis.coordinates.ids)
+    tree = cKDTree(coordinates[:, :2])
+    indices = tree.query_ball_point(coordinate[:2], vicinity)
+    z_values = [(id, coord[2]) for id, coord in zip(ids[indices], coordinates[indices])]
+
+    if extremum == "max":
+        extremum_z_atom = max(z_values, key=lambda item: item[1])
+    else:
+        extremum_z_atom = min(z_values, key=lambda item: item[1])
+
+    return extremum_z_atom[0]
+
+
+def height_check(z: float, z_extremum: float, depth: float, surface: SurfaceTypes):
+    return (z >= z_extremum - depth) if surface == SurfaceTypes.TOP else (z <= z_extremum + depth)
+
+
+def shadowing_check(z: float, neighbors_indices: List[int], surface: SurfaceTypes, coordinates: np.ndarray):
+    return not any(
+        (coordinates[n][2] > z if surface == SurfaceTypes.TOP else coordinates[n][2] < z) for n in neighbors_indices
+    )
+
+
+@decorator_handle_periodic_boundary_conditions(cutoff=0.1)
+def get_surface_atom_indices(
+    material: Material, surface: SurfaceTypes = SurfaceTypes.TOP, shadowing_radius: float = 2.5, depth: float = 5
+) -> List[int]:
+    """
+    Identify exposed atoms on the top or bottom surface of the material.
+
+    Args:
+        material (Material): Material object to get surface atoms from.
+        surface (SurfaceTypes): Specify "top" or "bottom" to detect the respective surface atoms.
+        shadowing_radius (float): Radius for atoms shadowing underlying from detecting as exposed.
+        depth (float): Depth from the surface to look for exposed atoms.
+
+    Returns:
+        List[int]: List of indices of exposed surface atoms.
+    """
+    new_material = material.clone()
+    new_material.to_cartesian()
+    coordinates = np.array(new_material.basis.coordinates.values)
+    ids = new_material.basis.coordinates.ids
+    kd_tree = cKDTree(coordinates)
+
+    z_extremum = np.max(coordinates[:, 2]) if surface == SurfaceTypes.TOP else np.min(coordinates[:, 2])
+
+    exposed_atoms_indices = []
+    for idx, (x, y, z) in enumerate(coordinates):
+        if height_check(z, z_extremum, depth, surface):
+            neighbors_indices = kd_tree.query_ball_point([x, y, z], r=shadowing_radius)
+            if shadowing_check(z, neighbors_indices, surface, coordinates):
+                exposed_atoms_indices.append(ids[idx])
+
+    return exposed_atoms_indices
+
+
+def get_coordination_numbers(
+    material: Material,
+    indices: Optional[List[int]] = None,
+    cutoff: float = 3.0,
+) -> List[int]:
+    """
+    Calculate the coordination numbers of atoms in the material.
+
+    Args:
+        material (Material): Material object to calculate coordination numbers for.
+        indices (List[int]): List of atom indices to calculate coordination numbers for.
+        cutoff (float): The cutoff radius for identifying neighbors.
+
+    Returns:
+        List[int]: List of coordination numbers for each atom in the material.
+    """
+    new_material = material.clone()
+    new_material.to_cartesian()
+    if indices is not None:
+        new_material.basis.coordinates.filter_by_indices(indices)
+    coordinates = np.array(new_material.basis.coordinates.values)
+    kd_tree = cKDTree(coordinates)
+
+    coordination_numbers = []
+    for idx, (x, y, z) in enumerate(coordinates):
+        neighbors = kd_tree.query_ball_point([x, y, z], r=cutoff)
+        # Explicitly remove the atom itself from the list of neighbors
+        neighbors = [n for n in neighbors if n != idx]
+        coordination_numbers.append(len(neighbors))
+
+    return coordination_numbers
+
+
+@decorator_handle_periodic_boundary_conditions(cutoff=0.1)
+def get_undercoordinated_atom_indices(
+    material: Material,
+    indices: List[int],
+    cutoff: float = 3.0,
+    coordination_threshold: int = 3,
+) -> List[int]:
+    """
+    Identify undercoordinated atoms among the specified indices in the material.
+
+    Args:
+        material (Material): Material object to identify undercoordinated atoms in.
+        indices (List[int]): List of atom indices to check for undercoordination.
+        cutoff (float): The cutoff radius for identifying neighbors.
+        coordination_threshold (int): The coordination number threshold for undercoordination.
+
+    Returns:
+        List[int]: List of indices of undercoordinated atoms.
+    """
+    coordination_numbers = get_coordination_numbers(material, indices, cutoff)
+    undercoordinated_atoms_indices = [i for i, cn in enumerate(coordination_numbers) if cn <= coordination_threshold]
+    return undercoordinated_atoms_indices
+
+
+def get_optimal_displacements(
+    material: Material,
+    grid_size: Tuple[int, int] = (10, 10),
+    search_range: Tuple[float, float] = (-1.0, 1.0),
+    shadowing_radius: float = 2.5,
+) -> List[List[float]]:
+    """
+    Return all optimal displacements for the interface material by calculating
+    the norm of distances for each film atom to substrate atoms within a certain radius.
+    The displacement is done on a grid and the displacement that yields minimum norm is returned.
+
+    Args:
+        material (Material): The interface Material object.
+        grid_size (Tuple[int, int]): The size of the grid.
+        search_range (Tuple[float, float]): The range to search for optimal displacements.
+        shadowing_radius (float): The shadowing radius to detect the surface atoms, in Angstroms.
+
+    Returns:
+        List[List[float]]: The optimal displacements.
+    """
+    from .calculate import calculate_norm_of_distances
+    from .modify import displace_interface
+
+    x_values = np.linspace(search_range[0], search_range[1], grid_size[0])
+    y_values = np.linspace(search_range[0], search_range[1], grid_size[1])
+
+    norms = np.zeros(grid_size)
+
+    for i, x in enumerate(x_values):
+        for j, y in enumerate(y_values):
+            displaced_material = displace_interface(material, [x, y, 0], use_cartesian_coordinates=True)
+            norm = calculate_norm_of_distances(displaced_material, shadowing_radius)
+            norms[i, j] = norm
+
+    min_norm = np.min(norms)
+    min_positions = np.argwhere(norms == min_norm)
+
+    displacements_with_min_norm = [[x_values[pos[0]], y_values[pos[1]], 0] for pos in min_positions]
+    return displacements_with_min_norm

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

@@ -1,10 +1,18 @@
 from typing import Optional
 
+import numpy as np
+from mat3ra.made.tools.convert.utils import InterfacePartsEnum
+
 from ..material import Material
-from .analyze import get_surface_area
+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 .build.passivation.enums import SurfaceTypes
+from .convert import decorator_convert_material_args_kwargs_to_atoms, from_ase
 from .third_party import ASEAtoms, ASECalculator, ASECalculatorEMT
+from .utils import calculate_norm_of_distances_between_coordinates, decorator_handle_periodic_boundary_conditions
+
+# TODO: import from thrid_party
+from ase.calculators.calculator import all_changes
 
 
 @decorator_convert_material_args_kwargs_to_atoms
@@ -125,3 +133,99 @@ def calculate_interfacial_energy(
     surface_energy_film = calculate_surface_energy(film_slab, film_bulk, calculator)
     adhesion_energy = calculate_adhesion_energy(interface, substrate_slab, film_slab, calculator)
     return surface_energy_film + surface_energy_substrate - adhesion_energy
+
+
+@decorator_handle_periodic_boundary_conditions(cutoff=0.25)
+def calculate_norm_of_distances(material: Material, shadowing_radius: float = 2.5) -> float:
+    """
+    Calculate the norm of distances between interfacial gap facing atoms of the film and the substrate.
+
+    Args:
+        material (Material): The interface Material object.
+        shadowing_radius (float): The shadowing radius to detect the surface atoms, in Angstroms.
+
+    Returns:
+        float: The calculated norm.
+    """
+    film_material = material.clone()
+    substrate_material = material.clone()
+    film_atoms_basis = film_material.basis.filter_atoms_by_labels([int(InterfacePartsEnum.FILM)])
+    substrate_atoms_basis = substrate_material.basis.filter_atoms_by_labels([int(InterfacePartsEnum.SUBSTRATE)])
+
+    film_material.basis = film_atoms_basis
+    substrate_material.basis = substrate_atoms_basis
+    film_atoms_surface_indices = get_surface_atom_indices(
+        film_material, SurfaceTypes.BOTTOM, shadowing_radius=shadowing_radius
+    )
+    substrate_atoms_surface_indices = get_surface_atom_indices(
+        substrate_material, SurfaceTypes.TOP, shadowing_radius=shadowing_radius
+    )
+
+    film_atoms_surface_coordinates = film_material.basis.coordinates
+    film_atoms_surface_coordinates.filter_by_ids(film_atoms_surface_indices)
+    substrate_atoms_surface_coordinates = substrate_material.basis.coordinates
+    substrate_atoms_surface_coordinates.filter_by_ids(substrate_atoms_surface_indices)
+
+    film_coordinates_values = np.array(film_atoms_surface_coordinates.values)
+    substrate_coordinates_values = np.array(substrate_atoms_surface_coordinates.values)
+
+    return calculate_norm_of_distances_between_coordinates(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.
+
+    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, **kwargs):
+        super().__init__(**kwargs)
+        self.shadowing_radius = shadowing_radius
+        self.force_constant = force_constant
+
+    @decorator_convert_material_args_kwargs_to_atoms
+    def calculate(self, atoms: Optional[ASEAtoms] = None, properties=["energy"], system_changes=all_changes):
+        if atoms is None:
+            atoms = self.atoms.copy()
+
+        ASECalculator.calculate(self, atoms, properties, system_changes)
+        material = Material(from_ase(atoms))
+        norm = calculate_norm_of_distances(material, self.shadowing_radius)
+        self.results = {"energy": norm}
+
+        if "forces" in properties:
+            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))
+                    norm_plus = calculate_norm_of_distances(material_plus, self.shadowing_radius)
+
+                    # Finite difference approximation of the force
+                    forces[i, j] = -self.force_constant * (norm_plus - norm) / dx
+
+            self.results["forces"] = forces