Upload liquid_flow/physics_loss.py

liquid_flow/physics_loss.py

@@ -1,37 +1,11 @@
 """
 Physics-Informed Regularization for LiquidFlow.
+CORRECTED VERSION: fixed intensity tracking, proper buffer handling.
 
-From: "Physics-Informed Diffusion Models" (Bastek & Sun, ICLR 2025)
-and "PID: Physics-Informed Diffusion for IR Image Generation" (Mao et al., 2024)
-
-Physics losses act as TRAINING-ONLY regularizers — they don't affect
-inference speed. The pattern:
-
-1. During training: denoise to get x̂₀, compute physics residual, add to loss
-2. During inference: no change at all
-
-Implemented physics constraints for image generation:
-
-A. Total Variation (TV) — penalizes non-smooth outputs
-   L_TV = ||∇_x x̂₀||₁ + ||∇_y x̂₀||₁
-   → Enforces spatial smoothness, reduces artifacts
-
-B. Conservation of Intensity — mass conservation across image
-   L_cons = ||mean(x̂₀) - E[mean(x_ref)]||²
-   → Prevents intensity drift
-
-C. Spectral Regularizer — penalizes high-frequency noise
-   L_spec = ||FFT_high(x̂₀)||²
-   → Reduces checkerboard artifacts
-
-D. Gradient Magnitude Balance — prevents exploding gradients in dark regions
-   L_grad = ||∇x̂₀||² (Sobolev regularization)
-   → Stabilizes training in low-signal regions
-
-Pattern: L_total = L_diffusion + λ_TV * L_TV + λ_cons * L_cons + λ_spec * L_spec
-
-The virtual-observable paradigm (from PAD-Hand, 2026):
-Physics constraints are SOFT — they guide without requiring perfect satisfaction.
+Pattern from: Bastek & Sun (ICLR 2025)
+- Physics losses computed on estimated x̂₀ during training
+- Zero cost at inference
+- Acts as implicit regularizer against artifacts
 """
 
 import torch
@@ -41,209 +15,124 @@ import torch.nn.functional as F
 
 class PhysicsRegularizer(nn.Module):
     """
-    Physics-informed regularizer for image generation training.
-    
-    All losses are computed on the estimated clean sample x̂₀ during training.
-    They are ADDITIVE regularizers — just add to the diffusion loss.
+    Physics-informed regularizer for diffusion training.
     
-    Args:
-        tv_weight: Total Variation weight (default 0.01)
-        cons_weight: Conservation of intensity weight (default 0.001)
-        spec_weight: Spectral regularizer weight (default 0.01)
-        grad_weight: Gradient magnitude penalty weight (default 0.001)
+    Computed on estimated clean sample x̂₀ (DDIM one-step estimate).
+    All losses are differentiable through the noise predictor.
     """
     
-    def __init__(
-        self,
-        tv_weight=0.01,
-        cons_weight=0.001,
-        spec_weight=0.01,
-        grad_weight=0.001,
-    ):
+    def __init__(self, tv_weight=0.01, cons_weight=0.001, spec_weight=0.01, grad_weight=0.001):
         super().__init__()
         self.tv_weight = tv_weight
         self.cons_weight = cons_weight
         self.spec_weight = spec_weight
         self.grad_weight = grad_weight
         
-        # Running mean for intensity conservation
-        self.register_buffer('intensity_mean', torch.tensor(0.0))
-        self.register_buffer('intensity_count', torch.tensor(0))
-        self.intensity_alpha = 0.99  # EMA decay
+        # EMA intensity tracking
+        self.register_buffer('intensity_ema', torch.tensor(0.0))
+        self.register_buffer('step_count', torch.tensor(0, dtype=torch.long))
     
     def total_variation(self, x):
-        """
-        Total Variation loss on image batch x.
-        
-        L_TV = mean(|x_{i+1,j} - x_{i,j}| + |x_{i,j+1} - x_{i,j}|)
-        
-        Args:
-            x: [B, C, H, W] images
-        Returns:
-            tv_loss: scalar
-        """
+        """L1 total variation: encourages spatial smoothness."""
         diff_h = torch.abs(x[:, :, 1:, :] - x[:, :, :-1, :])
         diff_w = torch.abs(x[:, :, :, 1:] - x[:, :, :, :-1])
         return diff_h.mean() + diff_w.mean()
     
     def conservation_intensity(self, x):
-        """
-        Conservation of image intensity (mass).
-        
-        L_cons = (mean(x) - running_mean)^2
-        
-        This prevents the generator from drifting into producing
-        images that are too dark or too bright.
-        
-        Args:
-            x: [B, C, H, W] images
-        Returns:
-            cons_loss: scalar
-        """
+        """Penalize deviation from running mean intensity."""
         batch_mean = x.mean()
         
-        # Update running statistics
         if self.training:
             with torch.no_grad():
-                self.intensity_mean = (
-                    self.intensity_alpha * self.intensity_mean +
-                    (1 - self.intensity_alpha) * batch_mean.detach()
-                )
-        
-        # Conservation loss: penalize deviation from running mean
-        if self.intensity_count > 100:  # Only after some warmup
-            return ((batch_mean - self.intensity_mean) ** 2).mean()
-        return torch.tensor(0.0, device=x.device)
+                self.step_count += 1
+                alpha = min(0.99, 1.0 - 1.0 / (self.step_count.float() + 1))
+                self.intensity_ema = alpha * self.intensity_ema + (1 - alpha) * batch_mean
+        
+        # Only activate after warmup (100 steps)
+        if self.step_count > 100:
+            return (batch_mean - self.intensity_ema.detach()) ** 2
+        return torch.zeros(1, device=x.device, requires_grad=True).squeeze()
     
     def spectral_regularizer(self, x):
-        """
-        Spectral regularizer: penalize high-frequency content.
-        
-        Uses FFT and penalizes high-frequency components.
-        This prevents high-frequency artifacts (checkerboard patterns).
+        """Penalize high-frequency energy (anti-checkerboard)."""
+        B, C, H, W = x.shape
         
-        Args:
-            x: [B, C, H, W] images
-        Returns:
-            spec_loss: scalar
-        """
         # 2D FFT
-        x_fft = torch.fft.fft2(x)
-        x_fft_shift = torch.fft.fftshift(x_fft)
+        x_fft = torch.fft.rfft2(x, norm='ortho')
+        mag = torch.abs(x_fft)
         
-        # Create high-frequency mask (center is low frequency)
-        B, C, H, W = x.shape
-        h_center, w_center = H // 2, W // 2
+        # High-frequency mask: upper-right quadrant of frequency space
+        # For rfft2, output shape is [B, C, H, W//2+1]
+        freq_h = torch.arange(H, device=x.device).float()
+        freq_w = torch.arange(W // 2 + 1, device=x.device).float()
         
-        y, x_coord = torch.meshgrid(
-            torch.arange(H, device=x.device),
-            torch.arange(W, device=x.device),
-            indexing='ij'
-        )
-        dist = torch.sqrt((y - h_center) ** 2 + (x_coord - w_center) ** 2)
+        # Normalize frequencies to [0, 1]
+        freq_h = torch.min(freq_h, H - freq_h) / (H / 2)
+        freq_w = freq_w / (W / 2)
         
-        # High frequency: distance > quarter of image size
-        high_freq_mask = (dist > min(H, W) / 4).float()
+        # Distance from DC (center)
+        dist = torch.sqrt(freq_h.unsqueeze(1) ** 2 + freq_w.unsqueeze(0) ** 2)
         
-        # Penalize high-frequency magnitude
-        spec_mag = torch.abs(x_fft_shift)
-        high_freq_energy = (spec_mag * high_freq_mask.unsqueeze(0).unsqueeze(0)).mean()
+        # High frequency: distance > 0.5 (half Nyquist)
+        high_mask = (dist > 0.5).float()
         
-        return high_freq_energy
+        high_energy = (mag * high_mask.unsqueeze(0).unsqueeze(0)).mean()
+        return high_energy
     
     def gradient_penalty(self, x):
-        """
-        Sobolev gradient penalty.
-        
-        L_grad = ||∇x||² (mean squared gradient magnitude)
-        
-        This prevents the generator from creating regions where
-        gradients explode (common in GAN-like training).
-        For diffusion, this helps stabilize the noise prediction.
-        
-        Args:
-            x: [B, C, H, W] images
-        Returns:
-            grad_loss: scalar
-        """
+        """Sobolev L2 gradient penalty."""
         grad_h = x[:, :, 1:, :] - x[:, :, :-1, :]
         grad_w = x[:, :, :, 1:] - x[:, :, :, :-1]
-        
-        grad_mag = (grad_h ** 2).mean() + (grad_w ** 2).mean()
-        return grad_mag
+        return (grad_h ** 2).mean() + (grad_w ** 2).mean()
     
     def forward(self, x0_hat, x_ref=None):
         """
-        Compute total physics loss.
-        
         Args:
             x0_hat: Estimated clean image [B, C, H, W]
-            x_ref: Optional ground truth reference (for intensity tracking)
-        
+            x_ref: Ground truth (unused, kept for API compat)
         Returns:
-            total_loss: Combined physics regularizer (scalar)
-            loss_dict: Dict of individual losses
+            total_loss, loss_dict
         """
         losses = {}
+        total = torch.zeros(1, device=x0_hat.device, requires_grad=True).squeeze()
         
-        # Total Variation
         if self.tv_weight > 0:
-            losses['tv'] = self.total_variation(x0_hat)
+            tv = self.total_variation(x0_hat)
+            losses['tv'] = tv
+            total = total + self.tv_weight * tv
         
-        # Conservation of Intensity
         if self.cons_weight > 0:
-            losses['cons'] = self.conservation_intensity(x0_hat)
+            cons = self.conservation_intensity(x0_hat)
+            losses['cons'] = cons
+            total = total + self.cons_weight * cons
         
-        # Spectral Regularizer
         if self.spec_weight > 0:
-            losses['spec'] = self.spectral_regularizer(x0_hat)
+            spec = self.spectral_regularizer(x0_hat)
+            losses['spec'] = spec
+            total = total + self.spec_weight * spec
         
-        # Gradient Penalty
         if self.grad_weight > 0:
-            losses['grad'] = self.gradient_penalty(x0_hat)
-        
-        # Weighted sum
-        total = (
-            self.tv_weight * losses.get('tv', 0.0) +
-            self.cons_weight * losses.get('cons', 0.0) +
-            self.spec_weight * losses.get('spec', 0.0) +
-            self.grad_weight * losses.get('grad', 0.0)
-        )
+            grad = self.gradient_penalty(x0_hat)
+            losses['grad'] = grad
+            total = total + self.grad_weight * grad
         
         return total, losses
 
 
 class DDIMEstimator:
-    """
-    DDIM clean-sample estimator for physics loss computation.
-    
-    From the Bastek & Sun (ICLR 2025) pattern:
-    x̂₀ = (x_t - √(1-ᾱ_t) · ε_pred) / √(ᾱ_t)
-    
-    This provides an estimate of the clean sample at training time
-    without requiring full reverse diffusion.
-    """
+    """DDIM one-step clean sample estimation."""
     
     @staticmethod
     def estimate_x0(x_t, eps_pred, alpha_bar_t):
         """
-        Estimate clean sample from noisy sample and predicted noise.
+        x̂₀ = (x_t - √(1-ᾱ_t) · ε_pred) / √(ᾱ_t)
         
         Args:
-            x_t: Noisy sample [B, C, H, W]
-            eps_pred: Predicted noise [B, C, H, W]
-            alpha_bar_t: Cumulative product of alphas at timestep t [B]
-        
-        Returns:
-            x0_hat: Estimated clean sample [B, C, H, W]
+            x_t: [B, C, H, W]
+            eps_pred: [B, C, H, W]
+            alpha_bar_t: [B] — cumulative alpha at timestep t
         """
-        alpha_bar_t = alpha_bar_t.reshape(-1, 1, 1, 1)
-        x0_hat = (x_t - torch.sqrt(1 - alpha_bar_t) * eps_pred) / torch.sqrt(alpha_bar_t)
-        return x0_hat
-    
-    @staticmethod
-    def estimate_noise(x_t, x0_hat, alpha_bar_t):
-        """Reverse: estimate noise from clean sample."""
-        alpha_bar_t = alpha_bar_t.reshape(-1, 1, 1, 1)
-        eps_pred = (x_t - torch.sqrt(alpha_bar_t) * x0_hat) / torch.sqrt(1 - alpha_bar_t)
-        return eps_pred
+        a = alpha_bar_t.reshape(-1, 1, 1, 1)
+        x0_hat = (x_t - torch.sqrt(1 - a) * eps_pred) / (torch.sqrt(a) + 1e-8)
+        # Clamp to prevent extreme values early in training
+        return x0_hat.clamp(-5, 5)