Refactor of beamcon (#59)

* Refactor 2d * Start on 3d * 3D tests * 3D fixes * Update deps * No deps * Sanity * Absolute paths * Add cores limiter * Fix logic * Add some rand to test * Ruff * Formatting * Cleanup * Simplify smoother * Docs * Docs * Docs * Docs * Ruff * Version bump * Fix types * Add dels * All the executors * Update yaml * Add CLI arg * Docs * data types and free unused matrices in robust convolution * pass executor type to smooth_fits_cube * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * readme performance suggestion * Get logging working * ignore * Fixes * Bump python --------- Co-authored-by: Austin Shen <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
AlecThomson · Jul 31, 2024 · faa15e2 · faa15e2
1 parent d7d0ad8
commit faa15e2
Show file tree

Hide file tree

Showing 16 changed files with 1,469 additions and 1,181 deletions.
diff --git a/.gitignore b/.gitignore
@@ -141,3 +141,5 @@ mk_convolve
 beamlog.3d.astropy.txt
 beamlog.3d.robust.txt
 beamlog.3d.scipy.txt
+*.fits
+beamlog*
diff --git a/Dockerfile b/Dockerfile
@@ -9,7 +9,7 @@ RUN mkdir /tmp/numba_cache & chmod 777 /tmp/numba_cache & NUMBA_CACHE_DIR=/tmp/n
 ENV NUMBA_CACHE_DIR=/tmp/numba_cache
 COPY --chown=$MAMBA_USER:$MAMBA_USER . ./src
 RUN echo "Installing python and uv"
-RUN micromamba install python=3.8 uv -y -c conda-forge && \
+RUN micromamba install python=3.12 uv -y -c conda-forge && \
     micromamba clean --all --yes
 RUN echo "Installing RACS-tools"
 RUN micromamba run uv pip install ./src

diff --git a/Dockerfile.mpich b/Dockerfile.mpich
@@ -9,7 +9,7 @@ RUN mkdir /tmp/numba_cache & chmod 777 /tmp/numba_cache & NUMBA_CACHE_DIR=/tmp/n
 ENV NUMBA_CACHE_DIR=/tmp/numba_cache
 COPY --chown=$MAMBA_USER:$MAMBA_USER . ./src
 RUN echo "Installing python and uv"
-RUN micromamba install python=3.10 uv -y -c conda-forge && \
+RUN micromamba install python=3.12 uv -y -c conda-forge && \
     micromamba clean --all --yes
 RUN echo "Installing RACS-tools"
 RUN micromamba run uv pip install ./src[mpi]

diff --git a/README.md b/README.md
@@ -43,124 +43,102 @@ pip install git+https://github.com/AlecThomson/RACS-tools
 
 ```
 $ beamcon_2D -h
-usage: beamcon_2D [-h] [-p PREFIX] [-s SUFFIX] [-o OUTDIR] [--conv_mode {robust,scipy,astropy,astropy_fft}] [-v] [-d] [--bmaj BMAJ] [--bmin BMIN] [--bpa BPA] [--log LOG] [--logfile LOGFILE] [-c CUTOFF] [--circularise] [-t TOLERANCE] [-e EPSILON] [-n NSAMPS] [--ncores N_CORES | --mpi] infile [infile ...]
-
-    Smooth a field of 2D images to a common resolution.
-
-    - Parallelisation can run using multiprocessing or MPI.
-
-    - Default names of output files are /path/to/beamlog{infile//.fits/.{SUFFIX}.fits}
-
-    - By default, the smallest common beam will be automatically computed.
-    - Optionally, you can specify a target beam to use.
-
+usage: beamcon_2D [-h] [-p PREFIX] [-s SUFFIX] [-o OUTDIR] [--conv_mode {robust,scipy,astropy,astropy_fft}] [-v] [-d] [--bmaj BMAJ] [--bmin BMIN]
+                  [--bpa BPA] [--log LOG] [--logfile LOGFILE] [-c CUTOFF] [--circularise] [-t TOLERANCE] [-e EPSILON] [-n NSAMPS] [--ncores NCORES]
+                  [--executor {thread,process,mpi}]
+                  infile [infile ...]
 
+Smooth a field of 2D images to a common resolution. - Parallelisation can run using multiprocessing or MPI. - Default names of output files are
+/path/to/beamlog{infile//.fits/.{SUFFIX}.fits} - By default, the smallest common beam will be automatically computed. - Optionally, you can specify a
+target beam to use.
 
 positional arguments:
   infile                Input FITS image(s) to smooth (can be a wildcard) - beam info must be in header.
 
-optional arguments:
+options:
   -h, --help            show this help message and exit
   -p PREFIX, --prefix PREFIX
-                        Add prefix to output filenames.
+                        Add prefix to output filenames. (default: None)
   -s SUFFIX, --suffix SUFFIX
-                        Add suffix to output filenames [sm].
+                        Add suffix to output filenames [sm]. (default: sm)
   -o OUTDIR, --outdir OUTDIR
-                        Output directory of smoothed FITS image(s) [same as input file].
+                        Output directory of smoothed FITS image(s) [same as input file]. (default: None)
   --conv_mode {robust,scipy,astropy,astropy_fft}
-                        Which method to use for convolution [robust].
-                                'robust' computes the analytic FT of the convolving Gaussian.
-                                Note this mode can now handle NaNs in the data.
-                                Can also be 'scipy', 'astropy', or 'astropy_fft'.
-                                Note these other methods cannot cope well with small convolving beams.
-
-  -v, --verbosity       Increase output verbosity
-  -d, --dryrun          Compute common beam and stop [False].
-  --bmaj BMAJ           Target BMAJ (arcsec) to convolve to [None].
-  --bmin BMIN           Target BMIN (arcsec) to convolve to [None].
-  --bpa BPA             Target BPA (deg) to convolve to [None].
-  --log LOG             Name of beamlog file. If provided, save beamlog data to a file [None - not saved].
-  --logfile LOGFILE     Save logging output to file
+                        Which method to use for convolution [robust]. 'robust' computes the analytic FT of the convolving Gaussian. Note this mode can
+                        now handle NaNs in the data. Can also be 'scipy', 'astropy', or 'astropy_fft'. Note these other methods cannot cope well with
+                        small convolving beams. (default: robust)
+  -v, --verbosity       Increase output verbosity (default: 0)
+  -d, --dryrun          Compute common beam and stop [False]. (default: False)
+  --bmaj BMAJ           Target BMAJ (arcsec) to convolve to [None]. (default: None)
+  --bmin BMIN           Target BMIN (arcsec) to convolve to [None]. (default: None)
+  --bpa BPA             Target BPA (deg) to convolve to [None]. (default: None)
+  --log LOG             Name of beamlog file. If provided, save beamlog data to a file [None - not saved]. (default: None)
+  --logfile LOGFILE     Save logging output to file (default: None)
   -c CUTOFF, --cutoff CUTOFF
-                        Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit]
-  --circularise         Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0.
+                        Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit] (default: None)
+  --circularise         Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0. (default: False)
   -t TOLERANCE, --tolerance TOLERANCE
-                        tolerance for radio_beam.commonbeam.
+                        tolerance for radio_beam.commonbeam. (default: 0.0001)
   -e EPSILON, --epsilon EPSILON
-                        epsilon for radio_beam.commonbeam.
+                        epsilon for radio_beam.commonbeam. (default: 0.0005)
   -n NSAMPS, --nsamps NSAMPS
-                        nsamps for radio_beam.commonbeam.
-  --ncores N_CORES      Number of processes (uses multiprocessing).
-  --mpi                 Run with MPI.
+                        nsamps for radio_beam.commonbeam. (default: 200)
+  --ncores NCORES       Number of cores to use for parallelisation. If None, use all available cores. (default: None)
+  --executor {thread,process,mpi}
+                        Executor to use for parallelisation (default: thread)
 ```
 
 ```
 $ beamcon_3D -h
-usage: beamcon_3D [-h] [--uselogs] [--mode MODE] [--conv_mode {robust,scipy,astropy,astropy_fft}] [-v] [--logfile LOGFILE] [-d] [-p PREFIX] [-s SUFFIX] [-o OUTDIR] [--bmaj BMAJ] [--bmin BMIN] [--bpa BPA] [-c CUTOFF] [--circularise] [--ref_chan {first,last,mid}] [-t TOLERANCE] [-e EPSILON] [-n NSAMPS] infile [infile ...]
-
-    Smooth a field of 3D cubes to a common resolution.
-
-    - Parallelisation is done using MPI.
-
-    - Default names of output files are /path/to/beamlog{infile//.fits/.{SUFFIX}.fits}
-
-    - By default, the smallest common beam will be automatically computed.
-    - Optionally, you can specify a target beam to use.
-
-    - It is currently assumed that cubes will be 4D with a dummy Stokes axis.
-    - Iterating over Stokes axis is not yet supported.
-
+usage: beamcon_3D [-h] [--uselogs] [--mode MODE] [--conv_mode {robust,scipy,astropy,astropy_fft}] [-v] [--logfile LOGFILE] [-d] [-p PREFIX] [-s SUFFIX]
+                  [-o OUTDIR] [--bmaj BMAJ] [--bmin BMIN] [--bpa BPA] [-c CUTOFF] [--circularise] [--ref_chan {first,last,mid}] [-t TOLERANCE]
+                  [-e EPSILON] [-n NSAMPS] [--ncores NCORES] [--executor_type {thread,process,mpi}]
+                  infile [infile ...]
 
+Smooth a field of 3D cubes to a common resolution. - Default names of output files are /path/to/beamlog{infile//.fits/.{SUFFIX}.fits} - By default, the
+smallest common beam will be automatically computed. - Optionally, you can specify a target beam to use. - It is currently assumed that cubes will be
+4D with a dummy Stokes axis. - Iterating over Stokes axis is not yet supported.
 
 positional arguments:
-  infile                Input FITS image(s) to smooth (can be a wildcard)
-                                - CASA beamtable will be used if present i.e. if CASAMBM = T
-                                - Otherwise beam info must be in co-located beamlog files.
-                                - beamlog must have the name /path/to/beamlog{infile//.fits/.txt}
-
+  infile                Input FITS image(s) to smooth (can be a wildcard) - CASA beamtable will be used if present i.e. if CASAMBM = T - Otherwise beam
+                        info must be in co-located beamlog files. - beamlog must have the name /path/to/beamlog{infile//.fits/.txt}
 
-optional arguments:
+options:
   -h, --help            show this help message and exit
-  --uselogs             Get convolving information from previous run [False].
-  --mode MODE           Common resolution mode [natural].
-                                natural -- allow frequency variation.
-                                total -- smooth all plans to a common resolution.
-
+  --uselogs             Get convolving information from previous run [False]. (default: False)
+  --mode MODE           Common resolution mode [natural]. natural -- allow frequency variation. total -- smooth all plans to a common resolution.
+                        (default: natural)
   --conv_mode {robust,scipy,astropy,astropy_fft}
-                        Which method to use for convolution [robust].
-                                'robust' computes the analytic FT of the convolving Gaussian.
-                                Note this mode can now handle NaNs in the data.
-                                Can also be 'scipy', 'astropy', or 'astropy_fft'.
-                                Note these other methods cannot cope well with small convolving beams.
-
-  -v, --verbosity       Increase output verbosity
-  --logfile LOGFILE     Save logging output to file
-  -d, --dryrun          Compute common beam and stop [False].
+                        Which method to use for convolution [robust]. 'robust' computes the analytic FT of the convolving Gaussian. Note this mode can
+                        now handle NaNs in the data. Can also be 'scipy', 'astropy', or 'astropy_fft'. Note these other methods cannot cope well with
+                        small convolving beams. (default: robust)
+  -v, --verbosity       Increase output verbosity (default: 0)
+  --logfile LOGFILE     Save logging output to file (default: None)
+  -d, --dryrun          Compute common beam and stop. (default: False)
   -p PREFIX, --prefix PREFIX
-                        Add prefix to output filenames.
+                        Add prefix to output filenames. (default: None)
   -s SUFFIX, --suffix SUFFIX
-                        Add suffix to output filenames [{MODE}].
+                        Add suffix to output filenames [{MODE}]. (default: None)
   -o OUTDIR, --outdir OUTDIR
-                        Output directory of smoothed FITS image(s) [None - same as input].
-  --bmaj BMAJ           BMAJ to convolve to [max BMAJ from given image(s)].
-  --bmin BMIN           BMIN to convolve to [max BMAJ from given image(s)].
-  --bpa BPA             BPA to convolve to [0].
+                        Output directory of smoothed FITS image(s) [None - same as input]. (default: None)
+  --bmaj BMAJ           BMAJ to convolve to [max BMAJ from given image(s)]. (default: None)
+  --bmin BMIN           BMIN to convolve to [max BMAJ from given image(s)]. (default: None)
+  --bpa BPA             BPA to convolve to [0]. (default: None)
   -c CUTOFF, --cutoff CUTOFF
-                        Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit]
-  --circularise         Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0.
+                        Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit] (default: None)
+  --circularise         Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0. (default: False)
   --ref_chan {first,last,mid}
-                        Reference psf for header [None].
-                                    first  -- use psf for first frequency channel.
-                                    last -- use psf for the last frequency channel.
-                                    mid -- use psf for the centre frequency channel.
-                                    Will use the CRPIX channel if not set.
-
+                        Reference psf for header [None]. first -- use psf for first frequency channel. last -- use psf for the last frequency channel.
+                        mid -- use psf for the centre frequency channel. Will use the CRPIX channel if not set. (default: None)
   -t TOLERANCE, --tolerance TOLERANCE
-                        tolerance for radio_beam.commonbeam.
+                        tolerance for radio_beam.commonbeam. (default: 0.0001)
   -e EPSILON, --epsilon EPSILON
-                        epsilon for radio_beam.commonbeam.
+                        epsilon for radio_beam.commonbeam. (default: 0.0005)
   -n NSAMPS, --nsamps NSAMPS
-                        nsamps for radio_beam.commonbeam.
+                        nsamps for radio_beam.commonbeam. (default: 200)
+  --ncores NCORES       Number of cores to use for parallelisation. If None, use all available cores. (default: None)
+  --executor_type {thread,process,mpi}
+                        Executor type for parallelisation. (default: thread)
 ```
 
 ```
@@ -186,6 +164,19 @@ options:
 
 If finding a common beam fails, try tweaking the `tolerance`, `epsilon`, and `nsamps` parameters. See [radio-beam](https://radio-beam.readthedocs.io/en/latest/) for more details.
 
+## Performance
+
+Profiling for `beamcon_3D` suggests this program requires a minimum of ~15X the memory of a data cube slice per process to perform convolution to a common beam. So for a 800 MB slice (e.g. typical POSSUM cube) you would want to allow 15 GB memory per worker (I use 20 GB). Choose `ncores` appropriately given your machine memory availability and this limit to ensure optimal performance with multiprocessing.
+
+An example slurm header for `beamcon_3D`:
+
+```
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=<ncores>
+#SBATCH --mem-per-cpu=20G
+```
+
 ## Contributing
 Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
 

diff --git a/environment.yml b/environment.yml
@@ -5,16 +5,14 @@ channels:
 - defaults
 - pkgw-forge
 dependencies:
-- python=3.8
+- python=3.10
 - pip
 - numpy
 - scipy
 - astropy
-- gfortran
 - mpi4py
 - pip:
   - spectral_cube
   - radio_beam
-  - schwimmbad
   - tqdm
   - ./
diff --git a/pyproject.toml b/pyproject.toml
@@ -12,12 +12,11 @@ python = ">=3.8"
 numpy = "<2"
 astropy = ">=5"
 radio_beam = "*"
-schwimmbad = "*"
 scipy = "*"
 spectral_cube = ">=0.6.3"
 tqdm = "*"
 numba = "*"
-mpi4py = {version = "*", optional = true}
+mpi4py = {version = ">=3", optional = true}
 
 [tool.poetry.dev-dependencies]
 black = "*"

diff --git a/racs_tools/au2.py b/racs_tools/au2.py
@@ -162,6 +162,6 @@ def gauss_factor(
         * bmin2
         / math.sqrt(alpha * beta - 0.25 * gamma * gamma)
     )
-    fac = ((math.sqrt(dx1**2) * math.sqrt(dy1**2))) / amp
+    fac = (math.sqrt(dx1**2) * math.sqrt(dy1**2)) / amp
 
     return fac, amp, bmaj, bmin, np.degrees(bpa)