Update physics_dynamics.py

physics_dynamics.py

@@ -14,28 +14,6 @@ The model predicts how the zone-level forecast state evolves over time,
 constrained by an advection-diffusion PDE residual that prevents physically
 impossible predictions (e.g. precipitation materialising from nothing,
 uncertainty decreasing without new observations).
-
-
-Classes
--------
-ZoneStateTensor         — named container for the three observation arrays
-TemporalDynamicsModel   — core learned model (GRU encoder + MLP transition)
-PhysicsResidualLoss     — advection-diffusion PDE residual along time axis
-DynamicsTrainer         — offline pre-training on ERA5 sequences
-EnsembleDynamics        — N models for epistemic uncertainty quantification
-DynaRolloutBuffer       — generates synthetic transitions for PPO augmentation
-
-Usage
------
-    # 1. Pre-train on ERA5 sequences
-    trainer = DynamicsTrainer(n_zones=4, horizon_days=14)
-    trainer.train(era5_sequences)                  # list of ZoneStateTensor
-    trainer.save("dynamics_model.pt")
-
-    # 2. Load in training loop and generate synthetic rollouts
-    dynamics = TemporalDynamicsModel.load("dynamics_model.pt")
-    buffer = DynaRolloutBuffer(dynamics, n_synthetic_per_real=4)
-    # Pass buffer to custom PPO callback (see train_curriculum.py notes)
 """
 
 from __future__ import annotations
@@ -60,19 +38,6 @@ logger = logging.getLogger(__name__)
 
 @dataclass
 class ZoneStateTensor:
-    """
-    A single time-step of zone-level state, as tensors.
-
-    Mirrors the WeatherForecastEnv observation space:
-        precip:       [batch, n_zones, horizon_days]
-        uncertainty:  [batch, n_zones]
-        belief:       [batch, n_zones]
-
-    All values float32 in their natural ranges:
-        precip        [0, 500]  mm
-        uncertainty   [0, 1]
-        belief        [0, 1]
-    """
     precip: torch.Tensor       # [batch, n_zones, horizon_days]
     uncertainty: torch.Tensor  # [batch, n_zones]
     belief: torch.Tensor       # [batch, n_zones]
@@ -97,7 +62,6 @@ class ZoneStateTensor:
         )
 
     def flat(self) -> torch.Tensor:
-        """Flatten to [batch, n_zones * (horizon_days + 2)] for MLP input."""
         B, Z, H = self.precip.shape
         precip_flat = self.precip.reshape(B, Z * H)
         return torch.cat([precip_flat, self.uncertainty, self.belief], dim=-1)
@@ -113,7 +77,6 @@ class ZoneStateTensor:
         uncertainty: np.ndarray,
         belief: np.ndarray,
     ) -> "ZoneStateTensor":
-        """Construct from numpy arrays (adds batch dim if missing)."""
         if precip.ndim == 2:
             precip = precip[None]
         if uncertainty.ndim == 1:
@@ -132,36 +95,10 @@ class ZoneStateTensor:
 # ---------------------------------------------------------------------------
 
 class PhysicsResidualLoss(nn.Module):
-    """
-    Advection-diffusion PDE residual along the TEMPORAL dimension.
-
-    The physical intuition: as time advances by dt days, the precipitation
-    forecast at day t in the horizon should approximately equal the forecast
-    at day (t - dt) from the previous time step, shifted by advection and
-    smoothed by diffusion. This is the atmospheric forecast evolution equation.
-
-    Residual: ∂u/∂t + v ∂u/∂τ - D ∂²u/∂τ² = 0
-
-    where:
-        u  = precipitation forecast value
-        t  = real time (day-to-day model evolution)
-        τ  = forecast lead time (the horizon axis, days 0..H-1)
-        v  = advection speed (learnable)
-        D  = diffusion coefficient (learnable)
-
-    This is applied per zone independently (zones are not a spatial grid).
-
-    Args:
-        weight: Scalar multiplier for the physics loss term.
-                Start with 0.01-0.05; increase if predictions are unphysical.
-    """
 
     def __init__(self, weight: float = 0.01):
         super().__init__()
         self.weight = weight
-        # Learnable physics parameters — initialised to physically plausible values
-        # v: forecast advection ~1 day/day (forecast evolves with real time)
-        # D: diffusion smoothing ~0.1 (moderate smoothing of forecast errors)
         self.log_v = nn.Parameter(torch.tensor(0.0))   # exp(0) = 1.0
         self.log_D = nn.Parameter(torch.tensor(-2.3))  # exp(-2.3) ≈ 0.1
 
@@ -179,13 +116,6 @@ class PhysicsResidualLoss(nn.Module):
         u_next: torch.Tensor,      # [batch, n_zones, horizon_days]
         dt: float = 1.0,
     ) -> torch.Tensor:
-        """
-        Compute mean squared PDE residual.
-
-        u_current: forecast at time t
-        u_next:    predicted forecast at time t + dt
-        dt:        real-time step in days
-        """
         B, Z, H = u_current.shape
 
         # ∂u/∂t ≈ (u_next - u_current) / dt
@@ -219,27 +149,6 @@ class PhysicsResidualLoss(nn.Module):
 # ---------------------------------------------------------------------------
 
 class TemporalDynamicsModel(nn.Module):
-    """
-    Learned dynamics model: predicts next zone state from current state.
-
-    Architecture:
-        1. Per-zone GRU encoder compresses the forecast horizon sequence
-           into a latent zone embedding.
-        2. MLP transition model maps current latent → next latent.
-        3. MLP decoder reconstructs full next-state from latent.
-
-    Why GRU encoder (not FNO): Your forecast data is [n_zones, horizon_days]
-    where n_zones is small (2-4) and horizon_days is short (7-14). FNO is
-    designed for large spatial fields (64x64+). A GRU over the horizon axis
-    per zone is exact for this scale and directly compatible with your
-    existing GRUWeatherFeaturesExtractor architecture.
-
-    Args:
-        n_zones:      Number of geographic zones (matches env config)
-        horizon_days: Forecast horizon length (matches env config)
-        latent_dim:   Dimension of per-zone latent embedding
-        hidden_dim:   MLP hidden size for transition and decoder
-    """
 
     def __init__(
         self,
@@ -253,8 +162,6 @@ class TemporalDynamicsModel(nn.Module):
         self.horizon_days = horizon_days
         self.latent_dim = latent_dim
 
-        # --- Encoder: horizon sequence → zone latent ---
-        # Applied identically to each zone (weight sharing)
         self.precip_encoder = nn.GRU(
             input_size=1,
             hidden_size=latent_dim,
@@ -262,7 +169,6 @@ class TemporalDynamicsModel(nn.Module):
             batch_first=True,
         )
 
-        # Zone metadata (uncertainty + belief) → extra latent dims
         self.meta_encoder = nn.Sequential(
             nn.Linear(2, latent_dim),
             nn.Tanh(),
@@ -270,7 +176,6 @@ class TemporalDynamicsModel(nn.Module):
 
         zone_latent_dim = latent_dim * 2  # precip latent + meta latent
 
-        # --- Transition: current latent → next latent (all zones jointly) ---
         full_latent_dim = n_zones * zone_latent_dim
         self.transition = nn.Sequential(
             nn.Linear(full_latent_dim, hidden_dim),
@@ -280,7 +185,6 @@ class TemporalDynamicsModel(nn.Module):
             nn.Linear(hidden_dim, full_latent_dim),
         )
 
-        # --- Decoder: latent → next state components ---
         self.precip_decoder = nn.Sequential(
             nn.Linear(zone_latent_dim, hidden_dim),
             nn.SiLU(),
@@ -300,7 +204,6 @@ class TemporalDynamicsModel(nn.Module):
             nn.Sigmoid(),  # belief in [0, 1]
         )
 
-        # Physics loss module (parameters learned jointly with model)
         self.physics_loss = PhysicsResidualLoss(weight=0.01)
 
         logger.info(
@@ -309,16 +212,12 @@ class TemporalDynamicsModel(nn.Module):
         )
 
     def _encode(self, state: ZoneStateTensor) -> torch.Tensor:
-        """Encode zone state → latent. Returns [batch, n_zones, zone_latent_dim]."""
         B, Z, H = state.precip.shape
 
-        # Encode each zone's precipitation forecast sequence with the GRU
-        # Reshape to [batch * n_zones, horizon_days, 1] for GRU
         precip_seq = state.precip.reshape(B * Z, H, 1)
         _, h_n = self.precip_encoder(precip_seq)  # h_n: [1, B*Z, latent_dim]
         precip_latent = h_n.squeeze(0).reshape(B, Z, self.latent_dim)
 
-        # Encode per-zone metadata [uncertainty, belief]
         meta = torch.stack([state.uncertainty, state.belief], dim=-1)  # [B, Z, 2]
         meta_flat = meta.reshape(B * Z, 2)
         meta_latent = self.meta_encoder(meta_flat).reshape(B, Z, self.latent_dim)
@@ -331,30 +230,14 @@ class TemporalDynamicsModel(nn.Module):
         return_physics_loss: bool = True,
         dt: float = 1.0,
     ) -> Tuple[ZoneStateTensor, Optional[torch.Tensor]]:
-        """
-        Predict next state from current state.
-
-        Args:
-            current: Current zone state
-            return_physics_loss: Whether to compute and return the physics residual
-            dt: Real-time gap in days between current and next snapshot.
-                Must match pairing cadence (e.g. 5.0 for cache step_days=5).
-
-        Returns:
-            next_state: Predicted next zone state
-            physics_loss: PDE residual loss (None if return_physics_loss=False)
-        """
         B, Z, H = current.precip.shape
 
-        # Encode current state
         latent = self._encode(current)              # [B, Z, zone_latent_dim]
         latent_flat = latent.reshape(B, -1)         # [B, Z * zone_latent_dim]
 
-        # Transition in latent space (all zones jointly — captures inter-zone correlations)
         next_latent_flat = self.transition(latent_flat)
         next_latent = next_latent_flat.reshape(B, Z, -1)  # [B, Z, zone_latent_dim]
 
-        # Decode next state per zone
         next_latent_per_zone = next_latent.reshape(B * Z, -1)
 
         next_precip = self.precip_decoder(next_latent_per_zone).reshape(B, Z, H)
@@ -367,7 +250,6 @@ class TemporalDynamicsModel(nn.Module):
             belief=next_belief,
         )
 
-        # Physics residual loss on the precipitation forecast evolution
         phys_loss = None
         if return_physics_loss:
             phys_loss = self.physics_loss(
@@ -381,20 +263,11 @@ class TemporalDynamicsModel(nn.Module):
         initial: ZoneStateTensor,
         steps: int = 5,
     ) -> List[ZoneStateTensor]:
-        """
-        Generate a multi-step synthetic rollout.
-
-        Used by DynaRolloutBuffer to produce model-imagined transitions
-        for PPO augmentation. Gradients are not tracked here (inference only).
-
-        Returns list of states [s_0, s_1, ..., s_steps] where s_0 = initial.
-        """
         states = [initial]
         current = initial
         with torch.no_grad():
             for _ in range(steps):
                 next_state, _ = self.forward(current, return_physics_loss=False)
-                # Clamp to valid ranges
                 next_state = ZoneStateTensor(
                     precip=torch.clamp(next_state.precip, 0.0, 500.0),
                     uncertainty=torch.clamp(next_state.uncertainty, 0.0, 1.0),
@@ -438,23 +311,6 @@ class TemporalDynamicsModel(nn.Module):
 # ---------------------------------------------------------------------------
 
 class DynamicsTrainer:
-    """
-    Pre-trains TemporalDynamicsModel on sequences of zone states.
-
-    Training data format: list of consecutive (current, next) ZoneStateTensor
-    pairs extracted from ERA5 reanalysis or from environment rollouts.
-
-    Loss: data_loss + physics_loss
-        data_loss = MSE(predicted_next, actual_next) for all three components
-        physics_loss = advection-diffusion PDE residual (see PhysicsResidualLoss)
-
-    Args:
-        n_zones:         Must match your env config
-        horizon_days:    Must match your env config
-        physics_weight:  Weight for physics residual in total loss.
-                         Start at 0.01, increase to 0.1 if predictions violate physics.
-        device:          'cuda' if available, else 'cpu'
-    """
 
     def __init__(
         self,
@@ -487,24 +343,6 @@ class DynamicsTrainer:
         dts: Optional[List[float]] = None,
         default_dt: float = 1.0,
     ) -> dict:
-        """
-        Train the dynamics model on (current_state, next_state) pairs.
-
-        Args:
-            sequence_pairs: List of (current, next) ZoneStateTensor pairs.
-            epochs:         Training epochs
-            batch_size:     Batch size
-            lr:             Learning rate
-            val_split:      Fraction of data held out for validation
-            dts:            Optional per-pair real-time gaps in days (same
-                            length as sequence_pairs). When None, uses
-                            default_dt for every pair.
-            default_dt:     Fallback dt (days). Use 5.0 for historical
-                            cache pairs built at step_days=5.
-
-        Returns:
-            Training history dict with 'train_loss' and 'val_loss' lists.
-        """
         if not sequence_pairs:
             raise ValueError("sequence_pairs is empty — provide ERA5 data")
 
@@ -514,7 +352,6 @@ class DynamicsTrainer:
                 f"{len(sequence_pairs)}"
             )
 
-        # Build tensor dataset from pairs
         current_precips, current_uncerts, current_beliefs = [], [], []
         next_precips, next_uncerts, next_beliefs = [], [], []
         dt_list: List[float] = []
@@ -528,7 +365,6 @@ class DynamicsTrainer:
             next_beliefs.append(nxt.belief)
             dt_list.append(float(dts[i]) if dts is not None else float(default_dt))
 
-        # Stack along batch dimension
         cp = torch.cat(current_precips, dim=0)
         cu = torch.cat(current_uncerts, dim=0)
         cb = torch.cat(current_beliefs, dim=0)
@@ -538,8 +374,6 @@ class DynamicsTrainer:
         dt_t = torch.tensor(dt_list, dtype=torch.float32)
 
         N = cp.shape[0]
-        # Guard: we need at least 1 sample in the training split.
-        # When val_split=0 or N is too small, skip validation entirely.
         n_val = int(N * val_split) if val_split > 0 else 0
         if n_val >= N:
             n_val = max(0, N - 1)   # leave at least 1 sample for training
@@ -571,7 +405,6 @@ class DynamicsTrainer:
         history = {"train_loss": [], "val_loss": [], "physics_loss": [], "dt_mean": float(dt_t.mean())}
 
         for epoch in range(epochs):
-            # --- Train ---
             self.model.train()
             epoch_data_loss = 0.0
             epoch_phys_loss = 0.0
@@ -584,15 +417,11 @@ class DynamicsTrainer:
                 current = ZoneStateTensor(precip=cp_b, uncertainty=cu_b, belief=cb_b)
                 target  = ZoneStateTensor(precip=np_b, uncertainty=nu_b, belief=nb_b)
 
-                # Batch may mix dts; use batch mean (pairs are homogeneous
-                # when built with exact step_days only).
                 batch_dt = float(dt_b.mean().item())
                 pred, phys_loss = self.model(
                     current, return_physics_loss=True, dt=batch_dt,
                 )
 
-                # Data fidelity: MSE on all three components
-                # Normalise precip by max scale (500mm) to balance loss magnitudes
                 data_loss = (
                     F.mse_loss(pred.precip / 500.0, target.precip / 500.0)
                     + F.mse_loss(pred.uncertainty, target.uncertainty)
@@ -614,7 +443,6 @@ class DynamicsTrainer:
             avg_data = epoch_data_loss / len(train_loader)
             avg_phys = epoch_phys_loss / len(train_loader)
 
-            # --- Validate ---
             self.model.eval()
             val_loss = 0.0
             with torch.no_grad():
@@ -658,19 +486,6 @@ class DynamicsTrainer:
 # ---------------------------------------------------------------------------
 
 class EnsembleDynamics:
-    """
-    Ensemble of N dynamics models for epistemic uncertainty quantification.
-
-    Each model is trained with different random seed initialization.
-    Disagreement between models = epistemic uncertainty = regions where
-    the policy should be conservative (important for plasma/fusion applications
-    where high uncertainty = potentially dangerous operating regime).
-
-    Usage in PPO training:
-        ensemble = EnsembleDynamics(n_models=5, ...)
-        mean_next, uncertainty = ensemble.predict(current_state)
-        # Add uncertainty penalty to reward: reward -= uncertainty_weight * uncertainty
-    """
 
     def __init__(self, n_models: int = 5, **model_kwargs):
         self.models = [TemporalDynamicsModel(**model_kwargs) for _ in range(n_models)]
@@ -680,13 +495,6 @@ class EnsembleDynamics:
         self,
         current: ZoneStateTensor,
     ) -> Tuple[ZoneStateTensor, torch.Tensor]:
-        """
-        Return mean prediction and epistemic uncertainty.
-
-        uncertainty is a scalar tensor: mean std across all zones and features.
-        Use this to penalise the RL policy for actions that lead to high
-        uncertainty states (encourages conservative, well-characterised behaviour).
-        """
         all_precips, all_uncerts, all_beliefs = [], [], []
 
         for model in self.models:
@@ -707,10 +515,7 @@ class EnsembleDynamics:
             belief=belief_stack.mean(0),
         )
 
-        # Epistemic uncertainty: normalised std across ensemble members
         epistemic = (
-            # correction=0 avoids NaN when n_models=1 (Bessel correction
-            # would divide by zero with a single sample).
             (precip_stack.std(0, correction=0) / 500.0).mean()
             + uncert_stack.std(0, correction=0).mean()
             + belief_stack.std(0, correction=0).mean()
@@ -729,40 +534,6 @@ class EnsembleDynamics:
 # ---------------------------------------------------------------------------
 
 class DynaRolloutBuffer:
-    """
-    Generates synthetic (s, a, r, s') transitions for PPO augmentation.
-
-    Dyna-style model-based RL: use the learned dynamics model to generate
-    additional training transitions from states already in the replay buffer.
-    This improves sample efficiency without changing the PPO algorithm.
-
-    Integration with train_curriculum.py:
-    -------------------------------------
-    Add a DynaCallback to the PPO training loop:
-
-        class DynaCallback(BaseCallback):
-            def __init__(self, dynamics: TemporalDynamicsModel, n_synthetic: int = 4):
-                super().__init__()
-                self.dynamics = dynamics
-                self.n_synthetic = n_synthetic
-
-            def _on_rollout_end(self) -> None:
-                # After each real rollout, generate synthetic transitions
-                # and inject them into the rollout buffer before the update.
-                # (Implementation depends on SB3 internals — see notes below.)
-                pass
-
-    Note: Direct rollout buffer injection is not officially supported in SB3.
-    The practical approach is to use the dynamics model for reward shaping:
-    predict next state, measure surprise (|| actual - predicted ||), and add
-    a small exploration bonus for high-surprise transitions. This requires
-    no SB3 modifications and still improves sample efficiency.
-
-    Args:
-        dynamics:           Trained TemporalDynamicsModel
-        n_synthetic_steps:  Steps to roll out from each seed state
-        uncertainty_weight: Weight for epistemic uncertainty penalty in reward
-    """
 
     def __init__(
         self,
@@ -779,20 +550,10 @@ class DynaRolloutBuffer:
         obs_current: ZoneStateTensor,
         obs_actual_next: ZoneStateTensor,
     ) -> torch.Tensor:
-        """
-        Reward bonus for transitions that surprise the dynamics model.
-
-        High surprise = model uncertainty = exploration bonus.
-        This is the simplest Dyna integration: no SB3 modifications needed,
-        just add this to the reward in a step callback.
-
-        Returns scalar bonus in [0, ~1].
-        """
         self.dynamics.eval()
         with torch.no_grad():
             pred_next, _ = self.dynamics(obs_current, return_physics_loss=False)
 
-        # Normalised prediction error
         precip_err = F.mse_loss(
             pred_next.precip / 500.0,
             obs_actual_next.precip / 500.0,
@@ -807,9 +568,4 @@ class DynaRolloutBuffer:
         self,
         seed_state: ZoneStateTensor,
     ) -> List[ZoneStateTensor]:
-        """
-        Generate synthetic state sequence from a seed state.
-
-        Returns list of n_synthetic_steps + 1 states starting from seed_state.
-        """
-        return self.dynamics.rollout(seed_state, steps=self.n_synthetic_steps)
+        return self.dynamics.rollout(seed_state, steps=self.n_synthetic_steps)