diff --git a/src/py/mat3ra/made/tools/analyze.py b/src/py/mat3ra/made/tools/analyze.py index d7cf5f5b..889b6a88 100644 --- a/src/py/mat3ra/made/tools/analyze.py +++ b/src/py/mat3ra/made/tools/analyze.py @@ -376,40 +376,63 @@ def get_surface_atom_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) - print("neighbors_indices", type(neighbors_indices), neighbors_indices) if shadowing_check(z, neighbors_indices, surface, coordinates): exposed_atoms_indices.append(ids[idx]) return exposed_atoms_indices -def get_undercoordinated_atom_indices( - material: Material, surface: SurfaceTypes = SurfaceTypes.TOP, coordination_number: int = 3, cutoff: float = 3.0 +def get_coordination_numbers( + material: Material, + indices: Optional[List[int]] = None, + cutoff: float = 3.0, ) -> List[int]: """ - Identify undercoordinated atoms on the top or bottom surface of the material. + Calculate the coordination numbers of atoms in the material. Args: - material (Material): Material object to get undercoordinated atoms from. - surface (SurfaceTypes): Specify "top" or "bottom" to detect the respective surface atoms. - coordination_number (int): The coordination number to detect undercoordinated atoms. + 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 indices of undercoordinated surface atoms. + 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) - 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]) - - undercoordinated_atoms_indices = [] + coordination_numbers = [] for idx, (x, y, z) in enumerate(coordinates): - if z == z_extremum: - neighbors = kd_tree.query_ball_point([x, y, z], r=cutoff) - if len(neighbors) < coordination_number: - undercoordinated_atoms_indices.append(ids[idx]) + 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 + +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 diff --git a/src/py/mat3ra/made/tools/build/passivation/builders.py b/src/py/mat3ra/made/tools/build/passivation/builders.py index 6be044ce..576efac8 100644 --- a/src/py/mat3ra/made/tools/build/passivation/builders.py +++ b/src/py/mat3ra/made/tools/build/passivation/builders.py @@ -5,7 +5,12 @@ from pydantic import BaseModel from .enums import SurfaceTypes -from ...analyze import get_surface_atom_indices, get_undercoordinated_atom_indices, get_nearest_neighbors_atom_indices +from ...analyze import ( + get_surface_atom_indices, + get_undercoordinated_atom_indices, + get_nearest_neighbors_atom_indices, + get_coordination_numbers, +) from ...modify import translate_to_z_level from ...build import BaseBuilder from .configuration import ( @@ -117,14 +122,13 @@ def _get_passivant_coordinates( return (np.array(surface_atoms_coordinates) + np.array(passivant_bond_vector_crystal)).tolist() -class UndercoordinationPassivationBuilderParameters(BaseModel): +class UndercoordinationPassivationBuilderParameters(SurfacePassivationBuilderParameters): """ Parameters for the UndercoordinationPassivationBuilder. Args: coordination_threshold (int): The coordination threshold for undercoordination. """ - cutoff: float = 3.0 coordination_threshold: int = 3 @@ -135,43 +139,51 @@ class UndercoordinationPassivationBuilder(PassivationBuilder): Detects atoms with coordination number below a threshold and passivates them. """ - build_parameters: UndercoordinationPassivationBuilderParameters = UndercoordinationPassivationBuilderParameters() + _BuildParametersType = UndercoordinationPassivationBuilderParameters + _DefaultBuildParameters = UndercoordinationPassivationBuilderParameters() def create_passivated_material(self, configuration: PassivationConfiguration) -> Material: material = super().create_passivated_material(configuration) - passivant_coordinates_values = self._get_passivant_coordinates(material, configuration) + surface_atoms_indices = get_surface_atom_indices( + material=material, + surface=SurfaceTypes.TOP, + shadowing_radius=self.build_parameters.shadowing_radius, + depth=self.build_parameters.depth, + ) + undercoordinated_atoms_indices = get_undercoordinated_atom_indices( + material=material, + indices=surface_atoms_indices, + cutoff=self.build_parameters.shadowing_radius, + coordination_threshold=self.build_parameters.coordination_threshold, + ) + passivant_coordinates_values = self._get_passivant_coordinates( + material, configuration, undercoordinated_atoms_indices + ) return self._add_passivant_atoms(material, passivant_coordinates_values, configuration.passivant) - def _get_passivant_coordinates(self, material: Material, configuration: PassivationConfiguration): + def _get_passivant_coordinates( + self, material: Material, configuration: PassivationConfiguration, undercoordinated_atoms_indices: list + ): """ - Calculate the coordinates for placing passivants based on the specified edge type. + Calculate the coordinates for placing passivating atoms based on the specified edge type. Args: material (Material): Material to passivate. configuration (SurfacePassivationConfiguration): Configuration for passivation. + undercoordinated_atoms_indices (list): Indices of undercoordinated atoms. """ - undercoordinated_atoms_indices = get_undercoordinated_atom_indices( - material=material, - surface=SurfaceTypes.TOP, - coordination_number=self.build_parameters.coordination_threshold, - cutoff=self.build_parameters.cutoff, - ) - passivant_coordinates = [] - for idx in undercoordinated_atoms_indices: nearest_neighbors = get_nearest_neighbors_atom_indices( material=material, coordinate=material.basis.coordinates.get_element_value_by_index(idx), - cutoff=self.build_parameters.cutoff, + cutoff=self.build_parameters.shadowing_radius, ) - if nearest_neighbors is None: continue average_coordinate = np.mean( [material.basis.coordinates.get_element_value_by_index(i) for i in nearest_neighbors], axis=0 ) - bond_vector = material.basis.coordinates.get_element_value_by_index(idx) - average_coordinate bond_vector = bond_vector / np.linalg.norm(bond_vector) * configuration.bond_length passivant_bond_vector_crystal = material.basis.cell.convert_point_to_crystal(bond_vector) @@ -181,3 +193,20 @@ def _get_passivant_coordinates(self, material: Material, configuration: Passivat ) return passivant_coordinates + + def get_coordination_numbers(self, material: Material): + """ + Get the coordination numbers for all atoms in the material. + + Args: + material (Material): The material object. + + Returns: + set: The coordination numbers for all atoms in the material. + """ + + coordination_numbers = set( + get_coordination_numbers(material=material, cutoff=self.build_parameters.shadowing_radius) + ) + print("coordination numbers:", coordination_numbers) + return coordination_numbers