Merge pull request #62 from AlecThomson/sbid

SBID

CHANGELOG.md

@@ -7,6 +7,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### What's Changed
+
+* Allow SBID to passed as an argument
+    * This will enable a 'single field mode'
+    * Database queries / updates changes to support this
+* Unified ArgParse mode
+    * Much easier argument parsing
+    * Now reused amongst modules
+* Fixes to typing
+    * Much better use of `pathlib.Path` and `pandas`
+
 ## [2.1.7] - 2024-04-03
 ### What's Changed
 

arrakis/.default_config.yaml

@@ -0,0 +1,113 @@
+# options:
+# hosted-wsclean: docker://alecthomson/wsclean:latest # Docker or Singularity image for wsclean (default: docker://alecthomson/wsclean:latest)
+# local_wsclean: null # Path to local wsclean Singularity image (default: None)
+
+# pipeline arguments:
+# dask_config: null # Config file for Dask SlurmCLUSTER. (default: None)
+# imager_dask_config: null #Config  file for Dask SlurmCLUSTER. (default: None)
+imager_only: false #Only run the imager component of the pipeline.  (default: False)
+skip_imager: false #Skip imaging stage [False]. (default: False)
+skip_cutout: false #Skip cutout stage [False]. (default: False)
+skip_linmos: false #Skip LINMOS stage [False]. (default: False)
+skip_frion: false #Skip cleanup stage [False]. (default: False)
+skip_rmsynth: false #Skip RM Synthesis stage [False]. (default: False)
+skip_rmclean: false #Skip RM-CLEAN stage [False]. (default: False)
+skip_cat: false #Skip catalogue stage [False]. (default: False)
+skip_cleanup: false #Skip cleanup stage [False]. (default: False)
+
+# # generic null arguments:
+# sbid: null #SBID of observation. (default: None)
+stokes: [I,Q,U] # List of Stokes parameters to image (default: ['I', 'Q', 'U'])
+epoch: 0 # Epoch of observation. (default: 0)
+# host: null # Host of mongodb (probably $hostname -i). (default: None)
+# username: null # Username of mongodb. (default: None)
+password: # Password of mongodb. (default: None)
+# limit: null # Limit the number of islands to process. (default: None)
+database: false # Add data to MongoDB. (default: False)
+
+# imaging arguments:
+# temp_dir_wsclean: null # Temporary directory for WSClean to store intermediate files (default: None)
+# temp_dir_images: null # Temporary directory for to store intermediate image files (default: None)
+# psf_cutoff: null # Cutoff for smoothing in units of arcseconds.  (default: None)
+robust: -0.5 # ROBUST
+nchan: 36 # NCHAN
+pols: IQU # POLS
+size: 6144 # SIZE
+scale: 2.5 # SCALE
+mgain: 0.7 # MGAIN
+niter: 500_000 # NITER
+nmiter: 15 # NMITER
+auto_mask: 4 # AUTO_MASK
+auto_threshold: 1 # AUTO_THRESHOLD
+local_rms: true #
+local_rms_window: 60 # LOCAL_RMS_WINDOW
+force_mask_rounds: 8 # FORCE_MASK_ROUNDS
+gridder: wgridder # {direct-ft,idg,wgridder,tuned-wgridder,wstacking}
+# taper: null # TAPER
+minuv: 200 # MINUV
+# parallel: null # PARALLEL
+mpi: false #                 Use MPI (default: False)
+purge: false # Purge intermediate files (default: False)
+multiscale: false # Use multiscale clean (default: False)
+# multiscale_scale_bias: null # The multiscale scale bias term provided to wsclean.  (default: None)
+multiscale_scales: 0,2,4,8,16,32,64,12 # The scales used in the multiscale clean.  (default: 0,2,4,8,16,32,64,128)
+# absmem: null # ABSMEM       Absolute memory limit in GB (default: None)
+make_residual_cubes: false # Create residual cubes as well as cubes from restored images.  (default: False)
+ms_glob_pattern: scienceData*_averaged_cal.leakage.ms # The pattern used to search for measurement sets.  (default: scienceData*_averaged_cal.leakage.ms)
+data_column: CORRECTED_DATA # Which column in the measurement set to image.  (default: CORRECTED_DATA)
+no_mf_weighting: false # Do not use multi-frequency weighting.  (default: False)
+skip_fix_ms: false # Do not apply the ASKAP MS corrections from the package fixms.  (default: False)
+
+# cutout arguments:
+pad: 3 # Number of beamwidths to pad around source [3]. (default: 3)
+dryrun: false # Do a dry-run [False]. (default: False)
+
+# # linmos null arguments:
+# holofile: null #Path to holography image (default: None)
+yanda: 1.3.0 # Yandasoft version to pull from DockerHub [1.3.0]. (default: 1.3.0)
+# yanda_image: null #Path to an existing yandasoft singularity container image.  (default: None)
+
+# frion arguments:
+ionex_server: ftp://ftp.aiub.unibe.ch/CODE/ # IONEX server (default: ftp://ftp.aiub.unibe.ch/CODE/)
+ionex_prefix: codg # IONEX_PREFIX
+# ionex_formatter: null # IONEX formatter. (default: ftp.aiub.unibe.ch)
+# ionex_proxy_server: null # Proxy server. (default: None)
+ionex_predownload: false # Pre-download IONEX files. (default: False)
+
+# common rm arguments:
+dimension: 1d # How many dimensions for RMsynth '1d' or '3d'. (default: 1d)
+save_plots: false #          save the plots. (default: False)
+rm_verbose: false #          Verbose RMsynth/RMClean. (default: False)
+
+# rm-synth arguments:
+ion: false # Use ionospheric-corrected data. (default: False)
+# tt0: null # TT0 MFS image -- will be used for model of Stokes I -- also needs --tt1. (default: None)
+# tt1: null # TT1 MFS image -- will be used for model of Stokes I -- also needs --tt0. (default: None)
+validate: false # Run on Stokes I. (default: False)
+own_fit: false # Use own Stokes I fit function. (default: False)
+weight_type: # weighting (inverse) 'variance' or 'uniform' (all 1s). (default: variance)
+fit_function: # Stokes I fitting function: 'linear' or 'log' polynomials. (default: log)
+fit_rmsf: false # Fit a Gaussian to the RMSF (default: False)
+# phi_max: null # Absolute max Faraday depth sampled (in rad/m^2) (overrides NSAMPLES). (default: None)
+# dphi: null # Width of Faraday depth channel. (default: None)
+n_samples: 5 # Number of samples across the FWHM RMSF. (default: 5)
+poly_ord: 3 # polynomial order to fit to I spectrum. (default: 3)
+no_stokes_i: false # ignore the Stokes I spectrum. (default: False)
+show_plots: false # show the plots. (default: False)
+not_rmsf: false # Skip calculation of RMSF? (default: False)
+debug: false # turn on debugging messages & plots. (default: False)
+
+# rm-clean arguments:
+cutoff: -8 # CLEAN cutoff (+ve = absolute, -ve = sigma). (default: -3)
+max_iter: 10000 # maximum number of CLEAN iterations. (default: 10000)
+gain: 0.1 # CLEAN loop gain. (default: 0.1)
+# window: null # Further CLEAN in mask to this threshold. (default: None)
+
+# catalogue arguments:
+leakage_degree: 4 # Degree of leakage polynomial fit. (default: 4)
+leakage_bins: 16 # Number of bins for leakage fit. (default: 16)
+leakage_snr: 30 # SNR cut for leakage fit. (default: 30.0)
+# write: null # File to save table to. (default: None)
+
+# cleanup arguments:
+overwrite: false # Overwrite existing tarball (default: False)

arrakis/cleanup.py

@@ -1,17 +1,18 @@
 #!/usr/bin/env python3
 """DANGER ZONE: Purge directories of un-needed FITS files."""
+import argparse
 import logging
 import tarfile
 from pathlib import Path
-from typing import List, Union
+from typing import List
 
 import astropy.units as u
 import numpy as np
-from prefect import flow, get_run_logger, task, unmapped
+from prefect import flow, get_run_logger, task
 from tqdm.auto import tqdm
 
 from arrakis.logger import TqdmToLogger, UltimateHelpFormatter, logger
-from arrakis.utils.pipeline import logo_str
+from arrakis.utils.pipeline import generic_parser, logo_str
 
 logger.setLevel(logging.INFO)
 
@@ -53,6 +54,10 @@ def make_cutout_tarball(cutdir: Path, overwrite: bool = False) -> Path:
     with tarfile.open(tarball, "w") as tar:
         for cutout in tqdm(all_things, file=TQDM_OUT, desc="Tarballing cutouts"):
             tar.add(cutout, arcname=cutout.name)
+
+    logger.info(f"Tarball created: {tarball}")
+    logger.critical(f"Removing {cutdir}")
+    cutdir.rmdir()
     return tarball
 
 
@@ -97,16 +102,11 @@ def main(
     for to_purge in tqdm(to_purge_all, file=TQDM_OUT, desc="Purging big beams"):
         purged.append(purge_cubelet_beams(to_purge))
     logger.info(f"Files purged: {len(purged)}")
-    total_purge_size = np.sum([p.stat().st_size for p in purged]) * u.byte
-    logger.info(f"Total space freed: {total_purge_size.to(u.GB)}")
 
     logger.info("Cleanup done!")
 
 
-def cli():
-    """Command-line interface"""
-    import argparse
-
+def cleanup_parser(parent_parser: bool = False) -> argparse.ArgumentParser:
     # Help string to be shown using the -h option
     descStr = f"""
     {logo_str}
@@ -117,33 +117,39 @@ def cli():
     """
 
     # Parse the command line options
-    parser = argparse.ArgumentParser(
-        description=descStr, formatter_class=UltimateHelpFormatter
-    )
-    parser.add_argument(
-        "outdir",
-        metavar="outdir",
-        type=Path,
-        help="Directory containing cutouts (in subdir outdir/cutouts).",
+    cleanup_parser = argparse.ArgumentParser(
+        add_help=not parent_parser,
+        description=descStr,
+        formatter_class=UltimateHelpFormatter,
     )
+    parser = cleanup_parser.add_argument_group("cleanup arguments")
     parser.add_argument(
-        "-o",
         "--overwrite",
         dest="overwrite",
         action="store_true",
         help="Overwrite existing tarball",
     )
-    parser.add_argument(
-        "-v", dest="verbose", action="store_true", help="Verbose output"
-    )
 
+    return cleanup_parser
+
+
+def cli():
+    """Command-line interface"""
+    gen_parser = generic_parser(parent_parser=True)
+    clean_parser = cleanup_parser(parent_parser=True)
+    parser = argparse.ArgumentParser(
+        parents=[gen_parser, clean_parser],
+        formatter_class=UltimateHelpFormatter,
+        description=clean_parser.description,
+    )
     args = parser.parse_args()
+
     verbose = args.verbose
 
     if verbose:
         logger.setLevel(logging.DEBUG)
 
-    main(datadir=Path(args.outdir), overwrite=args.overwrite)
+    main(datadir=Path(args.datadir), overwrite=args.overwrite)
 
 
 if __name__ == "__main__":

arrakis/configs/petrichor.yaml

@@ -7,7 +7,7 @@ cluster_kwargs:
     memory: "128GiB"
     account: 'OD-217087'
     #queue: 'workq'
-    walltime: '1-00:00:00'
+    walltime: '0-06:00:00'
     job_extra_directives: ['--qos express']
     # interface for the workers
     interface: "ib0"