Spaces:
Sleeping
Sleeping
| 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 | |
| def test_invalid_activation_or_momentum_is_rejected(field, value) -> None: | |
| with pytest.raises(ValueError): | |
| replace(Hyperparameters(), **{field: value}).validate() | |