bfflearn 0.1.0.dev1__py3-none-any.whl

bff/qoi/data.py

@@ -0,0 +1,169 @@
+from dataclasses import dataclass, field
+from typing import Any, Mapping
+
+import numpy as np
+
+from ..io.utils import load_pt, save_pt
+
+
+@dataclass(slots=True)
+class QoI:
+    """One named quantity of interest produced for a single trajectory."""
+
+    name: str
+    values: np.ndarray
+    labels: tuple[str, ...] | None = None
+    values_per_label: int = 1
+    settings: dict[str, Any] = field(default_factory=dict)
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        values = np.asarray(self.values, dtype=float).reshape(-1)
+        labels = None if self.labels is None else tuple(self.labels)
+        values_per_label = int(self.values_per_label)
+
+        if values_per_label <= 0:
+            raise ValueError("'values_per_label' must be a positive integer.")
+        if labels is not None and len(values) != len(labels) * values_per_label:
+            raise ValueError(
+                "Number of values does not match labels * values_per_label in QoI."
+            )
+
+        self.values = values
+        self.labels = labels
+        self.values_per_label = values_per_label
+        self.settings = dict(self.settings)
+        self.metadata = dict(self.metadata)
+
+    @property
+    def n_values(self) -> int:
+        """Total number of numeric values stored in the QoI."""
+        return int(self.values.size)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Serialize the QoI to a JSON/PT-friendly mapping."""
+        return {
+            "name": self.name,
+            "values": self.values.tolist(),
+            "labels": None if self.labels is None else list(self.labels),
+            "values_per_label": self.values_per_label,
+            "settings": dict(self.settings),
+            "metadata": dict(self.metadata),
+        }
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "QoI":
+        """Rebuild a serialized QoI object."""
+        return cls(
+            name=str(data["name"]),
+            values=np.asarray(data["values"], dtype=float),
+            labels=None if data.get("labels") is None else tuple(data["labels"]),
+            values_per_label=int(data.get("values_per_label", 1)),
+            settings=dict(data.get("settings", {})),
+            metadata=dict(data.get("metadata", {})),
+        )
+
+
+@dataclass(slots=True)
+class QoIDataset:
+    """Training-ready dataset for one named quantity of interest."""
+
+    name: str
+    inputs: np.ndarray
+    outputs: np.ndarray
+    outputs_ref: np.ndarray
+    labels: tuple[str, ...] | None = None
+    values_per_label: int = 1
+    nuisance: float | None = None
+    settings: dict[str, Any] | None = None
+    metadata: dict[str, Any] | None = None
+
+    def __post_init__(self) -> None:
+        self.inputs = np.asarray(self.inputs, dtype=float)
+        self.outputs = np.asarray(self.outputs, dtype=float)
+        self.outputs_ref = np.asarray(self.outputs_ref, dtype=float).reshape(-1)
+        self.settings = self._coerce_mapping(self.settings)
+        self.metadata = self._coerce_mapping(self.metadata)
+        self.labels = None if self.labels is None else tuple(self.labels)
+        self.values_per_label = int(self.values_per_label)
+
+        if self.inputs.shape[0] != self.outputs.shape[0]:
+            raise ValueError(
+                f"Number of input samples ({self.inputs.shape[0]}) does not match "
+                f"number of output samples ({self.outputs.shape[0]})."
+            )
+
+        if self.outputs.shape[1] != self.outputs_ref.shape[0]:
+            raise ValueError(
+                f"Output dimension ({self.outputs.shape[1]}) does not match "
+                f"reference dimension ({self.outputs_ref.shape[0]})."
+            )
+        if self.values_per_label <= 0:
+            raise ValueError("'values_per_label' must be a positive integer.")
+        if (
+            self.labels is not None
+            and self.outputs_ref.size != len(self.labels) * self.values_per_label
+        ):
+            raise ValueError(
+                "Reference output size does not match labels * values_per_label."
+            )
+
+    @staticmethod
+    def _coerce_mapping(value: Any) -> dict[str, Any]:
+        return dict(value or {})
+
+    @property
+    def n_samples(self) -> int:
+        return int(self.inputs.shape[0])
+
+    @property
+    def n_observations(self) -> int:
+        if self.outputs_ref.size % self.values_per_label != 0:
+            raise ValueError(
+                "Reference output size must be divisible by "
+                "'values_per_label'."
+            )
+        return int(self.outputs_ref.size // self.values_per_label)
+
+    def to_dict(self) -> dict[str, Any]:
+        data = {
+            "name": self.name,
+            "inputs": self.inputs.tolist(),
+            "outputs": self.outputs.tolist(),
+            "outputs_ref": self.outputs_ref.tolist(),
+            "labels": None if self.labels is None else list(self.labels),
+            "values_per_label": self.values_per_label,
+            "settings": dict(self.settings),
+            "metadata": dict(self.metadata),
+        }
+        if self.nuisance is not None:
+            data["nuisance"] = self.nuisance
+        return data
+
+    @classmethod
+    def from_dict(cls, data: Mapping[str, Any]) -> "QoIDataset":
+        return cls(
+            name=str(data["name"]),
+            inputs=np.asarray(data["inputs"], dtype=float),
+            outputs=np.asarray(data["outputs"], dtype=float),
+            outputs_ref=np.asarray(data["outputs_ref"], dtype=float),
+            labels=None if data.get("labels") is None else tuple(data["labels"]),
+            values_per_label=int(data.get("values_per_label", 1)),
+            nuisance=data.get("nuisance"),
+            settings=dict(data.get("settings", {})),
+            metadata=dict(data.get("metadata", {})),
+        )
+
+    def write(self, fn_out: str) -> None:
+        save_pt(self.to_dict(), fn_out)
+
+    @classmethod
+    def load(cls, fn_in: str) -> "QoIDataset":
+        data = load_pt(fn_in)
+        return cls.from_dict(data)
+
+    def __repr__(self) -> str:
+        return (
+            f"QoIDataset(n_samples={self.n_samples}, "
+            f"n_observations={self.n_observations}, name={self.name!r})"
+        )

bff/qoi/hbonds.py

@@ -0,0 +1,255 @@
+from typing import Dict
+
+import MDAnalysis as mda
+import numpy as np
+from MDAnalysis.lib.distances import calc_angles, capped_distance
+
+from ..tools import get_unitcell
+from .data import QoI
+
+__all__ = ["compute_all_hbonds"]
+
+
+def count_hbonds(
+    universe: mda.Universe,
+    donor_indices: np.ndarray,
+    acceptor_indices: np.ndarray,
+    n_frames: int
+) -> Dict[str, float]:
+    """
+    Counts the average number of hydrogen bonds per hydrogen bond type.
+
+    Parameters
+    ----------
+    universe : MDAnalysis.Universe
+    donor_indices : np.ndarray
+    acceptor_indices : np.ndarray
+    n_frames : int
+        Number of trajectory frames used for the analysis.
+    threshold : float, optional
+        Minimum frequency threshold for a hydrogen bond type. Default is 0.1.
+
+    Returns
+    -------
+    results : dict
+        Dictionary containing the hydrogen bond types and their frequencies.
+    """
+
+    if n_frames <= 0 or donor_indices.size == 0 or acceptor_indices.size == 0:
+        return {}
+
+    hbonds = np.column_stack(
+        (
+            universe.atoms[donor_indices].resnames,
+            universe.atoms[donor_indices].types,
+            universe.atoms[acceptor_indices].resnames,
+            universe.atoms[acceptor_indices].types,
+        )
+    ).astype("str")
+
+    hbond_types, counts = np.unique(hbonds, return_counts=True, axis=0)
+
+    # Filter and format results
+    results = {}
+    for hb, count in zip(hbond_types, counts):
+        name = f"{hb[0]}({hb[1]}) to {hb[2]}({hb[3]})"
+        results[name] = count / n_frames
+
+    return results
+
+
+def _paired_donors_and_hydrogens(
+    universe: mda.Universe,
+    selection: str,
+    *,
+    hb_elements: set[str],
+) -> tuple[mda.AtomGroup, mda.AtomGroup]:
+    """Return donor and hydrogen groups with one-to-one indexing."""
+    hydrogens_all = universe.select_atoms(f"{selection} and element H")
+    donor_indices: list[int] = []
+    hydrogen_indices: list[int] = []
+
+    for hydrogen in hydrogens_all:
+        if not hydrogen.bonded_atoms:
+            continue
+        donor = hydrogen.bonded_atoms[0]
+        if donor.element not in hb_elements:
+            continue
+        donor_indices.append(donor.index)
+        hydrogen_indices.append(hydrogen.index)
+
+    if not donor_indices:
+        empty = universe.atoms[np.asarray([], dtype=int)]
+        return empty, empty
+
+    donors = universe.atoms[np.asarray(donor_indices, dtype=int)]
+    hydrogens = universe.atoms[np.asarray(hydrogen_indices, dtype=int)]
+    return donors, hydrogens
+
+
+def _all_possible_hbond_labels(
+    donors: mda.AtomGroup,
+    acceptors: mda.AtomGroup,
+) -> set[str]:
+    """Enumerate all donor/acceptor type combinations possible for this system."""
+    if len(donors) == 0 or len(acceptors) == 0:
+        return set()
+
+    donor_types = np.unique(
+        np.column_stack((donors.resnames, donors.types)).astype(str),
+        axis=0,
+    )
+    acceptor_types = np.unique(
+        np.column_stack((acceptors.resnames, acceptors.types)).astype(str),
+        axis=0,
+    )
+
+    return {
+        f"{d_res}({d_type}) to {a_res}({a_type})"
+        for d_res, d_type in donor_types
+        for a_res, a_type in acceptor_types
+    }
+
+
+def compute_hbonds(
+    universe: mda.Universe,
+    hydrogens: mda.AtomGroup,
+    donors: mda.AtomGroup,
+    acceptors: mda.AtomGroup,
+    distance_cutoff: float = 3.5,
+    angle_cutoff: float = 150,
+    start: int = 0,
+    step: int = 1,
+    stop: int = None,
+) -> tuple[np.ndarray, np.ndarray, int]:
+    """
+    Computes hydrogen bonds between donor and acceptor atoms.
+
+    Parameters
+    ----------
+    universe : MDAnalysis.Universe
+    hydrogens : MDAnalysis.AtomGroup
+    donors : MDAnalysis.AtomGroup
+    acceptors : MDAnalysis.AtomGroup
+    distance_cutoff : float, optional
+        Maximum distance between donor and acceptor atoms. Default is 3.0A.
+    angle_cutoff : float, optional
+        Minimum angle between donor-hydrogen-acceptor atoms.
+        Default is 150 degrees.
+    start : int, optional
+        Start frame for the analysis. Default is 0.
+    step : int, optional
+        Frame stride for the analysis. Default is 1.
+    stop : int, optional
+        End frame for the analysis. Default is None.
+
+    Returns
+    -------
+    donor_indices : np.ndarray
+    acceptor_indices : np.ndarray
+    n_frames : int
+    """
+
+    sl = slice(start, stop or len(universe.trajectory), step)
+    angle_cutoff = np.deg2rad(angle_cutoff)
+
+    donor_indices, acceptor_indices = [], []
+    n_frames = 0
+    for ts in universe.trajectory[sl]:
+        n_frames += 1
+        box = get_unitcell(universe, ts)
+        da_indices = capped_distance(
+            donors,
+            acceptors,
+            max_cutoff=distance_cutoff,
+            min_cutoff=1.0,
+            box=box,
+            return_distances=False,
+        )
+        if da_indices.size == 0:
+            continue
+
+        d = donors[da_indices[:, 0]]
+        h = hydrogens[da_indices[:, 0]]
+        a = acceptors[da_indices[:, 1]]
+
+        dha_angles = calc_angles(d, h, a, box=box)
+        hbond_indices = np.where(dha_angles > angle_cutoff)[0]
+        if hbond_indices.size == 0:
+            continue
+
+        donor_indices.extend(d[hbond_indices].indices)
+        acceptor_indices.extend(a[hbond_indices].indices)
+
+    return (
+        np.asarray(donor_indices, dtype=int),
+        np.asarray(acceptor_indices, dtype=int),
+        n_frames,
+    )
+
+
+def compute_all_hbonds(
+    universe: mda.Universe,
+    *,
+    mol_resname: str,
+    water_resname: str = "SOL",
+    distance_cutoff: float = 3.5,
+    angle_cutoff: float = 150,
+    hb_elements: set[str] = {"O", "N", "S"},
+    start: int = 0,
+    stop: int | None = None,
+    step: int = 1,
+) -> QoI:
+    """Compute all solute-water hydrogen-bond QoIs for one trajectory."""
+    selection_1 = f"resname {mol_resname}"
+    selection_2 = f"resname {water_resname}"
+
+    hb_elements = set(hb_elements)
+    hb_elements_str = " ".join(sorted(hb_elements))
+
+    selection_pairs = ((selection_1, selection_2), (selection_2, selection_1))
+    possible_labels: set[str] = set()
+    hbonds: dict[str, float] = {}
+    for sel_donors, sel_acceptors in selection_pairs:
+        donors, hydrogens = _paired_donors_and_hydrogens(
+            universe,
+            sel_donors,
+            hb_elements=hb_elements,
+        )
+        acceptors = universe.select_atoms(
+            f"{sel_acceptors} and element {hb_elements_str}"
+        )
+
+        possible_labels.update(_all_possible_hbond_labels(donors, acceptors))
+
+        if len(donors) == 0 or len(acceptors) == 0:
+            continue
+
+        donor_indices, acceptor_indices, n_frames = compute_hbonds(
+            universe,
+            hydrogens,
+            donors,
+            acceptors,
+            distance_cutoff=distance_cutoff,
+            angle_cutoff=angle_cutoff,
+            start=start,
+            stop=stop,
+            step=step,
+        )
+        hbonds.update(count_hbonds(universe, donor_indices, acceptor_indices, n_frames))
+
+    labels = tuple(sorted(possible_labels))
+    values = np.asarray([hbonds.get(label, 0.0) for label in labels], dtype=float)
+    metadata = {
+        "mol_resname": mol_resname,
+        "water_resname": water_resname,
+        "distance_cutoff": float(distance_cutoff),
+        "angle_cutoff": float(angle_cutoff),
+    }
+    return QoI(
+        name="hb",
+        values=values,
+        labels=labels,
+        values_per_label=1,
+        settings=metadata,
+    )

bff/qoi/rdf.py

@@ -0,0 +1,108 @@
+from typing import Tuple
+
+import MDAnalysis as mda
+import numpy as np
+from scipy.ndimage import gaussian_filter
+
+from ..tools import compute_distances, get_unitcell
+from .data import QoI
+
+
+def compute_rdf(
+    universe: mda.Universe,
+    atoms_ref: mda.AtomGroup,
+    atoms_sel: mda.AtomGroup,
+    r_range: Tuple[float, float] = (0, 10),
+    n_bins: int = 200,
+    pbc: bool = True,
+    start: int = None,
+    stop: int = None,
+    step: int = None,
+    smooth: bool = True,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Compute the radial distribution function between two atom groups."""
+    distances = compute_distances(
+        universe,
+        atoms_ref,
+        atoms_sel,
+        start=start,
+        stop=stop,
+        step=step,
+        pbc=pbc,
+    )
+    g, edges = np.histogram(distances.reshape(-1), range=r_range, bins=n_bins)
+    g = g.astype(np.float64)
+    r = 0.5 * (edges[1:] + edges[:-1])
+
+    shell_volumes = 4.0 / 3.0 * np.pi * (edges[1:] ** 3 - edges[:-1] ** 3)
+    if pbc:
+        volume = float(np.prod(np.asarray(get_unitcell(universe)[:3], dtype=float)))
+        norm = distances.size * (1.0 / volume) * shell_volumes
+        g /= norm
+    else:
+        g /= distances.size * shell_volumes
+
+    if smooth:
+        g = gaussian_filter(g, sigma=3)
+
+    return r, g
+
+
+def compute_all_rdfs(
+    universe: mda.Universe,
+    *,
+    mol_resname: str,
+    solvent_sel: str = "resname SOL HOH WAT and name O*",
+    r_range: Tuple[float, float] = (0, 10),
+    n_bins: int = 200,
+    pbc: bool = True,
+    start: int = 0,
+    stop: int | None = None,
+    step: int = 1,
+    smooth: bool = True,
+) -> QoI:
+    """Compute solvent RDF QoIs around all solute atom types."""
+    mol = universe.select_atoms(f"resname {mol_resname}")
+    mask = mol.masses > 0.5
+    mol_atomtypes = np.unique(mol[mask].types)
+
+    rdf_results: dict[str, np.ndarray] = {}
+    atoms_solvent = universe.select_atoms(solvent_sel)
+    for atomtype in mol_atomtypes:
+        atoms_reference = universe.select_atoms(f"type {atomtype}")
+        r, g = compute_rdf(
+            universe,
+            atoms_reference,
+            atoms_solvent,
+            r_range=r_range,
+            n_bins=n_bins,
+            pbc=pbc,
+            start=start,
+            stop=stop,
+            step=step,
+            smooth=smooth,
+        )
+        rdf_results[atomtype] = np.array([r, g])
+
+    atomtypes = tuple(sorted(rdf_results))
+    values = np.concatenate(
+        [
+            np.asarray(rdf_results[atomtype][1], dtype=float).reshape(-1)
+            for atomtype in atomtypes
+        ]
+    )
+    metadata = {
+        "mol_resname": mol_resname,
+        "solvent_sel": solvent_sel,
+        "r_range": tuple(r_range),
+        "n_bins": int(n_bins),
+        "pbc": bool(pbc),
+        "smooth": bool(smooth),
+    }
+    return QoI(
+        name="rdf",
+        values=values,
+        labels=atomtypes,
+        values_per_label=int(n_bins),
+        settings=metadata,
+    )