Update test_physics_dynamics.py

test_physics_dynamics.py

@@ -124,7 +124,6 @@ def batch_size() -> int:
 
 @pytest.fixture
 def state(n_zones, horizon_days, batch_size) -> ZoneStateTensor:
-    """Standard random state tensor for use across tests."""
     torch.manual_seed(0)
     return ZoneStateTensor(
         precip=torch.rand(batch_size, n_zones, horizon_days) * 20.0,
@@ -143,7 +142,6 @@ def model(n_zones, horizon_days) -> TemporalDynamicsModel:
 
 @pytest.fixture
 def tiny_pairs(n_zones, horizon_days) -> list:
-    """Small set of (current, next) pairs for trainer tests."""
     torch.manual_seed(7)
     pairs = []
     for _ in range(20):
@@ -302,36 +300,23 @@ class TestPhysicsResidualLoss:
         assert loss_fn(u, v).item() == pytest.approx(0.0)
 
     def test_exact_linear_advection_has_low_residual(self):
-        """
-        If u_next is u shifted by exactly v steps along the horizon axis,
-        the advection term v*∂u/∂τ should largely cancel ∂u/∂t,
-        producing a small residual (not exactly zero due to diffusion term
-        and boundary approximations, but significantly lower than random).
-        """
         loss_fn = PhysicsResidualLoss(weight=1.0)
-        # Fix v to a known value for this test
         with torch.no_grad():
             loss_fn.log_v.fill_(0.0)   # v = 1.0
             loss_fn.log_D.fill_(-10.0) # D ≈ 0 (nearly pure advection)
 
         B, Z, H = 1, 1, 20
         tau = torch.arange(H, dtype=torch.float32)
-        # Linear ramp: u = a * tau + b (∂u/∂τ = a, ∂²u/∂τ² = 0)
-        # Exact solution after dt=1: u_next = a*(tau+1) + b = u + a
         a = 2.0
         u = a * tau.unsqueeze(0).unsqueeze(0).expand(B, Z, H)
         u_next = u + a  # shift by a (= v * ∂u/∂τ = 1.0 * a)
 
         residual_advection = loss_fn(u, u_next)
 
-        # For comparison: random u_next should have a larger residual than
-        # the analytically correct solution.
         torch.manual_seed(0)
         u_random = torch.rand_like(u) * 40.0  # unrelated to u
         residual_random = loss_fn(u, u_random)
 
-        # The analytically correct solution should produce a strictly lower
-        # residual than a completely unrelated random prediction.
         assert residual_advection.item() < residual_random.item(), (
             f"Expected advection residual ({residual_advection.item():.4f}) < "
             f"random residual ({residual_random.item():.4f})"
@@ -611,14 +596,8 @@ class TestDynamicsTrainer:
             assert loss >= 0.0, f"Negative loss: {loss}"
 
     def test_loss_decreases_over_training(self, n_zones, horizon_days):
-        """
-        Loss should trend downward over sufficient epochs on a small fixed dataset.
-        We test that final loss < initial loss (not strictly monotonic — that is
-        not guaranteed with SGD). Seed is fixed for reproducibility.
-        """
         torch.manual_seed(0)
         np.random.seed(0)
-        # Build a more learnable target: next = current + small noise
         pairs = []
         for _ in range(30):
             curr = ZoneStateTensor(
@@ -664,21 +643,10 @@ class TestDynamicsTrainer:
             ),
         )
         trainer = DynamicsTrainer(n_zones=n_zones, horizon_days=horizon_days)
-        # val_split=0.0 → n_val=0, n_train=1 (after the guard fix).
-        # With val_split=0.1 and N=1 → n_val would be 1, n_train=0 → should raise.
-        # val_split=0.0 is the correct way to train on a tiny dataset.
         history = trainer.train([pair], epochs=1, batch_size=1, val_split=0.0)
         assert "train_loss" in history
         assert len(history["train_loss"]) == 1
 
-        # Verify that passing exactly 2 pairs with val_split=1.0 raises clearly
-        # (n_val = max(0, 2-1) = 1, n_train = 1 is fine — but val_split=1.0
-        # with the guard: n_val = int(2*1.0) = 2 >= N=2 → n_val = max(0,2-1)=1
-        # So we need val_split such that int(N*val_split) >= N with our guard.
-        # The guard sets n_val = max(0, N-1), so n_train = 1 always survives.
-        # The real failure case is N=1 with val_split=0.5+: int(1*0.5)=0 → fine.
-        # Actually to trigger the error we need a 0-sample training set which
-        # cannot happen with the guard. Verify graceful handling instead.
         history2 = trainer.train([pair, pair], epochs=1, batch_size=1, val_split=0.5)
         assert "train_loss" in history2
 
@@ -698,10 +666,6 @@ class TestDynamicsTrainer:
             n_zones=n_zones, horizon_days=horizon_days, physics_weight=0.0
         )
         history = trainer.train(tiny_pairs, epochs=2, batch_size=4)
-        # history["physics_loss"] records the RAW unweighted PDE residual,
-        # not the weighted contribution to the total loss. When physics_weight=0
-        # the residual is still computed and logged — it just doesn't affect
-        # the gradient. We verify it is finite and non-negative.
         for pl in history["physics_loss"]:
             assert np.isfinite(pl), f"Physics loss is not finite: {pl}"
             assert pl >= 0.0, f"Physics loss is negative: {pl}"
@@ -735,19 +699,12 @@ class TestEnsembleDynamics:
         assert epistemic.item() >= 0.0
 
     def test_single_model_ensemble_has_zero_uncertainty(self, state, n_zones, horizon_days):
-        """
-        With n_models=1 there is no variance — std of a single value is 0.
-        """
         torch.manual_seed(0)
         ens = EnsembleDynamics(n_models=1, n_zones=n_zones, horizon_days=horizon_days)
         _, epistemic = ens.predict(state)
         assert epistemic.item() == pytest.approx(0.0, abs=1e-6)
 
     def test_multi_model_ensemble_has_positive_uncertainty(self, state, n_zones, horizon_days):
-        """
-        With n_models=5 and randomly initialised weights, the ensemble members
-        will disagree, producing non-zero epistemic uncertainty.
-        """
         torch.manual_seed(99)
         ens = EnsembleDynamics(n_models=5, n_zones=n_zones, horizon_days=horizon_days)
         _, epistemic = ens.predict(state)
@@ -797,9 +754,7 @@ class TestDynaRolloutBuffer:
         assert bonus.item() >= 0.0
 
     def test_compute_surprise_bonus_bounded_by_uncertainty_weight(self, buffer, state):
-        """Bonus is clipped to [0, uncertainty_weight]."""
         n_zones, horizon_days = state.n_zones, state.horizon_days
-        # Make next_state very different to maximise surprise
         very_different = ZoneStateTensor(
             precip=torch.full_like(state.precip, 499.0),
             uncertainty=torch.ones_like(state.uncertainty),
@@ -809,12 +764,6 @@ class TestDynaRolloutBuffer:
         assert bonus.item() <= buffer.uncertainty_weight + 1e-6
 
     def test_identical_states_give_near_zero_bonus(self, buffer, state):
-        """
-        When current and next are identical, the model's prediction error
-        should be low, producing a near-zero surprise bonus.
-        Note: not exactly zero because the model doesn't predict the identity.
-        We just check it is lower than the maximum possible bonus.
-        """
         bonus_same = buffer.compute_surprise_bonus(state, state)
         very_different = ZoneStateTensor(
             precip=torch.full_like(state.precip, 499.0),
@@ -822,7 +771,6 @@ class TestDynaRolloutBuffer:
             belief=torch.zeros_like(state.belief),
         )
         bonus_diff = buffer.compute_surprise_bonus(state, very_different)
-        # Identical input should produce smaller or equal bonus than maximally different
         assert bonus_same.item() <= bonus_diff.item() + 1e-6
 
     def test_no_gradient_in_compute_surprise_bonus(self, buffer, state):
@@ -857,4 +805,4 @@ class TestDynaRolloutBuffer:
         )
         bonus = buf.compute_surprise_bonus(state, very_different)
         assert bonus.item() <= weight + 1e-6
-        assert bonus.item() >= 0.0
+        assert bonus.item() >= 0.0