From 12e821d1ef97fed5c9d74f778fea249657365697 Mon Sep 17 00:00:00 2001 From: Erik Osinga Date: Sun, 18 Feb 2024 18:47:41 -0500 Subject: [PATCH] Decrease RAM usage of RMtools (#109) * Decrease RAM usage of RMtools This decreases the peak RAM usage by about 40-50 percent when write_separate_RMSF=True as we don't need to keep everything in memory. Additionally, it creates the option to do the RM synthesis and RMSF calculation separately instead of in the same function which decreases the RAM usage of the POSSUM Polarimetry pipeline by about 20-25 percent throughout, while not losing any noticable CPU time. To make this change, I did have to make sure lambda0Sq_m2 is reported by util_RM.py and handled correctly by both do_RMsynth_1D and do_RMsynth_3D. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Alec Thomson (S&A, Kensington WA) Co-authored-by: Cameron Van Eck --- .github/workflows/python-package.yml | 3 +- RMtools_1D/do_RMsynth_1D.py | 2 +- RMtools_3D/do_RMsynth_3D.py | 384 ++++++++++++++----------- RMutils/util_RM.py | 16 +- tests/{QA_tests.py => QA_test.py} | 84 +++--- tests/{test_nufft.py => nufft_test.py} | 2 +- 6 files changed, 276 insertions(+), 215 deletions(-) rename tests/{QA_tests.py => QA_test.py} (78%) rename tests/{test_nufft.py => nufft_test.py} (99%) diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index d8d08b3..c94d617 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -37,5 +37,4 @@ jobs: - name: Test with pytest shell: bash -l {0} run: | - cd tests/ - pytest *.py + pytest diff --git a/RMtools_1D/do_RMsynth_1D.py b/RMtools_1D/do_RMsynth_1D.py index 9e4701c..f5ea486 100755 --- a/RMtools_1D/do_RMsynth_1D.py +++ b/RMtools_1D/do_RMsynth_1D.py @@ -324,7 +324,7 @@ def run_rmsynth( ) # Calculate the Rotation Measure Spread Function - RMSFArr, phi2Arr_radm2, fwhmRMSFArr, fitStatArr = get_rmsf_planes( + RMSFArr, phi2Arr_radm2, fwhmRMSFArr, fitStatArr, _ = get_rmsf_planes( lambdaSqArr_m2=lambdaSqArr_m2, phiArr_radm2=phiArr_radm2, weightArr=weightArr, diff --git a/RMtools_3D/do_RMsynth_3D.py b/RMtools_3D/do_RMsynth_3D.py index 0056853..bdbdc7d 100755 --- a/RMtools_3D/do_RMsynth_3D.py +++ b/RMtools_3D/do_RMsynth_3D.py @@ -34,6 +34,7 @@ # # # =============================================================================# +import gc import math as m import os import sys @@ -68,6 +69,7 @@ def run_rmsynth( fitRMSF=False, nBits=32, verbose=True, + not_rmsynth=False, not_rmsf=False, log=print, super_resolution=False, @@ -91,7 +93,8 @@ def run_rmsynth( fitRMSF (bool): Fit a Gaussian to the RMSF? nBits (int): Precision of floating point numbers. verbose (bool): Verbosity. - not_rmsf (bool): Just do RM synthesis and ignore RMSF? + not_rmsynth (bool): Just do RMSF and ignore RM synthesis? + not_rmsf (bool): Just do RM synthesis and ignore RMSF? -- one of these must be False log (function): Which logging function to use. Returns: @@ -104,6 +107,13 @@ def run_rmsynth( """ + if not_rmsynth and not_rmsf: + log( + "Err: both RM synthesis and RMSF computation not requested?\n" + + "Please make sure either not_rmsynth or not_rmsf is False" + ) + sys.exit() + # Sanity check on header dimensions if not str(dataQ.shape) == str(dataU.shape): @@ -185,19 +195,24 @@ def run_rmsynth( uArr = dataU # Perform RM-synthesis on the cube - FDFcube, lam0Sq_m2 = do_rmsynth_planes( - dataQ=qArr, - dataU=uArr, - lambdaSqArr_m2=lambdaSqArr_m2, - phiArr_radm2=phiArr_radm2, - weightArr=weightArr, - nBits=32, - verbose=verbose, - lam0Sq_m2=0 if super_resolution else None, - ) + if not not_rmsynth: + FDFcube, lam0Sq_m2 = do_rmsynth_planes( + dataQ=qArr, + dataU=uArr, + lambdaSqArr_m2=lambdaSqArr_m2, + phiArr_radm2=phiArr_radm2, + weightArr=weightArr, + nBits=32, + verbose=verbose, + lam0Sq_m2=0 if super_resolution else None, + ) + else: + # need lambda0 for RMSF calculation + lam0Sq_m2 = 0 if super_resolution else None + # Calculate the Rotation Measure Spread Function cube - if not_rmsf is not True: - RMSFcube, phi2Arr_radm2, fwhmRMSFCube, fitStatArr = get_rmsf_planes( + if not not_rmsf: + RMSFcube, phi2Arr_radm2, fwhmRMSFCube, fitStatArr, lam0Sq_m2 = get_rmsf_planes( lambdaSqArr_m2=lambdaSqArr_m2, phiArr_radm2=phiArr_radm2, weightArr=weightArr, @@ -238,10 +253,11 @@ def run_rmsynth( # Multiply the dirty FDF by Ifreq0 to recover the PI FDFcube *= Ifreq0Arr - if not_rmsf: + if not_rmsf: # only RMsynth dataArr = [FDFcube, phiArr_radm2, lam0Sq_m2, lambdaSqArr_m2] - - else: + elif not_rmsynth: # only RMSF + dataArr = [RMSFcube, phi2Arr_radm2, fwhmRMSFCube, fitStatArr, lam0Sq_m2] + else: # both have been computed dataArr = [ FDFcube, phiArr_radm2, @@ -264,17 +280,21 @@ def writefits( outDir="", nBits=32, write_seperate_FDF=True, + not_rmsynth=False, not_rmsf=False, + do_peakmaps=True, verbose=False, log=print, ): """Write data to disk in FITS Args: - dataArr (list): FDF and RMSF information - if not_rmsf: + dataArr (list): FDF and/or RMSF information + if not_rmsf: # only RMsynth dataArr = [FDFcube, phiArr_radm2, lam0Sq_m2, lambdaSqArr_m2] - else: + elif not_rmsynth: # only RMSF + dataArr = [RMSFcube, phi2Arr_radm2, fwhmRMSFCube, fitStatArr, lam0Sq_m2] + else: # both dataArr = [FDFcube, phiArr_radm2, RMSFcube, phi2Arr_radm2, fwhmRMSFCube,fitStatArr, lam0Sq_m2, lambdaSqArr_m2] headtemplate: FITS header template @@ -285,7 +305,9 @@ def writefits( outDir (str): Directory to save files. write_seperate_FDF (bool): Write Q, U, and PI separately? verbose (bool): Verbosity. - not_rmsf (bool): Just do RM synthesis and ignore RMSF? + not_rmsynth (bool): Just do RMSF and ignore RM synthesis? + not_rmsf (bool): Just do RM synthesis and ignore RMSF? -- one of these must be False + do_peakmaps (bool): Compute and write peak RM and peak intensity? log (function): Which logging function to use. @@ -311,9 +333,21 @@ def writefits( """ + if not_rmsynth and not_rmsf: + log( + "Err: both RM synthesis and RMSF computation not done?\n" + + "Please make sure either not_rmsynth or not_rmsf is False" + ) + sys.exit() + if not_rmsf: FDFcube, phiArr_radm2, lam0Sq_m2, lambdaSqArr_m2 = dataArr - + if verbose: + log("Saving the dirty FDF and ancillary FITS files.") + elif not_rmsynth: + RMSFcube, phi2Arr_radm2, fwhmRMSFCube, fitStatArr, lam0Sq_m2 = dataArr + if verbose: + log("Saving the RMSF and ancillary FITS files.") else: ( FDFcube, @@ -325,13 +359,13 @@ def writefits( lam0Sq_m2, lambdaSqArr_m2, ) = dataArr + if verbose: + log("Saving the dirty FDF, RMSF and ancillary FITS files.") # Default data typess dtFloat = "float" + str(nBits) dtComplex = "complex" + str(2 * nBits) - if verbose: - log("Saving the dirty FDF, RMSF and ancillary FITS files.") # Make a copy of the Q header and alter frequency-axis as Faraday depth header = headtemplate.copy() Ndim = header["NAXIS"] @@ -347,17 +381,7 @@ def writefits( pass # The try statement is needed for if the FITS header does not # have CTYPE keywords. - header["NAXIS" + str(freq_axis)] = phiArr_radm2.size header["CTYPE" + str(freq_axis)] = ("FDEP", "Faraday depth (linear)") - header["CDELT" + str(freq_axis)] = ( - np.diff(phiArr_radm2)[0], - "[rad/m^2] Coordinate increment at reference point", - ) - header["CRPIX" + str(freq_axis)] = phiArr_radm2.size // 2 + 1 - header["CRVAL" + str(freq_axis)] = ( - phiArr_radm2[phiArr_radm2.size // 2], - "[rad/m^2] Coordinate value at reference point", - ) header["CUNIT" + str(freq_axis)] = "rad/m^2" if not np.isfinite(lam0Sq_m2): lam0Sq_m2 = 0.0 @@ -379,68 +403,86 @@ def writefits( freq_axis - 1 ] # Remove frequency axis (since it's first in the array) output_axes.reverse() # To get into numpy order. - # Put frequency axis first, and reshape to add degenerate axes: - FDFcube = np.reshape(FDFcube, [FDFcube.shape[0]] + output_axes) - if not_rmsf is not True: - RMSFcube = np.reshape(RMSFcube, [RMSFcube.shape[0]] + output_axes) - - # Move Faraday depth axis to appropriate position to match header. - 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) + # Save the FDF + if not not_rmsynth: + header["NAXIS" + str(freq_axis)] = phiArr_radm2.size + header["CDELT" + str(freq_axis)] = ( + np.diff(phiArr_radm2)[0], + "[rad/m^2] Coordinate increment at reference point", + ) + header["CRPIX" + str(freq_axis)] = phiArr_radm2.size // 2 + 1 + header["CRVAL" + str(freq_axis)] = ( + phiArr_radm2[phiArr_radm2.size // 2], + "[rad/m^2] Coordinate value at reference point", + ) - fitsFileOut = outDir + "/" + prefixOut + "FDF_im_dirty.fits" - if verbose: - log("> %s" % fitsFileOut) - hdu1.writeto(fitsFileOut, output_verify="fix", overwrite=True) + # Put frequency axis first, and reshape to add degenerate axes: + FDFcube = np.reshape(FDFcube, [FDFcube.shape[0]] + output_axes) + # Move Faraday depth axis to appropriate position to match header. + FDFcube = np.moveaxis(FDFcube, 0, Ndim - freq_axis) - fitsFileOut = outDir + "/" + prefixOut + "FDF_tot_dirty.fits" - if verbose: - log("> %s" % fitsFileOut) - hdu2.writeto(fitsFileOut, output_verify="fix", overwrite=True) + if write_seperate_FDF: # more memory efficient as well + header = _setStokes(header, "Q") + hdu0 = pf.PrimaryHDU(FDFcube.real.astype(dtFloat), header) + fitsFileOut = outDir + "/" + prefixOut + "FDF_real_dirty.fits" + if verbose: + log("> %s" % fitsFileOut) + hdu0.writeto(fitsFileOut, output_verify="fix", overwrite=True) + del hdu0 + gc.collect() - else: - 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]) - hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True) - hduLst.close() + header = _setStokes(header, "U") + hdu1 = pf.PrimaryHDU(FDFcube.imag.astype(dtFloat), header) + fitsFileOut = outDir + "/" + prefixOut + "FDF_im_dirty.fits" + if verbose: + log("> %s" % fitsFileOut) + hdu1.writeto(fitsFileOut, output_verify="fix", overwrite=True) + del hdu1 + gc.collect() - # Header for outputs that are RM maps (peakRM, RMSF_FWHM) + 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_tot_dirty.fits" + if verbose: + log("> %s" % fitsFileOut) + hdu2.writeto(fitsFileOut, output_verify="fix", overwrite=True) + del hdu2 + gc.collect() + + else: + 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]) + hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True) + hduLst.close() # Save the RMSF - if not_rmsf is not True: + if not not_rmsf: + # Put frequency axis first, and reshape to add degenerate axes: + 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) + + # Header for outputs that are RMSF header["NAXIS" + str(freq_axis)] = phi2Arr_radm2.size header["CDELT" + str(freq_axis)] = ( np.diff(phi2Arr_radm2)[0], @@ -466,7 +508,7 @@ def writefits( "Axis left in to avoid FITS errors", ) rmheader["CUNIT" + str(freq_axis)] = "" - rmheader["CRVAL" + str(freq_axis)] = phiArr_radm2[0] + 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}"]: @@ -477,35 +519,40 @@ def writefits( "Axis left in to avoid FITS errors", ) - if write_seperate_FDF: + if write_seperate_FDF: # more memory efficient as well 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) + del hdu0 + gc.collect() + header = _setStokes(header, "U") + hdu1 = pf.PrimaryHDU(RMSFcube.imag.astype(dtFloat), header) fitsFileOut = outDir + "/" + prefixOut + "RMSF_im.fits" if verbose: log("> %s" % fitsFileOut) hdu1.writeto(fitsFileOut, output_verify="fix", overwrite=True) + del hdu1 + gc.collect() + 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) fitsFileOut = outDir + "/" + prefixOut + "RMSF_tot.fits" if verbose: log("> %s" % fitsFileOut) hdu2.writeto(fitsFileOut, output_verify="fix", overwrite=True) + del hdu2 + gc.collect() + hdu3 = pf.PrimaryHDU( + np.expand_dims(fwhmRMSFCube.astype(dtFloat), axis=0), rmheader + ) fitsFileOut = outDir + "/" + prefixOut + "RMSF_FWHM.fits" if verbose: log("> %s" % fitsFileOut) @@ -532,77 +579,94 @@ def writefits( hduLst.writeto(fitsFileOut, output_verify="fix", overwrite=True) hduLst.close() - # 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. - header["NAXIS" + str(freq_axis)] = 1 - header["CRVAL" + str(freq_axis)] = ( - phiArr_radm2[0], - "[rad/m^2] Coordinate value at reference point", - ) - if "DATAMAX" in header: - del header["DATAMAX"] - if "DATAMIN" in header: - del header["DATAMIN"] + if not not_rmsynth and do_peakmaps: + ## Note that peaks are computed from the sampled functions + ## might be better to fit the FDF and compute the peak. + ## See RMpeakfit_3D.py - # Generate peak maps: + # 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. + # Generate peak maps: + + maxPI, peakRM = create_peak_maps(FDFcube, phiArr_radm2, Ndim - freq_axis) + # Save a maximum polarised intensity map + if "BUNIT" in headtemplate: + 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)] = "" - maxPI, peakRM = create_peak_maps(FDFcube, phiArr_radm2, Ndim - freq_axis) - # Save a maximum polarised intensity map - if "BUNIT" in headtemplate: - 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)] = "" + # Header for output that are RM maps (peakRM, RMSF_FWHM, maxPI) + header["NAXIS" + str(freq_axis)] = 1 + header["CRVAL" + str(freq_axis)] = ( + phiArr_radm2[0], + "[rad/m^2] Coordinate value at reference point", + ) + if "DATAMAX" in header: + del header["DATAMAX"] + if "DATAMIN" in header: + del header["DATAMIN"] - 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}"] = ( + # Generate peak maps: + + 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)] = "" - fitsFileOut = outDir + "/" + prefixOut + "FDF_maxPI.fits" - if verbose: - log("> %s" % fitsFileOut) - pf.writeto( - fitsFileOut, - np.expand_dims(maxPI.astype(dtFloat), axis=0), - header, - overwrite=True, - output_verify="fix", - ) - # 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", - ) + 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, + np.expand_dims(maxPI.astype(dtFloat), axis=0), + header, + overwrite=True, + output_verify="fix", + ) + # 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", + ) -# #Cameron: I've removed the moment 1 map for now because I don't think it's properly/robustly defined. -# # Save an RM moment-1 map -# fitsFileOut = outDir + "/" + prefixOut + "FDF_mom1.fits" -# header["BUNIT"] = "rad/m^2" -# mom1FDFmap = (np.nansum(np.moveaxis(np.abs(FDFcube),FDFcube.ndim-freq_axis,FDFcube.ndim-1) * phiArr_radm2, FDFcube.ndim-1) -# /np.nansum(np.abs(FDFcube), FDFcube.ndim-freq_axis)) -# mom1FDFmap = mom1FDFmap.astype(dtFloat) -# if(verbose): log("> %s" % fitsFileOut) -# pf.writeto(fitsFileOut, mom1FDFmap, header, overwrite=True, -# output_verify="fix") + # #Cameron: I've removed the moment 1 map for now because I don't think it's properly/robustly defined. + # # Save an RM moment-1 map + # fitsFileOut = outDir + "/" + prefixOut + "FDF_mom1.fits" + # header["BUNIT"] = "rad/m^2" + # mom1FDFmap = (np.nansum(np.moveaxis(np.abs(FDFcube),FDFcube.ndim-freq_axis,FDFcube.ndim-1) * phiArr_radm2, FDFcube.ndim-1) + # /np.nansum(np.abs(FDFcube), FDFcube.ndim-freq_axis)) + # mom1FDFmap = mom1FDFmap.astype(dtFloat) + # if(verbose): log("> %s" % fitsFileOut) + # pf.writeto(fitsFileOut, mom1FDFmap, header, overwrite=True, + # output_verify="fix") def _setStokes(header, stokes): diff --git a/RMutils/util_RM.py b/RMutils/util_RM.py index 4b9c107..583f328 100644 --- a/RMutils/util_RM.py +++ b/RMutils/util_RM.py @@ -60,6 +60,7 @@ # # # =============================================================================# +import gc import math as m import sys @@ -182,6 +183,10 @@ def do_rmsynth_planes( KArr[KArr == np.inf] = 0 KArr = np.nan_to_num(KArr) + # Clean up one cube worth of memory + del weightCube + gc.collect() + # Do the RM-synthesis on each plane a = lambdaSqArr_m2 - lam0Sq_m2 FDFcube = ( @@ -378,6 +383,10 @@ def get_rmsf_planes( * KArr[..., None] ).T + # Clean up one cube worth of memory + del weightCube + gc.collect() + # Default to the analytical RMSF fwhmRMSFArr = np.ones((nPix), dtype=dtFloat) * fwhmRMSF statArr = np.ones((nPix), dtype="int") * (-1) @@ -411,7 +420,7 @@ def get_rmsf_planes( fwhmRMSFArr = np.reshape(fwhmRMSFArr, (old_data_shape[1], old_data_shape[2])) statArr = np.reshape(statArr, (old_data_shape[1], old_data_shape[2])) - return RMSFcube, phi2Arr, fwhmRMSFArr, statArr + return RMSFcube, phi2Arr, fwhmRMSFArr, statArr, lam0Sq_m2 # -----------------------------------------------------------------------------# @@ -2139,12 +2148,7 @@ def threeDnoise_get_rmsf_planes( weightArr = np.where(np.isnan(weightArr), 0.0, weightArr) nDims = weightArr.ndim - # Set the mask array (default to 1D, no masked channels) - # Sanity checks on array sizes - # if not weightArr.shape == lambdaSqArr_m2.shape: - # print("Err: wavelength^2 and weight arrays must be the same shape.") - # return None, None, None, None if not nDims <= 3: log("Err: mask dimensions must be <= 3.") return None, None, None, None diff --git a/tests/QA_tests.py b/tests/QA_test.py similarity index 78% rename from tests/QA_tests.py rename to tests/QA_test.py index 2b5430d..3506e53 100644 --- a/tests/QA_tests.py +++ b/tests/QA_test.py @@ -20,11 +20,16 @@ import shutil import subprocess import unittest +from pathlib import Path import numpy as np from astropy.io import fits as pf from scipy.ndimage import gaussian_filter +TEST_PATH = Path(__file__).parent.absolute() +ONED_PATH = TEST_PATH / "simdata" / "1D" +THREED_PATH = TEST_PATH / "simdata" / "3D" + def Faraday_thin_complex_polarization(freq_array, RM, Polint, initial_angle): """freq_array = channel frequencies in Hz @@ -45,25 +50,10 @@ def create_1D_data(freq_arr): noise_amplitude = 0.1 spectral_index = -0.7 error_estimate = 1 # Size of assumed errors as multiple of actual error. - - ## Random data generation is not used any more, since it caused different - ## results on different machines. - # pol_spectrum=Faraday_thin_complex_polarization(freq_arr,RM,fracpol,pol_angle_deg) - # I_spectrum=StokesI_midband*(freq_arr/np.median(freq_arr))**spectral_index - # rng=np.random.default_rng(20200422) - # noise_spectrum_I=rng.normal(scale=noise_amplitude,size=freq_arr.shape) - # noise_spectrum_Q=rng.normal(scale=noise_amplitude,size=freq_arr.shape) - # noise_spectrum_U=rng.normal(scale=noise_amplitude,size=freq_arr.shape) - # dIQU=np.ones_like(freq_arr)*noise_amplitude*error_estimate - - if not os.path.isdir("simdata/1D"): - os.makedirs("simdata/1D") - shutil.copy("RMsynth1D_testdata.dat", "simdata/1D/simsource.dat") - # np.savetxt('simdata/1D/simsource.dat', list(zip(freq_arr,I_spectrum+noise_spectrum_I, - # I_spectrum*pol_spectrum.real+noise_spectrum_Q, - # I_spectrum*pol_spectrum.imag+noise_spectrum_U, - # dIQU,dIQU,dIQU))) - with open("simdata/1D/sim_truth.txt", "w") as f: + if not ONED_PATH.exists(): + ONED_PATH.mkdir(parents=True) + shutil.copy(TEST_PATH / "RMsynth1D_testdata.dat", ONED_PATH / "simsource.dat") + with open(ONED_PATH / "sim_truth.txt", "w") as f: f.write("RM = {} rad/m^2\n".format(RM)) f.write("Intrsinsic polarization angle = {} deg\n".format(pol_angle_deg)) f.write("Fractional polarization = {} %\n".format(fracpol * 100.0)) @@ -169,20 +159,20 @@ def create_3D_data(freq_arr, N_side=100): header["BUNIT"] = "Jy/beam" - if not os.path.isdir("simdata/3D"): - os.makedirs("simdata/3D") + if not THREED_PATH.exists(): + THREED_PATH.mkdir(parents=True) pf.writeto( - "simdata/3D/Q_cube.fits", np.transpose(Q_cube), header=header, overwrite=True + THREED_PATH / "Q_cube.fits", np.transpose(Q_cube), header=header, overwrite=True ) pf.writeto( - "simdata/3D/U_cube.fits", np.transpose(U_cube), header=header, overwrite=True + THREED_PATH / "U_cube.fits", np.transpose(U_cube), header=header, overwrite=True ) - with open("simdata/3D/freqHz.txt", "w") as f: + with open(THREED_PATH / "freqHz.txt", "w") as f: for freq in freq_arr: f.write("{:}\n".format(freq)) - with open("simdata/3D/sim_truth.txt", "w") as f: + with open(THREED_PATH / "sim_truth.txt", "w") as f: f.write("Point source:\n") f.write("RM = {} rad/m^2\n".format(src_RM)) f.write("Intrsinsic polarization angle = {} deg\n".format(src_pol_angle_deg)) @@ -211,13 +201,14 @@ def setUp(self): def test_a1_1D_synth_runs(self): create_1D_data(self.freq_arr) returncode = subprocess.call( - "rmsynth1d simdata/1D/simsource.dat -l 600 -d 3 -S -i", shell=True + f"rmsynth1d '{(ONED_PATH/'simsource.dat').as_posix()}' -l 600 -d 3 -S -i", + shell=True, ) self.assertEqual(returncode, 0, "RMsynth1D failed to run.") def test_a2_1D_synth_values(self): - mDict = json.load(open("simdata/1D/simsource_RMsynth.json", "r")) - refDict = json.load(open("RMsynth1D_referencevalues.json", "r")) + mDict = json.load(open(ONED_PATH / "simsource_RMsynth.json", "r")) + refDict = json.load(open(TEST_PATH / "RMsynth1D_referencevalues.json", "r")) for key in mDict.keys(): if (key == "polyCoefferr") or key == "polyCoeffs": ref_values = refDict[key].split(",") @@ -242,11 +233,11 @@ def test_a2_1D_synth_values(self): def test_c_3D_synth(self): create_3D_data(self.freq_arr) returncode = subprocess.call( - "rmsynth3d simdata/3D/Q_cube.fits simdata/3D/U_cube.fits simdata/3D/freqHz.txt -l 300 -d 10", + f"rmsynth3d '{(THREED_PATH/'Q_cube.fits').as_posix()}' '{(THREED_PATH/'U_cube.fits').as_posix()}' '{(THREED_PATH/'freqHz.txt').as_posix()}' -l 300 -d 10", shell=True, ) self.assertEqual(returncode, 0, "RMsynth3D failed to run.") - header = pf.getheader("simdata/3D/FDF_tot_dirty.fits") + header = pf.getheader(THREED_PATH / "FDF_tot_dirty.fits") self.assertEqual(header["NAXIS"], 3, "Wrong number of axes in output?") self.assertEqual( (header["NAXIS1"], header["NAXIS2"]), @@ -258,16 +249,16 @@ def test_c_3D_synth(self): ) def test_b1_1D_clean(self): - if not os.path.exists("simdata/1D/simsource_RMsynth.dat"): + if not (ONED_PATH / "simsource_RMsynth.dat").exists(): self.skipTest("Could not test 1D clean; 1D synth failed first.") returncode = subprocess.call( - "rmclean1d simdata/1D/simsource.dat -n 11 -S", shell=True + f"rmclean1d '{(ONED_PATH/'simsource.dat').as_posix()}' -n 11 -S", shell=True ) self.assertEqual(returncode, 0, "RMclean1D failed to run.") def test_b2_1D_clean_values(self): - mDict = json.load(open("simdata/1D/simsource_RMclean.json", "r")) - refDict = json.load(open("RMclean1D_referencevalues.json", "r")) + mDict = json.load(open(ONED_PATH / "simsource_RMclean.json", "r")) + refDict = json.load(open(TEST_PATH / "RMclean1D_referencevalues.json", "r")) for key in mDict.keys(): self.assertTrue( np.abs((mDict[key] - refDict[key]) / refDict[key]) < 1e-3, @@ -275,37 +266,40 @@ def test_b2_1D_clean_values(self): ) def test_d_3D_clean(self): - if not os.path.exists("simdata/3D/FDF_tot_dirty.fits"): + if not (THREED_PATH / "FDF_tot_dirty.fits").exists(): self.skipTest("Could not test 3D clean; 3D synth failed first.") returncode = subprocess.call( - "rmclean3d simdata/3D/FDF_tot_dirty.fits simdata/3D/RMSF_tot.fits -n 10", + f"rmclean3d '{(THREED_PATH/'FDF_tot_dirty.fits').as_posix()}' '{(THREED_PATH/'RMSF_tot.fits').as_posix()}' -n 10", shell=True, ) self.assertEqual(returncode, 0, "RMclean3D failed to run.") # what else? def test_e_1Dsynth_fromFITS(self): - if not os.path.exists("simdata/3D/Q_cube.fits"): + if not (THREED_PATH / "Q_cube.fits").exists(): create_3D_data(self.freq_arr) returncode = subprocess.call( - "rmsynth1dFITS simdata/3D/Q_cube.fits simdata/3D/U_cube.fits 25 25 -l 600 -d 3 -S", + f"rmsynth1dFITS '{(THREED_PATH/'Q_cube.fits').as_posix()}' '{(THREED_PATH/'U_cube.fits').as_posix()}' 25 25 -l 600 -d 3 -S", shell=True, ) self.assertEqual(returncode, 0, "RMsynth1D_fromFITS failed to run.") def test_f1_QUfitting(self): - if not os.path.exists("simdata/1D/simsource.dat"): + if not (ONED_PATH / "simsource.dat").exists(): create_1D_data(self.freq_arr) - if not os.path.exists("models_ns"): - shutil.copytree("../RMtools_1D/models_ns", "models_ns") + + local_models = Path("models_ns") + if not local_models.exists(): + shutil.copytree(TEST_PATH / ".." / "RMtools_1D" / "models_ns", local_models) returncode = subprocess.call( - "qufit simdata/1D/simsource.dat --sampler nestle", shell=True + f"qufit '{(ONED_PATH/'simsource.dat').as_posix()}' --sampler nestle", + shell=True, ) self.assertEqual(returncode, 0, "QU fitting failed to run.") - shutil.rmtree("models_ns") + shutil.rmtree(local_models) def test_f2_QUfit_values(self): - mDict = json.load(open("simdata/1D/simsource_m1_nestle.json", "r")) + mDict = json.load(open(ONED_PATH / "simsource_m1_nestle.json", "r")) # The QU-fitting code has internal randomness that I can't control. So every run # will produce slightly different results. I want to assert that these differences # are below 1%. @@ -332,7 +326,7 @@ def test_f2_QUfit_values(self): if __name__ == "__main__": os.chdir(os.path.dirname(os.path.realpath(__file__))) - if os.path.exists("simdata"): + if (TEST_PATH / "simdata").exists(): shutil.rmtree("simdata") print("\nUnit tests running.") diff --git a/tests/test_nufft.py b/tests/nufft_test.py similarity index 99% rename from tests/test_nufft.py rename to tests/nufft_test.py index 795bfeb..b2231c4 100644 --- a/tests/test_nufft.py +++ b/tests/nufft_test.py @@ -417,7 +417,7 @@ def test_rmsf(): """Test the NUFFT RMSF routine agaist DFT.""" fake_data = make_fake_data() tick = time() - RMSFcube, phi2Arr, fwhmRMSFArr, statArr = get_rmsf_planes( + RMSFcube, phi2Arr, fwhmRMSFArr, statArr, _ = get_rmsf_planes( lambdaSqArr_m2=fake_data.lsq, phiArr_radm2=fake_data.phis, weightArr=fake_data.weights,