Skip to content

Commit

Permalink
Fix RMSF and add return vals
Browse files Browse the repository at this point in the history
  • Loading branch information
@AlecThomson
AlecThomson committed May 6, 2024
1 parent 724a078 commit 0db86fd
Show file tree
Hide file tree
Showing 2 changed files with 92 additions and 20 deletions.
41 changes: 30 additions & 11 deletions RMtools_1D/calculate_RMSF.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -174,7 +174,7 @@ def determine_RMSF_parameters(
-1 * phi_max / 2, phi_max / 2 + 1e-6, dphi
) # division by two accounts for how RMSF is always twice as wide as FDF.

RMSFcube, phi2Arr, fwhmRMSFArr, statArr = get_rmsf_planes(
rmsf_results = get_rmsf_planes(
lambda2_array,
phi_array,
weightArr=weights_array,
Expand All @@ -190,7 +190,11 @@ def determine_RMSF_parameters(
3.8 / (l2_max - l2_min)
)
)
print("Measured FWHM: {:.4g} rad m^-2".format(fwhmRMSFArr))
print(
"Measured FWHM: {:.4g} rad m^-2".format(
rmsf_results.fwhmRMSFArr
)
)
print("Theoretical largest FD scale probed: {:.4g} rad m^-2".format(np.pi / l2_min))
print(
"Theoretical maximum FD*: {:.4g} rad m^-2".format(
Expand All @@ -208,15 +212,21 @@ def determine_RMSF_parameters(
# finds the highest amplitude one, and calls that the first sidelobe.
try:
x = np.diff(
np.sign(np.diff(np.abs(RMSFcube[RMSFcube.size // 2 :])))
np.sign(
np.diff(
np.abs(rmsf_results.RMSFcube[rmsf_results.RMSFcube.size // 2 :])
)
)
) # -2=local max, +2=local min
y = (
1 + np.where(x == -2)[0]
) # indices of peaks, +1 is because of offset from double differencing
peaks = np.abs(RMSFcube[RMSFcube.size // 2 :])[y]
peaks = np.abs(rmsf_results.RMSFcube[rmsf_results.RMSFcube.size // 2 :])[y]
print(
"First sidelobe FD and amplitude: {:.4g} rad m^-2".format(
phi2Arr[phi2Arr.size // 2 :][y[np.argmax(peaks)]]
rmsf_results.phi2Arr[rmsf_results.phi2Arr.size // 2 :][
y[np.argmax(peaks)]
]
)
)
print(
Expand All @@ -235,9 +245,15 @@ def determine_RMSF_parameters(
plt.title("Simulated RMSF")
else:
plt.title(plotname)
plt.plot(phi2Arr, np.real(RMSFcube), "b-", label="Stokes Q")
plt.plot(phi2Arr, np.imag(RMSFcube), "r--", label="Stokes U")
plt.plot(phi2Arr, np.abs(RMSFcube), "k-", label="Amplitude")
plt.plot(
rmsf_results.phi2Arr, np.real(rmsf_results.RMSFcube), "b-", label="Stokes Q"
)
plt.plot(
rmsf_results.phi2Arr, np.imag(rmsf_results.RMSFcube), "r--", label="Stokes U"
)
plt.plot(
rmsf_results.phi2Arr, np.abs(rmsf_results.RMSFcube), "k-", label="Amplitude"
)
plt.legend()
plt.xlabel(r"Faraday depth (rad m$^{-2}$)")
plt.ylabel("RMSF (unitless)")
Expand All @@ -259,7 +275,7 @@ def determine_RMSF_parameters(
+ "# of channels: {:.4g}\n"
).format(
3.8 / (l2_max - l2_min),
fwhmRMSFArr,
rmsf_results.fwhmRMSFArr,
np.pi / l2_min,
np.sqrt(3.0) / dl2,
np.min(freq_array) / 1e9,
Expand All @@ -282,7 +298,10 @@ def determine_RMSF_parameters(
+ "First sidelobe FD and amplitude: {:.4g} rad m^-2\n"
+ " {:.4g} % of peak"
).format(
phi2Arr[phi2Arr.size // 2 :][y[np.argmax(peaks)]], np.max(peaks) * 100
rmsf_results.phi2Arr[rmsf_results.phi2Arr.size // 2 :][
y[np.argmax(peaks)]
],
np.max(peaks) * 100,
),
family="monospace",
horizontalalignment="left",
Expand All @@ -298,7 +317,7 @@ def determine_RMSF_parameters(
# ax.text(0.,0.22,'First sidelobe FD and amplitude: {:.4g} rad m^-2'.format(phi2Arr[phi2Arr.size//2:][y[np.argmax(peaks)]]))
# ax.text(0.,0.1,' {:.4g} % of peak'.format(np.max(peaks)*100))

if plotfile != None:
if plotfile is not None:
plt.savefig(plotfile, bbox_inches="tight")
else:
plt.show()
Expand Down
71 changes: 62 additions & 9 deletions RMutils/util_RM.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@
import gc
import math as m
import sys
from typing import NamedTuple

import finufft
import numpy as np
Expand All @@ -86,6 +87,15 @@


# -----------------------------------------------------------------------------#
class RMsynthResults(NamedTuple):
"""Results of the RM-synthesis calculation"""

FDFcube: np.ndarray
"""The Faraday dispersion function cube"""
lam0Sq_m2: float
"""The reference lambda^2 value"""


def do_rmsynth_planes(
dataQ,
dataU,
Expand All @@ -96,7 +106,7 @@ def do_rmsynth_planes(
nBits=32,
eps=1e-6,
log=print,
):
) -> RMsynthResults:
"""Perform RM-synthesis on Stokes Q and U cubes (1,2 or 3D). This version
of the routine loops through spectral planes and is faster than the pixel-
by-pixel code. This version also correctly deals with isolated clumps of
Expand Down Expand Up @@ -214,10 +224,25 @@ def do_rmsynth_planes(
# Remove redundant dimensions in the FDF array
FDFcube = np.squeeze(FDFcube)

return FDFcube, lam0Sq_m2
return RMsynthResults(FDFcube=FDFcube, lam0Sq_m2=lam0Sq_m2)


# -----------------------------------------------------------------------------#
class RMSFResults(NamedTuple):
"""Results of the RMSF calculation"""

RMSFcube: np.ndarray
"""The RMSF cube"""
phi2Arr: np.ndarray
"""The (double length) Faraday depth array"""
fwhmRMSFArr: np.ndarray
"""The FWHM of the RMSF main lobe"""
statArr: np.ndarray
"""The status of the RMSF fit"""
lam0Sq_m2: float
"""The reference lambda^2 value"""


def get_rmsf_planes(
lambdaSqArr_m2,
phiArr_radm2,
Expand All @@ -231,7 +256,7 @@ def get_rmsf_planes(
eps=1e-6,
verbose=False,
log=print,
):
) -> RMSFResults:
"""Calculate the Rotation Measure Spread Function from inputs. This version
returns a cube (1, 2 or 3D) of RMSF spectra based on the shape of a
boolean mask array, where flagged data are True and unflagged data False.
Expand Down Expand Up @@ -423,10 +448,29 @@ def get_rmsf_planes(
fwhmRMSFArr = np.reshape(fwhmRMSFArr, (old_data_shape[1], old_data_shape[2]))
statArr = np.reshape(statArr, (old_data_shape[1], old_data_shape[2]))

return RMSFcube, phi2Arr, fwhmRMSFArr, statArr, lam0Sq_m2
return RMSFResults(
RMSFcube=RMSFcube,
phi2Arr=phi2Arr,
fwhmRMSFArr=fwhmRMSFArr,
statArr=statArr,
lam0Sq_m2=lam0Sq_m2,
)


# -----------------------------------------------------------------------------#
class RMCleanResults(NamedTuple):
"""Results of the RM-CLEAN calculation"""

cleanFDF: np.ndarray
"""The cleaned Faraday dispersion function cube"""
ccArr: np.ndarray
"""The clean components cube"""
iterCountArr: np.ndarray
"""The number of iterations for each pixel"""
residFDF: np.ndarray
"""The residual Faraday dispersion function cube"""


def do_rmclean_hogbom(
dirtyFDF,
phiArr_radm2,
Expand All @@ -444,7 +488,7 @@ def do_rmclean_hogbom(
chunksize=None,
log=print,
window=0,
):
) -> RMCleanResults:
"""Perform Hogbom CLEAN on a cube of complex Faraday dispersion functions
given a cube of rotation measure spread functions.
Expand Down Expand Up @@ -593,10 +637,19 @@ def do_rmclean_hogbom(
iterCountArr = np.squeeze(iterCountArr)
residFDF = np.squeeze(residFDF)

return cleanFDF, ccArr, iterCountArr, residFDF
return RMCleanResults(cleanFDF, ccArr, iterCountArr, residFDF)


# -----------------------------------------------------------------------------#
class CleanLoopResults(NamedTuple):
"""Results of the RM-CLEAN loop"""

cleanFDF: np.ndarray
"""The cleaned Faraday dispersion function cube"""
residFDF: np.ndarray
"""The residual Faraday dispersion function cube"""
ccArr: np.ndarray
"""The clean components cube"""


class RMcleaner:
Expand Down Expand Up @@ -630,10 +683,10 @@ def __init__(
self.nbits = nbits
self.window = window

def cleanloop(self, args):
def cleanloop(self, args) -> CleanLoopResults:
return self._cleanloop(*args)

def _cleanloop(self, yi, xi, dirtyFDF):
def _cleanloop(self, yi, xi, dirtyFDF) -> CleanLoopResults:
dirtyFDF = dirtyFDF[:, yi, xi]
# Initialise arrays to hold the residual FDF, clean components, clean FDF
residFDF = dirtyFDF.copy()
Expand Down Expand Up @@ -724,7 +777,7 @@ def _cleanloop(self, yi, xi, dirtyFDF):
residFDF = np.squeeze(residFDF)
ccArr = np.squeeze(ccArr)

return cleanFDF, residFDF, ccArr
return CleanLoopResults(cleanFDF=cleanFDF, residFDF=residFDF, ccArr=ccArr)


# -----------------------------------------------------------------------------#
Expand Down

0 comments on commit 0db86fd

Please sign in to comment.