1D RM-CLEAN

[1]:
from __future__ import annotations

import astropy.units as u
import matplotlib.pyplot as plt
import numpy as np
from astropy.visualization import quantity_support
from numpy.typing import NDArray
from rm_lite.tools_1d import rmclean, rmsynth

plt.rcParams["figure.dpi"] = 150

_ = quantity_support()
rng = np.random.default_rng(42)

First, we’ll start with the same data s the 1D RM-synth example. Keeping it simple, we’ll use a Faraday simple source with a bit of noise.

I’ll skip the details here, see the RM-synth page for more info.

[2]:
from rm_lite.utils.fitting import power_law
from rm_lite.utils.synthesis import faraday_simple_spectrum, freq_to_lambda2


def faraday_slab_spectrum(
    freq_arr_hz: NDArray[np.float64],
    frac_pol: float,
    psi0_deg: float,
    rm_radm2: float,
    delta_rm_radm2: float,
) -> NDArray[np.complex128]:
    lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz)

    return (
        frac_pol
        * np.exp(2j * (np.deg2rad(psi0_deg) + rm_radm2 * lambda_sq_arr_m2))
        * (
            np.sin(delta_rm_radm2 * lambda_sq_arr_m2)
            / (delta_rm_radm2 * lambda_sq_arr_m2)
        )
    )


def faraday_gaussian_spectrum(
    freq_arr_hz: NDArray[np.float64],
    frac_pol: float,
    psi0_deg: float,
    rm_radm2: float,
    sigma_rm_radm2: float,
):
    lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz)
    rm_term = np.exp(2j * (np.deg2rad(psi0_deg) + rm_radm2 * lambda_sq_arr_m2))
    depol_term = np.exp(-2.0 * sigma_rm_radm2**2 * lambda_sq_arr_m2**2)
    return frac_pol * rm_term * depol_term


bw_low = 288
bw_mid = 144
bw_high = 288
low = np.linspace(943.5 - bw_low / 2, 943.5 + bw_low / 2, 36) * u.MHz
mid = np.linspace(1367.5 - bw_mid / 2, 1367.5 + bw_mid / 2, 9) * u.MHz
high = np.linspace(1655.5 - bw_high / 2, 1655.5 + bw_high / 2, 9) * u.MHz
freqs = np.concatenate([low, mid, high])
freq_hz = freqs.to(u.Hz).value

delta_rm_radm2 = 30
rm_radm2 = 100
frac_pol = 0.7
psi0_deg = 10

complex_data_noiseless = faraday_simple_spectrum(
    freq_arr_hz=freq_hz,
    frac_pol=frac_pol,
    psi0_deg=psi0_deg,
    rm_radm2=rm_radm2,
)
stokes_i_flux = 2.0
spectral_index = -0.7
rms_noise = 0.1

stokes_i_model = power_law(order=1)
stokes_i_noiseless = stokes_i_model(
    freq_hz / (np.mean(freq_hz)), stokes_i_flux, spectral_index
)
stokes_i_noise = rng.normal(0, rms_noise, size=freq_hz.size)
stokes_i_noisy = stokes_i_noiseless + stokes_i_noise


stokes_q_noise = rng.normal(0, rms_noise, size=freq_hz.size)
stokes_u_noise = rng.normal(0, rms_noise, size=freq_hz.size)
complex_noise = stokes_q_noise + 1j * stokes_u_noise

complex_flux = complex_data_noiseless * stokes_i_noiseless
complex_data_noisy = complex_data_noiseless + complex_noise
rm_syth_results = rmsynth.run_rmsynth(
    freq_arr_hz=freq_hz,
    complex_pol_arr=complex_data_noisy,
    complex_pol_error=np.ones_like(complex_data_noiseless) * rms_noise,
    do_fit_rmsf=True,
    n_samples=100,
)
WARNING rmsynth.run_rmsynth: Stokes I array/errors or model not provided. No fractional polarization will be calculated.
INFO synthesis.rmsynth_nufft: Running RM-synthesis using the NUFFTs over 2001 Faraday depth channels.
INFO synthesis.rmsynth_nufft: NUFFT complete in 0.00853 seconds.
INFO synthesis.get_rmsf_nufft: Fitting main lobe in each RMSF spectrum.
INFO rmsynth._run_rmsynth: RM-synthesis completed in 21.10ms.

Let’s remind ourselves of the what the dirty spectrum looks like

[3]:
fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rm_syth_results

phi_arr_radm2 = fdf_arrs["phi_arr_radm2"].to_numpy()
fdf_dirty_arr = fdf_arrs["fdf_dirty_complex_arr"].to_numpy().astype(complex)

fig, ax = plt.subplots()

ax.plot(
    phi_arr_radm2,
    fdf_dirty_arr.real,
    color="tab:red",
    label="Stokes Q",
)
ax.plot(
    phi_arr_radm2,
    fdf_dirty_arr.imag,
    color="tab:blue",
    label="Stokes U",
)
ax.plot(
    phi_arr_radm2,
    np.abs(fdf_dirty_arr),
    color="k",
    label="Polarized intensity",
)

ax.errorbar(
    fdf_parameters["peak_rm_fit"],
    fdf_parameters["peak_pi_fit"],
    xerr=fdf_parameters["peak_rm_fit_error"],
    yerr=fdf_parameters["peak_pi_error"],
    fmt="o",
    lw=1,
    color="red",
    mfc="none",
    label="Fitted peak",
)

ax.set(
    xlabel=rf"$\phi$ / {u.rad / u.m**2:latex_inline}",
    ylabel="Flux density",
    title="Dirty FDF",
)
ax.legend()
[3]:
<matplotlib.legend.Legend at 0x791536623ce0>
../_images/examples_rmclean_1d_5_1.png

We can use the convenience function rmclean.run_rmclean_from_synth to directly take the outputs of 1D RM-synthesis for RM-CLEAN.

Inspired by WSClean, rm_lite implements a deep cleaning mode with auto-masking and auto-thresholding based on the noise profile.

[4]:
help(rmclean.run_rmclean_from_synth)
Help on function run_rmclean_from_synth in module rm_lite.tools_1d.rmclean:

run_rmclean_from_synth(rm_synth_1d_results: 'RMSynth1DResults', auto_mask: 'float' = 7, auto_threshold: 'float' = 1, max_iter: 'int' = 10000, gain: 'float' = 0.1, mask_arr: 'NDArray[np.bool_] | None' = None) -> 'RMClean1DResults'
    Run RM-CLEAN on the results of RM-synth.

    Args:
        rm_synth_1d_results (RMSynth1DResults): Results from RM-synth.
        auto_mask (float, optional): Masking threshold in SNR. Defaults to 7.
        auto_threshold (float, optional): Cleaning threshold in SNR. Defaults to 1.
        max_iter (int, optional): Maximum CLEAN iterations. Defaults to 10_000.
        gain (float, optional): CLEAN gain. Defaults to 0.1.
        mask_arr (NDArray[np.bool_] | None, optional): Optional mask array. Defaults to None.

    Returns:
        RMClean1DResults: RM-CLEAN results: `fdf_parameters`, `fdf_arrs`, `clean_parameters`.

[5]:
rmclean_results = rmclean.run_rmclean_from_synth(
    rm_synth_1d_results=rm_syth_results, auto_mask=10, auto_threshold=0.5
)
INFO rmclean.run_rmclean_from_synth: Theoretical noise: TheoreticalNoise(fdf_error_noise=0.006804138174397718, fdf_q_noise=0.013608276348795436, fdf_u_noise=0.0)
INFO rmclean.run_rmclean_from_synth: Auto mask: 0.07, Auto threshold: 0.00, Max iterations: 10000, Gain: 0.1
INFO clean.minor_cycle: Starting initial minor loop...
INFO clean.minor_loop: Starting minor loop... 1135 pixels in the mask
WARNING clean.minor_loop: All channels masked. Exiting loop...performed 23 iterations
INFO clean.minor_cycle: Initial loop complete. Starting deep clean...
INFO clean.minor_loop: Starting minor loop... 1 pixels in the mask
INFO clean.minor_loop: Threshold reached. Exiting loop...performed 51 iterations
[6]:
rmclean_results
[6]:
RMClean1DResults(fdf_parameters=shape: (1, 30)
┌───────────┬───────────┬───────────┬───────────┬───┬───────────┬───────────┬───────────┬──────────┐
│ fdf_error ┆ peak_pi_f ┆ peak_pi_e ┆ peak_pi_f ┆ … ┆ frac_pol_ ┆ sigma_add ┆ sigma_add ┆ sigma_ad │
│ _mad      ┆ it        ┆ rror      ┆ it_debias ┆   ┆ error     ┆ ---       ┆ _minus    ┆ d_plus   │
│ ---       ┆ ---       ┆ ---       ┆ ---       ┆   ┆ ---       ┆ f64       ┆ ---       ┆ ---      │
│ f64       ┆ f64       ┆ f64       ┆ f64       ┆   ┆ f64       ┆           ┆ f64       ┆ f64      │
╞═══════════╪═══════════╪═══════════╪═══════════╪═══╪═══════════╪═══════════╪═══════════╪══════════╡
│ NaN       ┆ 0.707165  ┆ 0.006804  ┆ 0.70709   ┆ … ┆ NaN       ┆ NaN       ┆ NaN       ┆ NaN      │
└───────────┴───────────┴───────────┴───────────┴───┴───────────┴───────────┴───────────┴──────────┘, fdf_arrs=shape: (2_001, 5)
┌───────────────┬────────────────────┬────────────────────┬────────────────────┬───────────────────┐
│ phi_arr_radm2 ┆ fdf_dirty_complex_ ┆ fdf_clean_complex_ ┆ fdf_model_complex_ ┆ fdf_residual_comp │
│ ---           ┆ arr                ┆ arr                ┆ arr                ┆ lex_arr           │
│ f64           ┆ ---                ┆ ---                ┆ ---                ┆ ---               │
│               ┆ object             ┆ object             ┆ object             ┆ object            │
╞═══════════════╪════════════════════╪════════════════════╪════════════════════╪═══════════════════╡
│ -336.725921   ┆ (-0.02201736851616 ┆ (-0.00396349771314 ┆ 0j                 ┆ (-0.0039634977131 │
│               ┆ 564+0.001648…      ┆ 7674-0.00653…      ┆                    ┆ 47674-0.00653…    │
│ -336.389195   ┆ (-0.02299530749011 ┆ (-0.00402686834948 ┆ 0j                 ┆ (-0.0040268683494 │
│               ┆ 624+0.001460…      ┆ 3787-0.00620…      ┆                    ┆ 83787-0.00620…    │
│ -336.052469   ┆ (-0.02395130376532 ┆ (-0.00408745753013 ┆ 0j                 ┆ (-0.0040874575301 │
│               ┆ 1848+0.00127…      ┆ 4169-0.00584…      ┆                    ┆ 34169-0.00584…    │
│ -335.715743   ┆ (-0.02488403640965 ┆ (-0.00414520259663 ┆ 0j                 ┆ (-0.0041452025966 │
│               ┆ 5887+0.00108…      ┆ 6997-0.00548…      ┆                    ┆ 36997-0.00548…    │
│ -335.379017   ┆ (-0.02579219967194 ┆ (-0.00420004152863 ┆ 0j                 ┆ (-0.0042000415286 │
│               ┆ 7413+0.00089…      ┆ 2685-0.00510…      ┆                    ┆ 32685-0.00510…    │
│ …             ┆ …                  ┆ …                  ┆ …                  ┆ …                 │
│ 335.379017    ┆ (-0.01154286845827 ┆ (-0.01048137848740 ┆ 0j                 ┆ (-0.0104813784874 │
│               ┆ 928+0.035471…      ┆ 5562+0.02410…      ┆                    ┆ 05562+0.02410…    │
│ 335.715743    ┆ (-0.01217000631392 ┆ (-0.01102325518601 ┆ 0j                 ┆ (-0.0110232551860 │
│               ┆ 1827+0.03526…      ┆ 0507+0.02413…      ┆                    ┆ 10507+0.02413…    │
│ 336.052469    ┆ (-0.01275396594827 ┆ (-0.01156730407346 ┆ 0j                 ┆ (-0.0115673040734 │
│               ┆ 4977+0.03504…      ┆ 0578+0.02415…      ┆                    ┆ 60578+0.02415…    │
│ 336.389195    ┆ (-0.01329477347709 ┆ (-0.01211291641524 ┆ 0j                 ┆ (-0.0121129164152 │
│               ┆ 4648+0.03483…      ┆ 6514+0.02417…      ┆                    ┆ 46514+0.02417…    │
│ 336.725921    ┆ (-0.01379251637084 ┆ (-0.01265950183469 ┆ 0j                 ┆ (-0.0126595018346 │
│               ┆ 0727+0.03461…      ┆ 4034+0.02418…      ┆                    ┆ 94034+0.02418…    │
└───────────────┴────────────────────┴────────────────────┴────────────────────┴───────────────────┘, clean_parameters=shape: (1, 3)
┌──────────┬───────────┬────────┐
│ mask     ┆ threshold ┆ n_iter │
│ ---      ┆ ---       ┆ ---    │
│ f64      ┆ f64       ┆ i64    │
╞══════════╪═══════════╪════════╡
│ 0.068041 ┆ 0.003402  ┆ 51     │
└──────────┴───────────┴────────┘)
[7]:
# Extract the DataFrames
clean_arrs = rmclean_results.fdf_arrs
clean_fdf_params = rmclean_results.fdf_parameters
clean_params = rmclean_results.clean_parameters
[8]:
# Extract some arrays
phi_arr_radm2 = clean_arrs["phi_arr_radm2"].to_numpy()
fdf_dirty_arr = clean_arrs["fdf_dirty_complex_arr"].to_numpy().astype(complex)
fdf_clean_arr = clean_arrs["fdf_clean_complex_arr"].to_numpy().astype(complex)
fdf_model_arr = clean_arrs["fdf_model_complex_arr"].to_numpy().astype(complex)
fdf_residual_arr = clean_arrs["fdf_residual_complex_arr"].to_numpy().astype(complex)

Let’s look at how the CLEAN went. Note that the deep clean allows us to model the entire simple component.

[9]:
fig, ax = plt.subplots()

ax.plot(
    phi_arr_radm2,
    np.abs(fdf_dirty_arr),
    color="k",
    label="Dirty FDF",
    alpha=0.5,
)
ax.plot(
    phi_arr_radm2,
    np.abs(fdf_clean_arr),
    color="k",
    label="Clean FDF",
)
ax.step(
    phi_arr_radm2,
    np.abs(fdf_model_arr),
    color="tab:red",
    label="Model FDF",
    where="mid",
)
ax.plot(
    phi_arr_radm2,
    np.abs(fdf_residual_arr),
    color="tab:blue",
    label="Residual FDF",
    alpha=0.5,
)

ax.errorbar(
    clean_fdf_params["peak_rm_fit"],
    clean_fdf_params["peak_pi_fit"],
    xerr=clean_fdf_params["peak_rm_fit_error"],
    yerr=clean_fdf_params["peak_pi_error"],
    fmt="o",
    lw=1,
    color="red",
    mfc="none",
    label="Fitted peak",
)

ax.axhline(
    y=clean_params["threshold"][0],
    color="tab:blue",
    label="Threshold",
    linestyle="--",
)
ax.axhline(
    y=clean_params["mask"][0],
    color="tab:orange",
    label="Mask",
    linestyle="--",
)

ax.set(
    xlabel=rf"$\phi$ / {u.rad / u.m**2:latex_inline}",
    ylabel="Flux density",
    title="Dirty FDF",
)
ax.legend()
[9]:
<matplotlib.legend.Legend at 0x791536544fb0>
../_images/examples_rmclean_1d_13_1.png

Now let’s take a look at the more complex example.

[10]:
delta_rm_radm2 = 30
rm_radm2 = 100
frac_pol = 0.5
psi0_deg = 10
complex_data_noiseless = faraday_slab_spectrum(
    freq_arr_hz=freq_hz,
    frac_pol=frac_pol,
    psi0_deg=psi0_deg,
    rm_radm2=rm_radm2,
    delta_rm_radm2=delta_rm_radm2,
)


stokes_i_flux = 1.0
spectral_index = -0.7
rms_noise = 0.1


stokes_i_model = power_law(order=1)
stokes_i_noiseless = stokes_i_model(
    freq_hz / (np.mean(freq_hz)), stokes_i_flux, spectral_index
)
stokes_i_noise = rng.normal(0, rms_noise, size=freq_hz.size)
stokes_i_noisy = stokes_i_noiseless + stokes_i_noise


stokes_q_noise = rng.normal(0, rms_noise, size=freq_hz.size)
stokes_u_noise = rng.normal(0, rms_noise, size=freq_hz.size)
complex_noise = stokes_q_noise + 1j * stokes_u_noise

complex_flux = complex_data_noiseless * stokes_i_noiseless
complex_data_noisy = complex_data_noiseless + complex_noise
[11]:
rm_syth_results = rmsynth.run_rmsynth(
    freq_arr_hz=freq_hz,
    complex_pol_arr=complex_data_noisy,
    complex_pol_error=np.ones_like(complex_data_noiseless) * rms_noise,
    do_fit_rmsf=True,
    n_samples=100,
)
WARNING rmsynth.run_rmsynth: Stokes I array/errors or model not provided. No fractional polarization will be calculated.
INFO synthesis.rmsynth_nufft: Running RM-synthesis using the NUFFTs over 2001 Faraday depth channels.
INFO synthesis.rmsynth_nufft: NUFFT complete in 0.00598 seconds.
INFO synthesis.get_rmsf_nufft: Fitting main lobe in each RMSF spectrum.
INFO rmsynth._run_rmsynth: RM-synthesis completed in 14.69ms.
[12]:
fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rm_syth_results

phi_arr_radm2 = fdf_arrs["phi_arr_radm2"].to_numpy()
fdf_dirty_arr = fdf_arrs["fdf_dirty_complex_arr"].to_numpy().astype(complex)

fig, ax = plt.subplots()

ax.plot(
    phi_arr_radm2,
    fdf_dirty_arr.real,
    color="tab:red",
    label="Stokes Q",
)
ax.plot(
    phi_arr_radm2,
    fdf_dirty_arr.imag,
    color="tab:blue",
    label="Stokes U",
)
ax.plot(
    phi_arr_radm2,
    np.abs(fdf_dirty_arr),
    color="k",
    label="Polarized intensity",
)

ax.errorbar(
    fdf_parameters["peak_rm_fit"],
    fdf_parameters["peak_pi_fit"],
    xerr=fdf_parameters["peak_rm_fit_error"],
    yerr=fdf_parameters["peak_pi_error"],
    fmt="o",
    lw=1,
    color="red",
    mfc="none",
    label="Fitted peak",
)

ax.set(
    xlabel=rf"$\phi$ / {u.rad / u.m**2:latex_inline}",
    ylabel="Flux density",
    title="Dirty FDF",
)
ax.legend()
[12]:
<matplotlib.legend.Legend at 0x79153408f6e0>
../_images/examples_rmclean_1d_17_1.png
[13]:
rmclean_results = rmclean.run_rmclean_from_synth(
    rm_synth_1d_results=rm_syth_results, auto_mask=10, auto_threshold=0.5
)

clean_arrs = rmclean_results.fdf_arrs
clean_fdf_params = rmclean_results.fdf_parameters
clean_params = rmclean_results.clean_parameters


phi_arr_radm2 = clean_arrs["phi_arr_radm2"].to_numpy()
fdf_dirty_arr = clean_arrs["fdf_dirty_complex_arr"].to_numpy().astype(complex)
fdf_clean_arr = clean_arrs["fdf_clean_complex_arr"].to_numpy().astype(complex)
fdf_model_arr = clean_arrs["fdf_model_complex_arr"].to_numpy().astype(complex)
fdf_residual_arr = clean_arrs["fdf_residual_complex_arr"].to_numpy().astype(complex)
INFO rmclean.run_rmclean_from_synth: Theoretical noise: TheoreticalNoise(fdf_error_noise=0.006804138174397718, fdf_q_noise=0.013608276348795436, fdf_u_noise=0.0)
INFO rmclean.run_rmclean_from_synth: Auto mask: 0.07, Auto threshold: 0.00, Max iterations: 10000, Gain: 0.1
INFO clean.minor_cycle: Starting initial minor loop...
INFO clean.minor_loop: Starting minor loop... 449 pixels in the mask
WARNING clean.minor_loop: All channels masked. Exiting loop...performed 18 iterations
INFO clean.minor_cycle: Initial loop complete. Starting deep clean...
INFO clean.minor_loop: Starting minor loop... 14 pixels in the mask
INFO clean.minor_loop: Threshold reached. Exiting loop...performed 140 iterations
[14]:
fig, ax = plt.subplots()

ax.plot(
    phi_arr_radm2,
    np.abs(fdf_dirty_arr),
    color="k",
    label="Dirty FDF",
    alpha=0.5,
)
ax.plot(
    phi_arr_radm2,
    np.abs(fdf_clean_arr),
    color="k",
    label="Clean FDF",
)
ax.step(
    phi_arr_radm2,
    np.abs(fdf_model_arr),
    color="tab:red",
    label="Model FDF",
    where="mid",
)
ax.plot(
    phi_arr_radm2,
    np.abs(fdf_residual_arr),
    color="tab:blue",
    label="Residual FDF",
    alpha=0.5,
)

ax.errorbar(
    clean_fdf_params["peak_rm_fit"],
    clean_fdf_params["peak_pi_fit"],
    xerr=clean_fdf_params["peak_rm_fit_error"],
    yerr=clean_fdf_params["peak_pi_error"],
    fmt="o",
    lw=1,
    color="red",
    mfc="none",
    label="Fitted peak",
)
ax.axhline(
    y=clean_params["threshold"][0],
    color="tab:blue",
    label="Threshold",
    linestyle="--",
)
ax.axhline(
    y=clean_params["mask"][0],
    color="tab:orange",
    label="Mask",
    linestyle="--",
)

ax.set(
    xlabel=rf"$\phi$ / {u.rad / u.m**2:latex_inline}",
    ylabel="Flux density",
    title="CLEANed FDF",
)
ax.legend()
[14]:
<matplotlib.legend.Legend at 0x79152d7ab710>
../_images/examples_rmclean_1d_19_1.png
[ ]: