diff --git a/.gitignore b/.gitignore index 9612d13..e304708 100644 --- a/.gitignore +++ b/.gitignore @@ -152,6 +152,13 @@ tests/simdata tests/models_ns *egg-info .DS_Store -*.ipynb scratch.ipynb rm_lite/_version.py +complexity.ipynb +fitting.ipynb +scal_activate.ipynb +scatcch.ipynb +rm_lite/clean.ipynb +rm_lite/ms_test.ipynb +rm_lite/multiscale.ipynb +rm_lite/Untitled-1.ipynb diff --git a/docs/examples/rmsynth_1d.ipynb b/docs/examples/rmsynth_1d.ipynb new file mode 100644 index 0000000..6565575 --- /dev/null +++ b/docs/examples/rmsynth_1d.ipynb @@ -0,0 +1,576 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# 1D RM-synthesis" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import astropy.units as u\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "from astropy.visualization import quantity_support\n", + "from numpy.typing import NDArray\n", + "from rm_lite.tools_1d import rmsynth\n", + "\n", + "plt.rcParams[\"figure.dpi\"] = 150\n", + "\n", + "_ = quantity_support()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "First generate some synthetic data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from rm_lite.utils.synthesis import faraday_simple_spectrum, freq_to_lambda2\n", + "\n", + "\n", + "def faraday_slab_spectrum(\n", + " freq_arr_hz: NDArray[np.float64],\n", + " frac_pol: float,\n", + " psi0_deg: float,\n", + " rm_radm2: float,\n", + " delta_rm_radm2: float,\n", + ") -> NDArray[np.complex128]:\n", + " lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz)\n", + "\n", + " return (\n", + " frac_pol\n", + " * np.exp(2j * (np.deg2rad(psi0_deg) + rm_radm2 * lambda_sq_arr_m2))\n", + " * (\n", + " np.sin(delta_rm_radm2 * lambda_sq_arr_m2)\n", + " / (delta_rm_radm2 * lambda_sq_arr_m2)\n", + " )\n", + " )\n", + "\n", + "\n", + "def faraday_gaussian_spectrum(\n", + " freq_arr_hz: NDArray[np.float64],\n", + " frac_pol: float,\n", + " psi0_deg: float,\n", + " rm_radm2: float,\n", + " sigma_rm_radm2: float,\n", + "):\n", + " lambda_sq_arr_m2 = freq_to_lambda2(freq_arr_hz)\n", + " rm_term = np.exp(2j * (np.deg2rad(psi0_deg) + rm_radm2 * lambda_sq_arr_m2))\n", + " depol_term = np.exp(-2.0 * sigma_rm_radm2**2 * lambda_sq_arr_m2**2)\n", + " return frac_pol * rm_term * depol_term" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here we'll simulate RACS-all frequency coverage" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "bw_low = 288\n", + "bw_mid = 144\n", + "bw_high = 288\n", + "low = np.linspace(943.5 - bw_low / 2, 943.5 + bw_low / 2, 36) * u.MHz\n", + "mid = np.linspace(1367.5 - bw_mid / 2, 1367.5 + bw_mid / 2, 9) * u.MHz\n", + "high = np.linspace(1655.5 - bw_high / 2, 1655.5 + bw_high / 2, 9) * u.MHz\n", + "freqs = np.concatenate([low, mid, high])\n", + "freq_hz = freqs.to(u.Hz).value" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we make a Faraday simple spectrum with a single RM component. We will use the following parameters:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "delta_rm_radm2 = 30\n", + "rm_radm2 = 100\n", + "frac_pol = 0.5\n", + "psi0_deg = 10\n", + "complex_data_noiseless = faraday_simple_spectrum(\n", + " freq_arr_hz=freq_hz,\n", + " frac_pol=frac_pol,\n", + " psi0_deg=psi0_deg,\n", + " rm_radm2=rm_radm2,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots()\n", + "ax.plot(\n", + " freq_hz, np.real(complex_data_noiseless), \".\", label=\"Stokes Q\", color=\"tab:red\"\n", + ")\n", + "ax.plot(\n", + " freq_hz, np.imag(complex_data_noiseless), \".\", label=\"Stokes U\", color=\"tab:blue\"\n", + ")\n", + "ax.legend()\n", + "ax.set(\n", + " xlabel=rf\"$\\nu$ / {u.Hz:latex_inline}\",\n", + " ylabel=\"Flux density\",\n", + " title=\"Stokes Q and U\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can run RM-synthesis by calling `rmsynth.run_rmsynth`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "help(rmsynth.run_rmsynth)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rmsynth.run_rmsynth(\n", + " freq_arr_hz=freq_hz,\n", + " complex_pol_arr=complex_data_noiseless,\n", + " complex_pol_error=np.zeros_like(complex_data_noiseless),\n", + " do_fit_rmsf=True,\n", + " n_samples=100,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The output values are Polars dataframes that can be inspected easily" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fdf_parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fdf_arrs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "rmsf_arrs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Since we provided no Stokes $I$ data, the stokes I model will just be unity with 0 error. The `flag_arr` array tells us which channels were not used in RM-synthesis or model fitting" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "stokes_i_arrs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "We can also easily visualise the data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phi_arr_radm2 = fdf_arrs[\"phi_arr_radm2\"].to_numpy()\n", + "fdf_dirty_arr = fdf_arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex)\n", + "\n", + "fig, ax = plt.subplots()\n", + "x1, x2, y1, y2 = 95, 105, 0.45, 0.55 # subregion of the original image\n", + "axins = ax.inset_axes(\n", + " (0.9, 0.6, 0.4, 0.4), xlim=(x1, x2), ylim=(y1, y2), xticklabels=[], yticklabels=[]\n", + ")\n", + "for _ax in [ax, axins]:\n", + " _ax.plot(\n", + " phi_arr_radm2,\n", + " fdf_dirty_arr.real,\n", + " color=\"tab:red\",\n", + " label=\"Stokes Q\",\n", + " )\n", + " _ax.plot(\n", + " phi_arr_radm2,\n", + " fdf_dirty_arr.imag,\n", + " color=\"tab:blue\",\n", + " label=\"Stokes U\",\n", + " )\n", + " _ax.plot(\n", + " phi_arr_radm2,\n", + " np.abs(fdf_dirty_arr),\n", + " color=\"k\",\n", + " label=\"Polarized intensity\",\n", + " )\n", + "\n", + " _ax.errorbar(\n", + " fdf_parameters[\"peak_rm_fit\"],\n", + " fdf_parameters[\"peak_pi_fit\"],\n", + " xerr=fdf_parameters[\"peak_rm_fit_error\"],\n", + " yerr=fdf_parameters[\"peak_pi_error\"],\n", + " fmt=\"o\",\n", + " lw=1,\n", + " color=\"red\",\n", + " mfc=\"none\",\n", + " label=\"Fitted peak\",\n", + " )\n", + "\n", + "ax.set(\n", + " xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n", + " ylabel=\"Flux density\",\n", + " title=\"Dirty FDF\",\n", + " # xlim=[50, 150],\n", + ")\n", + "ax.indicate_inset_zoom(axins, edgecolor=\"black\")\n", + "ax.legend()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phi2_arr_radm2 = rmsf_arrs[\"phi2_arr_radm2\"].to_numpy()\n", + "rmsf_arr = rmsf_arrs[\"rmsf_complex_arr\"].to_numpy().astype(complex)\n", + "\n", + "fig, ax = plt.subplots()\n", + "ax.plot(\n", + " phi2_arr_radm2,\n", + " rmsf_arr.real,\n", + " color=\"tab:red\",\n", + " label=\"Stokes Q\",\n", + ")\n", + "ax.plot(\n", + " phi2_arr_radm2,\n", + " rmsf_arr.imag,\n", + " color=\"tab:blue\",\n", + " label=\"Stokes U\",\n", + ")\n", + "ax.plot(\n", + " phi2_arr_radm2,\n", + " np.abs(rmsf_arr),\n", + " color=\"k\",\n", + " label=\"Polarized intensity\",\n", + ")\n", + "ax.legend()\n", + "ax.set(\n", + " xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n", + " ylabel=\"RMSF\",\n", + " title=\"RMSF\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now lets do a more complex example. We'll add noise and a Stokes $I$ spectrum" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from rm_lite.utils.fitting import power_law" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "delta_rm_radm2 = 30\n", + "rm_radm2 = 100\n", + "frac_pol = 0.5\n", + "psi0_deg = 10\n", + "complex_data_noiseless = faraday_slab_spectrum(\n", + " freq_arr_hz=freq_hz,\n", + " frac_pol=frac_pol,\n", + " psi0_deg=psi0_deg,\n", + " rm_radm2=rm_radm2,\n", + " delta_rm_radm2=delta_rm_radm2,\n", + ")\n", + "\n", + "\n", + "stokes_i_flux = 1.0\n", + "spectral_index = -0.7\n", + "rms_noise = 0.1\n", + "\n", + "rng = np.random.default_rng()\n", + "\n", + "stokes_i_model = power_law(order=1)\n", + "stokes_i_noiseless = stokes_i_model(\n", + " freq_hz / (np.mean(freq_hz)), stokes_i_flux, spectral_index\n", + ")\n", + "stokes_i_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n", + "stokes_i_noisy = stokes_i_noiseless + stokes_i_noise\n", + "\n", + "\n", + "stokes_q_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n", + "stokes_u_noise = rng.normal(0, rms_noise, size=freq_hz.size)\n", + "complex_noise = stokes_q_noise + 1j * stokes_u_noise\n", + "\n", + "complex_flux = complex_data_noiseless * stokes_i_noiseless\n", + "complex_data_noisy = complex_data_noiseless + complex_noise" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we enable Stokes $I$ model fitting through providing the data, and enabling `fit_order`. If `fit_order<0` an iterative fit will be performed." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fdf_parameters, fdf_arrs, rmsf_arrs, stokes_i_arrs = rmsynth.run_rmsynth(\n", + " freq_arr_hz=freq_hz,\n", + " complex_pol_arr=complex_data_noisy,\n", + " complex_pol_error=np.ones_like(complex_data_noiseless)\n", + " * (rms_noise + rms_noise * 1j),\n", + " stokes_i_arr=stokes_i_noisy,\n", + " stokes_i_error_arr=np.ones_like(stokes_i_noisy) * rms_noise,\n", + " do_fit_rmsf=True,\n", + " n_samples=100,\n", + " fit_order=-3,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fdf_parameters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fdf_arrs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "stokes_i_arrs" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots()\n", + "ax.plot(freq_hz, stokes_i_noiseless, label=\"Input model\")\n", + "ax.plot(freq_hz, stokes_i_noisy, \".\", label=\"Noisy data\")\n", + "ax.plot(\n", + " stokes_i_arrs[\"freq_arr_hz\"],\n", + " stokes_i_arrs[\"stokes_i_model_arr\"],\n", + " \"k--\",\n", + " label=\"Fitted model\",\n", + ")\n", + "ax.fill_between(\n", + " stokes_i_arrs[\"freq_arr_hz\"],\n", + " stokes_i_arrs[\"stokes_i_model_arr\"] - stokes_i_arrs[\"stokes_i_model_error\"],\n", + " stokes_i_arrs[\"stokes_i_model_arr\"] + stokes_i_arrs[\"stokes_i_model_error\"],\n", + " alpha=0.3,\n", + " color=\"k\",\n", + " label=\"Fitted model error\",\n", + ")\n", + "ax.legend()\n", + "ax.set(\n", + " xlabel=rf\"$\\nu$ / {u.Hz:latex_inline}\",\n", + " ylabel=\"Flux density\",\n", + " title=\"Stokes I\",\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phi_arr_radm2 = fdf_arrs[\"phi_arr_radm2\"].to_numpy()\n", + "fdf_dirty_arr = fdf_arrs[\"fdf_dirty_complex_arr\"].to_numpy().astype(complex)\n", + "\n", + "fig, ax = plt.subplots()\n", + "\n", + "ax.plot(\n", + " phi_arr_radm2,\n", + " fdf_dirty_arr.real,\n", + " color=\"tab:red\",\n", + " label=\"Stokes Q\",\n", + ")\n", + "ax.plot(\n", + " phi_arr_radm2,\n", + " fdf_dirty_arr.imag,\n", + " color=\"tab:blue\",\n", + " label=\"Stokes U\",\n", + ")\n", + "ax.plot(\n", + " phi_arr_radm2,\n", + " np.abs(fdf_dirty_arr),\n", + " color=\"k\",\n", + " label=\"Polarized intensity\",\n", + ")\n", + "\n", + "ax.errorbar(\n", + " fdf_parameters[\"peak_rm_fit\"],\n", + " fdf_parameters[\"peak_pi_fit\"],\n", + " xerr=fdf_parameters[\"peak_rm_fit_error\"],\n", + " yerr=fdf_parameters[\"peak_pi_error\"],\n", + " fmt=\"o\",\n", + " lw=1,\n", + " color=\"red\",\n", + " mfc=\"none\",\n", + " label=\"Fitted peak\",\n", + ")\n", + "\n", + "ax.set(\n", + " xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n", + " ylabel=\"Flux density\",\n", + " title=\"Dirty FDF\",\n", + ")\n", + "ax.legend()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "phi2_arr_radm2 = rmsf_arrs[\"phi2_arr_radm2\"].to_numpy()\n", + "rmsf_arr = rmsf_arrs[\"rmsf_complex_arr\"].to_numpy().astype(complex)\n", + "\n", + "fig, ax = plt.subplots()\n", + "ax.plot(\n", + " phi2_arr_radm2,\n", + " rmsf_arr.real,\n", + " color=\"tab:red\",\n", + " label=\"Stokes Q\",\n", + ")\n", + "ax.plot(\n", + " phi2_arr_radm2,\n", + " rmsf_arr.imag,\n", + " color=\"tab:blue\",\n", + " label=\"Stokes U\",\n", + ")\n", + "ax.plot(\n", + " phi2_arr_radm2,\n", + " np.abs(rmsf_arr),\n", + " color=\"k\",\n", + " label=\"Polarized intensity\",\n", + ")\n", + "ax.legend()\n", + "ax.set(\n", + " xlabel=rf\"$\\phi$ / {u.rad / u.m**2:latex_inline}\",\n", + " ylabel=\"RMSF\",\n", + " title=\"RMSF\",\n", + ")" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "rm-lite", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.9" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/docs/usage.md b/docs/usage.md index f3a5f80..a9c0d2b 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -7,6 +7,6 @@ examples should work as written. ```{toctree} :maxdepth: 1 :caption: Example notebooks: -examples/rmsyth_1d +examples/rmsynth_1d.ipynb ``` diff --git a/pyproject.toml b/pyproject.toml index b50ad9d..04cd22a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,8 @@ dev = [ "pytest-cov", "nox", "nbconvert", + "jupyter", + "ipython" ] docs = [ "sphinx>=7.0",