Exclude unpopulated lag bins from plateau and half-power lag

eddy/structurefunction.py

@@ -1630,6 +1630,47 @@ def max_lag_x(self):
     def max_lag_y(self):
         return self.S2.shape[1] // 2
 
+    @property
+    def counts_x(self):
+        """Pair counts on the radial slice, aligned with :attr:`S2_x`."""
+        return self.counts[self.max_lag_x:, self.max_lag_y]
+
+    @property
+    def counts_y(self):
+        """Pair counts on the azimuthal slice, aligned with :attr:`S2_y`."""
+        return self.counts[self.max_lag_x, self.max_lag_y:]
+
+    def _populated_slice(self, axis):
+        """One on-axis slice, with unpopulated lag bins masked to ``nan``.
+
+        :func:`_s2_kernel` leaves a lag bin with no finite pairs at its
+        initialized ``0.0``, indistinguishable from a genuine ``S_2 = 0``.
+        Any statistic that reduces over a whole slice (:meth:`plateau`) or
+        hunts for a level crossing (:meth:`half_power_lag`) has to drop those
+        bins explicitly: an annulus whose lag range runs off the grid -- or
+        into a deprojection mask -- otherwise reports a plateau pulled towards
+        zero, and with it a correspondingly short correlation scale. The GRF
+        fit already masks the same bins (``counts > 0`` in :meth:`fit_GRF`),
+        so this keeps the model-free and fitted paths consistent.
+
+        Returns:
+            tuple: ``(lags, S2)``, the lag axis and a NaN-masked *copy* of the
+            slice (``S2_x`` / ``S2_y`` themselves are never mutated). Objects
+            carrying no real pair counts (``counts`` defaulted to ones, as
+            :func:`gaussian_beam_s2` does) are returned unmasked.
+        """
+        if axis == "x":
+            lags, s2, cnt = self.lags_x, self.S2_x, self.counts_x
+        elif axis == "y":
+            lags, s2, cnt = self.lags_y, self.S2_y, self.counts_y
+        else:
+            raise ValueError("axis must be 'x' or 'y', got {!r}.".format(axis))
+        s2 = np.array(s2, dtype=float)
+        cnt = np.asarray(cnt)
+        if cnt.shape == s2.shape:
+            s2[cnt <= 0] = np.nan
+        return np.asarray(lags), s2
+
     @property
     def extent(self):
         """Matplotlib ``extent`` for ``imshow(S2)``: (l_y_min, l_y_max,
@@ -1955,7 +1996,10 @@ def plateau(self, frac=0.5, stat="median"):
         radial and azimuthal slices and returns a robust statistic,
         ``median`` by default, so a single noisy outlier (common in a
         denoised / low-pair-count slice) does not set the level the way
-        ``np.nanmax`` would.
+        ``np.nanmax`` would. Lag bins with no contributing pairs are
+        excluded (see :meth:`_populated_slice`) -- pooling them would drag
+        the plateau towards zero on any annulus whose lag range leaves the
+        grid or the deprojection mask.
 
         Args:
             frac (float): Fraction of the lag range treated as "large lag".
@@ -1966,10 +2010,11 @@ def plateau(self, frac=0.5, stat="median"):
             float: the plateau estimate (``nan`` if no finite outer cells).
         """
         vals = []
-        for lags, s2 in ((self.lags_x, self.S2_x), (self.lags_y, self.S2_y)):
+        for axis in ("x", "y"):
+            lags, s2 = self._populated_slice(axis)
             if lags.size == 0:
                 continue
-            sel = np.asarray(s2)[lags >= frac * lags[-1]]
+            sel = s2[lags >= frac * lags[-1]]
             vals.append(sel[np.isfinite(sel)])
         pooled = np.concatenate(vals) if vals else np.array([])
         if pooled.size == 0:
@@ -1983,7 +2028,10 @@ def half_power_lag(self, axis="x", level=0.5, plateau=None):
         crosses half its plateau (``= 1.18 ell`` for a Gaussian kernel, but
         no Gaussian assumption is made). Located by linear interpolation of
         the FIRST upward crossing, so it is robust to non-monotonic wiggles
-        at larger lag.
+        at larger lag. Unpopulated lag bins are excluded rather than read as
+        ``S_2 = 0`` (see :meth:`_populated_slice`); if the bin below the
+        crossing is itself empty the crossing bin's own lag is returned,
+        there being nothing to interpolate from.
 
         Args:
             axis ({'x', 'y'}): ``'x'`` -> radial slice ``S2_x`` (lag in the
@@ -1998,26 +2046,21 @@ def half_power_lag(self, axis="x", level=0.5, plateau=None):
             float: the crossing lag in that axis's units (``nan`` if the
             slice never reaches the level within the lag range).
         """
-        if axis == "x":
-            lags, s2 = self.lags_x, self.S2_x
-        elif axis == "y":
-            lags, s2 = self.lags_y, self.S2_y
-        else:
-            raise ValueError("axis must be 'x' or 'y', got {!r}.".format(axis))
+        lags, s2 = self._populated_slice(axis)
         if plateau is None:
             plateau = self.plateau()
         target = level * plateau
         if not np.isfinite(target) or target <= 0:
             return float("nan")
-        above = np.isfinite(s2) & (np.asarray(s2) >= target)
+        above = np.isfinite(s2) & (s2 >= target)
         if not above.any():
             return float("nan")
         i = int(np.argmax(above))                  # first crossing index
         if i == 0:
             return float(lags[0])
         x0, x1, y0, y1 = lags[i - 1], lags[i], s2[i - 1], s2[i]
-        if y1 == y0:
-            return float(x1)
+        if not np.isfinite(y0) or y1 == y0:
+            return float(x1)                       # nothing to interpolate from
         return float(x0 + (target - y0) * (x1 - x0) / (y1 - y0))
 
     def reliability_weight(self, kind="counts"):
@@ -2048,9 +2091,7 @@ def reliability_weight(self, kind="counts"):
             these are quick single-annulus proxies.
         """
         if kind == "counts":
-            cx = self.counts[self.max_lag_x:, self.max_lag_y]
-            cy = self.counts[self.max_lag_x, self.max_lag_y:]
-            w = np.nansum(cx) + np.nansum(cy)
+            w = np.nansum(self.counts_x) + np.nansum(self.counts_y)
             return float(w) if np.isfinite(w) else 0.0
         if kind == "neff":
             ell_r = self.half_power_lag("x")
@@ -2918,7 +2959,7 @@ def calculate_modal_power(self, modes=(1,), axis=None, p0=None):
             * ``frac_of_data_total``: ``sum_m frac_of_data``, ``(N_ref,)``.
             * ``popt``, ``perr``: the raw fit outputs.
         """
-        popt, perr = self.fit_spiral(modes=modes, axis=axis, p0=p0)
+        popt, perr, _ = self.fit_spiral(modes=modes, axis=axis, p0=p0)
         offset = popt[:, 0]
         amps = popt[:, 1:]
         power = amps ** 2

tests/test_structurefunction.py

@@ -272,6 +272,109 @@ def test_measure_heuristics_returns_finite_scalars():
     assert T1a_raw == pytest.approx(2.0 * sig_hat ** 2)
 
 
+def _empty_the_largest_lags(res, keep_x, keep_y=None):
+    """Emulate ``_s2_kernel`` on an annulus whose lag range runs off the grid:
+    the largest lags get no pairs, so ``counts`` is 0 and ``S_2`` is left at
+    its initialized 0.0 -- indistinguishable, without ``counts``, from a
+    genuine ``S_2 = 0``. ``keep_y=None`` leaves the azimuthal slice fully
+    populated, which is the real geometry: a ring near the edge of the map
+    loses its large *radial* lags while still spanning 360 deg."""
+    cx, cy = res.max_lag_x, res.max_lag_y
+    res.counts[cx + keep_x:, cy] = 0
+    res.S2_x[keep_x:] = 0.0
+    if keep_y is not None:
+        res.counts[cx, cy + keep_y:] = 0
+        res.S2_y[keep_y:] = 0.0
+
+
+def _old_pooled_plateau(res, frac=0.5):
+    """The plateau eddy computed before the ``counts`` mask: outer-lag median
+    with unpopulated bins read as genuine zeros."""
+    vals = []
+    for lags, s2 in ((res.lags_x, res.S2_x), (res.lags_y, res.S2_y)):
+        sel = np.asarray(s2)[lags >= frac * lags[-1]]
+        vals.append(sel[np.isfinite(sel)])
+    return float(np.median(np.concatenate(vals)))
+
+
+def test_plateau_and_half_power_lag_ignore_unpopulated_bins():
+    """An annulus whose radial lags run off the grid must still measure the
+    right plateau from its (intact) azimuthal slice, and must report an
+    unmeasurable radial scale as NaN rather than inventing a short one.
+
+    Before the ``counts`` mask the empty bins entered ``plateau``'s outer-lag
+    median as zeros, halving it; every half-power lag measured against that
+    level then came out spuriously short.
+    """
+    field = _smoothed_polar_field(n_r=100, n_phi=400)
+    dx, dy = 0.05, 3.0
+    x_axis = 0.5 + np.arange(field.shape[0]) * dx
+    stack = StructureFunction2DStack.from_array(
+        field, ref_rs=[1.5], x_axis=x_axis, dx=dx, dy=dy,
+        ref_band=0.05, max_lag_x=40, max_lag_y=40, n_bins=25,
+    )
+    res = stack.results[0]
+    plateau_before = res.plateau()
+    ell_r_before = res.half_power_lag("x")
+    ell_phi_before = res.half_power_lag("y")
+    assert np.isfinite([plateau_before, ell_r_before, ell_phi_before]).all()
+
+    # Radial slice dies past lag 5*dx = 0.25", i.e. short of ell_r itself.
+    _empty_the_largest_lags(res, keep_x=6)
+
+    # Plateau still comes off the intact azimuthal slice, ...
+    assert res.plateau() == pytest.approx(plateau_before, rel=0.1)
+    # ... the azimuthal scale is unmoved, ...
+    assert res.half_power_lag("y") == pytest.approx(ell_phi_before, rel=0.1)
+    # ... and the radial scale, no longer reachable, is NaN not a short lag.
+    assert np.isnan(res.half_power_lag("x"))
+
+    # The regression being guarded: reading the empty bins as zeros halves the
+    # plateau, and both lags measured against that level come out short.
+    old_plateau = _old_pooled_plateau(res)
+    assert old_plateau < 0.6 * plateau_before
+    assert res.half_power_lag("y", plateau=old_plateau) < 0.75 * ell_phi_before
+    assert res.half_power_lag("x", plateau=old_plateau) < 0.75 * ell_r_before
+
+    # The emptied bins are excluded, not read as zeros ...
+    _, s2_x = res._populated_slice("x")
+    assert np.all(np.isnan(s2_x[6:])) and np.all(np.isfinite(s2_x[:6]))
+    # ... and the stored slice itself is never mutated by the masking.
+    assert np.all(res.S2_x[6:] == 0.0)
+
+
+def test_heuristics_unchanged_when_outer_annuli_lose_their_largest_lags():
+    """Stationary field, and the outermost annuli lose their largest lags the
+    way real edge annuli do. Plateaus, ``ell_r`` and the radial trend T3 must
+    all be unmoved -- the failure mode was ell_r shrinking with radius, which
+    manufactures a T3 < 0 for a field that has no radial trend at all."""
+    field = _smoothed_polar_field(n_r=120, n_phi=140)
+    dx, dy = 0.05, 3.0
+    x_axis = 0.5 + np.arange(field.shape[0]) * dx
+    kw = dict(ref_rs=np.linspace(1.0, 5.5, 10), x_axis=x_axis, dx=dx, dy=dy,
+              ref_band=0.05, max_lag_x=25, max_lag_y=30, n_bins=25)
+    pristine = StructureFunction2DStack.from_array(field, **kw)
+    damaged = StructureFunction2DStack.from_array(field, **kw)
+    for res in damaged.results[-4:]:
+        _empty_the_largest_lags(res, keep_x=18, keep_y=22)
+
+    np.testing.assert_allclose(damaged.plateaus(), pristine.plateaus(),
+                               rtol=0.2)
+    T1b_p, T3_p = pristine.measure_heuristics()[1], pristine.measure_heuristics()[4]
+    T1b_d, T3_d = damaged.measure_heuristics()[1], damaged.measure_heuristics()[4]
+    assert T1b_d == pytest.approx(T1b_p, rel=0.05)
+    assert T3_d == pytest.approx(T3_p, abs=0.05)
+
+
+def test_populated_slice_rejects_bad_axis():
+    """``_populated_slice`` keeps ``half_power_lag``'s axis validation."""
+    field = _smoothed_polar_field()
+    res = StructureFunction2D.from_array(field, dx=0.05, dy=3.0,
+                                         max_lag_x=15, max_lag_y=20)
+    with pytest.raises(ValueError, match="axis must be"):
+        res.half_power_lag(axis="z")
+
+
 def test_measure_heuristics_all_zero_weights_raises():
     """A stack with no measurable correlation scale (constant field -> zero
     S_2 -> NaN half-power lags -> zero reliability weights) must raise a