Example #1
0
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)
Example #2
0
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)
Example #3
0
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
Example #4
0
# ##############################################################################
# _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
Example #5
0
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
Example #6
0
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')
Example #7
0
# ##############################################################################
# _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
    ----------
Example #8
0
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