diff --git a/rm_lite/tools_1d/rmsynth.py b/rm_lite/tools_1d/rmsynth.py index b9929cb..35b24c4 100644 --- a/rm_lite/tools_1d/rmsynth.py +++ b/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() diff --git a/rm_lite/tools_3d/rmsynth.py b/rm_lite/tools_3d/rmsynth.py index 4bb1c4a..418ea90 100644 --- a/rm_lite/tools_3d/rmsynth.py +++ b/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, diff --git a/rm_lite/utils/fitting.py b/rm_lite/utils/fitting.py index 5b9ae0f..6901769 100644 --- a/rm_lite/utils/fitting.py +++ b/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]) diff --git a/rm_lite/utils/synthesis.py b/rm_lite/utils/synthesis.py index 0811c5c..bc813d5 100644 --- a/rm_lite/utils/synthesis.py +++ b/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: diff --git a/tests/test_options_defaults.py b/tests/test_options_defaults.py index 8c0feec..b5a0c59 100644 --- a/tests/test_options_defaults.py +++ b/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", ), diff --git a/tests/test_synthesis_utils.py b/tests/test_synthesis_utils.py index 4f5c553..5a47a51 100644 --- a/tests/test_synthesis_utils.py +++ b/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) diff --git a/tests/test_tools_3d_dask.py b/tests/test_tools_3d_dask.py index 9986984..280da69 100644 --- a/tests/test_tools_3d_dask.py +++ b/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,