diff --git a/ipsuite/__init__.pyi b/ipsuite/__init__.pyi index 96a01394..7f0008ec 100644 --- a/ipsuite/__init__.pyi +++ b/ipsuite/__init__.pyi @@ -28,6 +28,18 @@ from .analysis import ( StressHistogram, ) +# Atom Selection +from .atom_selection import ( + ElementTypeSelection, + FixAtomsConstraint, + FixBondLengthConstraint, + FixBondLengthsConstraint, + LayerSelection, + RadialSelection, + SurfaceSelection, + ZPositionSelection, +) + # Base imports from .base import Flatten @@ -56,6 +68,8 @@ from .calculators import ( # Configuration Generation from .configuration_generation import ( + AddAdsorbate, + BuildSurface, MultiPackmol, Packmol, Smiles2Atoms, @@ -110,6 +124,9 @@ from .dynamics import ( # Geometry from .geometry import BarycenterMapping +# Interfaces +from .interfaces import AtomConstraint, AtomSelector + # Models from .models import ( GAP, @@ -151,8 +168,10 @@ __all__ = [ "ThresholdSelection", "FilterOutlier", # Configuration Generation - "Packmol", + "AddAdsorbate", + "BuildSurface", "MultiPackmol", + "Packmol", "Smiles2Atoms", "Smiles2Conformers", "Smiles2Gromacs", @@ -232,4 +251,16 @@ __all__ = [ # Calc "ApplyCalculator", "WrapModifier", + # Interfaces + "AtomConstraint", + "AtomSelector", + # Atom Selection + "ElementTypeSelection", + "FixAtomsConstraint", + "FixBondLengthConstraint", + "FixBondLengthsConstraint", + "LayerSelection", + "RadialSelection", + "SurfaceSelection", + "ZPositionSelection", ] diff --git a/ipsuite/atom_selection/__init__.py b/ipsuite/atom_selection/__init__.py new file mode 100644 index 00000000..31ca2904 --- /dev/null +++ b/ipsuite/atom_selection/__init__.py @@ -0,0 +1,21 @@ +"""Module for selecting atoms within individual ASE Atoms objects.""" + +from .constraints import ( + FixAtomsConstraint, +) +from .selections import ( + ElementTypeSelection, + LayerSelection, + RadialSelection, + SurfaceSelection, + ZPositionSelection, +) + +__all__ = [ + "ElementTypeSelection", + "LayerSelection", + "RadialSelection", + "SurfaceSelection", + "ZPositionSelection", + "FixAtomsConstraint", +] diff --git a/ipsuite/atom_selection/constraints.py b/ipsuite/atom_selection/constraints.py new file mode 100644 index 00000000..43f72370 --- /dev/null +++ b/ipsuite/atom_selection/constraints.py @@ -0,0 +1,141 @@ +"""Constraint classes that use atom selections.""" + +import dataclasses + +import ase +from ase.constraints import FixAtoms, FixBondLength, FixBondLengths + +from ipsuite.interfaces import AtomConstraint, AtomSelector + + +@dataclasses.dataclass +class FixAtomsConstraint(AtomConstraint): + """Fix atoms in a selection during optimization or dynamics. + + Parameters + ---------- + selection : AtomSelector + Atom selection strategy that determines which atoms to fix. + """ + + selection: AtomSelector + + def get_constraint(self, atoms: ase.Atoms) -> FixAtoms: + """Get the FixAtoms constraint for the selected atoms. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to apply constraint to. + + Returns + ------- + FixAtoms + ASE constraint that fixes the selected atoms. + """ + indices = self.selection.select(atoms) + if not indices: + raise ValueError("No atoms selected for FixAtoms constraint") + return FixAtoms(indices=indices) + + +@dataclasses.dataclass +class FixBondLengthConstraint(AtomConstraint): + """Fix bond length between two selected atoms. + + Parameters + ---------- + atom1_selection : AtomSelector + Selection for the first atom. Must select exactly one atom. + atom2_selection : AtomSelector + Selection for the second atom. Must select exactly one atom. + """ + + atom1_selection: AtomSelector + atom2_selection: AtomSelector + + def get_constraint(self, atoms: ase.Atoms) -> FixBondLength: + """Get the FixBondLength constraint for the selected atom pair. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to apply constraint to. + + Returns + ------- + FixBondLength + ASE constraint that fixes the bond length between selected atoms. + + Raises + ------ + ValueError + If either selection doesn't select exactly one atom. + """ + indices1 = self.atom1_selection.select(atoms) + indices2 = self.atom2_selection.select(atoms) + + if len(indices1) != 1: + raise ValueError( + f"atom1_selection must select exactly 1 atom, got {len(indices1)}" + ) + if len(indices2) != 1: + raise ValueError( + f"atom2_selection must select exactly 1 atom, got {len(indices2)}" + ) + + return FixBondLength(indices1[0], indices2[0]) + + +@dataclasses.dataclass +class FixBondLengthsConstraint(AtomConstraint): + """Fix multiple bond lengths using atom selections. + + Parameters + ---------- + bond_pairs : list[tuple[AtomSelector, AtomSelector]] + List of tuples, each containing two AtomSelector objects that should + each select exactly one atom to form a bond pair. + """ + + bond_pairs: list[tuple[AtomSelector, AtomSelector]] + + def get_constraint(self, atoms: ase.Atoms) -> FixBondLengths: + """Get the FixBondLengths constraint for multiple selected atom pairs. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to apply constraint to. + + Returns + ------- + FixBondLengths + ASE constraint that fixes bond lengths between selected atom pairs. + + Raises + ------ + ValueError + If any selection doesn't select exactly one atom. + """ + pairs = [] + + for i, (sel1, sel2) in enumerate(self.bond_pairs): + indices1 = sel1.select(atoms) + indices2 = sel2.select(atoms) + + if len(indices1) != 1: + raise ValueError( + f"Bond pair {i}, atom1 selection must select exactly 1 atom, got {len(indices1)}" + ) + if len(indices2) != 1: + raise ValueError( + f"Bond pair {i}, atom2 selection must select exactly 1 atom, got {len(indices2)}" + ) + + pairs.append([indices1[0], indices2[0]]) + + if not pairs: + raise ValueError("No bond pairs provided for FixBondLengths constraint") + + return FixBondLengths(pairs) diff --git a/ipsuite/atom_selection/selections.py b/ipsuite/atom_selection/selections.py new file mode 100644 index 00000000..0e0ca58c --- /dev/null +++ b/ipsuite/atom_selection/selections.py @@ -0,0 +1,258 @@ +"""Concrete atom selection implementations.""" + +import dataclasses +import typing as t + +import ase +import numpy as np + +from ipsuite.interfaces import AtomSelector + + +@dataclasses.dataclass +class ElementTypeSelection(AtomSelector): + """Select atoms based on element types. + + Parameters + ---------- + elements : list[str] + List of element symbols to select (e.g., ['H', 'O', 'C']). + """ + + elements: list[str] + + def select(self, atoms: ase.Atoms) -> list[int]: + """Select atoms based on the provided element types. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to select from. + + Returns + ------- + list[int] + Indices of atoms with the specified elements. + """ + symbols = atoms.get_chemical_symbols() + return [i for i, symbol in enumerate(symbols) if symbol in self.elements] + + +@dataclasses.dataclass +class ZPositionSelection(AtomSelector): + """Select atoms based on their Z-coordinate position. + + Parameters + ---------- + z_min : float | None, optional + Minimum Z-coordinate (inclusive). If None, no lower bound. + z_max : float | None, optional + Maximum Z-coordinate (inclusive). If None, no upper bound. + """ + + z_min: float | None = None + z_max: float | None = None + + def select(self, atoms: ase.Atoms) -> list[int]: + """Select atoms within the Z-coordinate range. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to select from. + + Returns + ------- + list[int] + Indices of atoms within the specified Z range. + """ + positions = atoms.get_positions() + z_coords = positions[:, 2] + + mask = np.ones(len(atoms), dtype=bool) + + if self.z_min is not None: + mask &= z_coords >= self.z_min + if self.z_max is not None: + mask &= z_coords <= self.z_max + + return np.where(mask)[0].tolist() + + +@dataclasses.dataclass +class RadialSelection(AtomSelector): + """Select atoms within a radial distance from a center point. + + Parameters + ---------- + center : tuple[float, float, float] | str + Center point for radial selection. Can be coordinates (x, y, z) + or 'com' for center of mass, 'geometric' for geometric center. + radius : float + Selection radius in Angstroms. + """ + + center: tuple[float, float, float] | t.Literal["com", "geometric"] + radius: float + + def _get_center_position(self, atoms: ase.Atoms) -> np.ndarray: + """Get the center position for radial selection. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure. + + Returns + ------- + np.ndarray + Center position coordinates. + """ + if isinstance(self.center, tuple): + return np.array(self.center) + elif self.center == "com": + return atoms.get_center_of_mass() + elif self.center == "geometric": + return atoms.get_positions().mean(axis=0) + else: + raise ValueError(f"Unknown center type: {self.center}") + + def select(self, atoms: ase.Atoms) -> list[int]: + """Select atoms within radial distance from center. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to select from. + + Returns + ------- + list[int] + Indices of atoms within the specified radius. + """ + center_pos = self._get_center_position(atoms) + positions = atoms.get_positions() + + distances = np.linalg.norm(positions - center_pos, axis=1) + mask = distances <= self.radius + + return np.where(mask)[0].tolist() + + +@dataclasses.dataclass +class LayerSelection(AtomSelector): + """Select atoms from specific layers in a slab. + + Parameters + ---------- + layer_indices : list[int] + Layer indices to select (0 = bottom layer, -1 = top layer). + tolerance : float, optional + Z-coordinate tolerance for grouping atoms into layers, by default 0.5. + """ + + layer_indices: list[int] + tolerance: float = 0.5 + + def _identify_layers(self, atoms: ase.Atoms) -> list[list[int]]: + """Identify atomic layers based on Z-coordinates. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure. + + Returns + ------- + list[list[int]] + List of lists, where each inner list contains atom indices + for atoms in that layer (sorted by Z-coordinate). + """ + positions = atoms.get_positions() + z_coords = positions[:, 2] + + # Sort atoms by Z-coordinate + sorted_indices = np.argsort(z_coords) + sorted_z = z_coords[sorted_indices] + + # Group atoms into layers based on tolerance + layers = [] + current_layer = [sorted_indices[0]] + current_z = sorted_z[0] + + for i in range(1, len(sorted_indices)): + if sorted_z[i] - current_z <= self.tolerance: + current_layer.append(sorted_indices[i]) + else: + layers.append(current_layer) + current_layer = [sorted_indices[i]] + current_z = sorted_z[i] + + layers.append(current_layer) # Add final layer + return layers + + def select(self, atoms: ase.Atoms) -> list[int]: + """Select atoms from specified layers. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to select from. + + Returns + ------- + list[int] + Indices of atoms in the specified layers. + """ + layers = self._identify_layers(atoms) + selected_indices = [] + + for layer_idx in self.layer_indices: + if -len(layers) <= layer_idx < len(layers): + selected_indices.extend(layers[layer_idx]) + + return sorted(selected_indices) + + +@dataclasses.dataclass +class SurfaceSelection(AtomSelector): + """Select surface atoms (atoms with fewer neighbors than bulk). + + Parameters + ---------- + cutoff : float, optional + Cutoff distance for neighbor counting, by default 3.0. + min_neighbors : int, optional + Minimum number of neighbors for an atom to be considered bulk, + by default 8. Atoms with fewer neighbors are considered surface atoms. + """ + + cutoff: float = 3.0 + min_neighbors: int = 8 + + def select(self, atoms: ase.Atoms) -> list[int]: + """Select surface atoms based on neighbor count. + + Parameters + ---------- + atoms : ase.Atoms + Atomic structure to select from. + + Returns + ------- + list[int] + Indices of surface atoms. + """ + positions = atoms.get_positions() + surface_indices = [] + + for i in range(len(atoms)): + # Calculate distances to all other atoms + distances = np.linalg.norm(positions - positions[i], axis=1) + # Count neighbors within cutoff (excluding self) + neighbor_count = np.sum((distances > 0) & (distances <= self.cutoff)) + + if neighbor_count < self.min_neighbors: + surface_indices.append(i) + + return surface_indices diff --git a/ipsuite/configuration_generation/__init__.pyi b/ipsuite/configuration_generation/__init__.pyi index 1c1ace55..8df732ac 100644 --- a/ipsuite/configuration_generation/__init__.pyi +++ b/ipsuite/configuration_generation/__init__.pyi @@ -3,6 +3,7 @@ from .gmx import Smiles2Gromacs from .packmol import MultiPackmol, Packmol from .smiles_to_atoms import Smiles2Atoms, Smiles2Conformers +from .surface_builder import AddAdsorbate, BuildSurface __all__ = [ "Smiles2Atoms", @@ -10,4 +11,6 @@ __all__ = [ "Smiles2Conformers", "MultiPackmol", "Smiles2Gromacs", + "BuildSurface", + "AddAdsorbate", ] diff --git a/ipsuite/configuration_generation/surface_builder.py b/ipsuite/configuration_generation/surface_builder.py new file mode 100644 index 00000000..a2a6204d --- /dev/null +++ b/ipsuite/configuration_generation/surface_builder.py @@ -0,0 +1,377 @@ +import typing as t +from pathlib import Path + +import ase +import h5py +import numpy as np +import znh5md +import zntrack +from ase.build import add_adsorbate, bulk, surface + +from ipsuite import base + + +class BuildSurface(base.IPSNode): + """Build crystal surfaces using ase.build.surface functionality. + + This node creates crystal surfaces from bulk structures using Miller indices + to define the surface orientation and adds vacuum layers for surface calculations. + + Parameters + ---------- + lattice : str + Chemical symbol for the bulk lattice (e.g., 'Au', 'Pt', 'Cu'). + indices : tuple[int, int, int] + Miller indices defining the surface orientation (e.g., (1, 1, 1), (1, 0, 0)). + layers : int + Number of equivalent atomic layers in the slab. + vacuum : float, optional + Thickness of vacuum layer in Angstroms, by default 10.0. + lattice_constant : float | None, optional + Custom lattice constant in Angstroms, by default None (uses ASE default). + crystal_structure : str | None, optional + Crystal structure type ('fcc', 'bcc', 'hcp', 'diamond', 'zincblende'), + by default None (uses ASE default). + + Attributes + ---------- + frames : list[ase.Atoms] + List containing the generated surface structure. + """ + + # Surface parameters + lattice: str = zntrack.params() + indices: tuple[int, int, int] = zntrack.params() + layers: int = zntrack.params() + + # Optional parameters + vacuum: float = zntrack.params(10.0) + lattice_constant: float | None = zntrack.params(None) + crystal_structure: t.Literal["fcc", "bcc", "hcp", "diamond", "zincblende"] | None = ( + zntrack.params(None) + ) + + frames_path: Path = zntrack.outs_path(zntrack.nwd / "frames.h5") + + def run(self) -> None: + """Generate the surface structure and save to frames file.""" + # Create bulk structure + if self.crystal_structure and self.lattice_constant: + # Create bulk structure with custom parameters + bulk_atoms = bulk( + self.lattice, self.crystal_structure, a=self.lattice_constant, cubic=True + ) + else: + # Use default bulk structure + bulk_atoms = self.lattice + + # Create surface + surface_atoms = surface( + lattice=bulk_atoms, indices=self.indices, layers=self.layers + ) + + # Add vacuum + surface_atoms.center(vacuum=self.vacuum, axis=2) + + # Save to frames file + io = znh5md.IO(filename=self.frames_path) + io.append(surface_atoms) + + @property + def frames(self) -> list[ase.Atoms]: + """Get the generated surface structure. + + Returns + ------- + list[ase.Atoms] + List containing the surface structure. + """ + with self.state.fs.open(self.frames_path, "rb") as f: + with h5py.File(f) as file: + return znh5md.IO(file_handle=file)[:] + + +class AddAdsorbate(base.IPSNode): + """Add adsorbate molecules to a surface slab with collision detection. + + This node places adsorbate molecules on a surface at specified heights with + automatic positioning to avoid collisions between adsorbates. + + Parameters + ---------- + slab : list[ase.Atoms] + List containing surface slab structures. + slab_idx : int, optional + Index to select which slab to use from the slab list. Defaults to -1 (last). + data : list[list[ase.Atoms]] + List of lists of adsorbate molecules. Each inner list represents different + conformations/structures of the same adsorbate species. + data_index : list[int | None] | None, optional + Indices to select from each adsorbate list in data. If None (default), + the last element (-1 index) is selected from each list. If provided, + must have same length as data list. + height : list[float] + Heights (in Angstroms) at which to place each adsorbate above the surface. + Length must match the number of adsorbates to be placed. + excluded_radius : list[float] + Exclusion radius (in Angstroms) around each adsorbate position where + no other adsorbate can be placed. Length must match height list. + + Attributes + ---------- + frames : list[ase.Atoms] + List containing the surface with adsorbed molecules. + """ + + slab: list[ase.Atoms] = zntrack.deps() + slab_idx: int = zntrack.params(-1) + data: list[list[ase.Atoms]] = zntrack.deps() + data_index: list[int | None] | None = zntrack.params(None) + height: list[float] = zntrack.params() + excluded_radius: list[float] = zntrack.params() + + frames_path: Path = zntrack.outs_path(zntrack.nwd / "frames.h5") + + def _get_com_position(self, atoms: ase.Atoms) -> np.ndarray: + """Calculate center of mass for adsorbate positioning. + + Parameters + ---------- + atoms : ase.Atoms + Adsorbate molecule. + + Returns + ------- + np.ndarray + Center of mass coordinates [x, y, z]. + """ + masses = atoms.get_masses() + positions = atoms.get_positions() + return np.average(positions, weights=masses, axis=0) + + def _get_surface_center(self, slab: ase.Atoms) -> tuple[float, float]: + """Get the center position of the surface in x-y plane. + + Parameters + ---------- + slab : ase.Atoms + Surface slab structure. + + Returns + ------- + tuple[float, float] + Center coordinates (x, y) of the surface. + """ + cell = slab.get_cell() + return (cell[0, 0] / 2, cell[1, 1] / 2) + + def _check_collision( + self, + new_pos: tuple[float, float], + existing_positions: list[tuple[float, float]], + new_radius: float, + existing_radii: list[float], + ) -> bool: + """Check if new adsorbate position collides with existing ones. + + Parameters + ---------- + new_pos : tuple[float, float] + Proposed (x, y) position for new adsorbate. + existing_positions : list[tuple[float, float]] + List of (x, y) positions of already placed adsorbates. + new_radius : float + Exclusion radius for the new adsorbate. + existing_radii : list[float] + Exclusion radii for existing adsorbates. + + Returns + ------- + bool + True if collision detected, False otherwise. + """ + for i, (ex_x, ex_y) in enumerate(existing_positions): + distance = np.sqrt((new_pos[0] - ex_x) ** 2 + (new_pos[1] - ex_y) ** 2) + min_distance = max(new_radius, existing_radii[i]) + if distance < min_distance: + return True + return False + + def _find_valid_position( + self, + slab: ase.Atoms, + radius: float, + existing_positions: list[tuple[float, float]], + existing_radii: list[float], + ) -> tuple[float, float]: + """Find a valid position for adsorbate that doesn't collide with existing ones. + + Parameters + ---------- + slab : ase.Atoms + Surface slab structure. + radius : float + Exclusion radius for the new adsorbate. + existing_positions : list[tuple[float, float]] + Positions of existing adsorbates. + existing_radii : list[float] + Exclusion radii of existing adsorbates. + + Returns + ------- + tuple[float, float] + Valid (x, y) position for the adsorbate. + + Raises + ------ + ValueError + If no valid position can be found. + """ + cell = slab.get_cell() + max_x, max_y = cell[0, 0], cell[1, 1] + + # Try different positions with increasing distance from existing adsorbates + for attempt in range(1000): # Limit attempts to prevent infinite loops + # Use systematic grid search with some randomness + if attempt < 100: + # Systematic grid search + grid_size = int(np.sqrt(attempt + 1)) + i, j = attempt % grid_size, attempt // grid_size + x = (i + 0.5) * max_x / (grid_size + 1) + y = (j + 0.5) * max_y / (grid_size + 1) + else: + # Random positions for remaining attempts + x = np.random.uniform(radius, max_x - radius) + y = np.random.uniform(radius, max_y - radius) + + if not self._check_collision( + (x, y), existing_positions, radius, existing_radii + ): + return (x, y) + + raise ValueError( + f"Could not find valid position for adsorbate with radius {radius}. " + "Consider reducing excluded_radius values or using a larger surface." + ) + + def _get_selected_adsorbates(self) -> list[ase.Atoms]: + """Extract the selected adsorbates from the nested data structure. + + Returns + ------- + list[ase.Atoms] + List of selected adsorbate molecules. + + Raises + ------ + ValueError + If data_index is provided but has wrong length. + """ + if self.data_index is None: + # Use -1 index (last element) from each list + return [adsorbate_list[-1] for adsorbate_list in self.data] + + if len(self.data_index) != len(self.data): + raise ValueError( + f"Length of data_index ({len(self.data_index)}) must match " + f"length of data ({len(self.data)})" + ) + + selected_adsorbates = [] + for i, (adsorbate_list, index) in enumerate(zip(self.data, self.data_index)): + if index is None: + # Use -1 index for this specific adsorbate + selected_adsorbates.append(adsorbate_list[-1]) + else: + # Use specified index + try: + selected_adsorbates.append(adsorbate_list[index]) + except IndexError: + raise ValueError( + f"Index {index} is out of bounds for adsorbate list {i} " + f"with length {len(adsorbate_list)}" + ) + + return selected_adsorbates + + def run(self) -> None: + """Add adsorbates to the surface slab.""" + if len(self.height) != len(self.excluded_radius): + raise ValueError("Length of height and excluded_radius lists must match") + + # Get selected adsorbates from nested structure + adsorbates = self._get_selected_adsorbates() + + n_adsorbates = min(len(adsorbates), len(self.height)) + if n_adsorbates == 0: + raise ValueError("No adsorbates to place") + + # Start with the selected surface slab + slab = self.slab[self.slab_idx].copy() + + # Track positions and radii of placed adsorbates + placed_positions = [] + placed_radii = [] + + for i in range(n_adsorbates): + adsorbate = adsorbates[i] + height = self.height[i] + radius = self.excluded_radius[i] + + if i == 0: + # Place first adsorbate at surface center + position = self._get_surface_center(slab) + else: + # Find valid position for subsequent adsorbates + position = self._find_valid_position( + slab, radius, placed_positions, placed_radii + ) + + # Add adsorbate to slab + # Use mol_index=0 instead of None to avoid array broadcasting issues + add_adsorbate( + slab=slab, + adsorbate=adsorbate, + height=height, + position=position, + offset=None, + mol_index=0, # Use first atom as reference instead of COM + ) + + # Clean up info dict to ensure JSON serializability + if "adsorbate_info" in slab.info: + # Convert numpy int64 to regular int for JSON compatibility + adsorbate_info = slab.info["adsorbate_info"] + if isinstance(adsorbate_info, dict): + for key, value in adsorbate_info.items(): + if ( + hasattr(value, "item") and value.size == 1 + ): # numpy scalar with single element + adsorbate_info[key] = value.item() + elif isinstance(value, np.integer): + adsorbate_info[key] = int(value) + elif isinstance(value, np.ndarray): + adsorbate_info[key] = ( + value.tolist() + ) # Convert arrays to lists + + # Track this adsorbate's position and radius + placed_positions.append(position) + placed_radii.append(radius) + + # Save result + io = znh5md.IO(filename=self.frames_path) + io.append(slab) + + @property + def frames(self) -> list[ase.Atoms]: + """Get the surface with adsorbed molecules. + + Returns + ------- + list[ase.Atoms] + List containing the surface with adsorbates. + """ + with self.state.fs.open(self.frames_path, "rb") as f: + with h5py.File(f) as file: + return znh5md.IO(file_handle=file)[:] diff --git a/ipsuite/interfaces.py b/ipsuite/interfaces.py index eefca0b9..4aabec21 100644 --- a/ipsuite/interfaces.py +++ b/ipsuite/interfaces.py @@ -3,6 +3,7 @@ import ase from ase.calculators.calculator import Calculator +from ase.constraints import FixConstraint from ase.optimize.optimize import Dynamics T = TypeVar("T", covariant=True) @@ -61,6 +62,52 @@ class ProcessAtoms(Protocol): frames: list[ase.Atoms] +class AtomSelector(Protocol): + """Protocol for selecting atoms within a single ASE Atoms object. + + This interface defines the contract for selecting atoms based on various + criteria within an individual frame/structure. + """ + + def select(self, atoms: ase.Atoms) -> list[int]: + """Select atoms based on the implemented criteria. + + Parameters + ---------- + atoms : ase.Atoms + The atomic structure to select from. + + Returns + ------- + list[int] + List of atom indices that match the selection criteria. + """ + ... + + +class AtomConstraint(Protocol): + """Protocol for applying constraints to selected atoms. + + This interface defines how to apply ASE constraints to atoms selected + by AtomSelector instances. + """ + + def get_constraint(self, atoms: ase.Atoms) -> FixConstraint: + """Get the ASE constraint object for the selected atoms. + + Parameters + ---------- + atoms : ase.Atoms + The atomic structure to apply constraints to. + + Returns + ------- + FixConstraint + ASE constraint object (e.g., FixAtoms, FixBondLengths, etc.) + """ + ... + + # Collection of complex type hints ATOMS_LST = list[ase.Atoms] UNION_ATOMS_OR_ATOMS_LST = Union[ATOMS_LST, List[ATOMS_LST]] @@ -73,6 +120,8 @@ class ProcessAtoms(Protocol): "HasAtoms", "HasSelectedConfigurations", "ProcessAtoms", + "AtomSelector", + "AtomConstraint", "ATOMS_LST", "UNION_ATOMS_OR_ATOMS_LST", "HasOrIsAtoms", @@ -81,4 +130,4 @@ class ProcessAtoms(Protocol): def interfaces() -> dict[str, list[str]]: """Return a dictionary of available interfaces.""" - return {"ipsuite.interfaces": __all__} + return {"ipsuite": __all__} diff --git a/pyproject.toml b/pyproject.toml index 94169dd4..374e1f33 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ dependencies = [ "seaborn>=0.13.2", "uncertainty-toolbox>=0.1.1", "znh5md>=0.4.8", - "zntrack>=0.8.7", + "zntrack>=0.8.9", ] diff --git a/tests/integration/test_atom_selection_integration.py b/tests/integration/test_atom_selection_integration.py new file mode 100644 index 00000000..5aedbea3 --- /dev/null +++ b/tests/integration/test_atom_selection_integration.py @@ -0,0 +1,140 @@ +"""Integration tests for atom selection with ASEGeoOpt.""" + +import ipsuite as ips + + +def test_asegeopt_with_atom_selection_constraints(proj_path): + """Test ASEGeoOpt using atom selection constraints.""" + with ips.Project() as proj: + # Create a small surface + surface = ips.BuildSurface( + lattice="Cu", + indices=(1, 1, 1), + layers=3, + vacuum=10.0 + ) + + # Create constraint to fix bottom layer atoms during relaxation + bottom_layer_selector = ips.LayerSelection(layer_indices=[0], tolerance=0.5) + fix_bottom_constraint = ips.FixAtomsConstraint(selection=bottom_layer_selector) + + # Set up MACE-MP model for forces + model = ips.MACEMPModel() + + # Run geometry optimization with constraint + relaxed_surface = ips.ASEGeoOpt( + data=surface.frames, + model=model, + constraints=[fix_bottom_constraint], # Use new constraint system + optimizer="BFGS", + run_kwargs={"fmax": 0.2, "steps": 20}, + sampling_rate=5, + maxstep=20 + ) + + proj.repro() + + # Verify surface was created + initial_surface = surface.frames[0] + assert len(initial_surface) > 0 + assert all(symbol == "Cu" for symbol in initial_surface.get_chemical_symbols()) + + # Verify optimization trajectory was generated + trajectory = relaxed_surface.frames + assert len(trajectory) > 0 + + # Check that final structure has same composition + final_surface = trajectory[-1] + assert len(final_surface) == len(initial_surface) + assert all(symbol == "Cu" for symbol in final_surface.get_chemical_symbols()) + + +def test_element_selection_constraint(proj_path): + """Test constraint using element-based selection.""" + with ips.Project() as proj: + # Create surface with adsorbate + surface = ips.BuildSurface( + lattice="Pt", + indices=(1, 1, 1), + layers=3, + vacuum=12.0 + ) + + water = ips.Smiles2Atoms(smiles="O") # H2O + + adsorbed_system = ips.AddAdsorbate( + slab=surface.frames, + data=[water.frames], + height=[2.0], + excluded_radius=[3.0] + ) + + # Create constraint to fix only Pt atoms (substrate) + pt_selector = ips.ElementTypeSelection(elements=['Pt']) + fix_substrate = ips.FixAtomsConstraint(selection=pt_selector) + + model = ips.MACEMPModel() + + # Optimize only the adsorbate, keeping substrate fixed + optimized = ips.ASEGeoOpt( + data=adsorbed_system.frames, + model=model, + constraints=[fix_substrate], + optimizer="FIRE", + run_kwargs={"fmax": 0.3, "steps": 15}, + maxstep=15 + ) + + proj.repro() + + # Verify the system was set up correctly + initial_system = adsorbed_system.frames[0] + symbols = initial_system.get_chemical_symbols() + + # Should have Pt + water atoms + assert 'Pt' in symbols + assert 'H' in symbols + assert 'O' in symbols + + # Verify optimization completed + trajectory = optimized.frames + assert len(trajectory) > 0 + + +def test_multiple_selection_types(proj_path): + """Test using multiple different selection types together.""" + with ips.Project() as proj: + # Create larger surface for more complex selections + surface = ips.BuildSurface( + lattice="Au", + indices=(1, 1, 1), + layers=4, + vacuum=15.0 + ) + + # Create multiple constraint types + # 1. Fix bottom 2 layers using layer selection + bottom_layers = ips.LayerSelection(layer_indices=[0, 1], tolerance=0.5) + fix_bottom = ips.FixAtomsConstraint(selection=bottom_layers) + + # 2. Could add more constraints here for complex scenarios + + model = ips.MACEMPModel() + + optimized = ips.ASEGeoOpt( + data=surface.frames, + model=model, + constraints=[fix_bottom], # Can extend with more constraints + optimizer="FIRE", + run_kwargs={"fmax": 0.2, "steps": 10}, + maxstep=10 + ) + + proj.repro() + + # Verify it worked + trajectory = optimized.frames + assert len(trajectory) > 0 + + final_atoms = trajectory[-1] + assert all(symbol == "Au" for symbol in final_atoms.get_chemical_symbols()) diff --git a/tests/integration/test_surface_builder.py b/tests/integration/test_surface_builder.py new file mode 100644 index 00000000..df46d3f7 --- /dev/null +++ b/tests/integration/test_surface_builder.py @@ -0,0 +1,370 @@ +import ipsuite as ips +import pytest +import numpy as np + + +def test_BuildSurface_basic(proj_path): + """Test basic BuildSurface functionality with default parameters.""" + with ips.Project() as proj: + surface = ips.BuildSurface( + lattice="Au", + indices=(1, 1, 1), + layers=5 + ) + + proj.repro() + + assert len(surface.frames) == 1 + atoms = surface.frames[0] + + # Check that we have a gold surface + assert all(symbol == "Au" for symbol in atoms.get_chemical_symbols()) + + # Check that vacuum was added (cell should be larger in z-direction) + cell = atoms.get_cell() + assert cell[2, 2] > 10.0 # Should have vacuum space + + +def test_BuildSurface_custom_parameters(proj_path): + """Test BuildSurface with custom lattice constant and crystal structure.""" + with ips.Project() as proj: + surface = ips.BuildSurface( + lattice="Cu", + indices=(1, 0, 0), + layers=3, + vacuum=15.0, + lattice_constant=3.6, + crystal_structure="fcc" + ) + + proj.repro() + + assert len(surface.frames) == 1 + atoms = surface.frames[0] + + # Check that we have a copper surface + assert all(symbol == "Cu" for symbol in atoms.get_chemical_symbols()) + + # Check that custom vacuum was applied + cell = atoms.get_cell() + assert cell[2, 2] > 15.0 # Should have at least 15 Å vacuum + + +def test_BuildSurface_different_indices(proj_path): + """Test BuildSurface with different Miller indices.""" + surfaces = [] + + with ips.Project() as proj: + # Test different surface orientations + surface_100 = ips.BuildSurface(lattice="Pt", indices=(1, 0, 0), layers=4) + surface_110 = ips.BuildSurface(lattice="Pt", indices=(1, 1, 0), layers=4) + surface_111 = ips.BuildSurface(lattice="Pt", indices=(1, 1, 1), layers=4) + surfaces = [surface_100, surface_110, surface_111] + + proj.repro() + + # All should generate valid surfaces + for surface in surfaces: + assert len(surface.frames) == 1 + atoms = surface.frames[0] + assert all(symbol == "Pt" for symbol in atoms.get_chemical_symbols()) + assert len(atoms) > 0 + + +def test_BuildSurface_with_mace_relaxation(proj_path): + """Test BuildSurface followed by MACE-MP relaxation.""" + with ips.Project() as proj: + # Build surface + surface = ips.BuildSurface( + lattice="Cu", + indices=(1, 1, 1), + layers=3, + vacuum=12.0 + ) + + # Set up MACE-MP model + model = ips.MACEMPModel() + + # Relax the surface using ASEGeoOpt + relaxed_surface = ips.ASEGeoOpt( + data=surface.frames, + model=model, + optimizer="FIRE", + run_kwargs={"fmax": 0.1, "steps": 50}, + sampling_rate=10, + maxstep=50 + ) + + proj.repro() + + # Verify the surface was built correctly + initial_atoms = surface.frames[0] + assert all(symbol == "Cu" for symbol in initial_atoms.get_chemical_symbols()) + assert len(initial_atoms) > 0 + + # Verify relaxation trajectory was generated + trajectory = relaxed_surface.frames + assert len(trajectory) > 0 + + # Check that final structure still has correct composition + final_atoms = trajectory[-1] + assert all(symbol == "Cu" for symbol in final_atoms.get_chemical_symbols()) + assert len(final_atoms) == len(initial_atoms) + + +def test_AddAdsorbate_basic(proj_path): + """Test basic AddAdsorbate functionality with single adsorbate.""" + with ips.Project() as proj: + # Create surface + surface = ips.BuildSurface( + lattice="Pt", + indices=(1, 1, 1), + layers=4, + vacuum=15.0 + ) + + # Create adsorbate molecules + water = ips.Smiles2Atoms(smiles="O") # H2O + + # Add adsorbate to surface (nested structure) + adsorbed_system = ips.AddAdsorbate( + slab=surface.frames, + data=[water.frames], # Wrap in list for nested structure + height=[2.0], + excluded_radius=[3.0] + ) + + proj.repro() + + # Verify surface was created + surface_atoms = surface.frames[0] + assert all(symbol == "Pt" for symbol in surface_atoms.get_chemical_symbols()) + + # Verify adsorbate was added + final_system = adsorbed_system.frames[0] + symbols = final_system.get_chemical_symbols() + + # Should have Pt atoms + water atoms + pt_count = symbols.count("Pt") + o_count = symbols.count("O") + h_count = symbols.count("H") + + assert pt_count > 0 # Original surface atoms + assert o_count == 1 # One oxygen from water + assert h_count == 2 # Two hydrogens from water + + # Total atoms should be surface + water + assert len(final_system) == len(surface_atoms) + 3 # Pt atoms + H2O + + +def test_AddAdsorbate_multiple_adsorbates(proj_path): + """Test AddAdsorbate with multiple adsorbates and collision detection.""" + with ips.Project() as proj: + # Create surface + surface = ips.BuildSurface( + lattice="Cu", + indices=(1, 0, 0), + layers=3, + vacuum=12.0 + ) + + # Create different adsorbate molecules + co = ips.Smiles2Atoms(smiles="[C-]#[O+]") # CO + h2 = ips.Smiles2Atoms(smiles="[H][H]") # H2 + + # Add multiple adsorbates with different heights and exclusion radii + adsorbed_system = ips.AddAdsorbate( + slab=surface.frames, + data=[co.frames, h2.frames], # Nested structure + height=[1.5, 3.0], + excluded_radius=[2.5, 2.0] + ) + + proj.repro() + + # Verify multiple adsorbates were added + final_system = adsorbed_system.frames[0] + symbols = final_system.get_chemical_symbols() + + cu_count = symbols.count("Cu") + c_count = symbols.count("C") + o_count = symbols.count("O") + h_count = symbols.count("H") + + assert cu_count > 0 # Original surface + assert c_count == 1 # One carbon from CO + assert o_count == 1 # One oxygen from CO + assert h_count == 2 # Two hydrogens from H2 + + # Verify positioning: CO should be closer to surface than H2 (1.5 vs 3.0 Å) + positions = final_system.get_positions() + z_coords = positions[:, 2] + + # Find CO and H2 z-coordinates (approximate check) + surface_top_z = np.max([z for i, z in enumerate(z_coords) + if symbols[i] == "Cu"]) + + # Both adsorbates should be above the surface + co_z = z_coords[symbols.index("C")] # CO carbon position + h2_z = np.mean([z_coords[i] for i, s in enumerate(symbols) if s == "H"]) # H2 average + + assert co_z > surface_top_z + assert h2_z > surface_top_z + assert co_z < h2_z # CO should be closer to surface + + +def test_AddAdsorbate_collision_detection_error(proj_path): + """Test that AddAdsorbate raises error when adsorbates can't be placed without collision.""" + with ips.Project() as proj: + # Create small surface + surface = ips.BuildSurface( + lattice="Au", + indices=(1, 1, 1), + layers=2, + vacuum=10.0, + lattice_constant=3.0 # Small lattice for small surface + ) + + # Try to place too many large adsorbates + water1 = ips.Smiles2Atoms(smiles="O") + water2 = ips.Smiles2Atoms(smiles="O") + water3 = ips.Smiles2Atoms(smiles="O") + + # Very large exclusion radii that should cause collision + adsorbed_system = ips.AddAdsorbate( + slab=surface.frames, + data=[water1.frames, water2.frames, water3.frames], # Nested structure + height=[2.0, 2.0, 2.0], + excluded_radius=[10.0, 10.0, 10.0] # Unrealistically large radii + ) + + # This should raise an error during repro due to collision detection + with pytest.raises(ValueError, match="Could not find valid position"): + proj.repro() + + +def test_AddAdsorbate_data_index_functionality(proj_path): + """Test AddAdsorbate with data_index parameter for selecting specific conformers.""" + with ips.Project() as proj: + # Create surface + surface = ips.BuildSurface( + lattice="Pt", + indices=(1, 1, 1), + layers=3, + vacuum=12.0 + ) + + # Create multiple conformers of the same molecule + water_conformers = ips.Smiles2Conformers(smiles="O", numConfs=5) + co_molecule = ips.Smiles2Atoms(smiles="[C-]#[O+]") + + # Test with data_index=None (default, uses -1 index) + adsorbed_system_default = ips.AddAdsorbate( + slab=surface.frames, + data=[water_conformers.frames, co_molecule.frames], + data_index=None, # Use last conformer of water, last (only) CO + height=[2.0, 2.5], + excluded_radius=[2.0, 2.0] + ) + + # Test with specific data_index + adsorbed_system_custom = ips.AddAdsorbate( + slab=surface.frames, + data=[water_conformers.frames, co_molecule.frames], + data_index=[0, None], # Use first water conformer, default CO + height=[2.0, 2.5], + excluded_radius=[2.0, 2.0] + ) + + proj.repro() + + # Both should have same atom counts but potentially different structures + system_default = adsorbed_system_default.frames[0] + system_custom = adsorbed_system_custom.frames[0] + + # Same composition in both cases + assert len(system_default) == len(system_custom) + + symbols_default = system_default.get_chemical_symbols() + symbols_custom = system_custom.get_chemical_symbols() + + # Should have same element counts + for element in set(symbols_default + symbols_custom): + assert symbols_default.count(element) == symbols_custom.count(element) + + # Should have Pt + water + CO + pt_count = symbols_default.count("Pt") + h_count = symbols_default.count("H") + o_count = symbols_default.count("O") + c_count = symbols_default.count("C") + + assert pt_count > 0 # Surface atoms + assert h_count == 2 # Water hydrogens + assert o_count == 2 # Water oxygen + CO oxygen + assert c_count == 1 # CO carbon + + +def test_AddAdsorbate_slab_idx_functionality(proj_path): + """Test AddAdsorbate with slab_idx parameter for selecting specific slab.""" + with ips.Project() as proj: + # Create multiple surfaces + surface1 = ips.BuildSurface( + lattice="Au", + indices=(1, 1, 1), + layers=3, + vacuum=10.0 + ) + + surface2 = ips.BuildSurface( + lattice="Pt", + indices=(1, 0, 0), + layers=4, + vacuum=12.0 + ) + + # Combine surfaces into a list + combined_surfaces = surface1.frames + surface2.frames + + # Create adsorbate + water = ips.Smiles2Atoms(smiles="O") + + # Test with slab_idx=0 (first slab - Au) + adsorbed_au = ips.AddAdsorbate( + slab=combined_surfaces, + slab_idx=0, + data=[water.frames], + height=[2.0], + excluded_radius=[2.0] + ) + + # Test with slab_idx=-1 (default, last slab - Pt) + adsorbed_pt = ips.AddAdsorbate( + slab=combined_surfaces, + slab_idx=-1, + data=[water.frames], + height=[2.0], + excluded_radius=[2.0] + ) + + proj.repro() + + # Check that correct surfaces were used + au_system = adsorbed_au.frames[0] + pt_system = adsorbed_pt.frames[0] + + au_symbols = au_system.get_chemical_symbols() + pt_symbols = pt_system.get_chemical_symbols() + + # Au system should have Au atoms + assert "Au" in au_symbols + assert "Pt" not in au_symbols + + # Pt system should have Pt atoms + assert "Pt" in pt_symbols + assert "Au" not in pt_symbols + + # Both should have water + assert au_symbols.count("H") == 2 + assert au_symbols.count("O") == 1 + assert pt_symbols.count("H") == 2 + assert pt_symbols.count("O") == 1 diff --git a/tests/unit_tests/atom_selection/test_selections.py b/tests/unit_tests/atom_selection/test_selections.py new file mode 100644 index 00000000..56ac30e2 --- /dev/null +++ b/tests/unit_tests/atom_selection/test_selections.py @@ -0,0 +1,191 @@ +"""Tests for atom selection implementations.""" + +import numpy as np +import pytest +from ase import Atoms + +import ipsuite as ips + + +def test_element_type_selection(): + """Test ElementTypeSelection with different element types.""" + # Create test atoms with mixed elements + atoms = Atoms(['H', 'H', 'O', 'C', 'C', 'N'], + positions=np.random.rand(6, 3) * 10) + + # Test single element selection + h_selector = ips.ElementTypeSelection(elements=['H']) + h_indices = h_selector.select(atoms) + assert h_indices == [0, 1] + + # Test multiple element selection + multi_selector = ips.ElementTypeSelection(elements=['O', 'N']) + multi_indices = multi_selector.select(atoms) + assert set(multi_indices) == {2, 5} + + # Test no matches + empty_selector = ips.ElementTypeSelection(elements=['Pt']) + empty_indices = empty_selector.select(atoms) + assert empty_indices == [] + + +def test_z_position_selection(): + """Test ZPositionSelection with various Z-coordinate ranges.""" + # Create atoms at different Z positions + positions = np.array([ + [0, 0, 1.0], + [0, 0, 2.5], + [0, 0, 4.0], + [0, 0, 5.5], + [0, 0, 7.0] + ]) + atoms = Atoms(['H'] * 5, positions=positions) + + # Test Z range selection + mid_selector = ips.ZPositionSelection(z_min=2.0, z_max=5.0) + mid_indices = mid_selector.select(atoms) + assert set(mid_indices) == {1, 2} + + # Test only minimum + min_selector = ips.ZPositionSelection(z_min=5.0) + min_indices = min_selector.select(atoms) + assert set(min_indices) == {3, 4} + + # Test only maximum + max_selector = ips.ZPositionSelection(z_max=3.0) + max_indices = max_selector.select(atoms) + assert set(max_indices) == {0, 1} + + +def test_radial_selection(): + """Test RadialSelection with different center types.""" + # Create atoms in a simple pattern + positions = np.array([ + [0, 0, 0], # at origin + [1, 0, 0], # 1 unit away + [0, 2, 0], # 2 units away + [3, 0, 0], # 3 units away + [0, 0, 4] # 4 units away + ]) + atoms = Atoms(['C'] * 5, positions=positions) + + # Test radial selection from specific center + center_selector = ips.RadialSelection(center=(0, 0, 0), radius=2.5) + center_indices = center_selector.select(atoms) + assert set(center_indices) == {0, 1, 2} # Within 2.5 units of origin + + # Test center of mass (should be at average position) + com_selector = ips.RadialSelection(center='com', radius=1.0) + com_indices = com_selector.select(atoms) + # COM is at (0.8, 0.4, 0.8), so check which atoms are within 1.0 unit + com = atoms.get_center_of_mass() + expected = [] + for i, pos in enumerate(positions): + if np.linalg.norm(pos - com) <= 1.0: + expected.append(i) + assert set(com_indices) == set(expected) + + +def test_layer_selection(): + """Test LayerSelection for identifying atomic layers.""" + # Create a simple 3-layer slab + positions = np.array([ + # Bottom layer (z=0) + [0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], + # Middle layer (z=2) + [0, 0, 2], [1, 0, 2], [0, 1, 2], [1, 1, 2], + # Top layer (z=4) + [0, 0, 4], [1, 0, 4], [0, 1, 4], [1, 1, 4] + ]) + atoms = Atoms(['Pt'] * 12, positions=positions) + + # Test bottom layer selection + bottom_selector = ips.LayerSelection(layer_indices=[0]) + bottom_indices = bottom_selector.select(atoms) + assert set(bottom_indices) == {0, 1, 2, 3} + + # Test top layer selection + top_selector = ips.LayerSelection(layer_indices=[-1]) + top_indices = top_selector.select(atoms) + assert set(top_indices) == {8, 9, 10, 11} + + # Test multiple layers + multi_selector = ips.LayerSelection(layer_indices=[0, -1]) + multi_indices = multi_selector.select(atoms) + assert set(multi_indices) == {0, 1, 2, 3, 8, 9, 10, 11} + + +def test_surface_selection(): + """Test SurfaceSelection for identifying surface atoms.""" + # Create a simple bulk-like structure with surface atoms + # 2x2x2 cube with one atom removed to create a surface + positions = np.array([ + [0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], # bottom layer + [0, 0, 1], [1, 0, 1], [0, 1, 1], # top layer (missing one atom) + ]) + atoms = Atoms(['Cu'] * 7, positions=positions) + + # All atoms should be surface atoms due to low coordination + surface_selector = ips.SurfaceSelection(cutoff=1.5, min_neighbors=6) + surface_indices = surface_selector.select(atoms) + + # In this small cluster, all atoms have < 6 neighbors within 1.5 Å + assert len(surface_indices) == 7 # All atoms are surface atoms + + +def test_fix_atoms_constraint(): + """Test FixAtomsConstraint integration with element selection.""" + atoms = Atoms(['Pt', 'Pt', 'H', 'O'], + positions=np.random.rand(4, 3) * 5) + + # Create constraint to fix Pt atoms + pt_selector = ips.ElementTypeSelection(elements=['Pt']) + constraint = ips.FixAtomsConstraint(selection=pt_selector) + + # Get ASE constraint + ase_constraint = constraint.get_constraint(atoms) + + # Check that correct atoms are fixed + assert set(ase_constraint.index) == {0, 1} + + +def test_fix_atoms_constraint_empty_selection(): + """Test FixAtomsConstraint with empty selection raises error.""" + atoms = Atoms(['H', 'O'], positions=np.random.rand(2, 3) * 5) + + # Try to fix non-existent atoms + empty_selector = ips.ElementTypeSelection(elements=['Pt']) + constraint = ips.FixAtomsConstraint(selection=empty_selector) + + with pytest.raises(ValueError, match="No atoms selected"): + constraint.get_constraint(atoms) + + +def test_selection_combination_workflow(): + """Test realistic workflow combining surface building with atom selection.""" + # This would be an integration test showing the full workflow + # Create a surface, select atoms, and apply constraints + + # Create test surface-like structure + positions = np.array([ + # Bottom layer (bulk-like) + [0, 0, 0], [2, 0, 0], [1, 1.7, 0], + # Top layer (surface) + [0, 0, 2], [2, 0, 2], [1, 1.7, 2] + ]) + atoms = Atoms(['Pt'] * 6, positions=positions) + + # Select bottom layer atoms + bottom_selector = ips.LayerSelection(layer_indices=[0], tolerance=0.5) + bottom_indices = bottom_selector.select(atoms) + assert set(bottom_indices) == {0, 1, 2} + + # Create constraint to fix bottom layer + fix_constraint = ips.FixAtomsConstraint(selection=bottom_selector) + ase_constraint = fix_constraint.get_constraint(atoms) + + # Verify constraint fixes the right atoms + assert set(ase_constraint.index) == {0, 1, 2} + + # This constraint could now be used with ASEGeoOpt for surface relaxation + # where only the surface layer atoms are allowed to move diff --git a/uv.lock b/uv.lock index e577b1f4..cf7d3f68 100644 --- a/uv.lock +++ b/uv.lock @@ -891,7 +891,7 @@ wheels = [ [[package]] name = "dvc" -version = "3.59.1" +version = "3.61.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, @@ -937,9 +937,9 @@ dependencies = [ { name = "voluptuous" }, { name = "zc-lockfile" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ec/3b/fb2256d3637dcf42cf8351356a6edda15173f62fea852fc3b31b2a883736/dvc-3.59.1.tar.gz", hash = "sha256:76c9cadc97a458321845424ff91489e8a2c5c6eaf36e52a79d8609beaed78cd6", size = 652872, upload-time = "2025-02-15T11:14:45.715Z" } +sdist = { url = "https://files.pythonhosted.org/packages/72/0e/5552098e86ca6da066010952e983e5fa5f01642db2d256ca38f7a66b1786/dvc-3.61.0.tar.gz", hash = "sha256:313c92d25c19352b0ec4b77f2ed68666e15093735752417601cf6e88d7598dbe", size = 658712, upload-time = "2025-07-07T11:23:57.853Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/41/65/46cbbe4907a14d458be93e4f3fdc6b25011623aa705ef0d9bbce11b53509/dvc-3.59.1-py3-none-any.whl", hash = "sha256:445c66f5d8304cf868ecff7b6d9428bff7879cff9545f0ac5a896572e14e14dd", size = 457702, upload-time = "2025-02-15T11:14:43.216Z" }, + { url = "https://files.pythonhosted.org/packages/93/e2/30a3985732b3ebe2aad2e3a974b85c1c91af56f5c0a11bc545cca4bddc12/dvc-3.61.0-py3-none-any.whl", hash = "sha256:f78c36a39a585060320300a5e3c975ffee109a573d9ebc37bbe541e15a797cfc", size = 461036, upload-time = "2025-07-07T11:23:55.76Z" }, ] [[package]] @@ -1523,7 +1523,7 @@ requires-dist = [ { name = "torch-dftd", marker = "extra == 'd3'", specifier = ">=0.5.1" }, { name = "uncertainty-toolbox", specifier = ">=0.1.1" }, { name = "znh5md", specifier = ">=0.4.8" }, - { name = "zntrack", specifier = ">=0.8.7" }, + { name = "zntrack", specifier = ">=0.8.9" }, ] provides-extras = ["d3"] @@ -4489,7 +4489,7 @@ wheels = [ [[package]] name = "zntrack" -version = "0.8.7" +version = "0.8.9" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "dvc" }, @@ -4501,7 +4501,7 @@ dependencies = [ { name = "znflow" }, { name = "znjson" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b0/8a/a48a9beeb0817bf5a48153275559ed5ba2c2d65b892d000c278a3ce48653/zntrack-0.8.7.tar.gz", hash = "sha256:fc51eb7a96cf5d3baf24f413895748e7de1404ba931928267f5fcb2e140ab85b", size = 345873, upload-time = "2025-05-15T14:40:22.564Z" } +sdist = { url = "https://files.pythonhosted.org/packages/bb/f1/7ee60e2bfa3a5fb424459f0d8b818f9e48c4848ce848bd729c790c5c1b2c/zntrack-0.8.9.tar.gz", hash = "sha256:c7c8f7dd554fec761c5f120eef672479c252a1b4ce3749ce72f7a73e58f6ef54", size = 467387, upload-time = "2025-08-08T13:27:54.892Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b7/e3/8540c753dc6d30ca1cec0789b431921961017eadd80b5f85a6f32a02bd26/zntrack-0.8.7-py3-none-any.whl", hash = "sha256:2eb2e8f4c9ea7d066a02ab28597e7ca21f1983035aad73c9f624aaba34b68316", size = 61030, upload-time = "2025-05-15T14:40:21.498Z" }, + { url = "https://files.pythonhosted.org/packages/5a/81/08d56f60f6c31510dec674e888953b109169659d742df839470d3cf1e073/zntrack-0.8.9-py3-none-any.whl", hash = "sha256:afb20f042b1a4d32f693b06eb2ad6ce6273ab6d52a9601a5cc18c964d041471f", size = 68314, upload-time = "2025-08-08T13:27:53.465Z" }, ]