WIP: Weighting

.github/workflows/ci.yml

@@ -44,7 +44,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+        python-version: ["3.10", "3.11", "3.12", "3.13"]
         runs-on: [ubuntu-latest]
 
     steps:

.gitignore

@@ -163,4 +163,8 @@ rm_lite/ms_test.ipynb
 rm_lite/multiscale.ipynb
 rm_lite/Untitled-1.ipynb
 2d.ipynb
+depol.ipynb
+nufft_test.ipynb
+robust.ipynb
 *.ipynb
+dragons_120.23_9.59.txt

.pre-commit-config.yaml

@@ -32,13 +32,6 @@ repos:
       - id: rst-directive-colons
       - id: rst-inline-touching-normal
 
-  - repo: https://github.com/pre-commit/mirrors-prettier
-    rev: "v4.0.0-alpha.8"
-    hooks:
-      - id: prettier
-        types_or: [yaml, markdown, html, css, scss, javascript, json]
-        args: [--prose-wrap=always]
-
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: "v0.11.7"
     hooks:

docs/examples/rmclean_2d.ipynb

@@ -96,7 +96,7 @@
     "\n",
     "\n",
     "for time_step, (rm_radm2, frac_pol, psi0_deg) in enumerate(\n",
-    "    zip(rm_time, frac_pol_time, psi0_time)\n",
+    "    zip(rm_time, frac_pol_time, psi0_time, strict=False)\n",
     "):\n",
     "    complex_data_noiseless = faraday_simple_spectrum(\n",
     "        freq_hz,\n",

docs/examples/rmsynth_2d.ipynb

@@ -97,7 +97,7 @@
     "\n",
     "\n",
     "for time_step, (rm_radm2, frac_pol, psi0_deg) in enumerate(\n",
-    "    zip(rm_time, frac_pol_time, psi0_time)\n",
+    "    zip(rm_time, frac_pol_time, psi0_time, strict=False)\n",
     "):\n",
     "    complex_data_noiseless = faraday_simple_spectrum(\n",
     "        freq_hz,\n",

pyproject.toml

@@ -11,7 +11,7 @@ authors = [
 ]
 license.file = "LICENSE"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 classifiers = [
     "Development Status :: 1 - Planning",
     "Intended Audience :: Science/Research",
@@ -21,10 +21,10 @@ classifiers = [
     "Programming Language :: Python",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3 :: Only",
-    "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Topic :: Scientific/Engineering",
     "Typing :: Typed",
 ]
@@ -42,6 +42,7 @@ dependencies = [
     "corner",
     "polars",
     "sigfig",
+    "pandas",
 ]
 
 [project.optional-dependencies]
@@ -118,7 +119,7 @@ isort.required-imports = ["from __future__ import annotations"]
 
 [tool.mypy]
 files = ["rm_tools", "tests"]
-python_version = "3.9"
+python_version = "3.10"
 warn_unused_configs = true
 strict = true
 enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]

rm_lite/tools_1d/rmsynth.py

@@ -12,6 +12,7 @@
 
 from rm_lite.utils.logging import logger
 from rm_lite.utils.synthesis import (
+    WEIGHT_TYPES,
     FDFOptions,
     StokesData,
     compute_rmsynth_params,
@@ -77,11 +78,12 @@ def run_rmsynth(
     phi_max_radm2: float | None = None,
     d_phi_radm2: float | None = None,
     n_samples: float | None = 10.0,
-    weight_type: Literal["variance", "uniform"] = "variance",
+    weight_type: WEIGHT_TYPES = "uniform",
     do_fit_rmsf: bool = False,
     do_fit_rmsf_real: bool = False,
     fit_function: Literal["log", "linear"] = "log",
     fit_order: int = 2,
+    robust: float | None = None,
     ignore_stokes_i: bool = False,
 ) -> RMSynth1DResults:
     """Run RM-synthesis on 1D data
@@ -97,11 +99,13 @@ def run_rmsynth(
         phi_max_radm2 (float | None, optional): Maximum Faraday depth. Defaults to None.
         d_phi_radm2 (float | None, optional): Spacing in Faraday depth. Defaults to None.
         n_samples (float | None, optional): Number of samples across the RMSF. Defaults to 10.0.
-        weight_type ("variance", "uniform", optional): Type of weighting. Defaults to "variance".
+        weight_type ("variance", "uniform", "natural", "briggs"): Weighting type. Defaults to "uniform".
         do_fit_rmsf (bool, optional): Fit the RMSF main lobe. Defaults to False.
         do_fit_rmsf_real (bool, optional): The the real part of the RMSF. Defaults to False.
         fit_function ("log" | "linear", optional): _description_. Defaults to "log".
         fit_order (int, optional): Polynomial fit order. Defaults to 2. Negative values will iterate until the fit is good.
+        robust (float | None, optional): Briggs robustness parameter. Defaults to None.
+        cell_size_m2 (float | None, optional): Cell size in m^2. Defaults to None.
 
     Returns:
         RMSynth1DResults:
@@ -126,6 +130,7 @@ def run_rmsynth(
         weight_type=weight_type,
         do_fit_rmsf=do_fit_rmsf,
         do_fit_rmsf_real=do_fit_rmsf_real,
+        robust=robust,
     )
 
     if (

rm_lite/utils/fitting.py

@@ -289,7 +289,7 @@ def static_fit(
         )
 
     errors = np.sqrt(np.diag(pcov))
-    fit_vals = [sf.round(p, e) for p, e in zip(popt, errors)]
+    fit_vals = [sf.round(p, e) for p, e in zip(popt, errors, strict=False)]
     logger.info(f"Fit results: {fit_vals}")
 
     return FitResult(

rm_lite/utils/multiscale.py

@@ -3,8 +3,9 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from functools import partial
-from typing import Callable, Literal
+from typing import Literal
 
 import numpy as np
 from numpy.typing import NDArray

rm_lite/utils/synthesis.py

@@ -8,6 +8,7 @@
 
 import finufft
 import numpy as np
+import pandas as pd
 import polars as pl
 from astropy.constants import c as speed_of_light
 from astropy.stats import mad_std
@@ -114,6 +115,9 @@
         return TheoreticalNoise(**as_dict)
 
 
+WEIGHT_TYPES = Literal["variance", "uniform", "natural", "briggs"]
+
+
 class FDFOptions(NamedTuple):
     """Options for RM-synthesis"""
 
@@ -123,12 +127,14 @@
     """ Faraday depth resolution """
     n_samples: float | None = 10.0
     """ Number of samples """
-    weight_type: Literal["variance", "uniform"] = "variance"
+    weight_type: WEIGHT_TYPES = "uniform"
     """ Weight type """
     do_fit_rmsf: bool = False
     """ Fit RMSF """
     do_fit_rmsf_real: bool = False
     """ Fit real part of the RMSF """
+    robust: float | None = None
+    """ Briggs robust parameter """
 
 
 def calc_mom2_fdf(
@@ -442,6 +448,94 @@
     )
 
 
+def uniform_weight(
+    freq_hz: NDArray[np.float64],
+) -> NDArray[np.float64]:
+    return np.ones_like(freq_hz) / len(freq_hz)
+
+
+def noise_weight(
+    noise: NDArray[np.float64],
+) -> NDArray[np.float64]:
+    weights = 1 / noise**2
+    weights /= np.nansum(weights)
+    return weights
+
+
+class WaveSqBins(NamedTuple):
+    categories: pd.Categorical
+    """Categories for the bins"""
+    bins: NDArray[np.intp]
+    """The bins in wavelength^2 space"""
+    bin_counts: NDArray[np.intp]
+    """The counts in each bin per channel"""
+
+
+def _get_wave_sq_bins(
+    freq_hz: NDArray[np.float64],
+    phi_max_radm2: float,
+) -> WaveSqBins:
+    lambda_sq_arr_m2 = freq_to_lambda2(freq_hz)
+
+    # Compute the 'cell n_bins' in lambda^2 space
+    rmsf_properties = get_fwhm_rmsf(lambda_sq_arr_m2)
+
+    cell_size_m2 = np.sqrt(3) / phi_max_radm2
+
+    n_bins = int(rmsf_properties.lambda_sq_range_m2 / cell_size_m2)
+
+    categories, bins = pd.cut(
+        lambda_sq_arr_m2,
+        bins=n_bins,
+        include_lowest=True,
+        retbins=True,
+    )
+    bin_idx = categories.codes
+    bin_counts = np.bincount(bin_idx)
+
+    # repeat the bin_counts to match the original array
+    bin_counts = bin_counts[bin_idx]
+    return WaveSqBins(categories, bins, bin_counts)
+
+
+def natural_weight(
+    freq_hz: NDArray[np.float64],
+    phi_max_radm2: float,
+    noise: NDArray[np.float64] | None = None,
+) -> NDArray[np.float64]:
+    wave_sq_bins = _get_wave_sq_bins(
+        freq_hz=freq_hz,
+        phi_max_radm2=phi_max_radm2,
+    )
+
+    weight = wave_sq_bins.bin_counts.astype(np.float64)
+    weight /= np.nansum(weight)
+    if noise is None:
+        return weight
+    var_weight = noise_weight(noise)
+    weight *= var_weight
+    weight /= np.nansum(weight)
+    return weight
+
+
+def briggs_weight(
+    robust: float,
+    freq_hz: NDArray[np.float64],
+    phi_max_radm2: float,
+    noise: NDArray[np.float64] | None = None,
+) -> NDArray[np.float64]:
+    # Briggs robust factor - CASA called it f**2
+    eff_squared = (5 * 10 ** -float(robust)) ** 2
+    nat_weights = natural_weight(
+        freq_hz=freq_hz,
+        noise=noise,
+        phi_max_radm2=phi_max_radm2,
+    )
+    weights = nat_weights / (1 + nat_weights * eff_squared)
+    weights /= np.nansum(weights)
+    return weights
+
+
 def compute_rmsynth_params(
     freq_arr_hz: NDArray[np.float64],
     complex_pol_arr: NDArray[np.complex128],
@@ -493,13 +587,34 @@
         f"phi = {phi_arr_radm2[0]:0.2f} to {phi_arr_radm2[-1]:0.2f} by {d_phi_radm2:0.2f} ({len(phi_arr_radm2)} chans)."
     )
 
-    # Calculate the weighting as 1/sigma^2 or all 1s (uniform)
-    if fdf_options.weight_type == "variance":
-        if (real_qu_error == 0).all():
-            real_qu_error = np.ones(len(real_qu_error))
-        weight_arr = 1.0 / real_qu_error**2
-    else:
-        weight_arr = np.ones_like(freq_arr_hz)
+    if (real_qu_error == 0).all():
+        real_qu_error = np.ones(len(real_qu_error))
+
+    logger.info(f"Computing weights for type {fdf_options.weight_type}...")
+    match fdf_options.weight_type:
+        case "variance":
+            weight_arr = noise_weight(noise=real_qu_error)
+        case "uniform":
+            weight_arr = uniform_weight(freq_hz=freq_arr_hz)
+        case "natural":
+            weight_arr = natural_weight(
+                freq_hz=freq_arr_hz,
+                noise=real_qu_error,
+                phi_max_radm2=phi_max_radm2,
+            )
+        case "briggs":
+            if fdf_options.robust is None:
+                msg = "Briggs weighting requires a robust parameter."
+                raise ValueError(msg)
+            weight_arr = briggs_weight(
+                robust=fdf_options.robust,
+                freq_hz=freq_arr_hz,
+                noise=real_qu_error,
+                phi_max_radm2=phi_max_radm2,
+            )
+        case _:
+            msg = f"Unknown weighting type: {fdf_options.weight_type}"  # type: ignore[unreachable]
+            raise ValueError(msg)
 
     logger.debug(f"Weighting type: {fdf_options.weight_type}")
 

tests/test_nufft.py

@@ -427,6 +427,7 @@ def test_rmsf():
             ["phi2Arr", "RMSFcube", "fwhmRMSFArr"],
             [phi2Arr, RMSFcube, fwhmRMSFArr],
             [phi2Arr_old, RMSFcube_old, fwhmRMSFArr_old],
+            strict=False,
         ):
             msg = f"Testing {name}"
             logger.info(msg)