Skip to content

Commit

Permalink
Improved FITS headers of 3D tools output.
Browse files Browse the repository at this point in the history
  • Loading branch information
@Cameron-Van-Eck
Cameron-Van-Eck committed Nov 21, 2023
1 parent 6b98aec commit efac24d
Show file tree
Hide file tree
Showing 4 changed files with 158 additions and 28 deletions.
7 changes: 5 additions & 2 deletions RMtools_3D/RMpeakfit_3D.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,15 +160,18 @@ def save_maps(map_dict, prefix_path,FDFheader):
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['HISTORY'] = "Polarization peak maps created with RM-Tools RMpeakfit_3D"

#Set flux unit from FITS header if possible
if 'BUNIT' in FDFheader:
flux_unit=FDFheader['BUNIT']
else:
flux_unit=''



#Dictionary of units for peak fitting output parameters (for FITS headers)
unit_dict={"dFDFcorMAD": flux_unit, "dFDFrms": flux_unit,
"phiPeakPIchan_rm2": 'rad/m^2', "dPhiPeakPIchan_rm2": 'rad/m^2',
Expand Down
70 changes: 54 additions & 16 deletions RMtools_3D/do_RMclean_3D.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@

from RMutils.util_RM import do_rmclean_hogbom
from RMutils.util_RM import fits_make_lin_axis
from RMtools_3D.do_RMsynth_3D import _setStokes

C = 2.997924538e8 # Speed of light [m/s]

Expand Down Expand Up @@ -190,54 +191,84 @@ def writefits(cleanFDF, ccArr, iterCountArr, residFDF, headtemp, nBits=32,
# Default data types
dtFloat = "float" + str(nBits)
dtComplex = "complex" + str(2*nBits)
header=headtemp.copy()

header['HISTORY']="RM-CLEAN3D output from RM-Tools"



if outDir=='': #To prevent code breaking if file is in current directory
outDir='.'
# Save the clean FDF
if not write_separate_FDF:
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(cleanFDF.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.ImageHDU(cleanFDF.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.ImageHDU(np.abs(cleanFDF).astype(dtFloat), header)


fitsFileOut = outDir + "/" + prefixOut + "FDF_clean.fits"
if(verbose): log("> %s" % fitsFileOut)
hdu0 = pf.PrimaryHDU(cleanFDF.real.astype(dtFloat), headtemp)
hdu1 = pf.ImageHDU(cleanFDF.imag.astype(dtFloat), headtemp)
hdu2 = pf.ImageHDU(np.abs(cleanFDF).astype(dtFloat), headtemp)
hduLst = pf.HDUList([hdu0, hdu1, hdu2])
hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True)
hduLst.close()
else:
hdu0 = pf.PrimaryHDU(cleanFDF.real.astype(dtFloat), headtemp)

header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(cleanFDF.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.PrimaryHDU(cleanFDF.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.PrimaryHDU(np.abs(cleanFDF).astype(dtFloat), header)

fitsFileOut = outDir + "/" + prefixOut + "FDF_clean_real.fits"
hdu0.writeto(fitsFileOut, output_verify="fix", overwrite=True)
if (verbose): log("> %s" % fitsFileOut)
hdu1 = pf.PrimaryHDU(cleanFDF.imag.astype(dtFloat), headtemp)
fitsFileOut = outDir + "/" + prefixOut + "FDF_clean_im.fits"
hdu1.writeto(fitsFileOut, output_verify="fix", overwrite=True)
if (verbose): log("> %s" % fitsFileOut)
hdu2 = pf.PrimaryHDU(np.abs(cleanFDF).astype(dtFloat), headtemp)
fitsFileOut = outDir + "/" + prefixOut + "FDF_clean_tot.fits"
hdu2.writeto(fitsFileOut, output_verify="fix", overwrite=True)
if (verbose): log("> %s" % fitsFileOut)




if not write_separate_FDF:
#Save the complex clean components as another file.
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(ccArr.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.ImageHDU(ccArr.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.ImageHDU(np.abs(ccArr).astype(dtFloat), header)

fitsFileOut = outDir + "/" + prefixOut + "FDF_CC.fits"
if (verbose): log("> %s" % fitsFileOut)
hdu0 = pf.PrimaryHDU(ccArr.real.astype(dtFloat), headtemp)
hdu1 = pf.ImageHDU(ccArr.imag.astype(dtFloat), headtemp)
hdu2 = pf.ImageHDU(np.abs(ccArr).astype(dtFloat), headtemp)
hduLst = pf.HDUList([hdu0, hdu1, hdu2])
hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True)
hduLst.close()
else:
hdu0 = pf.PrimaryHDU(ccArr.real.astype(dtFloat), headtemp)
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(ccArr.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.PrimaryHDU(ccArr.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.PrimaryHDU(np.abs(ccArr).astype(dtFloat), header)


fitsFileOut = outDir + "/" + prefixOut + "FDF_CC_real.fits"
hdu0.writeto(fitsFileOut, output_verify="fix", overwrite=True)
if (verbose): log("> %s" % fitsFileOut)
hdu1 = pf.PrimaryHDU(ccArr.imag.astype(dtFloat), headtemp)
fitsFileOut = outDir + "/" + prefixOut + "FDF_CC_im.fits"
hdu1.writeto(fitsFileOut, output_verify="fix", overwrite=True)
if (verbose): log("> %s" % fitsFileOut)
hdu2 = pf.PrimaryHDU(np.abs(ccArr).astype(dtFloat), headtemp)
fitsFileOut = outDir + "/" + prefixOut + "FDF_CC_tot.fits"
hdu2.writeto(fitsFileOut, output_verify="fix", overwrite=True)
if (verbose): log("> %s" % fitsFileOut)
Expand All @@ -246,17 +277,24 @@ def writefits(cleanFDF, ccArr, iterCountArr, residFDF, headtemp, nBits=32,
#don't try to remove the FD axis, but just make it degenerate.

if headtemp['NAXIS'] > 2:
headtemp["NAXIS3"] = 1
header["NAXIS3"] = 1
header['CTYPE3'] = ('DEGENERATE','Axis left in to avoid FITS errors')
header['CUNIT3'] = ''
if headtemp['NAXIS'] == 4:
headtemp["NAXIS4"] = 1
header["NAXIS4"] = 1
header['CTYPE4'] = ('DEGENERATE','Axis left in to avoid FITS errors')
header['CUNIT4'] = ''


# Save the iteration count mask
fitsFileOut = outDir + "/" + prefixOut + "CLEAN_nIter.fits"
if (verbose): log("> %s" % fitsFileOut)
headtemp["BUNIT"] = "Iterations"
header["BUNIT"] = "Iterations"
if 'STOKES' in header:
del header['STOKES']
hdu0 = pf.PrimaryHDU(np.expand_dims(iterCountArr.astype(dtFloat),
axis=tuple(range(headtemp['NAXIS']-iterCountArr.ndim))),
headtemp)
header)
hduLst = pf.HDUList([hdu0])
hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True)
hduLst.close()
Expand Down
82 changes: 77 additions & 5 deletions RMtools_3D/do_RMsynth_3D.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -321,11 +321,19 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",
FDFcube=np.moveaxis(FDFcube,0,Ndim-freq_axis)
if not_rmsf is not True: RMSFcube=np.moveaxis(RMSFcube,0,Ndim-freq_axis)

if 'BUNIT' in header:
header['BUNIT']=header['BUNIT']+'/RMSF'


if(write_seperate_FDF):
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(FDFcube.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.PrimaryHDU(FDFcube.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.PrimaryHDU(np.abs(FDFcube).astype(dtFloat), header)

fitsFileOut = outDir + "/" + prefixOut + "FDF_real_dirty.fits"
if(verbose): log("> %s" % fitsFileOut)
hdu0.writeto(fitsFileOut, output_verify="fix", overwrite=True)
Expand All @@ -339,10 +347,15 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",
hdu2.writeto(fitsFileOut, output_verify="fix", overwrite=True)

else:
# Save the dirty FDF
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(FDFcube.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.ImageHDU(FDFcube.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.ImageHDU(np.abs(FDFcube).astype(dtFloat), header)

# Save the dirty FDF
fitsFileOut = outDir + "/" + prefixOut + "FDF_dirty.fits"
if(verbose): log("> %s" % fitsFileOut)
hduLst = pf.HDUList([hdu0, hdu1, hdu2])
Expand All @@ -358,24 +371,42 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",
header["CDELT"+str(freq_axis)] = (np.diff(phi2Arr_radm2)[0], '[rad/m^2] Coordinate increment at reference point')
header["CRPIX"+str(freq_axis)] = phi2Arr_radm2.size//2+1
header["CRVAL"+str(freq_axis)] = (phi2Arr_radm2[phi2Arr_radm2.size//2], '[rad/m^2] Coordinate value at reference point')
header['BUNIT']=''


header["DATAMAX"] = np.max(fwhmRMSFCube) + 1
header["DATAMIN"] = np.max(fwhmRMSFCube) - 1
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)] = phiArr_radm2[0]
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')




if(write_seperate_FDF):
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(RMSFcube.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.PrimaryHDU(RMSFcube.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.PrimaryHDU(np.abs(RMSFcube).astype(dtFloat), header)
hdu3 = pf.PrimaryHDU(np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0),
rmheader)


fitsFileOut = outDir + "/" + prefixOut + "RMSF_real.fits"
if(verbose): log("> %s" % fitsFileOut)
hdu0.writeto(fitsFileOut, output_verify="fix", overwrite=True)
Expand All @@ -393,11 +424,18 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",
hdu3.writeto(fitsFileOut, output_verify="fix", overwrite=True)

else:
fitsFileOut = outDir + "/" + prefixOut + "RMSF.fits"
header=_setStokes(header, 'Q')
hdu0 = pf.PrimaryHDU(RMSFcube.real.astype(dtFloat), header)
header=_setStokes(header, 'U')
hdu1 = pf.ImageHDU(RMSFcube.imag.astype(dtFloat), header)
header=_setStokes(header, 'PI') #Sets Stokes axis to zero, which is a non-standard value.
del header['STOKES']
hdu2 = pf.ImageHDU(np.abs(RMSFcube).astype(dtFloat), header)
hdu3 = pf.ImageHDU(fwhmRMSFCube.astype(dtFloat), rmheader)
hdu3 = pf.ImageHDU(np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0),
rmheader)


fitsFileOut = outDir + "/" + prefixOut + "RMSF.fits"
hduLst = pf.HDUList([hdu0, hdu1, hdu2, hdu3])
if(verbose): log("> %s" % fitsFileOut)
hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True)
Expand All @@ -419,6 +457,19 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",

maxPI,peakRM = create_peak_maps(FDFcube,phiArr_radm2,Ndim-freq_axis)
# Save a maximum polarised intensity map
header['BUNIT']=headtemplate['BUNIT']
header["NAXIS"+str(freq_axis)] = 1
header['CTYPE'+str(freq_axis)] = ('DEGENERATE','Axis left in to avoid FITS errors')
header['CUNIT'+str(freq_axis)] = ''

stokes_axis=None
for axis in range(1,header['NAXIS']+1):
if 'STOKES' in header[f'CTYPE{axis}']:
stokes_axis=axis
if stokes_axis is not None:
header[f'CTYPE{stokes_axis}']=('DEGENERATE','Axis left in to avoid FITS errors')


fitsFileOut = outDir + "/" + prefixOut + "FDF_maxPI.fits"
if(verbose): log("> %s" % fitsFileOut)
pf.writeto(fitsFileOut,
Expand All @@ -428,6 +479,7 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",
# Save a peak RM map
fitsFileOut = outDir + "/" + prefixOut + "FDF_peakRM.fits"
header["BUNIT"] = "rad/m^2"
header['BTYPE'] = 'FDEP'
if(verbose): log("> %s" % fitsFileOut)
pf.writeto(fitsFileOut, np.expand_dims(peakRM,axis=0), header, overwrite=True,
output_verify="fix")
Expand All @@ -444,6 +496,26 @@ def writefits(dataArr, headtemplate, fitRMSF=False, prefixOut="", outDir="",
# output_verify="fix")


def _setStokes(header,stokes):
"""Check if header has Stokes axis. If so, set to correct numerical value
(if IQUV). If not a valid Stokes parameter, sets to zero.
Adds Stokes keyword regardless. Returns new, updated header.
"""
stokes_dict={'I':1, 'Q':2, 'U':3, 'V':4}

outheader=header.copy()
stokes_axis=None
for axis in range(1,header['NAXIS']+1):
if 'STOKES' in header[f'CTYPE{axis}']:
stokes_axis=axis

if stokes_axis is not None:
outheader[f'CRPIX{stokes_axis}']=1.0
outheader[f'CRVAL{stokes_axis}']=stokes_dict.get(stokes.upper(),0)
outheader['STOKES']=stokes

return outheader

def create_peak_maps(FDFcube,phiArr_radm2,phi_axis=0):
"""Finds the location and amplitude of the highest peak in the FDF (pixelwise)
and returns maps of those parameters. Does not fit the peak, only finds
Expand Down
27 changes: 22 additions & 5 deletions RMtools_3D/do_fitIcube.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -267,6 +267,9 @@ def cube_noise(datacube, header, freqArr_Hz, threshold=-5):
return rmsArr, mskSrc





def savefits_mask(data, header, outDir, prefixOut):

""" Save the derived mask to a fits file
Expand All @@ -283,7 +286,7 @@ def savefits_mask(data, header, outDir, prefixOut):
del headMask["BUNIT"]

mskArr = np.where(data > 0, 1.0, np.nan)
MaskfitsFile = os.path.join(outDir,prefixOut + "_mask.fits")
MaskfitsFile = os.path.join(outDir,prefixOut + "mask.fits")
print("> %s" % MaskfitsFile)
pf.writeto(MaskfitsFile, mskArr, headMask, output_verify="fix",
overwrite=True)
Expand All @@ -302,13 +305,15 @@ def savefits_Coeffs(data, dataerr, header, polyOrd, outDir, prefixOut):
"""

headcoeff = strip_fits_dims(header=header, minDim=2)
del headcoeff["BUNIT"]
headcoeff["BUNIT"]=''
if 'BTYPE' in headcoeff:
del headcoeff['BTYPE']

for i in range(np.abs(polyOrd)+1):
outname = os.path.join(outDir,prefixOut + '_coeff'+str(i) + '.fits')
outname = os.path.join(outDir,prefixOut + 'coeff'+str(i) + '.fits')
pf.writeto(outname, data[i], headcoeff, overwrite=True)

outname = os.path.join(outDir,prefixOut + '_coeff'+str(i) + 'err.fits')
outname = os.path.join(outDir,prefixOut + 'coeff'+str(i) + 'err.fits')
pf.writeto(outname, dataerr[i], headcoeff, overwrite=True)


Expand Down Expand Up @@ -339,7 +344,7 @@ def savefits_model_I(data, header, outDir, prefixOut):
while len(headModelCube) < (36 * 4 - 1):
headModelCube.append()

fitsModelFile = os.path.join(outDir ,prefixOut + "_model.i.fits")
fitsModelFile = os.path.join(outDir ,prefixOut + "model.i.fits")
headModelCube.tofile(fitsModelFile, overwrite=True)
with open(fitsModelFile, "rb+") as f:
f.seek(len(headModelCube.tostring()) + (nVoxels*int(nBits/8)) - 1)
Expand Down Expand Up @@ -497,6 +502,13 @@ def make_model_I(datacube, header, freqArr_Hz, polyOrd=2,
coeffs_error[5-k,x,y] = l


header['HISTORY'] = "Stokes I model fitted by RM-Tools"
if polyOrd < 0:
header['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"


print("Saving mask image.")
savefits_mask(data=mskSrc, header=header, outDir=outDir, prefixOut=prefixOut)

Expand All @@ -508,6 +520,11 @@ def make_model_I(datacube, header, freqArr_Hz, polyOrd=2,
savefits_model_I(data=modelIcube, header=header,
outDir=outDir, prefixOut=prefixOut)


np.savetxt(os.path.join(outDir, prefixOut + "noise.dat"), rms_Arr)



return modelIcube


Expand Down

0 comments on commit efac24d

Please sign in to comment.