Fix imodelrescale, issue #142

RMtools_3D/RMpeakfit_3D.py

@@ -13,10 +13,10 @@
 
 import astropy.io.fits as pf
 import numpy as np
-from tqdm.auto import tqdm, trange
+from tqdm.auto import trange
 
 from RMtools_3D.do_RMsynth_3D import readFitsCube, readFreqFile
-from RMutils.util_misc import interp_images
+from RMutils.util_misc import interp_images, remove_header_third_fourth_axis
 from RMutils.util_RM import fits_make_lin_axis, measure_FDF_parms
 
 C = 2.997924538e8  # Speed of light [m/s]
@@ -172,14 +172,9 @@ def save_maps(map_dict, prefix_path, FDFheader):
     """
     # Set up generic FITS header
     product_header = FDFheader.copy()
-    product_header["NAXIS"] = 2
     # Remove extra axes:
-    if "NAXIS3" in product_header:
-        delete_FITSheader_axis(product_header, 3)
-    if "NAXIS4" in product_header:
-        delete_FITSheader_axis(product_header, 4)
-    if "STOKES" in product_header:
-        del product_header["STOKES"]
+    product_header = remove_header_third_fourth_axis(product_header)
+
     product_header["HISTORY"] = (
         "Polarization peak maps created with RM-Tools RMpeakfit_3D"
     )

RMtools_3D/do_RMclean_3D.py

@@ -455,7 +455,7 @@ def main():
         "fitsRMSF",
         metavar="RMSF.fits",
         nargs=1,
-        help="FITS cube containing the RMSF and FWHM image.\n(Cans be any of the RMSF output cubes from do_RMsynth_3D.py)",
+        help="FITS cube containing the RMSF and FWHM image.\n(Cans be any of the RMSF output cubes (so not _FWHM.fits!) from do_RMsynth_3D.py)",
     )
     parser.add_argument(
         "-c",

RMtools_3D/do_RMsynth_3D.py

@@ -44,7 +44,7 @@
 import astropy.table as at
 import numpy as np
 
-from RMutils.util_misc import interp_images
+from RMutils.util_misc import interp_images, remove_header_third_fourth_axis
 from RMutils.util_RM import do_rmsynth_planes, get_rmsf_planes
 
 if sys.version_info.major == 2:
@@ -474,9 +474,11 @@ def writefits(
             hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True)
             hduLst.close()
 
-    # Save the RMSF
+    # Save the RMSF: real, im, tot (4D) and FWHM (2D)
     if not not_rmsf:
-        # Put frequency axis first, and reshape to add degenerate axes:
+        # Create header for RMSF_{real,imag,tot}.fits
+
+        # Put frequency axis first, and reshape to add degenerate Stokes axis:
         RMSFcube = np.reshape(RMSFcube, [RMSFcube.shape[0]] + output_axes)
         # Move Faraday depth axis to appropriate position to match header.
         RMSFcube = np.moveaxis(RMSFcube, 0, Ndim - freq_axis)
@@ -494,29 +496,12 @@ def writefits(
         )
         header["BUNIT"] = ""
 
-        rmheader = header.copy()
-        rmheader["BUNIT"] = "rad/m^2"
-        if "BTYPE" in rmheader:
-            del rmheader["BTYPE"]
-        # Because there can be problems with different axes having different FITS keywords,
-        # don't try to remove the FD axis, but just make it degenerate.
-        # Also requires np.expand_dims to set the correct NAXIS.
-        rmheader["NAXIS" + str(freq_axis)] = 1
-        rmheader["CTYPE" + str(freq_axis)] = (
-            "DEGENERATE",
-            "Axis left in to avoid FITS errors",
-        )
-        rmheader["CUNIT" + str(freq_axis)] = ""
-        rmheader["CRVAL" + str(freq_axis)] = 0  # doesnt mean anything
-        stokes_axis = None
-        for axis in range(1, rmheader["NAXIS"] + 1):
-            if "STOKES" in rmheader[f"CTYPE{axis}"]:
-                stokes_axis = axis
-        if stokes_axis is not None:
-            rmheader[f"CTYPE{stokes_axis}"] = (
-                "DEGENERATE",
-                "Axis left in to avoid FITS errors",
-            )
+        # Create header for RMSF_FHWM.fits
+        rmsffwhm_header = header.copy()
+        rmsffwhm_header["BUNIT"] = "rad/m^2"
+        rmsffwhm_header.pop("BTYPE", None)
+        # Remove 3rd and 4th axis, RMSF_FWHM is a plane, like the 2D peak maps
+        rmsffwhm_header = remove_header_third_fourth_axis(rmsffwhm_header)
 
         if write_seperate_FDF:  # more memory efficient as well
             header = _setStokes(header, "Q")
@@ -550,7 +535,9 @@ def writefits(
             gc.collect()
 
             hdu3 = pf.PrimaryHDU(
-                np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0), rmheader
+                # np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0), rmsffwhm_header
+                fwhmRMSFCube.astype(dtFloat),
+                rmsffwhm_header,
             )
             fitsFileOut = outDir + "/" + prefixOut + "RMSF_FWHM.fits"
             if verbose:
@@ -568,7 +555,9 @@ def writefits(
             del header["STOKES"]
             hdu2 = pf.ImageHDU(np.abs(RMSFcube).astype(dtFloat), header)
             hdu3 = pf.ImageHDU(
-                np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0), rmheader
+                # np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0), rmsffwhm_header
+                fwhmRMSFCube.astype(dtFloat),
+                rmsffwhm_header,
             )
 
             fitsFileOut = outDir + "/" + prefixOut + "RMSF.fits"
@@ -588,6 +577,11 @@ def writefits(
         # Also requires np.expand_dims to set the correct NAXIS.
         # Generate peak maps:
 
+        ## Erik: THIS NOW INCONSISTENT WITH RMSF_FWHM (2D). PEAK MAPS STILL USE np.expand_dims
+        ## but since do_peakmaps shouldnt be used anyways, I did not update the code below
+        ## it still produces OK data though, just inconsistent in dimensions
+        log("WARNING: dimensions of these peak maps are not 2D")
+
         maxPI, peakRM = create_peak_maps(FDFcube, phiArr_radm2, Ndim - freq_axis)
         # Save a maximum polarised intensity map
         if "BUNIT" in headtemplate:

RMtools_3D/do_fitIcube.py

@@ -48,7 +48,12 @@
 from RMtools_3D.do_RMsynth_3D import readFitsCube
 from RMtools_3D.make_freq_file import get_freq_array
 from RMutils.util_FITS import strip_fits_dims
-from RMutils.util_misc import MAD, calculate_StokesI_model, fit_StokesI_model
+from RMutils.util_misc import (
+    MAD,
+    calculate_StokesI_model,
+    fit_StokesI_model,
+    remove_header_third_fourth_axis,
+)
 
 # -----------------------------------------------------------------------------#
 
@@ -288,52 +293,60 @@ def cube_noise(datacube, header, freqArr_Hz, threshold=-5):
     # start = time.time()
     for i in range(nChan):
         dataPlane = datacube[i]
-        if threshold > 0:
-            idxSky = np.where(dataPlane < threshold)  # replaced cutoff with threshold
-        else:
-            idxSky = np.where(dataPlane)
-
-        # Pass 1
-        rmsTmp = MAD(dataPlane[idxSky])
-        medTmp = np.nanmedian(dataPlane[idxSky])
+        if np.isnan(dataPlane).all():
+            # If this plane is fully flagged, dont have to calculate
+            medSky[i] = np.nan
+            rmsArr[i] = np.nan
 
-        # Pass 2: use a fixed 3-sigma cutoff to mask off emission
-        idxSky = np.where(dataPlane < medTmp + rmsTmp * 3)
-        medSky[i] = np.nanmedian(dataPlane[idxSky])
-        rmsArr[i] = MAD(dataPlane[idxSky])
-
-        # When building final emission mask treat +ve threshold as absolute
-        # values and negative threshold as sigma values
-        if threshold > 0:
-            idxSrc = np.where(dataPlane > threshold)
         else:
-            idxSrc = np.where(dataPlane > medSky[i] - 1 * rmsArr[i] * threshold)
-
-        mskSrc[idxSrc] += 1
+            if threshold > 0:
+                idxSky = np.where(
+                    dataPlane < threshold
+                )  # replaced cutoff with threshold
+            else:
+                idxSky = np.where(dataPlane)
+
+            # Pass 1
+            rmsTmp = MAD(dataPlane[idxSky])
+            medTmp = np.nanmedian(dataPlane[idxSky])
+
+            # Pass 2: use a fixed 3-sigma cutoff to mask off emission
+            idxSky = np.where(dataPlane < medTmp + rmsTmp * 3)
+            medSky[i] = np.nanmedian(dataPlane[idxSky])
+            rmsArr[i] = MAD(dataPlane[idxSky])
+
+            # When building final emission mask treat +ve threshold as absolute
+            # values and negative threshold as sigma values
+            if threshold > 0:
+                idxSrc = np.where(dataPlane > threshold)
+            else:
+                idxSrc = np.where(dataPlane > medSky[i] - 1 * rmsArr[i] * threshold)
+
+            mskSrc[idxSrc] += 1
 
     # end = time.time()
     # print(' For loop masking takes %.3fs'%(end-start))
     return rmsArr, mskSrc
 
 
-def savefits_mask(data, header, outDir, prefixOut):
+def savefits_mask(data, header, outDir, prefixOut, dtFloat):
     """Save the derived mask to a fits file
 
     data:  2D data defining the mask.
     header: header to describe the mask
     outDir: directory to save the mask fits data
     prefixOut: prefix to use on the output name
+    dtFloat: type to use for output file
     """
 
-    headMask = strip_fits_dims(header=header, minDim=2)
+    headMask = remove_header_third_fourth_axis(header=header)
     headMask["DATAMAX"] = 1
     headMask["DATAMIN"] = 0
     if "BUNIT" in headMask:
         del headMask["BUNIT"]
 
-    mskArr = np.where(data > 0, 1.0, np.nan)
+    mskArr = np.where(data > 0, 1.0, np.nan).astype(dtFloat)
     MaskfitsFile = os.path.join(outDir, prefixOut + "mask.fits")
-    print("> %s" % MaskfitsFile)
     pf.writeto(MaskfitsFile, mskArr, headMask, output_verify="fix", overwrite=True)
 
 
@@ -500,11 +513,11 @@ def make_model_I(
     modelIcube[:] = np.nan
     results = []
 
-    coeffs = np.array([mskArr] * 6)
-    coeffs_error = np.array([mskArr] * 6)
-    reffreq = np.squeeze(np.array([mskArr]))
+    coeffs = np.array([mskArr] * 6, dtype=dtFloat)
+    coeffs_error = np.array([mskArr] * 6, dtype=dtFloat)
+    reffreq = np.squeeze(np.array([mskArr], dtype=dtFloat))
 
-    covars = np.array([[mskArr] * 6] * 6)
+    covars = np.array([[mskArr] * 6] * 6, dtype=dtFloat)
     datacube = np.squeeze(datacube)
 
     # Select only the spectra with emission
@@ -549,7 +562,7 @@ def make_model_I(
         with mp.Pool(processes=num_cores) as pool:
             results = pool.map(func, args_list)
 
-    headcoeff = strip_fits_dims(header=header, minDim=2)
+    headcoeff = remove_header_third_fourth_axis(header=header.copy())
     del headcoeff["BUNIT"]
 
     endTime = time.time()
@@ -579,7 +592,13 @@ def make_model_I(
 
     if verbose:
         print("Saving mask image.")
-    savefits_mask(data=mskSrc, header=headcoeff, outDir=outDir, prefixOut=prefixOut)
+    savefits_mask(
+        data=mskSrc,
+        header=headcoeff,
+        outDir=outDir,
+        prefixOut=prefixOut,
+        dtFloat=dtFloat,
+    )
 
     if verbose:
         print("Saving model I coefficients.")

RMtools_3D/rescale_I_model_3D.py

@@ -210,7 +210,7 @@ def read_data(basename):
 
     # Get coefficient maps (without knowing how many there are)
     # Reverse index order to match RM-Tools internal ordering (highest to lowest polynomial order)
-    coeffs = np.zeros((6, *old_reffreq_map.shape))
+    coeffs = np.zeros((6, *old_reffreq_map.shape), dtype=covar_map.dtype)
     for i in range(6):
         try:  # Keep trying higher orders
             data = pf.getdata(basename + f"coeff{i}.fits")
@@ -267,9 +267,9 @@ def rescale_I_model_3D(
         new_errors (3D array): maps of new parameter uncertainties (highest to lowest)
     """
 
-    # Initialize output arrays:
-    new_coeffs = np.zeros_like(coeffs)
-    new_errors = np.zeros_like(coeffs)
+    # Initialize output arrays, keep dtype consistent
+    new_coeffs = np.zeros_like(coeffs, dtype=coeffs.dtype)
+    new_errors = np.zeros_like(coeffs, dtype=coeffs.dtype)
     rs = old_reffreq_map.shape[
         1
     ]  # Get the length of a row, for array indexing later on.
@@ -321,7 +321,10 @@ def write_new_parameters(
         del out_header["BUNIT"]
 
     # Work out highest order of polynomial:
-    max_order = np.sum(np.any(new_coeffs != 0.0, axis=(1, 2))) - 1
+    # if any of the 6 possible coeff planes contain non-zero and non-nan values, it's a 'good' plane.
+    max_order = (
+        np.sum(np.any((new_coeffs != 0.0) & (~np.isnan(new_coeffs)), axis=(1, 2))) - 1
+    )
 
     for i in range(max_order + 1):
         pf.writeto(

RMutils/util_FITS.py

@@ -275,7 +275,7 @@ def get_beam_from_header(header):
         bpa = None
 
     # Or try history for AIPS CLEAN Beam
-    beamHistStr = "AIPS\s+CLEAN\sBMAJ=\s+(\S+)\s+BMIN=\s+(\S+)\s+BPA=\s+(\S+)"
+    beamHistStr = r"AIPS\s+CLEAN\sBMAJ=\s+(\S+)\s+BMIN=\s+(\S+)\s+BPA=\s+(\S+)"
     bmHistPat = re.compile(beamHistStr)
     if bmaj is None or bmin is None:
         try: