fix def loss

training/s2f_trainer.py

@@ -16,12 +16,17 @@ if S2F_ROOT not in sys.path:
 
 from models.s2f_model import create_settings_channels
 from utils.substrate_settings import compute_settings_normalization
-from utils.metrics import calculate_psnr, calculate_ssim_tensor, calculate_pearson_correlation
+from utils.metrics import (
+    WFMRMELoss,
+    calculate_psnr,
+    calculate_ssim_tensor,
+    calculate_pearson_correlation,
+)
 from scipy.stats import pearsonr
 
 
 class S2FLoss(nn.Module):
-    """S2F loss: reconstruction (L1) + GAN + optional force consistency."""
+    """S2F loss: reconstruction (WFM-RME by default) + GAN + optional force consistency."""
     def __init__(self, lambda_L1=100.0, lambda_gan=1.0, lambda_force=1.0,
                  gan_mode='vanilla', custom_loss=None, use_force_consistency=False,
                  force_consistency_target='mean'):
@@ -32,7 +37,7 @@ class S2FLoss(nn.Module):
         self.gan_mode = gan_mode
         self.use_force_consistency = use_force_consistency
         self.force_consistency_target = force_consistency_target
-        self.reconstruction_loss = custom_loss if custom_loss is not None else nn.L1Loss()
+        self.reconstruction_loss = custom_loss if custom_loss is not None else WFMRMELoss()
         self.force_consistency_loss = nn.MSELoss() if use_force_consistency else None
         self.gan_loss = nn.BCEWithLogitsLoss() if gan_mode == 'vanilla' else nn.MSELoss()
 

utils/metrics.py

@@ -106,6 +106,36 @@ def _force_mag_wfm(f):
     return np.sqrt(fx**2 + fy**2)
 
 
+def _force_magnitude_tensor(x: torch.Tensor) -> torch.Tensor:
+    """Per-pixel force magnitude: (B,1,H,W) uses channel 0 as magnitude; (B,2+,H,W) uses sqrt(fx^2+fy^2)."""
+    if x.dim() != 4:
+        raise ValueError(f"Expected 4D tensor, got shape {tuple(x.shape)}")
+    c = x.size(1)
+    if c == 1:
+        return x[:, 0]
+    if c >= 2:
+        return torch.sqrt(x[:, 0].pow(2) + x[:, 1].pow(2))
+    raise ValueError(f"Expected at least 1 channel, got {c}")
+
+
+class WFMRMELoss(nn.Module):
+    """Weighted force-magnitude relative error; matches ``wfm_relative_magnitude_error`` (differentiable)."""
+
+    def __init__(self, eps: float = 1e-8):
+        super().__init__()
+        self.eps = eps
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        mag_t = _force_magnitude_tensor(target)
+        mag_p = _force_magnitude_tensor(pred)
+        if mag_t.shape != mag_p.shape:
+            raise ValueError(f"Shape mismatch after magnitude: {mag_t.shape} vs {mag_p.shape}")
+        fbar = mag_t.mean().clamp_min(self.eps)
+        w = mag_t / fbar
+        rel = (mag_p - mag_t).abs() / (mag_t + self.eps)
+        return (w * rel).mean()
+
+
 def wfm_correlation(y_true, y_pred, mode="magnitude"):
     """Pearson correlation between prediction and ground truth (magnitude mode for heatmaps)."""
     t = _ensure_shape_wfm(_to_numpy_wfm(y_true))