"""Data, image, weight, and noise cubes used throughout :mod:`ps_eor`.
The main analysis container is :class:`CartDataCube`: a complex visibility
array shaped ``(frequency, spatial mode)``, together with its frequency and UV
coordinates, image metadata, and optional :class:`CartWeightCube`. Keeping
those objects together lets filtering, foreground fitting, and power-spectrum
estimation preserve the geometry and statistical weights of the data.
For calibrated, frequency-dependent FITS images, the usual entry point is
:meth:`CartDataCube.load_from_fits_image_and_psf`. It Fourier-transforms the
images, converts Jy/PSF to Kelvin when requested, uses the matching PSF FITS
files for the normalization and visibility weights, and restricts the result
to the requested baseline range::
import numpy as np
from ps_eor import datacube, psutil
image_files = psutil.sort_by_fits_key(image_files, 'CRVAL3')
psf_files = psutil.sort_by_fits_key(psf_files, 'CRVAL3')
cube = datacube.CartDataCube.load_from_fits_image_and_psf(
image_files,
psf_files,
umin=50,
umax=250,
theta_fov=np.deg2rad(4),
int_time=10,
total_time=10 * 3600,
)
cube.save('visibilities.h5')
Here ``umin`` and ``umax`` are in wavelengths, ``theta_fov`` is in radians,
and observing times are in seconds. When no PSF image is available,
:meth:`CartDataCube.load_from_fits_image` provides an unweighted alternative.
Saved cubes retain their weights and metadata. Frequency slicing returns a new
cube, while methods named ``filter_*`` generally modify the cube in place::
cube = datacube.CartDataCube.load('visibilities.h5')
subband = cube.get_slice(120e6, 130e6)
subband.filter_uvrange(60, 200)
Use :meth:`DataCube.copy` before an in-place operation when the original cube
must be retained.
"""
import re
import warnings
from typing import ClassVar
import astropy.constants as const
import astropy.io.fits as pf
import astropy.wcs as pywcs
import matplotlib.pyplot as plt
import numpy as np
import scipy.interpolate
import tables
from scipy.signal import get_window # noqa: F401 - re-exported as datacube.get_window
from scipy.stats import binned_statistic
from . import __version__, psutil
def _fmhz(freqs, precision=3):
return np.round(freqs * 1e-6, precision)
[docs]
def get_common_idx(cube1, cube2):
"""Return frequency and spatial-mode masks shared by two cubes."""
freqs = np.intersect1d(_fmhz(cube1.freqs), _fmhz(cube2.freqs))
idx1 = np.isin(_fmhz(cube1.freqs), freqs)
idx2 = np.isin(_fmhz(cube2.freqs), freqs)
a = cube1.get_unique_xy()
b = cube2.get_unique_xy()
z = np.intersect1d(a, b)
idx1_uv = np.isin(a, z)
idx2_uv = np.isin(b, z)
return idx1, idx1_uv, idx2, idx2_uv
[docs]
def get_common_cube(cube1, cube2, only_frequency=False):
"""Return two cubes restricted to their common coordinates."""
idx1, idx1_uv, idx2, idx2_uv = get_common_idx(cube1, cube2)
c1 = cube1.get_slice_from_idx(idx1)
c2 = cube2.get_slice_from_idx(idx2)
if not only_frequency:
c1.filter_uv_from_index(idx1_uv)
c2.filter_uv_from_index(idx2_uv)
return c1, c2
[docs]
def concatenate_datacubes(cubes):
"""Concatenate Cartesian cubes with a shared UV grid along frequency."""
assert len(cubes) > 1
assert np.all([isinstance(cube, DataCube) for cube in cubes])
assert np.all([np.allclose(cubes[0].uu, cube.uu) for cube in cubes[1:]])
assert np.all([np.allclose(cubes[0].vv, cube.vv) for cube in cubes[1:]])
data = np.concatenate([c.data for c in cubes])
weights_data = np.concatenate([c.weights.data for c in cubes])
freqs = np.concatenate([c.freqs for c in cubes])
weights = CartWeightCube(weights_data, cubes[0].uu, cubes[0].vv, freqs, cubes[0].meta)
return CartDataCube(data, cubes[0].uu, cubes[0].vv, freqs, cubes[0].meta, weights=weights)
[docs]
class Mask:
"""A sum of image-domain masks evaluated from cube metadata."""
def __init__(self, masks=None):
if masks is None:
masks = []
self.masks = masks
def __mul__(self, other):
return MaskProd(self, other)
def __add__(self, other):
return Mask(self.masks + other.masks)
[docs]
def generate(self, meta_data):
"""Evaluate the combined mask on the grid described by ``meta_data``.
Returns:
ndarray: the mask image, shape ``meta_data.shape``.
"""
return np.sum([m.generate(meta_data) for m in self.masks], axis=0)
[docs]
def get_power(self, meta_data):
"""Return the mean squared mask response."""
return (self.generate(meta_data) ** 2).mean()
[docs]
def get_area(self, meta_data, normalize=False):
"""Effective area of the mask: its mean response.
Args:
meta_data: the image geometry to evaluate on.
normalize (bool): divide by the peak response.
Returns:
float: the (optionally peak-normalized) mean response.
"""
mask = self.generate(meta_data)
area = mask.mean()
if normalize:
area = area / float(mask.max())
return area
[docs]
class MaskProd(Mask):
"""Pairwise product of two mask collections."""
def __init__(self, m1, m2):
self.m1s = m1.masks
self.m2s = m2.masks
if len(self.m1s) == 1 and len(self.m2s):
self.m1s = [self.m1s[0]] * len(self.m2s)
elif len(self.m2s) == 1 and len(self.m1s):
self.m2s = [self.m2s[0]] * len(self.m1s)
def __add__(self, other):
raise NotImplementedError()
def __mul__(self, other):
raise NotImplementedError()
[docs]
def generate(self, meta_data):
"""Evaluate and sum the paired mask products."""
if len(self.m1s) == len(self.m2s):
return np.sum([m1.generate(meta_data) * m2.generate(meta_data)
for m1, m2 in zip(self.m1s, self.m2s, strict=False)], axis=0)
elif len(self.m1s) == 1 and len(self.m2s) > 1:
m1_map = self.m1s[0].generate(meta_data)
return np.sum([m1_map * m2.generate(meta_data) for m2 in self.m2s], axis=0)
elif len(self.m2s) == 1 and len(self.m1s) > 1:
m2_map = self.m2s[0].generate(meta_data)
return np.sum([m2_map * m1.generate(meta_data) for m1 in self.m1s], axis=0)
else:
raise NotImplementedError()
[docs]
class WindowFunction(Mask):
"""Two-dimensional spectral window applied in the image plane."""
def __init__(self, name, circular=True):
self.name = name
self.circular = circular
Mask.__init__(self, [self])
def __str__(self):
return f'WindowFunction({self.name}, circular={self.circular})'
[docs]
@staticmethod
def parse_winfct_str(s):
"""Parse a SciPy window name and optional numeric argument."""
s = s.strip()
if ',' in s:
s, o = s.strip('() ').split(',')
return (re.sub(r'\W+', '', s), float(o))
else:
return s
[docs]
def generate_window(self, nx):
"""Generate an ``nx × nx`` window.""" # noqa: RUF002
return psutil.generate_2d_window(self.name, nx, circular=self.circular)
[docs]
def generate(self, meta_data):
"""Generate the window at the image size given by ``meta_data.shape``."""
return self.generate_window(meta_data.shape[0])
[docs]
class BasePrimaryBeam(Mask):
"""Base class for frequency-dependent primary-beam masks."""
name = 'pb'
def __init__(self, masks, freq=None):
self.freq = freq
Mask.__init__(self, masks)
[docs]
def set_freq(self, freq):
"""Set the default evaluation frequency in Hz."""
self.freq = freq
[docs]
def get_freq(self, freq=None):
"""Return an explicit frequency or the configured default."""
if freq is not None:
return freq
return self.freq
[docs]
class PrimaryBeam(BasePrimaryBeam):
"""Analytic primary beam defined by an effective aperture diameter."""
def __init__(self, antenna_diameter, alpha_tapering, beam_type, freq=None):
self.antenna_diameter = antenna_diameter
self.alpha_tapering = alpha_tapering
self.beam_type = beam_type
BasePrimaryBeam.__init__(self, [self], freq=freq)
def __str__(self):
return (f'PrimaryBeam({self.antenna_diameter}m, alpha={self.alpha_tapering}, '
f'{self.beam_type})')
[docs]
@staticmethod
def from_name(name):
"""Construct a primary beam by name.
``name`` is either a registered telescope (e.g. ``'lofar_hba'``,
``'nenufar'``) or a custom ``ant_<diameter>_<alpha>_<beam_type>`` string.
Returns:
PrimaryBeam: the matching beam.
"""
if name.startswith('ant_'):
_, diameter, alpha, beam_type = name.split('_')
return PrimaryBeam(float(diameter), float(alpha), beam_type)
klasses = BasePrimaryBeam.__subclasses__()
[klasses.extend(k.__subclasses__()) for k in klasses[:]]
for klass in klasses:
if hasattr(klass, 'name') and klass.name == name:
return klass()
raise ValueError(f'No primary beam with name: {name}')
[docs]
def get_fwhm(self, freq=None):
"""Return the beam FWHM in radians."""
freq = self.get_freq(freq)
assert freq is not None
lamb = const.c.value / freq
return self.alpha_tapering * lamb / self.antenna_diameter
[docs]
def generate_beam(self, fwhm, res, shape):
"""Render the beam on a Cartesian grid.
Args:
fwhm: beam full-width at half-maximum, in radians.
res: pixel size, in radians.
shape: image shape ``(nx, ny)``.
Returns:
ndarray: the beam image.
"""
return psutil.get_beam_cart(res, tuple(shape), self.beam_type, fwhm, n_sidelobe=None)
[docs]
def generate(self, meta_data, freq=None):
"""Render the beam on the grid given by ``meta_data`` at ``freq`` (Hz;
default the configured frequency).
Returns:
ndarray: the beam image, shape ``meta_data.shape``.
"""
fwhm = self.get_fwhm(freq=freq)
return self.generate_beam(fwhm, meta_data.res, meta_data.shape)
[docs]
class LofarHBAPrimaryBeam(PrimaryBeam):
"""Analytic beam approximation for a LOFAR HBA station."""
name = 'lofar_hba'
def __init__(self):
PrimaryBeam.__init__(self, 30.75, 1.25, 'gaussian')
[docs]
class AartfaacA12HBAPrimaryBeam(PrimaryBeam):
"""Analytic beam approximation for an AARTFAAC-12 HBA tile."""
name = 'a12_hba'
def __init__(self):
PrimaryBeam.__init__(self, 5, 1.02, 'gaussian')
[docs]
class LofarLBAInnerPrimaryBeam(PrimaryBeam):
"""Analytic beam approximation for LOFAR LBA inner mode."""
name = 'lofar_lba_inner'
def __init__(self):
PrimaryBeam.__init__(self, 32.25, 1.1, 'gaussian')
[docs]
class LofarLBAOuterPrimaryBeam(PrimaryBeam):
"""Analytic beam approximation for LOFAR LBA outer mode."""
name = 'lofar_lba_outer'
def __init__(self):
PrimaryBeam.__init__(self, 81.34, 1.1, 'gaussian')
[docs]
class SkaLowPrimaryBeam(PrimaryBeam):
"""Analytic beam approximation for an SKA-Low station."""
name = 'ska_low'
def __init__(self):
PrimaryBeam.__init__(self, 38, 1.0335, 'bessel')
[docs]
class NenuFARPrimaryBeam(PrimaryBeam):
"""Analytic beam approximation for a NenuFAR mini-array."""
name = 'nenufar'
def __init__(self):
PrimaryBeam.__init__(self, 25, 1.02, 'gaussian')
[docs]
class NoPrimaryBeam(BasePrimaryBeam):
"""Unit primary beam."""
name = 'no_pb'
def __init__(self):
BasePrimaryBeam.__init__(self, [self])
[docs]
def generate(self, meta_data, freq=None):
"""Return a unit-response image."""
return np.ones(meta_data.shape)
[docs]
class ImageCube:
"""Base frequency-dependent image cube."""
def __init__(self, image_cube, freqs, meta):
self.data = image_cube.real
self.freqs = freqs
self.meta = meta
[docs]
def get_slice(self, freq_start, freq_end):
"""The channels within ``[freq_start, freq_end]`` (Hz).
Returns:
CartImageCube: the frequency-sliced cube.
"""
freq_slice = psutil.get_freq_slice(self.freqs, freq_start, freq_end)
return CartImageCube(self.data[freq_slice], self.freqs[freq_slice], self.meta.copy())
[docs]
def get_freq(self, freq):
"""The single channel at or just above ``freq`` (Hz).
Returns:
CartImageCube: a one-channel cube.
"""
i = np.nonzero(self.freqs >= freq)[0][0]
return CartImageCube(self.data[i:i + 1], self.freqs[i:i + 1], self.meta.copy())
[docs]
class CartImageCube(ImageCube):
"""Frequency cube of Cartesian sky images."""
def __init__(self, image_cube, freqs, meta):
"""Create a Cartesian image cube.
Args:
image_cube: Array with shape ``(n_freqs, nx, ny)``.
freqs: Frequencies in Hz.
meta: Image geometry and observing metadata.
"""
ImageCube.__init__(self, image_cube, freqs, meta)
[docs]
def trim(self, new_theta_fov):
"""Trim the image in place to ``new_theta_fov`` radians."""
n = new_theta_fov / self.meta.res
nx, ny = self.meta.shape
i = int((nx - n) / 2.)
if i > 0:
self.data = self.data[:, i:nx - i, i:ny - i]
self.meta.slice(i, nx - i, i, ny - i)
[docs]
def apply_window_function(self, win_fct, add_to_meta=True):
"""Apply an image-plane window in place."""
win_mask = win_fct.generate(self.meta)
self.data = self.data * win_mask
if add_to_meta:
win_fct.to_meta(self.meta)
[docs]
def ft(self, umin, umax):
"""Fourier transform image cube and return a CartDataCube.
Args:
umin (float): Min U in wavelength
umax (float): Max U in wavelength
Returns:
CartDataCube: a new visibility cube.
"""
uu, vv, ft_cube = psutil.ft_cart_cube(self.data, self.meta.res, umin, umax)
return CartDataCube(ft_cube, uu, vv, self.freqs, self.meta.copy())
[docs]
def save_to_fits(self, fname, overwrite=True):
"""Write the image cube to a FITS file (in Kelvin).
Args:
fname: output file path.
overwrite (bool): replace an existing file.
"""
hdu = pf.PrimaryHDU(self.data[None].real)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
hdu.header.update(self.meta.to_header(add_origin=True))
hdu.header["BUNIT"] = "K"
hdu.header["BTYPE"] = "Intensity"
hdu.header["CRVAL3"] = self.freqs[0]
hdu.header["CDELT3"] = self.meta.freq_width
hdu.writeto(fname, overwrite=overwrite)
[docs]
def plot(self, fmhz='med', action_fct=None, theta_lines=None, ax=None, title=None,
auto_scale_quantiles=None, **kargs):
"""Show one channel (or a frequency reduction) as a sky image.
Args:
fmhz: channel to show, in MHz, or ``'med'`` / ``'first'`` / ``'last'``
(ignored if ``action_fct`` is given).
action_fct: a reducer ``f(data, axis=0)`` over frequency instead of
one channel.
theta_lines (list): angular radii (deg) at which to draw guide circles.
ax: matplotlib Axes (new figure if None).
auto_scale_quantiles: ``(lo, hi)`` quantiles for the colour limits.
**kargs: forwarded to imshow.
"""
if theta_lines is None:
theta_lines = []
if ax is None:
_fig, ax = plt.subplots()
d = self.data
if action_fct is None:
if fmhz == 'med':
fmhz = (self.freqs[0] + (self.freqs[-1] - self.freqs[0]) / 2.) * 1e-6
elif fmhz == 'first':
fmhz = self.freqs[0] * 1e-6
elif fmhz == 'last':
fmhz = self.freqs[-1] * 1e-6
i = np.argmin(abs(self.freqs - fmhz * 1e6))
d = d[i]
else:
d = action_fct(d, axis=0)
theta_max = np.clip(0.5 * self.meta.theta_fov, 0, 1)
show_degrees = theta_max < 0.4
if auto_scale_quantiles is not None:
kargs['vmin'] = np.quantile(d, auto_scale_quantiles[0])
kargs['vmax'] = np.quantile(d, auto_scale_quantiles[1])
psutil.plot_cart_map(d, theta_max, ax=ax, title=title, theta_lines=theta_lines,
show_degrees=show_degrees, **kargs)
[docs]
def plot_slice(self, ax=None, **kargs):
"""Plot the central spatial slice as a function of frequency."""
if ax is None:
_fig, ax = plt.subplots()
islice = self.data[:, :, self.meta.shape[0] // 2]
cbs = psutil.ColorbarSetting(psutil.ColorbarOutterPosition(width='3%'))
theta_max = np.clip(0.5 * self.meta.theta_fov, 0, 1)
show_degrees = theta_max < 0.4
if show_degrees:
theta_max = np.degrees(theta_max)
extent = np.array([-theta_max, theta_max, min(self.freqs) * 1e-6, max(self.freqs) * 1e-6])
im_mappable = ax.imshow(islice.real, extent=extent, aspect='auto', **kargs)
cbs.add_colorbar(im_mappable, ax)
ax.set_ylabel('Freqs (MHz)')
if show_degrees:
ax.set_xlabel('m (l = 0) [deg]')
else:
ax.set_xlabel('m (l = 0)')
[docs]
def copy(self):
"""Return an independent deep copy of this image cube."""
return CartImageCube(self.data.copy(), self.freqs, self.meta.copy())
[docs]
class GriddedCartDataCube:
"""Visibility cube stored on a regular Cartesian UV grid."""
def __init__(self, g_vis, g_uu, g_vv, freqs, meta):
"""Create a gridded visibility cube.
Args:
g_vis (n_freqs, n_u, n_v): gridded visibilities cube array
g_uu (n_u, n_v): U in lambda
g_vv (n_u, n_v): V in lambda
freqs (n_freqs): Frequencies in Hz
meta: Image geometry and observing metadata.
"""
self.data = g_vis
self.g_uu = g_uu
self.g_vv = g_vv
self.freqs = freqs
self.meta = meta
[docs]
def image(self, res=None, low_memory=False):
"""Fourier-transform the gridded visibilities to the image plane.
Args:
res: output pixel size, in radians (default: set by the uv grid).
low_memory (bool): transform one channel at a time.
Returns:
CartImageCube: the sky image cube.
"""
du = abs(self.g_uu[0, 1] - self.g_uu[0, 0])
s = None
meta = self.meta.copy()
data = self.data
if res is not None:
nu = int(1 / res / du)
s = (nu, nu)
meta = ImageMetaData.from_res(res, s, **meta.kargs)
data = np.array([psutil.resize(k, s) for k in data])
def ft_fct(data, axes):
return psutil.vis_to_img(data, axes=axes)
image_cube = ft_fct(data, (1, 2))
return CartImageCube(image_cube, self.freqs, meta)
[docs]
class DataCubeCombiner:
"""Incrementally combine visibility cubes using configurable weights."""
def __init__(self, umin, umax, weighting_mode='uv', inhomogeneous=False, w_square=False):
self.cube = None
self.weighting_mode = weighting_mode
self.total_weights = None
self.total_weights_psf = None
self.total_time = 0
self.umin = umin
self.umax = umax
self.inhomogeneous = inhomogeneous
self.w_square = w_square
self.night_ids = []
self.freqs_n_nights = []
def _get_weights(self, cube, with_uv_scale=True):
"""Combining weights for ``cube`` under the current ``weighting_mode``."""
if self.weighting_mode == 'uv':
w = np.median(cube.weights.get(with_uv_scale=with_uv_scale), axis=0)
elif self.weighting_mode == 'full':
w = cube.weights.get(with_uv_scale=with_uv_scale)
elif self.weighting_mode == 'global':
w = np.median(cube.weights.get(with_uv_scale=with_uv_scale))
elif self.weighting_mode == 'none':
w = 1
else:
raise AssertionError(f"'{self.weighting_mode}' incorrect")
if self.w_square:
w = w ** 2
return np.atleast_2d(w)
def _add(self, a, b, idx1_d1, idx1_d2):
"""Add ``b`` into the ``idx1_d1`` x ``idx1_d2`` sub-block of ``a`` in place."""
a_1 = a[idx1_d1]
a_1[:, idx1_d2] = a_1[:, idx1_d2] + b
a[idx1_d1] = a_1
return a
[docs]
def add(self, cube, night_id):
"""Accumulate one cube and its observing-night identifier."""
self.night_ids.append(night_id)
if self.cube is None:
if self.inhomogeneous:
cube = cube.make_full_cube(self.umin, self.umax)
weights = self._get_weights(cube)
self.cube = cube.new_with_data(weights * cube.data)
self.total_weights = weights
self.total_weights_psf = cube.weights.get(with_uv_scale=False)
self.total_time = self.cube.meta.total_time
self.freqs_n_nights = np.ones(len(self.cube.freqs))
else:
idx1, idx1_uv, idx2, idx2_uv = get_common_idx(self.cube, cube)
cube = cube.get_slice_from_idx(idx2)
cube.filter_uv_from_index(idx2_uv)
weights = self._get_weights(cube)
idx1_w = idx1 if np.squeeze(weights).ndim == 2 else slice(None)
idx1_w_uv = idx1_uv if np.squeeze(weights).ndim >= 1 else slice(None)
if not self.inhomogeneous:
self.cube = self.cube.get_slice_from_idx(idx1)
self.cube.filter_uv_from_index(idx1_uv)
self.cube.data = self.cube.data + weights * cube.data
self.cube.weights = self.cube.weights + cube.weights
self.total_weights_psf = self.total_weights_psf[idx1][:, idx1_uv] \
+ cube.weights.get(with_uv_scale=False)
self.total_weights = self.total_weights[idx1_w][:, idx1_w_uv] + weights
self.freqs_n_nights = self.freqs_n_nights[idx1] + 1
else:
self.cube.data = self._add(self.cube.data, weights * cube.data, idx1, idx1_uv)
self.cube.weights.data = self._add(self.cube.weights.data, cube.weights.data, idx1,
idx1_uv)
self.total_weights_psf = self._add(
self.total_weights_psf, cube.weights.get(with_uv_scale=False), idx1, idx1_uv)
self.total_weights = self._add(self.total_weights, weights, idx1_w, idx1_w_uv)
self.freqs_n_nights[idx1] += 1
self.total_time += cube.meta.total_time
[docs]
def get(self, min_n_nights=None):
"""The inverse-variance-weighted combination of all added cubes.
Args:
min_n_nights: drop channels covered by fewer than this many nights.
Returns:
CartDataCube: the combined cube (``None`` if nothing was added).
"""
if self.cube is None:
return self.cube
cube = self.cube.new_with_data(
np.divide(self.cube.data, self.total_weights, where=self.total_weights != 0),
weights=self.cube.weights)
cube.meta.set('PETOTTIM', self.total_time)
cube.weights.meta.set('PETOTTIM', self.total_time)
cube.meta.set('NIGHTS', ','.join(self.night_ids))
cube.meta.set('PECMODE', self.weighting_mode)
cube.meta.set('PECHOMO', not self.inhomogeneous)
tw_psf_m = self.total_weights_psf.mean(axis=0)
uv_scale = abs(np.divide(cube.weights.data.mean(axis=0), tw_psf_m, where=tw_psf_m != 0))
cube.weights.uv_scale = uv_scale[None, :]
cube.weights.data = self.total_weights_psf
cube.weights.freqs_n_nights = self.freqs_n_nights
if min_n_nights is not None:
print('Filter:', cube.freqs[self.freqs_n_nights < min_n_nights])
cube.filter_outliers(self.freqs_n_nights < min_n_nights)
idx_zero = cube.weights.data.mean(axis=1) != 0
cube = cube.get_slice_from_idx(idx_zero)
return cube
[docs]
class DataCube:
"""Base class for frequency-dependent data and optional weights."""
def __init__(self, data, freqs, weights=None):
self.freqs = freqs
self.data = data
self.weights = weights
def __add_sub__(self, other, sub=False):
if not np.allclose(self.freqs, other.freqs):
raise TypeError('The two DataCube must have the same frequencies channels')
if sub:
weights = psutil.safe_sum(self.weights, other.weights)
return self.new_with_data(self.data - other.data, weights=weights)
weights = psutil.safe_sum(self.weights, other.weights)
return self.new_with_data(self.data + other.data, weights=weights)
def __sub__(self, other):
return self.__add_sub__(other, True)
def __add__(self, other):
return self.__add_sub__(other, False)
def __mul__(self, other):
assert psutil.is_number(other)
weights = None
if self.weights is not None:
weights = other * self.weights
return self.new_with_data(other * self.data, weights=weights)
def __rmul__(self, other):
return self.__mul__(other)
[docs]
def get_unique_xy(self):
"""The unique (u, v) coordinates of the cube's modes. Subclass hook."""
raise NotImplementedError()
[docs]
def set_weights(self, weights_cube):
"""Attach a copy of ``weights_cube``."""
if weights_cube is None:
self.weights = None
else:
self.weights = weights_cube.copy()
[docs]
def new_with_data(self, data, weights=None, freqs=None):
"""Build a cube of the same type/geometry with new ``data`` (and,
optionally, ``weights`` / ``freqs``). Subclass hook."""
raise NotImplementedError()
[docs]
def get_slice_from_idx(self, idx_freqs):
"""The channels selected by ``idx_freqs`` (a mask, index array, or slice).
Returns:
DataCube: the frequency-selected cube (same type as ``self``).
"""
weights = None
if self.weights is not None:
weights = self.weights.get_slice_from_idx(idx_freqs)
return self.new_with_data(self.data[idx_freqs], weights=weights,
freqs=self.freqs[idx_freqs])
[docs]
def get_slice(self, freq_start, freq_end):
"""The channels within ``[freq_start, freq_end]`` (Hz).
Returns:
DataCube: the frequency-sliced cube.
"""
idx_freqs = psutil.get_freq_slice(self.freqs, freq_start, freq_end)
return self.get_slice_from_idx(idx_freqs)
[docs]
def get_freq(self, freq):
"""The single channel at or just above ``freq`` (Hz).
Returns:
DataCube: a one-channel cube.
"""
i = np.nonzero(self.freqs >= freq)[0][0]
return self.get_slice_from_idx(slice(i, i + 1))
[docs]
def filter_uv_from_index(self, idx_uv):
"""Filter spatial modes in place using ``idx_uv``."""
self.data = self.data[:, idx_uv]
if self.weights is not None:
self.weights.filter_uv_from_index(idx_uv)
[docs]
def make_diff_cube(self):
"""Adjacent-channel differences (times sqrt(0.5)) -- a signal-free,
noise-like cube used to estimate the noise level.
Returns:
DataCube: the difference cube (one fewer channel).
"""
data = np.sqrt(0.5) * np.diff(self.data, axis=0)
weights = None
if self.weights is not None:
weights = self.weights.new_with_data(self.weights.data[:-1], freqs=self.freqs[:-1])
return self.new_with_data(data, weights=weights, freqs=self.freqs[:-1])
[docs]
def make_diff_cube_interp(self):
"""Like :meth:`make_diff_cube` but padded back to the original number of
channels (last difference repeated).
Returns:
DataCube: the difference cube (same channel count).
"""
data = np.sqrt(0.5) * np.diff(self.data, axis=0)
data = np.vstack([data, data[-2:-1]])
return self.new_with_data(data)
[docs]
def copy(self):
"""Return an independent cube copy."""
w = None
if self.weights is not None:
w = self.weights.copy()
return self.new_with_data(self.data.copy(), weights=w)
[docs]
@staticmethod
def load(filename):
"""Load a data cube from HDF5, dispatching on the stored format.
Returns:
CartDataCube or SphDataCube: whichever the file holds.
"""
with tables.open_file(filename, 'r') as h5_file:
if hasattr(h5_file.root, 'ft_cube'):
return CartDataCube.load(filename)
elif hasattr(h5_file.root, 'alm_cube'):
from .sphcube import SphDataCube
return SphDataCube.load(filename)
else:
raise ValueError(f'File {filename} is not of a supported format')
[docs]
class CartDataCube(DataCube):
"""Complex visibilities on a non-gridded Cartesian UV plane."""
def __init__(self, data, uu, vv, freqs, meta, weights=None):
"""Create a Cartesian visibility cube.
Args:
data (n_freqs, n_vis): visibilities cube array
uu (n_vis): U in lambda
vv (n_vis): V in lambda
freqs (n_freqs): Frequencies in Hz
meta: Geometry and observing metadata of the source images.
"""
self.uu = uu
self.vv = vv
self.meta = meta
self.ru = np.sqrt(uu ** 2 + vv ** 2)
DataCube.__init__(self, data, freqs, weights)
def __add_sub__(self, other, sub=False):
if not np.allclose(self.meta.theta_fov, other.meta.theta_fov):
print(f'Warning: The two datacubes does not have the same FoV {self.meta.theta_fov} vs '
f'{other.meta.theta_fov}')
elif not np.allclose(self.meta.win_fct_power, other.meta.win_fct_power):
print(f'Warning: The two datacubes does not have the same window fct '
f'{self.meta.win_fct_power} vs {other.meta.win_fct_power}')
elif not np.allclose(self.meta.chan_width, other.meta.chan_width):
print(f'Warning: The two datacubes does not have the same channel width '
f'{self.meta.chan_width} vs {other.meta.chan_width}')
return DataCube.__add_sub__(self, other, sub=sub)
[docs]
def get_unique_xy(self):
"""A per-mode key encoding (u, v), used to match modes between cubes.
Returns:
ndarray: one float per mode, combining rounded ``uu`` and ``vv``.
"""
return np.round(self.uu, decimals=2) + 1e-6 * np.round(self.vv, decimals=2)
[docs]
@staticmethod
def load_from_fits(files, umin, umax, convert_jy2k=True):
"""Build a visibility cube from UV-plane FITS files (already gridded uv,
not images).
Args:
files: UV-plane FITS files, in frequency order.
umin, umax: baseline range kept, in wavelengths.
convert_jy2k (bool): convert from Jy to Kelvin.
Returns:
CartDataCube: the visibility cube.
"""
freqs = []
data_cube = []
pr = psutil.progress_report(len(files))
for i, file in enumerate(files):
pr(i)
hdu = pf.open(file)[0]
data = hdu.data.squeeze()
shape = data.shape
du = hdu.header['CDELT1']
freq = hdu.header['CRVAL3']
fov = 1 / du
res = fov / shape[0]
if convert_jy2k:
lamb = const.c.value / freq
jy2k = ((1e-26 * lamb ** 2) / (2 * const.k_B.value))
data = data * jy2k
uu, vv, idx = psutil.get_ungrid_vis_idx(shape, res, umin, umax)
data_cube.append(data[idx])
freqs.append(freq)
meta = ImageMetaData.from_res(res, shape, **hdu.header)
return CartDataCube(np.array(data_cube), uu, vv, np.array(freqs), meta)
[docs]
@staticmethod
def load_from_fits_image(files, umin, umax, theta_fov, imager_scale_factor=None,
convert_jy2k=True, compat_wscnormf='old_normpsf',
int_time=None, total_time=None,
window_function=None, data_dtype=np.complex128):
"""Build a visibility cube from frequency-dependent FITS images (no PSF).
Images are optionally trimmed and windowed, converted from Jy/PSF to
Kelvin, Fourier-transformed, and restricted to ``[umin, umax]``. Use
:meth:`load_from_fits_image_and_psf` when a PSF is available (it also
derives weights).
Args:
files: FITS image files, in frequency order.
umin, umax: baseline range kept, in wavelengths.
theta_fov: output field of view, in radians.
imager_scale_factor: explicit PSF area (pixel units); by default
taken from ``WSCNORMF`` or a Gaussian-beam approximation.
convert_jy2k (bool): convert Jy/PSF images to Kelvin.
compat_wscnormf (str): interpretation of legacy ``WSCNORMF`` metadata.
int_time, total_time: integration and total observing time (s),
stored in the metadata.
window_function: optional image-plane :class:`WindowFunction`.
Returns:
CartDataCube: the visibility cube (unweighted).
"""
ft_cube = []
freqs = []
pr = psutil.progress_report(len(files))
omega_gauss_warning = False
compat_warning = False
idx_ft = None
meta = None
if theta_fov is not None:
theta_fov = 2 * np.sin(theta_fov / 2.)
for i, file in enumerate(files):
pr(i)
hdu = pf.open(file)[0]
header = hdu.header
data = hdu.data.astype(data_dtype)
res = abs(np.radians(header['CDELT1']))
freq_start = header['CRVAL3']
df = header['CDELT3']
nf = header['NAXIS3']
freq_end = freq_start + (nf - 1) * df
fits_freqs = np.linspace(freq_start, freq_end, nf)
lamb = const.c.value / fits_freqs
cart_map = np.squeeze(data)
if cart_map.ndim == 2:
cart_map = cart_map[None, :, :]
_, nx, ny = cart_map.shape
if convert_jy2k:
if 'WSCNORMF' in header and imager_scale_factor is None:
if compat_wscnormf == 'old_normpsf':
imager_scale_factor = header['WSCNORMF']
if not compat_warning:
print('Warning: using WSCNORMF as obtained by normpsf < 16/06/2017')
elif compat_wscnormf == 'old_wsclean':
imager_scale_factor = (nx * ny) / header['WSCNORMF'] / 4.
if not compat_warning:
print('Warning: using WSCNORMF as obtained by wsclean < 16/06/2017')
else:
imager_scale_factor = (nx * ny) / header['WSCNORMF']
if imager_scale_factor is not None:
omega = imager_scale_factor * res ** 2
else:
bmaj = header['BMAJ']
bmin = header['BMIN']
omega = np.radians(bmaj) * np.radians(bmin) * np.pi / (4 * np.log(2))
imager_scale_factor = omega / (res ** 2)
if not omega_gauss_warning:
omega_gauss_warning = True
print(f'Warning: WSCNORMF not found, using Gaussian approx. of the PSF '
f'({imager_scale_factor:.2f})')
jy2k = ((1e-26 * lamb ** 2) / (2 * const.k_B.value))
jypsf2K = jy2k / omega
cart_map = cart_map * jypsf2K[:, None, None]
if theta_fov is not None:
n = theta_fov / res
i = int((nx - n) / 2.)
if i > 0:
cart_map = cart_map[:, i:-i, i:-i]
if window_function is not None:
mask = window_function.generate_window(cart_map.shape[1])
cart_map = cart_map * mask[None, :, :]
if meta is None:
meta = ImageMetaData.from_header(header, (nx, ny))
if theta_fov is not None and i > 0:
meta.slice(i, nx - i, i, ny - i)
if window_function is not None:
window_function.to_meta(meta)
if idx_ft is None:
uu, vv, idx_ft = psutil.get_ungrid_vis_idx(cart_map.shape[1:], res, umin, umax)
ft = psutil.img_to_vis(cart_map, axes=(1, 2))
uu = uu.flatten()
vv = vv.flatten()
for freq, ft_slice in zip(fits_freqs, ft, strict=False):
ft_cube.append(ft_slice[idx_ft])
freqs.append(freq)
if int_time is not None and 'PEINTTIM' not in meta:
meta.set('PEINTTIM', int_time)
if total_time is not None and 'PETOTTIM' not in meta:
meta.set('PETOTTIM', total_time)
return CartDataCube(np.array(ft_cube), uu, vv, np.array(freqs), meta)
[docs]
@staticmethod
def load_from_fits_image_and_psf(
files, files_psf, umin, umax, theta_fov, int_time=None, total_time=None,
convert_jy2k=True, min_weight_ratio=0.01, trim_method='before', use_wscnormf=False,
compat_wscnormf='old_normpsf', window_function=None, abs_min_weight=0.5,
data_dtype=np.complex128):
"""Build a visibility cube from matched image + PSF FITS files.
The standard loader when a per-frequency PSF is available: it divides the
imaged visibilities by the PSF (undoing the dirty-image sampling response)
and derives per-mode weights from that PSF, giving a properly weighted cube.
Args:
files: image FITS files, in frequency order.
files_psf: matching PSF FITS files, one per image.
umin, umax: baseline range kept, in wavelengths.
theta_fov: output field of view, in radians.
int_time, total_time: integration and total observing time (s),
stored in the metadata for later noise / SEFD estimation.
convert_jy2k (bool): convert Jy/PSF images to Kelvin.
min_weight_ratio (float): drop modes whose weight is below this
fraction of the per-mode maximum (0 disables).
abs_min_weight (float): absolute floor applied to that threshold.
trim_method (str): ``'before'`` trims to ``theta_fov`` before the PSF
division, ``'after'`` trims afterwards.
use_wscnormf (bool): normalise from the image ``WSCNORMF`` metadata
instead of the PSF (weights still come from the PSF files).
window_function: optional image-plane :class:`WindowFunction`.
Returns:
CartDataCube: the PSF-weighted visibility cube.
"""
if use_wscnormf:
ft_I_cube = CartDataCube.load_from_fits_image(
files, umin, umax, theta_fov, compat_wscnormf=compat_wscnormf,
convert_jy2k=convert_jy2k, total_time=total_time, int_time=int_time,
window_function=window_function, data_dtype=data_dtype)
weight_cube = CartWeightCube.load_from_fits_psf(
files_psf, umin, umax, int_time, total_time, theta_fov=theta_fov,
output_psf_cube=False, window_function=window_function, data_dtype=data_dtype)
ft_I_cube.set_weights(weight_cube)
return ft_I_cube
b_theta_fov = None
if trim_method in ['b', 'before']:
b_theta_fov = theta_fov
ft_I_cube = CartDataCube.load_from_fits_image(
files, umin, umax, b_theta_fov, imager_scale_factor=1, convert_jy2k=convert_jy2k,
total_time=total_time, int_time=int_time, window_function=window_function,
data_dtype=data_dtype)
ft_psf_cube, weight_cube = CartWeightCube.load_from_fits_psf(
files_psf, umin, umax, int_time, total_time, theta_fov=b_theta_fov,
output_psf_cube=True, window_function=window_function)
f = 1 / float(ft_I_cube.meta.shape[0] ** 2)
with np.errstate(divide='ignore', invalid='ignore'):
d_over_psf = np.where(ft_psf_cube.data != 0,
np.divide(ft_I_cube.data, ft_psf_cube.data), 0)
ft_I_rw_cube = ft_I_cube.new_with_data(d_over_psf * f, weights=weight_cube)
if min_weight_ratio > 0:
min_weight = min_weight_ratio * ft_I_rw_cube.weights.get().min(axis=0).max()
min_weight = np.max([abs_min_weight, min_weight])
ft_I_rw_cube.filter_min_weight(min_weight, replace=False)
if trim_method in ['after', 'a']:
ft_I_rw_cube = ft_I_rw_cube.reduce_fov(theta_fov, umin=umin, umax=umax)
return ft_I_rw_cube
[docs]
@staticmethod
def load_from_hd5(h5_group):
"""Deprecated alias for :meth:`load_from_h5`."""
warnings.warn('load_from_hd5 is deprecated; use load_from_h5.', DeprecationWarning,
stacklevel=2)
return CartDataCube.load_from_h5(h5_group)
[docs]
@staticmethod
def load_from_h5(h5_group):
"""Load a Cartesian visibility cube from an open HDF5 group."""
ft_cube = h5_group.data.read()
freqs = h5_group.freqs.read()
uu = h5_group.uu.read()
vv = h5_group.vv.read()
attrs = h5_group.data.attrs
header = {k: psutil.safe_decode_bytes(attrs[k]) for k in attrs._f_list() if k[0].isupper()}
if 'WCSAXES' in attrs:
meta = ImageMetaData.from_header(header, attrs.shape)
else:
shape = (attrs.nx, attrs.ny) if 'shape' not in attrs else attrs.shape
meta = ImageMetaData.from_res(attrs.res, shape, **header)
return CartDataCube(ft_cube, uu, vv, freqs, meta)
[docs]
@staticmethod
def load(filename):
"""Load a Cartesian visibility cube (and its weights, if present) from HDF5.
Returns:
CartDataCube: the loaded cube.
"""
with tables.open_file(filename, 'r') as h5_file:
cart_cube = CartDataCube.load_from_h5(h5_file.root.ft_cube)
if 'weights' in h5_file.root:
cart_cube.weights = CartWeightCube.load_from_h5(h5_file.root.weights)
return cart_cube
[docs]
@staticmethod
def join_cubes(cubes):
"""Concatenate cubes along frequency, keeping the first copy of any
overlapping channels.
Returns:
CartDataCube: the joined cube.
"""
j_cube = cubes[0]
for cube in cubes[1:]:
idx_new = ~np.isin(_fmhz(cube.freqs), _fmhz(j_cube.freqs))
j_cube.data = np.vstack([j_cube.data, cube.data[idx_new]])
j_cube.freqs = np.concatenate([j_cube.freqs, cube.freqs[idx_new]])
if cubes[0].weights is not None:
weights = CartWeightCube.join_cubes([c.weights for c in cubes])
else:
weights = None
return CartDataCube(j_cube.data, j_cube.uu, j_cube.vv, j_cube.freqs, j_cube.meta,
weights=weights)
[docs]
def save_to_hd5(self, h5_file, group):
"""Deprecated alias for :meth:`save_to_h5`."""
warnings.warn('save_to_hd5 is deprecated; use save_to_h5.', DeprecationWarning,
stacklevel=2)
self.save_to_h5(h5_file, group)
[docs]
def save_to_h5(self, h5_file, group):
"""Write this cube into an open HDF5 group."""
h5_file.create_array(group, 'data', self.data, "Visibilities (K)")
h5_file.create_array(group, 'freqs', self.freqs, "Frequencies (Hz)")
h5_file.create_array(group, 'uu', self.uu, "U (lambda)")
h5_file.create_array(group, 'vv', self.vv, "V (lambda)")
with warnings.catch_warnings():
warnings.simplefilter("ignore")
for key, value in self.meta.items(add_origin=True):
group.data.attrs[key] = value
[docs]
def save(self, filename):
"""Save the cube and optional weights to HDF5."""
with tables.open_file(filename, 'w') as h5_file:
group = h5_file.create_group("/", 'ft_cube', 'Visibilty cube (n_freqs, n_vis)')
self.save_to_h5(h5_file, group)
if self.weights is not None:
group = h5_file.create_group("/", 'weights', 'Weights cube (n_freqs, n_vis)')
self.weights.save_to_h5(h5_file, group)
[docs]
def new_with_data(self, data, weights=None, freqs=None, interpolate_weights=True):
"""A new cube with new ``data`` on this cube's uv geometry.
Args:
data (n_freqs, n_vis): the new visibilities.
weights: weights to attach (default: this cube's).
freqs: frequencies for ``data`` (default: this cube's).
interpolate_weights (bool): interpolate the weights onto ``freqs``
when the frequency grids differ.
Returns:
CartDataCube: the new cube.
"""
if freqs is None:
freqs = self.freqs
if weights is None:
weights = self.weights
assert data.shape[0] == len(freqs)
assert data.shape[1] == len(self.uu)
if weights is not None:
weights = weights.copy()
if weights.data.shape[0] != len(freqs):
w = scipy.interpolate.interp1d(weights.freqs, weights.data, bounds_error=False,
fill_value='extrapolate', axis=0)(freqs)
weights = weights.new_with_data(w, freqs=freqs)
return CartDataCube(data, self.uu, self.vv, freqs, self.meta.copy(), weights=weights)
[docs]
def make_full_cube(self, umin, umax, output_idx=False):
"""Place the data on a complete frequency + uv grid, filling gaps with zero.
Args:
umin, umax: baseline range of the target uv grid, in wavelengths.
output_idx (bool): also return the index arrays mapping this cube
into the full grid.
Returns:
CartDataCube: the gap-filled cube (or ``(cube, idx1, idx1_uv)`` when
``output_idx``).
"""
freqs = np.array(sorted(np.concatenate((self.freqs, psutil.get_freqs_gaps(self.freqs)))))
uu, vv, _ = psutil.get_ungrid_vis_idx(self.meta.shape, self.meta.res, umin, umax)
weights = None
if self.weights is not None:
weights = self.weights.make_full_cube(umin, umax)
cube = CartDataCube(np.zeros((len(freqs), len(uu)), dtype=np.complex128),
uu, vv, freqs, self.meta.copy(), weights=weights)
idx1, idx1_uv, _idx2, _idx2_uv = get_common_idx(cube, self)
data_f = cube.data[idx1]
data_f[:, idx1_uv] = self.data
cube.data[idx1] = data_f
if output_idx:
return cube, idx1, idx1_uv
return cube
[docs]
def estimate_sefd(self, sefd_jansky=True):
"""Overall system-equivalent flux density from this cube and its weights.
Returns:
float: the SEFD, in Jy (``sefd_jansky=True``) or Kelvin.
"""
assert self.weights is not None
return self.weights.estimate_sefd(self, sefd_jansky=sefd_jansky)
[docs]
def estimate_uv_sefd(self, sefd_jansky=True):
"""SEFD per uv sample (see :meth:`estimate_sefd` for units).
Returns:
CartDataCube: the per-mode SEFD.
"""
assert self.weights is not None
return self.weights.estimate_uv_sefd(self, sefd_jansky=sefd_jansky)
[docs]
def estimate_freqs_sefd(self, sefd_jansky=True, sefd_poly_fit_deg=0):
"""SEFD per frequency channel (see :meth:`estimate_sefd` for units).
Args:
sefd_jansky (bool): return Jy (else Kelvin).
sefd_poly_fit_deg (int): if > 0, smooth the SEFD with a polynomial
of this degree in log-log space.
Returns:
ndarray: the SEFD per channel, shape ``(n_freqs,)``.
"""
assert self.weights is not None
sefd = self.weights.estimate_freqs_sefd(self, sefd_jansky=sefd_jansky)
if sefd_poly_fit_deg > 0:
sefd_fct = np.poly1d(np.polyfit(np.log(self.freqs), np.log(sefd), sefd_poly_fit_deg))
sefd = np.exp(sefd_fct(np.log(self.freqs)))
return sefd
[docs]
def get_hermitian_index(self):
"""Indices pairing each mode with its Hermitian conjugate (-u, -v).
Returns:
tuple: ``(idx1, idx2)`` -- the positive-half modes and their
conjugate partners.
"""
idx1_a = (self.uu > 0)
idx1_b = ((self.uu == 0) & (self.vv > 0))
idx2_a = psutil.get_selection_index(self.uu, self.vv, - self.uu[idx1_a],
- self.vv[idx1_a], True)
idx2_b = psutil.get_selection_index(self.uu, self.vv, self.uu[idx1_b],
- self.vv[idx1_b], True)
idx1 = np.concatenate([np.where(idx1_a)[0], np.where(idx1_b)[0]])
idx2 = np.concatenate([idx2_a, idx2_b])
return idx1, idx2
[docs]
def regrid(self):
"""Place the visibilities on their regular 2-D uv grid.
Returns:
GriddedCartDataCube: the gridded cube.
"""
nx, ny = self.meta.shape
g_uu, g_vv, idx = psutil.get_regrid_vis_idx(self.uu, self.vv, self.meta.res,
self.meta.shape)
g_data = np.zeros((self.freqs.size, g_uu.size), dtype=np.complex128)
g_data[:, idx] = self.data
g_data = g_data.reshape((self.freqs.size, nx, ny))
return GriddedCartDataCube(g_data, g_uu, g_vv, self.freqs, self.meta.copy())
[docs]
def image(self):
"""Fourier-transform the visibilities to the image plane.
Returns:
CartImageCube: the image cube.
"""
return self.regrid().image()
[docs]
def reduce_fov(self, new_fov, low_memory=False, umin=None, umax=None):
"""A copy with a smaller image-plane field of view.
Regrids, images, trims to ``new_fov`` (radians), and Fourier-transforms
back. A ``new_fov`` larger than the current one returns ``self``.
Args:
new_fov: target field of view, in radians.
low_memory (bool): process one channel at a time.
umin, umax: baseline range of the output (defaults to this cube's).
Returns:
CartDataCube: the reduced-FoV cube (weights reduced too if present).
"""
if new_fov >= self.meta.theta_fov:
return self
if umin is None:
umin = np.round(self.ru.min(), 2)
if umax is None:
umax = np.round(self.ru.max(), 2)
if self.weights is not None:
new_weights = self.weights.reduce_fov(new_fov, low_memory=low_memory, umin=umin,
umax=umax)
if not low_memory:
img_cube = self.regrid().image()
img_cube.trim(new_fov)
new_cube = img_cube.ft(umin, umax)
else:
trimmed_cubes = []
for freq in self.freqs:
img_cube = self.get_freq(freq).regrid().image()
img_cube.trim(new_fov)
trimmed_cubes.append(img_cube.ft(umin, umax))
data = np.array([c.data[0] for c in trimmed_cubes])
freqs = np.array([c.freqs[0] for c in trimmed_cubes])
new_cube = CartDataCube(data, trimmed_cubes[0].uu, trimmed_cubes[0].vv, freqs,
trimmed_cubes[0].meta)
if self.weights is not None:
new_cube.set_weights(new_weights)
return new_cube
[docs]
def apply_window_function(self, win_fct, umin=None, umax=None, add_to_meta=True):
"""A copy with the image-plane window function ``win_fct`` applied.
Args:
win_fct: the :class:`WindowFunction` to apply (or its name).
umin, umax: baseline range of the output (defaults to this cube's).
add_to_meta (bool): record the window in the output metadata.
Returns:
CartDataCube: the windowed cube.
"""
if umin is None:
umin = np.round(self.ru.min(), 2)
if umax is None:
umax = np.round(self.ru.max(), 2)
if self.weights is not None:
new_weights = self.weights.apply_window_function(win_fct, umin=umin, umax=umax,
add_to_meta=add_to_meta)
img_cube = self.regrid().image()
img_cube.apply_window_function(win_fct, add_to_meta=add_to_meta)
new_cube = img_cube.ft(umin, umax)
if self.weights is not None:
new_cube.set_weights(new_weights)
return new_cube
[docs]
def filter_uvrange(self, umin, umax):
"""Keep UV samples within ``[umin, umax]`` in place."""
idx_uv = (self.ru >= umin) & (self.ru <= umax)
self.filter_uv_from_index(idx_uv)
[docs]
def filter_outliers(self, idx_outliers):
"""Remove channels marked by ``idx_outliers`` in place."""
self.freqs = self.freqs[~idx_outliers]
self.data = self.data[~idx_outliers]
if self.weights is not None:
self.weights.filter_outliers(idx_outliers)
[docs]
def filter_freqs_from_other(self, other):
"""Keep only channels also present in ``other``."""
idx1 = np.isin(_fmhz(self.freqs), _fmhz(other.freqs))
self.filter_outliers(~idx1)
[docs]
def filter_nan(self):
"""Remove channels containing any NaN visibility."""
idx_nan = np.any(np.isnan(self.data), 1)
if len(self.freqs[idx_nan]) > 0:
print('SB with NaN:', self.freqs[idx_nan])
self.filter_outliers(idx_nan)
[docs]
def filter_min_weight(self, min_weight, replace=False, verbose=True):
"""Discard low-weight uv data in place.
Args:
min_weight: weight threshold below which data is removed.
replace (bool): if True, zero individual samples below the threshold
(keeping the mode); if False, drop whole modes whose median
weight is below it.
verbose (bool): print how many modes/samples were filtered.
"""
if self.weights is not None and min_weight > 0:
if replace:
idx_uv = self.weights.get() >= min_weight
n_filt = np.sum(~idx_uv)
n_tot = float(len(self.uu) * len(self.freqs))
if verbose:
print(f'Filtering {n_filt} visibilities ({n_filt / n_tot * 100:.2f} %)')
self.data[~idx_uv] = 0
self.weights.data[~idx_uv] = 0
else:
idx_uv = np.median(self.weights.get(), axis=0) >= min_weight
n_filt = np.sum(~idx_uv)
n_tot = float(len(self.uu))
if verbose:
print(f'Filtering {n_filt} modes ({n_filt / n_tot * 100:.2f} %)')
self.filter_uv_from_index(idx_uv)
[docs]
def filter_sefd_uv(self, max_sefd, min_sefd=0):
"""Keep, in place, only uv modes whose mean estimated SEFD is within
``[min_sefd, max_sefd]``.
Returns:
ndarray: boolean mask of the kept modes.
"""
sefd = self.estimate_uv_sefd().data.mean(axis=0)
idx_uv = (sefd >= min_sefd) & (sefd <= max_sefd)
n_filt = np.sum(~idx_uv)
n_tot = float(len(self.uu))
print(f'Filtering {n_filt} modes ({n_filt // n_tot * 100:.2f} %)')
self.filter_uv_from_index(idx_uv)
return idx_uv
[docs]
def filter_uv_from_index(self, idx_uv):
"""Keep only the uv modes selected by ``idx_uv`` (a boolean mask or
index array), dropping the rest in place (data, weights, and uu/vv/ru)."""
DataCube.filter_uv_from_index(self, idx_uv)
self.uu = self.uu[idx_uv]
self.vv = self.vv[idx_uv]
self.ru = self.ru[idx_uv]
[docs]
def average_freqs(self, n_freqs):
"""Weight-average consecutive groups of ``n_freqs`` channels together.
Returns:
CartDataCube: the down-sampled cube (weights summed accordingly).
"""
if (not isinstance(n_freqs, (int, np.integer))
or isinstance(n_freqs, (bool, np.bool_))):
raise TypeError('n_freqs must be a positive integer')
if n_freqs <= 0:
raise ValueError('n_freqs must be a positive integer')
if len(self.freqs) % n_freqs != 0:
raise ValueError(
f'{len(self.freqs)} frequency channels cannot be divided '
f'into complete groups of {n_freqs}')
freqs = self.freqs
bins = np.arange(len(freqs) // n_freqs)
digi = np.repeat(bins, n_freqs)[:len(freqs)]
new_freqs = np.array([freqs[digi == k].mean() for k in bins])
w = self.weights.data if self.weights is not None else np.ones_like(self.data)
new_data = np.array([(self.data[digi == k, :] * w[digi == k, :]).sum(axis=0) /
w[digi == k, :].sum(axis=0) for k in bins])
new_weight_cube = None
if self.weights is not None:
new_w = np.array([w[digi == k, :].sum(axis=0) for k in bins])
new_weight_cube = CartWeightCube(new_w, self.weights.uu, self.weights.vv, new_freqs,
self.weights.meta)
avg_cube = self.new_with_data(new_data, weights=new_weight_cube, freqs=new_freqs)
avg_cube.meta = self.meta.copy()
avg_cube.meta.average_freqs(n_freqs)
return avg_cube
[docs]
def average_same_uv(self):
"""Weight-average visibilities sharing the same (u, v) into one mode each.
Returns:
CartDataCube: the cube with unique uv coordinates.
"""
def binsum(x, y):
return np.array([np.bincount(x, k) for k in y])
x = np.round(self.uu, decimals=2) + 1e-6 * np.round(self.vv, decimals=2)
_x_u, idx, idx_r = np.unique(x, return_index=True, return_inverse=True)
w = binsum(idx_r, self.weights.get())
w_psf = binsum(idx_r, self.weights.data.real)
dw = binsum(idx_r, self.weights.get() * self.data.real) + 1j * \
binsum(idx_r, self.weights.get() * self.data.imag)
d = dw / w
uv_scale = w.mean(axis=0) / w_psf.mean(axis=0)
uu = self.uu[idx]
vv = self.vv[idx]
weights = CartWeightCube(w_psf, uu, vv, self.freqs, self.meta, uv_scale=uv_scale[None, :])
cube = CartDataCube(d, uu, vv, self.freqs, self.meta, weights=weights)
return cube
[docs]
def plot_uv(self, fmhz='med', action_fct=None, uv_lines=None, ax=None,
apply_uv_scale=False, title=None, **kargs):
"""Scatter-plot the visibilities on the uv plane.
Args:
fmhz: channel to show, in MHz, or one of ``'med'`` / ``'first'`` /
``'last'`` (ignored if ``action_fct`` is given).
action_fct: a reducer ``f(data, axis=0)`` applied over frequency
instead of picking one channel (e.g. ``np.mean``).
uv_lines (list): baseline lengths at which to draw guide circles.
ax: matplotlib Axes (new figure if None).
apply_uv_scale (bool): plot the uv-scaled (weighted) data.
**kargs: forwarded to the matplotlib scatter call.
"""
if uv_lines is None:
uv_lines = [50, 100, 150, 200, 250]
if ax is None:
_fig, ax = plt.subplots()
d = self.data
if apply_uv_scale:
d = self.get()
if action_fct is None:
if fmhz == 'med':
fmhz = (self.freqs[0] + (self.freqs[-1] - self.freqs[0]) / 2.) * 1e-6
elif fmhz == 'first':
fmhz = self.freqs[0] * 1e-6
elif fmhz == 'last':
fmhz = self.freqs[-1] * 1e-6
i = np.argmin(abs(self.freqs - fmhz * 1e6))
d = d[i]
else:
d = action_fct(d, axis=0)
cbs = psutil.ColorbarSetting(psutil.ColorbarOutterPosition())
im_mappable = ax.scatter(self.uu, self.vv, c=d.real, **kargs)
cbs.add_colorbar(im_mappable, ax)
ax.set_xlabel('U (lambda)')
ax.set_ylabel('V (lambda)')
for uv in uv_lines:
ax.add_artist(plt.Circle([0, 0], uv, ls='--', fc=None, ec=psutil.lblack, fill=False))
if title is not None:
ax.set_title(title)
[docs]
class CartDataCubeMeter(CartDataCube):
"""Cartesian cube whose UV coordinates are stored in metres."""
[docs]
def get_cube(self, mfreq):
"""Convert to a wavelength-coordinate :class:`CartDataCube` at reference
frequency ``mfreq`` (Hz), dividing the metre coordinates by that wavelength.
Returns:
CartDataCube: the wavelength-coordinate cube.
"""
lamb = const.c.value / mfreq
weights = None
if self.weights is not None:
weights = self.weights.get_cube(mfreq)
return CartDataCube(self.data, self.uu / lamb, self.vv / lamb,
self.freqs, self.meta, weights=weights)
[docs]
def get_baseline(self, mfreq, baseline):
"""The single mode at baseline length ``baseline`` (metres), as a
wavelength-coordinate cube at reference frequency ``mfreq`` (Hz).
Returns:
CartDataCube: a one-mode cube (``None`` if no such baseline).
"""
i = np.nonzero(np.round(baseline, 2) == np.round(self.ru, 2))[0]
if len(i) == 0:
print(f'No baseline with length {baseline} m')
return None
cube = self.get_cube(mfreq)
cube.uu = cube.uu[i[0]:i[0] + 1]
cube.vv = cube.vv[i[0]:i[0] + 1]
cube.ru = cube.ru[i[0]:i[0] + 1]
cube.data = cube.data[:, i[0]:i[0] + 1]
if self.weights is not None:
cube.weights = self.weights.get_baseline(mfreq, baseline)
cube.weights.uv_scale = cube.weights.uv_scale[:, i[0]:i[0] + 1]
return cube
[docs]
def new_with_data(self, data, weights=None, freqs=None):
"""A new CartDataCubeMeter of the same geometry with new ``data`` (and
optionally ``weights`` / ``freqs``)."""
if freqs is None:
freqs = self.freqs
if weights is None:
weights = self.weights
assert data.shape[0] == len(freqs)
assert data.shape[1] == len(self.uu)
if weights is not None:
weights = weights.copy()
return CartDataCubeMeter(data, self.uu, self.vv, freqs, self.meta.copy(), weights=weights)
[docs]
@staticmethod
def load(filename):
"""Load a CartDataCubeMeter (and its weights) from an HDF5 file."""
cube = CartDataCube.load(filename)
weights = None
if cube.weights is not None:
weights = CartWeightsCubeMeter(cube.weights.data, cube.uu, cube.vv, cube.freqs,
cube.meta, cube.weights.uv_scale)
return CartDataCubeMeter(cube.data, cube.uu, cube.vv, cube.freqs, cube.meta,
weights=weights)
[docs]
class MultiNightsCube:
"""Collection of compatible cubes indexed by observing night."""
def __init__(self, cubes=None, nights=None, inhomogeneous=False):
if cubes is None:
cubes = []
nights = []
self.cubes = cubes
self.nights = nights
self.inhomogeneous = inhomogeneous
def __iter__(self):
yield from self.cubes
[docs]
def concat(self):
"""Stack all nights into one cube along the visibility axis (each mode's
night recorded in ``cube.origin``).
Returns:
CartDataCube: the concatenated cube.
"""
data = np.hstack([c.data for c in self.cubes])
weights_data = np.hstack([c.weights.data for c in self.cubes])
uu = np.tile(self.uu, len(self.nights))
vv = np.tile(self.vv, len(self.nights))
weights = CartWeightCube(weights_data, uu, vv, self.freqs, self.meta)
cube = CartDataCube(data, uu, vv, self.freqs, self.meta, weights=weights)
cube.origin = np.repeat(self.nights, len(self.uu))
return cube
@property
def data(self):
"""All nights' data stacked along a third axis, ``(n_freqs, n_vis, n_nights)``."""
return np.dstack([c.data for c in self.cubes])
@property
def uu(self):
"""Shared u coordinates (from the first night)."""
return self.cubes[0].uu
@property
def vv(self):
"""Shared v coordinates (from the first night)."""
return self.cubes[0].vv
@property
def ru(self):
"""Shared baseline lengths (from the first night)."""
return self.cubes[0].ru
@property
def freqs(self):
"""Shared frequencies (from the first night)."""
return self.cubes[0].freqs
@property
def meta(self):
"""Shared image metadata (from the first night)."""
return self.cubes[0].meta
[docs]
def get_slice_from_idx(self, idx_freqs):
"""Restrict every night to the channels ``idx_freqs``, in place.
Returns:
MultiNightsCube: self, for chaining.
"""
for i in np.arange(len(self.cubes)):
self.cubes[i] = self.cubes[i].get_slice_from_idx(idx_freqs)
return self
[docs]
def get_slice(self, freq_start, freq_end):
"""Return all nights restricted to the given frequency slice."""
idx_freqs = psutil.get_freq_slice(self.freqs, freq_start, freq_end)
return self.get_slice_from_idx(idx_freqs)
[docs]
def add(self, cube, night):
"""Add a cube, intersecting coordinates when required."""
if len(self.cubes) > 0 and not self.inhomogeneous:
self.cubes[0], cube = get_common_cube(self.cubes[0], cube)
self.cubes.append(cube)
self.nights.append(night)
[docs]
def done(self):
"""Restrict all stored cubes to the coordinates of the first night."""
for i in np.arange(len(self.cubes) - 1) + 1:
_, self.cubes[i] = get_common_cube(self.cubes[0], self.cubes[i])
[docs]
class MultiDataInfo:
"""Observation start, end, and duration values loaded from a text table."""
def __init__(self, filename):
self.d = {}
for n, s, e, d in np.loadtxt(filename, str):
self.d[n] = (float(s), float(e), float(d))
[docs]
def start(self, night):
"""Start time recorded for ``night``."""
return self.d[night][0]
[docs]
def end(self, night):
"""End time recorded for ``night``."""
return self.d[night][1]
[docs]
def duration(self, night):
"""Observation duration recorded for ``night``."""
return self.d[night][2]
[docs]
class NoiseStdCube(CartDataCube):
"""Full complex thermal-noise standard deviation on a visibility grid.
``data`` is real and non-negative and represents
``sqrt(E[abs(noise) ** 2])``. For circular complex Gaussian noise, each
of the real and imaginary components therefore has variance
``data ** 2 / 2``.
"""
def __init__(self, data, uu, vv, freqs, meta, weights=None):
data = np.asarray(data)
if np.iscomplexobj(data):
raise ValueError('NoiseStdCube data must be a real full complex standard deviation')
if np.any(np.isnan(data)):
raise ValueError('NoiseStdCube data must not contain NaN')
if np.any(data < 0):
raise ValueError('NoiseStdCube data must be non-negative')
CartDataCube.__init__(self, data, uu, vv, freqs, meta, weights=weights)
def __add_sub__(self, other, sub=False):
if not isinstance(other, NoiseStdCube):
raise TypeError('NoiseStdCube object can not be combined with another type.')
if not np.allclose(self.freqs, other.freqs):
raise TypeError('The two NoiseStdCube objects must have the same frequency channels')
if not (np.allclose(self.uu, other.uu) and np.allclose(self.vv, other.vv)):
raise TypeError('The two NoiseStdCube objects must have the same uv coordinates')
return self.new_with_data(np.sqrt(self.variance + other.variance))
@property
def variance(self):
"""Full complex noise variance ``E[abs(noise) ** 2]``."""
return self.data ** 2
def __mul__(self, other):
if not psutil.is_number(other) or not np.isreal(other):
raise TypeError('NoiseStdCube can only be scaled by a real number')
return self.new_with_data(abs(other) * self.data)
def __rmul__(self, other):
return self.__mul__(other)
[docs]
def make_diff_cube(self):
"""Return an unchanged copy; channel differencing is not defined here."""
return self.copy()
[docs]
def make_diff_cube_interp(self):
"""Return an unchanged copy; channel differencing is not defined here."""
return self.copy()
[docs]
def generate_noise_cube(self, hermitian=True, rng=None):
"""Draw a complex Gaussian noise realization at this per-mode std.
Args:
hermitian (bool): enforce ``V(-u,-v) = conj(V(u,v))`` so the image
transform is real.
Returns:
CartDataCube: one noise realization.
"""
if np.iscomplexobj(self.data):
std_real = self.data.real
std_imag = self.data.imag
else:
std_real = std_imag = self.data.real / np.sqrt(2)
rng = np.random.default_rng(rng)
data = (rng.normal(scale=std_real) + 1j * rng.normal(scale=std_imag))
if hermitian:
idx1, idx2 = self.get_hermitian_index()
data[:, idx2] = np.conj(data[:, idx1])
weights = None if self.weights is None else self.weights.copy()
return CartDataCube(
data, self.uu, self.vv, self.freqs, self.meta.copy(),
weights=weights)
[docs]
def new_with_data(self, data, weights=None, freqs=None):
"""A new NoiseStdCube of this geometry with new ``data`` (a real,
non-negative standard deviation).
Returns:
NoiseStdCube: the new cube.
"""
if freqs is None:
freqs = self.freqs
if weights is None:
weights = self.weights
assert data.shape[0] == len(freqs)
assert data.shape[1] == len(self.uu)
if weights is not None:
weights = weights.copy()
return NoiseStdCube(data, self.uu, self.vv, freqs, self.meta.copy(), weights=weights)
[docs]
@staticmethod
def load(filename):
"""Load a noise-standard-deviation cube from HDF5 (converting legacy data).
Returns:
NoiseStdCube: the loaded cube.
"""
d = CartDataCube.load(filename)
data = d.data
if np.iscomplexobj(data):
warnings.warn(
'Loading a legacy complex NoiseStdCube; converting its '
'real/imaginary component standard deviations to a real full '
'complex standard deviation.',
UserWarning, stacklevel=2,
)
data = np.hypot(data.real, data.imag)
return NoiseStdCube(data, d.uu, d.vv, d.freqs, d.meta, weights=d.weights)
[docs]
class CartWeightCube(CartDataCube):
"""Visibility sampling weights with an optional UV-dependent scale."""
def __init__(self, weight_cube, uu, vv, freqs, meta, uv_scale=None, freqs_n_nights=None):
CartDataCube.__init__(self, weight_cube, uu, vv, freqs, meta)
if uv_scale is None:
self.unscale()
else:
self.uv_scale = uv_scale
if freqs_n_nights is None:
self.freqs_n_nights = np.ones(len(self.freqs))
else:
self.freqs_n_nights = freqs_n_nights
def __add_sub__(self, other, sub=False):
a = self.copy_with_applied_uv_scale()
b = other.copy_with_applied_uv_scale()
return CartDataCube.__add_sub__(a, b, sub=sub)
[docs]
def get(self, with_uv_scale=True):
"""The effective (absolute) weights.
Args:
with_uv_scale (bool): fold in ``uv_scale`` (the per-mode noise
scaling); False returns the raw PSF weights.
Returns:
ndarray: the weights, shape ``(n_freqs, n_vis)``.
"""
if with_uv_scale:
return abs(self.data * self.uv_scale)
return abs(self.data)
[docs]
def copy_with_applied_uv_scale(self):
"""Return a copy with ``uv_scale`` absorbed into stored weights."""
new = self.new_with_data(self.data * self.uv_scale)
new.unscale()
return new
[docs]
def filter_uv_from_index(self, idx_uv):
"""Keep only the uv modes ``idx_uv`` (as
:meth:`CartDataCube.filter_uv_from_index`), also slicing ``uv_scale``."""
CartDataCube.filter_uv_from_index(self, idx_uv)
self.uv_scale = self.uv_scale[:, idx_uv]
[docs]
def get_slice_from_idx(self, idx):
"""Select the channels ``idx``, carrying the per-channel night counts.
Returns:
CartWeightCube: the sliced weights.
"""
cube = CartDataCube.get_slice_from_idx(self, idx)
cube.freqs_n_nights = self.freqs_n_nights.copy()[idx]
return cube
[docs]
def filter_outliers(self, idx_outliers):
"""Remove channels marked by ``idx_outliers`` in place."""
CartDataCube.filter_outliers(self, idx_outliers)
self.freqs_n_nights = self.freqs_n_nights[~idx_outliers]
[docs]
def make_full_cube(self, umin, umax):
"""Expand onto the full ``[umin, umax]`` uv grid (as
:meth:`CartDataCube.make_full_cube`), filling absent ``uv_scale`` with 1.
Returns:
CartWeightCube: the gridded weights.
"""
cube, _idx1, idx1_uv = CartDataCube.make_full_cube(self, umin, umax, output_idx=True)
uv_scale = np.ones((1, cube.data.shape[1]))
uv_scale[:, idx1_uv] = self.uv_scale
return CartWeightCube(cube.data, cube.uu, cube.vv, cube.freqs, cube.meta, uv_scale=uv_scale)
[docs]
def reduce_fov(self, new_theta_fov, low_memory=False, umin=None, umax=None):
"""Re-grid to a smaller field of view (as :meth:`CartDataCube.reduce_fov`).
Returns:
CartWeightCube: the re-gridded weights.
"""
new_cube = CartDataCube.reduce_fov(self, new_theta_fov, low_memory=low_memory, umin=umin,
umax=umax)
return CartWeightCube(new_cube.data, new_cube.uu, new_cube.vv, new_cube.freqs,
new_cube.meta)
[docs]
def apply_window_function(self, win_fct, umin=None, umax=None, mc_n_samples=2000,
add_to_meta=True):
"""Apply a spatial window function (as :meth:`CartDataCube.apply_window_function`).
Returns:
CartWeightCube: the windowed weights.
"""
new_cube = CartDataCube.apply_window_function(self, win_fct, umin=umin, umax=umax,
add_to_meta=add_to_meta)
return CartWeightCube(new_cube.data, new_cube.uu, new_cube.vv, new_cube.freqs,
new_cube.meta)
[docs]
@staticmethod
def load_from_fits_psf(
files, umin, umax, int_time=None, total_time=None, theta_fov=None, low_memory=False,
output_psf_cube=False, window_function=None, data_dtype=np.complex128):
"""Build a weight cube from frequency-dependent PSF FITS files.
The Fourier transform of the PSF is the uv sampling, which sets the
per-mode weights. Used by :meth:`CartDataCube.load_from_fits_image_and_psf`.
Args:
files: PSF FITS files, in frequency order.
umin, umax: baseline range kept, in wavelengths.
int_time, total_time: integration and total observing time (s),
required if not already in the FITS metadata.
theta_fov: output field of view, in radians.
low_memory (bool): process one channel at a time.
output_psf_cube (bool): also return the (unweighted) PSF cube.
window_function: optional image-plane :class:`WindowFunction`.
Returns:
CartWeightCube: the weights (or ``(psf_cube, weights)`` when
``output_psf_cube``).
"""
cart_cube = CartDataCube.load_from_fits_image(
files, umin, umax, theta_fov, convert_jy2k=False, window_function=window_function,
data_dtype=data_dtype)
if 'PEINTTIM' not in cart_cube.meta:
assert int_time is not None
cart_cube.meta.set('PEINTTIM', int_time)
if 'PETOTTIM' not in cart_cube.meta:
assert total_time is not None
cart_cube.meta.set('PETOTTIM', total_time)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
header = pf.getheader(files[0])
if 'WSCENVIS' in header:
w_key = 'WSCENVIS'
elif 'WSCNVIS' in header:
w_key = 'WSCNVIS'
elif 'WEIGHT' in header:
w_key = 'WEIGHT'
else:
print('Error: no normalization factor found in header.')
return None
cart_cube.meta.set('PEWKEY', w_key)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
n_vis = np.array([pf.getheader(file)[w_key] for file in files])
# The UV plane contains each measured visibility and its transpose.
weights = cart_cube.data * n_vis[:, None] * 2
weight_cube = CartWeightCube(weights, cart_cube.uu, cart_cube.vv, cart_cube.freqs,
cart_cube.meta)
if output_psf_cube:
return cart_cube, weight_cube
return weight_cube
[docs]
@staticmethod
def load(filename):
"""Load a weight cube from HDF5."""
d = CartDataCube.load(filename)
return CartWeightCube(d.data, d.uu, d.vv, d.freqs, d.meta)
[docs]
@staticmethod
def load_from_hd5(h5_group):
"""Deprecated alias for :meth:`load_from_h5`."""
warnings.warn('load_from_hd5 is deprecated; use load_from_h5.', DeprecationWarning,
stacklevel=2)
return CartWeightCube.load_from_h5(h5_group)
[docs]
@staticmethod
def load_from_h5(h5_group):
"""Load a weight cube from an open HDF5 group."""
d = CartDataCube.load_from_h5(h5_group)
uv_scale = None
if hasattr(h5_group, 'uv_scale'):
uv_scale = h5_group.uv_scale.read()
freqs_n_nights = None
if hasattr(h5_group, 'freqs_n_nights'):
freqs_n_nights = h5_group.freqs_n_nights.read()
return CartWeightCube(d.data, d.uu, d.vv, d.freqs, d.meta, uv_scale=uv_scale,
freqs_n_nights=freqs_n_nights)
[docs]
def save_to_hd5(self, h5_file, group):
"""Deprecated alias for :meth:`save_to_h5`."""
warnings.warn('save_to_hd5 is deprecated; use save_to_h5.', DeprecationWarning,
stacklevel=2)
self.save_to_h5(h5_file, group)
[docs]
def save_to_h5(self, h5_file, group):
"""Write weights and scaling metadata into an HDF5 group."""
CartDataCube.save_to_h5(self, h5_file, group)
h5_file.create_array(group, 'uv_scale', self.uv_scale, "UV scale")
h5_file.create_array(group, 'freqs_n_nights', self.freqs_n_nights, "Number of nights")
[docs]
@staticmethod
def from_noise_cube(noise_cube, delta_u):
"""Build inverse-variance weights from a noise cube.
Bins the noise by baseline length (steps of ``delta_u`` wavelengths),
measures its MAD per bin, and returns weights proportional to
``1 / noise_scale**2`` (normalised to unit mean).
Returns:
CartWeightCube: the derived weights.
"""
d_ru = np.arange(noise_cube.ru.min(), noise_cube.ru.max(), delta_u)
d_ru = np.concatenate([[0], d_ru[1:-1], [np.inf]])
noise_scale = np.ones_like(noise_cube.ru, dtype=float)
for ru_min, ru_max in psutil.pairwise(d_ru):
idx = (noise_cube.ru >= ru_min) & (noise_cube.ru <= ru_max)
noise_scale[idx] = psutil.mad(noise_cube.data[:, idx])
noise_scale = noise_scale / noise_scale.mean()
noise_scale = np.repeat(noise_scale[None, :], len(noise_cube.freqs), axis=0)
# TODO: scale the noise_cube better
meta = noise_cube.meta.copy()
meta.set('PEINTTIM', 10)
meta.set('PETOTTIM', 10)
return CartWeightCube(1 / noise_scale ** 2, noise_cube.uu, noise_cube.vv,
noise_cube.freqs, meta)
[docs]
def estimate_sefd(self, noise_cube, sefd_jansky=True, axis=None):
"""SEFD estimated from a noise realization and these weights.
Args:
noise_cube: a noise :class:`CartDataCube` on the same geometry.
sefd_jansky (bool): return Jy (else Kelvin).
axis: reduce over this axis (0 -> per uv mode, 1 -> per channel,
None -> a single overall value).
Returns:
The SEFD (MAD-averaged over ``axis``); scalar or array per ``axis``.
"""
df = self.meta.chan_width
int_time = self.meta.int_time
w = abs(self.get_slice(noise_cube.freqs[0], noise_cube.freqs[-1]).data)
w[w == 0] = np.nan
sefd = noise_cube.data * (2 * df * int_time * 1) ** 0.5 * w ** 0.5
# Kelvin noise is normalized to the image field of view.
if sefd_jansky:
lamb = const.c.value / noise_cube.freqs
fov = (self.meta.shape[0] * self.meta.res) ** 2
jy2k = ((1e-26 * lamb ** 2) / (2 * const.k_B.value)) / fov
sefd = sefd / jy2k[:, None]
# Correct for spatial coherence introduced by the image window.
sefd = sefd / self.meta.win_fct_power ** .5
return psutil.mad(sefd, axis=axis)
[docs]
def estimate_uv_sefd(self, noise_cube, sefd_jansky=True):
"""SEFD per uv sample (see :meth:`estimate_sefd`).
Returns:
CartDataCube: the per-mode SEFD (single averaged channel).
"""
sefd = self.estimate_sefd(noise_cube, sefd_jansky=sefd_jansky, axis=0)
return CartDataCube(sefd[None, :], self.uu, self.vv, np.array([self.freqs.mean()]),
self.meta.copy())
[docs]
def estimate_freqs_sefd(self, noise_cube, sefd_jansky=True):
"""SEFD per frequency channel (see :meth:`estimate_sefd`).
Returns:
ndarray: the SEFD per channel, shape ``(n_freqs,)``.
"""
return self.estimate_sefd(noise_cube, sefd_jansky=sefd_jansky, axis=1)
[docs]
def scale_with_noise_cube(self, noise_cube, sefd_poly_fit_deg=0, sefd_filter_n_bins=0,
expected_sefd=None, scale_freqs=False):
"""Set ``uv_scale`` (in place) so the weights reflect the SEFD measured
from ``noise_cube`` -- down-weighting noisier uv modes.
Args:
noise_cube: a noise :class:`CartDataCube` on the same geometry.
sefd_poly_fit_deg (int): if > 0, smooth the radial scale with a
log-log polynomial of this degree.
sefd_filter_n_bins (int): if > 0, smooth the scale over this many
baseline-length bins instead.
expected_sefd: reference SEFD the scale normalises to (default: the
median measured SEFD).
scale_freqs (bool): also apply a per-channel SEFD scaling.
"""
sefd_uv = self.estimate_uv_sefd(noise_cube).data
if expected_sefd is None:
expected_sefd = np.nanmedian(sefd_uv)
self.uv_scale = (expected_sefd / sefd_uv) ** 2
if sefd_poly_fit_deg > 0:
mask = np.isfinite(self.uv_scale[0])
uv_fct = np.poly1d(np.polyfit(np.log(self.ru[mask]), np.log(self.uv_scale[0, mask]),
sefd_poly_fit_deg))
self.uv_scale = np.exp(uv_fct(np.log(self.ru))[None, :])
if sefd_filter_n_bins > 0:
m_sefd, _, _ = binned_statistic(self.ru, self.uv_scale, bins=sefd_filter_n_bins)
m_ru, _, _ = binned_statistic(self.ru, self.ru, bins=sefd_filter_n_bins)
self.uv_scale = scipy.interpolate.interp1d(
m_ru, m_sefd, bounds_error=False, kind='slinear',
fill_value='extrapolate')(self.ru)
if scale_freqs:
sefd_freqs = self.estimate_freqs_sefd(noise_cube)
self.uv_scale = self.uv_scale * (np.nanmedian(sefd_freqs) / sefd_freqs)[:, None] ** 2
self.uv_scale[~np.isfinite(self.uv_scale)] = 1
[docs]
def random_scale(self, max_ratio=2, hermitian=True, rng=None):
"""Randomly perturb ``uv_scale`` in place (log-uniform up to ``max_ratio``).
Args:
max_ratio: largest multiplicative perturbation.
hermitian (bool): keep conjugate modes consistent.
"""
rng = np.random.default_rng(rng)
shape = self.uv_scale.shape
self.uv_scale = 10 ** (np.log10(max_ratio) / 0.5 * (rng.random(shape) - 0.5))
if hermitian:
idx1, idx2 = self.get_hermitian_index()
self.uv_scale[:, idx1] = self.uv_scale[:, idx2]
[docs]
def unscale(self):
"""Reset ``uv_scale`` to one."""
self.uv_scale = np.ones((1, len(self.uu)))
[docs]
def simulate_noise(self, sefd, time, hermitian=True, weights_uncertainity_ratio=None,
sefd_jansky=True, fake_apply_win_fct=False, rng=None):
"""Draw a thermal-noise visibility cube for a given SEFD and integration.
Noise per mode follows the radiometer equation, scaled by these weights.
Args:
sefd: system-equivalent flux density (Jy if ``sefd_jansky``, else K).
time: total integration time, in seconds.
hermitian (bool): make the realization Hermitian (real image).
weights_uncertainity_ratio: if set, also perturb the weights randomly
by up to this ratio (see :meth:`random_scale`).
sefd_jansky (bool): interpret ``sefd`` in Jy (else Kelvin).
fake_apply_win_fct (bool): scale as if a spatial taper were applied.
Returns:
CartDataCube: the simulated noise cube.
"""
df = self.meta.chan_width
int_time = self.meta.int_time
total_time = self.meta.total_time
f = time / total_time
noise_rms = np.atleast_1d(sefd / (2 * df * int_time * f) ** 0.5)[:, None]
# Split full complex variance equally between real and imaginary parts.
noise_rms = np.sqrt(0.5) * noise_rms
if fake_apply_win_fct:
# Scale noise_rms as if a spatial tapering had been applied.
noise_rms = noise_rms * self.meta.win_fct_power ** .5
if sefd_jansky:
lamb = const.c.value / self.freqs
fov = (self.meta.shape[0] * self.meta.res) ** 2
jy2k = ((1e-26 * lamb ** 2) / (2 * const.k_B.value)) / fov
noise_rms = jy2k[:, None] * noise_rms
rng = np.random.default_rng(rng)
w = self.get() ** 0.5
s = self.data.shape
if weights_uncertainity_ratio is not None:
w_scale = 10 ** (np.log10(weights_uncertainity_ratio) / 0.5
* (rng.random(w.shape) - 0.5))
if hermitian:
idx1, idx2 = self.get_hermitian_index()
w_scale[idx1] = w_scale[idx2]
w = w * w_scale
if hermitian:
idx1, idx2 = self.get_hermitian_index()
hs = (self.data.shape[0], len(idx1))
noise_data = np.zeros_like(self.data, dtype=np.complex128)
noise_data[:, idx1] = noise_rms / w[:, idx1] * \
(rng.standard_normal(hs) + 1j * rng.standard_normal(hs))
noise_data[:, idx2] = np.conj(noise_data[:, idx1])
else:
noise_data = noise_rms / w * (rng.standard_normal(s) + 1j * rng.standard_normal(s))
weights = self.new_with_data(self.data * f)
noise_cube = CartDataCube(noise_data, self.uu, self.vv, self.freqs.copy(),
self.meta.copy(), weights=weights)
if not fake_apply_win_fct:
noise_cube.meta.remove('PEWINFCT')
noise_cube.weights.meta.remove('PEWINFCT')
return noise_cube
[docs]
def get_noise_std_cube(self, sefd, time, fake_apply_win_fct=False):
"""The *expected* per-mode thermal-noise standard deviation (not a
realization -- see :meth:`simulate_noise` for that).
Args:
sefd: system-equivalent flux density, in Jy.
time: total integration time, in seconds.
fake_apply_win_fct (bool): scale as if a spatial taper were applied.
Returns:
NoiseStdCube: the per-mode noise standard deviation.
"""
df = self.meta.chan_width
int_time = self.meta.int_time
total_time = self.meta.total_time
f = time / total_time
noise_rms = np.atleast_1d(sefd / (2 * df * int_time * f) ** 0.5)
lamb = const.c.value / self.freqs
fov = (self.meta.shape[0] * self.meta.res) ** 2
jy2k = ((1e-26 * lamb ** 2) / (2 * const.k_B.value)) / fov
noise_rms = (jy2k * noise_rms)[:, None]
w = self.get() ** 0.5
noise_rms = noise_rms / w
if fake_apply_win_fct:
# Scale noise_rms as if a spatial tapering had been applied.
noise_rms = noise_rms * self.meta.win_fct_power ** .5
weights = self.new_with_data(self.data * f)
noise_cube = NoiseStdCube(noise_rms, self.uu, self.vv, self.freqs.copy(),
self.meta.copy(), weights=weights)
if not fake_apply_win_fct:
noise_cube.meta.remove('PEWINFCT')
noise_cube.weights.meta.remove('PEWINFCT')
return noise_cube
[docs]
def new_with_data(self, data, weights=None, freqs=None):
"""A new CartWeightCube of this geometry with new weight ``data`` (and,
optionally, ``freqs``); ``uv_scale`` is carried over.
Returns:
CartWeightCube: the new weight cube.
"""
if freqs is None:
freqs = self.freqs
assert data.shape[0] == len(freqs)
assert data.shape[1] == len(self.uu)
assert data.shape[1] == self.uv_scale.shape[1]
return CartWeightCube(data, self.uu, self.vv, freqs, self.meta.copy(),
uv_scale=self.uv_scale.copy())
[docs]
@staticmethod
def join_cubes(cubes):
"""Concatenate weight cubes over frequency (uv_scale folded into the data).
Returns:
CartWeightCube: the joined weights.
"""
j_cube = CartDataCube.join_cubes([c.copy_with_applied_uv_scale() for c in cubes])
return CartWeightCube(j_cube.data, j_cube.uu, j_cube.vv, j_cube.freqs, j_cube.meta.copy(),
uv_scale=None)
[docs]
def copy(self):
"""Return an independent deep copy of this weight cube."""
return CartWeightCube(self.data.copy(), self.uu.copy(), self.vv.copy(), self.freqs.copy(),
self.meta.copy(), uv_scale=self.uv_scale.copy(),
freqs_n_nights=self.freqs_n_nights.copy())
[docs]
class CartWeightsCubeMeter(CartWeightCube, CartDataCubeMeter):
"""Weight-cube counterpart of :class:`CartDataCubeMeter`."""
[docs]
def get_cube(self, mfreq):
"""A :class:`CartWeightCube` with uv in wavelengths at frequency ``mfreq``
(Hz), converting from the metre coordinates held here."""
lamb = const.c.value / mfreq
return CartWeightCube(self.data, self.uu / lamb, self.vv / lamb,
self.freqs, self.meta, uv_scale=self.uv_scale)
[docs]
def new_with_data(self, data, weights=None, freqs=None):
"""A new CartWeightsCubeMeter of the same geometry with new ``data``."""
if freqs is None:
freqs = self.freqs
assert data.shape[0] == len(freqs)
assert data.shape[1] == len(self.uu)
assert data.shape[1] == self.uv_scale.shape[1]
return CartWeightsCubeMeter(data, self.uu, self.vv, freqs, self.meta.copy(),
uv_scale=self.uv_scale.copy())
[docs]
def copy(self):
"""Return an independent deep copy of this weight cube."""
return CartWeightsCubeMeter(
self.data.copy(), self.uu.copy(), self.vv.copy(), self.freqs.copy(),
self.meta.copy(), uv_scale=self.uv_scale.copy(),
freqs_n_nights=self.freqs_n_nights.copy())