Skip to content

Commit

Permalink
Catch custom models in 1D RMsynth (#126)
Browse files Browse the repository at this point in the history
* I love the smell of Rust in the morning

* Paren typo

* Not optional vals
  • Loading branch information
AlecThomson authored Apr 11, 2024
1 parent 5cb7a01 commit 6d475fb
Show file tree
Hide file tree
Showing 6 changed files with 245 additions and 190 deletions.
25 changes: 16 additions & 9 deletions RMtools_1D/clean_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
import numpy as np

from RMtools_1D.do_RMsynth_1D import readFile as read_freqFile
from RMutils.util_misc import calculate_StokesI_model
from RMutils.util_misc import FitResult, calculate_StokesI_model


def calculate_QU_model(freqArr, phiArr, CCArr, lambdaSq_0, Iparms=None):
Expand All @@ -39,15 +39,22 @@ def calculate_QU_model(freqArr, phiArr, CCArr, lambdaSq_0, Iparms=None):
a = lambdaSqArr_m2 - lambdaSq_0
quarr = np.sum(CCArr[:, np.newaxis] * np.exp(2.0j * np.outer(phiArr, a)), axis=0)

fitDict = {}
# TODO: Pass in fit function, which is currently not output by rmsynth1d
fitDict["fit_function"] = "log"
if Iparms is not None:
fitDict["p"] = Iparms
else:
fitDict["p"] = [0, 0, 0, 0, 0, 1]
fitDict["reference_frequency_Hz"] = C / np.sqrt(lambdaSq_0)
StokesI_model = calculate_StokesI_model(fitDict, freqArr)
fit_result = FitResult(
params=Iparms if Iparms is not None else [0, 0, 0, 0, 0, 1],
fitStatus=0,
chiSq=0.0,
chiSqRed=0.0,
AIC=0,
polyOrd=0,
nIter=0,
reference_frequency_Hz=C / np.sqrt(lambdaSq_0),
dof=0,
pcov=None,
perror=None,
fit_function="log",
)
StokesI_model = calculate_StokesI_model(fit_result, freqArr)

QUarr = StokesI_model * quarr

Expand Down
21 changes: 9 additions & 12 deletions RMtools_1D/do_QUfit_1D_mnest.py
Original file line number Diff line number Diff line change
Expand Up @@ -324,12 +324,12 @@ def run_qufit(
verbose=True,
fit_function=fit_function,
)
(IModArr, qArr, uArr, dqArr, duArr, IfitDict) = dataArr
(IModArr, qArr, uArr, dqArr, duArr, fit_result) = dataArr

# Plot the data and the Stokes I model fit
print("Plotting the input data and spectral index fit.")
freqHirArr_Hz = np.linspace(freqArr_Hz[0], freqArr_Hz[-1], 10000)
IModHirArr = calculate_StokesI_model(IfitDict, freqHirArr_Hz)
IModHirArr = calculate_StokesI_model(fit_result, freqHirArr_Hz)
specFig = plt.figure(facecolor="w", figsize=(10, 6))
plot_Ipqu_spectra_fig(
freqArr_Hz=freqArr_Hz,
Expand Down Expand Up @@ -487,7 +487,6 @@ def run_qufit(

# Create a save dictionary and store final p in values
outFile = f"{prefixOut}_m{modelNum}_{sampler}.json"
# IfitDict["p"] = toscalar(IfitDict["p"].tolist())
saveDict = {
"parNames": toscalar(parNames),
"labels": toscalar(labels),
Expand All @@ -506,14 +505,12 @@ def run_qufit(
"ln(EVIDENCE) ": toscalar(lnEvidence),
"dLn(EVIDENCE)": toscalar(dLnEvidence),
"nFree": toscalar(nFree),
"Imodel": ",".join([str(x.astype(np.float32)) for x in IfitDict["p"]]),
"Imodel_errs": ",".join(
[str(x.astype(np.float32)) for x in IfitDict["perror"]]
),
"IfitChiSq": toscalar(IfitDict["chiSq"]),
"IfitChiSqRed": toscalar(IfitDict["chiSqRed"]),
"IfitPolyOrd": toscalar(IfitDict["polyOrd"]),
"Ifitfreq0": toscalar(IfitDict["reference_frequency_Hz"]),
"Imodel": ",".join([str(x.astype(np.float32)) for x in fit_result.params]),
"Imodel_errs": ",".join([str(x.astype(np.float32)) for x in fit_result.perror]),
"IfitChiSq": toscalar(fit_result.chiSq),
"IfitChiSqRed": toscalar(fit_result.chiSqRed),
"IfitPolyOrd": toscalar(fit_result.polyOrd),
"Ifitfreq0": toscalar(fit_result.reference_frequency_Hz),
}

for k, v in saveDict.items():
Expand Down Expand Up @@ -564,7 +561,7 @@ def run_qufit(
# Plot the data and best-fitting model
lamSqHirArr_m2 = np.linspace(lamSqArr_m2[0], lamSqArr_m2[-1], 10000)
freqHirArr_Hz = C / np.sqrt(lamSqHirArr_m2)
IModArr = calculate_StokesI_model(IfitDict, freqHirArr_Hz)
IModArr = calculate_StokesI_model(fit_result, freqHirArr_Hz)
pDict = {k: v for k, v in zip(parNames, p)}
quModArr = model(pDict, lamSqHirArr_m2)
model_dict = {
Expand Down
39 changes: 17 additions & 22 deletions RMtools_1D/do_RMsynth_1D.py
Original file line number Diff line number Diff line change
Expand Up @@ -189,7 +189,7 @@ def run_rmsynth(
dQUArr = (dQArr + dUArr) / 2.0

# Fit the Stokes I spectrum and create the fractional spectra
IModArr, qArr, uArr, dqArr, duArr, fitDict = create_frac_spectra(
IModArr, qArr, uArr, dqArr, duArr, fit_result = create_frac_spectra(
freqArr=freqArr_Hz,
IArr=IArr,
QArr=QArr,
Expand All @@ -212,7 +212,7 @@ def run_rmsynth(
log("Plotting the input data and spectral index fit.")
freqHirArr_Hz = np.linspace(freqArr_Hz[0], freqArr_Hz[-1], 10000)
if modStokesI is None:
IModHirArr = calculate_StokesI_model(fitDict, freqHirArr_Hz)
IModHirArr = calculate_StokesI_model(fit_result, freqHirArr_Hz)
elif modStokesI is not None:
modStokesI_interp = interp1d(freqArr_Hz, modStokesI)
IModHirArr = modStokesI_interp(freqHirArr_Hz)
Expand Down Expand Up @@ -356,16 +356,19 @@ def run_rmsynth(
# Stokes I reference frequency.

if lam0Sq_m2 == 0: # Rudnick-Cotton adapatation
freq0_Hz = fitDict["reference_frequency_Hz"]
freq0_Hz = fit_result.reference_frequency_Hz
else: # standard RM-synthesis
freq0_Hz = C / m.sqrt(lam0Sq_m2)
fitDict = renormalize_StokesI_model(fitDict, freq0_Hz)
if modStokesI is None:
fit_result = renormalize_StokesI_model(fit_result, freq0_Hz)
else:
fit_result = fit_result.with_options(reference_frequency_Hz=freq0_Hz)

# Set Ifreq0 (Stokes I at reference frequency) from either supplied model
# (interpolated as required) or fit model, as appropriate.
# Multiply the dirty FDF by Ifreq0 to recover the PI
if modStokesI is None:
Ifreq0 = calculate_StokesI_model(fitDict, freq0_Hz)
Ifreq0 = calculate_StokesI_model(fit_result, freq0_Hz)
elif modStokesI is not None:
modStokesI_interp = interp1d(freqArr_Hz, modStokesI)
Ifreq0 = modStokesI_interp(freq0_Hz)
Expand All @@ -388,13 +391,15 @@ def run_rmsynth(
lam0Sq=lam0Sq_m2,
)
mDict["Ifreq0"] = toscalar(Ifreq0)
mDict["polyCoeffs"] = ",".join([str(x.astype(np.float32)) for x in fitDict["p"]])
mDict["polyCoeffs"] = ",".join(
[str(x.astype(np.float32)) for x in fit_result.params]
)
mDict["polyCoefferr"] = ",".join(
[str(x.astype(np.float32)) for x in fitDict["perror"]]
[str(x.astype(np.float32)) for x in fit_result.perror]
)
mDict["polyOrd"] = fitDict["polyOrd"]
mDict["IfitStat"] = fitDict["fitStatus"]
mDict["IfitChiSqRed"] = fitDict["chiSqRed"]
mDict["polyOrd"] = fit_result.polyOrd
mDict["IfitStat"] = fit_result.fitStatus
mDict["IfitChiSqRed"] = fit_result.chiSqRed
mDict["fit_function"] = fit_function
mDict["lam0Sq_m2"] = toscalar(lam0Sq_m2)
mDict["freq0_Hz"] = toscalar(freq0_Hz)
Expand All @@ -403,9 +408,9 @@ def run_rmsynth(
mDict["dFDFth"] = toscalar(dFDFth)
mDict["units"] = units

if (fitDict["fitStatus"] >= 128) and verbose:
if (fit_result.fitStatus >= 128) and verbose:
log("WARNING: Stokes I model contains negative values!")
elif (fitDict["fitStatus"] >= 64) and verbose:
elif (fit_result.fitStatus >= 64) and verbose:
log("Caution: Stokes I model has low signal-to-noise.")

# Add information on nature of channels:
Expand Down Expand Up @@ -520,16 +525,6 @@ def run_rmsynth(
units=units,
)

# Use the custom navigation toolbar
# try:
# fdfFig.canvas.toolbar.pack_forget()
# CustomNavbar(fdfFig.canvas, fdfFig.canvas.toolbar.window)
# except Exception:
# pass

# Display the figure
# fdfFig.show()

# Pause if plotting enabled
if showPlots:
plt.show()
Expand Down
16 changes: 8 additions & 8 deletions RMtools_1D/rmtools_bwdepol.py
Original file line number Diff line number Diff line change
Expand Up @@ -1068,7 +1068,7 @@ def run_adjoint_rmsynth(
dQUArr = (dQArr + dUArr) / 2.0

# Fit the Stokes I spectrum and create the fractional spectra
IModArr, qArr, uArr, dqArr, duArr, fitDict = create_frac_spectra(
IModArr, qArr, uArr, dqArr, duArr, fit_result = create_frac_spectra(
freqArr=freqArr_GHz,
IArr=IArr,
QArr=QArr,
Expand All @@ -1086,7 +1086,7 @@ def run_adjoint_rmsynth(
if verbose:
log("Plotting the input data and spectral index fit.")
freqHirArr_Hz = np.linspace(freqArr_Hz[0], freqArr_Hz[-1], 10000)
IModHirArr = poly5(fitDict["p"])(freqHirArr_Hz / 1e9)
IModHirArr = poly5(fit_result.params)(freqHirArr_Hz / 1e9)
specFig = plt.figure(figsize=(12.0, 8))
plot_Ipqu_spectra_fig(
freqArr_Hz=freqArr_Hz,
Expand Down Expand Up @@ -1215,7 +1215,7 @@ def run_adjoint_rmsynth(
# Determine the Stokes I value at lam0Sq_m2 from the Stokes I model
# Multiply the dirty FDF by Ifreq0 to recover the PI
freq0_Hz = C / m.sqrt(lam0Sq_m2)
Ifreq0 = poly5(fitDict["p"])(freq0_Hz / 1e9)
Ifreq0 = poly5(fit_result.params)(freq0_Hz / 1e9)
dirtyFDF *= Ifreq0 # FDF is in fracpol units initially, convert back to flux

# Calculate the theoretical noise in the FDF !!
Expand All @@ -1240,19 +1240,19 @@ def run_adjoint_rmsynth(
)

mDict["Ifreq0"] = toscalar(Ifreq0)
mDict["polyCoeffs"] = ",".join([str(x) for x in fitDict["p"]])
mDict["IfitStat"] = fitDict["fitStatus"]
mDict["IfitChiSqRed"] = fitDict["chiSqRed"]
mDict["polyCoeffs"] = ",".join([str(x) for x in fit_result.params])
mDict["IfitStat"] = fit_result.fitStatus
mDict["IfitChiSqRed"] = fit_result.chiSqRed
mDict["lam0Sq_m2"] = toscalar(lam0Sq_m2)
mDict["freq0_Hz"] = toscalar(freq0_Hz)
mDict["fwhmRMSF"] = toscalar(fwhmRMSF)
mDict["dQU"] = toscalar(nanmedian(dQUArr))
# mDict["dFDFth"] = toscalar(dFDFth)
mDict["units"] = units

if fitDict["fitStatus"] >= 128:
if fit_result.fitStatus >= 128:
log("WARNING: Stokes I model contains negative values!")
elif fitDict["fitStatus"] >= 64:
elif fit_result.fitStatus >= 64:
log("Caution: Stokes I model has low signal-to-noise.")

# Add information on nature of channels:
Expand Down
20 changes: 10 additions & 10 deletions RMtools_3D/do_fitIcube.py
Original file line number Diff line number Diff line change
Expand Up @@ -409,23 +409,23 @@ def fit_spectra_I(
nDetectPix: the total number of pixels to be fit.
"""

pixFitDict = fit_StokesI_model(
pix_fit_result = fit_StokesI_model(
freqArr_Hz, Ispectrum, rms_Arr, polyOrd=polyOrd, fit_function=fit_function
)

pixImodel = calculate_StokesI_model(pixFitDict, freqArr_Hz)
pixImodel = calculate_StokesI_model(pix_fit_result, freqArr_Hz)

outs = dict()

outs["I"] = pixImodel.astype("float32")
outs["coeffs"] = pixFitDict["p"].astype("float32")
outs["coeffs_err"] = pixFitDict["perror"].astype("float32")
outs["chiSq"] = pixFitDict["chiSq"]
outs["chiSqRed"] = pixFitDict["chiSqRed"]
outs["nIter"] = pixFitDict["nIter"]
outs["AIC"] = pixFitDict["AIC"]
outs["covar"] = pixFitDict["pcov"]
outs["reference_frequency_Hz"] = pixFitDict["reference_frequency_Hz"]
outs["coeffs"] = pix_fit_result.params.astype("float32")
outs["coeffs_err"] = pix_fit_result.perror.astype("float32")
outs["chiSq"] = pix_fit_result.chiSq
outs["chiSqRed"] = pix_fit_result.chiSqRed
outs["nIter"] = pix_fit_result.nIter
outs["AIC"] = pix_fit_result.AIC
outs["covar"] = pix_fit_result.pcov
outs["reference_frequency_Hz"] = pix_fit_result.reference_frequency_Hz

return outs

Expand Down
Loading

0 comments on commit 6d475fb

Please sign in to comment.