From dcfc49403c7563f5541eac5d2cbba179cc4e88a6 Mon Sep 17 00:00:00 2001 From: Thorsten Kurth Date: Tue, 14 Jul 2026 03:13:47 -0700 Subject: [PATCH] added changes --- makani/utils/constraints.py | 19 +++++++++++++------ tests/test_constraints.py | 34 +++++++++++++++++++++++++++++++--- 2 files changed, 44 insertions(+), 9 deletions(-) diff --git a/makani/utils/constraints.py b/makani/utils/constraints.py index 5d430bca..3a489a54 100644 --- a/makani/utils/constraints.py +++ b/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: diff --git a/tests/test_constraints.py b/tests/test_constraints.py index a6610455..a7337e7d 100644 --- a/tests/test_constraints.py +++ b/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):