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()