From e5183135613a173e6ccb206e7078785e04226a24 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Thu, 27 Oct 2022 15:56:25 -0400 Subject: [PATCH 01/21] Added obscured_airy_disk() function and added PBCorrectedCube to SimpleNet, currently can only handle one channel --- src/mpol/images.py | 117 +++++++++++++++++++++++++++++++++++++++- src/mpol/precomposed.py | 6 +++ 2 files changed, 122 insertions(+), 1 deletion(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index 63e35df1..bd2ec9b7 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -1,6 +1,7 @@ r"""The ``images`` module provides the core functionality of MPoL via :class:`mpol.images.ImageCube`.""" import numpy as np +from scipy.special import j1 import torch import torch.fft # to avoid conflicts with old torch.fft *function* from torch import nn @@ -279,6 +280,120 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul.close() +class PBCorrectedCube(nn.Module): + r""" + A ImageCube that has been corrected for the expected primary beam. Currently only implemented for ALMA. + + Args: + cell_size (float): the width of a pixel [arcseconds] + npix (int): the number of pixels per image side + coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. + nchan (int): the number of channels in the image + """ + def __init__( + self, + cell_size=None, + npix=None, + coords=None, + chan_freqs=None, + telescope="ALMA_12" + ): + nchan = len(np.at_least_1d(chan_freqs)) + super().__init__() + _setup_coords(self, cell_size, npix, coords, nchan) + + self.telescope_to_radius = {"ALMA_12": 6.} # in meters + self.telescope_to_blocked_radius = {"ALMA_12": 0.375} + + assert (telescope in self.telescope_to_radius) and (telescope in self.telescope_to_blocked_radius) + + self.pb_mask = self.obscured_airy_disk(chan_freqs, telescope) + + + def forward(self, cube): + r"""Args: + cube (torch.double tensor, of shape ``(nchan, npix, npix)``): a prepacked image cube, for example, from ImageCube.forward() + + Returns: + (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. + """ + return self.pb_mask * cube + + def obscured_airy_disk(self, telescope): + r""" + Generates airy disk primary beam correction mask assuming some is obscured + by secondary reflector. + """ + + R = telescope_to_radius[telescope] + R_obs = telescope_to_blocked_radius[telescope] + + r_coords = np.sqrt(self.packed_x_centers_2D**2 + self.packed_y_centers_2D**2) # units of arcsec + r_coords_rads = r_coords * np.pi/648000. # radians + wl = 2.998e8 / self.chan_freqs[0] # TODO: handle multiple channels + x = np.pi * R * r_coords_rads / wl + + eps = R_obs / R + j1_x = scipy.special.j1(x) + j1_epsx = scipy.special.j1(eps * x) + mask = 4. * (j1_x / x - eps * j1_epsx / x)**2 / (1. - eps**2)**2 + + return mask + + @property + def sky_cube(self): + """ + The image cube arranged as it would appear on the sky. + + Returns: + torch.double : 3D image cube of shape ``(nchan, npix, npix)`` + + """ + return utils.packed_cube_to_sky_cube(self.cube) + + def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): + """ + Export the image cube to a FITS file. + + Args: + fname (str): the name of the FITS file to export to. + overwrite (bool): if the file already exists, overwrite? + header_kwargs (dict): Extra keyword arguments to write to the FITS header. + + Returns: + None + """ + + try: + from astropy import wcs + from astropy.io import fits + except ImportError: + print( + "Please install the astropy package to use FITS export functionality." + ) + + w = wcs.WCS(naxis=2) + + w.wcs.crpix = np.array([1, 1]) + w.wcs.cdelt = ( + np.array([self.coords.cell_size, self.coords.cell_size]) / 3600 + ) # decimal degrees + w.wcs.ctype = ["RA---TAN", "DEC--TAN"] + + header = w.to_header() + + # add in the kwargs to the header + if header_kwargs is not None: + for k, v in header_kwargs.items(): + header[k] = v + + hdu = fits.PrimaryHDU(self.sky_cube.detach().cpu().numpy(), header=header) + + hdul = fits.HDUList([hdu]) + hdul.writeto(fname, overwrite=overwrite) + + hdul.close() + class FourierCube(nn.Module): r""" @@ -292,7 +407,7 @@ class FourierCube(nn.Module): def __init__(self, cell_size=None, npix=None, coords=None): - super().__init__() + super().__init__(cell_size) # we don't want to bother with the nchan argument here, so # we don't use the convenience method _setup_coords diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index 1849f3ce..e2ed48cc 100644 --- a/src/mpol/precomposed.py +++ b/src/mpol/precomposed.py @@ -50,6 +50,11 @@ def __init__( self.icube = images.ImageCube( coords=self.coords, nchan=self.nchan, passthrough=True ) + + self.pbcube = images.PBCorrectedCube( + coords = self.coords, nchan=self.nchan, telescope="ALMA_12" + ) + self.fcube = images.FourierCube(coords=self.coords) def forward(self): @@ -61,5 +66,6 @@ def forward(self): x = self.bcube.forward() x = self.conv_layer(x) x = self.icube.forward(x) + x = self.pbcube.forward(x) vis = self.fcube.forward(x) return vis From 1149d735af7668f9b37f0fa25639617941dcd6f1 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sat, 26 Nov 2022 18:40:34 -0500 Subject: [PATCH 02/21] Modified gridding checks to include optional frequency array. --- src/mpol/gridding.py | 43 ++++++++++++++++++++++++++++++++++++++----- 1 file changed, 38 insertions(+), 5 deletions(-) diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index dc731ac7..84d28ae9 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -6,7 +6,7 @@ from .datasets import GriddedDataset -def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=None): +def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=None, freq=None): """ Check that all data inputs are the same shape, the weights are positive, and the data_re and data_im are floats. @@ -17,7 +17,11 @@ def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=N assert ( uu.ndim == 2 or uu.ndim == 1 ), "Input data vectors should be either 1D or 2D numpy arrays." + shape = uu.shape + + if freq is not None: + assert len(u) == len(freq), "uu must have same number of channels as freq array." for a in [vv, weight, data_re, data_im]: assert ( @@ -40,11 +44,36 @@ def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=N data_re = np.atleast_2d(data_re) data_im = np.atleast_2d(data_im) - return uu, vv, weight, data_re, data_im + return uu, vv, weight, data_re, data_im, freq # expand to 2d with complex conjugates +def _check_freq_1d(freq=None): + """ + Check that the frequency input array contains only positive floats. + + If the user supplied a float, convert to a 1D array. If no frequency array + was supplied, simply skip. + """ + if freq is None: + return freq + + assert ( + np.isscalar(freq) or freq.ndim == 1 + ), "Input data vectors should be either None, scalar, or 1D array." + + assert np.all(freq > 0.0), "Not all frequencies are positive, check inputs." + + if np.isscalar(freq): + freq = np.atleast_1d(freq) + + assert (freq.dtype == np.single) or ( + freq.dtype == np.double + ), "freq should be type single or double" + + return freq + class Gridder: r""" The Gridder object uses desired image dimensions (via the ``cell_size`` and ``npix`` arguments) to define a corresponding Fourier plane grid as a :class:`.GridCoords` object. A pre-computed :class:`.GridCoords` can be supplied in lieu of ``cell_size`` and ``npix``, but all three arguments should never be supplied at once. For more details on the properties of the grid that is created, see the :class:`.GridCoords` documentation. @@ -77,13 +106,17 @@ def __init__( weight=None, data_re=None, data_im=None, + freq=None, ): + # check frequency array is 1d or None, expand if not + freq = _check_freq_1d(freq) + # check everything should be 2d, expand if not - uu, vv, weight, data_re, data_im = _check_data_inputs_2d( - uu, vv, weight, data_re, data_im + uu, vv, weight, data_re, data_im, freq = _check_data_inputs_2d( + uu, vv, weight, data_re, data_im, freq ) - + # setup the coordinates object nchan = len(uu) _setup_coords(self, cell_size, npix, coords, nchan) From da9fdad893aebb876ceb02a73ec8f59697ad672e Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sat, 26 Nov 2022 18:48:03 -0500 Subject: [PATCH 03/21] Added class attribute chan_freq, changed naming to better differentiate between channel and spatial frequencies --- src/mpol/gridding.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index 84d28ae9..ebeb4019 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -106,15 +106,15 @@ def __init__( weight=None, data_re=None, data_im=None, - freq=None, + chan_freq=None, ): # check frequency array is 1d or None, expand if not - freq = _check_freq_1d(freq) + chan_freq = _check_freq_1d(chan_freq) # check everything should be 2d, expand if not - uu, vv, weight, data_re, data_im, freq = _check_data_inputs_2d( - uu, vv, weight, data_re, data_im, freq + uu, vv, weight, data_re, data_im, chan_freq = _check_data_inputs_2d( + uu, vv, weight, data_re, data_im, chan_freq ) # setup the coordinates object @@ -133,6 +133,7 @@ def __init__( self.weight = np.concatenate([weight, weight], axis=1) self.data_re = np.concatenate([data_re, data_re], axis=1) self.data_im = np.concatenate([data_im, -data_im], axis=1) + self.chan_freq = chan_freq # figure out which visibility cell each datapoint lands in, so that # we can later assign it the appropriate robust weight for that cell From 7d8a9d0e0ce3d68b7ec7c5d767bbd31e151f6d9b Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 00:53:17 -0500 Subject: [PATCH 04/21] Set up framework for PrimaryBeamCube and dish types --- src/mpol/images.py | 77 ++++++++++++++++++++++++----------------- src/mpol/precomposed.py | 14 ++++++-- 2 files changed, 57 insertions(+), 34 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index bd2ec9b7..0a922f6a 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -8,7 +8,7 @@ from . import utils from .coordinates import GridCoords -from .gridding import _setup_coords +from .gridding import _setup_coords, _check_freq_1d class BaseCube(nn.Module): @@ -280,34 +280,55 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul.close() -class PBCorrectedCube(nn.Module): +class PrimaryBeamCube(nn.Module): r""" - A ImageCube that has been corrected for the expected primary beam. Currently only implemented for ALMA. + A ImageCube representing the primary beam of a described dish type. Currently can correct for a + uniform or center-obscured dish. The forward() method multiplies an image cube by this primary beam mask. Args: cell_size (float): the width of a pixel [arcseconds] npix (int): the number of pixels per image side coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. nchan (int): the number of channels in the image + dish_type (string): the type of dish to correct for. Either 'uniform' or 'obscured'. + dish_radius (float): the radius of the dish (in meters) + dish_kwargs (dict): any additional arguments needed for special dish types. Currently only uses: + dish_obscured_radius (float): the radius of the obscured portion of the dish """ def __init__( self, cell_size=None, npix=None, coords=None, + nchan=None, chan_freqs=None, - telescope="ALMA_12" + dish_type=None, + dish_radius=None, + **dish_kwargs, ): - nchan = len(np.at_least_1d(chan_freqs)) super().__init__() _setup_coords(self, cell_size, npix, coords, nchan) - self.telescope_to_radius = {"ALMA_12": 6.} # in meters - self.telescope_to_blocked_radius = {"ALMA_12": 0.375} + _check_freq_1d(chan_freqs) + assert (chan_freqs is None) or (len(chan_freqs) == nchan), "Length of chan_freqs must be equal to nchan" + + assert (dish_type is None) or (dish_type in ["uniform", "obscured"]), "Provided dish_type must be 'uniform' or 'obscured'" - assert (telescope in self.telescope_to_radius) and (telescope in self.telescope_to_blocked_radius) + self.default_mask = nn.Parameter( + torch.full( + (self.nchan, self.coords.npix, self.coords.npix), + fill_value=1.0, + requires_grad=False, + dtype=torch.double, + ) + ) - self.pb_mask = self.obscured_airy_disk(chan_freqs, telescope) + if dish_type is None: + self.pb_mask = self.default_mask + elif dish_type == "uniform": + self.pb_mask = self.uniform_mask(chan_freqs, dish_radius) + elif dish_type == "obscured": + self.pb_mask = self.obscured_mask(chan_freqs, dish_radius, **dish_kwargs) def forward(self, cube): @@ -317,43 +338,35 @@ def forward(self, cube): Returns: (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. """ - return self.pb_mask * cube + return torch.mul(self.pb_mask, cube) - def obscured_airy_disk(self, telescope): + def uniform_mask(self, chan_freqs, dish_radius): r""" - Generates airy disk primary beam correction mask assuming some is obscured - by secondary reflector. + Generates airy disk primary beam correction mask. """ - - R = telescope_to_radius[telescope] - R_obs = telescope_to_blocked_radius[telescope] - - r_coords = np.sqrt(self.packed_x_centers_2D**2 + self.packed_y_centers_2D**2) # units of arcsec - r_coords_rads = r_coords * np.pi/648000. # radians - wl = 2.998e8 / self.chan_freqs[0] # TODO: handle multiple channels - x = np.pi * R * r_coords_rads / wl - - eps = R_obs / R - j1_x = scipy.special.j1(x) - j1_epsx = scipy.special.j1(eps * x) - mask = 4. * (j1_x / x - eps * j1_epsx / x)**2 / (1. - eps**2)**2 - - return mask + return self.default_mask + + def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): + r""" + Generates airy disk primary beam correction mask. + """ + assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" + return self.default_mask @property def sky_cube(self): """ - The image cube arranged as it would appear on the sky. + The primary beam mask arranged as it would appear on the sky. Returns: torch.double : 3D image cube of shape ``(nchan, npix, npix)`` """ - return utils.packed_cube_to_sky_cube(self.cube) + return utils.packed_cube_to_sky_cube(self.pb_mask) def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): """ - Export the image cube to a FITS file. + Export the primary beam cube to a FITS file. Args: fname (str): the name of the FITS file to export to. @@ -387,7 +400,7 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): for k, v in header_kwargs.items(): header[k] = v - hdu = fits.PrimaryHDU(self.sky_cube.detach().cpu().numpy(), header=header) + hdu = fits.PrimaryHDU(self.pb_mask.detach().cpu().numpy(), header=header) hdul = fits.HDUList([hdu]) hdul.writeto(fname, overwrite=overwrite) diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index e2ed48cc..a856cea9 100644 --- a/src/mpol/precomposed.py +++ b/src/mpol/precomposed.py @@ -36,7 +36,12 @@ def __init__( coords=None, nchan=None, base_cube=None, + chan_freqs=None, + dish_type=None, + dish_radius=None, + **dish_kwargs, ): + super().__init__() _setup_coords(self, cell_size, npix, coords, nchan) @@ -51,8 +56,13 @@ def __init__( coords=self.coords, nchan=self.nchan, passthrough=True ) - self.pbcube = images.PBCorrectedCube( - coords = self.coords, nchan=self.nchan, telescope="ALMA_12" + self.pbcube = images.PrimaryBeamCube( + coords = self.coords, + nchan=self.nchan, + chan_freqs=chan_freqs, + dish_type=dish_type, + dish_radius=dish_radius, + **dish_kwargs ) self.fcube = images.FourierCube(coords=self.coords) From c40f1b15cdb8a53f488577bb0b9953168a058269 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 02:03:53 -0500 Subject: [PATCH 05/21] First attempt at both versions of primary beam corrections implemented --- src/mpol/images.py | 34 +++++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index 0a922f6a..80e980b2 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -338,20 +338,48 @@ def forward(self, cube): Returns: (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. """ - return torch.mul(self.pb_mask, cube) + return torch.mul(self.pbmask, cube) def uniform_mask(self, chan_freqs, dish_radius): r""" Generates airy disk primary beam correction mask. """ - return self.default_mask + assert dish_radius > 0., "Dish radius must be positive" + ratio = 2. * dish_radius / chan_freqs + ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] + r_cube = np.tile(r_2D,(self.nchan,1,1)) + + r_normed_cube = np.pi * r_cube /R_cube + + mask = np.ones((self.nchan, self.npix, self.npix)) + norm_factor = (2. * j1(1e-5) / 1e-5)**2 + mask[r_normed_cube > 0.] = (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor + return torch.tensor(mask) + def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): r""" Generates airy disk primary beam correction mask. """ assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" - return self.default_mask + assert dish_radius > 0., "Dish radius must be positive" + assert dish_obscured_radius > 0., "Obscured dish radius must be positive" + assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" + + ratio = 2. * dish_radius / chan_freqs + ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] + r_cube = np.tile(r_2D,(self.nchan,1,1)) + + eps = dish_obscured_radius / dish_radius + r_normed_cube = np.pi * r_cube /R_cube + + mask = np.ones((self.nchan, self.npix, self.npix)) + norm_factor = (j1(1e-5) / 1e-5 - eps*j1(eps*1e-5)/1e-5)**2 + mask[r_normed_cube > 0.] = (j1(r_normed_cube) / r_normed_cube + - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor + return torch.tensor(mask) @property def sky_cube(self): From 0ec986240cb33d9458914139ad37cf65b7f85b94 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 16:41:25 -0500 Subject: [PATCH 06/21] Fixed typo in diameter to wavelength ratio calculation --- src/mpol/images.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index 80e980b2..9af1e128 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -345,7 +345,7 @@ def uniform_mask(self, chan_freqs, dish_radius): Generates airy disk primary beam correction mask. """ assert dish_radius > 0., "Dish radius must be positive" - ratio = 2. * dish_radius / chan_freqs + ratio = 2. * dish_radius * chan_freqs / 2.998e8 ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] r_cube = np.tile(r_2D,(self.nchan,1,1)) @@ -367,7 +367,7 @@ def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **ex assert dish_obscured_radius > 0., "Obscured dish radius must be positive" assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" - ratio = 2. * dish_radius / chan_freqs + ratio = 2. * dish_radius * chan_freqs / 2.998e8 ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] r_cube = np.tile(r_2D,(self.nchan,1,1)) From c60d0b7a87fef2a7d725452d429873a829224f15 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 20:28:47 -0500 Subject: [PATCH 07/21] Fixed miscellaneous errors, including missing arcsec to radian conversions --- src/mpol/images.py | 42 ++++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index 9af1e128..41e8a9f8 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -340,21 +340,25 @@ def forward(self, cube): """ return torch.mul(self.pbmask, cube) + def uniform_mask(self, chan_freqs, dish_radius): r""" Generates airy disk primary beam correction mask. """ assert dish_radius > 0., "Dish radius must be positive" - ratio = 2. * dish_radius * chan_freqs / 2.998e8 - ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] - r_cube = np.tile(r_2D,(self.nchan,1,1)) - - r_normed_cube = np.pi * r_cube /R_cube - - mask = np.ones((self.nchan, self.npix, self.npix)) + ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 + + ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec + r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians + r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) + + r_normed_cube = np.pi * r_cube * ratio_cube + norm_factor = (2. * j1(1e-5) / 1e-5)**2 - mask[r_normed_cube > 0.] = (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor + mask = np.where(r_normed_cube > 0., + (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor, + 1.) return torch.tensor(mask) @@ -367,18 +371,20 @@ def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **ex assert dish_obscured_radius > 0., "Obscured dish radius must be positive" assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" - ratio = 2. * dish_radius * chan_freqs / 2.998e8 - ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] - r_cube = np.tile(r_2D,(self.nchan,1,1)) + ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 + ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec + r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians + r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) eps = dish_obscured_radius / dish_radius - r_normed_cube = np.pi * r_cube /R_cube + r_normed_cube = np.pi * r_cube * ratio_cube - mask = np.ones((self.nchan, self.npix, self.npix)) norm_factor = (j1(1e-5) / 1e-5 - eps*j1(eps*1e-5)/1e-5)**2 - mask[r_normed_cube > 0.] = (j1(r_normed_cube) / r_normed_cube - - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor + mask = np.where(r_normed_cube > 0., + (j1(r_normed_cube) / r_normed_cube + - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, + 1.) return torch.tensor(mask) @property @@ -448,7 +454,7 @@ class FourierCube(nn.Module): def __init__(self, cell_size=None, npix=None, coords=None): - super().__init__(cell_size) + super().__init__() # we don't want to bother with the nchan argument here, so # we don't use the convenience method _setup_coords From 3065dc66f549d11c0666a1bac34cb9aa61b12d7b Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Wed, 30 Nov 2022 11:06:43 -0500 Subject: [PATCH 08/21] Normed factor IS 1 for the uniform case, no need to calculate it --- src/mpol/images.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index 41e8a9f8..65f2121f 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -355,9 +355,8 @@ def uniform_mask(self, chan_freqs, dish_radius): r_normed_cube = np.pi * r_cube * ratio_cube - norm_factor = (2. * j1(1e-5) / 1e-5)**2 mask = np.where(r_normed_cube > 0., - (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor, + (2. * j1(r_normed_cube) / r_normed_cube)**2, 1.) return torch.tensor(mask) @@ -380,7 +379,7 @@ def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **ex eps = dish_obscured_radius / dish_radius r_normed_cube = np.pi * r_cube * ratio_cube - norm_factor = (j1(1e-5) / 1e-5 - eps*j1(eps*1e-5)/1e-5)**2 + norm_factor = (1.-eps**2)**2 mask = np.where(r_normed_cube > 0., (j1(r_normed_cube) / r_normed_cube - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, From 3d1de04665d93030a305de630999d534541a90af Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Thu, 27 Oct 2022 15:56:25 -0400 Subject: [PATCH 09/21] Added obscured_airy_disk() function and added PBCorrectedCube to SimpleNet, currently can only handle one channel --- src/mpol/images.py | 117 ++++++++++++++++++++++++++++++++++++++++ src/mpol/precomposed.py | 7 +++ 2 files changed, 124 insertions(+) diff --git a/src/mpol/images.py b/src/mpol/images.py index 0c54013d..a1bc7b5d 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -3,6 +3,7 @@ from __future__ import annotations import numpy as np +from scipy.special import j1 import torch import torch.fft # to avoid conflicts with old torch.fft *function* from torch import nn @@ -303,3 +304,119 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul.writeto(fname, overwrite=overwrite) hdul.close() + + +class PBCorrectedCube(nn.Module): + r""" + A ImageCube that has been corrected for the expected primary beam. Currently only implemented for ALMA. + + Args: + cell_size (float): the width of a pixel [arcseconds] + npix (int): the number of pixels per image side + coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. + nchan (int): the number of channels in the image + """ + def __init__( + self, + cell_size=None, + npix=None, + coords=None, + chan_freqs=None, + telescope="ALMA_12" + ): + nchan = len(np.at_least_1d(chan_freqs)) + super().__init__() + + #_setup_coords(self, cell_size, npix, coords, nchan) TODO: update this + + self.telescope_to_radius = {"ALMA_12": 6.} # in meters + self.telescope_to_blocked_radius = {"ALMA_12": 0.375} + + assert (telescope in self.telescope_to_radius) and (telescope in self.telescope_to_blocked_radius) + + self.pb_mask = self.obscured_airy_disk(chan_freqs, telescope) + + + def forward(self, cube): + r"""Args: + cube (torch.double tensor, of shape ``(nchan, npix, npix)``): a prepacked image cube, for example, from ImageCube.forward() + + Returns: + (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. + """ + return self.pb_mask * cube + + def obscured_airy_disk(self, telescope): + r""" + Generates airy disk primary beam correction mask assuming some is obscured + by secondary reflector. + """ + + R = telescope_to_radius[telescope] + R_obs = telescope_to_blocked_radius[telescope] + + r_coords = np.sqrt(self.packed_x_centers_2D**2 + self.packed_y_centers_2D**2) # units of arcsec + r_coords_rads = r_coords * np.pi/648000. # radians + wl = 2.998e8 / self.chan_freqs[0] # TODO: handle multiple channels + x = np.pi * R * r_coords_rads / wl + + eps = R_obs / R + j1_x = scipy.special.j1(x) + j1_epsx = scipy.special.j1(eps * x) + mask = 4. * (j1_x / x - eps * j1_epsx / x)**2 / (1. - eps**2)**2 + + return mask + + @property + def sky_cube(self): + """ + The image cube arranged as it would appear on the sky. + + Returns: + torch.double : 3D image cube of shape ``(nchan, npix, npix)`` + + """ + return utils.packed_cube_to_sky_cube(self.cube) + + def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): + """ + Export the image cube to a FITS file. + + Args: + fname (str): the name of the FITS file to export to. + overwrite (bool): if the file already exists, overwrite? + header_kwargs (dict): Extra keyword arguments to write to the FITS header. + + Returns: + None + """ + + try: + from astropy import wcs + from astropy.io import fits + except ImportError: + print( + "Please install the astropy package to use FITS export functionality." + ) + + w = wcs.WCS(naxis=2) + + w.wcs.crpix = np.array([1, 1]) + w.wcs.cdelt = ( + np.array([self.coords.cell_size, self.coords.cell_size]) / 3600 + ) # decimal degrees + w.wcs.ctype = ["RA---TAN", "DEC--TAN"] + + header = w.to_header() + + # add in the kwargs to the header + if header_kwargs is not None: + for k, v in header_kwargs.items(): + header[k] = v + + hdu = fits.PrimaryHDU(self.sky_cube.detach().cpu().numpy(), header=header) + + hdul = fits.HDUList([hdu]) + hdul.writeto(fname, overwrite=overwrite) + + hdul.close() diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index 149fae73..253bbf74 100644 --- a/src/mpol/precomposed.py +++ b/src/mpol/precomposed.py @@ -50,12 +50,17 @@ def __init__( self.icube = images.ImageCube( coords=self.coords, nchan=self.nchan, passthrough=True ) + self.pbcube = images.PBCorrectedCube( + coords = self.coords, nchan=self.nchan, telescope="ALMA_12" + ) self.fcube = fourier.FourierCube(coords=self.coords) + @classmethod def from_image_properties(cls, cell_size, npix, nchan, base_cube): coords = GridCoords(cell_size, npix) return cls(coords, nchan, base_cube) + def forward(self): r""" @@ -66,5 +71,7 @@ def forward(self): x = self.bcube() x = self.conv_layer(x) x = self.icube(x) + x = self.pbcube(x) vis = self.fcube(x) + return vis From 5551ff69a01f557ef51077424c2570c908216a16 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sat, 26 Nov 2022 18:40:34 -0500 Subject: [PATCH 10/21] Modified gridding checks to include optional frequency array. --- src/mpol/gridding.py | 45 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 40 insertions(+), 5 deletions(-) diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index ae75fb15..f5891322 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -14,7 +14,7 @@ from .datasets import GriddedDataset -def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=None): +def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=None, freq=None): """ Check that all data inputs are the same shape, the weights are positive, and the data_re and data_im are floats. @@ -34,6 +34,9 @@ def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=N "All dataset inputs must be the same input shape and size." ) + if freq is not None: # TODO:change to wrongdimensionerror + assert len(uu) == len(freq), "uu must have same number of channels as freq array." + if np.any(weight <= 0.0): raise ValueError("Not all thermal weights are positive, check inputs.") @@ -53,7 +56,7 @@ def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=N # check to see that uu, vv and data do not contain Hermitian pairs verify_no_hermitian_pairs(uu, vv, data_re + 1.0j * data_im) - return uu, vv, weight, data_re, data_im + return uu, vv, weight, data_re, data_im, freq def verify_no_hermitian_pairs(uu, vv, data, test_vis=5, test_channel=0): @@ -144,6 +147,33 @@ def verify_no_hermitian_pairs(uu, vv, data, test_vis=5, test_channel=0): return False +def _check_freq_1d(freq=None): + """ + Check that the frequency input array contains only positive floats. + + If the user supplied a float, convert to a 1D array. If no frequency array + was supplied, simply skip. + + """ + if freq is None: + return freq + + assert ( + np.isscalar(freq) or freq.ndim == 1 + ), "Input data vectors should be either None, scalar, or 1D array." + + assert np.all(freq > 0.0), "Not all frequencies are positive, check inputs." + + if np.isscalar(freq): + freq = np.atleast_1d(freq) + + assert (freq.dtype == np.single) or ( + freq.dtype == np.double + ), "freq should be type single or double" + + return freq + + class GridderBase: r""" This class is not designed to be used directly, but rather to be subclassed. @@ -172,13 +202,18 @@ def __init__( weight=None, data_re=None, data_im=None, + freq=None, ): + + # check frequency array is 1d or None, expand if not + freq = _check_freq_1d(freq) + # check everything should be 2d, expand if not # also checks data does not contain Hermitian pairs - uu, vv, weight, data_re, data_im = _check_data_inputs_2d( - uu, vv, weight, data_re, data_im + uu, vv, weight, data_re, data_im, freq = _check_data_inputs_2d( + uu, vv, weight, data_re, data_im, freq ) - + # setup the coordinates object self.coords = coords self.nchan = len(uu) From ad3653a60b84d1e384f5671be681c24001807991 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sat, 26 Nov 2022 18:48:03 -0500 Subject: [PATCH 11/21] Added class attribute chan_freq, changed naming to better differentiate between channel and spatial frequencies --- src/mpol/gridding.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index f5891322..327b5438 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -202,16 +202,16 @@ def __init__( weight=None, data_re=None, data_im=None, - freq=None, + chan_freq=None, ): # check frequency array is 1d or None, expand if not - freq = _check_freq_1d(freq) + chan_freq = _check_freq_1d(chan_freq) # check everything should be 2d, expand if not # also checks data does not contain Hermitian pairs - uu, vv, weight, data_re, data_im, freq = _check_data_inputs_2d( - uu, vv, weight, data_re, data_im, freq + uu, vv, weight, data_re, data_im, chan_freq = _check_data_inputs_2d( + uu, vv, weight, data_re, data_im, chan_freq ) # setup the coordinates object @@ -228,6 +228,7 @@ def __init__( self.weight = weight self.data_re = data_re self.data_im = data_im + self.chan_freq = chan_freq # and register cell indices against data self._create_cell_indices() From d3852b2d0ec15e7d5ef107e5dfec15934c60ac80 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 00:53:17 -0500 Subject: [PATCH 12/21] Set up framework for PrimaryBeamCube and dish types --- src/mpol/images.py | 80 ++++++++++++++++++++++++----------------- src/mpol/precomposed.py | 14 ++++++-- 2 files changed, 60 insertions(+), 34 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index a1bc7b5d..e6a20df2 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -11,6 +11,9 @@ from . import utils from .coordinates import GridCoords +from .gridding import _check_freq_1d + + class BaseCube(nn.Module): r""" @@ -305,36 +308,57 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul.close() - -class PBCorrectedCube(nn.Module): + +class PrimaryBeamCube(nn.Module): r""" - A ImageCube that has been corrected for the expected primary beam. Currently only implemented for ALMA. + A ImageCube representing the primary beam of a described dish type. Currently can correct for a + uniform or center-obscured dish. The forward() method multiplies an image cube by this primary beam mask. Args: cell_size (float): the width of a pixel [arcseconds] npix (int): the number of pixels per image side coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. nchan (int): the number of channels in the image + dish_type (string): the type of dish to correct for. Either 'uniform' or 'obscured'. + dish_radius (float): the radius of the dish (in meters) + dish_kwargs (dict): any additional arguments needed for special dish types. Currently only uses: + dish_obscured_radius (float): the radius of the obscured portion of the dish """ def __init__( self, cell_size=None, npix=None, coords=None, + nchan=None, chan_freqs=None, - telescope="ALMA_12" + dish_type=None, + dish_radius=None, + **dish_kwargs, ): - nchan = len(np.at_least_1d(chan_freqs)) super().__init__() #_setup_coords(self, cell_size, npix, coords, nchan) TODO: update this - self.telescope_to_radius = {"ALMA_12": 6.} # in meters - self.telescope_to_blocked_radius = {"ALMA_12": 0.375} + _check_freq_1d(chan_freqs) + assert (chan_freqs is None) or (len(chan_freqs) == nchan), "Length of chan_freqs must be equal to nchan" - assert (telescope in self.telescope_to_radius) and (telescope in self.telescope_to_blocked_radius) + assert (dish_type is None) or (dish_type in ["uniform", "obscured"]), "Provided dish_type must be 'uniform' or 'obscured'" - self.pb_mask = self.obscured_airy_disk(chan_freqs, telescope) + self.default_mask = nn.Parameter( + torch.full( + (self.nchan, self.coords.npix, self.coords.npix), + fill_value=1.0, + requires_grad=False, + dtype=torch.double, + ) + ) + + if dish_type is None: + self.pb_mask = self.default_mask + elif dish_type == "uniform": + self.pb_mask = self.uniform_mask(chan_freqs, dish_radius) + elif dish_type == "obscured": + self.pb_mask = self.obscured_mask(chan_freqs, dish_radius, **dish_kwargs) def forward(self, cube): @@ -344,43 +368,35 @@ def forward(self, cube): Returns: (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. """ - return self.pb_mask * cube + return torch.mul(self.pb_mask, cube) - def obscured_airy_disk(self, telescope): + def uniform_mask(self, chan_freqs, dish_radius): r""" - Generates airy disk primary beam correction mask assuming some is obscured - by secondary reflector. + Generates airy disk primary beam correction mask. """ - - R = telescope_to_radius[telescope] - R_obs = telescope_to_blocked_radius[telescope] - - r_coords = np.sqrt(self.packed_x_centers_2D**2 + self.packed_y_centers_2D**2) # units of arcsec - r_coords_rads = r_coords * np.pi/648000. # radians - wl = 2.998e8 / self.chan_freqs[0] # TODO: handle multiple channels - x = np.pi * R * r_coords_rads / wl - - eps = R_obs / R - j1_x = scipy.special.j1(x) - j1_epsx = scipy.special.j1(eps * x) - mask = 4. * (j1_x / x - eps * j1_epsx / x)**2 / (1. - eps**2)**2 - - return mask + return self.default_mask + + def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): + r""" + Generates airy disk primary beam correction mask. + """ + assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" + return self.default_mask @property def sky_cube(self): """ - The image cube arranged as it would appear on the sky. + The primary beam mask arranged as it would appear on the sky. Returns: torch.double : 3D image cube of shape ``(nchan, npix, npix)`` """ - return utils.packed_cube_to_sky_cube(self.cube) + return utils.packed_cube_to_sky_cube(self.pb_mask) def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): """ - Export the image cube to a FITS file. + Export the primary beam cube to a FITS file. Args: fname (str): the name of the FITS file to export to. @@ -414,7 +430,7 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): for k, v in header_kwargs.items(): header[k] = v - hdu = fits.PrimaryHDU(self.sky_cube.detach().cpu().numpy(), header=header) + hdu = fits.PrimaryHDU(self.pb_mask.detach().cpu().numpy(), header=header) hdul = fits.HDUList([hdu]) hdul.writeto(fname, overwrite=overwrite) diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index 253bbf74..39a00cdd 100644 --- a/src/mpol/precomposed.py +++ b/src/mpol/precomposed.py @@ -35,7 +35,12 @@ def __init__( coords=None, nchan=1, base_cube=None, + chan_freqs=None, + dish_type=None, + dish_radius=None, + **dish_kwargs, ): + super().__init__() self.coords = coords @@ -50,8 +55,13 @@ def __init__( self.icube = images.ImageCube( coords=self.coords, nchan=self.nchan, passthrough=True ) - self.pbcube = images.PBCorrectedCube( - coords = self.coords, nchan=self.nchan, telescope="ALMA_12" + self.pbcube = images.PrimaryBeamCube( + coords = self.coords, + nchan=self.nchan, + chan_freqs=chan_freqs, + dish_type=dish_type, + dish_radius=dish_radius, + **dish_kwargs ) self.fcube = fourier.FourierCube(coords=self.coords) From 2006878fd1cd3acd0dcb13474fd7bb83dfef6f10 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 02:03:53 -0500 Subject: [PATCH 13/21] First attempt at both versions of primary beam corrections implemented --- src/mpol/images.py | 34 +++++++++++++++++++++++++++++++--- 1 file changed, 31 insertions(+), 3 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index e6a20df2..eae083f4 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -368,20 +368,48 @@ def forward(self, cube): Returns: (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. """ - return torch.mul(self.pb_mask, cube) + return torch.mul(self.pbmask, cube) def uniform_mask(self, chan_freqs, dish_radius): r""" Generates airy disk primary beam correction mask. """ - return self.default_mask + assert dish_radius > 0., "Dish radius must be positive" + ratio = 2. * dish_radius / chan_freqs + ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] + r_cube = np.tile(r_2D,(self.nchan,1,1)) + + r_normed_cube = np.pi * r_cube /R_cube + + mask = np.ones((self.nchan, self.npix, self.npix)) + norm_factor = (2. * j1(1e-5) / 1e-5)**2 + mask[r_normed_cube > 0.] = (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor + return torch.tensor(mask) + def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): r""" Generates airy disk primary beam correction mask. """ assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" - return self.default_mask + assert dish_radius > 0., "Dish radius must be positive" + assert dish_obscured_radius > 0., "Obscured dish radius must be positive" + assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" + + ratio = 2. * dish_radius / chan_freqs + ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] + r_cube = np.tile(r_2D,(self.nchan,1,1)) + + eps = dish_obscured_radius / dish_radius + r_normed_cube = np.pi * r_cube /R_cube + + mask = np.ones((self.nchan, self.npix, self.npix)) + norm_factor = (j1(1e-5) / 1e-5 - eps*j1(eps*1e-5)/1e-5)**2 + mask[r_normed_cube > 0.] = (j1(r_normed_cube) / r_normed_cube + - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor + return torch.tensor(mask) @property def sky_cube(self): From 5eb95eefe9598c8f2250e7275480fbb4cefab8b8 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 16:41:25 -0500 Subject: [PATCH 14/21] Fixed typo in diameter to wavelength ratio calculation --- src/mpol/images.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index eae083f4..d2be5cb2 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -375,7 +375,7 @@ def uniform_mask(self, chan_freqs, dish_radius): Generates airy disk primary beam correction mask. """ assert dish_radius > 0., "Dish radius must be positive" - ratio = 2. * dish_radius / chan_freqs + ratio = 2. * dish_radius * chan_freqs / 2.998e8 ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] r_cube = np.tile(r_2D,(self.nchan,1,1)) @@ -397,7 +397,7 @@ def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **ex assert dish_obscured_radius > 0., "Obscured dish radius must be positive" assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" - ratio = 2. * dish_radius / chan_freqs + ratio = 2. * dish_radius * chan_freqs / 2.998e8 ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] r_cube = np.tile(r_2D,(self.nchan,1,1)) From 02651fba81b15107515570a13c946aa38b4a3117 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Sun, 27 Nov 2022 20:28:47 -0500 Subject: [PATCH 15/21] Fixed miscellaneous errors, including missing arcsec to radian conversions --- src/mpol/images.py | 42 ++++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 18 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index d2be5cb2..49a5dba0 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -370,21 +370,25 @@ def forward(self, cube): """ return torch.mul(self.pbmask, cube) + def uniform_mask(self, chan_freqs, dish_radius): r""" Generates airy disk primary beam correction mask. """ assert dish_radius > 0., "Dish radius must be positive" - ratio = 2. * dish_radius * chan_freqs / 2.998e8 - ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] - r_cube = np.tile(r_2D,(self.nchan,1,1)) - - r_normed_cube = np.pi * r_cube /R_cube - - mask = np.ones((self.nchan, self.npix, self.npix)) + ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 + + ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec + r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians + r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) + + r_normed_cube = np.pi * r_cube * ratio_cube + norm_factor = (2. * j1(1e-5) / 1e-5)**2 - mask[r_normed_cube > 0.] = (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor + mask = np.where(r_normed_cube > 0., + (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor, + 1.) return torch.tensor(mask) @@ -397,18 +401,20 @@ def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **ex assert dish_obscured_radius > 0., "Obscured dish radius must be positive" assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" - ratio = 2. * dish_radius * chan_freqs / 2.998e8 - ratio_cube = np.tile(ratio,(1,self.npix,self.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # [arcsec] - r_cube = np.tile(r_2D,(self.nchan,1,1)) + ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 + ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec + r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians + r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) eps = dish_obscured_radius / dish_radius - r_normed_cube = np.pi * r_cube /R_cube + r_normed_cube = np.pi * r_cube * ratio_cube - mask = np.ones((self.nchan, self.npix, self.npix)) norm_factor = (j1(1e-5) / 1e-5 - eps*j1(eps*1e-5)/1e-5)**2 - mask[r_normed_cube > 0.] = (j1(r_normed_cube) / r_normed_cube - - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor + mask = np.where(r_normed_cube > 0., + (j1(r_normed_cube) / r_normed_cube + - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, + 1.) return torch.tensor(mask) @property @@ -463,4 +469,4 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul = fits.HDUList([hdu]) hdul.writeto(fname, overwrite=overwrite) - hdul.close() + hdul.close() \ No newline at end of file From b7d9a78796a12b95ee69854cfc80343d932693e0 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Wed, 30 Nov 2022 11:06:43 -0500 Subject: [PATCH 16/21] Normed factor IS 1 for the uniform case, no need to calculate it --- src/mpol/images.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index 49a5dba0..f951d81c 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -385,9 +385,8 @@ def uniform_mask(self, chan_freqs, dish_radius): r_normed_cube = np.pi * r_cube * ratio_cube - norm_factor = (2. * j1(1e-5) / 1e-5)**2 mask = np.where(r_normed_cube > 0., - (2. * j1(r_normed_cube) / r_normed_cube)**2 / norm_factor, + (2. * j1(r_normed_cube) / r_normed_cube)**2, 1.) return torch.tensor(mask) @@ -410,7 +409,7 @@ def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **ex eps = dish_obscured_radius / dish_radius r_normed_cube = np.pi * r_cube * ratio_cube - norm_factor = (j1(1e-5) / 1e-5 - eps*j1(eps*1e-5)/1e-5)**2 + norm_factor = (1.-eps**2)**2 mask = np.where(r_normed_cube > 0., (j1(r_normed_cube) / r_normed_cube - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, From 5a7d7f0be240f6a06f25c4946f45417dd49f9a5b Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Thu, 23 Mar 2023 16:44:23 -0400 Subject: [PATCH 17/21] Rebase + removed _setup_coords function for PrimaryBeamCorrected --- src/mpol/images.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index f951d81c..a0f6b209 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -326,10 +326,8 @@ class PrimaryBeamCube(nn.Module): """ def __init__( self, - cell_size=None, - npix=None, - coords=None, - nchan=None, + coords, + nchan=1, chan_freqs=None, dish_type=None, dish_radius=None, @@ -360,7 +358,15 @@ def __init__( elif dish_type == "obscured": self.pb_mask = self.obscured_mask(chan_freqs, dish_radius, **dish_kwargs) - + @classmethod + def from_image_properties( + cls, cell_size, npix, nchan=1, + chan_freqs=None, dish_type=None, + dish_radius=None, **dish_kwargs + ) -> ImageCube: + coords = GridCoords(cell_size, npix) + return cls(coords, nchan, chan_freqs, dish_type, dish_radius, **dish_kwargs) + def forward(self, cube): r"""Args: cube (torch.double tensor, of shape ``(nchan, npix, npix)``): a prepacked image cube, for example, from ImageCube.forward() From 503acd36628e4f4196fea343a53511c37ed8a94a Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Thu, 23 Mar 2023 17:00:15 -0400 Subject: [PATCH 18/21] Removed last merge conflict from gridding --- src/mpol/gridding.py | 30 ------------------------------ 1 file changed, 30 deletions(-) diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index 8ab0964d..9d11be81 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -62,37 +62,7 @@ def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=N return uu, vv, weight, data_re, data_im, freq -def _check_freq_1d(freq=None): - """ - Check that the frequency input array contains only positive floats. - - If the user supplied a float, convert to a 1D array. If no frequency array - was supplied, simply skip. - -<<<<<<< HEAD def verify_no_hermitian_pairs(uu, vv, data, test_vis=5, test_channel=0): -======= - """ - if freq is None: - return freq - - assert ( - np.isscalar(freq) or freq.ndim == 1 - ), "Input data vectors should be either None, scalar, or 1D array." - - assert np.all(freq > 0.0), "Not all frequencies are positive, check inputs." - - if np.isscalar(freq): - freq = np.atleast_1d(freq) - - assert (freq.dtype == np.single) or ( - freq.dtype == np.double - ), "freq should be type single or double" - - return freq - -class Gridder: ->>>>>>> 3065dc66f549d11c0666a1bac34cb9aa61b12d7b r""" Check that the dataset does not contain Hermitian pairs. Because the sky brightness :math:`I_\nu` is real, the visibility function :math:`\mathcal{V}` is Hermitian, meaning that From cf042b8f2041f24241766fbf9a3ce1354818cd8a Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Thu, 23 Mar 2023 17:06:39 -0400 Subject: [PATCH 19/21] Removed duplicated PrimaryBeamCube class --- src/mpol/images.py | 160 --------------------------------------------- 1 file changed, 160 deletions(-) diff --git a/src/mpol/images.py b/src/mpol/images.py index a381df3f..ad564528 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -308,167 +308,7 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul.close() -class PrimaryBeamCube(nn.Module): - r""" - A ImageCube representing the primary beam of a described dish type. Currently can correct for a - uniform or center-obscured dish. The forward() method multiplies an image cube by this primary beam mask. - - Args: - cell_size (float): the width of a pixel [arcseconds] - npix (int): the number of pixels per image side - coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. - nchan (int): the number of channels in the image - dish_type (string): the type of dish to correct for. Either 'uniform' or 'obscured'. - dish_radius (float): the radius of the dish (in meters) - dish_kwargs (dict): any additional arguments needed for special dish types. Currently only uses: - dish_obscured_radius (float): the radius of the obscured portion of the dish - """ - def __init__( - self, - cell_size=None, - npix=None, - coords=None, - nchan=None, - chan_freqs=None, - dish_type=None, - dish_radius=None, - **dish_kwargs, - ): - super().__init__() - _setup_coords(self, cell_size, npix, coords, nchan) - - _check_freq_1d(chan_freqs) - assert (chan_freqs is None) or (len(chan_freqs) == nchan), "Length of chan_freqs must be equal to nchan" - - assert (dish_type is None) or (dish_type in ["uniform", "obscured"]), "Provided dish_type must be 'uniform' or 'obscured'" - - self.default_mask = nn.Parameter( - torch.full( - (self.nchan, self.coords.npix, self.coords.npix), - fill_value=1.0, - requires_grad=False, - dtype=torch.double, - ) - ) - - if dish_type is None: - self.pb_mask = self.default_mask - elif dish_type == "uniform": - self.pb_mask = self.uniform_mask(chan_freqs, dish_radius) - elif dish_type == "obscured": - self.pb_mask = self.obscured_mask(chan_freqs, dish_radius, **dish_kwargs) - - - def forward(self, cube): - r"""Args: - cube (torch.double tensor, of shape ``(nchan, npix, npix)``): a prepacked image cube, for example, from ImageCube.forward() - - Returns: - (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. - """ - return torch.mul(self.pbmask, cube) - - - def uniform_mask(self, chan_freqs, dish_radius): - r""" - Generates airy disk primary beam correction mask. - """ - assert dish_radius > 0., "Dish radius must be positive" - ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 - - ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec - r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians - r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) - - r_normed_cube = np.pi * r_cube * ratio_cube - - mask = np.where(r_normed_cube > 0., - (2. * j1(r_normed_cube) / r_normed_cube)**2, - 1.) - return torch.tensor(mask) - - - def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): - r""" - Generates airy disk primary beam correction mask. - """ - assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" - assert dish_radius > 0., "Dish radius must be positive" - assert dish_obscured_radius > 0., "Obscured dish radius must be positive" - assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" - - ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 - ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec - r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians - r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) - - eps = dish_obscured_radius / dish_radius - r_normed_cube = np.pi * r_cube * ratio_cube - norm_factor = (1.-eps**2)**2 - mask = np.where(r_normed_cube > 0., - (j1(r_normed_cube) / r_normed_cube - - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, - 1.) - return torch.tensor(mask) - - @property - def sky_cube(self): - """ - The primary beam mask arranged as it would appear on the sky. - - Returns: - torch.double : 3D image cube of shape ``(nchan, npix, npix)`` - - """ - return utils.packed_cube_to_sky_cube(self.pb_mask) - - def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): - """ - Export the primary beam cube to a FITS file. - - Args: - fname (str): the name of the FITS file to export to. - overwrite (bool): if the file already exists, overwrite? - header_kwargs (dict): Extra keyword arguments to write to the FITS header. - - Returns: - None - """ - - try: - from astropy import wcs - from astropy.io import fits - except ImportError: - print( - "Please install the astropy package to use FITS export functionality." - ) - - w = wcs.WCS(naxis=2) - - w.wcs.crpix = np.array([1, 1]) - w.wcs.cdelt = ( - np.array([self.coords.cell_size, self.coords.cell_size]) / 3600 - ) # decimal degrees - w.wcs.ctype = ["RA---TAN", "DEC--TAN"] - - header = w.to_header() - - # add in the kwargs to the header - if header_kwargs is not None: - for k, v in header_kwargs.items(): - header[k] = v - - hdu = fits.PrimaryHDU(self.pb_mask.detach().cpu().numpy(), header=header) - - hdul = fits.HDUList([hdu]) - hdul.writeto(fname, overwrite=overwrite) - - hdul.close() - - class PrimaryBeamCube(nn.Module): r""" A ImageCube representing the primary beam of a described dish type. Currently can correct for a From 45639dd783d897a12e67672463f151cecb2cf91e Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Thu, 23 Mar 2023 17:13:18 -0400 Subject: [PATCH 20/21] Small typo in gridding.py --- src/mpol/gridding.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index 9d11be81..ed28120a 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -212,7 +212,8 @@ def __init__( chan_freq = _check_freq_1d(chan_freq) # check everything should be 2d, expand if not - # also checks data does not contain Hermitian pairsuu, vv, weight, data_re, data_im, chan_freq = _check_data_inputs_2d( + # also checks data does not contain Hermitian pairs + uu, vv, weight, data_re, data_im, chan_freq = _check_data_inputs_2d( uu, vv, weight, data_re, data_im, chan_freq ) From e276817df872f176c1bb8bf9ca51ae95b9fb4620 Mon Sep 17 00:00:00 2001 From: Kaylee de Soto Date: Mon, 10 Apr 2023 15:21:33 -0400 Subject: [PATCH 21/21] Corrected bugs in primary_beam.py, moved to separate file, added default test case --- src/mpol/gridding.py | 6 +- src/mpol/images.py | 170 ----------------------------------- src/mpol/precomposed.py | 4 +- src/mpol/primary_beam.py | 184 ++++++++++++++++++++++++++++++++++++++ test/conftest.py | 9 ++ test/primary_beam_test.py | 16 ++++ 6 files changed, 212 insertions(+), 177 deletions(-) create mode 100644 src/mpol/primary_beam.py create mode 100644 test/primary_beam_test.py diff --git a/src/mpol/gridding.py b/src/mpol/gridding.py index ed28120a..facb234b 100644 --- a/src/mpol/gridding.py +++ b/src/mpol/gridding.py @@ -29,10 +29,6 @@ def _check_data_inputs_2d(uu=None, vv=None, weight=None, data_re=None, data_im=N "Input data vectors should be either 1D or 2D numpy arrays." ) - if freq is not None: # TODO: remove assertion - assert len(uu) == len(freq), "uu must have same number of channels as freq array." - - if not all(array.shape == uu.shape for array in [vv, weight, data_re, data_im]): raise WrongDimensionError( "All dataset inputs must be the same input shape and size." @@ -626,7 +622,7 @@ def __init__( ): # check everything should be 2d, expand if not # also checks data does not contain Hermitian pairs - uu, vv, weight, data_re, data_im = _check_data_inputs_2d( + uu, vv, weight, data_re, data_im, freq = _check_data_inputs_2d( uu, vv, weight, data_re, data_im ) diff --git a/src/mpol/images.py b/src/mpol/images.py index ad564528..855b15e8 100644 --- a/src/mpol/images.py +++ b/src/mpol/images.py @@ -11,8 +11,6 @@ from . import utils from .coordinates import GridCoords -from .gridding import _check_freq_1d - class BaseCube(nn.Module): @@ -307,171 +305,3 @@ def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): hdul.writeto(fname, overwrite=overwrite) hdul.close() - - -class PrimaryBeamCube(nn.Module): - r""" - A ImageCube representing the primary beam of a described dish type. Currently can correct for a - uniform or center-obscured dish. The forward() method multiplies an image cube by this primary beam mask. - - Args: - cell_size (float): the width of a pixel [arcseconds] - npix (int): the number of pixels per image side - coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. - nchan (int): the number of channels in the image - dish_type (string): the type of dish to correct for. Either 'uniform' or 'obscured'. - dish_radius (float): the radius of the dish (in meters) - dish_kwargs (dict): any additional arguments needed for special dish types. Currently only uses: - dish_obscured_radius (float): the radius of the obscured portion of the dish - """ - def __init__( - self, - coords, - nchan=1, - chan_freqs=None, - dish_type=None, - dish_radius=None, - **dish_kwargs, - ): - super().__init__() - - #_setup_coords(self, cell_size, npix, coords, nchan) TODO: update this - - _check_freq_1d(chan_freqs) - assert (chan_freqs is None) or (len(chan_freqs) == nchan), "Length of chan_freqs must be equal to nchan" - - assert (dish_type is None) or (dish_type in ["uniform", "obscured"]), "Provided dish_type must be 'uniform' or 'obscured'" - - self.default_mask = nn.Parameter( - torch.full( - (self.nchan, self.coords.npix, self.coords.npix), - fill_value=1.0, - requires_grad=False, - dtype=torch.double, - ) - ) - - if dish_type is None: - self.pb_mask = self.default_mask - elif dish_type == "uniform": - self.pb_mask = self.uniform_mask(chan_freqs, dish_radius) - elif dish_type == "obscured": - self.pb_mask = self.obscured_mask(chan_freqs, dish_radius, **dish_kwargs) - - @classmethod - def from_image_properties( - cls, cell_size, npix, nchan=1, - chan_freqs=None, dish_type=None, - dish_radius=None, **dish_kwargs - ) -> ImageCube: - coords = GridCoords(cell_size, npix) - return cls(coords, nchan, chan_freqs, dish_type, dish_radius, **dish_kwargs) - - def forward(self, cube): - r"""Args: - cube (torch.double tensor, of shape ``(nchan, npix, npix)``): a prepacked image cube, for example, from ImageCube.forward() - - Returns: - (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. - """ - return torch.mul(self.pbmask, cube) - - - def uniform_mask(self, chan_freqs, dish_radius): - r""" - Generates airy disk primary beam correction mask. - """ - assert dish_radius > 0., "Dish radius must be positive" - ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 - - ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec - r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians - r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) - - r_normed_cube = np.pi * r_cube * ratio_cube - - mask = np.where(r_normed_cube > 0., - (2. * j1(r_normed_cube) / r_normed_cube)**2, - 1.) - return torch.tensor(mask) - - - def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): - r""" - Generates airy disk primary beam correction mask. - """ - assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" - assert dish_radius > 0., "Dish radius must be positive" - assert dish_obscured_radius > 0., "Obscured dish radius must be positive" - assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" - - ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 - ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) - r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec - r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians - r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) - - eps = dish_obscured_radius / dish_radius - r_normed_cube = np.pi * r_cube * ratio_cube - - norm_factor = (1.-eps**2)**2 - mask = np.where(r_normed_cube > 0., - (j1(r_normed_cube) / r_normed_cube - - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, - 1.) - return torch.tensor(mask) - - @property - def sky_cube(self): - """ - The primary beam mask arranged as it would appear on the sky. - - Returns: - torch.double : 3D image cube of shape ``(nchan, npix, npix)`` - - """ - return utils.packed_cube_to_sky_cube(self.pb_mask) - - def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): - """ - Export the primary beam cube to a FITS file. - - Args: - fname (str): the name of the FITS file to export to. - overwrite (bool): if the file already exists, overwrite? - header_kwargs (dict): Extra keyword arguments to write to the FITS header. - - Returns: - None - """ - - try: - from astropy import wcs - from astropy.io import fits - except ImportError: - print( - "Please install the astropy package to use FITS export functionality." - ) - - w = wcs.WCS(naxis=2) - - w.wcs.crpix = np.array([1, 1]) - w.wcs.cdelt = ( - np.array([self.coords.cell_size, self.coords.cell_size]) / 3600 - ) # decimal degrees - w.wcs.ctype = ["RA---TAN", "DEC--TAN"] - - header = w.to_header() - - # add in the kwargs to the header - if header_kwargs is not None: - for k, v in header_kwargs.items(): - header[k] = v - - hdu = fits.PrimaryHDU(self.pb_mask.detach().cpu().numpy(), header=header) - - hdul = fits.HDUList([hdu]) - hdul.writeto(fname, overwrite=overwrite) - - hdul.close() diff --git a/src/mpol/precomposed.py b/src/mpol/precomposed.py index ce27d8a7..f07af11e 100644 --- a/src/mpol/precomposed.py +++ b/src/mpol/precomposed.py @@ -2,7 +2,7 @@ from mpol.coordinates import GridCoords -from . import fourier, images +from . import fourier, images, primary_beam class SimpleNet(torch.nn.Module): @@ -56,7 +56,7 @@ def __init__( coords=self.coords, nchan=self.nchan, passthrough=True ) - self.pbcube = images.PrimaryBeamCube( + self.pbcube = primary_beam.PrimaryBeamCube( coords = self.coords, nchan=self.nchan, chan_freqs=chan_freqs, diff --git a/src/mpol/primary_beam.py b/src/mpol/primary_beam.py new file mode 100644 index 00000000..255c2916 --- /dev/null +++ b/src/mpol/primary_beam.py @@ -0,0 +1,184 @@ +r"""The ``primary_beam`` module provides the core functionality of MPoL via :class:`mpol.fourier.PrimaryBeamCube`.""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.fft # to avoid conflicts with old torch.fft *function* +import torchkbnufft +from torch import nn + +from . import utils +from .coordinates import GridCoords + +from .gridding import _check_freq_1d + +class PrimaryBeamCube(nn.Module): + r""" + A ImageCube representing the primary beam of a described dish type. Currently can correct for a + uniform or center-obscured dish. The forward() method multiplies an image cube by this primary beam mask. + + Args: + cell_size (float): the width of a pixel [arcseconds] + npix (int): the number of pixels per image side + coords (GridCoords): an object already instantiated from the GridCoords class. If providing this, cannot provide ``cell_size`` or ``npix``. + nchan (int): the number of channels in the image + dish_type (string): the type of dish to correct for. Either 'uniform' or 'obscured'. + dish_radius (float): the radius of the dish (in meters) + dish_kwargs (dict): any additional arguments needed for special dish types. Currently only uses: + dish_obscured_radius (float): the radius of the obscured portion of the dish + """ + def __init__( + self, + coords, + nchan=1, + chan_freqs=None, + dish_type=None, + dish_radius=None, + **dish_kwargs, + ): + super().__init__() + + #_setup_coords(self, cell_size, npix, coords, nchan) TODO: update this + + _check_freq_1d(chan_freqs) + assert (chan_freqs is None) or (len(chan_freqs) == nchan), "Length of chan_freqs must be equal to nchan" + + assert (dish_type is None) or (dish_type in ["uniform", "obscured"]), "Provided dish_type must be 'uniform' or 'obscured'" + + self.coords = coords + self.nchan = nchan + + self.default_mask = nn.Parameter( + torch.full( + (self.nchan, self.coords.npix, self.coords.npix), + fill_value=1.0, + requires_grad=False, + dtype=torch.double, + ) + ) + + if dish_type is None: + self.pb_mask = self.default_mask + elif dish_type == "uniform": + self.pb_mask = self.uniform_mask(chan_freqs, dish_radius) + elif dish_type == "obscured": + self.pb_mask = self.obscured_mask(chan_freqs, dish_radius, **dish_kwargs) + + @classmethod + def from_image_properties( + cls, cell_size, npix, nchan=1, + chan_freqs=None, dish_type=None, + dish_radius=None, **dish_kwargs + ) -> ImageCube: + coords = GridCoords(cell_size, npix) + return cls(coords, nchan, chan_freqs, dish_type, dish_radius, **dish_kwargs) + + def forward(self, cube): + r"""Args: + cube (torch.double tensor, of shape ``(nchan, npix, npix)``): a prepacked image cube, for example, from ImageCube.forward() + + Returns: + (torch.complex tensor, of shape ``(nchan, npix, npix)``): the FFT of the image cube, in packed format. + """ + return torch.mul(self.pb_mask, cube) + + + def uniform_mask(self, chan_freqs, dish_radius): + r""" + Generates airy disk primary beam correction mask. + """ + assert dish_radius > 0., "Dish radius must be positive" + ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 + + ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec + r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians + r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) + + r_normed_cube = np.pi * r_cube * ratio_cube + + mask = np.where(r_normed_cube > 0., + (2. * j1(r_normed_cube) / r_normed_cube)**2, + 1.) + return torch.tensor(mask) + + + def obscured_mask(self, chan_freqs, dish_radius, dish_obscured_radius=None, **extra_kwargs): + r""" + Generates airy disk primary beam correction mask. + """ + assert dish_obscured_radius is not None, "Obscured dish requires kwarg 'dish_obscured_radius'" + assert dish_radius > 0., "Dish radius must be positive" + assert dish_obscured_radius > 0., "Obscured dish radius must be positive" + assert dish_radius > dish_obscured_radius, "Primary dish radius must be greater than obscured radius" + + ratio = 2. * dish_radius * np.array([[chan_freqs]]).T / 2.998e8 + ratio_cube = np.tile(ratio,(1,self.coords.npix,self.coords.npix)) + r_2D = np.sqrt(self.coords.packed_x_centers_2D**2 + self.coords.packed_y_centers_2D**2) # arcsec + r_2D_rads = r_2D * np.pi / 180. / 60. / 60. # radians + r_cube = np.tile(r_2D_rads,(self.nchan,1,1)) + + eps = dish_obscured_radius / dish_radius + r_normed_cube = np.pi * r_cube * ratio_cube + + norm_factor = (1.-eps**2)**2 + mask = np.where(r_normed_cube > 0., + (j1(r_normed_cube) / r_normed_cube + - eps*j1(eps*r_normed_cube) / r_normed_cube)**2 / norm_factor, + 1.) + return torch.tensor(mask) + + @property + def sky_cube(self): + """ + The primary beam mask arranged as it would appear on the sky. + + Returns: + torch.double : 3D image cube of shape ``(nchan, npix, npix)`` + + """ + return utils.packed_cube_to_sky_cube(self.pb_mask) + + def to_FITS(self, fname="cube.fits", overwrite=False, header_kwargs=None): + """ + Export the primary beam cube to a FITS file. + + Args: + fname (str): the name of the FITS file to export to. + overwrite (bool): if the file already exists, overwrite? + header_kwargs (dict): Extra keyword arguments to write to the FITS header. + + Returns: + None + """ + + try: + from astropy import wcs + from astropy.io import fits + except ImportError: + print( + "Please install the astropy package to use FITS export functionality." + ) + + w = wcs.WCS(naxis=2) + + w.wcs.crpix = np.array([1, 1]) + w.wcs.cdelt = ( + np.array([self.coords.cell_size, self.coords.cell_size]) / 3600 + ) # decimal degrees + w.wcs.ctype = ["RA---TAN", "DEC--TAN"] + + header = w.to_header() + + # add in the kwargs to the header + if header_kwargs is not None: + for k, v in header_kwargs.items(): + header[k] = v + + hdu = fits.PrimaryHDU(self.pb_mask.detach().cpu().numpy(), header=header) + + hdul = fits.HDUList([hdu]) + hdul.writeto(fname, overwrite=overwrite) + + hdul.close() \ No newline at end of file diff --git a/test/conftest.py b/test/conftest.py index 2686b89c..51823962 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -1,6 +1,7 @@ import numpy as np import pytest from astropy.utils.data import download_file +import torch from mpol import coordinates, gridding @@ -53,6 +54,14 @@ def coords(): return coordinates.GridCoords(cell_size=0.005, npix=800) +@pytest.fixture +def unit_cube(coords): + nchan = 4 + input_cube = torch.full( + (nchan, coords.npix, coords.npix), fill_value=1.0, dtype=torch.double + ) + return input_cube + @pytest.fixture def averager(mock_visibility_data, coords): uu, vv, weight, data_re, data_im = mock_visibility_data diff --git a/test/primary_beam_test.py b/test/primary_beam_test.py new file mode 100644 index 00000000..dcea5a22 --- /dev/null +++ b/test/primary_beam_test.py @@ -0,0 +1,16 @@ +import matplotlib.pyplot as plt +import torch +from pytest import approx + +from mpol import primary_beam, images, utils +from mpol.constants import * + +def test_no_dish_correction(coords, unit_cube): + # Tests layer when no PB correction is applied (passthrough layer) + pblayer = primary_beam.PrimaryBeamCube(coords=coords) + out_cube = pblayer(unit_cube) + + assert torch.equal(unit_cube, out_cube) + + + \ No newline at end of file