class BeamWeights(array.LabeledMatrix):
"""
Beamforming coefficients.
Examples
--------
.. testsetup::
from pypeline.phased_array.beamforming import BeamWeights
import numpy as np
import pandas as pd
.. doctest::
>>> N_antenna, N_beam = 5, 3
>>> data = 1j * np.arange(N_antenna * N_beam).reshape(N_antenna, N_beam)
>>> ant_idx = pd.MultiIndex.from_product([range(N_antenna), [0]],
... names=['STATION_ID', 'ANTENNA_ID'])
>>> beam_idx = pd.Index(range(3), name='BEAM_ID')
>>> W = BeamWeights(data, ant_idx, beam_idx)
>>> W.data
array([[0. +0.j, 0. +1.j, 0. +2.j],
[0. +3.j, 0. +4.j, 0. +5.j],
[0. +6.j, 0. +7.j, 0. +8.j],
[0. +9.j, 0.+10.j, 0.+11.j],
[0.+12.j, 0.+13.j, 0.+14.j]])
"""
@chk.check(
dict(data=chk.accept_any(chk.has_reals, chk.has_complex,
sparse.isspmatrix),
ant_idx=instrument.is_antenna_index,
beam_idx=is_beam_index))
def __init__(self, data, ant_idx, beam_idx):
"""
Parameters
----------
data : :py:class:`~numpy.ndarray`
(N_antenna, N_beam) beamforming weights.
ant_idx
(N_antenna,) index.
beam_idx
(N_beam,) index.
"""
N_antenna, N_beam = len(ant_idx), len(beam_idx)
if not chk.has_shape((N_antenna, N_beam))(data):
raise ValueError('Parameters[data, ant_idx, beam_idx] are not '
'consistent.')
super().__init__(data=data, row_idx=ant_idx, col_idx=beam_idx)
class GramMatrix(array.LabeledMatrix):
"""
Gram coefficients.
Examples
--------
.. testsetup::
import numpy as np
import pandas as pd
from pypeline.phased_array.util.gram import GramMatrix
.. doctest::
>>> N_beam = 5
>>> beam_idx = pd.Index(range(N_beam), name='BEAM_ID')
>>> G = GramMatrix(np.eye(N_beam), beam_idx)
>>> G.data
array([[1., 0., 0., 0., 0.],
[0., 1., 0., 0., 0.],
[0., 0., 1., 0., 0.],
[0., 0., 0., 1., 0.],
[0., 0., 0., 0., 1.]])
"""
@chk.check(
dict(data=chk.accept_any(chk.has_reals, chk.has_complex),
beam_idx=beamforming.is_beam_index))
def __init__(self, data, beam_idx):
"""
Parameters
----------
data : :py:class:`~numpy.ndarray`
(N_beam, N_beam) Gram coefficients.
beam_idx
(N_beam,) index.
"""
data = np.array(data, copy=False)
N_beam = len(beam_idx)
if not chk.has_shape((N_beam, N_beam))(data):
raise ValueError('Parameters[data, beam_idx] are not consistent.')
if not np.allclose(data, data.conj().T):
raise ValueError('Parameter[data] must be hermitian symmetric.')
super().__init__(data, beam_idx, beam_idx)
class Wishart(Distribution):
"""
`Wishart <https://en.wikipedia.org/wiki/Wishart_distribution>`_ distribution.
"""
@chk.check(dict(V=chk.accept_any(chk.has_reals, chk.has_complex),
n=chk.is_integer))
def __init__(self, V, n):
"""
Parameters
----------
V : :py:class:`~numpy.ndarray`
(p, p) positive-semidefinite Hermitian scale matrix.
n : int
degrees of freedom.
"""
super().__init__()
V = np.array(V)
p = len(V)
if not (chk.has_shape([p, p])(V) and np.allclose(V, V.conj().T)):
raise ValueError('Parameter[V] must be hermitian symmetric.')
if not (n > p):
raise ValueError(f'Parameter[n] must be greater than {p}.')
self._V = V
self._p = p
self._n = n
Vq = linalg.sqrtm(V)
_, R = linalg.qr(Vq)
self._L = R.conj().T
@property
def mean(self):
"""
Mean of the distribution.
"""
if self._mean is None:
self._mean = self._n * self._V
return self._mean
@chk.check('x', chk.accept_any(chk.has_reals, chk.has_complex))
def pdf(self, x):
"""
Density of the distribution at sample points.
Parameters
----------
x : :py:class:`~numpy.ndarray`
(N, p, p) values at which to determine the pdf.
Returns
-------
pdf : :py:class:`~numpy.ndarray`
(N,) densities.
"""
x = np.array(x, copy=False)
if x.ndim == 2:
x = x[np.newaxis]
elif x.ndim == 3:
pass
else:
raise ValueError('Parameter[x] must have shape (N, p, p).')
N = len(x)
if not (chk.has_shape([N, self._p, self._p])(x) and
np.allclose(x, x.conj().transpose(0, 2, 1))):
raise ValueError('Parameter[x] must be hermitian symmetric.')
if np.linalg.matrix_rank(self._V) < self._p:
raise linalg.LinAlgError('Wishart density is not defined when '
'scale matrix V is singular.')
# Determinants: real-valued since (V,X) are Hermitian.
Vs, Vl = np.linalg.slogdet(self._V)
dV = np.real(Vs * np.exp(Vl))
Xs, Xl = np.linalg.slogdet(x)
dX = np.real(Xs * np.exp(Xl))
# Trace term
A = np.linalg.solve(self._V, x)
trA = np.trace(A, axis1=1, axis2=2).real
num = (np.float_power(dX, (self._n - self._p - 1) / 2) *
np.exp(-trA / 2))
den = (np.float_power(2, self._n * self._p / 2) *
np.float_power(dV, self._n / 2) *
np.exp(special.multigammaln(self._n / 2, self._p)))
pdf = num / den
return pdf
@chk.check('N_sample', chk.is_integer)
def __call__(self, N_sample=1):
"""
Generate random samples.
Parameters
----------
N_sample : int
Number of samples to generate.
Returns
-------
x : :py:class:`~numpy.ndarray`
(N_sample, p, p) samples.
Notes
-----
The Wishart estimate is obtained using the `Bartlett Decomposition`_.
.. _Bartlett Decomposition: https://en.wikipedia.org/wiki/Wishart_distribution#Bartlett_decomposition
"""
if N_sample < 1:
raise ValueError('Parameter[N_sample] must be positive.')
A = np.zeros((N_sample, self._p, self._p))
diag_idx = np.diag_indices(self._p)
df = (self._n * np.ones((N_sample, 1)) - np.arange(self._p))
A[:, diag_idx[0], diag_idx[1]] = np.sqrt(stats.chi2.rvs(df=df))
tril_idx = np.tril_indices(self._p, k=-1)
size = (N_sample, self._p * (self._p - 1) // 2)
A[:, tril_idx[0], tril_idx[1]] = stats.norm.rvs(size=size)
W = self._L @ A
X = W @ W.conj().transpose(0, 2, 1)
return X
# ############################################################################## # _fourier.py # =========== # Author : Sepand KASHANI [[email protected]] # ############################################################################## import cmath import numpy as np import scipy.fftpack as fftpack import pypeline.util.argcheck as chk import pypeline.util.array as array @chk.check(dict(x=chk.accept_any(chk.has_reals, chk.has_complex), T=chk.is_real, T_c=chk.is_real, N_FS=chk.is_odd, axis=chk.is_integer)) def ffs(x, T, T_c, N_FS, axis=-1): r""" Fourier Series coefficients from signal samples. Parameters ---------- x : :py:class:`~numpy.ndarray` (..., N_s, ...) function values at sampling points specified by :py:func:`~pypeline.util.math.fourier.ffs_sample`. T : float Function period. T_c : float
class EqualAngleInterpolator(core.Block):
r"""
Interpolate order-limited zonal function from Equal-Angle samples.
Computes :math:`f(r) = \sum_{q, l} \alpha_{q} f(r_{q, l}) K_{N}(\langle r, r_{q, l} \rangle)`, where :math:`r_{q, l} \in \mathbb{S}^{2}` are points from an Equal-Angle sampling scheme, :math:`K_{N}(\cdot)` is the spherical Dirichlet kernel of order :math:`N`, and the :math:`\alpha_{q}` are scaling factors tailored to an Equal-Angle sampling scheme.
Examples
--------
Let :math:`\gamma_{N}(r): \mathbb{S}^{2} \to \mathbb{R}` be the order-:math:`N` approximation of :math:`\gamma(r) = \delta(r - r_{0})`:
.. math::
\gamma_{N}(r) = \frac{N + 1}{4 \pi} \frac{P_{N + 1}(\langle r, r_{0} \rangle) - P_{N}(\langle r, r_{0} \rangle)}{\langle r, r_{0} \rangle -1}.
As :math:`\gamma_{N}` is order-limited, it can be exactly reconstructed from it's samples on an order-:math:`N` Equal-Angle grid:
.. testsetup::
import numpy as np
from pypeline.util.math.sphere import ea_sample, EqualAngleInterpolator, pol2cart
from pypeline.util.math.func import SphericalDirichlet
def gammaN(r, r0, N):
similarity = np.tensordot(r0, r, axes=1)
d_func = SphericalDirichlet(N)
return d_func(similarity) * (N + 1) / (4 * np.pi)
.. doctest::
# \gammaN Parameters
>>> N = 3
>>> r0 = np.array([0, 0, 1])
# Interpolate \gammaN from it's samples
>>> colat, lon = ea_sample(N)
>>> r = np.stack(pol2cart(1, colat, lon), axis=0)
>>> g_samples = gammaN(r, r0, N)
>>> q, l = np.meshgrid(np.arange(colat.size),
... np.arange(lon.size),
... indexing='ij')
>>> ea_interp = EqualAngleInterpolator(q.reshape(-1),
... l.reshape(-1),
... g_samples.reshape(-1),
... N)
# Compare with exact solution at off-sample positions
>>> colat, lon = ea_sample(2 * N) # denser grid
>>> g_interp = ea_interp(colat, lon)
>>> g_exact = gammaN(pol2cart(1, colat, lon), r0, N)
>>> np.allclose(g_interp, g_exact)
True
"""
@chk.check(
dict(q=chk.has_integers,
l=chk.has_integers,
f=chk.accept_any(chk.has_reals, chk.has_complex),
N=chk.is_integer,
approximate_kernel=chk.is_boolean))
def __init__(self, q, l, f, N, approximate_kernel=False):
r"""
Parameters
----------
q : :py:class:`~numpy.ndarray`
(N_s,) polar indices of an order-`N` Equal-Angle grid.
l : :py:class:`~numpy.ndarray`
(N_s,) azimuthal indices of an order-`N` Equal-Angle grid.
f : :py:class:`~numpy.ndarray`
(N_s,) samples of the zonal function at data-points. (float or complex)
:math:`L`-dimensional zonal functions are also supported by supplying an (N_s, L) array instead.
N : int
Order of the reconstructed zonal function.
approximate_kernel : bool
If :py:obj:`True`, pass the `approx` option to :py:class:`~pypeline.util.math.func.SphericalDirichlet`.
Notes
-----
If :math:`f(r)` only takes non-negligeable values when :math:`r \in \mathcal{S} \subset \mathbb{S}^{2}`, then the runtime of :py:meth:`~pypeline.util.math.sphere.EqualAngleInterpolator.__call__` can be significantly reduced by only supplying the triplets (`q`, `l`, `f`) that belong to :math:`\mathcal{S}`.
"""
super().__init__()
colat_sph, _ = ea_sample(N)
_2N2 = colat_sph.size
q, l = np.array(q), np.array(l)
if not ((q.shape == l.shape) and chk.has_shape((q.size, ))(q)):
raise ValueError("Parameter[q, l] must be 1D and of equal length.")
if not all(np.all(0 <= _) and np.all(_ < _2N2) for _ in [q, l]):
raise ValueError(f"Parameter[q, l] must contain entries in "
f"{{0, ..., 2N + 1}}.")
self._N = N
self._q = q
self._l = l
N_s = q.size
f = np.array(f, copy=False)
if (f.ndim == 1) and (len(f) == N_s):
self._L = 1
elif (f.ndim == 2) and (len(f) == N_s):
self._L = f.shape[1]
else:
raise ValueError(
"Parameter[f] must have shape (N_s,) or (N_s, L).")
f = f.reshape(N_s, self._L)
_2m1 = np.reshape(2 * np.r_[:N + 1] + 1, (1, N + 1))
alpha = (
np.sin(colat_sph) / _2N2 *
np.sum(np.sin(_2m1 * colat_sph) / _2m1, axis=1, keepdims=True))
self._weight = f * alpha[q]
self._kernel_func = func.SphericalDirichlet(N, approximate_kernel)
@chk.check(
dict(colat=chk.accept_any(chk.is_real, chk.has_reals),
lon=chk.accept_any(chk.is_real, chk.has_reals)))
def __call__(self, colat, lon):
"""
Interpolate function samples at order `N`.
Parameters
----------
colat : float or :py:class:`~numpy.ndarray`
Polar/Zenith angle [rad].
lon : float or :py:class:`~numpy.ndarray`
Longitude angle [rad].
Returns
-------
:py:class:`~numpy.ndarray`
(L, ...) function values at specified coordinates.
"""
r = np.stack(sph.pol2cart(1, colat, lon), axis=0)
sh_kern = (1, ) + r.shape[1:]
sh_weight = (self._L, ) + (1, ) * len(r.shape[1:])
colat_sph, lon_sph = ea_sample(self._N)
f_interp = np.zeros((self._L, ) + r.shape[1:],
dtype=self._weight.dtype)
with tqdm.tqdm(total=len(self._weight)) as pbar:
for w, t, p in zip(self._weight, colat_sph[self._q, 0],
lon_sph[0, self._l]):
similarity = np.tensordot(np.stack(sph.pol2cart(1, t, p),
axis=-1),
r,
axes=1)
kernel = self._kernel_func(similarity)
f_interp += kernel.reshape(sh_kern) * w.reshape(sh_weight)
pbar.update()
return f_interp
class SphericalImage:
"""
Wrapper around :py:class:`SphericalImageContainer_float32` and
:py:class:`SphericalImageContainer_float64` to enable advanced functionality.
Main features:
* import images from FITS format;
* export images to FITS format;
* advanced 2D plotting based on `Matplotlib <https://matplotlib.org/>`_;
* view exported images with `DS9 <http://ds9.si.edu/site/Home.html>`_.
Examples
--------
.. doctest::
import numpy as np
import pypeline.phased_array.util.grid as grid
import pypeline.phased_array.util.io.image as image
import pypeline.util.math.sphere as sph
import pypeline.util.math.stat as stat
# grid settings =======================
direction = np.array(sph.eq2cart(1,
lat=np.radians(30),
lon=np.radians(20)))
FoV = np.radians(60)
N_height, N_width = 256, 384
px_grid = grid.uniform_grid(direction, FoV, size=[N_height, N_width])
# data settings =======================
beta0, a0 = 0.7, [1, 1, 1]
beta1, a1 = 0.9, [0, 0, 1]
kent0 = stat.Kent(k=stat.Kent.min_scale(FoV, beta0) * 2,
beta=beta0,
g1=direction,
a=a0)
kent1 = stat.Kent(k=stat.Kent.min_scale(FoV, beta1) * 2,
beta=beta1,
g1=direction,
a=a1)
data0 = (kent0
.pdf(px_grid.reshape(3, N_height * N_width).T)
.reshape(N_height, N_width))
data1 = (kent1
.pdf(px_grid.reshape(3, N_height * N_width).T)
.reshape(N_height, N_width))
data = np.stack([data0, data1], axis=0)
# Image creation ======================
I_container = image.SphericalImageContainer_float64(data, px_grid)
I = image.SphericalImage(I_container)
Data IO:
.. doctest::
I.to_fits('test.fits') # save to FITS
I2 = image.from_fits('test.fits') # load from FITS
Interactive plotting:
.. doctest::
I.draw() # AEQD projection by default, all layers.
.. image:: _img/sphericalimage_aeqd_example.png
.. doctest::
I.draw(index=0, projection='GNOM') # Only show first data slice.
.. image:: _img/sphericalimage_gnom_example.png
.. doctest::
I.draw(index=1, projection='LCC', data_kwargs=dict(cmap='jet'))
.. image:: _img/sphericalimage_lcc_example.png
"""
@chk.check('container',
chk.is_instance(im_cpp.SphericalImageContainer_float32,
im_cpp.SphericalImageContainer_float64))
def __init__(self, container):
"""
Parameters
----------
container: :py:class:`~pypeline_phased_array_util_io_image_pybind11.SphericalImageContainer_float32` or :py:class:`~pypeline_phased_array_util_io_image_pybind11.SphericalImageContainer_float64`
Bare container holding spherical image data.
"""
self._container = container
@property
def image(self):
"""
Returns
-------
:py:class:`~numpy.ndarray`
(N_image, ...) data cube.
"""
return self._container.image
@property
def grid(self):
"""
Returns
-------
:py:class:`~numpy.ndarray`
(3, ...) Cartesian coordinates of the sky on which the data points are defined.
"""
return self._container.grid
@chk.check('file_name', chk.is_instance(str))
def to_fits(self, file_name):
"""
Save image to FITS file.
Parameters
----------
file_name : str
Name of file.
Notes
-----
* :py:class:`~pypeline.phased_array.util.io.image.SphericalImage` subclasses that write WCS information assume the grid is specified in ICRS.
If this is not the case, rotate the grid accordingly before calling :py:meth:`~pypeline.phased_array.util.io.image.SphericalImage.to_fits`.
* Data cubes are stored in a secondary IMAGE frame and can be viewed with DS9 using::
$ ds9 <FITS_file>.fits[IMAGE]
WCS information is only available in external FITS viewers if using :py:class:`~pypeline.phased_array.util.io.image.EqualAngleImage`.
"""
primary_hdu = self._PrimaryHDU()
image_hdu = self._ImageHDU()
hdulist = fits.HDUList([primary_hdu, image_hdu])
hdulist.writeto(file_name, overwrite=True)
def _PrimaryHDU(self):
"""
Generate primary Header Descriptor Unit (HDU) for FITS export.
Returns
-------
hdu : :py:class:`~astropy.io.fits.PrimaryHDU`
"""
metadata = dict(IMG_TYPE=(self.__class__.__name__,
'SphericalImage subclass'), )
# grid: stored as angles to reduce file size.
_, colat, lon = sph.cart2pol(*self.grid)
coordinates = np.stack([np.degrees(colat), np.degrees(lon)], axis=0)
hdu = fits.PrimaryHDU(data=coordinates)
for k, v in metadata.items():
hdu.header[k] = v
return hdu
def _ImageHDU(self):
"""
Generate image Header Descriptor Unit (HDU) for FITS export.
Returns
-------
hdu : :py:class:`~astropy.io.fits.ImageHDU`
"""
hdu = fits.ImageHDU(data=self.image, name='IMAGE')
return hdu
@classmethod
@chk.check(
dict(primary_hdu=chk.is_instance(fits.PrimaryHDU),
image_hdu=chk.is_instance(fits.ImageHDU)))
def _from_fits(cls, primary_hdu, image_hdu):
"""
Load image from Header Descriptor Units.
Parameters
----------
primary_hdu : :py:class:`~astropy.io.fits.PrimaryHDU`
image_hdu : :py:class:`~astropy.io.fits.ImageHDU`
Returns
-------
I : :py:class:`~pypeline.phased_array.util.io.image.SphericalImage`
"""
# PrimaryHDU: grid specification.
colat, lon = primary_hdu.data
x, y, z = sph.pol2cart(1, np.radians(colat), np.radians(lon))
grid = np.stack([x, y, z], axis=0)
# ImageHDU: extract data cube.
image = image_hdu.data
# Make sure (image, grid) have the same dtype to work with SphericalImageContainer_floatxx().
image = image.astype(grid.dtype)
if grid.dtype == np.dtype(np.float32):
I_container = im_cpp.SphericalImageContainer_float32(image, grid)
else: # float64 mode
I_container = im_cpp.SphericalImageContainer_float64(image, grid)
I = cls(I_container)
return I
@property
def shape(self):
"""
Returns
-------
tuple
Shape of data cube.
"""
return self.image.shape
@chk.check(
dict(index=chk.accept_any(chk.is_integer, chk.has_integers,
chk.is_instance(slice)),
projection=chk.is_instance(str),
catalog=chk.allow_None(chk.is_instance(sky.SkyEmission)),
show_gridlines=chk.is_boolean,
show_colorbar=chk.is_boolean,
ax=chk.allow_None(chk.is_instance(axes.Axes)),
data_kwargs=chk.allow_None(chk.is_instance(dict)),
grid_kwargs=chk.allow_None(chk.is_instance(dict)),
catalog_kwargs=chk.allow_None(chk.is_instance(dict))))
def draw(self,
index=slice(None),
projection='AEQD',
catalog=None,
show_gridlines=True,
show_colorbar=True,
ax=None,
data_kwargs=None,
grid_kwargs=None,
catalog_kwargs=None):
"""
Plot spherical image using a 2D projection.
Parameters
----------
index : int, array-like(int), slice
Slices of the data-cube to show.
If multiple layers are provided, they are summed together.
projection : str
Plot projection.
Must be one of (case-insensitive):
* AEQD: `Azimuthal Equi-Distant <https://en.wikipedia.org/wiki/Azimuthal_equidistant_projection>`_; (default)
* LAEA: `Lambert Equal-Area <https://en.wikipedia.org/wiki/Lambert_azimuthal_equal-area_projection>`_;
* LCC: `Lambert Conformal Conic <https://en.wikipedia.org/wiki/Lambert_conformal_conic_projection>`_;
* ROBIN: `Robinson <https://en.wikipedia.org/wiki/Robinson_projection>`_;
* GNOM: `Gnomonic <https://en.wikipedia.org/wiki/Gnomonic_projection>`_;
* HEALPIX: `Hierarchical Equal-Area Pixelisation <https://en.wikipedia.org/wiki/HEALPix>`_.
Notes
-----
* (AEQD, LAEA, LCC, GNOM) are recommended for mapping portions of the sphere.
* LCC breaks down when mapping polar regions.
* GNOM breaks down when mapping large FoVs.
* (ROBIN, HEALPIX) are recommended for mapping the entire sphere.
catalog : :py:class:`~pypeline.phased_array.util.data_gen.sky.SkyEmission`
Source catalog to overlay on top of images. (Default: no overlay)
show_gridlines : bool
Show RA/DEC gridlines. (Default: True)
It is possible for the gridlines to be plotted on the wrong range
when the grid crosses the 180W/E meridian.
show_colorbar : bool
Show colorbar. (Default: True)
ax : :py:class:`~matplotlib.axes.Axes`
Axes to draw on.
If :py:obj:`None`, a new axes is used.
data_kwargs : dict
Keyword arguments related to data-cube visualization.
Accepted keys are:
* :py:meth:`~matplotlib.axes.Axes.contourf` options.
* :py:meth:`~matplotlib.axes.Axes.tricontourf` options.
grid_kwargs : dict
Keyword arguments related to grid visualization.
Accepted keys are:
* N_parallel : int
Number declination lines to show in viewable region. (Default: 3)
* N_meridian : int
Number of right-ascension lines to show in viewable region. (Default: 3)
* polar_plot : bool
Correct RA/DEC gridlines when mapping polar regions. (Default: False)
When mapping polar regions, meridian lines may be doubled at 180W/E, making it seem like a meridian line is missing.
Setting `polar_plot` to :py:obj:`True` redistributes the meridians differently to correct the issue.
This option only makes sense when mapping polar regions, and will produce incorrect gridlines otherwise.
* ticks : bool
Add RA/DEC labels next to gridlines. (Default: False)
TODO: change to True once implemented
catalog_kwargs : dict
Keyword arguments related to catalog visualization.
Accepted keys are:
* :py:meth:`~matplotlib.axes.Axes.scatter` options.
Returns
-------
ax : :py:class:`~matplotlib.axes.Axes`
"""
if ax is None:
fig, ax = plt.subplots()
proj = self._draw_projection(projection)
scm = self._draw_data(index, data_kwargs, proj, ax)
cbar = self._draw_colorbar(show_colorbar, scm, ax) # noqa: F841
self._draw_gridlines(show_gridlines, grid_kwargs, proj, ax)
self._draw_catalog(catalog, catalog_kwargs, proj, ax)
self._draw_beautify(proj, ax)
return ax
@chk.check('projection', chk.is_instance(str))
def _draw_projection(self, projection):
"""
Setup :py:class:`pyproj.Proj` object to do (lon,lat) <-> (x,y) transforms.
Parameters
----------
projection : str
`projection` parameter given to :py:meth:`draw`.
Returns
-------
proj : :py:class:`pyproj.Proj`
"""
# Most projections can be provided a point in space around which distortions are minimized.
# We choose this point to approximately map to the center of the grid when appropriate.
# (approximate since it is not always a spherical cap.)
if self._container.is_gridded: # (3, N_height, N_width) grid
grid_dir = np.mean(self.grid, axis=(1, 2))
else: # (3, N_points) grid
grid_dir = np.mean(self.grid, axis=1)
_, grid_lat, grid_lon = sph.cart2eq(*grid_dir)
grid_lat = coord.Angle(grid_lat * u.rad).to_value(u.deg)
grid_lon = coord.Angle(grid_lon * u.rad).wrap_at(180 * u.deg).to_value(
u.deg)
p_name = projection.lower()
if p_name == 'lcc':
# Lambert Conformal Conic
proj = pyproj.Proj(proj='lcc', lon_0=grid_lon, lat_0=grid_lat, R=1)
elif p_name == 'aeqd':
# Azimuthal Equi-Distant
proj = pyproj.Proj(proj='aeqd',
lon_0=grid_lon,
lat_0=grid_lat,
R=1)
elif p_name == 'laea':
# Lambert Equal-Area
proj = pyproj.Proj(proj='laea',
lon_0=grid_lon,
lat_0=grid_lat,
R=1)
elif p_name == 'robin':
# Robinson
proj = pyproj.Proj(proj='robin', lon_0=grid_lon, R=1)
elif p_name == 'gnom':
# Gnomonic
proj = pyproj.Proj(proj='gnom',
lon_0=grid_lon,
lat_0=grid_lat,
R=1)
elif p_name == 'healpix':
# Hierarchical Equal-Area Pixelisation
proj = pyproj.Proj(proj='healpix',
lon_0=grid_lon,
lat_0=grid_lat,
R=1)
else:
raise ValueError('Parameter[projection] is not a valid projection '
'specifier.')
return proj
@chk.check(
dict(index=chk.accept_any(chk.is_integer, chk.has_integers,
chk.is_instance(slice)),
data_kwargs=chk.allow_None(chk.is_instance(dict)),
projection=chk.is_instance(pyproj.Proj),
ax=chk.is_instance(axes.Axes)))
def _draw_data(self, index, data_kwargs, projection, ax):
"""
Contour plot of data.
Parameters
----------
index : int, array-like(int), slice
`index` parameter given to :py:meth:`draw`.
data_kwargs : dict
`data_kwargs` parameter given to :py:meth:`draw`.
projection : :py:class:`~pyproj.Proj`
PyProj projection object.
ax : :py:class:`~matplotlib.axes.Axes`
Axes to plot on.
Returns
-------
scm : :py:class:`~matplotlib.cm.ScalarMappable`
"""
if data_kwargs is None:
data_kwargs = dict()
N_image = self.shape[0]
if chk.is_integer(index):
index = np.array([index], dtype=int)
elif chk.has_integers(index):
index = np.array(index, dtype=int)
else: # slice()
index = np.arange(N_image, dtype=int)[index]
if index.size == 0:
raise ValueError('No data-cube slice chosen.')
if not np.all((0 <= index) & (index < N_image)):
raise ValueError('Parameter[index] is out of bounds.')
data = np.sum(self.image[index], axis=0)
# Transform (lon,lat) to (x,y).
# Some projections have unmappable regions or exhibit singularities at certain points.
# These regions are colored white in contour plots by replacing their incorrect value (1e30) with NaN.
_, grid_lat, grid_lon = sph.cart2eq(*self.grid)
grid_lat = coord.Angle(grid_lat * u.rad).to_value(u.deg)
grid_lon = coord.Angle(grid_lon * u.rad).wrap_at(180 * u.deg).to_value(
u.deg)
grid_x, grid_y = projection(grid_lon, grid_lat, errcheck=False)
grid_x[np.isclose(grid_x, 1e30)] = np.nan
grid_y[np.isclose(grid_y, 1e30)] = np.nan
# Colormap choice
if 'cmap' in data_kwargs:
obj = data_kwargs.pop('cmap')
if chk.is_instance(str)(obj):
cmap = cm.get_cmap(obj)
else:
cmap = obj
else:
cmap = plot.cmap('matthieu-custom-sky', N=38)
if self._container.is_gridded:
scm = ax.contourf(grid_x,
grid_y,
data,
cmap.N,
cmap=cmap,
**data_kwargs)
else:
triangulation = tri.Triangulation(grid_x, grid_y)
scm = ax.tricontourf(triangulation,
data,
cmap.N,
cmap=cmap,
**data_kwargs)
# Show coordinates in status bar
def sexagesimal_coords(x, y):
lon, lat = projection(x, y, errcheck=False, inverse=True)
lon = (coord.Angle(lon * u.deg).wrap_at(180 * u.deg).to_string(
unit=u.hourangle, sep='hms'))
lat = (coord.Angle(lat * u.deg).to_string(unit=u.degree,
sep='dms'))
msg = f'RA: {lon}, DEC: {lat}'
return msg
ax.format_coord = sexagesimal_coords
return scm
@chk.check(
dict(show_colorbar=chk.is_boolean,
scm=chk.is_instance(cm.ScalarMappable),
ax=chk.is_instance(axes.Axes)))
def _draw_colorbar(self, show_colorbar, scm, ax):
"""
Attach colorbar.
Parameters
----------
show_colorbar : bool
`show_colorbar` parameter given to :py:meth:`draw`.
scm : :py:class:`~matplotlib.cm.ScalarMappable`
Intensity scale.
ax : :py:class:`~matplotlib.axes.Axes`
Axes to plot on.
Returns
-------
cbar : :py:class:`~matplotlib.colorbar.Colorbar`
"""
if show_colorbar:
cbar = plot.colorbar(scm, ax)
else:
cbar = None
return cbar
@chk.check(
dict(show_gridlines=chk.is_boolean,
grid_kwargs=chk.allow_None(chk.is_instance(dict)),
projection=chk.is_instance(pyproj.Proj),
ax=chk.is_instance(axes.Axes)))
def _draw_gridlines(self, show_gridlines, grid_kwargs, projection, ax):
"""
Plot Right-Ascension / Declination lines.
Parameters
----------
show_gridlines : bool
`show_gridlines` parameter given to :py:meth:`draw`.
grid_kwargs : dict
`grid_kwargs` parameter given to :py:meth:`draw`.
projection : :py:class:`pyproj.Proj`
PyProj projection object.
ax : :py:class:`~matplotlib.axes.Axes`
Axes to plot on.
"""
if grid_kwargs is None:
grid_kwargs = dict()
if 'N_parallel' in grid_kwargs:
N_parallel = grid_kwargs.pop('N_parallel')
if not (chk.is_integer(N_parallel) and (N_parallel >= 3)):
raise ValueError('Value[N_parallel] must be at least 3.')
else:
N_parallel = 3
if 'N_meridian' in grid_kwargs:
N_meridian = grid_kwargs.pop('N_meridian')
if not (chk.is_integer(N_meridian) and (N_meridian >= 3)):
raise ValueError('Value[N_meridian] must be at least 3.')
else:
N_meridian = 3
if 'polar_plot' in grid_kwargs:
polar_plot = grid_kwargs.pop('polar_plot')
if not chk.is_boolean(polar_plot):
raise ValueError('Value[polar_plot] must be boolean.')
else:
polar_plot = False
if 'ticks' in grid_kwargs:
show_ticks = grid_kwargs.pop('ticks')
if not chk.is_boolean(show_ticks):
raise ValueError('Value[ticks] must be boolean.')
else:
# TODO: change to True once implemented.
show_ticks = False
plot_style = dict(alpha=0.5, color='k', linewidth=1, linestyle='solid')
plot_style.update(grid_kwargs)
_, grid_lat, grid_lon = sph.cart2eq(*self.grid)
grid_lat = coord.Angle(grid_lat * u.rad).to_value(u.deg)
grid_lon = coord.Angle(grid_lon * u.rad).wrap_at(180 * u.deg).to_value(
u.deg)
# RA curves
meridian = dict()
dec_span = np.linspace(grid_lat.min(), grid_lat.max(), 200)
if polar_plot:
ra = np.linspace(-180, 180, N_meridian, endpoint=False)
else:
ra = np.linspace(grid_lon.min(), grid_lon.max(), N_meridian)
for _ in ra:
ra_span = _ * np.ones_like(dec_span)
# Transform (lon,lat) to (x,y).
# Some projections have unmappable regions or exhibit singularities at certain points.
# These regions are colored white in contour plots by replacing their incorrect value (1e30) with NaN.
grid_x, grid_y = projection(ra_span, dec_span, errcheck=False)
grid_x[np.isclose(grid_x, 1e30)] = np.nan
grid_y[np.isclose(grid_y, 1e30)] = np.nan
if show_gridlines:
mer = ax.plot(grid_x, grid_y, **plot_style)[0]
meridian[_] = mer
# DEC curves
parallel = dict()
ra_span = np.linspace(grid_lon.min(), grid_lon.max(), 200)
if polar_plot:
dec = np.linspace(grid_lat.min(), grid_lat.max(), N_parallel + 1)
else:
dec = np.linspace(grid_lat.min(), grid_lat.max(), N_parallel)
for _ in dec:
dec_span = _ * np.ones_like(ra_span)
# Transform (lon,lat) to (x,y).
# Some projections have unmappable regions or exhibit singularities at certain points.
# These regions are colored white in contour plots by replacing their incorrect value (1e30) with NaN.
grid_x, grid_y = projection(ra_span, dec_span, errcheck=False)
grid_x[np.isclose(grid_x, 1e30)] = np.nan
grid_y[np.isclose(grid_y, 1e30)] = np.nan
if show_gridlines:
par = ax.plot(grid_x, grid_y, **plot_style)[0]
parallel[_] = par
# LAT/LON ticks
if show_gridlines and show_ticks:
raise NotImplementedError('Not yet implemented.')
@chk.check(
dict(catalog=chk.allow_None(chk.is_instance(sky.SkyEmission)),
projection=chk.is_instance(pyproj.Proj),
ax=chk.is_instance(axes.Axes)))
def _draw_catalog(self, catalog, catalog_kwargs, projection, ax):
"""
Overlay catalog on top of map.
Parameters
----------
catalog : :py:class:`~pypeline.phased_array.util.data_gen.sky.SkyEmission`
`catalog` parameter given to :py:meth:`draw`.
catalog_kwargs : dict
`catalog_kwargs` parameter given to :py:meth:`draw`.
projection : :py:class:`pyproj.Proj`
PyProj projection object.
ax : :py:class:`~matplotlib.axes.Axes`
Axes to plot on.
"""
if catalog is not None:
_, c_lat, c_lon = sph.cart2eq(*catalog.xyz.T)
c_lat = coord.Angle(c_lat * u.rad).to_value(u.deg)
c_lon = coord.Angle(c_lon * u.rad).wrap_at(180 * u.deg).to_value(
u.deg)
c_x, c_y = projection(c_lon, c_lat, errcheck=False)
c_x[np.isclose(c_x, 1e30)] = np.nan
c_y[np.isclose(c_y, 1e30)] = np.nan
if catalog_kwargs is None:
catalog_kwargs = dict()
plot_style = dict(s=400, facecolors='none', edgecolors='w')
plot_style.update(catalog_kwargs)
ax.scatter(c_x, c_y, **plot_style)
@chk.check(
dict(projection=chk.is_instance(pyproj.Proj),
ax=chk.is_instance(axes.Axes)))
def _draw_beautify(self, projection, ax):
"""
Format plot.
Parameters
----------
projection : :py:class:`pyproj.Proj`
PyProj projection object.
ax : :py:class:`~matplotlib.axes.Axes`
Axes to draw on.
"""
ax.axis('off')
ax.axis('equal')
# ############################################################################## # _linalg.py # ========== # Author : Sepand KASHANI [[email protected]] # ############################################################################## import numpy as np import scipy.linalg as linalg import pypeline.util.argcheck as chk @chk.check( dict(A=chk.accept_any(chk.has_reals, chk.has_complex), B=chk.allow_None(chk.accept_any(chk.has_reals, chk.has_complex)), tau=chk.is_real, N=chk.allow_None(chk.is_integer))) def eigh(A, B=None, tau=1, N=None): """ Solve a generalized eigenvalue problem. Finds :math:`(D, V)`, solution of the generalized eigenvalue problem .. math:: A V = B V D. This function is a wrapper around :py:func:`scipy.linalg.eigh` that adds energy truncation and extra output formats. Parameters ----------
class SphericalDirichlet(core.Block):
r"""
Parameterized spherical Dirichlet kernel.
Examples
--------
.. testsetup::
import numpy as np
from pypeline.util.math.func import SphericalDirichlet
.. doctest::
>>> N = 4
>>> f = SphericalDirichlet(N)
>>> sample_points = np.linspace(-1, 1, 25).reshape(5, 5) # any shape
>>> amplitudes = f(sample_points)
>>> np.around(amplitudes, 2)
array([[ 1. , 0.2 , -0.25, -0.44, -0.44],
[-0.32, -0.13, 0.07, 0.26, 0.4 ],
[ 0.47, 0.46, 0.38, 0.22, 0. ],
[-0.24, -0.48, -0.67, -0.76, -0.68],
[-0.37, 0.27, 1.3 , 2.84, 5. ]])
When only interested in kernel values close to 1, the approximation method provides significant speedups, at the cost of approximation error in values far from 1:
.. doctest::
N = 11
f_exact = SphericalDirichlet(N)
f_approx = SphericalDirichlet(N, approx=True)
x = np.linspace(-1, 1, 2000)
e_y = f_exact(x)
a_y = f_approx(x)
rel_err = np.abs((e_y - a_y) / e_y)
fig, ax = plt.subplots(nrows=2)
ax[0].plot(x, e_y, 'r')
ax[0].plot(x, a_y, 'b')
ax[0].legend(['exact', 'approx'])
ax[0].set_title('Dirichlet Kernel')
ax[1].plot(x, rel_err)
ax[1].set_title('Relative Error (Exact vs. Approx)')
fig.show()
.. image:: _img/sph_dirichlet_example.png
Notes
-----
The spherical Dirichlet function :math:`K_{N}(t): [-1, 1] \to \mathbb{R}` is defined as:
.. math:: K_{N}(t) = \frac{P_{N+1}(t) - P_{N}(t)}{t - 1},
where :math:`P_{N}(t)` is the `Legendre polynomial <https://en.wikipedia.org/wiki/Legendre_polynomials>`_ of order :math:`N`.
"""
@chk.check(dict(N=chk.is_integer, approx=chk.is_boolean))
def __init__(self, N, approx=False):
"""
Parameters
----------
N : int
Kernel order.
approx : bool
Approximate kernel using cubic-splines.
This method provides extremely reliable estimates of :math:`K_{N}(t)` in the vicinity of 1 where the function's main sidelobes are found.
Values outside the vicinity smoothly converge to 0.
Only works for `N` greater than 10.
"""
super().__init__()
if N < 0:
raise ValueError("Parameter[N] must be non-negative.")
self._N = N
if (approx is True) and (N <= 10):
raise ValueError('Cannot use approximation method if '
'Parameter[N] <= 10.')
self._approx = approx
if approx is True: # Fit cubic-spline interpolator.
N_samples = 10**3
# Find interval LHS after which samples will be evaluated exactly.
theta_max = np.pi
while True:
x = np.linspace(0, theta_max, N_samples)
cx = np.cos(x)
cy = self._exact_kernel(cx)
zero_cross = np.diff(np.sign(cy))
N_cross = np.abs(np.sign(zero_cross)).sum()
if N_cross > 10:
theta_max /= 2
else:
break
window = func.Tukey(T=2 - 2 * np.cos(2 * theta_max),
beta=1,
alpha=0.5)
x = np.r_[np.linspace(np.cos(theta_max * 2),
np.cos(theta_max),
N_samples,
endpoint=False),
np.linspace(np.cos(theta_max), 1, N_samples)]
y = self._exact_kernel(x) * window(x)
self.__cs_interp = interpolate.interp1d(x,
y,
kind='cubic',
bounds_error=False,
fill_value=0)
@chk.check('x', chk.accept_any(chk.is_real, chk.has_reals))
def __call__(self, x):
r"""
Sample the order-N spherical Dirichlet kernel.
Parameters
----------
x : float or :py:class:`~numpy.ndarray`
Values at which to compute :math:`K_{N}(x)`.
Returns
-------
K_N(x) : :py:class:`~numpy.ndarray`
"""
if chk.is_scalar(x):
x = np.array([x], dtype=float)
else:
x = np.array(x, copy=False, dtype=float)
if not np.all((-1 <= x) & (x <= 1)):
raise ValueError('Parameter[x] must lie in [-1, 1].')
if self._approx is True:
f = self._approx_kernel
else:
f = self._exact_kernel
amplitude = f(x)
return amplitude
# @chk.check('x', chk.accept_any(chk.is_real, chk.has_reals))
def _exact_kernel(self, x):
amplitude = (sp.eval_legendre(self._N + 1, x) -
sp.eval_legendre(self._N, x))
with warnings.catch_warnings():
# The kernel is so condensed near 1 at high N that np.isclose()
# does a terrible job at letting us manually treat values close to
# the upper limit.
# The best way to implement K_N(t) is to let the floating point
# division fail and then replace NaNs.
warnings.simplefilter(action='ignore', category=RuntimeWarning)
amplitude /= x - 1
amplitude[np.isnan(amplitude)] = self._N + 1
return amplitude
# @chk.check('x', chk.accept_any(chk.is_real, chk.has_reals))
def _approx_kernel(self, x):
amplitude = self.__cs_interp(x)
return amplitude