lowered floor of softclamp

makani/utils/constraints.py

@@ -40,11 +40,16 @@ class NonNegativeConstraint(nn.Module):
         sits at the physical-zero floor (e.g. stratospheric specific humidity q50), that
         negative-side gradient drives the prediction ever more negative until the gradient
         vanishes, a self-reinforcing collapse to 0 with no way back.
-      - "softplus": a leaky blend ``leak*x + (1-leak)*eps*softplus(x/eps)``. Monotonic
+      - "softplus": a leaky blend ``leak*x + (1-leak)*eps*(softplus(x/eps) - ln2)``. Monotonic
         (gradient > 0 everywhere) so a below-target prediction is always pushed up, with a
-        gradient floor of ``leak`` so an already-collapsed channel can still recover. It is
-        identity on the bulk (large x), matching "silu" spectrally, and only lifts the floor
-        slightly (physical zero -> ~(1-leak)*eps*ln2), which nudges values off the dead floor.
+        gradient floor of ``leak`` so an already-collapsed channel can still recover. The
+        ``-ln2`` constant pins the floor to a fixed point at physical zero (w(0)=0), matching
+        the eval hard-clamp floor: without it the raw ``softplus`` sits ~(1-leak)*eps*ln2 above
+        physical zero, biasing genuinely-dry channels (e.g. stratospheric q50) upward. The loss
+        then drives raw predictions negative to cancel that bias, and the inference hard clamp
+        flattens them to 0 -- a collapse routed through the floor mismatch rather than a negative
+        gradient. It is identity minus a negligible ~(1-leak)*eps*ln2 on the bulk, matching "silu"
+        spectrally.
 
     Eval/inference mode (both): hard clamp, guaranteeing x_raw >= 0 before any
     downstream conservation corrections.
@@ -92,8 +97,10 @@ def forward(self, x):
             w_shifted = w + offset if offset is not None else w
             if self.mode == "silu":
                 w = w_shifted * torch.sigmoid(w_shifted / self.eps)
-            else:  # "softplus": monotonic leaky blend (no collapse-inducing negative-gradient dip)
-                w = self.leak * w_shifted + (1.0 - self.leak) * self.eps * F.softplus(w_shifted / self.eps)
+            else:  # "softplus": monotonic leaky blend (no collapse-inducing negative-gradient dip),
+                   # with the constant softplus(0)=ln2 subtracted so physical zero is a fixed point
+                   # (w(0)=0), matching the eval hard-clamp floor instead of sitting ~(1-leak)*eps*ln2 above it.
+                w = self.leak * w_shifted + (1.0 - self.leak) * self.eps * (F.softplus(w_shifted / self.eps) - math.log(2.0))
             if offset is not None:
                 w = w - offset
         else:

tests/test_constraints.py

@@ -416,13 +416,41 @@ def test_train_slightly_negative_not_zeroed(self, mode):
 
     @parameterized.expand([("silu",), ("softplus",)])
     def test_train_large_positive_identity(self, mode):
-        """Training mode: large positive values pass through essentially unchanged."""
+        """Training mode: on the bulk the clamp has unit slope, so it preserves spatial
+        structure (only the l=0 DC mode may shift). silu is exact identity there; softplus
+        is identity up to the constant (1-leak)*eps*ln2 introduced to pin the floor to physical
+        zero -- a DC offset the decoder bias absorbs, so we check the slope, not the offset."""
         B, C, H, W = 2, len(self.ALL_CHANNELS), 4, 4
         c = self._make(eps=0.1, mode=mode)
         c.train()
-        x = torch.ones(B, C, H, W, device=self.device) * 5.0
+        x1 = torch.ones(B, C, H, W, device=self.device) * 5.0
+        y1 = c(x1)
+
+        with self.subTest("bulk unit slope"):
+            # two bulk inputs one unit apart: unit slope <=> y2 - y1 == x2 - x1 == 1
+            y2 = c(x1 + 1.0)
+            slope = (y2 - y1)[:, self.CLAMP_IDX, :, :]
+            self.assertTrue(compare_tensors("bulk unit slope", slope, torch.ones_like(slope), atol=1e-3))
+
+        if mode == "silu":
+            with self.subTest("exact passthrough (silu has zero DC offset)"):
+                self.assertTrue(compare_tensors("large positive passthrough",
+                                                y1[:, self.CLAMP_IDX, :, :], x1[:, self.CLAMP_IDX, :, :], atol=1e-3))
+
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_floor_fixed_point(self, mode):
+        """The training soft clamp must map physical zero to physical zero (w(0)=0 in the
+        shifted space). Otherwise the training floor sits above the eval hard-clamp floor,
+        biasing genuinely-dry channels (e.g. q50) upward; the loss then drives raw predictions
+        negative to cancel the bias and the inference hard clamp flattens them to 0."""
+        c = self._make(names_to_clamp=["q850"], mode=mode)  # no normalization -> input is the shifted space
+        c.train()
+        C = len(self.ALL_CHANNELS)
+        x = torch.full((1, C, 1, 1), 5.0, device=self.device)
+        x[:, 1, 0, 0] = 0.0  # physical zero for the clamped channel
         y = c(x)
-        self.assertTrue(compare_tensors("large positive passthrough", y[:, self.CLAMP_IDX, :, :], x[:, self.CLAMP_IDX, :, :], atol=1e-3))
+        self.assertAlmostEqual(y[0, 1, 0, 0].item(), 0.0, places=5,
+                               msg=f"{mode} training floor must be a fixed point at physical zero")
 
     @parameterized.expand([("silu",), ("softplus",)])
     def test_train_gradient_flows(self, mode):