diffbio 0.1.0__py3-none-any.whl

diffbio/sources/dti.py

@@ -0,0 +1,359 @@
+"""Deterministic drug-target interaction sources and paired-input helpers."""
+
+from __future__ import annotations
+
+import csv
+from collections.abc import Iterator
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal, cast
+
+import jax.numpy as jnp
+import numpy as np
+from flax import nnx
+
+from datarax.core.config import StructuralConfig
+from datarax.core.data_source import DataSourceModule
+from datarax.typing import Element
+
+DTI_DATASET_CONTRACT_KEYS = (
+    "pair_ids",
+    "protein_ids",
+    "protein_sequences",
+    "drug_ids",
+    "drug_smiles",
+    "targets",
+    "task_type",
+    "dataset_provenance",
+)
+_DTI_PROVENANCE_KEYS = (
+    "dataset_name",
+    "split",
+    "source_type",
+    "source_path",
+    "seed",
+    "task_type",
+    "n_pairs",
+    "promotion_eligible",
+    "biological_validation",
+)
+
+
+@dataclass(frozen=True, kw_only=True)
+class DTISourceConfig(StructuralConfig):
+    """Configuration for deterministic DTI sources."""
+
+    dataset_name: Literal["davis", "biosnap"]
+    split: Literal["train", "valid", "test"] = "train"
+    data_dir: Path | None = None
+    seed: int = 42
+    use_synthetic_fallback: bool = True
+
+
+def validate_dti_dataset(data: dict[str, Any]) -> None:
+    """Validate the shared DTI paired-input source contract."""
+    missing_keys = [key for key in DTI_DATASET_CONTRACT_KEYS if key not in data]
+    if missing_keys:
+        raise ValueError(f"DTI dataset is missing required keys: {missing_keys}")
+
+    lengths = [
+        len(data["pair_ids"]),
+        len(data["protein_ids"]),
+        len(data["protein_sequences"]),
+        len(data["drug_ids"]),
+        len(data["drug_smiles"]),
+        int(np.asarray(data["targets"]).shape[0]),
+    ]
+    if len(set(lengths)) != 1:
+        raise ValueError("All DTI paired-input fields must have the same length.")
+
+    task_type = str(data["task_type"])
+    if task_type not in {"affinity_regression", "binary_interaction"}:
+        raise ValueError(f"Unsupported DTI task_type: {task_type!r}")
+
+    targets = np.asarray(data["targets"])
+    if targets.ndim != 1:
+        raise ValueError("DTI targets must be a rank-1 array.")
+    if task_type == "binary_interaction":
+        unique_values = set(np.asarray(targets, dtype=np.int32).tolist())
+        if not unique_values.issubset({0, 1}):
+            raise ValueError("Binary DTI targets must contain only 0/1 labels.")
+
+    provenance = data["dataset_provenance"]
+    if not isinstance(provenance, dict):
+        raise ValueError("DTI dataset_provenance must be a dict.")
+    missing_provenance_keys = [key for key in _DTI_PROVENANCE_KEYS if key not in provenance]
+    if missing_provenance_keys:
+        raise ValueError(
+            f"DTI dataset_provenance is missing required keys: {missing_provenance_keys}"
+        )
+    if provenance["task_type"] != task_type:
+        raise ValueError("DTI dataset_provenance.task_type must match task_type.")
+    if int(provenance["n_pairs"]) != lengths[0]:
+        raise ValueError("DTI dataset_provenance.n_pairs must match paired field length.")
+
+
+def deterministic_dti_split(
+    n_items: int,
+    *,
+    seed: int,
+    train_frac: float = 2.0 / 3.0,
+    valid_frac: float = 1.0 / 6.0,
+) -> dict[str, np.ndarray]:
+    """Build a deterministic train/valid/test split over interaction indices."""
+    indices = np.arange(n_items, dtype=np.int32)
+    rng = np.random.default_rng(seed)
+    rng.shuffle(indices)
+
+    n_train = int(round(n_items * train_frac))
+    n_valid = int(round(n_items * valid_frac))
+    n_train = min(max(n_train, 1), n_items - 2)
+    n_valid = min(max(n_valid, 1), n_items - n_train - 1)
+    n_test = n_items - n_train - n_valid
+
+    return {
+        "train": np.sort(indices[:n_train]),
+        "valid": np.sort(indices[n_train : n_train + n_valid]),
+        "test": np.sort(indices[n_train + n_valid : n_train + n_valid + n_test]),
+    }
+
+
+def build_paired_dti_batch(
+    data: dict[str, Any],
+    *,
+    indices: np.ndarray,
+) -> dict[str, Any]:
+    """Slice a DTI dataset into one aligned paired batch."""
+    validate_dti_dataset(data)
+    batch_indices = np.asarray(indices, dtype=np.int32)
+    provenance = dict(data["dataset_provenance"])
+    provenance["source_n_pairs"] = int(provenance["n_pairs"])
+    provenance["n_pairs"] = int(batch_indices.shape[0])
+    batch = {
+        "pair_ids": [data["pair_ids"][index] for index in batch_indices],
+        "protein_ids": [data["protein_ids"][index] for index in batch_indices],
+        "protein_sequences": [data["protein_sequences"][index] for index in batch_indices],
+        "drug_ids": [data["drug_ids"][index] for index in batch_indices],
+        "drug_smiles": [data["drug_smiles"][index] for index in batch_indices],
+        "targets": jnp.asarray(np.asarray(data["targets"])[batch_indices]),
+        "task_type": data["task_type"],
+        "dataset_provenance": provenance,
+    }
+    validate_dti_dataset(batch)
+    return batch
+
+
+class _BaseDTISource(DataSourceModule):
+    """Shared DTI source implementation for dataset-specific wrappers."""
+
+    _data: list[Element] = nnx.data()
+
+    def __init__(
+        self,
+        config: DTISourceConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        super().__init__(config, rngs=rngs, name=name)
+        self._data = self._load_dataset()
+        self._current_idx = 0
+
+    def _load_dataset(self) -> list[Element]:
+        records = self._load_records()
+        split_indices = deterministic_dti_split(len(records), seed=self.config.seed)[
+            self.config.split
+        ]
+        return [records[int(index)] for index in split_indices]
+
+    def _load_records(self) -> list[Element]:
+        dataset_path = _resolve_dataset_path(cast(DTISourceConfig, self.config))
+        if dataset_path is not None:
+            return _load_records_from_table(dataset_path, dataset_name=self.config.dataset_name)
+        if not self.config.use_synthetic_fallback:
+            raise FileNotFoundError(
+                f"No local {self.config.dataset_name} table found under {self.config.data_dir!r}."
+            )
+        return _build_synthetic_records(dataset_name=self.config.dataset_name)
+
+    def load(self) -> dict[str, Any]:
+        """Load the paired DTI payload for the configured split."""
+        data = {
+            "pair_ids": [element.data["pair_id"] for element in self._data],
+            "protein_ids": [element.data["protein_id"] for element in self._data],
+            "protein_sequences": [element.data["protein_sequence"] for element in self._data],
+            "drug_ids": [element.data["drug_id"] for element in self._data],
+            "drug_smiles": [element.data["drug_smiles"] for element in self._data],
+            "targets": jnp.asarray([element.data["target"] for element in self._data]),
+            "task_type": self._data[0].data["task_type"],
+        }
+        data["dataset_provenance"] = _build_dataset_provenance(
+            cast(DTISourceConfig, self.config),
+            task_type=str(data["task_type"]),
+            n_pairs=len(data["pair_ids"]),
+        )
+        validate_dti_dataset(data)
+        return data
+
+    def __len__(self) -> int:
+        """Return the number of interactions in the configured split."""
+        return len(self._data)
+
+    def __getitem__(self, index: int) -> Element | None:
+        """Return one interaction element or ``None`` if out of range."""
+        if index < 0 or index >= len(self):
+            return None
+        return self._data[index]
+
+    def __iter__(self) -> Iterator[Element]:
+        """Iterate over paired DTI elements."""
+        return iter(self._data)
+
+
+class DavisDTISource(_BaseDTISource):
+    """Deterministic affinity-regression source for the Davis DTI task."""
+
+
+class BioSNAPDTISource(_BaseDTISource):
+    """Deterministic binary-interaction source for the BioSNAP DTI task."""
+
+
+def _resolve_dataset_path(config: DTISourceConfig) -> Path | None:
+    """Return the first matching local dataset table if it exists."""
+    if config.data_dir is None:
+        return None
+
+    candidates = [
+        Path(config.data_dir) / f"{config.dataset_name}.csv",
+        Path(config.data_dir) / f"{config.dataset_name}.tsv",
+    ]
+    for candidate in candidates:
+        if candidate.exists():
+            return candidate
+    return None
+
+
+def _build_dataset_provenance(
+    config: DTISourceConfig,
+    *,
+    task_type: str,
+    n_pairs: int,
+) -> dict[str, Any]:
+    """Build one provenance record for the loaded DTI split."""
+    dataset_path = _resolve_dataset_path(config)
+    source_type = "local_table" if dataset_path is not None else "synthetic_scaffold"
+    biological_validation = (
+        "local_table_unverified" if dataset_path is not None else "contract_validation_only"
+    )
+    return {
+        "dataset_name": config.dataset_name,
+        "split": config.split,
+        "source_type": source_type,
+        "source_path": None if dataset_path is None else str(dataset_path),
+        "seed": config.seed,
+        "task_type": task_type,
+        "n_pairs": int(n_pairs),
+        "promotion_eligible": False,
+        "biological_validation": biological_validation,
+    }
+
+
+def _load_records_from_table(
+    dataset_path: Path,
+    *,
+    dataset_name: Literal["davis", "biosnap"],
+) -> list[Element]:
+    """Load paired DTI records from a local CSV/TSV table."""
+    delimiter = "\t" if dataset_path.suffix == ".tsv" else ","
+    task_type = "affinity_regression" if dataset_name == "davis" else "binary_interaction"
+    target_columns = ("target", "affinity", "label")
+
+    with dataset_path.open("r", encoding="utf-8", newline="") as handle:
+        reader = csv.DictReader(handle, delimiter=delimiter)
+        records: list[Element] = []
+        for row_index, row in enumerate(reader):
+            target_value = None
+            for column_name in target_columns:
+                if column_name in row and row[column_name] not in {"", None}:
+                    target_value = float(row[column_name])
+                    break
+            if target_value is None:
+                raise ValueError(f"Missing target column in {dataset_path} row {row_index}.")
+
+            records.append(
+                _build_element(
+                    pair_id=row.get("pair_id", f"{dataset_name}_{row_index}"),
+                    protein_id=row["protein_id"],
+                    protein_sequence=row["protein_sequence"],
+                    drug_id=row["drug_id"],
+                    drug_smiles=row["drug_smiles"],
+                    target=target_value,
+                    task_type=task_type,
+                )
+            )
+    return records
+
+
+def _build_synthetic_records(
+    *,
+    dataset_name: Literal["davis", "biosnap"],
+) -> list[Element]:
+    """Build deterministic fallback paired DTI records."""
+    proteins = (
+        ("P0", "MKTAYI"),
+        ("P1", "MNNQKLI"),
+        ("P2", "MPEPTIDER"),
+    )
+    drugs = (
+        ("D0", "CCO"),
+        ("D1", "CCN"),
+        ("D2", "CCCl"),
+        ("D3", "c1ccccc1"),
+    )
+
+    records: list[Element] = []
+    for protein_index, (protein_id, protein_sequence) in enumerate(proteins):
+        for drug_index, (drug_id, drug_smiles) in enumerate(drugs):
+            if dataset_name == "davis":
+                target = 5.5 + 0.45 * protein_index + 0.3 * drug_index
+                task_type = "affinity_regression"
+            else:
+                target = float((protein_index + drug_index) % 2 == 0)
+                task_type = "binary_interaction"
+
+            records.append(
+                _build_element(
+                    pair_id=f"{dataset_name}_{protein_id}_{drug_id}",
+                    protein_id=protein_id,
+                    protein_sequence=protein_sequence,
+                    drug_id=drug_id,
+                    drug_smiles=drug_smiles,
+                    target=target,
+                    task_type=task_type,
+                )
+            )
+    return records
+
+
+def _build_element(
+    *,
+    pair_id: str,
+    protein_id: str,
+    protein_sequence: str,
+    drug_id: str,
+    drug_smiles: str,
+    target: float,
+    task_type: str,
+) -> Element:
+    """Build one DTI interaction element."""
+    return Element(
+        data={
+            "pair_id": pair_id,
+            "protein_id": protein_id,
+            "protein_sequence": protein_sequence,
+            "drug_id": drug_id,
+            "drug_smiles": drug_smiles,
+            "target": target,
+            "task_type": task_type,
+        },
+    )

diffbio/sources/embeddings.py

@@ -0,0 +1,203 @@
+"""Embedding-artifact sources built on Datarax source primitives.
+
+This module keeps file-format parsing local to DiffBio because the supported
+artifacts are biology-specific, but the actual source abstraction follows the
+same Datarax source model used elsewhere in the repository. The canonical
+runtime substrate is therefore:
+
+1. file decoding in one place
+2. source semantics via Datarax ``MemorySource``
+3. biology-specific alignment handled by specialized source subclasses
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, cast
+
+import jax.numpy as jnp
+import numpy as np
+from datarax.sources import MemorySource, MemorySourceConfig
+from flax import nnx
+
+
+def _require_torch() -> Any:
+    """Import torch, raising a clear error if it is not installed."""
+    try:
+        import torch  # noqa: PLC0415  # pyright: ignore[reportMissingImports]
+
+        return torch
+    except ImportError as err:
+        raise ImportError(
+            "PyTorch is required to load .pt embedding files. "
+            "Install with: uv pip install 'diffbio[torch-io]'"
+        ) from err
+
+
+@dataclass(frozen=True, slots=True)
+class EmbeddingArtifactPayload:
+    """Canonical embedding matrix plus optional artifact metadata arrays."""
+
+    embeddings: np.ndarray
+    metadata: dict[str, np.ndarray]
+
+
+def _coerce_pt_value(value: Any, *, field_name: str) -> np.ndarray:
+    """Convert a supported PyTorch payload field to a NumPy array."""
+    if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "numpy"):
+        return np.asarray(value.detach().cpu().numpy())
+
+    if isinstance(value, np.ndarray):
+        return value
+
+    if isinstance(value, (list, tuple)):
+        return np.asarray(value)
+
+    raise TypeError(
+        "Expected PyTorch artifact field "
+        f"'{field_name}' to be a tensor, NumPy array, list, or tuple, "
+        f"but received {type(value).__name__}."
+    )
+
+
+def _load_npz_payload(path: Path) -> EmbeddingArtifactPayload:
+    """Load the canonical array and metadata arrays from a NumPy archive."""
+    with np.load(path) as archive:
+        embedding_key = "embeddings" if "embeddings" in archive else next(iter(archive.files), None)
+        if embedding_key is None:
+            raise ValueError(f"Embedding archive is empty: {path}")
+
+        metadata = {key: np.asarray(archive[key]) for key in archive.files if key != embedding_key}
+        return EmbeddingArtifactPayload(
+            embeddings=np.asarray(archive[embedding_key], dtype=np.float32),
+            metadata=metadata,
+        )
+
+
+def _load_pt_payload(path: Path) -> EmbeddingArtifactPayload:
+    """Load a PyTorch embedding artifact plus optional metadata fields."""
+    torch = _require_torch()
+    payload = torch.load(path, map_location="cpu", weights_only=True)
+
+    if isinstance(payload, dict):
+        if "embeddings" not in payload:
+            raise ValueError(
+                "PyTorch embedding artifacts stored as mappings must include an 'embeddings' entry."
+            )
+        metadata = {
+            str(key): _coerce_pt_value(value, field_name=str(key))
+            for key, value in payload.items()
+            if key != "embeddings"
+        }
+        return EmbeddingArtifactPayload(
+            embeddings=np.asarray(
+                _coerce_pt_value(payload["embeddings"], field_name="embeddings"),
+                dtype=np.float32,
+            ),
+            metadata=metadata,
+        )
+
+    if hasattr(payload, "numpy"):
+        return EmbeddingArtifactPayload(
+            embeddings=np.asarray(payload.numpy(), dtype=np.float32),
+            metadata={},
+        )
+
+    raise TypeError(
+        "Expected a PyTorch tensor or mapping in the embedding artifact, "
+        f"but received {type(payload).__name__}."
+    )
+
+
+def load_embedding_artifact(path: Path | str) -> EmbeddingArtifactPayload:
+    """Load a canonical embedding artifact plus optional metadata arrays."""
+    resolved_path = Path(path)
+    if not resolved_path.exists():
+        raise FileNotFoundError(f"Embedding file not found: {resolved_path}")
+
+    suffix = resolved_path.suffix.lower()
+    if suffix == ".npy":
+        return EmbeddingArtifactPayload(
+            embeddings=np.asarray(np.load(resolved_path), dtype=np.float32),
+            metadata={},
+        )
+    if suffix == ".npz":
+        return _load_npz_payload(resolved_path)
+    if suffix == ".pt":
+        return _load_pt_payload(resolved_path)
+
+    raise ValueError(
+        f"Unsupported embedding file extension '{suffix}'. Use .npy, .npz, or .pt format."
+    )
+
+
+@dataclass(frozen=True)
+class EmbeddingArtifactSourceConfig(MemorySourceConfig):
+    """Configuration for eager artifact-backed embedding sources."""
+
+    file_path: str | None = None
+
+    def __post_init__(self) -> None:
+        """Validate the artifact path and delegate common source validation."""
+        super().__post_init__()
+
+        if self.file_path is None:
+            raise ValueError("file_path is required for EmbeddingArtifactSourceConfig")
+
+        resolved_path = Path(self.file_path)
+        if not resolved_path.exists():
+            raise FileNotFoundError(f"Embedding file not found: {resolved_path}")
+
+        if resolved_path.suffix.lower() not in {".npy", ".npz", ".pt"}:
+            raise ValueError(
+                "Embedding artifact sources only support .npy, .npz, or .pt files, "
+                f"got '{resolved_path.suffix}'."
+            )
+
+
+class EmbeddingArtifactSource(MemorySource):
+    """Eager Datarax-style source for external embedding artifacts."""
+
+    config: EmbeddingArtifactSourceConfig  # pyright: ignore[reportIncompatibleVariableOverride]
+    _artifact_metadata: dict[str, np.ndarray] = nnx.data()
+
+    def __init__(
+        self,
+        config: EmbeddingArtifactSourceConfig,
+        *,
+        rngs: nnx.Rngs | None = None,
+        name: str | None = None,
+    ) -> None:
+        """Load an embedding artifact into a Datarax ``MemorySource``."""
+        file_path = config.file_path
+        if file_path is None:
+            raise ValueError("file_path is required for EmbeddingArtifactSource")
+
+        payload = load_embedding_artifact(file_path)
+        source_name = name or f"EmbeddingArtifactSource({file_path})"
+        data = {"embeddings": jnp.asarray(payload.embeddings, dtype=jnp.float32)}
+
+        super().__init__(config, data=data, rngs=rngs, name=source_name)
+
+        self._artifact_metadata = payload.metadata
+        object.__setattr__(self, "_artifact_path", Path(file_path))
+
+    def load(self) -> dict[str, Any]:
+        """Return the eager in-memory payload exposed by the source."""
+        return dict(cast(dict[str, Any], self.data))
+
+    @property
+    def embeddings(self) -> jnp.ndarray:
+        """Full embedding matrix loaded from the artifact."""
+        return cast(jnp.ndarray, cast(dict[str, Any], self.data)["embeddings"])
+
+    @property
+    def artifact_metadata(self) -> dict[str, np.ndarray]:
+        """Auxiliary metadata arrays loaded from the artifact."""
+        return dict(self._artifact_metadata)
+
+    @property
+    def artifact_path(self) -> Path:
+        """Resolved path to the backing embedding artifact."""
+        return self._artifact_path