adaptivepy-sampling 0.1.0__py3-none-any.whl

adaptivepy/clustering/sklearn_kmeans.py

@@ -0,0 +1,93 @@
+"""KMeans clustering via scikit-learn."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+import numpy as np
+from sklearn.cluster import KMeans
+
+from adaptivepy.clustering.base import Clusterer
+
+
+class SklearnKMeansClusterer(Clusterer):
+    """Wrap ``sklearn.cluster.KMeans`` for AdaptivePy.
+
+    Parameters
+    ----------
+    n_clusters : int
+        Number of clusters.
+    random_state : int or None
+        Random seed passed to KMeans.
+    **kwargs
+        Additional keyword arguments forwarded to ``KMeans``.
+    """
+
+    def __init__(
+        self,
+        n_clusters: int,
+        random_state: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.n_clusters = n_clusters
+        self.random_state = random_state
+        self._extra_kwargs = kwargs
+        self._model: Optional[KMeans] = None
+
+    def fit(self, X: np.ndarray) -> "SklearnKMeansClusterer":
+        """Fit KMeans on the provided feature matrix.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Feature matrix of shape ``(n_samples, n_features)``.
+
+        Returns
+        -------
+        SklearnKMeansClusterer
+            Fitted clusterer.
+        """
+        self._model = KMeans(
+            n_clusters=self.n_clusters,
+            random_state=self.random_state,
+            n_init=10,
+            **self._extra_kwargs,
+        )
+        self._model.fit(X)
+        return self
+
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        """Predict cluster labels for ``X``.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Feature matrix.
+
+        Returns
+        -------
+        np.ndarray
+            Cluster labels.
+
+        Raises
+        ------
+        RuntimeError
+            If ``fit`` has not been called.
+        """
+        if self._model is None:
+            raise RuntimeError("Clusterer must be fitted before predict.")
+        return self._model.predict(X)
+
+    @property
+    def cluster_centers_(self) -> Optional[np.ndarray]:
+        """Return KMeans cluster centers."""
+        if self._model is None:
+            return None
+        return self._model.cluster_centers_
+
+    @property
+    def model(self) -> Any:
+        """Return the fitted ``KMeans`` instance."""
+        if self._model is None:
+            raise RuntimeError("Clusterer must be fitted before accessing model.")
+        return self._model

adaptivepy/clustering/sklearn_minibatch.py

@@ -0,0 +1,94 @@
+"""MiniBatchKMeans clustering via scikit-learn."""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+import numpy as np
+from sklearn.cluster import MiniBatchKMeans
+
+from adaptivepy.clustering.base import Clusterer
+
+
+class SklearnMiniBatchClusterer(Clusterer):
+    """Wrap ``sklearn.cluster.MiniBatchKMeans`` for large datasets.
+
+    Parameters
+    ----------
+    n_clusters : int
+        Number of clusters.
+    random_state : int or None
+        Random seed passed to MiniBatchKMeans.
+    **kwargs
+        Additional keyword arguments forwarded to ``MiniBatchKMeans``.
+    """
+
+    def __init__(
+        self,
+        n_clusters: int,
+        random_state: Optional[int] = None,
+        **kwargs: Any,
+    ) -> None:
+        self.n_clusters = n_clusters
+        self.random_state = random_state
+        self._extra_kwargs = kwargs
+        self._model: Optional[MiniBatchKMeans] = None
+
+    def fit(self, X: np.ndarray) -> "SklearnMiniBatchClusterer":
+        """Fit MiniBatchKMeans on the provided feature matrix.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Feature matrix of shape ``(n_samples, n_features)``.
+
+        Returns
+        -------
+        SklearnMiniBatchClusterer
+            Fitted clusterer.
+        """
+        self._model = MiniBatchKMeans(
+            n_clusters=self.n_clusters,
+            random_state=self.random_state,
+            n_init=3,
+            batch_size=min(1024, max(256, X.shape[0] // 10)),
+            **self._extra_kwargs,
+        )
+        self._model.fit(X)
+        return self
+
+    def predict(self, X: np.ndarray) -> np.ndarray:
+        """Predict cluster labels for ``X``.
+
+        Parameters
+        ----------
+        X : np.ndarray
+            Feature matrix.
+
+        Returns
+        -------
+        np.ndarray
+            Cluster labels.
+
+        Raises
+        ------
+        RuntimeError
+            If ``fit`` has not been called.
+        """
+        if self._model is None:
+            raise RuntimeError("Clusterer must be fitted before predict.")
+        return self._model.predict(X)
+
+    @property
+    def cluster_centers_(self) -> Optional[np.ndarray]:
+        """Return MiniBatchKMeans cluster centers."""
+        if self._model is None:
+            return None
+        return self._model.cluster_centers_
+
+    @property
+    def model(self) -> Any:
+        """Return the fitted ``MiniBatchKMeans`` instance."""
+        if self._model is None:
+            raise RuntimeError("Clusterer must be fitted before accessing model.")
+        return self._model

adaptivepy/config/__init__.py

@@ -0,0 +1,17 @@
+"""Configuration package for AdaptivePy."""
+
+from adaptivepy.config.schema import (
+    ClusteringConfig,
+    RunConfig,
+    SeedSelectionConfig,
+    config_to_dict,
+    load_config,
+)
+
+__all__ = [
+    "ClusteringConfig",
+    "RunConfig",
+    "SeedSelectionConfig",
+    "config_to_dict",
+    "load_config",
+]

adaptivepy/config/schema.py

@@ -0,0 +1,196 @@
+"""Configuration schema and validation for AdaptivePy runs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+
+DEFAULT_CLUSTERING_METHOD = "kmeans"
+DEFAULT_SEED_SELECTION = "nearest_center"
+DEFAULT_N_SEEDS = 10
+DEFAULT_RANDOM_SEED = 42
+
+
+@dataclass
+class ClusteringConfig:
+    """Clustering backend configuration.
+
+    Attributes
+    ----------
+    method : str
+        Clustering method name: ``kmeans``, ``minibatch_kmeans``, or
+        ``regular_space``.
+    n_clusters : int
+        Number of clusters to fit.
+    params : dict
+        Additional keyword arguments passed to the clusterer.
+    """
+
+    method: str = DEFAULT_CLUSTERING_METHOD
+    n_clusters: int = 10
+    params: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class SeedSelectionConfig:
+    """Frame-level seed selection configuration.
+
+    Attributes
+    ----------
+    method : str
+        Selection method: ``nearest_center`` or ``random_frame``.
+    """
+
+    method: str = DEFAULT_SEED_SELECTION
+
+
+@dataclass
+class RunConfig:
+    """Full configuration for an adaptive sampling run.
+
+    Attributes
+    ----------
+    features_dir : Path
+        Directory containing ``*.npy`` feature files.
+    output_dir : Path
+        Directory where results are written.
+    trajectories_dir : Path or None
+        Optional directory containing coordinate trajectories.
+    topology : Path or None
+        Topology file required when trajectories are provided.
+    clustering : ClusteringConfig
+        Clustering settings.
+    policies : list of str
+        Policy names to evaluate in parallel.
+    n_seeds : int
+        Number of seed frames to select per policy.
+    seed_selection : SeedSelectionConfig
+        Frame selection method within chosen clusters.
+    random_seed : int
+        Global random seed for reproducibility.
+    write_pdbs : bool
+        Whether to write PDB files when trajectories are available.
+    """
+
+    features_dir: Path
+    output_dir: Path
+    trajectories_dir: Optional[Path] = None
+    topology: Optional[Path] = None
+    clustering: ClusteringConfig = field(default_factory=ClusteringConfig)
+    policies: List[str] = field(default_factory=lambda: ["least_counts"])
+    n_seeds: int = DEFAULT_N_SEEDS
+    seed_selection: SeedSelectionConfig = field(default_factory=SeedSelectionConfig)
+    random_seed: int = DEFAULT_RANDOM_SEED
+    write_pdbs: bool = True
+
+
+def load_config(path: str | Path) -> RunConfig:
+    """Load and parse a YAML run configuration file.
+
+    Parameters
+    ----------
+    path : str or Path
+        Path to the YAML configuration file.
+
+    Returns
+    -------
+    RunConfig
+        Parsed and validated configuration object.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the configuration file does not exist.
+    ValueError
+        If required fields are missing or invalid.
+    """
+    config_path = Path(path)
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Configuration file not found: {config_path}")
+
+    with config_path.open("r", encoding="utf-8") as handle:
+        raw = yaml.safe_load(handle) or {}
+
+    if "features_dir" not in raw:
+        raise ValueError("Configuration must specify 'features_dir'.")
+    if "output_dir" not in raw:
+        raise ValueError("Configuration must specify 'output_dir'.")
+
+    features_dir = Path(raw["features_dir"])
+    output_dir = Path(raw["output_dir"])
+    trajectories_dir = (
+        Path(raw["trajectories_dir"]) if raw.get("trajectories_dir") else None
+    )
+    topology = Path(raw["topology"]) if raw.get("topology") else None
+
+    if trajectories_dir is not None and topology is None:
+        raise ValueError(
+            "When 'trajectories_dir' is provided, 'topology' must also be set."
+        )
+
+    clustering_raw = raw.get("clustering", {})
+    clustering = ClusteringConfig(
+        method=clustering_raw.get("method", DEFAULT_CLUSTERING_METHOD),
+        n_clusters=int(clustering_raw.get("n_clusters", 10)),
+        params=dict(clustering_raw.get("params", {})),
+    )
+
+    seed_raw = raw.get("seed_selection", {})
+    seed_selection = SeedSelectionConfig(
+        method=seed_raw.get("method", DEFAULT_SEED_SELECTION),
+    )
+
+    policies = raw.get("policies", ["least_counts"])
+    if isinstance(policies, str):
+        policies = [policies]
+
+    return RunConfig(
+        features_dir=features_dir,
+        output_dir=output_dir,
+        trajectories_dir=trajectories_dir,
+        topology=topology,
+        clustering=clustering,
+        policies=policies,
+        n_seeds=int(raw.get("n_seeds", DEFAULT_N_SEEDS)),
+        seed_selection=seed_selection,
+        random_seed=int(raw.get("random_seed", DEFAULT_RANDOM_SEED)),
+        write_pdbs=bool(raw.get("write_pdbs", True)),
+    )
+
+
+def config_to_dict(config: RunConfig) -> Dict[str, Any]:
+    """Convert a :class:`RunConfig` to a plain dictionary for serialization.
+
+    Parameters
+    ----------
+    config : RunConfig
+        Configuration object to serialize.
+
+    Returns
+    -------
+    dict
+        YAML-serializable configuration dictionary.
+    """
+    result: Dict[str, Any] = {
+        "features_dir": str(config.features_dir),
+        "output_dir": str(config.output_dir),
+        "clustering": {
+            "method": config.clustering.method,
+            "n_clusters": config.clustering.n_clusters,
+            "params": config.clustering.params,
+        },
+        "policies": list(config.policies),
+        "n_seeds": config.n_seeds,
+        "seed_selection": {"method": config.seed_selection.method},
+        "random_seed": config.random_seed,
+        "write_pdbs": config.write_pdbs,
+    }
+    if config.trajectories_dir is not None:
+        result["trajectories_dir"] = str(config.trajectories_dir)
+    if config.topology is not None:
+        result["topology"] = str(config.topology)
+    return result

adaptivepy/io/__init__.py

@@ -0,0 +1,27 @@
+"""Input/output utilities for AdaptivePy."""
+
+from adaptivepy.io.loader import (
+    list_feature_files,
+    list_trajectory_files,
+    load_features,
+    validate_dataset,
+    validate_feature_trajectory_mapping,
+)
+from adaptivepy.io.trajectory import (
+    build_trajectory_map,
+    extract_frame,
+    load_trajectory,
+    validate_trajectory_frame_counts,
+)
+
+__all__ = [
+    "build_trajectory_map",
+    "extract_frame",
+    "list_feature_files",
+    "list_trajectory_files",
+    "load_features",
+    "load_trajectory",
+    "validate_dataset",
+    "validate_feature_trajectory_mapping",
+    "validate_trajectory_frame_counts",
+]

adaptivepy/io/loader.py

@@ -0,0 +1,267 @@
+"""Feature loading and dataset validation."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+from adaptivepy.models import Dataset, FrameRecord
+
+logger = logging.getLogger(__name__)
+
+
+def list_feature_files(features_dir: Path) -> List[Path]:
+    """List ``*.npy`` feature files in a directory, sorted by name.
+
+    Parameters
+    ----------
+    features_dir : Path
+        Directory containing feature arrays.
+
+    Returns
+    -------
+    list of Path
+        Sorted paths to feature files.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the directory does not exist.
+    ValueError
+        If no ``*.npy`` files are found.
+    """
+    features_dir = Path(features_dir)
+    if not features_dir.is_dir():
+        raise FileNotFoundError(f"Features directory not found: {features_dir}")
+
+    files = sorted(features_dir.glob("*.npy"))
+    if not files:
+        raise ValueError(f"No .npy feature files found in {features_dir}")
+    return files
+
+
+def list_trajectory_files(trajectories_dir: Path) -> List[Path]:
+    """List coordinate trajectory files supported by mdtraj.
+
+    Parameters
+    ----------
+    trajectories_dir : Path
+        Directory containing trajectory files.
+
+    Returns
+    -------
+    list of Path
+        Sorted paths to trajectory files (``.xtc``, ``.dcd``, ``.trr``).
+
+    Raises
+    ------
+    FileNotFoundError
+        If the directory does not exist.
+    ValueError
+        If no supported trajectory files are found.
+    """
+    trajectories_dir = Path(trajectories_dir)
+    if not trajectories_dir.is_dir():
+        raise FileNotFoundError(
+            f"Trajectories directory not found: {trajectories_dir}"
+        )
+
+    extensions = ("*.xtc", "*.dcd", "*.trr", "*.nc", "*.pdb")
+    files: List[Path] = []
+    for pattern in extensions:
+        files.extend(trajectories_dir.glob(pattern))
+    files = sorted(set(files), key=lambda p: p.name)
+
+    if not files:
+        raise ValueError(
+            f"No supported trajectory files found in {trajectories_dir}"
+        )
+    return files
+
+
+def _stem(path: Path) -> str:
+    """Return the filename stem without extension."""
+    return path.stem
+
+
+def validate_feature_trajectory_mapping(
+    feature_files: List[Path],
+    trajectory_files: Optional[List[Path]] = None,
+) -> None:
+    """Ensure feature and trajectory filenames match one-to-one.
+
+    Parameters
+    ----------
+    feature_files : list of Path
+        Feature ``.npy`` file paths.
+    trajectory_files : list of Path or None
+        Optional coordinate trajectory file paths.
+
+    Raises
+    ------
+    ValueError
+        If stems do not match exactly between features and trajectories.
+    """
+    if trajectory_files is None:
+        return
+
+    feature_stems = {_stem(f) for f in feature_files}
+    traj_stems = {_stem(t) for t in trajectory_files}
+
+    missing_traj = feature_stems - traj_stems
+    missing_features = traj_stems - feature_stems
+
+    if missing_traj or missing_features:
+        messages = []
+        if missing_traj:
+            messages.append(
+                f"Features without matching trajectories: {sorted(missing_traj)}"
+            )
+        if missing_features:
+            messages.append(
+                f"Trajectories without matching features: {sorted(missing_features)}"
+            )
+        raise ValueError("; ".join(messages))
+
+
+def load_features(features_dir: Path) -> Dataset:
+    """Load feature arrays from disk and build a :class:`Dataset`.
+
+    Each ``*.npy`` file must have shape ``(n_frames, n_features)``. Per-frame
+    ``FrameRecord`` objects are created while preserving trajectory identity.
+
+    Parameters
+    ----------
+    features_dir : Path
+        Directory containing feature files.
+
+    Returns
+    -------
+    Dataset
+        Loaded dataset with concatenated feature matrix and frame records.
+
+    Raises
+    ------
+    ValueError
+        If feature arrays have inconsistent dimensionality.
+    """
+    feature_files = list_feature_files(features_dir)
+    frames: List[FrameRecord] = []
+    feature_blocks: List[np.ndarray] = []
+    traj_index_map: Dict[int, Tuple[int, int]] = {}
+    traj_names: List[str] = []
+    global_offset = 0
+    n_features: Optional[int] = None
+
+    for traj_id, feature_path in enumerate(feature_files):
+        features = np.load(feature_path)
+        if features.ndim != 2:
+            raise ValueError(
+                f"Feature file {feature_path} must be 2D (n_frames, n_features), "
+                f"got shape {features.shape}"
+            )
+
+        if n_features is None:
+            n_features = features.shape[1]
+        elif features.shape[1] != n_features:
+            raise ValueError(
+                f"Inconsistent feature dimension in {feature_path}: "
+                f"expected {n_features}, got {features.shape[1]}"
+            )
+
+        n_frames = features.shape[0]
+        start_idx = global_offset
+        end_idx = global_offset + n_frames
+
+        for frame_id in range(n_frames):
+            global_index = global_offset + frame_id
+            frames.append(
+                FrameRecord(
+                    traj_id=traj_id,
+                    frame_id=frame_id,
+                    features=features[frame_id],
+                    global_index=global_index,
+                )
+            )
+
+        feature_blocks.append(features)
+        traj_index_map[traj_id] = (start_idx, end_idx)
+        traj_names.append(_stem(feature_path))
+        global_offset = end_idx
+
+        logger.info(
+            "Loaded %s: %d frames, %d features",
+            feature_path.name,
+            n_frames,
+            n_features,
+        )
+
+    feature_matrix = (
+        np.vstack(feature_blocks) if feature_blocks else np.empty((0, 0))
+    )
+
+    return Dataset(
+        frames=frames,
+        feature_matrix=feature_matrix,
+        traj_index_map=traj_index_map,
+        traj_names=traj_names,
+    )
+
+
+def validate_dataset(
+    dataset: Dataset,
+    trajectory_files: Optional[List[Path]] = None,
+) -> None:
+    """Run consistency checks on a loaded dataset.
+
+    Parameters
+    ----------
+    dataset : Dataset
+        Dataset to validate.
+    trajectory_files : list of Path or None
+        Optional trajectory files for cross-validation.
+
+    Raises
+    ------
+    ValueError
+        If internal consistency checks fail.
+    """
+    if dataset.feature_matrix is None or len(dataset.frames) == 0:
+        raise ValueError("Dataset is empty.")
+
+    n_frames, n_features = dataset.feature_matrix.shape
+    if n_frames != len(dataset.frames):
+        raise ValueError(
+            "Feature matrix row count does not match number of frame records."
+        )
+
+    for record in dataset.frames:
+        if record.features.shape != (n_features,):
+            raise ValueError(
+                f"Frame ({record.traj_id}, {record.frame_id}) has invalid "
+                f"feature shape {record.features.shape}."
+            )
+
+    if trajectory_files is not None:
+        feature_stems = set(dataset.traj_names)
+        traj_stems = {path.stem for path in trajectory_files}
+        missing_traj = feature_stems - traj_stems
+        missing_features = traj_stems - feature_stems
+        if missing_traj or missing_features:
+            messages = []
+            if missing_traj:
+                messages.append(
+                    f"Features without matching trajectories: {sorted(missing_traj)}"
+                )
+            if missing_features:
+                messages.append(
+                    f"Trajectories without matching features: {sorted(missing_features)}"
+                )
+            raise ValueError("; ".join(messages))
+        if len(trajectory_files) != len(dataset.traj_names):
+            raise ValueError(
+                "Number of trajectory files must match number of feature files."
+            )