bfflearn 0.1.0.dev1__py3-none-any.whl

bff/domain/bias.py

@@ -0,0 +1,149 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Mapping
+
+PathLike = str | Path
+
+
+@dataclass(frozen=True, slots=True)
+class BiasSpec:
+    """Bias specification for one simulation window.
+
+    Parameters
+    ----------
+    kind
+        Bias type. Supported values are ``"none"``, ``"colvars"``, and
+        ``"plumed"``.
+    colvars_file
+        Path to the user-supplied COLVARS input file.
+    plumed_file
+        Path to the user-supplied PLUMED input file.
+    """
+
+    kind: str = "none"
+    colvars_file: Path | None = None
+    plumed_file: Path | None = None
+
+    def __post_init__(self) -> None:
+        colvars_file = None if self.colvars_file is None else Path(self.colvars_file)
+        plumed_file = None if self.plumed_file is None else Path(self.plumed_file)
+
+        if self.kind not in {"none", "colvars", "plumed"}:
+            raise ValueError(
+                f"Unsupported bias kind {self.kind!r}. "
+                "Supported values are 'none', 'colvars', and 'plumed'."
+            )
+        if colvars_file is not None and plumed_file is not None:
+            raise ValueError(
+                "Bias specification must define at most one of 'colvars_file' "
+                "and 'plumed_file'."
+            )
+        if self.kind == "colvars" and colvars_file is None:
+            raise ValueError("COLVARS bias requires 'colvars_file'.")
+        if self.kind == "plumed" and plumed_file is None:
+            raise ValueError("PLUMED bias requires 'plumed_file'.")
+        if self.kind == "none" and (
+            colvars_file is not None or plumed_file is not None
+        ):
+            raise ValueError("Unbiased systems cannot define bias input files.")
+        if colvars_file is not None and not colvars_file.exists():
+            raise FileNotFoundError(f"COLVARS file not found: {colvars_file}")
+        if plumed_file is not None and not plumed_file.exists():
+            raise FileNotFoundError(f"PLUMED file not found: {plumed_file}")
+
+        object.__setattr__(self, "colvars_file", colvars_file)
+        object.__setattr__(self, "plumed_file", plumed_file)
+
+    @classmethod
+    def from_any(
+        cls,
+        value: "BiasSpec | Mapping[str, Any] | PathLike | None",
+        *,
+        base_dir: Path | None = None,
+    ) -> "BiasSpec":
+        """Normalize a raw config value into a ``BiasSpec``.
+
+        Parameters
+        ----------
+        value
+            Existing ``BiasSpec`` instance, raw mapping, or ``None``.
+        base_dir
+            Base directory used to resolve relative COLVARS file paths.
+
+        Returns
+        -------
+        BiasSpec
+            Normalized bias specification.
+        """
+        if isinstance(value, cls):
+            return value
+        if value is None:
+            return cls()
+        if isinstance(value, (str, Path)):
+            if base_dir is not None:
+                value = Path(value)
+                if not value.is_absolute():
+                    value = (base_dir / value).resolve()
+            return cls.load(value)
+        if not isinstance(value, Mapping):
+            raise ValueError(f"Invalid bias specification: {value!r}")
+
+        colvars_raw = value.get("colvars_file", value.get("fn_colvars"))
+        plumed_raw = value.get("plumed_file", value.get("fn_plumed"))
+        colvars_file = None
+        plumed_file = None
+        if colvars_raw is not None:
+            colvars_file = Path(colvars_raw)
+            if base_dir is not None and not colvars_file.is_absolute():
+                colvars_file = (base_dir / colvars_file).resolve()
+        if plumed_raw is not None:
+            plumed_file = Path(plumed_raw)
+            if base_dir is not None and not plumed_file.is_absolute():
+                plumed_file = (base_dir / plumed_file).resolve()
+
+        if value.get("kind") is not None:
+            kind = str(value["kind"])
+        elif colvars_file is not None:
+            kind = "colvars"
+        elif plumed_file is not None:
+            kind = "plumed"
+        else:
+            kind = "none"
+        return cls(
+            kind=kind,
+            colvars_file=colvars_file,
+            plumed_file=plumed_file,
+        )
+
+    @property
+    def is_biased(self) -> bool:
+        """Whether this specification defines any simulation bias."""
+        return self.kind != "none"
+
+    @property
+    def input_file(self) -> Path | None:
+        """Bias input file for the configured engine, if any."""
+        return self.colvars_file if self.kind == "colvars" else self.plumed_file
+
+    @property
+    def input_filename(self) -> str | None:
+        """Canonical filename used when copying a bias input file."""
+        if self.kind == "colvars":
+            return "bias.colvars.dat"
+        if self.kind == "plumed":
+            return "bias.plumed.dat"
+        return None
+
+    @classmethod
+    def load(cls, fn_in: PathLike) -> "BiasSpec":
+        """Load a bias specification from a direct bias-input file."""
+        fn_in = Path(fn_in).resolve()
+        name = fn_in.name
+        if name.endswith(".colvars.dat"):
+            return cls(kind="colvars", colvars_file=fn_in)
+        if name.endswith(".plumed.dat"):
+            return cls(kind="plumed", plumed_file=fn_in)
+        raise ValueError(
+            f"Unsupported bias file {fn_in}. Expected '.colvars.dat' "
+            "or '.plumed.dat'."
+        )

bff/domain/specs.py

@@ -0,0 +1,347 @@
+from dataclasses import InitVar, dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from scipy.stats.qmc import LatinHypercube
+
+from ..io.utils import load_yaml, save_yaml
+
+try:
+    import torch
+except ModuleNotFoundError:
+    torch = None
+
+
+PathLike = Union[str, Path]
+ArrayLike = Union[np.ndarray, Any]
+
+
+def _is_torch_tensor(values: object) -> bool:
+    return torch is not None and isinstance(values, torch.Tensor)
+
+
+@dataclass(frozen=True)
+class Bounds:
+    """Named parameter bounds with a stable sorted order."""
+
+    by_name: Mapping[str, Tuple[float, float]]
+    _items: tuple[tuple[str, tuple[float, float]], ...] = field(
+        init=False,
+        repr=False,
+    )
+
+    def __post_init__(self) -> None:
+        items = tuple(sorted(self.by_name.items()))
+        for name, (lower, upper) in items:
+            if lower > upper:
+                raise ValueError(
+                    f"Lower bound {lower} is greater than upper bound {upper} "
+                    f"for parameter {name!r}."
+                )
+        object.__setattr__(self, "_items", items)
+
+    @property
+    def names(self) -> np.ndarray:
+        return np.array([name for name, _ in self._items], dtype=str)
+
+    @property
+    def array(self) -> np.ndarray:
+        return np.asarray([bounds for _, bounds in self._items], dtype=float)
+
+    @property
+    def lower(self) -> np.ndarray:
+        return self.array[:, 0]
+
+    @property
+    def upper(self) -> np.ndarray:
+        return self.array[:, 1]
+
+    @property
+    def n_params(self) -> int:
+        return len(self._items)
+
+    def get(self, name: str) -> tuple[float, float]:
+        return self.by_name[name]
+
+    def index(self, name: str) -> int:
+        names = self.names
+        mask = names == name
+        if not np.any(mask):
+            raise ValueError(f"Parameter {name!r} not found in bounds.")
+        return int(np.argwhere(mask).ravel()[0])
+
+    def without(self, name: str) -> "Bounds":
+        items = {
+            key: value for key, value in self.by_name.items() if key != name
+        }
+        return Bounds(items)
+
+    def to_dict(self) -> dict[str, list[float]]:
+        return {
+            name: [float(lower), float(upper)]
+            for name, (lower, upper) in self._items
+        }
+
+
+@dataclass(frozen=True)
+class Specs:
+    """Force-field specification data."""
+
+    source: InitVar[dict[str, Any] | PathLike]
+
+    mol_resname: str = field(init=False)
+    bounds: Bounds = field(init=False)
+    total_charge: float = field(init=False)
+    constraint_charge: float = field(init=False)
+    implicit_atoms: tuple[str, ...] = field(init=False)
+
+    def __post_init__(self, source: dict[str, Any] | PathLike) -> None:
+        if isinstance(source, dict):
+            data = dict(source)
+        elif isinstance(source, (str, Path)):
+            data = load_yaml(source)
+        else:
+            raise TypeError(f"Unsupported source type: {type(source)}")
+
+        required = {"bounds", "implicit_atoms", "charge_target"}
+        missing = required - set(data)
+        if missing:
+            missing_list = ", ".join(sorted(repr(key) for key in missing))
+            raise ValueError(f"Missing required specs field(s): {missing_list}")
+
+        implicit_atoms = data["implicit_atoms"]
+        if isinstance(implicit_atoms, str):
+            implicit_atoms = implicit_atoms.split()
+
+        object.__setattr__(self, "mol_resname", str(data.get("mol_resname", "")))
+        object.__setattr__(self, "bounds", Bounds(data["bounds"]))
+        object.__setattr__(self, "total_charge", float(data.get("total_charge", 0.0)))
+        object.__setattr__(self, "constraint_charge", float(data["charge_target"]))
+        object.__setattr__(self, "implicit_atoms", tuple(implicit_atoms))
+
+    @classmethod
+    def load(cls, source: PathLike) -> "Specs":
+        return cls(source)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "mol_resname": self.mol_resname,
+            "bounds": self.bounds.to_dict(),
+            "total_charge": self.total_charge,
+            "charge_target": self.constraint_charge,
+            "implicit_atoms": list(self.implicit_atoms),
+        }
+
+    def write(self, fn_out: PathLike) -> None:
+        save_yaml(self.to_dict(), fn_out)
+
+    @property
+    def implicit_param(self) -> str:
+        return f"charge {' '.join(self.implicit_atoms)}"
+
+    @property
+    def implicit_param_index(self) -> int:
+        return self.bounds.index(self.implicit_param)
+
+    @property
+    def explicit_bounds(self) -> Bounds:
+        return self.bounds.without(self.implicit_param)
+
+    def parameter_names(self, *, explicit_only: bool = False) -> tuple[str, ...]:
+        bounds = self.explicit_bounds if explicit_only else self.bounds
+        return tuple(str(name) for name in bounds.names)
+
+    def parameter_dict(
+        self,
+        values: Sequence[float] | np.ndarray | Any,
+        *,
+        explicit_only: bool = False,
+    ) -> dict[str, float]:
+        array = (
+            values.detach().cpu().numpy()
+            if _is_torch_tensor(values)
+            else np.asarray(values, dtype=float)
+        ).reshape(-1)
+        names = self.parameter_names(explicit_only=explicit_only)
+        if array.size != len(names):
+            raise ValueError(
+                f"Expected {len(names)} parameter values, got {array.size}."
+            )
+        return {name: float(value) for name, value in zip(names, array)}
+
+    @property
+    def explicit_charge_coefficients(self) -> np.ndarray:
+        coeffs: list[int] = []
+        for param in self.explicit_bounds.names:
+            atoms = param.split()[1:]
+            coeffs.append(len(atoms) if param.startswith("charge") else 0)
+        return np.asarray(coeffs, dtype=int)
+
+    @property
+    def constraint_matrix(self) -> np.ndarray:
+        return self.explicit_charge_coefficients
+
+    def implicit_charge(
+        self,
+        values: Sequence[float] | np.ndarray | Any,
+    ) -> np.ndarray:
+        array = (
+            values.detach().cpu().numpy()
+            if _is_torch_tensor(values)
+            else np.asarray(values, dtype=float)
+        )
+        array = np.atleast_2d(array)
+        n_explicit = self.explicit_bounds.n_params
+        if array.shape[1] < n_explicit:
+            raise ValueError(
+                f"Expected at least {n_explicit} explicit parameter columns, "
+                f"got {array.shape[1]}."
+            )
+        explicit = array[:, :n_explicit]
+        explicit_charge = explicit @ self.constraint_matrix
+        return (self.constraint_charge - explicit_charge) / len(self.implicit_atoms)
+
+    def with_implicit_charge(
+        self,
+        values: Sequence[float] | np.ndarray | Any,
+    ) -> np.ndarray:
+        array = (
+            values.detach().cpu().numpy()
+            if _is_torch_tensor(values)
+            else np.asarray(values, dtype=float)
+        )
+        array = np.atleast_2d(array)
+        n_explicit = self.explicit_bounds.n_params
+        if array.shape[1] < n_explicit:
+            raise ValueError(
+                f"Expected at least {n_explicit} explicit parameter columns, "
+                f"got {array.shape[1]}."
+            )
+
+        implicit = self.implicit_charge(array)[:, None]
+        insert_at = self.implicit_param_index
+        return np.concatenate(
+            [
+                array[:, :insert_at],
+                implicit,
+                array[:, insert_at:],
+            ],
+            axis=1,
+        )
+
+
+@dataclass(frozen=True)
+class ChargeConstraint:
+    """Charge-aware validity check derived from ``Specs``."""
+
+    specs: Specs
+
+    def __init__(self, specs: Specs | dict[str, Any] | PathLike) -> None:
+        object.__setattr__(
+            self,
+            "specs",
+            specs if isinstance(specs, Specs) else Specs(specs),
+        )
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.specs, name)
+
+    @property
+    def explicit_bounds(self) -> np.ndarray:
+        return self.specs.explicit_bounds.array
+
+    @property
+    def implicit_bounds(self) -> tuple[float, float]:
+        return self.specs.bounds.get(self.specs.implicit_param)
+
+    @property
+    def n_params(self) -> int:
+        return self.specs.explicit_bounds.n_params
+
+    @staticmethod
+    def _to_2d_array(values: ArrayLike) -> np.ndarray:
+        if _is_torch_tensor(values):
+            arr = values.detach().cpu().numpy()
+        else:
+            arr = np.asarray(values, dtype=float)
+        return np.atleast_2d(arr)
+
+    def is_valid(self, values: ArrayLike) -> np.ndarray:
+        x = self._to_2d_array(values)
+        if x.shape[1] != self.n_params:
+            raise ValueError(
+                f"Input array must have shape (n_samples, {self.n_params}), "
+                f"but got {x.shape}."
+            )
+
+        explicit_lower, explicit_upper = self.explicit_bounds.T
+        implicit_lower, implicit_upper = self.implicit_bounds
+
+        in_explicit_bounds = ((x >= explicit_lower) & (x <= explicit_upper)).all(axis=1)
+        explicit_total_charge = np.sum(
+            x * self.specs.explicit_charge_coefficients,
+            axis=1,
+        )
+        implicit_charge = self.specs.constraint_charge - explicit_total_charge
+        in_implicit_bounds = (
+            (implicit_charge >= implicit_lower) & (implicit_charge <= implicit_upper)
+        )
+        return in_explicit_bounds & in_implicit_bounds
+
+    def __call__(self, values: ArrayLike) -> np.ndarray:
+        return self.is_valid(values)
+
+
+class RandomParamsGenerator:
+    """Latin-hypercube parameter generator with optional validity filter."""
+
+    def __init__(
+        self,
+        bounds: np.ndarray,
+        constraint: Optional[Callable[[ArrayLike], np.ndarray]] = None,
+    ) -> None:
+        if bounds.ndim != 2 or bounds.shape[1] != 2:
+            raise ValueError("Bounds must have shape (n_params, 2).")
+
+        self.bounds = np.asarray(bounds, dtype=float)
+        self.constraint = constraint
+        self.n_generated = 0
+        self.sampler = LatinHypercube(self.bounds.shape[0])
+
+    def __call__(self, n: int) -> np.ndarray:
+        if n < 0:
+            raise ValueError("Number of samples must be non-negative.")
+        if n == 0:
+            return np.empty((0, self.bounds.shape[0]), dtype=float)
+
+        lower, upper = self.bounds.T
+        if self.constraint is None:
+            unit_samples = self.sampler.random(n)
+            self.n_generated += n
+            return unit_samples * (upper - lower) + lower
+
+        collected: list[np.ndarray] = []
+        n_valid = 0
+        attempts = 0
+        while n_valid < n:
+            batch_size = max(2 * (n - n_valid), 1)
+            unit_samples = self.sampler.random(batch_size)
+            self.n_generated += batch_size
+            attempts += 1
+
+            samples = unit_samples * (upper - lower) + lower
+            mask = np.asarray(self.constraint(samples), dtype=bool).reshape(-1)
+            valid = samples[mask]
+            if valid.size > 0:
+                collected.append(valid)
+                n_valid += len(valid)
+                attempts = 0
+
+            if attempts >= 1000:
+                raise RuntimeError(
+                    "Failed to generate valid parameter samples within 1000 "
+                    "consecutive Latin-hypercube batches."
+                )
+
+        return np.vstack(collected)[:n]

bff/domain/trainset.py

@@ -0,0 +1,201 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+
+from ..io.utils import extract_train_dir, prepare_path
+from .specs import Specs
+
+PathLike = Union[str, Path]
+
+
+@dataclass(frozen=True, slots=True)
+class SimulationSystem:
+    """One staged simulation system inside a trainset directory."""
+
+    system_id: str
+    fn_topol: Path
+    fn_coord: Path
+
+
+@dataclass(frozen=True, slots=True)
+class TrajectorySet:
+    """One analyzed sample or reference set spanning one or more trajectories."""
+
+    sample_id: str
+    fn_topol: tuple[Path, ...]
+    fn_coord: tuple[Path, ...]
+    fn_trj: tuple[Path, ...]
+    params: np.ndarray | None = None
+
+    def __post_init__(self) -> None:
+        if not (len(self.fn_topol) == len(self.fn_coord) == len(self.fn_trj)):
+            raise ValueError(
+                "TrajectorySet requires the same number of topology, coordinate, "
+                "and trajectory files."
+            )
+        if self.params is not None:
+            object.__setattr__(
+                self,
+                "params",
+                np.asarray(self.params, dtype=float).reshape(-1),
+            )
+
+
+@dataclass(slots=True)
+class TrainSetInfo:
+    """Completed simulation samples and staged systems for one trainset."""
+
+    train_dir: Path
+    specs: Specs
+    systems: list[SimulationSystem]
+    samples: list[TrajectorySet]
+
+    @classmethod
+    def from_dir(
+        cls,
+        train_dir: PathLike,
+        *,
+        strict: bool = True,
+    ) -> "TrainSetInfo":
+        prepared_dir = prepare_path(Path(train_dir).resolve())
+        specs_data, systems_data, samples_data = extract_train_dir(prepared_dir)
+        if specs_data is None:
+            raise ValueError("Training set is missing specs information.")
+
+        specs = Specs(specs_data)
+        systems = [
+            SimulationSystem(
+                system_id=str(system["system_id"]),
+                fn_topol=prepared_dir / system["topology"],
+                fn_coord=prepared_dir / system["coordinates"],
+            )
+            for system in systems_data
+        ]
+
+        completed_samples: list[TrajectorySet] = []
+        sample_issues: list[str] = []
+        for sample_id in sorted((samples_data or {})):
+            sample = samples_data[sample_id]
+            if not isinstance(sample, dict):
+                sample_issues.append(
+                    f"sample {sample_id!r}: record must be a mapping, got "
+                    f"{type(sample).__name__}."
+                )
+                continue
+
+            if sample.get("status") != "completed":
+                continue
+            outputs = sample.get("outputs", [])
+            if "params" not in sample:
+                sample_issues.append(
+                    f"sample {sample_id!r}: completed sample is missing 'params'."
+                )
+                continue
+            if not isinstance(outputs, list):
+                sample_issues.append(
+                    f"sample {sample_id!r}: 'outputs' must be a list."
+                )
+                continue
+            if any(not isinstance(output, dict) for output in outputs):
+                sample_issues.append(
+                    f"sample {sample_id!r}: all outputs must be mappings."
+                )
+                continue
+            if any("system_id" not in output for output in outputs):
+                sample_issues.append(
+                    f"sample {sample_id!r}: each output must define 'system_id'."
+                )
+                continue
+            outputs = sorted(outputs, key=lambda item: item["system_id"])
+            if len(outputs) != len(systems):
+                sample_issues.append(
+                    f"sample {sample_id!r}: expected {len(systems)} outputs, "
+                    f"found {len(outputs)}."
+                )
+                continue
+            if any(
+                output.get("trajectory") is None
+                for output in outputs
+            ):
+                sample_issues.append(
+                    f"sample {sample_id!r}: completed outputs must all define "
+                    "trajectory files."
+                )
+                continue
+
+            fn_trj = tuple(prepared_dir / output["trajectory"] for output in outputs)
+            if not all(path.exists() for path in fn_trj):
+                missing = [
+                    str(path.relative_to(prepared_dir))
+                    for path in fn_trj
+                    if not path.exists()
+                ]
+                sample_issues.append(
+                    f"sample {sample_id!r}: missing trajectory file(s): "
+                    + ", ".join(missing)
+                )
+                continue
+
+            completed_samples.append(
+                TrajectorySet(
+                    sample_id=str(sample_id),
+                    fn_topol=tuple(system.fn_topol for system in systems),
+                    fn_coord=tuple(system.fn_coord for system in systems),
+                    fn_trj=fn_trj,
+                    params=np.asarray(sample["params"], dtype=float),
+                )
+            )
+
+        if strict and sample_issues:
+            issues = "\n".join(f"- {issue}" for issue in sample_issues)
+            raise ValueError(
+                "Training set contains invalid completed sample records:\n"
+                f"{issues}"
+            )
+        if strict and not completed_samples:
+            raise ValueError(
+                f"No completed training samples found in {prepared_dir}. "
+                "Check that validation/simulation finished and samples.yaml "
+                "contains completed outputs."
+            )
+
+        return cls(
+            train_dir=prepared_dir,
+            specs=specs,
+            systems=systems,
+            samples=completed_samples,
+        )
+
+    @property
+    def sample_ids(self) -> list[str]:
+        return [sample.sample_id for sample in self.samples]
+
+    @property
+    def hashes(self) -> list[str]:
+        return self.sample_ids
+
+    @property
+    def inputs(self) -> np.ndarray:
+        if not self.samples:
+            return np.empty((0, 0), dtype=float)
+        return np.asarray([sample.params for sample in self.samples], dtype=float)
+
+    @property
+    def n_samples(self) -> int:
+        return len(self.samples)
+
+    @property
+    def fn_topol(self) -> list[Path]:
+        return [system.fn_topol for system in self.systems]
+
+    @property
+    def fn_coord(self) -> list[Path]:
+        return [system.fn_coord for system in self.systems]
+
+    @property
+    def fn_trj(self) -> list[list[Path]]:
+        if not self.samples:
+            return []
+        return [list(sample.fn_trj) for sample in self.samples]