From a397fbc52f00fe6d0195a77c677e93e35c5c48ab Mon Sep 17 00:00:00 2001 From: Thomas Williams Date: Thu, 4 Jun 2026 16:26:28 +0100 Subject: [PATCH] Keeps Stokes throughout postprocessing This PR fixes the unnecessary slowdowns that come with re-adding degenerate axes throughout postprocessing. CASA now typically crashes if the images don't have 4 axes in, which it didn't use to. So we now spend a lot of time ripping them out and putting them back in, which can take a long time for big cubes. Now, we just keep the Stokes axis in throughout, dropping when we export to fits at the end. This also finally fixes the feather/feather_before_mosaic logic. Now, if feather is True and feather_before_mosaic is False, it will feather together individual mosaic tiles, as well as mosaicking the interferometric and singledish data before feathering those together for the final mosaic of mosaics. If feather_before mosaic is True, then it will instead mosaic together the feathered individual tiles. Note that if feather_before_mosaic is True, it will still create the full singledish image, but will not use this for the final mosaic. - Fix bug with has_memory_issue call - Ensure coordinate axes between interferometric/singledish data match up when staging in postprocessing - Fix pixperbeam in trim cube to actually respect min_pixperbeam - Remove dropdegaxes calls, since we don't use them any more - ccr_dropdeg has now been renamed to ccr_importfits, and modified to either import a fits file or just copy over - Fix logic with feather_before_mosaic --- CHANGELOG.md | 1 + phangsPipeline/casaCubeRoutines.py | 82 +++++------- phangsPipeline/casaFeatherRoutines.py | 66 +++------- phangsPipeline/casaMosaicRoutines.py | 29 +++-- phangsPipeline/handlerPostprocess.py | 177 ++++++++++++-------------- 5 files changed, 147 insertions(+), 208 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c6e3de8..0551ca15 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Refactored logging, and added tests (#338). - Updated bespoke sdintimaging task, to align with latest CASA version (#347). - If we don't have any model flux, then overwrite minimum number of major cycles (#359). +- Keep all 4 axes throughout postprocessing, to avoid slowdowns with re-adding degenerate axes (#353). ### Fixed diff --git a/phangsPipeline/casaCubeRoutines.py b/phangsPipeline/casaCubeRoutines.py index 25718f33..bd99f3ae 100644 --- a/phangsPipeline/casaCubeRoutines.py +++ b/phangsPipeline/casaCubeRoutines.py @@ -102,20 +102,27 @@ def check_getchunk_putchunk_memory_issue( list_to_return.append(cube_mask) if return_shape: list_to_return.append(cube_shape) - return tuple(list_to_return) + + list_to_return = tuple(list_to_return) + + # If we're only returning one thing, remove + # the list for forward-proofing + if len(list_to_return) == 1: + list_to_return = list_to_return[0] + + return list_to_return #endregion #region Copying, scaling, etc. -def copy_dropdeg( +def copy_importfits( infile=None, outfile=None, overwrite=False ): """ - Copy using imsubimage to drop degenerate axes. Optionally handle - overwriting and importing from FITS. + Copy file, handling overwriting and importing from FITS. """ if os.path.isdir(outfile): @@ -124,7 +131,6 @@ def copy_dropdeg( return(False) os.system('rm -rf '+outfile) - used_temp_outfile = False if (infile[-4:] == 'FITS') and os.path.isfile(infile): logger.info("Importing from FITS.") temp_outfile = outfile+'.temp' @@ -135,24 +141,15 @@ def copy_dropdeg( os.system('rm -rf '+temp_outfile) casaStuff.importfits(fitsimage=infile, - imagename=temp_outfile, - zeroblanks=False, - overwrite=overwrite) - used_temp_outfile = True - - if used_temp_outfile: - casaStuff.imsubimage( - imagename=temp_outfile, - outfile=outfile, - dropdeg=True) - os.system('rm -rf '+temp_outfile) + imagename=outfile, + zeroblanks=False, + overwrite=overwrite, + ) + else: - casaStuff.imsubimage( - imagename=infile, - outfile=outfile, - dropdeg=True) + os.system("cp -r "+infile+" "+outfile) - return(True) + return True def get_mask(infile, huge_cube_workaround=True): @@ -236,20 +233,6 @@ def copy_mask(infile, outfile, huge_cube_workaround=True): Copy a mask from infile to outfile. Includes a switch for large cubes, where getchunk/putchunk can segfault """ - #if huge_cube_workaround: - # os.system('rm -rf ' + outfile + '/mask0') - # os.system('cp -r ' + infile + '/mask0' + ' ' + outfile + '/mask0') - #else: - # myia = au.createCasaTool(casaStuff.iatool) - # myia.open(infile) - # mask = myia.getchunk(getmask=True) - # myia.close() - # - # myia = au.createCasaTool(casaStuff.iatool) - # myia.open(outfile) - # mask = myia.putregion(pixelmask=mask) - # myia.close() - mask = get_mask(infile) # use putregion to update pixel mask @@ -266,8 +249,10 @@ def copy_mask(infile, outfile, huge_cube_workaround=True): raise Exception('Error! The infile and outfile have different dimensions! Cannot copy mask.') has_memory_issue = check_getchunk_putchunk_memory_issue( - outfile, myia=myia, - huge_cube_workaround=huge_cube_workaround) + outfile, + myia=myia, + huge_cube_workaround=huge_cube_workaround, + ) if not has_memory_issue: # getchunk was successful, no memory issue myia.putregion(pixelmask=mask) @@ -557,15 +542,18 @@ def trim_cube( return(False) bmaj = hdr['restoringbeam']['major']['value'] - pix_per_beam = bmaj*1.0 / pixel_as*1.0 + pixperbeam = bmaj*1.0 / pixel_as*1.0 + + # Calculate the rebinning factor, and rebin if >1 + rebin_factor = int(np.floor(pixperbeam / min_pixperbeam)) - if pix_per_beam > 6: + if rebin_factor > 1: casaStuff.imrebin( imagename=infile, outfile=outfile+'.temp', - factor=[2,2,1], + factor=[rebin_factor, rebin_factor, 1], crop=True, - dropdeg=True, + dropdeg=False, overwrite=overwrite, ) just_copy = False @@ -575,17 +563,6 @@ def trim_cube( outfile_ext = ".temp" - # If we don't have the Stokes I axis, add this in - if "Stokes" not in hdr["axisnames"]: - myia = au.createCasaTool(casaStuff.iatool) - myia.open(outfile + '.temp') - os.system('rm -rf '+outfile + '.temp_deg') - deg_im = myia.adddegaxes(outfile=outfile + '.temp_deg', stokes='I', overwrite=True) - deg_im.done() - myia.close() - outfile_ext = ".temp_deg" - - # This should either be .temp or .temp_deg (check if it works for 2d cases) mask = get_mask(outfile + outfile_ext, huge_cube_workaround=True) # Figure out the extent of the image inside the cube @@ -617,7 +594,6 @@ def trim_cube( ) os.system('rm -rf '+outfile+'.temp') - os.system('rm -rf '+outfile+'.temp_deg') return(True) diff --git a/phangsPipeline/casaFeatherRoutines.py b/phangsPipeline/casaFeatherRoutines.py index b7928cc1..5b81c28a 100644 --- a/phangsPipeline/casaFeatherRoutines.py +++ b/phangsPipeline/casaFeatherRoutines.py @@ -22,7 +22,6 @@ def prep_sd_for_feather( sdfile_out=None, interf_file=None, do_import=True, - do_dropdeg=True, do_align=True, do_checkunits=True, overwrite=False): @@ -42,11 +41,6 @@ def prep_sd_for_feather( do_import (default True) : If True and the infile is a FITS file, then import the data from FITS. - do_dropdeg (default True) : if True then pare degenerate axes from - the signle dish file. In general this is a good idea for - postprocessing, where mixing and matching degenerate axes causes - many CASA routines to fail. - do_align (default True) : If True then align the single dish data to the astrometric grid of the the interferometer file. @@ -96,37 +90,7 @@ def prep_sd_for_feather( overwrite=overwrite) current_infile = current_outfile - # Drop degenerate axes using a call to imsubimage - - if do_dropdeg: - if current_infile == current_outfile: - if os.path.isdir(tempfile_name) or os.path.isfile(tempfile_name): - if overwrite: - os.system('rm -rf '+tempfile_name) - else: - logger.error("Temp file needed but exists and overwrite=False - "+tempfile_name) - return(None) - os.system('cp -r '+current_infile+' '+tempfile_name) - current_infile = tempfile_name - os.system('rm -rf '+current_outfile) - - if os.path.isdir(current_outfile) or os.path.isfile(current_outfile): - if overwrite: - os.system('rm -rf '+current_outfile) - else: - logger.error("Output file needed exists and overwrite=False - "+current_outfile) - return(None) - - casaStuff.imsubimage( - imagename=current_infile, - outfile=current_outfile, - overwrite=overwrite, - dropdeg=True) - - current_infile = current_outfile - # Align the single dish data to the astrometric grid of the interferometric data - if do_align: if current_infile == current_outfile: if os.path.isdir(tempfile_name) or os.path.isfile(tempfile_name): @@ -148,7 +112,25 @@ def prep_sd_for_feather( interpolation='cubic', overwrite=overwrite) - current_infile = current_outfile + # Get out and match axis orders from interferometric images + myia = au.createCasaTool(casaStuff.iatool) + myia.open(interf_file) + interf_cs = myia.coordsys() + myia.close() + + myia = au.createCasaTool(casaStuff.iatool) + myia.open(current_outfile) + sd_cs = myia.coordsys() + myia.close() + + if not sd_cs.names() == interf_cs.names(): + casaStuff.imtrans( + imagename=current_outfile, + outfile=current_outfile + ".reorder", + order=interf_cs.names(), + ) + os.system("rm -rf "+current_outfile) + os.system("mv "+current_outfile+".reorder "+current_outfile) # Check the units on the singledish file and convert from K to Jy/beam if needed. @@ -422,21 +404,15 @@ def feather_two_cubes( myia.adddegaxes(outfile=current_sd_file, stokes='I', overwrite=True) myia.close() - # Call feather, followed by an imsubimage to deal with degenerate - # axis stuff. - + # Call feather if overwrite: os.system('rm -rf '+out_file) os.system('rm -rf '+out_file+'.temp') - casaStuff.feather(imagename=out_file+'.temp', + casaStuff.feather(imagename=out_file, highres=current_interf_file, lowres=current_sd_file, sdfactor=1.0, lowpassfiltersd=False) - casaStuff.imsubimage(imagename=out_file+'.temp', - outfile=out_file, - dropdeg=True) - os.system('rm -rf '+out_file+'.temp') # If we apodized, now divide out the common kernel. diff --git a/phangsPipeline/casaMosaicRoutines.py b/phangsPipeline/casaMosaicRoutines.py index 38caa59b..b3e61643 100644 --- a/phangsPipeline/casaMosaicRoutines.py +++ b/phangsPipeline/casaMosaicRoutines.py @@ -219,8 +219,6 @@ def calculate_mosaic_extent( # Loop over input files and calculate RA and Dec coordinates of # the corners. - myia = au.createCasaTool(casaStuff.iatool) - for this_infile in infile_list: this_hdr = casaStuff.imhead(this_infile) @@ -266,10 +264,10 @@ def calculate_mosaic_extent( dec_list.append(trc['coords'][:, 1]) dec_list.append(brc['coords'][:, 1]) - freq_list.append(blc['coords'][:, 2]) - freq_list.append(tlc['coords'][:, 2]) - freq_list.append(trc['coords'][:, 2]) - freq_list.append(brc['coords'][:, 2]) + freq_list.append(blc['coords'][:, -1]) + freq_list.append(tlc['coords'][:, -1]) + freq_list.append(trc['coords'][:, -1]) + freq_list.append(brc['coords'][:, -1]) # Get the minimum and maximum RA and Declination. @@ -439,6 +437,11 @@ def build_common_header( target_hdr = casaStuff.imregrid(template_file, template='get') + # Get out the spectral key, which can change + spec_key = 'spectral2' + if spec_key not in target_hdr['csys']: + spec_key = 'spectral1' + # Get the pixel scale. This makes some assumptions. We could put a # lot of general logic here, but we are usually working in a # case where this works. @@ -457,7 +460,7 @@ def build_common_header( target_hdr['csys']['direction0']['crval'][0] = ra_ctr_in_rad target_hdr['csys']['direction0']['crval'][1] = dec_ctr_in_rad - target_hdr['csys']['spectral1']['wcs']['crval'] = freq_ctr + target_hdr['csys'][spec_key]['wcs']['crval'] = freq_ctr # Calculate the size of the image in pixels and set the central # pixel coordinate for the RA and Dec axis. @@ -470,7 +473,7 @@ def build_common_header( dec_axis_size = np.ceil(delta_dec / dec_pix_in_as) + 1 new_dec_ctr_pix = (dec_axis_size + 1)/2.0 - freq_pix_in_hz = np.abs(target_hdr['csys']['spectral1']['wcs']['cdelt']) + freq_pix_in_hz = np.abs(target_hdr['csys'][spec_key]['wcs']['cdelt']) freq_axis_size = np.ceil(delta_freq / freq_pix_in_hz) + 1 # +1 or the 1-indexing new_freq_ctr_pix = (freq_axis_size + 1) / 2.0 @@ -490,11 +493,12 @@ def build_common_header( target_hdr['csys']['direction0']['crpix'][0] = new_ra_ctr_pix target_hdr['csys']['direction0']['crpix'][1] = new_dec_ctr_pix - target_hdr['csys']['spectral1']['wcs']['crpix'] = new_freq_ctr_pix + target_hdr['csys'][spec_key]['wcs']['crpix'] = new_freq_ctr_pix target_hdr['shap'][0] = int(ra_axis_size) target_hdr['shap'][1] = int(dec_axis_size) - target_hdr['shap'][2] = int(freq_axis_size) + target_hdr['shap'][-1] = int(freq_axis_size) + return(target_hdr) def common_grid_for_mosaic( @@ -1356,13 +1360,12 @@ def mosaic_aligned_data( cur_maskfile = copy.deepcopy(local_maskfile) myia.close() - # Strip out any degenerate axes and create the final output file. - + # Create the final output file. casaStuff.imsubimage(imagename=temp_file, outfile=local_outfile, mask='"'+cur_maskfile+'"', overwrite=overwrite, - dropdeg=True, + dropdeg=False, ) # Remove any temp Stokes files we've made along the way diff --git a/phangsPipeline/handlerPostprocess.py b/phangsPipeline/handlerPostprocess.py index caec2b69..27173432 100644 --- a/phangsPipeline/handlerPostprocess.py +++ b/phangsPipeline/handlerPostprocess.py @@ -11,6 +11,7 @@ call any of the CASA-specific routines). Right now, just avoid direct calls to CASA from this class. """ +import copy import logging import os @@ -379,7 +380,6 @@ def task_stage_interf_data( logger.info(str(target) + " , " + str(product) + " , " + str(config)) logger.info("&%&%&%&%&%&%&%&%&%&%&%&%&%&%") logger.info("") - # logger.info("Using ccr.copy_dropdeg.") logger.info("Staging " + outfile) # Move the cubes to the postprocess directory, trimming along the way @@ -387,8 +387,6 @@ def task_stage_interf_data( if not self._dry_run: os.system('rm -rf ' + outdir + outfile) os.system('rm -rf ' + outdir + outfile + ".temp") - os.system('rm -rf ' + outdir + outfile + ".temp_deg") - # os.system('cp -r ' + indir + infile + ' ' + outdir + outfile) ccr.trim_cube( infile=indir + infile, @@ -398,10 +396,6 @@ def task_stage_interf_data( pad=1, rebin=False, ) - # ccr.copy_dropdeg( - # infile=indir+infile, - # outfile=outdir+outfile, - # overwrite=True) # in case of merged datasets with non-identical frequency setups imaged with per-plane beam, # some edge channels will have much coarser beam, we trim these edge channels here. @@ -414,49 +408,6 @@ def task_stage_interf_data( return () - def task_remove_degenerate_axes( - self, - target=None, - product=None, - config=None, - imaging_method='tclean', - extra_ext='', - check_files=True, - ): - """ - Remove - degenerate axes for target, product, config combination in postprocessing directory. - """ - - # Generate file names - - file_dir = self._kh.get_postprocess_dir_for_target(target) - fname_dict = self._fname_dict(target=target, config=config, product=product, extra_ext=extra_ext, - imaging_method=imaging_method) - - # Copy the primary beam and the interferometric imaging - - logger.info("Dropping degenerate axes from postprocess image/pb files.") - - for this_tag in ['orig', 'pb', 'pbcorr', 'pbcorr_round']: - - file_name = fname_dict[this_tag] - - # Check input file existence - if check_files: - if not (os.path.isdir(file_dir + file_name)): - logger.warning("Missing " + file_dir + file_name) - continue - - if not self._dry_run: - ccr.copy_dropdeg(file_dir + file_name, file_dir + file_name + '_nodeg', overwrite=True) - - os.system('rm -rf ' + file_dir + file_name) - os.system('cp -r ' + file_dir + file_name + '_nodeg ' + file_dir + file_name) - os.system('rm -rf ' + file_dir + file_name + '_nodeg') - - return () - def task_pbcorr( self, target=None, @@ -650,7 +601,6 @@ def task_stage_singledish( sdfile_out=outdir + outfile, interf_file=tempdir + template, do_import=True, - do_dropdeg=True, do_align=True, do_checkunits=True, overwrite=True) @@ -914,9 +864,10 @@ def task_feather( logger.info("Copying from " + interf_weight_file) logger.info("Copying to " + out_weight_file) if not self._dry_run: - ccr.copy_dropdeg(infile=indir + interf_weight_file, - outfile=outdir + out_weight_file, - overwrite=True) + ccr.copy_importfits(infile=indir + interf_weight_file, + outfile=outdir + out_weight_file, + overwrite=True, + ) return () def task_rename_sdintimaging(self, @@ -1576,48 +1527,54 @@ def recipe_prep_one_target( return () # Call tasks - self.task_stage_interf_data( - target=target, config=config, product=product, + target=target, + config=config, + product=product, check_files=check_files, imaging_method=imaging_method, trim_coarse_beam_edge_channels=trim_coarse_beam_edge_channels, ) self.task_pbcorr( - target=target, config=config, product=product, + target=target, + config=config, + product=product, check_files=check_files, imaging_method=imaging_method ) self.task_round_beam( - target=target, config=config, product=product, - check_files=check_files, - imaging_method=imaging_method - ) - - self.task_remove_degenerate_axes( - target=target, config=config, product=product, + target=target, + config=config, + product=product, check_files=check_files, imaging_method=imaging_method ) if has_singledish and imaging_method not in ['sdintimaging']: self.task_stage_singledish( - target=target, config=config, product=product, + target=target, + config=config, + product=product, check_files=check_files ) if is_part_of_mosaic: self.task_make_interf_weight( - target=target, config=config, product=product, - check_files=check_files, scale_by_noise=True, + target=target, + config=config, + product=product, + check_files=check_files, + scale_by_noise=True, imaging_method=imaging_method ) if is_part_of_mosaic and has_singledish and imaging_method not in ['sdintimaging']: self.task_make_singledish_weight( - target=target, config=config, product=product, + target=target, + config=config, + product=product, check_files=check_files, ) @@ -1895,34 +1852,47 @@ def loop_postprocess( self.looper(do_targets=True, do_products=True, do_configs=True): has_imaging = False + imaging_method_prep = copy.deepcopy(imaging_method) + fname_imaging_dir = "" if imaging_method == 'sdintimaging': fname_dict = self._fname_dict( - target=this_target, product=this_product, config=this_config, - imaging_method=imaging_method) + target=this_target, + product=this_product, + config=this_config, + imaging_method=imaging_method, + ) imaging_dir = self._kh.get_imaging_dir_for_target(this_target) - has_imaging = os.path.isdir(imaging_dir + fname_dict['orig']) + fname_imaging_dir = os.path.join(imaging_dir, fname_dict['orig']) + has_imaging = os.path.isdir(fname_imaging_dir) if has_imaging: imaging_method_prep = 'sdintimaging' if not has_imaging: fname_dict = self._fname_dict( - target=this_target, product=this_product, config=this_config) + target=this_target, + product=this_product, + config=this_config, + ) imaging_dir = self._kh.get_imaging_dir_for_target(this_target) - has_imaging = os.path.isdir(imaging_dir + fname_dict['orig']) + fname_imaging_dir = os.path.join(imaging_dir, fname_dict['orig']) + has_imaging = os.path.isdir(fname_imaging_dir) if has_imaging: imaging_method_prep = 'tclean' if not has_imaging: logger.debug("Skipping " + this_target + " because it lacks imaging.") - logger.debug(imaging_dir + fname_dict['orig']) + logger.debug(fname_imaging_dir) continue self.recipe_prep_one_target( - target=this_target, product=this_product, config=this_config, + target=this_target, + product=this_product, + config=this_config, check_files=True, trim_coarse_beam_edge_channels=trim_coarse_beam_edge_channels, - imaging_method=imaging_method_prep) + imaging_method=imaging_method_prep, + ) # Feather the interferometer configuration data that has # single dish imaging. We'll return to feather mosaicked @@ -1955,11 +1925,6 @@ def loop_postprocess( has_imaging = os.path.isdir(imaging_dir + fname_dict['orig']) has_singledish = self._kh.has_singledish(target=this_target, product=this_product) - is_part_of_mosaic = self._kh.is_target_in_mosaic(this_target) - if is_part_of_mosaic and not feather_before_mosaic: - logger.debug("Skipping " + this_target + " because feather_before_mosaic is False.") - continue - if not has_imaging: logger.debug("Skipping " + this_target + " because it lacks imaging.") logger.debug(imaging_dir + fname_dict['orig']) @@ -1971,15 +1936,24 @@ def loop_postprocess( if feather_apod: self.task_feather( - target=this_target, product=this_product, config=this_config, - apodize=True, apod_ext='pb', extra_ext_out='_apod', check_files=True, + target=this_target, + product=this_product, + config=this_config, + apodize=True, + apod_ext='pb', + extra_ext_out='_apod', + check_files=True, copy_weights=True, ) if feather_noapod: self.task_feather( - target=this_target, product=this_product, config=this_config, - apodize=False, extra_ext_out='', check_files=True, + target=this_target, + product=this_product, + config=this_config, + apodize=False, + extra_ext_out='', + check_files=True, copy_weights=True, ) @@ -1996,10 +1970,6 @@ def loop_postprocess( if not is_mosaic: continue - if feather_before_mosaic: - logger.debug("Skipping " + this_target + " because feather_before_mosaic is True.") - continue - # Mosaic the interferometer data and the # single dish data (need to verify if parts # have single dish, enforce the same @@ -2022,6 +1992,11 @@ def loop_postprocess( if not is_mosaic: continue + # Skip if we're not feathering before mosaic + if do_feather and not feather_before_mosaic: + logger.debug("Skipping " + this_target + " because feather_before_mosaic is False.") + continue + if feather_apod: self.recipe_mosaic_one_target( target=this_target, product=this_product, config=this_config, @@ -2038,10 +2013,9 @@ def loop_postprocess( extra_ext_out='', ) - # This round of feathering targets only mosaicked data. All - # other data have been feathered above already. - - if do_feather and feather_before_mosaic: + # This round of feathering targets only mosaicked data, and only if we haven't feathered before mosaicking. + # All other data have been feathered above already. + if do_feather and not feather_before_mosaic: # N.B. if using sdintimaging this will just crash out since it hasn't staged any singledish. This is # intended! @@ -2056,14 +2030,23 @@ def loop_postprocess( if feather_apod: self.task_feather( - target=this_target, product=this_product, config=this_config, - apodize=True, apod_ext='pb', extra_ext_out='_apod', check_files=True, + target=this_target, + product=this_product, + config=this_config, + apodize=True, + apod_ext='pb', + extra_ext_out='_apod', + check_files=True, ) if feather_noapod: self.task_feather( - target=this_target, product=this_product, config=this_config, - apodize=False, extra_ext_out='', check_files=True, + target=this_target, + product=this_product, + config=this_config, + apodize=False, + extra_ext_out='', + check_files=True, ) # Trim and downsample the data, convert to Kelvin, etc.