wiring up softplus properly

makani/utils/constraints.py

@@ -31,11 +31,22 @@ class NonNegativeConstraint(nn.Module):
     zero sits at x_norm = -bias/scale. The offset = bias/scale is precomputed
     in the constructor so the forward pass is cheap.
 
-    Training mode: smooth multiplicative approximation x*sigmoid(x/eps) applied
-    in the shifted (physical-zero-centered) space so gradients flow for slightly
-    negative values.
-
-    Eval/inference mode: hard clamp, guaranteeing x_raw >= 0 before any
+    Training mode: a smooth soft clamp applied in the shifted (physical-zero-centered)
+    space so gradients flow for slightly negative values. Two shapes are available via
+    ``mode``:
+
+      - "silu" (default): ``x * sigmoid(x/eps)``. Smooth, but *non-monotonic* -- it dips
+        below zero for x < 0 with a negative gradient there. For a channel whose target
+        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
+        (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.
+
+    Eval/inference mode (both): hard clamp, guaranteeing x_raw >= 0 before any
     downstream conservation corrections.
 
     Args:
@@ -45,11 +56,17 @@ class NonNegativeConstraint(nn.Module):
         bias:           normalization bias tensor, shape (1, C, 1, 1) or None.
         scale:          normalization scale tensor, shape (1, C, 1, 1) or None.
         eps:            transition width for the soft clamp (normalized units).
+        mode:           "silu" (default, legacy) or "softplus" (monotonic, recommended).
+        leak:           negative-side gradient floor for the "softplus" mode (ignored for "silu").
     """
 
-    def __init__(self, channel_names, names_to_clamp, bias=None, scale=None, eps=0.1):
+    def __init__(self, channel_names, names_to_clamp, bias=None, scale=None, eps=0.1, mode="silu", leak=0.02):
         super().__init__()
+        if mode not in ("silu", "softplus"):
+            raise ValueError(f"NonNegativeConstraint mode must be 'silu' or 'softplus', got {mode!r}")
         self.eps = eps
+        self.mode = mode
+        self.leak = leak
 
         # resolve names to indices, skipping any not present in channel_names
         chan_idx = [channel_names.index(n) for n in names_to_clamp if n in channel_names]
@@ -73,7 +90,10 @@ def forward(self, x):
         if self.training:
             # shift so physical zero maps to 0, apply smooth clamp, shift back
             w_shifted = w + offset if offset is not None else w
-            w = w_shifted * torch.sigmoid(w_shifted / self.eps)
+            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)
             if offset is not None:
                 w = w - offset
         else:

tests/test_constraints.py

@@ -365,76 +365,83 @@ def _make(self, names_to_clamp=None, means=None, stds=None, **kwargs):
         return c.to(self.device)
 
     # --- eval / hard clamp ---
-    def test_eval_hard_clamp_no_normalization(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_eval_hard_clamp_no_normalization(self, mode):
         """Eval mode: constrained channels are >= 0; unconstrained channels unchanged."""
         B, C, H, W = 2, len(self.ALL_CHANNELS), 8, 8
-        c = self._make()
+        c = self._make(mode=mode)
         c.eval()
         x = torch.randn(B, C, H, W, device=self.device)
         y = c(x)
         self.assertTrue((y[:, self.CLAMP_IDX, :, :] >= 0).all().item())
         unconstrained = [i for i in range(C) if i not in self.CLAMP_IDX]
         self.assertTrue(compare_tensors("unconstrained channels", y[:, unconstrained, :, :], x[:, unconstrained, :, :]))
 
-    def test_eval_hard_clamp_with_normalization(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_eval_hard_clamp_with_normalization(self, mode):
         """Eval mode: x_raw = y_norm * scale + bias >= 0 after clamping."""
         B, C, H, W = 2, len(self.ALL_CHANNELS), 6, 6
         means = torch.tensor([0.0, 5.0, 270.0, 3.0, 250.0])
         stds  = torch.tensor([1.0, 2.0,  10.0, 1.5,   8.0])
-        c = self._make(means=means, stds=stds)
+        c = self._make(means=means, stds=stds, mode=mode)
         c.eval()
         x = torch.randn(B, C, H, W, device=self.device) * 3.0
         y = c(x)
         for i, ci in enumerate(self.CLAMP_IDX):
             x_raw = y[:, ci, :, :] * stds[ci].item() + means[ci].item()
             self.assertTrue((x_raw >= -1e-6).all().item(), f"channel {self.ALL_CHANNELS[ci]} has negative physical values")
 
-    def test_eval_positive_input_unchanged(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_eval_positive_input_unchanged(self, mode):
         """Eval mode: values already above physical zero are not modified."""
         B, C, H, W = 2, len(self.ALL_CHANNELS), 4, 4
         means = torch.tensor([0.0, 1.0, 270.0, 2.0, 250.0])
         stds  = torch.ones(len(self.ALL_CHANNELS))
-        c = self._make(means=means, stds=stds)
+        c = self._make(means=means, stds=stds, mode=mode)
         c.eval()
         x = torch.ones(B, C, H, W, device=self.device) * 5.0
         y = c(x)
         self.assertTrue(compare_tensors("positive inputs unchanged", y, x))
 
     # --- training / soft clamp ---
-    def test_train_slightly_negative_not_zeroed(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_slightly_negative_not_zeroed(self, mode):
         """Training mode: slightly negative values are not exactly zeroed (gradient path open)."""
         B, C, H, W = 1, len(self.ALL_CHANNELS), 4, 4
-        c = self._make(names_to_clamp=["q850"])
+        c = self._make(names_to_clamp=["q850"], mode=mode)
         c.train()
         x = torch.full((B, C, H, W), -0.05, device=self.device)
         y = c(x)
         self.assertFalse((y[:, [1], :, :] == 0).all().item())
 
-    def test_train_large_positive_identity(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_large_positive_identity(self, mode):
         """Training mode: large positive values pass through essentially unchanged."""
         B, C, H, W = 2, len(self.ALL_CHANNELS), 4, 4
-        c = self._make(eps=0.1)
+        c = self._make(eps=0.1, mode=mode)
         c.train()
         x = torch.ones(B, C, H, W, device=self.device) * 5.0
         y = c(x)
         self.assertTrue(compare_tensors("large positive passthrough", y[:, self.CLAMP_IDX, :, :], x[:, self.CLAMP_IDX, :, :], atol=1e-3))
 
-    def test_train_gradient_flows(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_gradient_flows(self, mode):
         """Training mode: gradient is nonzero for slightly negative inputs."""
         B, C, H, W = 1, len(self.ALL_CHANNELS), 4, 4
-        c = self._make()
+        c = self._make(mode=mode)
         c.train()
         x = torch.full((B, C, H, W), -0.2, device=self.device, requires_grad=True)
         c(x).sum().backward()
         self.assertIsNotNone(x.grad)
         self.assertFalse((x.grad[:, self.CLAMP_IDX, :, :] == 0).all().item())
 
-    def test_train_normalization_offset(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_normalization_offset(self, mode):
         """Training mode: with normalization, the clamp boundary is at physical zero."""
         B, C, H, W = 1, len(self.ALL_CHANNELS), 4, 4
         means = torch.tensor([0.0, 4.0, 270.0, 6.0, 250.0])
         stds  = torch.tensor([1.0, 2.0,  10.0, 3.0,   8.0])
-        c = self._make(means=means, stds=stds, eps=0.01)
+        c = self._make(means=means, stds=stds, eps=0.01, mode=mode)
         c.train()
         # set constrained channels to physical zero in normalized space
         x = torch.zeros(B, C, H, W, device=self.device)
@@ -447,10 +454,11 @@ def test_train_normalization_offset(self):
                                             x_raw, torch.zeros_like(x_raw), atol=0.1))
 
     # --- mode switching ---
-    def test_train_eval_switch(self):
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_eval_switch(self, mode):
         """Switching train/eval changes hard vs soft clamping on the same instance."""
         B, C, H, W = 1, len(self.ALL_CHANNELS), 4, 4
-        c = self._make(names_to_clamp=["q850"])
+        c = self._make(names_to_clamp=["q850"], mode=mode)
         x = torch.full((B, C, H, W), -1.0, device=self.device)
         c.train()
         y_train = c(x)
@@ -460,6 +468,27 @@ def test_train_eval_switch(self):
         self.assertTrue(compare_tensors("hard clamp to zero", y_eval[:, [1], :, :],
                                         torch.zeros_like(y_eval[:, [1], :, :])))
 
+    # --- soft-clamp shape: the reason "softplus" mode exists ---
+    @parameterized.expand([("silu",), ("softplus",)])
+    def test_train_soft_clamp_shape(self, mode):
+        """softplus is monotonic (grad > 0 everywhere), so a below-target channel is always
+        pushed up. silu is non-monotonic: it has a negative-gradient dip that can drive a
+        near-zero channel (e.g. q50) deeper negative until the gradient vanishes."""
+        c = self._make(names_to_clamp=["q850"], mode=mode)  # no normalization -> input is the shifted space
+        c.train()
+        C = len(self.ALL_CHANNELS)
+        # sweep the clamped channel (index 1) across negative values; keep the rest positive
+        x = torch.full((1, C, 1, 40), 5.0, device=self.device)
+        x[:, 1, 0, :] = torch.linspace(-1.0, 0.5, 40, device=self.device)
+        x = x.requires_grad_(True)
+        c(x).sum().backward()
+        # d(sum output)/dx is the elementwise soft-clamp derivative on this channel
+        g = x.grad[:, 1, 0, :]
+        if mode == "softplus":
+            self.assertTrue((g > 0).all().item(), "softplus soft clamp must be monotonic (grad > 0)")
+        else:
+            self.assertTrue((g < 0).any().item(), "silu soft clamp has a non-monotonic negative-gradient dip")
+
 
 @parameterized_class(("device",), _devices)
 class TestHydrostaticBalanceProjection(unittest.TestCase):