scdlkit 0.1.0__py3-none-any.whl

scdlkit/__init__.py

@@ -0,0 +1,20 @@
+"""Public package surface for scDLKit."""
+
+from scdlkit.data import PreparedData, prepare_data
+from scdlkit.evaluation.compare import BenchmarkResult, compare_models
+from scdlkit.models import BaseModel, create_model
+from scdlkit.runner import TaskRunner
+from scdlkit.training import Trainer
+
+__all__ = [
+    "BaseModel",
+    "BenchmarkResult",
+    "PreparedData",
+    "TaskRunner",
+    "Trainer",
+    "compare_models",
+    "create_model",
+    "prepare_data",
+]
+
+__version__ = "0.1.0"

scdlkit/data/__init__.py

@@ -0,0 +1,6 @@
+"""Data preparation utilities."""
+
+from scdlkit.data.prepare import prepare_data, transform_adata
+from scdlkit.data.schemas import PreparedData, SplitData
+
+__all__ = ["PreparedData", "SplitData", "prepare_data", "transform_adata"]

scdlkit/data/datasets.py

@@ -0,0 +1,30 @@
+"""PyTorch datasets backed by dense or sparse matrices."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+from scipy import sparse
+from torch.utils.data import Dataset
+
+from scdlkit.data.schemas import SplitData
+
+
+class AnnDataset(Dataset[dict[str, torch.Tensor]]):
+    """Dataset that converts rows to dense float32 on access."""
+
+    def __init__(self, split: SplitData):
+        self.split = split
+
+    def __len__(self) -> int:
+        return len(self.split)
+
+    def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
+        row = self.split.X[index]
+        x = row.toarray().ravel() if sparse.issparse(row) else np.asarray(row).ravel()
+        sample: dict[str, torch.Tensor] = {"x": torch.as_tensor(x, dtype=torch.float32)}
+        if self.split.labels is not None:
+            sample["y"] = torch.as_tensor(int(self.split.labels[index]), dtype=torch.long)
+        if self.split.batches is not None:
+            sample["batch"] = torch.as_tensor(int(self.split.batches[index]), dtype=torch.long)
+        return sample

scdlkit/data/prepare.py

@@ -0,0 +1,262 @@
+"""AnnData preparation and transformation utilities."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+from anndata import AnnData
+from scipy import sparse
+from sklearn.preprocessing import LabelEncoder, StandardScaler
+
+from scdlkit.data.schemas import PreparedData, SplitData
+from scdlkit.data.splits import build_splits
+
+
+def _require_scanpy() -> Any:
+    try:
+        import scanpy as sc
+    except ImportError as exc:
+        msg = "scanpy-backed preprocessing requires `pip install scdlkit[scanpy]`."
+        raise ImportError(msg) from exc
+    return sc
+
+
+def _encode_obs(values: np.ndarray | None) -> tuple[np.ndarray | None, dict[str, int] | None]:
+    if values is None:
+        return None, None
+    encoder = LabelEncoder()
+    encoded = encoder.fit_transform(values.astype(str))
+    mapping = {label: int(index) for index, label in enumerate(encoder.classes_)}
+    return encoded, mapping
+
+
+def _transform_obs(values: np.ndarray | None, mapping: dict[str, int] | None) -> np.ndarray | None:
+    if values is None or mapping is None:
+        return None
+    encoded = np.empty(values.shape[0], dtype=int)
+    for index, value in enumerate(values.astype(str)):
+        if value not in mapping:
+            msg = f"Encountered unseen label '{value}' during transform."
+            raise ValueError(msg)
+        encoded[index] = mapping[value]
+    return encoded
+
+
+def _extract_matrix(adata: AnnData, layer: str) -> Any:
+    if layer == "X":
+        return adata.X
+    if layer not in adata.layers:
+        msg = f"Layer '{layer}' not found in AnnData.layers."
+        raise ValueError(msg)
+    return adata.layers[layer]
+
+
+def _to_split_data(
+    x_matrix: Any,
+    indices: np.ndarray,
+    *,
+    labels: np.ndarray | None,
+    batches: np.ndarray | None,
+    obs_names: list[str],
+) -> SplitData:
+    return SplitData(
+        X=x_matrix[indices],
+        labels=labels[indices] if labels is not None else None,
+        batches=batches[indices] if batches is not None else None,
+        obs_names=[obs_names[index] for index in indices],
+    )
+
+
+def _prepare_matrix(
+    adata: AnnData,
+    *,
+    layer: str,
+    use_hvg: bool,
+    n_top_genes: int,
+    normalize: bool,
+    log1p: bool,
+    scale: bool,
+) -> tuple[AnnData, Any, list[str], StandardScaler | None]:
+    working = adata
+    if normalize or log1p or use_hvg:
+        sc = _require_scanpy()
+        if normalize:
+            sc.pp.normalize_total(working)
+        if log1p:
+            sc.pp.log1p(working)
+        if use_hvg:
+            sc.pp.highly_variable_genes(working, n_top_genes=n_top_genes, subset=True)
+    x_matrix = _extract_matrix(working, layer)
+    scaler: StandardScaler | None = None
+    if scale:
+        scaler = StandardScaler(with_mean=not sparse.issparse(x_matrix))
+        x_matrix = scaler.fit_transform(x_matrix)
+        if sparse.issparse(x_matrix):
+            x_matrix = x_matrix.tocsr()
+    feature_names = working.var_names.astype(str).tolist()
+    return working, x_matrix, feature_names, scaler
+
+
+def prepare_data(
+    adata: AnnData,
+    *,
+    layer: str = "X",
+    use_hvg: bool = False,
+    n_top_genes: int = 2000,
+    normalize: bool = False,
+    log1p: bool = False,
+    scale: bool = False,
+    label_key: str | None = None,
+    batch_key: str | None = None,
+    val_size: float = 0.15,
+    test_size: float = 0.15,
+    batch_aware_split: bool = False,
+    random_state: int = 42,
+    copy: bool = True,
+) -> PreparedData:
+    """Prepare AnnData splits and preprocessing metadata."""
+
+    working = adata.copy() if copy else adata
+    working, x_matrix, feature_names, scaler = _prepare_matrix(
+        working,
+        layer=layer,
+        use_hvg=use_hvg,
+        n_top_genes=n_top_genes,
+        normalize=normalize,
+        log1p=log1p,
+        scale=scale,
+    )
+    labels_raw = (
+        working.obs[label_key].astype(str).to_numpy()
+        if label_key is not None and label_key in working.obs
+        else None
+    )
+    if label_key is not None and labels_raw is None:
+        msg = f"label_key '{label_key}' not found in adata.obs."
+        raise ValueError(msg)
+    batches_raw = (
+        working.obs[batch_key].astype(str).to_numpy()
+        if batch_key is not None and batch_key in working.obs
+        else None
+    )
+    if batch_key is not None and batches_raw is None:
+        msg = f"batch_key '{batch_key}' not found in adata.obs."
+        raise ValueError(msg)
+    labels, label_encoder = _encode_obs(labels_raw)
+    batches, batch_encoder = _encode_obs(batches_raw)
+
+    split_indices = build_splits(
+        working.n_obs,
+        val_size=val_size,
+        test_size=test_size,
+        random_state=random_state,
+        stratify=labels if label_key is not None else None,
+        groups=batches if batch_aware_split and batch_key is not None else None,
+    )
+    obs_names = working.obs_names.astype(str).tolist()
+    train = _to_split_data(
+        x_matrix,
+        split_indices.train,
+        labels=labels,
+        batches=batches,
+        obs_names=obs_names,
+    )
+    val = (
+        _to_split_data(
+            x_matrix,
+            split_indices.val,
+            labels=labels,
+            batches=batches,
+            obs_names=obs_names,
+        )
+        if split_indices.val.size
+        else None
+    )
+    test = (
+        _to_split_data(
+            x_matrix,
+            split_indices.test,
+            labels=labels,
+            batches=batches,
+            obs_names=obs_names,
+        )
+        if split_indices.test.size
+        else None
+    )
+    preprocessing = {
+        "layer": layer,
+        "use_hvg": use_hvg,
+        "n_top_genes": n_top_genes,
+        "normalize": normalize,
+        "log1p": log1p,
+        "scale": scale,
+        "scaler": scaler,
+        "feature_names": feature_names,
+        "label_key": label_key,
+        "batch_key": batch_key,
+        "batch_aware_split": batch_aware_split,
+    }
+    return PreparedData(
+        train=train,
+        val=val,
+        test=test,
+        input_dim=int(x_matrix.shape[1]),
+        feature_names=feature_names,
+        label_encoder=label_encoder,
+        batch_encoder=batch_encoder,
+        preprocessing=preprocessing,
+    )
+
+
+def transform_adata(
+    adata: AnnData,
+    preprocessing: dict[str, Any],
+    *,
+    label_encoder: dict[str, int] | None = None,
+    batch_encoder: dict[str, int] | None = None,
+    copy: bool = True,
+) -> SplitData:
+    """Transform new AnnData using stored preprocessing metadata."""
+
+    working = adata.copy() if copy else adata
+    working, x_matrix, _, _ = _prepare_matrix(
+        working,
+        layer=preprocessing["layer"],
+        use_hvg=preprocessing["use_hvg"],
+        n_top_genes=preprocessing["n_top_genes"],
+        normalize=preprocessing["normalize"],
+        log1p=preprocessing["log1p"],
+        scale=False,
+    )
+    feature_names = preprocessing["feature_names"]
+    if list(working.var_names.astype(str)) != feature_names:
+        missing = sorted(set(feature_names) - set(working.var_names.astype(str)))
+        if missing:
+            msg = f"AnnData is missing required features: {missing[:5]}"
+            raise ValueError(msg)
+        working = working[:, feature_names].copy()
+        x_matrix = _extract_matrix(working, preprocessing["layer"])
+    scaler = preprocessing.get("scaler")
+    if scaler is not None:
+        x_matrix = scaler.transform(x_matrix)
+        if sparse.issparse(x_matrix):
+            x_matrix = x_matrix.tocsr()
+    labels = None
+    if preprocessing["label_key"] is not None and preprocessing["label_key"] in working.obs:
+        labels = _transform_obs(
+            working.obs[preprocessing["label_key"]].astype(str).to_numpy(),
+            label_encoder,
+        )
+    batches = None
+    if preprocessing["batch_key"] is not None and preprocessing["batch_key"] in working.obs:
+        batches = _transform_obs(
+            working.obs[preprocessing["batch_key"]].astype(str).to_numpy(),
+            batch_encoder,
+        )
+    return SplitData(
+        X=x_matrix,
+        labels=labels,
+        batches=batches,
+        obs_names=working.obs_names.astype(str).tolist(),
+    )

scdlkit/data/schemas.py

@@ -0,0 +1,35 @@
+"""Dataclasses for prepared datasets and preprocessing metadata."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+
+@dataclass(slots=True)
+class SplitData:
+    """One dataset split with optional encoded labels and batches."""
+
+    X: Any
+    labels: np.ndarray | None = None
+    batches: np.ndarray | None = None
+    obs_names: list[str] = field(default_factory=list)
+
+    def __len__(self) -> int:
+        return int(self.X.shape[0])
+
+
+@dataclass(slots=True)
+class PreparedData:
+    """Prepared train/validation/test splits and metadata."""
+
+    train: SplitData
+    val: SplitData | None
+    test: SplitData | None
+    input_dim: int
+    feature_names: list[str]
+    label_encoder: dict[str, int] | None
+    batch_encoder: dict[str, int] | None
+    preprocessing: dict[str, Any]

scdlkit/data/splits.py

@@ -0,0 +1,90 @@
+"""Split helpers for prepared AnnData workflows."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+from sklearn.model_selection import GroupShuffleSplit, train_test_split
+
+
+@dataclass(slots=True)
+class SplitIndices:
+    train: np.ndarray
+    val: np.ndarray
+    test: np.ndarray
+
+
+def build_splits(
+    n_samples: int,
+    *,
+    val_size: float,
+    test_size: float,
+    random_state: int,
+    stratify: np.ndarray | None = None,
+    groups: np.ndarray | None = None,
+) -> SplitIndices:
+    """Create train/validation/test indices."""
+
+    all_indices = np.arange(n_samples)
+    if val_size < 0 or test_size < 0 or val_size + test_size >= 1:
+        msg = "val_size and test_size must be >= 0 and sum to less than 1"
+        raise ValueError(msg)
+
+    holdout_fraction = val_size + test_size
+    if holdout_fraction == 0:
+        return SplitIndices(
+            train=all_indices,
+            val=np.array([], dtype=int),
+            test=np.array([], dtype=int),
+        )
+
+    if groups is not None:
+        splitter = GroupShuffleSplit(
+            n_splits=1,
+            test_size=holdout_fraction,
+            random_state=random_state,
+        )
+        train_idx, holdout_idx = next(splitter.split(all_indices, groups=groups))
+    else:
+        train_idx, holdout_idx = train_test_split(
+            all_indices,
+            test_size=holdout_fraction,
+            random_state=random_state,
+            stratify=stratify,
+        )
+
+    if test_size == 0:
+        return SplitIndices(
+            train=np.sort(train_idx),
+            val=np.sort(holdout_idx),
+            test=np.array([], dtype=int),
+        )
+
+    holdout_stratify = stratify[holdout_idx] if stratify is not None else None
+    if val_size == 0:
+        return SplitIndices(
+            train=np.sort(train_idx),
+            val=np.array([], dtype=int),
+            test=np.sort(holdout_idx),
+        )
+
+    test_fraction = test_size / holdout_fraction
+    if groups is not None:
+        holdout_groups = groups[holdout_idx]
+        splitter = GroupShuffleSplit(
+            n_splits=1,
+            test_size=test_fraction,
+            random_state=random_state,
+        )
+        val_rel, test_rel = next(splitter.split(holdout_idx, groups=holdout_groups))
+    else:
+        val_rel, test_rel = train_test_split(
+            np.arange(holdout_idx.size),
+            test_size=test_fraction,
+            random_state=random_state,
+            stratify=holdout_stratify,
+        )
+    val_idx = holdout_idx[val_rel]
+    test_idx = holdout_idx[test_rel]
+    return SplitIndices(train=np.sort(train_idx), val=np.sort(val_idx), test=np.sort(test_idx))

scdlkit/evaluation/__init__.py

@@ -0,0 +1,6 @@
+"""Evaluation utilities."""
+
+from scdlkit.evaluation.compare import BenchmarkResult, compare_models
+from scdlkit.evaluation.evaluator import evaluate_predictions
+
+__all__ = ["BenchmarkResult", "compare_models", "evaluate_predictions"]

scdlkit/evaluation/compare.py

@@ -0,0 +1,69 @@
+"""Compare multiple models on the same AnnData workflow."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import pandas as pd
+
+from scdlkit.evaluation.report import save_markdown_report
+from scdlkit.utils import ensure_directory
+
+
+@dataclass(slots=True)
+class BenchmarkResult:
+    """Collected results from comparing multiple models."""
+
+    metrics_frame: pd.DataFrame
+    runners: dict[str, Any]
+    output_paths: dict[str, str] = field(default_factory=dict)
+
+
+def compare_models(
+    adata: Any,
+    *,
+    models: list[str],
+    task: str,
+    shared_kwargs: dict[str, Any] | None = None,
+    output_dir: str | None = None,
+) -> BenchmarkResult:
+    """Train and evaluate several models with shared configuration."""
+
+    from scdlkit.runner import TaskRunner
+    from scdlkit.visualization.compare import plot_model_comparison
+
+    shared = dict(shared_kwargs or {})
+    records: list[dict[str, Any]] = []
+    runners: dict[str, TaskRunner] = {}
+    output_paths: dict[str, str] = {}
+    for model_name in models:
+        runner = TaskRunner(model=model_name, task=task, **shared)
+        runner.fit(adata)
+        metrics = runner.evaluate()
+        scalar_metrics = {k: v for k, v in metrics.items() if isinstance(v, (int, float))}
+        records.append({"model": model_name, **scalar_metrics})
+        runners[model_name] = runner
+
+    metrics_frame = pd.DataFrame.from_records(records).sort_values("model").reset_index(drop=True)
+    if output_dir is not None:
+        directory = ensure_directory(output_dir)
+        csv_path = directory / "benchmark_metrics.csv"
+        md_path = directory / "benchmark_report.md"
+        png_path = directory / "benchmark_comparison.png"
+        metrics_frame.to_csv(csv_path, index=False)
+        fig, _ = plot_model_comparison(metrics_frame)
+        fig.savefig(png_path, dpi=150, bbox_inches="tight")
+        report_lines = ["## Compared models", "", *[f"- `{name}`" for name in models]]
+        save_markdown_report(
+            {"num_models": len(models), "task": task},
+            path=md_path,
+            title="Benchmark Report",
+            extra_sections=report_lines,
+        )
+        output_paths = {
+            "metrics_csv": str(csv_path),
+            "report_md": str(md_path),
+            "comparison_png": str(png_path),
+        }
+    return BenchmarkResult(metrics_frame=metrics_frame, runners=runners, output_paths=output_paths)

scdlkit/evaluation/evaluator.py

@@ -0,0 +1,34 @@
+"""Task-aware evaluation entrypoints."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+
+from scdlkit.evaluation.metrics import (
+    classification_metrics,
+    reconstruction_metrics,
+    representation_metrics,
+)
+
+
+def evaluate_predictions(task: str, predictions: dict[str, np.ndarray]) -> dict[str, Any]:
+    """Evaluate model predictions for a task."""
+
+    if task == "classification":
+        if "y" not in predictions:
+            msg = "Classification evaluation requires encoded labels."
+            raise ValueError(msg)
+        return classification_metrics(predictions["y"], predictions["logits"])
+
+    metrics = reconstruction_metrics(predictions["x"], predictions["reconstruction"])
+    if task == "representation":
+        metrics.update(
+            representation_metrics(
+                predictions.get("latent", np.empty((predictions["x"].shape[0], 0))),
+                predictions.get("y"),
+                predictions.get("batch"),
+            )
+        )
+    return metrics

scdlkit/evaluation/metrics.py

@@ -0,0 +1,83 @@
+"""Metric helpers for reconstruction, representation, and classification."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+
+import numpy as np
+from scipy.stats import pearsonr, spearmanr
+from sklearn.cluster import KMeans
+from sklearn.metrics import (
+    accuracy_score,
+    adjusted_rand_score,
+    confusion_matrix,
+    f1_score,
+    normalized_mutual_info_score,
+    silhouette_score,
+)
+from sklearn.neighbors import NearestNeighbors
+
+
+def _safe_correlation(
+    func: Callable[[np.ndarray, np.ndarray], tuple[float, float]],
+    y_true: np.ndarray,
+    y_pred: np.ndarray,
+) -> float:
+    flat_true = np.ravel(y_true)
+    flat_pred = np.ravel(y_pred)
+    if np.std(flat_true) == 0 or np.std(flat_pred) == 0:
+        return 0.0
+    corr, _ = func(flat_true, flat_pred)
+    return 0.0 if math.isnan(corr) else float(corr)
+
+
+def reconstruction_metrics(y_true: np.ndarray, y_pred: np.ndarray) -> dict[str, float]:
+    error = y_true - y_pred
+    return {
+        "mse": float(np.mean(error**2)),
+        "mae": float(np.mean(np.abs(error))),
+        "pearson": _safe_correlation(pearsonr, y_true, y_pred),
+        "spearman": _safe_correlation(spearmanr, y_true, y_pred),
+    }
+
+
+def knn_label_consistency(latent: np.ndarray, labels: np.ndarray, n_neighbors: int = 10) -> float:
+    if len(np.unique(labels)) < 2 or latent.shape[0] <= 1:
+        return 0.0
+    neighbors = min(n_neighbors + 1, latent.shape[0])
+    knn = NearestNeighbors(n_neighbors=neighbors)
+    knn.fit(latent)
+    indices = knn.kneighbors(latent, return_distance=False)[:, 1:]
+    votes = labels[indices]
+    majority = np.array([np.bincount(row).argmax() for row in votes])
+    return float(np.mean(majority == labels))
+
+
+def representation_metrics(
+    latent: np.ndarray,
+    labels: np.ndarray | None,
+    batches: np.ndarray | None,
+) -> dict[str, float]:
+    metrics: dict[str, float] = {}
+    unique_labels = np.unique(labels) if labels is not None else np.array([])
+    unique_batches = np.unique(batches) if batches is not None else np.array([])
+    if labels is not None and latent.shape[0] > len(unique_labels) and len(unique_labels) > 1:
+        metrics["silhouette"] = float(silhouette_score(latent, labels))
+        metrics["knn_label_consistency"] = knn_label_consistency(latent, labels)
+        kmeans = KMeans(n_clusters=len(unique_labels), random_state=42, n_init="auto")
+        clusters = kmeans.fit_predict(latent)
+        metrics["ari"] = float(adjusted_rand_score(labels, clusters))
+        metrics["nmi"] = float(normalized_mutual_info_score(labels, clusters))
+    if batches is not None and latent.shape[0] > len(unique_batches) and len(unique_batches) > 1:
+        metrics["batch_silhouette"] = float(silhouette_score(latent, batches))
+    return metrics
+
+
+def classification_metrics(y_true: np.ndarray, logits: np.ndarray) -> dict[str, object]:
+    predicted = logits.argmax(axis=1)
+    return {
+        "accuracy": float(accuracy_score(y_true, predicted)),
+        "macro_f1": float(f1_score(y_true, predicted, average="macro")),
+        "confusion_matrix": confusion_matrix(y_true, predicted).tolist(),
+    }

scdlkit/evaluation/report.py

@@ -0,0 +1,40 @@
+"""Report export helpers."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import pandas as pd
+
+
+def _scalar_metrics(metrics: dict[str, Any]) -> dict[str, float]:
+    return {key: float(value) for key, value in metrics.items() if isinstance(value, (int, float))}
+
+
+def save_metrics_table(metrics: dict[str, Any], path: str | Path) -> Path:
+    """Write scalar metrics to CSV."""
+
+    output = Path(path)
+    frame = pd.DataFrame([_scalar_metrics(metrics)])
+    frame.to_csv(output, index=False)
+    return output
+
+
+def save_markdown_report(
+    metrics: dict[str, Any],
+    *,
+    path: str | Path,
+    title: str,
+    extra_sections: list[str] | None = None,
+) -> Path:
+    """Write a markdown report with scalar and structured metrics."""
+
+    output = Path(path)
+    lines = [f"# {title}", "", "## Metrics", ""]
+    for key, value in metrics.items():
+        lines.append(f"- **{key}**: {value}")
+    if extra_sections:
+        lines.extend(["", *extra_sections])
+    output.write_text("\n".join(lines), encoding="utf-8")
+    return output