From 0df1683a958a584416424a16c35ae58612722163 Mon Sep 17 00:00:00 2001 From: Fabian Zills Date: Thu, 7 Aug 2025 14:10:16 +0200 Subject: [PATCH 1/6] add surface builder node --- ipsuite/configuration_generation/__init__.pyi | 2 + .../surface_builder.py | 93 +++++++++++++++++++ tests/integration/test_surface_builder.py | 69 ++++++++++++++ 3 files changed, 164 insertions(+) create mode 100644 ipsuite/configuration_generation/surface_builder.py create mode 100644 tests/integration/test_surface_builder.py diff --git a/ipsuite/configuration_generation/__init__.pyi b/ipsuite/configuration_generation/__init__.pyi index 1c1ace55..1ae48f57 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 BuildSurface __all__ = [ "Smiles2Atoms", @@ -10,4 +11,5 @@ __all__ = [ "Smiles2Conformers", "MultiPackmol", "Smiles2Gromacs", + "BuildSurface", ] diff --git a/ipsuite/configuration_generation/surface_builder.py b/ipsuite/configuration_generation/surface_builder.py new file mode 100644 index 00000000..efc7fab4 --- /dev/null +++ b/ipsuite/configuration_generation/surface_builder.py @@ -0,0 +1,93 @@ +from pathlib import Path +import typing as t + +import ase +import h5py +import znh5md +import zntrack +from ase.build import 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)[:] \ No newline at end of file diff --git a/tests/integration/test_surface_builder.py b/tests/integration/test_surface_builder.py new file mode 100644 index 00000000..aaba28f2 --- /dev/null +++ b/tests/integration/test_surface_builder.py @@ -0,0 +1,69 @@ +import ipsuite as ips + + +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 \ No newline at end of file From 397b32f6296a337e16671c230679025611486161 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 7 Aug 2025 12:33:23 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../surface_builder.py | 39 +++++++++---------- tests/integration/test_surface_builder.py | 18 ++++----- 2 files changed, 27 insertions(+), 30 deletions(-) diff --git a/ipsuite/configuration_generation/surface_builder.py b/ipsuite/configuration_generation/surface_builder.py index efc7fab4..bd0ec1ce 100644 --- a/ipsuite/configuration_generation/surface_builder.py +++ b/ipsuite/configuration_generation/surface_builder.py @@ -1,5 +1,5 @@ -from pathlib import Path import typing as t +from pathlib import Path import ase import h5py @@ -12,10 +12,10 @@ 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 @@ -29,25 +29,27 @@ class BuildSurface(base.IPSNode): 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'), + 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) - + 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: @@ -56,25 +58,20 @@ def run(self) -> None: 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 + 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 + 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) @@ -82,7 +79,7 @@ def run(self) -> None: @property def frames(self) -> list[ase.Atoms]: """Get the generated surface structure. - + Returns ------- list[ase.Atoms] @@ -90,4 +87,4 @@ def frames(self) -> list[ase.Atoms]: """ with self.state.fs.open(self.frames_path, "rb") as f: with h5py.File(f) as file: - return znh5md.IO(file_handle=file)[:] \ No newline at end of file + return znh5md.IO(file_handle=file)[:] diff --git a/tests/integration/test_surface_builder.py b/tests/integration/test_surface_builder.py index aaba28f2..593dda1f 100644 --- a/tests/integration/test_surface_builder.py +++ b/tests/integration/test_surface_builder.py @@ -5,8 +5,8 @@ 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), + lattice="Au", + indices=(1, 1, 1), layers=5 ) @@ -14,10 +14,10 @@ def test_BuildSurface_basic(proj_path): 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 @@ -39,10 +39,10 @@ def test_BuildSurface_custom_parameters(proj_path): 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 @@ -51,11 +51,11 @@ def test_BuildSurface_custom_parameters(proj_path): 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_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] @@ -66,4 +66,4 @@ def test_BuildSurface_different_indices(proj_path): assert len(surface.frames) == 1 atoms = surface.frames[0] assert all(symbol == "Pt" for symbol in atoms.get_chemical_symbols()) - assert len(atoms) > 0 \ No newline at end of file + assert len(atoms) > 0 From 1d4745428978ea153c4880b3ef1603f0cef87465 Mon Sep 17 00:00:00 2001 From: Fabian Zills Date: Thu, 7 Aug 2025 14:46:49 +0200 Subject: [PATCH 3/6] test relax --- ipsuite/__init__.pyi | 2 ++ tests/integration/test_surface_builder.py | 42 +++++++++++++++++++++++ 2 files changed, 44 insertions(+) diff --git a/ipsuite/__init__.pyi b/ipsuite/__init__.pyi index 96a01394..41d48fe4 100644 --- a/ipsuite/__init__.pyi +++ b/ipsuite/__init__.pyi @@ -61,6 +61,7 @@ from .configuration_generation import ( Smiles2Atoms, Smiles2Conformers, Smiles2Gromacs, + BuildSurface ) # Configuration Selection @@ -156,6 +157,7 @@ __all__ = [ "Smiles2Atoms", "Smiles2Conformers", "Smiles2Gromacs", + "BuildSurface", # Data "AddData", "AddDataH5MD", diff --git a/tests/integration/test_surface_builder.py b/tests/integration/test_surface_builder.py index 593dda1f..5d4d21dc 100644 --- a/tests/integration/test_surface_builder.py +++ b/tests/integration/test_surface_builder.py @@ -1,4 +1,5 @@ import ipsuite as ips +import pytest def test_BuildSurface_basic(proj_path): @@ -67,3 +68,44 @@ def test_BuildSurface_different_indices(proj_path): 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) From 42def7cf14e0a63250c0473fdec5e8501d149b2b Mon Sep 17 00:00:00 2001 From: Fabian Zills Date: Thu, 7 Aug 2025 17:52:38 +0200 Subject: [PATCH 4/6] selection / constraint API and bunch of Agent created testing --- ipsuite/__init__.pyi | 35 ++- ipsuite/atom_selection/__init__.py | 21 ++ ipsuite/atom_selection/constraints.py | 134 +++++++++ ipsuite/atom_selection/selections.py | 260 +++++++++++++++++ ipsuite/configuration_generation/__init__.pyi | 3 +- .../surface_builder.py | 275 +++++++++++++++++- ipsuite/interfaces.py | 51 +++- .../test_atom_selection_integration.py | 140 +++++++++ tests/integration/test_surface_builder.py | 259 +++++++++++++++++ .../atom_selection/test_selections.py | 191 ++++++++++++ 10 files changed, 1363 insertions(+), 6 deletions(-) create mode 100644 ipsuite/atom_selection/__init__.py create mode 100644 ipsuite/atom_selection/constraints.py create mode 100644 ipsuite/atom_selection/selections.py create mode 100644 tests/integration/test_atom_selection_integration.py create mode 100644 tests/unit_tests/atom_selection/test_selections.py diff --git a/ipsuite/__init__.pyi b/ipsuite/__init__.pyi index 41d48fe4..ada508c9 100644 --- a/ipsuite/__init__.pyi +++ b/ipsuite/__init__.pyi @@ -56,12 +56,13 @@ from .calculators import ( # Configuration Generation from .configuration_generation import ( + AddAdsorbate, + BuildSurface, MultiPackmol, Packmol, Smiles2Atoms, Smiles2Conformers, Smiles2Gromacs, - BuildSurface ) # Configuration Selection @@ -127,6 +128,21 @@ from .models import ( from .project import Project from .version import __version__ +# Interfaces +from .interfaces import AtomConstraint, AtomSelector + +# Atom Selection +from .atom_selection import ( + ElementTypeSelection, + FixAtomsConstraint, + FixBondLengthConstraint, + FixBondLengthsConstraint, + LayerSelection, + RadialSelection, + SurfaceSelection, + ZPositionSelection, +) + # Update __all__ for lazy loading __all__ = [ "__version__", @@ -152,12 +168,13 @@ __all__ = [ "ThresholdSelection", "FilterOutlier", # Configuration Generation - "Packmol", + "AddAdsorbate", + "BuildSurface", "MultiPackmol", + "Packmol", "Smiles2Atoms", "Smiles2Conformers", "Smiles2Gromacs", - "BuildSurface", # Data "AddData", "AddDataH5MD", @@ -234,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..a05ff2c2 --- /dev/null +++ b/ipsuite/atom_selection/__init__.py @@ -0,0 +1,21 @@ +"""Module for selecting atoms within individual ASE Atoms objects.""" + +from .selections import ( + ElementTypeSelection, + LayerSelection, + RadialSelection, + SurfaceSelection, + ZPositionSelection, +) +from .constraints import ( + FixAtomsConstraint, +) + +__all__ = [ + "ElementTypeSelection", + "LayerSelection", + "RadialSelection", + "SurfaceSelection", + "ZPositionSelection", + "FixAtomsConstraint", +] \ No newline at end of file diff --git a/ipsuite/atom_selection/constraints.py b/ipsuite/atom_selection/constraints.py new file mode 100644 index 00000000..d7ab292f --- /dev/null +++ b/ipsuite/atom_selection/constraints.py @@ -0,0 +1,134 @@ +"""Constraint classes that use atom selections.""" + +import dataclasses +import typing as t + +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) \ No newline at end of file diff --git a/ipsuite/atom_selection/selections.py b/ipsuite/atom_selection/selections.py new file mode 100644 index 00000000..12e89ff7 --- /dev/null +++ b/ipsuite/atom_selection/selections.py @@ -0,0 +1,260 @@ +"""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 \ No newline at end of file diff --git a/ipsuite/configuration_generation/__init__.pyi b/ipsuite/configuration_generation/__init__.pyi index 1ae48f57..8df732ac 100644 --- a/ipsuite/configuration_generation/__init__.pyi +++ b/ipsuite/configuration_generation/__init__.pyi @@ -3,7 +3,7 @@ from .gmx import Smiles2Gromacs from .packmol import MultiPackmol, Packmol from .smiles_to_atoms import Smiles2Atoms, Smiles2Conformers -from .surface_builder import BuildSurface +from .surface_builder import AddAdsorbate, BuildSurface __all__ = [ "Smiles2Atoms", @@ -12,4 +12,5 @@ __all__ = [ "MultiPackmol", "Smiles2Gromacs", "BuildSurface", + "AddAdsorbate", ] diff --git a/ipsuite/configuration_generation/surface_builder.py b/ipsuite/configuration_generation/surface_builder.py index bd0ec1ce..7f093443 100644 --- a/ipsuite/configuration_generation/surface_builder.py +++ b/ipsuite/configuration_generation/surface_builder.py @@ -3,9 +3,10 @@ import ase import h5py +import numpy as np import znh5md import zntrack -from ase.build import bulk, surface +from ase.build import add_adsorbate, bulk, surface from ipsuite import base @@ -88,3 +89,275 @@ def frames(self) -> list[ase.Atoms]: 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..c5207c57 100644 --- a/ipsuite/interfaces.py +++ b/ipsuite/interfaces.py @@ -1,9 +1,10 @@ import pathlib -from typing import Dict, List, Protocol, TypeVar, Union +from typing import Any, Dict, List, Protocol, TypeVar, Union import ase from ase.calculators.calculator import Calculator from ase.optimize.optimize import Dynamics +from ase.constraints import FixConstraint 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", diff --git a/tests/integration/test_atom_selection_integration.py b/tests/integration/test_atom_selection_integration.py new file mode 100644 index 00000000..6f6eca38 --- /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()) \ No newline at end of file diff --git a/tests/integration/test_surface_builder.py b/tests/integration/test_surface_builder.py index 5d4d21dc..2a32a9d0 100644 --- a/tests/integration/test_surface_builder.py +++ b/tests/integration/test_surface_builder.py @@ -1,5 +1,6 @@ import ipsuite as ips import pytest +import numpy as np def test_BuildSurface_basic(proj_path): @@ -109,3 +110,261 @@ def test_BuildSurface_with_mace_relaxation(proj_path): 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..9aca1732 --- /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 \ No newline at end of file From b5df5243e7beab3eb14f3e5f2c6a0122b0db4342 Mon Sep 17 00:00:00 2001 From: Fabian Zills Date: Fri, 8 Aug 2025 15:28:54 +0200 Subject: [PATCH 5/6] bump zntrack release requirement --- ipsuite/interfaces.py | 2 +- pyproject.toml | 2 +- uv.lock | 14 +++++++------- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/ipsuite/interfaces.py b/ipsuite/interfaces.py index c5207c57..a70ecfd1 100644 --- a/ipsuite/interfaces.py +++ b/ipsuite/interfaces.py @@ -130,4 +130,4 @@ def get_constraint(self, atoms: ase.Atoms) -> FixConstraint: 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/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" }, ] From aabaeb727d0bd10a6d14531345aca1b873b6b41d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 8 Aug 2025 13:29:03 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- ipsuite/__init__.pyi | 34 +++--- ipsuite/atom_selection/__init__.py | 10 +- ipsuite/atom_selection/constraints.py | 51 +++++---- ipsuite/atom_selection/selections.py | 76 ++++++------ .../surface_builder.py | 108 ++++++++++-------- ipsuite/interfaces.py | 16 +-- .../test_atom_selection_integration.py | 44 +++---- tests/integration/test_surface_builder.py | 96 ++++++++-------- .../atom_selection/test_selections.py | 62 +++++----- 9 files changed, 258 insertions(+), 239 deletions(-) diff --git a/ipsuite/__init__.pyi b/ipsuite/__init__.pyi index ada508c9..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 @@ -112,6 +124,9 @@ from .dynamics import ( # Geometry from .geometry import BarycenterMapping +# Interfaces +from .interfaces import AtomConstraint, AtomSelector + # Models from .models import ( GAP, @@ -128,21 +143,6 @@ from .models import ( from .project import Project from .version import __version__ -# Interfaces -from .interfaces import AtomConstraint, AtomSelector - -# Atom Selection -from .atom_selection import ( - ElementTypeSelection, - FixAtomsConstraint, - FixBondLengthConstraint, - FixBondLengthsConstraint, - LayerSelection, - RadialSelection, - SurfaceSelection, - ZPositionSelection, -) - # Update __all__ for lazy loading __all__ = [ "__version__", @@ -169,7 +169,7 @@ __all__ = [ "FilterOutlier", # Configuration Generation "AddAdsorbate", - "BuildSurface", + "BuildSurface", "MultiPackmol", "Packmol", "Smiles2Atoms", @@ -256,7 +256,7 @@ __all__ = [ "AtomSelector", # Atom Selection "ElementTypeSelection", - "FixAtomsConstraint", + "FixAtomsConstraint", "FixBondLengthConstraint", "FixBondLengthsConstraint", "LayerSelection", diff --git a/ipsuite/atom_selection/__init__.py b/ipsuite/atom_selection/__init__.py index a05ff2c2..31ca2904 100644 --- a/ipsuite/atom_selection/__init__.py +++ b/ipsuite/atom_selection/__init__.py @@ -1,5 +1,8 @@ """Module for selecting atoms within individual ASE Atoms objects.""" +from .constraints import ( + FixAtomsConstraint, +) from .selections import ( ElementTypeSelection, LayerSelection, @@ -7,15 +10,12 @@ SurfaceSelection, ZPositionSelection, ) -from .constraints import ( - FixAtomsConstraint, -) __all__ = [ "ElementTypeSelection", - "LayerSelection", + "LayerSelection", "RadialSelection", "SurfaceSelection", "ZPositionSelection", "FixAtomsConstraint", -] \ No newline at end of file +] diff --git a/ipsuite/atom_selection/constraints.py b/ipsuite/atom_selection/constraints.py index d7ab292f..43f72370 100644 --- a/ipsuite/atom_selection/constraints.py +++ b/ipsuite/atom_selection/constraints.py @@ -1,7 +1,6 @@ """Constraint classes that use atom selections.""" import dataclasses -import typing as t import ase from ase.constraints import FixAtoms, FixBondLength, FixBondLengths @@ -23,12 +22,12 @@ class FixAtomsConstraint(AtomConstraint): 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 @@ -48,7 +47,7 @@ class FixBondLengthConstraint(AtomConstraint): ---------- atom1_selection : AtomSelector Selection for the first atom. Must select exactly one atom. - atom2_selection : AtomSelector + atom2_selection : AtomSelector Selection for the second atom. Must select exactly one atom. """ @@ -57,17 +56,17 @@ class FixBondLengthConstraint(AtomConstraint): 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 @@ -75,12 +74,16 @@ def get_constraint(self, atoms: ase.Atoms) -> FixBondLength: """ 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)}") + 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)}") - + raise ValueError( + f"atom2_selection must select exactly 1 atom, got {len(indices2)}" + ) + return FixBondLength(indices1[0], indices2[0]) @@ -99,36 +102,40 @@ class FixBondLengthsConstraint(AtomConstraint): 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)}") + 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)}") - + 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) \ No newline at end of file + + return FixBondLengths(pairs) diff --git a/ipsuite/atom_selection/selections.py b/ipsuite/atom_selection/selections.py index 12e89ff7..0e0ca58c 100644 --- a/ipsuite/atom_selection/selections.py +++ b/ipsuite/atom_selection/selections.py @@ -23,24 +23,22 @@ class ElementTypeSelection(AtomSelector): 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 - ] + return [i for i, symbol in enumerate(symbols) if symbol in self.elements] -@dataclasses.dataclass +@dataclasses.dataclass class ZPositionSelection(AtomSelector): """Select atoms based on their Z-coordinate position. @@ -48,7 +46,7 @@ class ZPositionSelection(AtomSelector): ---------- z_min : float | None, optional Minimum Z-coordinate (inclusive). If None, no lower bound. - z_max : float | None, optional + z_max : float | None, optional Maximum Z-coordinate (inclusive). If None, no upper bound. """ @@ -57,12 +55,12 @@ class ZPositionSelection(AtomSelector): 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] @@ -70,14 +68,14 @@ def select(self, atoms: ase.Atoms) -> list[int]: """ 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() @@ -88,23 +86,23 @@ class RadialSelection(AtomSelector): Parameters ---------- center : tuple[float, float, float] | str - Center point for radial selection. Can be coordinates (x, y, z) + 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'] + 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 @@ -112,21 +110,21 @@ def _get_center_position(self, atoms: ase.Atoms) -> np.ndarray: """ if isinstance(self.center, tuple): return np.array(self.center) - elif self.center == 'com': + elif self.center == "com": return atoms.get_center_of_mass() - elif self.center == 'geometric': + 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] @@ -134,10 +132,10 @@ def select(self, atoms: ase.Atoms) -> list[int]: """ 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() @@ -158,30 +156,30 @@ class LayerSelection(AtomSelector): 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 + 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]) @@ -189,18 +187,18 @@ def _identify_layers(self, atoms: ase.Atoms) -> list[list[int]]: 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] @@ -208,11 +206,11 @@ def select(self, atoms: ase.Atoms) -> list[int]: """ 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) @@ -234,12 +232,12 @@ class SurfaceSelection(AtomSelector): 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] @@ -247,14 +245,14 @@ def select(self, atoms: ase.Atoms) -> list[int]: """ 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 \ No newline at end of file + + return surface_indices diff --git a/ipsuite/configuration_generation/surface_builder.py b/ipsuite/configuration_generation/surface_builder.py index 7f093443..a2a6204d 100644 --- a/ipsuite/configuration_generation/surface_builder.py +++ b/ipsuite/configuration_generation/surface_builder.py @@ -93,10 +93,10 @@ def frames(self) -> list[ase.Atoms]: 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] @@ -116,30 +116,30 @@ class AddAdsorbate(base.IPSNode): 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 @@ -151,12 +151,12 @@ def _get_com_position(self, atoms: ase.Atoms) -> np.ndarray: 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] @@ -165,11 +165,15 @@ def _get_surface_center(self, slab: ase.Atoms) -> tuple[float, float]: 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: + 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] @@ -180,24 +184,28 @@ def _check_collision(self, new_pos: tuple[float, 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) + 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]: + 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 @@ -208,12 +216,12 @@ def _find_valid_position(self, slab: ase.Atoms, radius: 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 @@ -221,7 +229,7 @@ def _find_valid_position(self, slab: ase.Atoms, radius: float, """ 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 @@ -229,16 +237,18 @@ def _find_valid_position(self, slab: ase.Atoms, radius: float, # 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) + 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): + + 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." @@ -246,12 +256,12 @@ def _find_valid_position(self, slab: ase.Atoms, radius: float, 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 @@ -260,13 +270,13 @@ def _get_selected_adsorbates(self) -> list[ase.Atoms]: 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: @@ -281,33 +291,33 @@ def _get_selected_adsorbates(self) -> list[ase.Atoms]: 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) @@ -316,7 +326,7 @@ def run(self) -> None: 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( @@ -325,26 +335,30 @@ def run(self) -> None: height=height, position=position, offset=None, - mol_index=0 # Use first atom as reference instead of COM + 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: + if "adsorbate_info" in slab.info: # Convert numpy int64 to regular int for JSON compatibility - adsorbate_info = slab.info['adsorbate_info'] + 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 + 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 - + 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) @@ -352,7 +366,7 @@ def run(self) -> None: @property def frames(self) -> list[ase.Atoms]: """Get the surface with adsorbed molecules. - + Returns ------- list[ase.Atoms] diff --git a/ipsuite/interfaces.py b/ipsuite/interfaces.py index a70ecfd1..4aabec21 100644 --- a/ipsuite/interfaces.py +++ b/ipsuite/interfaces.py @@ -1,10 +1,10 @@ import pathlib -from typing import Any, Dict, List, Protocol, TypeVar, Union +from typing import Dict, List, Protocol, TypeVar, Union import ase from ase.calculators.calculator import Calculator -from ase.optimize.optimize import Dynamics from ase.constraints import FixConstraint +from ase.optimize.optimize import Dynamics T = TypeVar("T", covariant=True) @@ -64,19 +64,19 @@ class ProcessAtoms(Protocol): 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] @@ -87,19 +87,19 @@ def select(self, atoms: ase.Atoms) -> list[int]: 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 diff --git a/tests/integration/test_atom_selection_integration.py b/tests/integration/test_atom_selection_integration.py index 6f6eca38..5aedbea3 100644 --- a/tests/integration/test_atom_selection_integration.py +++ b/tests/integration/test_atom_selection_integration.py @@ -13,14 +13,14 @@ def test_asegeopt_with_atom_selection_constraints(proj_path): 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, @@ -38,11 +38,11 @@ def test_asegeopt_with_atom_selection_constraints(proj_path): 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) @@ -54,48 +54,48 @@ def test_element_selection_constraint(proj_path): with ips.Project() as proj: # Create surface with adsorbate surface = ips.BuildSurface( - lattice="Pt", + 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", + 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 @@ -111,16 +111,16 @@ def test_multiple_selection_types(proj_path): 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, @@ -131,10 +131,10 @@ def test_multiple_selection_types(proj_path): ) 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()) \ No newline at end of file + 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 index 2a32a9d0..df46d3f7 100644 --- a/tests/integration/test_surface_builder.py +++ b/tests/integration/test_surface_builder.py @@ -81,10 +81,10 @@ def test_BuildSurface_with_mace_relaxation(proj_path): 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, @@ -101,11 +101,11 @@ def test_BuildSurface_with_mace_relaxation(proj_path): 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 + + # 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()) @@ -122,10 +122,10 @@ def test_AddAdsorbate_basic(proj_path): 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, @@ -139,20 +139,20 @@ def test_AddAdsorbate_basic(proj_path): # 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") + 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 @@ -167,15 +167,15 @@ def test_AddAdsorbate_multiple_adsorbates(proj_path): 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 + data=[co.frames, h2.frames], # Nested structure height=[1.5, 3.0], excluded_radius=[2.5, 2.0] ) @@ -185,29 +185,29 @@ def test_AddAdsorbate_multiple_adsorbates(proj_path): # 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 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) + 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 @@ -224,12 +224,12 @@ def test_AddAdsorbate_collision_detection_error(proj_path): 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") + 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, @@ -253,11 +253,11 @@ def test_AddAdsorbate_data_index_functionality(proj_path): 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, @@ -266,7 +266,7 @@ def test_AddAdsorbate_data_index_functionality(proj_path): height=[2.0, 2.5], excluded_radius=[2.0, 2.0] ) - + # Test with specific data_index adsorbed_system_custom = ips.AddAdsorbate( slab=surface.frames, @@ -281,25 +281,25 @@ def test_AddAdsorbate_data_index_functionality(proj_path): # 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") + 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 h_count == 2 # Water hydrogens assert o_count == 2 # Water oxygen + CO oxygen assert c_count == 1 # CO carbon @@ -314,20 +314,20 @@ def test_AddAdsorbate_slab_idx_functionality(proj_path): layers=3, vacuum=10.0 ) - + surface2 = ips.BuildSurface( - lattice="Pt", + 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, @@ -336,7 +336,7 @@ def test_AddAdsorbate_slab_idx_functionality(proj_path): height=[2.0], excluded_radius=[2.0] ) - + # Test with slab_idx=-1 (default, last slab - Pt) adsorbed_pt = ips.AddAdsorbate( slab=combined_surfaces, @@ -351,20 +351,20 @@ def test_AddAdsorbate_slab_idx_functionality(proj_path): # 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 + + # 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("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 index 9aca1732..56ac30e2 100644 --- a/tests/unit_tests/atom_selection/test_selections.py +++ b/tests/unit_tests/atom_selection/test_selections.py @@ -10,19 +10,19 @@ 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'], + 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) @@ -34,24 +34,24 @@ def test_z_position_selection(): # Create atoms at different Z positions positions = np.array([ [0, 0, 1.0], - [0, 0, 2.5], + [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 + + # Test only maximum max_selector = ips.ZPositionSelection(z_max=3.0) max_indices = max_selector.select(atoms) assert set(max_indices) == {0, 1} @@ -68,12 +68,12 @@ def test_radial_selection(): [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) @@ -92,23 +92,23 @@ def test_layer_selection(): positions = np.array([ # Bottom layer (z=0) [0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], - # Middle layer (z=2) + # 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) @@ -124,27 +124,27 @@ def test_surface_selection(): [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_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'], + 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} @@ -152,11 +152,11 @@ def test_fix_atoms_constraint(): 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) @@ -165,27 +165,27 @@ 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) + # 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 \ No newline at end of file + # where only the surface layer atoms are allowed to move