Stokes I 3d frequency fix

RMtools_1D/do_RMsynth_1D.py

@@ -351,21 +351,27 @@ def run_rmsynth(
     if verbose:
         log("> RM-synthesis completed in %.2f seconds." % cputime)
 
-    # Determine the Stokes I value at lam0Sq_m2 from the Stokes I model
-    # This will break if lam0Sq_m2==0. Using the mean frequency in this case.
+    # Convert Stokes I model to polarization reference frequency. If lambda^2_0 is
+    # non-zero, use that as polarization reference frequency and adapt Stokes I model.
+    # If lambda^2_0 is zero, make polarization reference frequency equal to
+    # Stokes I reference frequency.
+
+    if lam0Sq_m2 == 0:  # Rudnick-Cotton adapatation
+        freq0_Hz = fitDict["reference_frequency_Hz"]
+    else:  # standard RM-synthesis
+        freq0_Hz = C / m.sqrt(lam0Sq_m2)
+        fitDict = renormalize_StokesI_model(fitDict, 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
-    freq0_Hz = C / m.sqrt(lam0Sq_m2) if lam0Sq_m2 > 0 else np.nanmean(freqArr_Hz)
     if modStokesI is None:
         Ifreq0 = calculate_StokesI_model(fitDict, freq0_Hz)
     elif modStokesI is not None:
         modStokesI_interp = interp1d(freqArr_Hz, modStokesI)
         Ifreq0 = modStokesI_interp(freq0_Hz)
     dirtyFDF *= Ifreq0  # FDF is in fracpol units initially, convert back to flux
 
-    # if modStokesI is None:
-    #     #Need to renormalize the Stokes I parameters here to the actual reference frequency.
-    #     fitDict=renormalize_StokesI_model(fitDict,freq0_Hz)
-
     # Calculate the theoretical noise in the FDF !!Old formula only works for wariance weights!
     weightArr = np.where(np.isnan(weightArr), 0.0, weightArr)
     dFDFth = np.abs(Ifreq0) * np.sqrt(
@@ -387,7 +393,6 @@ def run_rmsynth(
     mDict["polyCoefferr"] = ",".join(
         [str(x.astype(np.float32)) for x in fitDict["perror"]]
     )
-    mDict["poly_reffreq"] = fitDict["reference_frequency_Hz"]
     mDict["polyOrd"] = fitDict["polyOrd"]
     mDict["IfitStat"] = fitDict["fitStatus"]
     mDict["IfitChiSqRed"] = fitDict["chiSqRed"]

RMtools_3D/do_fitIcube.py

@@ -35,15 +35,13 @@
 # =============================================================================#
 
 import argparse
-import multiprocessing as mp
 import os
 import sys
 import time
 from functools import partial
 
 import astropy.io.fits as pf
 import numpy as np
-from tqdm.auto import tqdm
 from tqdm.contrib.concurrent import process_map
 
 from RMtools_3D.do_RMsynth_3D import readFitsCube
@@ -329,7 +327,8 @@ def savefits_mask(data, header, outDir, prefixOut):
     headMask = strip_fits_dims(header=header, minDim=2)
     headMask["DATAMAX"] = 1
     headMask["DATAMIN"] = 0
-    del headMask["BUNIT"]
+    if "BUNIT" in headMask:
+        del headMask["BUNIT"]
 
     mskArr = np.where(data > 0, 1.0, np.nan)
     MaskfitsFile = os.path.join(outDir, prefixOut + "mask.fits")
@@ -348,17 +347,16 @@ def savefits_Coeffs(data, dataerr, header, polyOrd, outDir, prefixOut):
     prefixOut: prefix to use on the output name
     """
 
-    headcoeff = strip_fits_dims(header=header, minDim=2)
-    headcoeff["BUNIT"] = ""
-    if "BTYPE" in headcoeff:
-        del headcoeff["BTYPE"]
+    header["BUNIT"] = ""
+    if "BTYPE" in header:
+        del header["BTYPE"]
 
     for i in range(np.abs(polyOrd) + 1):
         outname = os.path.join(outDir, prefixOut + "coeff" + str(i) + ".fits")
-        pf.writeto(outname, data[i], headcoeff, overwrite=True)
+        pf.writeto(outname, data[i], header, overwrite=True)
 
         outname = os.path.join(outDir, prefixOut + "coeff" + str(i) + "err.fits")
-        pf.writeto(outname, dataerr[i], headcoeff, overwrite=True)
+        pf.writeto(outname, dataerr[i], header, overwrite=True)
 
 
 def savefits_model_I(data, header, outDir, prefixOut):
@@ -422,10 +420,12 @@ def fit_spectra_I(
     outs["I"] = pixImodel.astype("float32")
     outs["coeffs"] = pixFitDict["p"].astype("float32")
     outs["coeffs_err"] = pixFitDict["perror"].astype("float32")
-    outs["chiSq"] = pixFitDict["chiSq"].astype("float32")
-    outs["chiSqRed"] = pixFitDict["chiSqRed"].astype("float32")
+    outs["chiSq"] = pixFitDict["chiSq"]
+    outs["chiSqRed"] = pixFitDict["chiSqRed"]
     outs["nIter"] = pixFitDict["nIter"]
-    outs["AIC"] = pixFitDict["AIC"].astype("float32")
+    outs["AIC"] = pixFitDict["AIC"]
+    outs["covar"] = pixFitDict["pcov"]
+    outs["reference_frequency_Hz"] = pixFitDict["reference_frequency_Hz"]
 
     return outs
 
@@ -501,7 +501,11 @@ def make_model_I(
 
     coeffs = np.array([mskArr] * 6)
     coeffs_error = np.array([mskArr] * 6)
+    reffreq = np.squeeze(np.array([mskArr]))
+
+    covars = np.array([[mskArr] * 6] * 6)
     datacube = np.squeeze(datacube)
+
     # Select only the spectra with emission
     srcData = np.rot90(datacube[:, mskSrc > 0])
 
@@ -547,36 +551,72 @@ def make_model_I(
     for _, an in enumerate(xy):
         i, x, y = an
 
-        modelIcube[:, x, y] = results[_]["I"]
+        modelIcube[:, x, y] = results[i]["I"]
+        reffreq[x, y] = results[i]["reference_frequency_Hz"]
+        covars[:, :, x, y] = results[i]["covar"]
 
         for k, j, l in zip(
-            range(len(coeffs)), results[_]["coeffs"], results[_]["coeffs_err"]
+            range(len(coeffs)), results[i]["coeffs"], results[i]["coeffs_err"]
         ):
             coeffs[5 - k, x, y] = j
             coeffs_error[5 - k, x, y] = l
 
-    header["HISTORY"] = "Stokes I model fitted by RM-Tools"
+    headcoeff["HISTORY"] = "Stokes I model fitted by RM-Tools"
     if polyOrd < 0:
-        header["HISTORY"] = (
+        headcoeff["HISTORY"] = (
             f"Fit model is dynamic order {fit_function}-polynomial, max order {-polyOrd}"
         )
     else:
-        header["HISTORY"] = f"Fit model is {polyOrd}-order {fit_function}-polynomial"
+        headcoeff["HISTORY"] = f"Fit model is {polyOrd}-order {fit_function}-polynomial"
 
-    print("Saving mask image.")
-    savefits_mask(data=mskSrc, header=header, outDir=outDir, prefixOut=prefixOut)
+    if verbose:
+        print("Saving mask image.")
+    savefits_mask(data=mskSrc, header=headcoeff, outDir=outDir, prefixOut=prefixOut)
 
-    print("Saving model I coefficients.")
+    if verbose:
+        print("Saving model I coefficients.")
     savefits_Coeffs(
         data=coeffs,
         dataerr=coeffs_error,
-        header=header,
+        header=headcoeff,
         polyOrd=polyOrd,
         outDir=outDir,
         prefixOut=prefixOut,
     )
 
-    print("Saving model I cube image. ")
+    # Save frequency map
+    head_freq = headcoeff.copy()
+    head_freq["BUNIT"] = "Hz"
+    if "BTYPE" in headcoeff:
+        del headcoeff["BTYPE"]
+
+    outname = os.path.join(outDir, prefixOut + "reffreq.fits")
+    pf.writeto(outname, reffreq, head_freq, overwrite=True)
+
+    # Save covariance maps -- these are necessary if/when converting the model
+    # reference frequency.
+    # Structure will be a single file as a 4D cube, with the 3rd and 4th dimensions
+    # iterating over the two axes of the covariance matrix.
+    head_covar = headcoeff.copy()
+    head_covar["NAXIS"] = 4
+    head_covar["NAXIS3"] = 6
+    head_covar["NAXIS4"] = 6
+    head_covar["CTYPE3"] = "INDEX"
+    head_covar["CTYPE4"] = "INDEX"
+    head_covar["CRVAL3"] = 0
+    head_covar["CRVAL4"] = 0
+    head_covar["CDELT3"] = 1
+    head_covar["CDELT4"] = 1
+    head_covar["CRPIX3"] = 1
+    head_covar["CRPIX4"] = 1
+    head_covar["CUNIT3"] = ""
+    head_covar["CUNIT4"] = ""
+
+    outname = os.path.join(outDir, prefixOut + "covariance.fits")
+    pf.writeto(outname, covars, head_covar, overwrite=True)
+
+    if verbose:
+        print("Saving model I cube image. ")
     savefits_model_I(data=modelIcube, header=header, outDir=outDir, prefixOut=prefixOut)
 
     np.savetxt(os.path.join(outDir, prefixOut + "noise.dat"), rms_Arr)