Fit the RMSF over a sized window cut at the main lobe

rm_lite/tools_1d/rmsynth.py

@@ -155,6 +155,7 @@ def run_rmsynth(
     lam_sq_0_m2: float | LamSq0Mode = "auto",
     do_fit_rmsf: bool = False,
     do_fit_rmsf_real: bool = False,
+    rmsf_fitting_size: float = 1.25,
     fit_function: Literal["log", "linear"] = "log",
     fit_order: int = 2,
     stokes_i_robust_loss: RobustLoss = "cauchy",
@@ -188,6 +189,7 @@ def run_rmsynth(
         robust (float | None, optional): Briggs robust parameter, required for weight_type='briggs'. Defaults to None.
         do_fit_rmsf (bool, optional): Fit the RMSF main lobe. Defaults to False.
         do_fit_rmsf_real (bool, optional): Fit only the real part of the RMSF. Defaults to False.
+        rmsf_fitting_size (float, optional): RMSF fit window in analytic FWHMs, cut at the main lobe's first minimum. Defaults to 1.25.
         fit_function ("log" | "linear", optional): RMSF fit function. Defaults to "log".
         fit_order (int, optional): Polynomial fit order. Defaults to 2. Negative values will iterate until the fit is good.
         stokes_i_robust_loss (RobustLoss, optional): Downweight channels far from
@@ -235,6 +237,7 @@ def run_rmsynth(
         lam_sq_0_m2=lam_sq_0_m2,
         do_fit_rmsf=do_fit_rmsf,
         do_fit_rmsf_real=do_fit_rmsf_real,
+        rmsf_fitting_size=rmsf_fitting_size,
         stokes_i_weighting=stokes_i_weighting,
         stokes_i_weight_alpha=stokes_i_weight_alpha,
     )
@@ -379,6 +382,7 @@ def _run_rmsynth(
         mask_arr=~no_nan_idx,
         do_fit_rmsf=fdf_options.do_fit_rmsf,
         do_fit_rmsf_real=fdf_options.do_fit_rmsf_real,
+        rmsf_fitting_size=fdf_options.rmsf_fitting_size,
     )
 
     tock = time.time()

rm_lite/tools_3d/rmsynth.py

@@ -55,7 +55,6 @@
     compute_theoretical_noise,
     derotate_to,
     error_from_weight,
-    get_fwhm_rmsf,
     get_rmsf_nufft,
     lam_sq_0_per_pixel,
     lambda2_to_freq,
@@ -81,7 +80,8 @@ class RMSynth3DResults(NamedTuple):
     phi_double_arr_radm2: NDArray[np.float64]
     """Double-length Faraday depth values in rad/m^2, for the RMSF."""
     fwhm_rmsf_radm2: float
-    """Analytic RMSF FWHM (per-pixel fitting is not performed in 3D)."""
+    """RMSF FWHM, one value for the cube: analytic, or fitted to `rmsf_arr` with
+    `do_fit_rmsf`."""
     lambda_sq_arr_m2: NDArray[np.float64]
     """Channel wavelength^2 values in m^2."""
     lam_sq_0_m2: float
@@ -309,11 +309,12 @@ def _summarise_weight(
 
 def _shared_rmsf(
     rmsynth_params: RMSynthParams,
+    fdf_options: FDFOptions,
     nthreads: int,
     log_level: int,
     dtype: np.dtype[np.complexfloating],
-) -> NDArray[np.complexfloating]:
-    """The single RMSF the whole cube shares, from the per-channel weights.
+) -> tuple[NDArray[np.complexfloating], float]:
+    """The single RMSF the whole cube shares, from the per-channel weights, and its FWHM.
 
     Every pixel whose flagged channels are the cube's flagged channels has this
     RMSF, and for the noise-based `weight_type`s a channel blank across the cube
@@ -328,12 +329,16 @@ def _shared_rmsf(
                 real_dtype(dtype), copy=False
             ),
             lam_sq_0_m2=rmsynth_params.lam_sq_0_m2,
-            do_fit_rmsf=False,
+            do_fit_rmsf=fdf_options.do_fit_rmsf,
+            rmsf_fitting_size=fdf_options.rmsf_fitting_size,
             nthreads=nthreads,
         )
     # RMSFResults.rmsf_cube is annotated NDArray[np.float64] but is complex at
     # runtime (built from a finufft complex output).
-    return np.asarray(rmsf_result.rmsf_cube, dtype=dtype)
+    return (
+        np.asarray(rmsf_result.rmsf_cube, dtype=dtype),
+        float(rmsf_result.fwhm_rmsf_arr),
+    )
 
 
 def target_chunk_mb_for_worker(
@@ -719,6 +724,8 @@ def rmsynth_3d(
     compute_model_error: bool = False,
     n_error_samples: int = 1000,
     per_pixel_rmsf: bool = False,
+    do_fit_rmsf: bool = False,
+    rmsf_fitting_size: float = 1.25,
     nufft_nthreads: int = 1,
     target_chunk_mb: float = DEFAULT_TARGET_CHUNK_MB,
     log_level: int = logging.WARNING,
@@ -802,6 +809,12 @@ def rmsynth_3d(
             per-channel `weight_arr` does not already say so; otherwise every
             pixel of it holds `rmsf_arr` at `2 * n_phi_double / n_phi` times the
             cost of the FDF cube. Defaults to False.
+        do_fit_rmsf (bool, optional): Take `fwhm_rmsf_radm2` from a Gaussian fit
+            to the main lobe of `rmsf_arr` rather than the analytic
+            3.8 / (lambda^2 range). One fit for the cube, per-pixel RMSFs or not.
+            Defaults to False.
+        rmsf_fitting_size (float, optional): Fit window in analytic FWHMs, cut
+            at the main lobe's first minimum. Defaults to 1.25.
         nufft_nthreads (int, optional): finufft OpenMP threads per chunk. Defaults
             to 1 so dask parallelises across chunks without oversubscribing finufft's
             own threads (the fast config on many chunks). Set to 0 (finufft default,
@@ -834,6 +847,8 @@ def rmsynth_3d(
         weight_type=weight_type,
         robust=robust,
         lam_sq_0_m2=lam_sq_0_m2,
+        do_fit_rmsf=do_fit_rmsf,
+        rmsf_fitting_size=rmsf_fitting_size,
         stokes_i_weighting=stokes_i_weighting,
         stokes_i_weight_alpha=stokes_i_weight_alpha,
     )
@@ -911,9 +926,10 @@ def rmsynth_3d(
     n_phi = rmsynth_params.phi_arr_radm2.shape[0]
     phi_double_arr_radm2 = make_double_phi_arr(rmsynth_params.phi_arr_radm2)
     n_phi_double = phi_double_arr_radm2.shape[0]
-    fwhm_rmsf_radm2 = get_fwhm_rmsf(rmsynth_params.lambda_sq_arr_m2).fwhm_rmsf_radm2
     fdf_dtype = complex_dtype(np.result_type(stokes_q.dtype, stokes_u.dtype))
-    rmsf_arr = _shared_rmsf(rmsynth_params, nufft_nthreads, log_level, fdf_dtype)
+    rmsf_arr, fwhm_rmsf_radm2 = _shared_rmsf(
+        rmsynth_params, fdf_options, nufft_nthreads, log_level, fdf_dtype
+    )
 
     stokes_q, stokes_u = _match_chunks_to_fdf(
         stokes_q, stokes_u, n_phi_double, fdf_dtype, target_chunk_mb
@@ -1265,6 +1281,8 @@ def rmsynth_3d_from_fits(
     compute_model_error: bool = False,
     n_error_samples: int = 1000,
     per_pixel_rmsf: bool = False,
+    do_fit_rmsf: bool = False,
+    rmsf_fitting_size: float = 1.25,
     nufft_nthreads: int = 1,
     target_chunk_mb: float = DEFAULT_TARGET_CHUNK_MB,
     convert_to_zarr: bool = False,
@@ -1317,6 +1335,8 @@ def rmsynth_3d_from_fits(
         compute_model_error (bool, optional): See `rmsynth_3d`. Defaults to False.
         n_error_samples (int, optional): See `rmsynth_3d`. Defaults to 1000.
         per_pixel_rmsf (bool, optional): See `rmsynth_3d`. Defaults to False.
+        do_fit_rmsf (bool, optional): See `rmsynth_3d`. Defaults to False.
+        rmsf_fitting_size (float, optional): See `rmsynth_3d`. Defaults to 1.25.
         nufft_nthreads (int, optional): See `rmsynth_3d`. Defaults to 1.
         target_chunk_mb (float, optional): Target per-chunk memory footprint
             in MB, see `read_cube_dask`. Defaults to 256.
@@ -1463,6 +1483,8 @@ def rmsynth_3d_from_fits(
         compute_model_error=compute_model_error,
         n_error_samples=n_error_samples,
         per_pixel_rmsf=per_pixel_rmsf,
+        do_fit_rmsf=do_fit_rmsf,
+        rmsf_fitting_size=rmsf_fitting_size,
         nufft_nthreads=nufft_nthreads,
         target_chunk_mb=target_chunk_mb,
         log_level=log_level,

rm_lite/utils/fitting.py

@@ -174,25 +174,32 @@ def fit_rmsf(
     rmsf_to_fit_arr: NDArray[np.float64],
     phi_double_arr_radm2: NDArray[np.float64],
     fwhm_rmsf_radm2: float,
+    fitting_size: float = 1.25,
 ) -> float:
-    rmsf_to_fit_arr = rmsf_to_fit_arr.copy()
-    rmsf_to_fit_arr /= np.nanmax(rmsf_to_fit_arr)
+    """FWHM of a Gaussian fitted to the RMSF main lobe.
+
+    The fit spans `fitting_size` times `fwhm_rmsf_radm2`, cut short at the
+    lobe's first minimum on each side so no sidelobe gets in (cf. WSClean's
+    `-beam-fitting-size`).
+    """
+    rmsf_to_fit_arr = rmsf_to_fit_arr / np.nanmax(rmsf_to_fit_arr)
     d_phi = phi_double_arr_radm2[1] - phi_double_arr_radm2[0]
-    mask = np.zeros_like(phi_double_arr_radm2, dtype=bool)
-    mask[np.argmax(rmsf_to_fit_arr)] = True
-    sigma_rmsf_radm2 = fwhm_to_sigma(fwhm_rmsf_radm2)
-    sigma_rmsf_arr_pix = sigma_rmsf_radm2 / d_phi
-    for i in np.where(mask)[0]:
-        # Clamped: a negative start wraps and empties the slice, leaving one
-        # point and a fit with no degrees of freedom.
-        start = max(0, int(i - sigma_rmsf_arr_pix / 2))
-        end = int(i + sigma_rmsf_arr_pix / 2)
-        mask[start : end + 2] = True
+    peak = int(np.argmax(rmsf_to_fit_arr))
+    half_width = max(1, round(fitting_size * fwhm_rmsf_radm2 / 2 / d_phi))
+    # Samples from the peak outwards until the profile stops falling.
+    right = rmsf_to_fit_arr[peak:]
+    left = rmsf_to_fit_arr[peak::-1]
+    to_right_min = int(np.argmax(np.append(np.diff(right) >= 0, True)))
+    to_left_min = int(np.argmax(np.append(np.diff(left) >= 0, True)))
+    # Stop one short of each minimum, but keep a point either side of the peak.
+    start = peak - max(1, min(half_width, to_left_min - 1))
+    end = peak + max(1, min(half_width, to_right_min - 1))
+    start, end = max(0, start), min(len(rmsf_to_fit_arr) - 1, end)
     popt, _ = optimize.curve_fit(
         unit_centred_gaussian,
-        phi_double_arr_radm2[mask],
-        rmsf_to_fit_arr[mask],
-        p0=[sigma_rmsf_radm2],
+        phi_double_arr_radm2[start : end + 1],
+        rmsf_to_fit_arr[start : end + 1],
+        p0=[fwhm_to_sigma(fwhm_rmsf_radm2)],
         bounds=([0], [np.inf]),
     )
     return sigma_to_fwhm(popt[0])

rm_lite/utils/synthesis.py

@@ -173,6 +173,8 @@ class FDFOptions:
     """ Fit RMSF """
     do_fit_rmsf_real: bool = False
     """ Fit real part of the RMSF """
+    rmsf_fitting_size: float = 1.25
+    """ Width of the RMSF fit window in analytic FWHMs, cut at the main lobe """
     lam_sq_0_m2: float | LamSq0Mode = "auto"
     """ Reference lambda^2 in m^2, or "auto"/"per_pixel" to derive one. The
     Stokes I reference frequency is derived from it, so the phase and flux
@@ -231,7 +233,7 @@ def __post_init__(self) -> None:
         if self.d_phi_radm2 is None and self.n_samples is None:
             msg = "Either d_phi_radm2 or n_samples must be provided."
             raise ValueError(msg)
-        for name in ("phi_max_radm2", "d_phi_radm2", "n_samples"):
+        for name in ("phi_max_radm2", "d_phi_radm2", "n_samples", "rmsf_fitting_size"):
             value = getattr(self, name)
             if value is not None and value <= 0:
                 msg = f"{name} must be positive, got {value}."
@@ -2090,6 +2092,7 @@ def get_rmsf_nufft(
     mask_arr: NDArray[np.bool_] | None = None,
     do_fit_rmsf: bool = False,
     do_fit_rmsf_real: bool = False,
+    rmsf_fitting_size: float = 1.25,
     eps: float = 1e-6,
     nthreads: int = 0,
     reuse_rmsf: bool = True,
@@ -2108,6 +2111,8 @@ def get_rmsf_nufft(
         mask_arr (Optional[NDArray[np.float64]], optional): Mask array. Defaults to None.
         do_fit_rmsf (bool, optional): Fit the RMSF with a Gaussian. Defaults to False.
         do_fit_rmsf_real (bool, optional): Fit the *real* part of the. Defaults to False.
+        rmsf_fitting_size (float, optional): Fit window in analytic FWHMs, cut at
+            the main lobe's first minimum. Defaults to 1.25.
         eps (float, optional): NUFFT tolerance. Defaults to 1e-6.
         nthreads (int, optional): finufft OpenMP threads. 0 uses finufft's default
             (all cores). Set to 1 when parallelising across chunks with dask, to
@@ -2282,6 +2287,7 @@ def get_rmsf_nufft(
                     ),
                     phi_double_arr_radm2=phi_double_arr_radm2,
                     fwhm_rmsf_radm2=fwhm_rmsf_radm2,
+                    fitting_size=rmsf_fitting_size,
                 )
                 fit_status = True
             except Exception as e:

tests/test_options_defaults.py

@@ -44,6 +44,7 @@ def param_defaults(func: Any) -> dict[str, Any]:
                 "weight_type": "weight_type",
                 "do_fit_rmsf": "do_fit_rmsf",
                 "do_fit_rmsf_real": "do_fit_rmsf_real",
+                "rmsf_fitting_size": "rmsf_fitting_size",
             },
             id="run_rmsynth-FDFOptions",
         ),
@@ -61,6 +62,8 @@ def param_defaults(func: Any) -> dict[str, Any]:
                 "d_phi_radm2": "d_phi_radm2",
                 "n_samples": "n_samples",
                 "weight_type": "weight_type",
+                "do_fit_rmsf": "do_fit_rmsf",
+                "rmsf_fitting_size": "rmsf_fitting_size",
             },
             id="rmsynth_3d-FDFOptions",
         ),

tests/test_synthesis_utils.py

@@ -10,6 +10,7 @@
 from rm_lite.utils.clean import restore_model
 from rm_lite.utils.fitting import (
     StokesIFitOptions,
+    fit_rmsf,
     fit_sampled_peak,
     gaussian,
     gaussian_integrand,
@@ -1098,3 +1099,54 @@ def test_moments_reject_unknown_units():
     fdf = gaussian(phi_arr, 1.0, 0.0, fwhm=20.0).astype(np.complex128)
     with pytest.raises(ValueError, match="fdf_units"):
         calc_faraday_moments(fdf, phi_arr, 20.0, fdf_units="jy/beam")  # type: ignore[arg-type]
+
+
+def _abs_rmsf(
+    lambda_sq_arr_m2: NDArray[np.float64],
+    weight_arr: NDArray[np.float64],
+    fwhm_radm2: float,
+    n_samples: float,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """|RMSF| on a double phi grid, summed directly."""
+    weight_arr = weight_arr / weight_arr.sum()
+    lam_sq_0_m2 = np.sum(weight_arr * lambda_sq_arr_m2)
+    phi_double = make_double_phi_arr(
+        make_phi_arr(20 * fwhm_radm2, fwhm_radm2 / n_samples)
+    )
+    kernel = np.exp(-2j * np.outer(phi_double, lambda_sq_arr_m2 - lam_sq_0_m2))
+    return phi_double, np.abs(kernel @ weight_arr)
+
+
+def test_fit_rmsf_is_exact_on_a_gaussian_rmsf():
+    # Gaussian weights in lambda^2 give exp(-2 phi^2 s^2): sigma_phi = 1 / (2 s).
+    lambda_sq = np.linspace(0.01, 0.2, 4000)
+    s = 0.02
+    weights = np.exp(-0.5 * ((lambda_sq - 0.1) / s) ** 2)
+    expected = 2 * np.sqrt(2 * np.log(2)) / (2 * s)
+    phi_double, rmsf = _abs_rmsf(lambda_sq, weights, expected, 10)
+    assert np.isclose(fit_rmsf(rmsf, phi_double, expected), expected, rtol=1e-3)
+
+
+@pytest.mark.parametrize("n_samples", [10, 100])
+def test_fit_rmsf_matches_the_boxcar_half_max(n_samples: int):
+    # A boxcar in lambda^2 gives |sinc|, which is at half max 3.791 / range apart.
+    lambda_sq = np.linspace(0.01, 0.2, 4000)
+    expected = 3.791 / np.ptp(lambda_sq)
+    phi_double, rmsf = _abs_rmsf(
+        lambda_sq, np.ones_like(lambda_sq), expected, n_samples
+    )
+    assert np.isclose(fit_rmsf(rmsf, phi_double, expected), expected, rtol=0.015)
+
+
+def test_fit_rmsf_window_stops_at_the_main_lobe():
+    # Past the first minimum the window cannot grow, so no sidelobe gets in.
+    lambda_sq = np.linspace(0.01, 0.2, 4000)
+    fwhm = 3.8 / np.ptp(lambda_sq)
+    phi_double, rmsf = _abs_rmsf(lambda_sq, np.ones_like(lambda_sq), fwhm, 10)
+    at_lobe = fit_rmsf(rmsf, phi_double, fwhm, fitting_size=2.0)
+    assert fit_rmsf(rmsf, phi_double, fwhm, fitting_size=50.0) == at_lobe
+
+
+def test_rmsf_fitting_size_must_be_positive():
+    with pytest.raises(ValueError, match="rmsf_fitting_size"):
+        FDFOptions(rmsf_fitting_size=0)

tests/test_tools_3d_dask.py

@@ -1819,6 +1819,35 @@ def test_shared_rmsf_is_what_the_per_pixel_cube_holds(
             np.testing.assert_allclose(cube[:, j, i], default.rmsf_arr, atol=1e-12)
 
 
+def test_rmsynth_3d_fits_the_shared_rmsf(
+    synthetic_cube: SyntheticCube, chunked: Callable[..., da.Array]
+):
+    """`do_fit_rmsf` swaps the analytic FWHM for a fit to `rmsf_arr`."""
+    q_dask = chunked(synthetic_cube.stokes_q, 3, 4)
+    u_dask = chunked(synthetic_cube.stokes_u, 3, 4)
+    analytic = rmsynth_3d(
+        q_dask, u_dask, synthetic_cube.freq_arr_hz, d_phi_radm2=D_PHI_RADM2
+    )
+    fitted = rmsynth_3d(
+        q_dask,
+        u_dask,
+        synthetic_cube.freq_arr_hz,
+        d_phi_radm2=D_PHI_RADM2,
+        do_fit_rmsf=True,
+        rmsf_fitting_size=1.0,
+    )
+    lambda_sq_range = np.ptp(analytic.lambda_sq_arr_m2)
+    assert analytic.fwhm_rmsf_radm2 == 3.8 / lambda_sq_range
+    expected = fitting_mod.fit_rmsf(
+        np.abs(fitted.rmsf_arr),
+        fitted.phi_double_arr_radm2,
+        analytic.fwhm_rmsf_radm2,
+        fitting_size=1.0,
+    )
+    assert np.isclose(fitted.fwhm_rmsf_radm2, expected)
+    assert fitted.fwhm_rmsf_radm2 != analytic.fwhm_rmsf_radm2
+
+
 @pytest.mark.filterwarnings("ignore: All channels masked")
 def test_rmclean_agrees_between_shared_and_per_pixel_rmsf(
     synthetic_cube: SyntheticCube,