Blame for tests/test_core.py

from __future__ import annotations

import pickle
from dataclasses import replace

import numpy as np
import pandas as pd
import pytest
import torch

from biolmnet.artifacts import load_bundle, save_bundle
from biolmnet.data import (
    BranchPriors,
    PreparedWorkspace,
    attach_embeddings_and_pathways,
    build_biological_mask,
    deterministic_gene_embeddings,
)
from biolmnet.model import BioMaskedLinear, BioLMNet
from biolmnet.training import (
    DEFAULT_VALIDATION_FRACTION,
    Hyperparameters,
    predict,
    split_train_validation,
    train,
)


def test_biological_mask_uses_pdi_and_undirected_ppi() -> None:
    pdi = pd.DataFrame(
        {"TF": ["A", "C", "outside"], "Target": ["B", "D", "A"]}
    )
    ppi = pd.DataFrame(
        {
            "protein1": ["A", "X", "B"],
            "protein2": ["X", "C", "Y"],
            "combined_score": [950, 950, 710],
        }
    )
    branch = build_biological_mask(["A", "B", "C", "D"], pdi, ppi)

    assert {"B", "D", "X"}.issubset(branch.hidden_genes)
    x_index = branch.hidden_genes.index("X")
    assert branch.biological_mask[0, x_index] > 0
    assert branch.biological_mask[2, x_index] > 0
    assert branch.pdi_edges == 2


def test_masked_linear_disconnects_unlisted_weights() -> None:
    mask = torch.tensor([[1.0, 0.0], [0.0, 1.0]])
    layer = BioMaskedLinear(mask, bias=False)
    values = torch.tensor([[2.0, 3.0]])
    baseline = layer(values).detach().clone()
    with torch.no_grad():
        layer.weight[0, 1] = 10_000
        layer.weight[1, 0] = -10_000
    changed = layer(values).detach()
    assert torch.allclose(baseline, changed)


def _tiny_workspace() -> PreparedWorkspace:
    rng = np.random.default_rng(7)
    sample_count = 48
    labels = np.repeat(np.array([0, 1]), sample_count // 2)
    gene_values = rng.normal(size=(sample_count, 4)).astype(np.float32)
    dna_values = rng.normal(size=(sample_count, 4)).astype(np.float32)
    gene_values[:, 0] += labels * 1.5
    dna_values[:, 1] -= labels * 1.2
    input_genes = ["A", "B", "C", "D"]
    hidden_genes = ["A", "B", "C"]
    biological = np.array(
        [
            [1.0, 1.0, 0.0],
            [0.0, 1.0, 1.0],
            [1.0, 0.0, 1.0],
            [0.0, 1.0, 0.0],
        ],
        dtype=np.float32,
    )
    pathway_mapping = pd.DataFrame(
        {
            "SYMBOL": ["A", "B", "B", "C"],
            "PathwayID": ["hsa1", "hsa1", "hsa2", "hsa2"],
        }
    )
    embeddings = deterministic_gene_embeddings(hidden_genes, dimensions=8)

    def branch() -> BranchPriors:
        value = BranchPriors(
            input_genes=input_genes.copy(),
            hidden_genes=hidden_genes.copy(),
            biological_mask=biological.copy(),
            pdi_edges=3,
            ppi_edges=4,
        )
        attach_embeddings_and_pathways(
            value,
            embeddings,
            pathway_mapping,
            precomputed_significant=True,
        )
        return value

    return PreparedWorkspace(
        gene_expression=gene_values,
        dna_methylation=dna_values,
        labels=labels,
        label_names=["control", "case"],
        gene_branch=branch(),
        dna_branch=branch(),
        source_name="unit test",
    )


def test_model_forward_probabilistic_shape() -> None:
    workspace = _tiny_workspace()
    gene = workspace.gene_branch
    dna = workspace.dna_branch
    model = BioLMNet(
        torch.from_numpy(gene.biological_mask),
        torch.from_numpy(dna.biological_mask),
        torch.from_numpy(gene.embeddings),
        torch.from_numpy(dna.embeddings),
        torch.from_numpy(gene.pathway_mask),
        torch.from_numpy(dna.pathway_mask),
        n_classes=2,
        projection_dim=4,
        fusion_dim=3,
        dropout=0.0,
    )
    logits = model(torch.randn(5, 4), torch.randn(5, 4))
    assert logits.shape == (5, 2)
    assert torch.allclose(
        model.gene_branch.pathway_attention.attention_weights().sum(dim=0),
        torch.ones(2),
    )


def test_duplicate_embedding_rows_do_not_break_pathway_attachment() -> None:
    branch = BranchPriors(
        input_genes=["A", "B"],
        hidden_genes=["A", "B", "C"],
        biological_mask=np.ones((2, 3), dtype=np.float32),
    )
    embeddings = pd.DataFrame(
        np.arange(16, dtype=np.float32).reshape(4, 4),
        index=["A", "B", "B", "C"],
    )
    pathway_mapping = pd.DataFrame(
        {
            "SYMBOL": ["A", "B", "C"],
            "PathwayID": ["hsa1", "hsa1", "hsa2"],
        }
    )

    attach_embeddings_and_pathways(
        branch,
        embeddings,
        pathway_mapping,
        precomputed_significant=True,
    )

    assert branch.embeddings.shape[0] == len(branch.hidden_genes)
    assert branch.biological_mask.shape[1] == len(branch.hidden_genes)
    assert branch.pathway_mask.shape[0] == len(branch.hidden_genes)


def test_training_artifact_roundtrip(tmp_path) -> None:
    workspace = _tiny_workspace()
    result = train(
        workspace,
        Hyperparameters(
            epochs=3,
            batch_size=8,
            projection_dim=4,
            fusion_dim=3,
            dropout=0.0,
            early_stopping_patience=3,
        ),
    )
    restored_from_process_boundary = pickle.loads(pickle.dumps(result.bundle))
    assert restored_from_process_boundary.label_names == ["control", "case"]
    artifact = save_bundle(result.bundle, tmp_path / "model.zip")
    restored = load_bundle(artifact)
    gene_frame = pd.DataFrame(
        workspace.gene_expression[:5], columns=restored.gene_features
    )
    dna_frame = pd.DataFrame(
        workspace.dna_methylation[:5], columns=restored.dna_features
    )
    before = predict(gene_frame, dna_frame, result.bundle)
    after = predict(gene_frame, dna_frame, restored)
    probability_columns = [column for column in before if column.startswith("P(")]
    np.testing.assert_allclose(
        before[probability_columns].to_numpy(),
        after[probability_columns].to_numpy(),
        atol=1e-6,
    )


def test_prediction_can_score_prepared_workspace_without_uploads() -> None:
    from app import run_prediction

    workspace = _tiny_workspace()
    result = train(
        workspace,
        Hyperparameters(
            epochs=3,
            batch_size=8,
            projection_dim=4,
            fusion_dim=3,
            dropout=0.0,
            early_stopping_patience=3,
        ),
    )

    _, status, output, _, download_path, predict_meta = run_prediction(
        result.bundle,
        workspace,
        None,
        False,
        True,
        None,
        None,
        {},
    )

    assert "Predicted 48 samples" in status
    assert len(output) == len(workspace.labels)
    assert download_path is not None
    assert predict_meta["n_samples"] == len(workspace.labels)


def test_zerogpu_duration_estimator_is_bounded_and_scales() -> None:
    from app import estimate_training_duration

    workspace = _tiny_workspace()
    common = (workspace, 16, 0.001, 0.01, 0.3, 64, 12, 0.2, "Adam", True)
    short = estimate_training_duration(common[0], 10, *common[1:])
    long = estimate_training_duration(common[0], 200, *common[1:])
    assert 30 <= short <= 300
    assert short <= long <= 300


_FAST = Hyperparameters(
    epochs=3,
    batch_size=8,
    projection_dim=4,
    fusion_dim=3,
    dropout=0.0,
    early_stopping_patience=3,
)


def test_default_split_is_75_training_25_validation() -> None:
    assert DEFAULT_VALIDATION_FRACTION == 0.25
    assert Hyperparameters().validation_fraction == 0.25
    labels = np.repeat(np.array([0, 1]), 50)
    train_index, validation_index = split_train_validation(
        labels, Hyperparameters().validation_fraction, Hyperparameters().seed
    )
    assert len(train_index) == 75
    assert len(validation_index) == 25


def test_split_partitions_every_sample_exactly_once() -> None:
    labels = np.repeat(np.array([0, 1, 2]), [40, 25, 18])
    train_index, validation_index = split_train_validation(labels, 0.25, 42)
    assert len(train_index) + len(validation_index) == len(labels)
    assert set(train_index).isdisjoint(validation_index)
    assert sorted(np.concatenate([train_index, validation_index])) == list(range(len(labels)))


def test_split_preserves_class_proportions() -> None:
    labels = np.repeat(np.array([0, 1, 2]), [120, 60, 20])  # 60 / 30 / 10
    train_index, validation_index = split_train_validation(labels, 0.25, 42)
    expected = np.array([0.6, 0.3, 0.1])
    for partition in (train_index, validation_index):
        observed = np.bincount(labels[partition], minlength=3) / len(partition)
        np.testing.assert_allclose(observed, expected, atol=0.02)


def test_split_is_reproducible_for_a_fixed_seed() -> None:
    labels = np.repeat(np.array([0, 1]), [30, 34])
    first = split_train_validation(labels, 0.25, 42)
    second = split_train_validation(labels, 0.25, 42)
    other_seed = split_train_validation(labels, 0.25, 7)
    np.testing.assert_array_equal(first[0], second[0])
    np.testing.assert_array_equal(first[1], second[1])
    assert set(first[1]) != set(other_seed[1])


def test_split_rejects_partitions_too_small_to_stratify() -> None:
    with pytest.raises(ValueError, match="at least two samples"):
        split_train_validation(np.array([0, 0, 0, 1]), 0.25, 42)
    # 4 classes x 2 samples: 25% leaves 2 validation samples for 4 classes.
    with pytest.raises(ValueError, match="one sample per class"):
        split_train_validation(np.repeat(np.arange(4), 2), 0.25, 42)


def test_training_uses_paired_rows_and_train_only_scaler() -> None:
    workspace = _tiny_workspace()
    hyperparameters = replace(_FAST, validation_fraction=0.25)
    train_index, validation_index = split_train_validation(
        workspace.labels, hyperparameters.validation_fraction, hyperparameters.seed
    )
    # Shift only the validation rows: a scaler fitted on all samples would move.
    workspace.gene_expression[validation_index] += 100.0
    workspace.dna_methylation[validation_index] -= 100.0

    result = train(workspace, hyperparameters)
    bundle = result.bundle

    assert bundle.metrics["training_samples"] == 36
    assert bundle.metrics["validation_samples"] == 12
    np.testing.assert_array_equal(
        result.validation_predictions["sample_row"].to_numpy(), validation_index
    )
    expected_labels = [workspace.label_names[value] for value in workspace.labels[validation_index]]
    assert result.validation_predictions["observed"].tolist() == expected_labels

    # Both omics scalers come from the same training rows only.
    np.testing.assert_allclose(
        bundle.gene_mean, workspace.gene_expression[train_index].mean(axis=0), rtol=1e-5
    )
    np.testing.assert_allclose(
        bundle.gene_scale, workspace.gene_expression[train_index].std(axis=0), rtol=1e-5
    )
    np.testing.assert_allclose(
        bundle.dna_mean, workspace.dna_methylation[train_index].mean(axis=0), rtol=1e-5
    )
    np.testing.assert_allclose(
        bundle.dna_scale, workspace.dna_methylation[train_index].std(axis=0), rtol=1e-5
    )

    # Validation rows are transformed with that training-fitted scaler: the
    # stored validation probabilities match predict(), which uses the bundle's scaler.
    rescored = predict(
        pd.DataFrame(workspace.gene_expression[validation_index], columns=bundle.gene_features),
        pd.DataFrame(workspace.dna_methylation[validation_index], columns=bundle.dna_features),
        bundle,
    )
    probability_columns = [f"P({label})" for label in workspace.label_names]
    np.testing.assert_allclose(
        result.validation_predictions[probability_columns].to_numpy(),
        rescored[probability_columns].to_numpy(),
        atol=1e-5,
    )


def test_class_weights_use_training_labels_only(monkeypatch) -> None:
    import biolmnet.training as training

    captured = {}
    original = training.nn.CrossEntropyLoss

    def spy(*args, weight=None, **kwargs):
        captured["weight"] = None if weight is None else weight.cpu().numpy()
        return original(*args, weight=weight, **kwargs)

    monkeypatch.setattr(training.nn, "CrossEntropyLoss", spy)
    workspace = _tiny_workspace()  # 24 control / 24 case: all-sample weights are [1, 1].
    # Force a lopsided split so training-only weights clearly differ from all-sample ones.
    train_index = np.concatenate([np.arange(0, 22), np.arange(24, 30)])
    validation_index = np.setdiff1d(np.arange(48), train_index)
    monkeypatch.setattr(
        training, "split_train_validation", lambda *args: (train_index, validation_index)
    )
    train(workspace, _FAST)

    counts = np.bincount(workspace.labels[train_index], minlength=2).astype(np.float32)
    assert counts.tolist() == [22.0, 6.0]
    np.testing.assert_allclose(captured["weight"], len(train_index) / (2 * counts), rtol=1e-6)


def test_results_page_receives_training_history_plot() -> None:
    from app import refresh_results

    assert refresh_results(None, {})[-1] is None

    result = train(_tiny_workspace(), _FAST)
    figure = refresh_results(result.bundle, {})[-1]
    traces = {trace.name: trace for trace in figure.data}
    assert set(traces) == {"Training", "Validation"}
    history = result.bundle.history
    assert list(traces["Training"].y) == [row["training_loss"] for row in history]
    assert list(traces["Validation"].y) == [row["validation_loss"] for row in history]

    # Artifacts saved without history still render the rest of the page.
    result.bundle.history = []
    assert refresh_results(result.bundle, {})[-1] is None


def test_activation_choices_reach_model_config_and_artifact(tmp_path) -> None:
    hyperparameters = replace(
        _FAST,
        biological_activation="tanh",
        projection_activation="relu",
        fusion_activation="sigmoid",
    )
    result = train(_tiny_workspace(), hyperparameters)
    branch = result.bundle.model.gene_branch
    assert isinstance(branch.biological_activation, torch.nn.Tanh)
    assert isinstance(branch.projection_activation, torch.nn.ReLU)
    assert isinstance(result.bundle.model.fusion_activation, torch.nn.Sigmoid)
    assert result.bundle.config["architecture"]["biological_activation"] == "tanh"

    restored = load_bundle(save_bundle(result.bundle, tmp_path / "model.zip"))
    assert isinstance(restored.model.gene_branch.biological_activation, torch.nn.Tanh)
    assert isinstance(restored.model.fusion_activation, torch.nn.Sigmoid)


def test_defaults_follow_lab_grid_configuration() -> None:
    defaults = Hyperparameters()
    assert (defaults.optimizer, defaults.learning_rate, defaults.weight_decay) == ("Adam", 0.01, 0.01)
    assert (defaults.batch_size, defaults.projection_dim, defaults.dropout, defaults.epochs) == (32, 64, 0.3, 200)
    assert (defaults.biological_activation, defaults.projection_activation, defaults.fusion_activation) == (
        "tanh",
        "relu",
        "tanh",
    )
    assert defaults.sgd_momentum == 0.03


def test_auto_genept_uses_lab_recommended_single_cell_file() -> None:
    from app import GENEPT_OPTIONS, _resolve_embedding_file

    auto = next(label for label, value in GENEPT_OPTIONS.items() if value == "auto")
    for dataset in ("BRCA", "COAD", "GBM", "scTrioseq2"):
        assert _resolve_embedding_file(auto, "BioLM-NET examples", dataset) == (
            "embedding_associations_cell_type_tissue_drug_pathway_openai_large.parquet"
        )


def test_sgd_momentum_reaches_optimizer(monkeypatch) -> None:
    import biolmnet.training as training

    captured = {}
    original = training.torch.optim.SGD

    def spy(params, **kwargs):
        captured.update(kwargs)
        return original(params, **kwargs)

    monkeypatch.setattr(training.torch.optim, "SGD", spy)
    train(_tiny_workspace(), replace(_FAST, optimizer="SGD", learning_rate=0.01, sgd_momentum=0.03))
    assert captured["momentum"] == 0.03


@pytest.mark.parametrize(
    "field, value",
    [("biological_activation", "softmax"), ("fusion_activation", "gelu"), ("sgd_momentum", 1.0)],
)
def test_invalid_activation_or_momentum_is_rejected(field, value) -> None:
    with pytest.raises(ValueError):
        replace(Hyperparameters(), **{field: value}).validate()