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