Skip to content

Commit faa15e2

Browse files
AlecThomsonaxshenpre-commit-ci[bot]
authored
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>
1 parent d7d0ad8 commit faa15e2

File tree

16 files changed

+1469
-1181
lines changed

16 files changed

+1469
-1181
lines changed

‎.gitignore‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -141,3 +141,5 @@ mk_convolve
141141
beamlog.3d.astropy.txt
142142
beamlog.3d.robust.txt
143143
beamlog.3d.scipy.txt
144+
*.fits
145+
beamlog*

‎Dockerfile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ RUN mkdir /tmp/numba_cache & chmod 777 /tmp/numba_cache & NUMBA_CACHE_DIR=/tmp/n
99
ENV NUMBA_CACHE_DIR=/tmp/numba_cache
1010
COPY --chown=$MAMBA_USER:$MAMBA_USER . ./src
1111
RUN echo "Installing python and uv"
12-
RUN micromamba install python=3.8 uv -y -c conda-forge && \
12+
RUN micromamba install python=3.12 uv -y -c conda-forge && \
1313
micromamba clean --all --yes
1414
RUN echo "Installing RACS-tools"
1515
RUN micromamba run uv pip install ./src

‎Dockerfile.mpich‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ RUN mkdir /tmp/numba_cache & chmod 777 /tmp/numba_cache & NUMBA_CACHE_DIR=/tmp/n
99
ENV NUMBA_CACHE_DIR=/tmp/numba_cache
1010
COPY --chown=$MAMBA_USER:$MAMBA_USER . ./src
1111
RUN echo "Installing python and uv"
12-
RUN micromamba install python=3.10 uv -y -c conda-forge && \
12+
RUN micromamba install python=3.12 uv -y -c conda-forge && \
1313
micromamba clean --all --yes
1414
RUN echo "Installing RACS-tools"
1515
RUN micromamba run uv pip install ./src[mpi]

‎README.md‎

Lines changed: 77 additions & 86 deletions
Original file line numberDiff line numberDiff line change
@@ -43,124 +43,102 @@ pip install git+https://github.com/AlecThomson/RACS-tools
4343

4444
```
4545
$ beamcon_2D -h
46-
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 ...]
47-
48-
Smooth a field of 2D images to a common resolution.
49-
50-
- Parallelisation can run using multiprocessing or MPI.
51-
52-
- Default names of output files are /path/to/beamlog{infile//.fits/.{SUFFIX}.fits}
53-
54-
- By default, the smallest common beam will be automatically computed.
55-
- Optionally, you can specify a target beam to use.
56-
46+
usage: beamcon_2D [-h] [-p PREFIX] [-s SUFFIX] [-o OUTDIR] [--conv_mode {robust,scipy,astropy,astropy_fft}] [-v] [-d] [--bmaj BMAJ] [--bmin BMIN]
47+
[--bpa BPA] [--log LOG] [--logfile LOGFILE] [-c CUTOFF] [--circularise] [-t TOLERANCE] [-e EPSILON] [-n NSAMPS] [--ncores NCORES]
48+
[--executor {thread,process,mpi}]
49+
infile [infile ...]
5750
51+
Smooth a field of 2D images to a common resolution. - Parallelisation can run using multiprocessing or MPI. - Default names of output files are
52+
/path/to/beamlog{infile//.fits/.{SUFFIX}.fits} - By default, the smallest common beam will be automatically computed. - Optionally, you can specify a
53+
target beam to use.
5854
5955
positional arguments:
6056
infile Input FITS image(s) to smooth (can be a wildcard) - beam info must be in header.
6157
62-
optional arguments:
58+
options:
6359
-h, --help show this help message and exit
6460
-p PREFIX, --prefix PREFIX
65-
Add prefix to output filenames.
61+
Add prefix to output filenames. (default: None)
6662
-s SUFFIX, --suffix SUFFIX
67-
Add suffix to output filenames [sm].
63+
Add suffix to output filenames [sm]. (default: sm)
6864
-o OUTDIR, --outdir OUTDIR
69-
Output directory of smoothed FITS image(s) [same as input file].
65+
Output directory of smoothed FITS image(s) [same as input file]. (default: None)
7066
--conv_mode {robust,scipy,astropy,astropy_fft}
71-
Which method to use for convolution [robust].
72-
'robust' computes the analytic FT of the convolving Gaussian.
73-
Note this mode can now handle NaNs in the data.
74-
Can also be 'scipy', 'astropy', or 'astropy_fft'.
75-
Note these other methods cannot cope well with small convolving beams.
76-
77-
-v, --verbosity Increase output verbosity
78-
-d, --dryrun Compute common beam and stop [False].
79-
--bmaj BMAJ Target BMAJ (arcsec) to convolve to [None].
80-
--bmin BMIN Target BMIN (arcsec) to convolve to [None].
81-
--bpa BPA Target BPA (deg) to convolve to [None].
82-
--log LOG Name of beamlog file. If provided, save beamlog data to a file [None - not saved].
83-
--logfile LOGFILE Save logging output to file
67+
Which method to use for convolution [robust]. 'robust' computes the analytic FT of the convolving Gaussian. Note this mode can
68+
now handle NaNs in the data. Can also be 'scipy', 'astropy', or 'astropy_fft'. Note these other methods cannot cope well with
69+
small convolving beams. (default: robust)
70+
-v, --verbosity Increase output verbosity (default: 0)
71+
-d, --dryrun Compute common beam and stop [False]. (default: False)
72+
--bmaj BMAJ Target BMAJ (arcsec) to convolve to [None]. (default: None)
73+
--bmin BMIN Target BMIN (arcsec) to convolve to [None]. (default: None)
74+
--bpa BPA Target BPA (deg) to convolve to [None]. (default: None)
75+
--log LOG Name of beamlog file. If provided, save beamlog data to a file [None - not saved]. (default: None)
76+
--logfile LOGFILE Save logging output to file (default: None)
8477
-c CUTOFF, --cutoff CUTOFF
85-
Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit]
86-
--circularise Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0.
78+
Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit] (default: None)
79+
--circularise Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0. (default: False)
8780
-t TOLERANCE, --tolerance TOLERANCE
88-
tolerance for radio_beam.commonbeam.
81+
tolerance for radio_beam.commonbeam. (default: 0.0001)
8982
-e EPSILON, --epsilon EPSILON
90-
epsilon for radio_beam.commonbeam.
83+
epsilon for radio_beam.commonbeam. (default: 0.0005)
9184
-n NSAMPS, --nsamps NSAMPS
92-
nsamps for radio_beam.commonbeam.
93-
--ncores N_CORES Number of processes (uses multiprocessing).
94-
--mpi Run with MPI.
85+
nsamps for radio_beam.commonbeam. (default: 200)
86+
--ncores NCORES Number of cores to use for parallelisation. If None, use all available cores. (default: None)
87+
--executor {thread,process,mpi}
88+
Executor to use for parallelisation (default: thread)
9589
```
9690

9791
```
9892
$ beamcon_3D -h
99-
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 ...]
100-
101-
Smooth a field of 3D cubes to a common resolution.
102-
103-
- Parallelisation is done using MPI.
104-
105-
- Default names of output files are /path/to/beamlog{infile//.fits/.{SUFFIX}.fits}
106-
107-
- By default, the smallest common beam will be automatically computed.
108-
- Optionally, you can specify a target beam to use.
109-
110-
- It is currently assumed that cubes will be 4D with a dummy Stokes axis.
111-
- Iterating over Stokes axis is not yet supported.
112-
93+
usage: beamcon_3D [-h] [--uselogs] [--mode MODE] [--conv_mode {robust,scipy,astropy,astropy_fft}] [-v] [--logfile LOGFILE] [-d] [-p PREFIX] [-s SUFFIX]
94+
[-o OUTDIR] [--bmaj BMAJ] [--bmin BMIN] [--bpa BPA] [-c CUTOFF] [--circularise] [--ref_chan {first,last,mid}] [-t TOLERANCE]
95+
[-e EPSILON] [-n NSAMPS] [--ncores NCORES] [--executor_type {thread,process,mpi}]
96+
infile [infile ...]
11397
98+
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
99+
smallest common beam will be automatically computed. - Optionally, you can specify a target beam to use. - It is currently assumed that cubes will be
100+
4D with a dummy Stokes axis. - Iterating over Stokes axis is not yet supported.
114101
115102
positional arguments:
116-
infile Input FITS image(s) to smooth (can be a wildcard)
117-
- CASA beamtable will be used if present i.e. if CASAMBM = T
118-
- Otherwise beam info must be in co-located beamlog files.
119-
- beamlog must have the name /path/to/beamlog{infile//.fits/.txt}
120-
103+
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
104+
info must be in co-located beamlog files. - beamlog must have the name /path/to/beamlog{infile//.fits/.txt}
121105
122-
optional arguments:
106+
options:
123107
-h, --help show this help message and exit
124-
--uselogs Get convolving information from previous run [False].
125-
--mode MODE Common resolution mode [natural].
126-
natural -- allow frequency variation.
127-
total -- smooth all plans to a common resolution.
128-
108+
--uselogs Get convolving information from previous run [False]. (default: False)
109+
--mode MODE Common resolution mode [natural]. natural -- allow frequency variation. total -- smooth all plans to a common resolution.
110+
(default: natural)
129111
--conv_mode {robust,scipy,astropy,astropy_fft}
130-
Which method to use for convolution [robust].
131-
'robust' computes the analytic FT of the convolving Gaussian.
132-
Note this mode can now handle NaNs in the data.
133-
Can also be 'scipy', 'astropy', or 'astropy_fft'.
134-
Note these other methods cannot cope well with small convolving beams.
135-
136-
-v, --verbosity Increase output verbosity
137-
--logfile LOGFILE Save logging output to file
138-
-d, --dryrun Compute common beam and stop [False].
112+
Which method to use for convolution [robust]. 'robust' computes the analytic FT of the convolving Gaussian. Note this mode can
113+
now handle NaNs in the data. Can also be 'scipy', 'astropy', or 'astropy_fft'. Note these other methods cannot cope well with
114+
small convolving beams. (default: robust)
115+
-v, --verbosity Increase output verbosity (default: 0)
116+
--logfile LOGFILE Save logging output to file (default: None)
117+
-d, --dryrun Compute common beam and stop. (default: False)
139118
-p PREFIX, --prefix PREFIX
140-
Add prefix to output filenames.
119+
Add prefix to output filenames. (default: None)
141120
-s SUFFIX, --suffix SUFFIX
142-
Add suffix to output filenames [{MODE}].
121+
Add suffix to output filenames [{MODE}]. (default: None)
143122
-o OUTDIR, --outdir OUTDIR
144-
Output directory of smoothed FITS image(s) [None - same as input].
145-
--bmaj BMAJ BMAJ to convolve to [max BMAJ from given image(s)].
146-
--bmin BMIN BMIN to convolve to [max BMAJ from given image(s)].
147-
--bpa BPA BPA to convolve to [0].
123+
Output directory of smoothed FITS image(s) [None - same as input]. (default: None)
124+
--bmaj BMAJ BMAJ to convolve to [max BMAJ from given image(s)]. (default: None)
125+
--bmin BMIN BMIN to convolve to [max BMAJ from given image(s)]. (default: None)
126+
--bpa BPA BPA to convolve to [0]. (default: None)
148127
-c CUTOFF, --cutoff CUTOFF
149-
Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit]
150-
--circularise Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0.
128+
Cutoff BMAJ value (arcsec) -- Blank channels with BMAJ larger than this [None -- no limit] (default: None)
129+
--circularise Circularise the final PSF -- Sets the BMIN = BMAJ, and BPA=0. (default: False)
151130
--ref_chan {first,last,mid}
152-
Reference psf for header [None].
153-
first -- use psf for first frequency channel.
154-
last -- use psf for the last frequency channel.
155-
mid -- use psf for the centre frequency channel.
156-
Will use the CRPIX channel if not set.
157-
131+
Reference psf for header [None]. first -- use psf for first frequency channel. last -- use psf for the last frequency channel.
132+
mid -- use psf for the centre frequency channel. Will use the CRPIX channel if not set. (default: None)
158133
-t TOLERANCE, --tolerance TOLERANCE
159-
tolerance for radio_beam.commonbeam.
134+
tolerance for radio_beam.commonbeam. (default: 0.0001)
160135
-e EPSILON, --epsilon EPSILON
161-
epsilon for radio_beam.commonbeam.
136+
epsilon for radio_beam.commonbeam. (default: 0.0005)
162137
-n NSAMPS, --nsamps NSAMPS
163-
nsamps for radio_beam.commonbeam.
138+
nsamps for radio_beam.commonbeam. (default: 200)
139+
--ncores NCORES Number of cores to use for parallelisation. If None, use all available cores. (default: None)
140+
--executor_type {thread,process,mpi}
141+
Executor type for parallelisation. (default: thread)
164142
```
165143

166144
```
@@ -186,6 +164,19 @@ options:
186164

187165
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.
188166

167+
## Performance
168+
169+
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.
170+
171+
An example slurm header for `beamcon_3D`:
172+
173+
```
174+
#SBATCH --nodes=1
175+
#SBATCH --ntasks-per-node=1
176+
#SBATCH --cpus-per-task=<ncores>
177+
#SBATCH --mem-per-cpu=20G
178+
```
179+
189180
## Contributing
190181
Pull requests are welcome. For major changes, please open an issue first to discuss what you would like to change.
191182

‎environment.yml‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,16 +5,14 @@ channels:
55
- defaults
66
- pkgw-forge
77
dependencies:
8-
- python=3.8
8+
- python=3.10
99
- pip
1010
- numpy
1111
- scipy
1212
- astropy
13-
- gfortran
1413
- mpi4py
1514
- pip:
1615
- spectral_cube
1716
- radio_beam
18-
- schwimmbad
1917
- tqdm
2018
- ./

‎pyproject.toml‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,11 @@ python = ">=3.8"
1212
numpy = "<2"
1313
astropy = ">=5"
1414
radio_beam = "*"
15-
schwimmbad = "*"
1615
scipy = "*"
1716
spectral_cube = ">=0.6.3"
1817
tqdm = "*"
1918
numba = "*"
20-
mpi4py = {version = "*", optional = true}
19+
mpi4py = {version = ">=3", optional = true}
2120

2221
[tool.poetry.dev-dependencies]
2322
black = "*"

‎racs_tools/au2.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,6 +162,6 @@ def gauss_factor(
162162
* bmin2
163163
/ math.sqrt(alpha * beta - 0.25 * gamma * gamma)
164164
)
165-
fac = ((math.sqrt(dx1**2) * math.sqrt(dy1**2))) / amp
165+
fac = (math.sqrt(dx1**2) * math.sqrt(dy1**2)) / amp
166166

167167
return fac, amp, bmaj, bmin, np.degrees(bpa)

0 commit comments

Comments
 (0)