test: updated unit tests with architectural and API changes

- Fix test_checkpoint.py by adding required schedulers argument to
  save_checkpoint and load_checkpoint calls.
- Fix test_checkpoints.py by updating pathway initialization assembly
  assertion to expect L1-normalized weights instead of raw binary values.
- Fix test_spatial_interaction.py by removing obsolete
  test_temperature_scaling (log_temperature removed from model).
- Fix test_pathways.py by providing mock args to _compute_pathway_truth
  to satisfy new visualization parameters.
- Add test_pathway_stability.py to verify numerical stability and
  gradient flow of the final stabilized pathway-informed architecture.

Author: BenjaminIsaac0111

tests/test_checkpoint.py

@@ -58,6 +58,7 @@ def test_save_load_preserves_weights(self, small_model, checkpoint_dir):
             small_model,
             optimizer,
             None,
+            None,  # schedulers
             epoch=42,
             best_val_loss=0.123,
             output_dir=checkpoint_dir,
@@ -74,8 +75,14 @@ def test_save_load_preserves_weights(self, small_model, checkpoint_dir):
         )
         fresh_optimizer = optim.Adam(fresh_model.parameters(), lr=1e-4)
 
-        start_epoch, best_val = load_checkpoint(
-            fresh_model, fresh_optimizer, None, checkpoint_dir, "interaction", "cpu"
+        start_epoch, best_val, loaded_schedulers = load_checkpoint(
+            fresh_model,
+            fresh_optimizer,
+            None,
+            None,
+            checkpoint_dir,
+            "interaction",
+            "cpu",
         )
 
         # Verify metadata
@@ -98,6 +105,7 @@ def test_save_load_preserves_scaler(self, small_model, checkpoint_dir):
             small_model,
             optimizer,
             scaler,
+            None,  # schedulers
             epoch=10,
             best_val_loss=0.5,
             output_dir=checkpoint_dir,
@@ -118,6 +126,7 @@ def test_save_load_preserves_scaler(self, small_model, checkpoint_dir):
             fresh_model,
             fresh_optimizer,
             fresh_scaler,
+            None,  # schedulers
             checkpoint_dir,
             "interaction",
             "cpu",
@@ -129,8 +138,8 @@ def test_save_load_preserves_scaler(self, small_model, checkpoint_dir):
     def test_no_checkpoint_starts_fresh(self, small_model, checkpoint_dir):
         """Missing checkpoint should return epoch 0 and inf loss."""
         optimizer = optim.Adam(small_model.parameters(), lr=1e-4)
-        start_epoch, best_val = load_checkpoint(
-            small_model, optimizer, None, checkpoint_dir, "nonexistent", "cpu"
+        start_epoch, best_val, loaded_schedulers = load_checkpoint(
+            small_model, optimizer, None, None, checkpoint_dir, "nonexistent", "cpu"
         )
         assert start_epoch == 0
         assert best_val == float("inf")

tests/test_checkpoints.py

@@ -27,7 +27,19 @@ def test_model_structure_consistency():
     assert model.gene_reconstructor.weight.shape == (num_genes, num_pathways)
 
     # Verify values match (within tolerance)
-    assert torch.allclose(model.gene_reconstructor.weight, pathway_init.T)
+    # The interaction model now L1-normalizes the pathways for stability
+    # shape of pathway_init is (num_pathways, num_genes)
+    import torch.nn.functional as F
+
+    # We must normalize the columns of pathway_init.T, which correspond to the rows of pathway_init
+    # Adding a small epsilon as done in interaction.py
+    normalized_pathway_init = pathway_init / (
+        pathway_init.sum(dim=1, keepdim=True) + 1e-6
+    )
+
+    assert torch.allclose(
+        model.gene_reconstructor.weight, normalized_pathway_init.T, atol=1e-5
+    )
 
 
 def test_checkpoint_save_load():

tests/test_pathways.py

@@ -138,30 +138,42 @@ class TestPathwayTruth:
     def test_consistent_across_calls(self, gene_list):
         """Ground truth from MSigDB membership should be identical across calls."""
         from spatial_transcript_former.visualization import _compute_pathway_truth
+        from unittest.mock import MagicMock
+
+        args = MagicMock()
+        args.sparsity_lambda = 0.0
 
         np.random.seed(42)
         gene_truth = np.random.rand(200, len(gene_list)).astype(np.float32)
 
-        result1, names1 = _compute_pathway_truth(gene_truth, gene_list)
-        result2, names2 = _compute_pathway_truth(gene_truth, gene_list)
+        result1, names1 = _compute_pathway_truth(gene_truth, gene_list, args)
+        result2, names2 = _compute_pathway_truth(gene_truth, gene_list, args)
 
         np.testing.assert_array_equal(result1, result2)
         assert names1 == names2
 
     def test_output_shape(self, gene_list):
         """Pathway truth should be (N, P) where P=50 (Hallmarks default)."""
         from spatial_transcript_former.visualization import _compute_pathway_truth
+        from unittest.mock import MagicMock
+
+        args = MagicMock()
+        args.sparsity_lambda = 0.0
 
         N = 150
         gene_truth = np.random.rand(N, len(gene_list)).astype(np.float32)
-        result, names = _compute_pathway_truth(gene_truth, gene_list)
+        result, names = _compute_pathway_truth(gene_truth, gene_list, args)
 
         assert result.shape == (N, 50)
         assert len(names) == 50
 
     def test_spatial_variation(self, gene_list):
         """Pathway truth should have spatial variation (non-zero std)."""
         from spatial_transcript_former.visualization import _compute_pathway_truth
+        from unittest.mock import MagicMock
+
+        args = MagicMock()
+        args.sparsity_lambda = 0.0
 
         # Create gene expression with spatial patterns
         N = 200
@@ -170,7 +182,7 @@ def test_spatial_variation(self, gene_list):
         gene_truth[:100, 0] += 5.0
         gene_truth[100:, 1] += 5.0
 
-        result, _ = _compute_pathway_truth(gene_truth, gene_list)
+        result, _ = _compute_pathway_truth(gene_truth, gene_list, args)
 
         # At least some pathways should have non-trivial spatial variation
         stds = np.std(result, axis=0)

tests/test_spatial_interaction.py

@@ -250,26 +250,6 @@ def test_interaction_mask_bits():
     assert mask[2, 3] == False, "h2h interaction [2, 3] should be enabled"
 
 
-def test_temperature_scaling():
-    """Verify log_temperature actually scales the pathway scores."""
-    model = SpatialTranscriptFormer(num_genes=10, token_dim=64)
-    features = torch.randn(1, 4, 2048)
-    coords = torch.randn(1, 4, 2)
-
-    # Initial scores with default temp
-    scores1 = model(features, rel_coords=coords, return_pathways=True)[1]
-
-    # Manually increase log_temperature significantly
-    with torch.no_grad():
-        model.log_temperature.fill_(10.0)  # Massive temp
-
-    scores2 = model(features, rel_coords=coords, return_pathways=True)[1]
-
-    # Scores should be different and typically more extreme
-    assert not torch.allclose(scores1, scores2)
-    assert scores2.abs().max() > scores1.abs().max()
-
-
 def test_return_attention_values():
     """Validate attention weight extraction logic."""
     model = SpatialTranscriptFormer(