def create_IQU_cube_data(dataPath, inCatFile, startFreq_Hz, endFreq_Hz, nChans, rmsNoise_mJy, beamMinFWHM_deg, beamMajFWHM_deg, beamPA_deg, pixScale_deg, xCent_deg, yCent_deg, nPixX, nPixY, coordSys="EQU", noiseTmpArr=None, flagRanges_Hz=[]): """ Create a set of Stokes I Q & U data-cubes containing polarised sources. """ # Sample frequency space freqArr_Hz = np.linspace(startFreq_Hz, endFreq_Hz, nChans) freqNoFlgArr_Hz = freqArr_Hz.copy() dFreq_Hz = (endFreq_Hz - startFreq_Hz) / (nChans - 1) print("\nSampling frequencies %.2f - %.2f MHz by %.2f MHz." % \ (freqArr_Hz[0]/1e6, freqArr_Hz[-1]/1e6, dFreq_Hz/1e6)) if len(flagRanges_Hz) > 0: print("Flagging frequency ranges:") print("> ", flagRanges_Hz) for i in range(len(freqArr_Hz)): for fRng in flagRanges_Hz: if freqArr_Hz[i] >= fRng[0] and freqArr_Hz[i] <= fRng[1]: freqArr_Hz[i] = np.nan # Create normalised noise array from a template or assume all ones. if noiseTmpArr is None: print("Assuming flat noise versus frequency curve.") noiseArr = np.ones(freqArr_Hz.shape, dtype="f4") else: print("Scaling noise curve by external template.") xp = noiseTmpArr[0] yp = noiseTmpArr[1] mDict = calc_stats(yp) yp /= mDict["median"] noiseArr = extrap(freqArr_Hz, xp, yp) # Check the catalogue file exists if not os.path.exists(inCatFile): print("Err: File does not exist '%s'." % inCatFile) sys.exit() catInLst = csv_read_to_list(inCatFile, doFloat=True) print("Found %d entries in the catalogue." % len(catInLst)) # Create the output directory path dirs = dataPath.rstrip("/").split("/") for i in range(1, len(dirs)): dirStr = "/".join(dirs[:i]) if not os.path.exists(dirStr): os.mkdir(dirStr) if os.path.exists(dataPath): print("\n ", end=' ') print("*** WARNING ***" * 5) print(" About to delete existing data directory!", end=' ') print("Previous results will be deleted.\n") print("Press <RETURN> to continue ...", end=' ') input() shutil.rmtree(dataPath, True) os.mkdir(dataPath) # Create simple HDUs in memmory print("Creating test data in memory.") hduI = create_simple_fits_hdu(shape=(1, nChans, nPixY, nPixX), freq_Hz=startFreq_Hz, dFreq_Hz=dFreq_Hz, xCent_deg=xCent_deg, yCent_deg=yCent_deg, beamMinFWHM_deg=beamMinFWHM_deg, beamMajFWHM_deg=beamMajFWHM_deg, beamPA_deg=beamPA_deg, pixScale_deg=pixScale_deg, stokes="I", system=coordSys) head2D = strip_fits_dims(header=hduI.header, minDim=2, forceCheckDims=4) wcs2D = pw.WCS(head2D) shape2D = (nPixY, nPixX) hduQ = hduI.copy() hduQ.header["CRVAL4"] = 2.0 hduU = hduI.copy() hduU.header["CRVAL4"] = 3.0 # Calculate some beam parameters gfactor = 2.0 * m.sqrt(2.0 * m.log(2.0)) beamMinSigma_deg = beamMinFWHM_deg / gfactor beamMajSigma_deg = beamMajFWHM_deg / gfactor beamMinSigma_pix = beamMinSigma_deg / pixScale_deg beamMajSigma_pix = beamMajSigma_deg / pixScale_deg beamPA_rad = m.radians(beamPA_deg) # Loop through the sources, calculate the spectra and pixel position spectraILst = [] spectraQLst = [] spectraULst = [] coordLst_deg = [] coordLst_pix = [] successCount = 0 for i in range(len(catInLst)): e = catInLst[i] modelType = int(e[0]) # Type 1 = multiple Burn depolarisation affected components if modelType == 1: # Parse the parameters of multiple components preLst, parmArr = split_repeat_lst(e[1:], 7, 4) # Create the model spectra from multiple thin components # modified by external depolarisation IArr_Jy, QArr_Jy, UArr_Jy = \ create_IQU_spectra_burn(freqArr_Hz = freqArr_Hz, fluxI = preLst[5]/1e3, # mJy->Jy SI = preLst[6], fracPolArr = parmArr[0], psi0Arr_deg = parmArr[1], RMArr_radm2 = parmArr[2], sigmaRMArr_radm2 = parmArr[3], freq0_Hz = freq0_Hz) # Type 2 = multiple internal depolarisation affected components elif modelType == 2: # Parse the parameters of multiple components preLst, parmArr = split_repeat_lst(e[1:], 7, 3) # Create the model spectra from multiple components # modified by internal Faraday depolarisation IArr_Jy, QArr_Jy, UArr_Jy = \ create_IQU_spectra_diff(freqArr_Hz = freqArr_Hz, fluxI = preLst[5]/1e3, # mJy->Jy SI = preLst[6], fracPolArr = parmArr[0], psi0Arr_deg = parmArr[1], RMArr_radm2 = parmArr[2], freq0_Hz = freq0_Hz) else: continue spectraILst.append(IArr_Jy) spectraQLst.append(QArr_Jy) spectraULst.append(UArr_Jy) coordLst_deg.append([preLst[0], preLst[1]]) [(x_pix, y_pix)] = wcs2D.wcs_world2pix([(preLst[0], preLst[1])], 0) coordLst_pix.append([x_pix, y_pix]) successCount += 1 # Loop through the frequency channels & insert the IQU planes print("Looping through %d frequency channels:" % nChans) progress(40, 0.0) for iChan in range(len(freqArr_Hz)): progress(40, (100.0 * (iChan + 1) / nChans)) for iSrc in range(len(spectraILst)): params = [ spectraILst[iSrc][iChan], # amplitude coordLst_pix[iSrc][0], # X centre (pix) coordLst_pix[iSrc][1], # Y centre beamMinSigma_pix, # width (sigma) beamMajSigma_pix, # height (sigma) beamPA_rad ] # PA (rad) W of N (clockwise) planeI = twodgaussian(params, shape2D).reshape((nPixY, nPixX)) params[0] = spectraQLst[iSrc][iChan] planeQ = twodgaussian(params, shape2D).reshape((nPixY, nPixX)) params[0] = spectraULst[iSrc][iChan] planeU = twodgaussian(params, shape2D).reshape((nPixY, nPixX)) hduI.data[0, iChan, :, :] += planeI hduQ.data[0, iChan, :, :] += planeQ hduU.data[0, iChan, :, :] += planeU # Add the noise rmsNoise_Jy = rmsNoise_mJy / 1e3 hduI.data[0, iChan, :, :] += ( np.random.normal(scale=rmsNoise_Jy, size=(nPixY, nPixX)) * noiseArr[iChan]) hduQ.data[0, iChan, :, :] += ( np.random.normal(scale=rmsNoise_Jy, size=(nPixY, nPixX)) * noiseArr[iChan]) hduU.data[0, iChan, :, :] += ( np.random.normal(scale=rmsNoise_Jy, size=(nPixY, nPixX)) * noiseArr[iChan]) # DEBUG if False: # Mask a sub=cube hduI.data[0, 20:50, 20:40, 20:40] = np.nan hduQ.data[0, 20:50, 20:40, 20:40] = np.nan hduU.data[0, 20:50, 20:40, 20:40] = np.nan # Mask a sub=cube hduI.data[0, 60:90, 50:70, 50:70] = np.nan hduQ.data[0, 60:90, 50:70, 50:70] = np.nan hduU.data[0, 60:90, 50:70, 50:70] = np.nan if False: # Mask full planes hduI.data[0, 50:60, :, :] = np.nan hduQ.data[0, 50:60, :, :] = np.nan hduU.data[0, 50:60, :, :] = np.nan if False: # Mask full spectra hduI.data[0, :, 20:60, 20:70] = np.nan hduQ.data[0, :, 20:60, 20:70] = np.nan hduU.data[0, :, 20:60, 20:70] = np.nan # Write to the FITS files print("Saving the FITS files to disk in directory '%s'" % dataPath) sys.stdout.flush() fitsFileOut = dataPath + "/StokesI.fits" print("> %s" % fitsFileOut) hduI.writeto(fitsFileOut, clobber=True) fitsFileOut = dataPath + "/StokesQ.fits" print("> %s" % fitsFileOut) hduQ.writeto(fitsFileOut, clobber=True) fitsFileOut = dataPath + "/StokesU.fits" print("> %s" % fitsFileOut) hduU.writeto(fitsFileOut, clobber=True) # Save a vector of frequency values freqFileOut = dataPath + "/freqs_Hz.dat" print("> %s" % freqFileOut) np.savetxt(freqFileOut, freqNoFlgArr_Hz) return successCount
def create_IQU_ascii_data(dataPath, inCatFile, startFreq_Hz, endFreq_Hz, nChans, rmsNoise_mJy, noiseTmpArr=None, flagRanges_Hz=[], freq0_Hz=None): """ Create a set of ASCII files containing Stokes I Q & U spectra. """ # Sample frequency space freqArr_Hz = np.linspace(startFreq_Hz, endFreq_Hz, nChans) freqNoFlgArr_Hz = freqArr_Hz.copy() dFreq_Hz = (endFreq_Hz - startFreq_Hz)/ (nChans-1) print("\nSampling frequencies %.2f - %.2f MHz by %.2f MHz." % \ (freqArr_Hz[0]/1e6, freqArr_Hz[-1]/1e6, dFreq_Hz/1e6)) if len(flagRanges_Hz)>0: print("Flagging frequency ranges:") print("> ", flagRanges_Hz) for i in range(len(freqArr_Hz)): for fRng in flagRanges_Hz: if freqArr_Hz[i]>=fRng[0] and freqArr_Hz[i]<=fRng[1]: freqArr_Hz[i]=np.nan # Create normalised noise array from a template or assume all ones. print("Input RMS noise is %.3g mJy" % rmsNoise_mJy) if noiseTmpArr is None: print("Assuming flat noise versus frequency curve.") noiseArr = np.ones(freqArr_Hz.shape, dtype="f8") else: print("Scaling noise by external template curve.") xp = noiseTmpArr[0] yp = noiseTmpArr[1] mDict = calc_stats(yp) yp /= mDict["median"] noiseArr = extrap(freqArr_Hz, xp, yp) # Check the catalogue file exists if not os.path.exists(inCatFile): print("Err: File does not exist '%s'." % inCatFile) sys.exit() catInLst = csv_read_to_list(inCatFile, doFloat=True) print("Found %d entries in the catalogue." % len(catInLst)) # Set the frequency at which the flux has been defined if freq0_Hz is None: freq0_Hz = startFreq_Hz print("Catalogue flux is defined at %.3f MHz" % (freq0_Hz/1e6)) # Create the output directory path dataPath = dataPath.rstrip("/") print("Creating test dataset in '%s/'" % dataPath) dirs = dataPath.split("/") for i in range(1, len(dirs)): dirStr = "/".join(dirs[:i]) if not os.path.exists(dirStr): os.mkdir(dirStr) if os.path.exists(dataPath): print("\n ", end=' ') print("*** WARNING ***" *5) print(" About to delete existing data directory!", end=' ') print("Previous results will be deleted.\n") print("Press <RETURN> to continue ...", end=' ') input() shutil.rmtree(dataPath, True) os.mkdir(dataPath) # Loop through the sources, calculate the spectra and save to disk successCount = 0 for i in range(len(catInLst)): e = catInLst[i] modelType = int(e[0]) # Type 1 = multiple Burn depolarisation affected components if modelType==1: # Parse the parameters of multiple components preLst, parmArr = split_repeat_lst(e[1:],7,4) # Create the model spectra from multiple thin components # modified by external depolarisation IArr_Jy, QArr_Jy, UArr_Jy = \ create_IQU_spectra_burn(freqArr_Hz = freqArr_Hz, fluxI = preLst[5]/1e3, # mJy->Jy SI = preLst[6], fracPolArr = parmArr[0], psi0Arr_deg = parmArr[1], RMArr_radm2 = parmArr[2], sigmaRMArr_radm2 = parmArr[3], freq0_Hz = freq0_Hz) # Type 2 = multiple internal depolarisation affected components elif modelType==2: # Parse the parameters of multiple components preLst, parmArr = split_repeat_lst(e[1:],7,3) # Create the model spectra from multiple components # modified by internal Faraday depolarisation IArr_Jy, QArr_Jy, UArr_Jy = \ create_IQU_spectra_diff(freqArr_Hz = freqArr_Hz, fluxI = preLst[5]/1e3, # mJy->Jy SI = preLst[6], fracPolArr = parmArr[0], psi0Arr_deg = parmArr[1], RMArr_radm2 = parmArr[2], freq0_Hz = freq0_Hz) else: continue # Add scatter to the data to simulate noise rmsNoise_Jy = rmsNoise_mJy/1e3 IArr_Jy += (np.random.normal(scale=rmsNoise_Jy, size=IArr_Jy.shape) * noiseArr) QArr_Jy += (np.random.normal(scale=rmsNoise_Jy, size=QArr_Jy.shape) * noiseArr) UArr_Jy += (np.random.normal(scale=rmsNoise_Jy, size=UArr_Jy.shape) * noiseArr) dIArr_Jy = noiseArr * rmsNoise_Jy dIArr_Jy *= np.random.normal(loc=1.0, scale=0.05, size=noiseArr.shape) dQArr_Jy = noiseArr * rmsNoise_Jy dQArr_Jy *= np.random.normal(loc=1.0, scale=0.05, size=noiseArr.shape) dUArr_Jy = noiseArr * rmsNoise_Jy dUArr_Jy *= np.random.normal(loc=1.0, scale=0.05, size=noiseArr.shape) # Save spectra to disk outFileName = "Source%d.dat" % (i+1) outFilePath = dataPath + "/" + outFileName print("> Writing ASCII file '%s' ..." % outFileName, end=' ') np.savetxt(outFilePath, np.column_stack((freqArr_Hz, IArr_Jy, QArr_Jy, UArr_Jy, dIArr_Jy, dQArr_Jy, dUArr_Jy))) print("done.") successCount += 1 return successCount