hillclimber 0.1.0a1__py3-none-any.whl → 0.1.0a2__py3-none-any.whl

hillclimber/interfaces.py

@@ -69,13 +69,52 @@ class CollectiveVariable(Protocol):
         ...
 
 
-class MetadynamicsBiasCollectiveVariable(Protocol):
-    """Protocol for metadata associated with a bias in PLUMED."""
+class BiasProtocol(Protocol):
+    """Protocol for individual bias potentials that act on collective variables.
+
+    Individual bias potentials (restraints, walls) operate on collective variables
+    and generate their own PLUMED commands via to_plumed(). These biases are added
+    to MetaDynamicsModel via the `actions` parameter.
+
+    This is distinct from MetadynamicsBias, which is used for configuring the
+    collective METAD command itself (via the `bias_cvs` parameter).
+    """
+
+    cv: CollectiveVariable
+
+    def to_plumed(self, atoms: ase.Atoms) -> list[str]:
+        """Generate PLUMED input strings for this bias.
+
+        Parameters
+        ----------
+        atoms : ase.Atoms
+            The atomic structure to use for generating the PLUMED string.
+
+        Returns
+        -------
+        list[str]
+            List of PLUMED command strings.
+        """
+        ...
+
+
+class MetadynamicsBias(Protocol):
+    """Protocol for metadynamics bias configuration objects.
+
+    This protocol defines the configuration structure for per-CV metadynamics
+    parameters (sigma, grid settings, etc.). These configuration objects are
+    used in the `bias_cvs` parameter of MetaDynamicsModel to build the collective
+    METAD command.
+
+    This is distinct from BiasProtocol, which is for individual bias potentials
+    (restraints, walls) that generate their own PLUMED commands and are added
+    via the `actions` parameter.
+    """
 
     cv: CollectiveVariable
     sigma: float | None = None
-    grid_min: float | None = None
-    grid_max: float | None = None
+    grid_min: float | str | None = None
+    grid_max: float | str | None = None
     grid_bin: int | None = None
 
 
@@ -84,7 +123,8 @@ __all__ = [
     "AtomSelector",
     "PlumedGenerator",
     "CollectiveVariable",
-    "MetadynamicsBiasCollectiveVariable",
+    "BiasProtocol",
+    "MetadynamicsBias",
 ]
 
 def interfaces() -> dict[str, list[str]]:

hillclimber/metadynamics.py

@@ -7,15 +7,14 @@ import zntrack
 from hillclimber.calc import NonOverwritingPlumed
 from hillclimber.interfaces import (
     CollectiveVariable,
-    MetadynamicsBiasCollectiveVariable,
     NodeWithCalculator,
     PlumedGenerator,
 )
 
 
 @dataclasses.dataclass
-class MetaDBiasCV(MetadynamicsBiasCollectiveVariable):
-    """Metadynamics bias on a collective variable.
+class MetadBias:
+    """Metadynamics bias configuration for a collective variable.
 
     Parameters
     ----------
@@ -115,7 +114,7 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
     ...     x2=pn.SMARTSSelector(pattern="CO[C:1]"),
     ...     prefix="d",
     ... )
-    >>> metad_cv1 = pn.MetaDBiasCV(
+    >>> metad_cv1 = pn.MetadBias(
     ...     cv=cv1, sigma=0.1, grid_min=0.0, grid_max=2.0, grid_bin=200
     ... )
     >>> model = pn.MetaDynamicsModel(
@@ -133,7 +132,7 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
     config: MetaDynamicsConfig = zntrack.deps()
     data: list[ase.Atoms] = zntrack.deps()
     data_idx: int = zntrack.params(-1)
-    bias_cvs: list[MetaDBiasCV] = zntrack.deps(default_factory=list)
+    bias_cvs: list[MetadBias] = zntrack.deps(default_factory=list)
     actions: list[PlumedGenerator] = zntrack.deps(default_factory=list)
     timestep: float = zntrack.params(1.0)  # in fs, default is 1 fs
     model: NodeWithCalculator = zntrack.deps()
@@ -240,10 +239,13 @@ class MetaDynamicsModel(zntrack.Node, NodeWithCalculator):
             )
 
         plumed_lines.append(f"metad: {' '.join(metad_parts)}")
-        # Temporary until https://github.com/zincware/ZnTrack/issues/936
-        from hillclimber.actions import PrintCVAction
-        lines = PrintCVAction(cvs=[x.cv for x in self.bias_cvs], stride=100).to_plumed(atoms)
-        plumed_lines.extend(lines)
+
+        # Add any additional actions (restraints, walls, print actions, etc.)
+        for action in self.actions:
+            action_lines = action.to_plumed(atoms)
+            plumed_lines.extend(action_lines)
+
+        # Add FLUSH if configured
         if self.config.flush is not None:
             plumed_lines.append(f"FLUSH STRIDE={self.config.flush}")
         return plumed_lines

hillclimber/opes.py

@@ -0,0 +1,342 @@
+"""OPES (On-the-fly Probability Enhanced Sampling) methods.
+
+This module provides classes for OPES enhanced sampling, a modern alternative
+to traditional metadynamics with improved convergence properties.
+"""
+
+import dataclasses
+from pathlib import Path
+
+import ase.units
+import zntrack
+
+from hillclimber.calc import NonOverwritingPlumed
+from hillclimber.interfaces import (
+    CollectiveVariable,
+    NodeWithCalculator,
+    PlumedGenerator,
+)
+
+
+@dataclasses.dataclass
+class OPESBias:
+    """OPES bias configuration for a collective variable.
+
+    Parameters
+    ----------
+    cv : CollectiveVariable
+        The collective variable to bias.
+    sigma : float | str, optional
+        Initial kernel width. Use "ADAPTIVE" for automatic adaptation
+        (recommended). If numeric, specifies the initial width.
+        Default: "ADAPTIVE".
+
+    Resources
+    ---------
+    - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD/
+
+    Notes
+    -----
+    The ADAPTIVE sigma option automatically adjusts kernel widths based on
+    CV fluctuations, which is usually the best choice. The automatic width
+    is measured every ADAPTIVE_SIGMA_STRIDE steps (default: 10×PACE).
+    """
+
+    cv: CollectiveVariable
+    sigma: float | str = "ADAPTIVE"
+
+
+@dataclasses.dataclass
+class OPESConfig:
+    """Configuration for OPES_METAD and OPES_METAD_EXPLORE.
+
+    OPES (On-the-fly Probability Enhanced Sampling) is a modern enhanced
+    sampling method that samples well-tempered target distributions.
+
+    Parameters
+    ----------
+    barrier : float
+        Highest free energy barrier to overcome (kJ/mol). This is the key
+        parameter that determines sampling efficiency.
+    pace : int, optional
+        Frequency of kernel deposition in MD steps (default: 500).
+    temp : float, optional
+        Temperature in Kelvin (default: 300.0). If -1, retrieved from MD engine.
+    explore_mode : bool, optional
+        If True, uses OPES_METAD_EXPLORE which estimates target distribution
+        directly (better exploration, slower reweighting convergence).
+        If False, uses OPES_METAD which estimates unbiased distribution
+        (faster convergence, less exploration). Default: False.
+    biasfactor : float, optional
+        Well-tempered gamma factor. If not specified, uses default behavior.
+        Set to inf for custom target distributions.
+    compression_threshold : float, optional
+        Merge kernels if closer than this threshold in sigma units (default: 1.0).
+    file : str, optional
+        File to store deposited kernels (default: "KERNELS").
+    adaptive_sigma_stride : int, optional
+        Steps between adaptive sigma measurements. If not set, uses 10×PACE.
+    sigma_min : float, optional
+        Minimum allowable sigma value for adaptive sigma.
+    state_wfile : str, optional
+        State file for writing exact restart information.
+    state_rfile : str, optional
+        State file for reading restart information.
+    state_wstride : int, optional
+        Frequency of STATE file writing in number of kernel depositions.
+    walkers_mpi : bool, optional
+        Enable multiple walker mode with MPI communication (default: False).
+    calc_work : bool, optional
+        Calculate and output accumulated work (default: False).
+    flush : int, optional
+        Frequency of flushing output files.
+
+    Resources
+    ---------
+    - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD/
+    - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD_EXPLORE/
+    - Invernizzi & Parrinello, J. Phys. Chem. Lett. 2020
+
+    Notes
+    -----
+    **When to use OPES_METAD vs OPES_METAD_EXPLORE:**
+
+    - OPES_METAD: Use when you want quick convergence of reweighted free energy.
+      Estimates unbiased distribution P(s).
+
+    - OPES_METAD_EXPLORE: Use for systems with unknown barriers or when testing
+      new CVs. Allows more exploration but slower reweighting convergence.
+      Estimates target distribution p^WT(s) directly.
+
+    Both methods converge to the same bias given enough time, but approach it
+    differently. OPES is more sensitive to degenerate CVs than standard METAD.
+    """
+
+    barrier: float  # kJ/mol
+    pace: int = 500
+    temp: float = 300.0
+    explore_mode: bool = False  # False=OPES_METAD, True=OPES_METAD_EXPLORE
+    biasfactor: float | None = None
+    compression_threshold: float = 1.0
+    file: str = "KERNELS"
+    adaptive_sigma_stride: int | None = None
+    sigma_min: float | None = None
+    state_wfile: str | None = None
+    state_rfile: str | None = None
+    state_wstride: int | None = None
+    walkers_mpi: bool = False
+    calc_work: bool = False
+    flush: int | None = None
+
+
+class OPESModel(zntrack.Node, NodeWithCalculator):
+    """OPES (On-the-fly Probability Enhanced Sampling) model.
+
+    Implements OPES_METAD and OPES_METAD_EXPLORE enhanced sampling methods.
+    OPES samples well-tempered target distributions and provides better
+    convergence properties than traditional metadynamics.
+
+    Parameters
+    ----------
+    config : OPESConfig
+        Configuration for the OPES simulation.
+    data : list[ase.Atoms]
+        Input data for simulation.
+    data_idx : int, optional
+        Index of data to use (default: -1).
+    bias_cvs : list[OPESBias], optional
+        Collective variables to bias (default: []).
+    actions : list[PlumedGenerator], optional
+        Additional actions like restraints, walls, print (default: []).
+    timestep : float, optional
+        Timestep in fs (default: 1.0).
+    model : NodeWithCalculator
+        Underlying force field model.
+
+    Examples
+    --------
+    >>> import hillclimber as hc
+    >>>
+    >>> # Define collective variables
+    >>> phi = hc.TorsionCV(...)
+    >>> psi = hc.TorsionCV(...)
+    >>>
+    >>> # OPES configuration (standard mode)
+    >>> config = hc.OPESConfig(
+    ...     barrier=40.0,  # kJ/mol
+    ...     pace=500,
+    ...     temp=300.0,
+    ...     explore_mode=False  # Use OPES_METAD
+    ... )
+    >>>
+    >>> # Bias configuration with adaptive sigma
+    >>> bias1 = hc.OPESBias(cv=phi, sigma="ADAPTIVE")
+    >>> bias2 = hc.OPESBias(cv=psi, sigma="ADAPTIVE")
+    >>>
+    >>> # Create OPES model
+    >>> opes = hc.OPESModel(
+    ...     config=config,
+    ...     bias_cvs=[bias1, bias2],
+    ...     data=data.frames,
+    ...     model=force_field,
+    ...     timestep=0.5
+    ... )
+    >>>
+    >>> # For exploration mode, set explore_mode=True
+    >>> explore_config = hc.OPESConfig(
+    ...     barrier=40.0,
+    ...     pace=500,
+    ...     explore_mode=True  # Use OPES_METAD_EXPLORE
+    ... )
+
+    Resources
+    ---------
+    - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD/
+    - https://www.plumed.org/doc-master/user-doc/html/OPES_METAD_EXPLORE/
+    - https://www.plumed.org/doc-master/user-doc/html/masterclass-22-03.html
+    - Invernizzi & Parrinello, J. Phys. Chem. Lett. 2020
+
+    Notes
+    -----
+    **Output Components:**
+    OPES provides several diagnostic outputs (accessible via PrintAction):
+    - opes.bias: Instantaneous bias potential value
+    - opes.rct: Convergence indicator (should flatten at convergence)
+    - opes.zed: Normalization estimate (should stabilize)
+    - opes.neff: Effective sample size
+    - opes.nker: Number of compressed kernels
+    - opes.work: Accumulated work (if calc_work=True)
+
+    **Advantages over Metadynamics:**
+    - Better convergence properties
+    - Automatic variance adaptation
+    - Lower systematic error
+    - More sensitive to degenerate CVs (helps identify CV problems)
+    """
+
+    config: OPESConfig = zntrack.deps()
+    data: list[ase.Atoms] = zntrack.deps()
+    data_idx: int = zntrack.params(-1)
+    bias_cvs: list[OPESBias] = zntrack.deps(default_factory=list)
+    actions: list[PlumedGenerator] = zntrack.deps(default_factory=list)
+    timestep: float = zntrack.params(1.0)
+    model: NodeWithCalculator = zntrack.deps()
+
+    figures: Path = zntrack.outs_path(zntrack.nwd / "figures", independent=True)
+
+    def run(self):
+        self.figures.mkdir(parents=True, exist_ok=True)
+        for cv in self.bias_cvs:
+            img = cv.cv.get_img(self.data[self.data_idx])
+            img.save(self.figures / f"{cv.cv.prefix}.png")
+
+    def get_calculator(
+        self, *, directory: str | Path | None = None, **kwargs
+    ) -> NonOverwritingPlumed:
+        if directory is None:
+            raise ValueError("Directory must be specified for PLUMED input files.")
+        directory = Path(directory)
+        directory.mkdir(parents=True, exist_ok=True)
+
+        lines = self.to_plumed(self.data[self.data_idx])
+        # replace FILE= with f"FILE={directory}/" inside config
+        lines = [line.replace("FILE=", f"FILE={directory}/") for line in lines]
+
+        # Write plumed input file
+        with (directory / "plumed.dat").open("w") as file:
+            for line in lines:
+                file.write(line + "\n")
+
+        kT = ase.units.kB * self.config.temp
+
+        return NonOverwritingPlumed(
+            calc=self.model.get_calculator(directory=directory),
+            atoms=self.data[self.data_idx],
+            input=lines,
+            timestep=float(self.timestep * ase.units.fs),
+            kT=float(kT),
+            log=(directory / "plumed.log").as_posix(),
+        )
+
+    def to_plumed(self, atoms: ase.Atoms) -> list[str]:
+        """Generate PLUMED input string for the OPES model."""
+        # check for duplicate CV prefixes
+        cv_labels = set()
+        for bias_cv in self.bias_cvs:
+            if bias_cv.cv.prefix in cv_labels:
+                raise ValueError(f"Duplicate CV prefix found: {bias_cv.cv.prefix}")
+            cv_labels.add(bias_cv.cv.prefix)
+
+        plumed_lines = []
+        all_labels = []
+
+        sigmas = []
+
+        plumed_lines.append(
+            f"UNITS LENGTH=A TIME={1 / (1000 * ase.units.fs)} ENERGY={ase.units.mol / ase.units.kJ}"
+        )
+
+        for bias_cv in self.bias_cvs:
+            labels, cv_str = bias_cv.cv.to_plumed(atoms)
+            plumed_lines.extend(cv_str)
+            all_labels.extend(labels)
+
+            # Collect sigma values
+            if isinstance(bias_cv.sigma, str):
+                sigmas.append(bias_cv.sigma)
+            else:
+                sigmas.append(str(bias_cv.sigma))
+
+        # Determine which OPES method to use
+        method_name = "OPES_METAD_EXPLORE" if self.config.explore_mode else "OPES_METAD"
+
+        # Build OPES command
+        opes_parts = [
+            f"opes: {method_name}",
+            f"ARG={','.join(all_labels)}",
+            f"PACE={self.config.pace}",
+            f"BARRIER={self.config.barrier}",
+            f"TEMP={self.config.temp}",
+        ]
+
+        # Add SIGMA (required parameter)
+        # If all sigmas are the same, use single value; otherwise comma-separated
+        if len(set(sigmas)) == 1:
+            opes_parts.append(f"SIGMA={sigmas[0]}")
+        else:
+            opes_parts.append(f"SIGMA={','.join(sigmas)}")
+
+        # Add FILE and COMPRESSION_THRESHOLD
+        opes_parts.append(f"FILE={self.config.file}")
+        opes_parts.append(f"COMPRESSION_THRESHOLD={self.config.compression_threshold}")
+
+        # Optional parameters
+        if self.config.biasfactor is not None:
+            opes_parts.append(f"BIASFACTOR={self.config.biasfactor}")
+        if self.config.adaptive_sigma_stride is not None:
+            opes_parts.append(f"ADAPTIVE_SIGMA_STRIDE={self.config.adaptive_sigma_stride}")
+        if self.config.sigma_min is not None:
+            opes_parts.append(f"SIGMA_MIN={self.config.sigma_min}")
+        if self.config.state_wfile is not None:
+            opes_parts.append(f"STATE_WFILE={self.config.state_wfile}")
+        if self.config.state_rfile is not None:
+            opes_parts.append(f"STATE_RFILE={self.config.state_rfile}")
+        if self.config.state_wstride is not None:
+            opes_parts.append(f"STATE_WSTRIDE={self.config.state_wstride}")
+        if self.config.walkers_mpi:
+            opes_parts.append("WALKERS_MPI")
+        if self.config.calc_work:
+            opes_parts.append("CALC_WORK")
+
+        plumed_lines.append(" ".join(opes_parts))
+
+        # Add any additional actions (restraints, walls, print actions, etc.)
+        for action in self.actions:
+            action_lines = action.to_plumed(atoms)
+            plumed_lines.extend(action_lines)
+
+        # Add FLUSH if configured
+        if self.config.flush is not None:
+            plumed_lines.append(f"FLUSH STRIDE={self.config.flush}")
+
+        return plumed_lines

hillclimber/selectors.py

@@ -7,6 +7,107 @@ import rdkit2ase
 from hillclimber.interfaces import AtomSelector
 
 
+# --- Indexable Selector Wrappers ---
+
+
+@dataclasses.dataclass
+class _GroupIndexedSelector(AtomSelector):
+    """Selector with group-level indexing applied.
+
+    This is an internal class created when you index a selector at the group level.
+    For example: water_sel[0] or water_sel[0:2]
+    """
+    selector: AtomSelector
+    group_index: int | slice | list[int]
+
+    def __getitem__(self, idx: int | slice | list[int]) -> AtomSelector:
+        """Atom-level indexing.
+
+        After group indexing, this applies atom-level indexing.
+        For example: water_sel[0:2][1:3] selects atoms 1-2 from groups 0-1.
+        """
+        return _AtomIndexedSelector(self, idx)
+
+    def __add__(self, other: AtomSelector) -> AtomSelector:
+        """Combine two selectors."""
+        return _CombinedSelector([self, other])
+
+    def select(self, atoms: ase.Atoms) -> list[list[int]]:
+        """Apply group indexing to the underlying selector."""
+        groups = self.selector.select(atoms)
+
+        # Apply group indexing (supports negative indices)
+        if isinstance(self.group_index, int):
+            return [groups[self.group_index]]  # Python handles negative indices
+        elif isinstance(self.group_index, slice):
+            return groups[self.group_index]  # Python handles negative indices in slices
+        else:  # list[int]
+            return [groups[i] for i in self.group_index]  # Negative indices work here too
+
+
+@dataclasses.dataclass
+class _AtomIndexedSelector(AtomSelector):
+    """Selector with both group and atom-level indexing applied.
+
+    This is an internal class created when you apply two levels of indexing.
+    For example: water_sel[0][0] or water_sel[0:2][1:3]
+    """
+    group_selector: _GroupIndexedSelector
+    atom_index: int | slice | list[int]
+
+    def __getitem__(self, idx) -> AtomSelector:
+        """Prevent three-level indexing."""
+        raise ValueError("Cannot index beyond 2 levels (group, then atom)")
+
+    def __add__(self, other: AtomSelector) -> AtomSelector:
+        """Combine two selectors."""
+        return _CombinedSelector([self, other])
+
+    def select(self, atoms: ase.Atoms) -> list[list[int]]:
+        """Apply atom-level indexing to each group."""
+        groups = self.group_selector.select(atoms)
+
+        # Apply atom-level indexing to each group (supports negative indices)
+        result = []
+        for group in groups:
+            if isinstance(self.atom_index, int):
+                result.append([group[self.atom_index]])  # Negative indices work
+            elif isinstance(self.atom_index, slice):
+                result.append(group[self.atom_index])  # Negative indices in slices work
+            else:  # list[int]
+                result.append([group[i] for i in self.atom_index])  # Negative indices work
+
+        return result
+
+
+@dataclasses.dataclass
+class _CombinedSelector(AtomSelector):
+    """Selector that combines multiple selectors.
+
+    This is an internal class created when you combine selectors with +.
+    For example: water_sel + ethanol_sel
+    """
+    selectors: list[AtomSelector]
+
+    def __getitem__(self, idx: int | slice | list[int]) -> AtomSelector:
+        """Group-level indexing on combined result."""
+        return _GroupIndexedSelector(self, idx)
+
+    def __add__(self, other: AtomSelector) -> AtomSelector:
+        """Combine with another selector."""
+        # Flatten if other is also a CombinedSelector
+        if isinstance(other, _CombinedSelector):
+            return _CombinedSelector(self.selectors + other.selectors)
+        return _CombinedSelector(self.selectors + [other])
+
+    def select(self, atoms: ase.Atoms) -> list[list[int]]:
+        """Concatenate all groups from all selectors."""
+        result = []
+        for selector in self.selectors:
+            result.extend(selector.select(atoms))
+        return result
+
+
 @dataclasses.dataclass
 class IndexSelector(AtomSelector):
     """Select atoms based on grouped indices.
@@ -23,6 +124,14 @@ class IndexSelector(AtomSelector):
     # mostly used for debugging
     indices: list[list[int]]
 
+    def __getitem__(self, idx: int | slice | list[int]) -> AtomSelector:
+        """Group-level indexing."""
+        return _GroupIndexedSelector(self, idx)
+
+    def __add__(self, other: AtomSelector) -> AtomSelector:
+        """Combine two selectors."""
+        return _CombinedSelector([self, other])
+
     def select(self, atoms: ase.Atoms) -> list[list[int]]:
         return self.indices
 
@@ -39,6 +148,14 @@ class SMILESSelector(AtomSelector):
 
     smiles: str
 
+    def __getitem__(self, idx: int | slice | list[int]) -> AtomSelector:
+        """Group-level indexing."""
+        return _GroupIndexedSelector(self, idx)
+
+    def __add__(self, other: AtomSelector) -> AtomSelector:
+        """Combine two selectors."""
+        return _CombinedSelector([self, other])
+
     def select(self, atoms: ase.Atoms) -> list[list[int]]:
         matches = rdkit2ase.match_substructure(atoms, smiles=self.smiles)
         return [list(match) for match in matches]
@@ -55,8 +172,8 @@ class SMARTSSelector(AtomSelector):
 
     Note
     ----
-    The selector is applied only to the first trajectory frame.  
-    Since indices can change during e.g. proton transfer, biasing specific groups (e.g. `[OH-]`) may fail.  
+    The selector is applied only to the first trajectory frame.
+    Since indices can change during e.g. proton transfer, biasing specific groups (e.g. `[OH-]`) may fail.
     In such cases, select all `[OH2]` and `[OH-]` groups and use CoordinationNumber CVs.
     Account for this method with all changes in chemical structure.
 
@@ -90,6 +207,14 @@ class SMARTSSelector(AtomSelector):
     pattern: str
     hydrogens: tp.Literal["include", "exclude", "isolated"] = "exclude"
 
+    def __getitem__(self, idx: int | slice | list[int]) -> AtomSelector:
+        """Group-level indexing."""
+        return _GroupIndexedSelector(self, idx)
+
+    def __add__(self, other: AtomSelector) -> AtomSelector:
+        """Combine two selectors."""
+        return _CombinedSelector([self, other])
+
     def select(self, atoms: ase.Atoms) -> list[list[int]]:
         return rdkit2ase.select_atoms_grouped(
             rdkit2ase.ase2rdkit(atoms), self.pattern, self.hydrogens