Source code for irfpy.util.polarframe

r"""Polar frame related module.

This module provides 3-D grid system together with the
data allocated on the grid.

- :class:`RegularGridContainer` can store values on a
  regularly gridded sphere coordinate system.
- :class:`RegularGridVolumeContainer` extends the
  :class:`RegularGridContainer` with volume support.

Similar functionality is found in :mod:`irfpy.util.gridsphere`
module, which is for 2-D mapping on the sphere.

In addition, frame conversion between RTP and XYZ is provided.

- :meth:`rtp2xyz` and :meth:`xyz2rtp`.

.. codeauthor:: Yoshifumi Futaana
"""

from irfpy.util import exception as ex
from irfpy.util import nptools

import numpy as np     # This will be replaced
import numpy as _np     # to _np.


[docs]class RegularGridContainer: r""" Container that preserves 3-D polar frame. Values at regularly gridded spherical coordinates (r, |theta|, |phi|). All should be "ascending order". """ def __init__(self, data, rgrid, tgrid, pgrid): """ Initialize the container. :param data: Data at the grid. It must be a numpy array with (nr, nt, np) or (nr, nt, np, data_dimension) array. ``data`` is just copied by reference (not deep-copied). :param rgrid: Grid in the r direction. Should be ascending. :param tgrid: Grid in the theta direction. Should be ascending. The angle is measured from the north pole. and thus the range is (0, 2|pi|). :param pgrid: Grid in the phi direction. Should be ascending. Range is user specific; so that you can consider (-|pi|, |pi|) or (0, 2 |pi|) or whatever else. It could be more than 1 rotation in principle, namely (0, 4 |pi|), but not well tested for such case. >>> rgrid = [0, 1, 2, 3, 4, 5] >>> tgrid = np.deg2rad([40, 60, 80, 100, 120]) >>> pgrid = np.deg2rad([-180, -90, 0, 90]) >>> data = np.arange(6 * 5 * 4).reshape([6, 5, 4]) >>> print(data.shape) (6, 5, 4) >>> reg_grid_val = RegularGridContainer(data, rgrid, tgrid, pgrid) """ if not nptools.isascending(rgrid): raise ex.IrfpyException('The given rgrid is not ascending order') if not nptools.isascending(tgrid): raise ex.IrfpyException('The given tgrid is not ascending order') if not nptools.isascending(pgrid): raise ex.IrfpyException('The given pgrid is not ascending order') rgrid = np.copy(rgrid) tgrid = np.copy(tgrid) pgrid = np.copy(pgrid) nr = rgrid.shape[0] nt = tgrid.shape[0] nph = pgrid.shape[0] if data.shape[0] != nr: raise ex.IrfpyException( 'Inconsistency of dimension. data.shape[0] != rgrid.shape[0]') if data.shape[1] != nt: raise ex.IrfpyException( 'Inconsistency of dimension. data.shape[1] != tgrid.shape[0]') if data.shape[2] != nph: raise ex.IrfpyException( 'Inconsistency of dimension. data.shape[1] != pgrid.shape[0]') self.rgrid = rgrid self.tgrid = tgrid self.pgrid = pgrid self.data = data @staticmethod def _getindex(ab, v): """ Find the index that v belongs to. The ``bisect_right`` is used. :param ab: Array defining the boundary :param v: Value (array-like) to examine to find the index corresponding to ``ab``. :return: Index array. If out-of-range, the value is masked. :rtype: np.ma.masked_array >>> ab = [0, 1, 2, 3, 4, 5] >>> v = 1.5 >>> print(RegularGridContainer._getindex(ab, v)) # 1.5 should be in the range ``[1, 2)`` 1 >>> v = 4 >>> print(RegularGridContainer._getindex(ab, v)) # 4 should be in the range ``[4, 5)`` 4 >>> v = -5 >>> print(RegularGridContainer._getindex(ab, v)) # Out of range -- >>> v = 5.5 >>> print(RegularGridContainer._getindex(ab, v)) # Out of range -- >>> v = 5 >>> print(RegularGridContainer._getindex(ab, v)) # Out of range -- In summary, >>> print(RegularGridContainer._getindex(ab, [-0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5.5])) # Array input allowed [-- 0 1 2 3 4 --] """ idx = _np.searchsorted(ab, v, side='right') - 1 return _np.ma.masked_outside(idx, 0, len(ab) - 2)
[docs] def get(self, rs, ts, ps, degrees=False): """ Return the value at the positions. Values are the data of the belonging grids :param rs: Array of r :param ts: Array of theta (radians) :param ps: Array of phi (radians) :keyword degrees: Use degrees for ts and ps. :return: The value at the position of (rs, ts, ps). All the shape rs, ts, ps should be the same. No broadcasting supported (currently). >>> rs = [0.5, 1.5, 2.5, 3.5, 4.5] >>> ts = np.deg2rad([30, 90, 150]) >>> ps = np.deg2rad([45, 135, 225, 315]) >>> value = np.arange(5 * 3 * 4).reshape([5, 3, 4]) >>> rgc = RegularGridContainer(value, rs, ts, ps) >>> print(rgc.get(0.5, 30, 45, degrees=True)) 0 >>> print(rgc.get(0.5, 30, 135, degrees=True)) 1 >>> print(rgc.get(0, 30, 45, degrees=True)) -- Arrays can be input >>> print(rgc.get([0, 0.5, 1.5, 6], [30, 30, 30, 30], [45, 45, 45, 45], degrees=True)) [-- 0 12 --] """ if degrees: ts = _np.deg2rad(ts) ps = _np.deg2rad(ps) irs = RegularGridContainer._getindex(self.rgrid, rs) its = RegularGridContainer._getindex(self.tgrid, ts) ips = RegularGridContainer._getindex(self.pgrid, ps) mask = (irs.mask | its.mask | ips.mask) return np.ma.masked_array(self.data[irs, its, ips], mask=mask)
def _shape(self): return (len(self.rgrid), len(self.tgrid), len(self.pgrid)) def _has_same_grid(self, other): return (self._shape() == other._shape())
[docs] @classmethod def zeros(cls, rgrid, tgrid, pgrid): """ Create a :class:`RegularGridContainer` with data of zero. :param rgrid: Array of r-grid. See :meth:`__init__` :param tgrid: Array of t-grid :param pgrid: Array of p-grid :return: Zero-data filled :class:`RegularGridContainer` Equivallent to ``RegularGridContainer(np.ma.zeros( [len(rgrid), len(tgrid), len(pgrid)]), rgrid, tgrid, pgrid)`` >>> rgrid = [1, 3, 5] >>> tgrid = np.deg2rad([30, 60, 90, 120, 150]) >>> pgrid = np.deg2rad([10, 100, 190, 280]) >>> zero_grid = RegularGridContainer.zeros(rgrid, tgrid, pgrid) >>> print(zero_grid.data.max()) 0.0 >>> print(zero_grid.data.min()) 0.0 """ if cls != RegularGridContainer: raise NotImplementedError( 'It is not supported to create zero-data grid object ' 'of {c} using {c}.zero function.'.format(c=cls.__name__)) data = np.ma.zeros([len(rgrid), len(tgrid), len(pgrid)]) return RegularGridContainer(data, rgrid, tgrid, pgrid)
def __add__(self, other): """ Operation + is defined. >>> rgrid = [0, 1, 2, 3, 4, 5] >>> tgrid = np.deg2rad([40, 60, 80, 100, 120]) >>> pgrid = np.deg2rad([-180, -90, 0, 90]) >>> data = np.arange(6 * 5 * 4).reshape([6, 5, 4]) # Maximum is 119. >>> gc1 = RegularGridContainer(data, rgrid, tgrid, pgrid) >>> gc2 = RegularGridContainer(data ** 2, rgrid, tgrid, pgrid) >>> gc3 = gc1 + gc2 >>> print(gc3.data.max()) # 119 + 119 ** 2 = 14280. 14280 """ if self.__class__ != RegularGridContainer: raise NotImplementedError( 'No support yet for operators in {c}'.format( c=self.__class__.__name__)) if not self._has_same_grid(other): raise ex.IrfpyException( 'Grid size is different. Cannot add ({}) and ({})'.format( str(self._shape()), str(other._shape()))) shape = self._shape() retobj = RegularGridContainer.zeros(self.rgrid, self.tgrid, self.pgrid) retobj.data = self.data + other.data return retobj def __sub__(self, other): """ Operation - is defined. >>> rgrid = [0, 1, 2, 3, 4, 5] >>> tgrid = np.deg2rad([40, 60, 80, 100, 120]) >>> pgrid = np.deg2rad([-180, -90, 0, 90]) >>> data = np.arange(6 * 5 * 4).reshape([6, 5, 4]) # Maximum is 119. >>> gc1 = RegularGridContainer(data, rgrid, tgrid, pgrid) >>> gc2 = RegularGridContainer(data ** 2, rgrid, tgrid, pgrid) >>> gc3 = gc2 - gc1 >>> print(gc3.data.max()) # 119 ** 2 - 119 = 14280. 14042 """ if self.__class__ != RegularGridContainer: raise NotImplementedError( 'No support yet for operators in {c}'.format( c=self.__class__.__name__)) if not self._has_same_grid(other): raise ex.IrfpyException( 'Grid size is different. ' 'Cannot add ({}) and ({})'.format( str(self._shape()), str(other._shape()))) shape = self._shape() retobj = RegularGridContainer.zeros(self.rgrid, self.tgrid, self.pgrid) retobj.data = self.data - other.data return retobj
[docs]class RegularGridVolumeContainer(RegularGridContainer): r""" Container of data in 3-D polar frame with volume information. The class supports, in addition to the data container of :class:`RegularGridContainer`, the volume of each bin. Each bin size can explicitly given by the parameters, or can be automatically calculated. .. note:: Operators and other functions have NOT yet supported. .. todo:: - Implement zeros function. - Implement __add__ and __sub__ functions. >>> rgrid = [10, 100, 1000, 10000] >>> tgrid = np.deg2rad([-30, -15, 0, 15, 30]) >>> pgrid = np.deg2rad([300, 350, 400, 450, 500, 550]) >>> RegularGridVolumeContainer.zeros(rgrid, tgrid, pgrid) # doctest: +IGNORE_EXCEPTION_DETAIL Traceback (most recent call last): File "<stdin>", line 1, in ? NotImplementedError: ... """ def __init__(self, data, rgrid, tgrid, pgrid, rgridsize='linear', tgridsize='linear', pgridsize='linear'): r""" :param data: Data, passed to :class:`RegularGridContainer` :param rgrid: Grid in r direction. Ascending. :param tgrid: Grid in |theta| direction. Ascending. Radians. :param pgrid: Grid in |phi| direction. Ascending. Radians. (If periodic condition will make the array non-ascending, the array should be sorted by using, for example, :func:`irfpy.util.nputils.sort_periodic`.) :param rgridsize: Size for each bin. Numpy array with the same dimension as `rgrid`. Automatic guess is available with 'linear' or 'log' options, where :func:`irfpy.util.nptools.guess_upper` and :func:`irfpy.util.nptools.guess_lower` is used. :param tgridsize: Identical to ``rgridsize``, but for theta direction. :param pgridsize: Identical to ``rgridsize``, but for phi direction. >>> rgrid = [10, 100, 1000, 10000] >>> tgrid = np.deg2rad([-30, -15, 0, 15, 30]) >>> pgrid = np.deg2rad([300, 350, 400, 450, 500, 550]) >>> data = np.arange(120).reshape([4, 5, 6]) >>> regGrid = RegularGridVolumeContainer(data, rgrid, tgrid, pgrid, rgridsize='log') >>> print(regGrid.rgridsize) # doctest: +NORMALIZE_WHITESPACE [2.84604989e+01 2.84604989e+02 2.84604989e+03 2.84604989e+04] >>> print(regGrid.tgridsize) # doctest: +NORMALIZE_WHITESPACE [0.26179939 0.26179939 0.26179939 0.26179939 0.26179939] >>> print(regGrid.pgridsize) # doctest: +NORMALIZE_WHITESPACE [0.87266463 0.87266463 0.87266463 0.87266463 0.87266463 0.87266463] """ RegularGridContainer.__init__(self, data, rgrid, tgrid, pgrid) if isinstance(rgridsize, str) and rgridsize == 'linear': self.rgridsize = nptools.guess_upper(rgrid) - nptools.guess_lower(rgrid) elif isinstance(rgridsize, str) and rgridsize == 'log': self.rgridsize = nptools.guess_upper(rgrid, log=True) - nptools.guess_lower(rgrid, log=True) else: if rgridsize.shape == self.rgrid.shape: self.rgridsize = rgridsize else: raise ex.IrfpyException('rgridsize should be "linear", "log" or array with the same shape as rgrid.') if isinstance(tgridsize, str) and tgridsize == 'linear': self.tgridsize = nptools.guess_upper(tgrid) - nptools.guess_lower(tgrid) elif isinstance(tgridsize, str) and tgridsize == 'log': self.tgridsize = nptools.guess_upper(tgrid, log=True) - nptools.guess_lower(tgrid, log=True) else: if tgridsize.shape == self.tgrid.shape: self.tgridsize = tgridsize else: raise ex.IrfpyException('tgridsize should be "linear", "log" or array with the same shape as tgrid.') if isinstance(pgridsize, str) and pgridsize == 'linear': self.pgridsize = nptools.guess_upper(pgrid) - nptools.guess_lower(pgrid) elif isinstance(pgridsize, str) and pgridsize == 'log': self.pgridsize = nptools.guess_upper(pgrid, log=True) - nptools.guess_lower(pgrid, log=True) else: if pgridsize.shape == self.pgrid.shape: self.pgridsize = pgridsize else: raise ex.IrfpyException('pgridsize should be "linear", "log" or array with the same shape as pgrid.')
[docs] def volume(self): r""" Return a matrix of volume for each grid. :return: The volume in (nr, nt, nph) shaped array. The volume of the each bin is .. math:: V = r^2 \sin\theta \Delta r\Delta\theta\Delta\phi >>> egrid_c, egrid_b, egrid_d = nptools.loggrids(1, 2, 80) >>> tgrid_c, tgrid_b, tgrid_d = nptools.lingrids(0, np.pi, 50) >>> pgrid_c, pgrid_b, pgrid_d = nptools.lingrids(-np.pi, np.pi, 90) >>> val = np.ones([80, 50, 90]) >>> vol = RegularGridVolumeContainer(val, egrid_c, tgrid_c, pgrid_c, rgridsize='log') >>> dvol = vol.volume() >>> print('{:.3e}'.format((vol.data * dvol).sum())) # This is theoretically the sphere of radius of r1 - that of r0, i.e. 4.1846e6 4.184e+06 """ r = self.rgrid[:, np.newaxis, np.newaxis] dr = self.rgridsize[:, np.newaxis, np.newaxis] th = self.tgrid[np.newaxis, :, np.newaxis] dth = self.tgridsize[np.newaxis, :, np.newaxis] dph = self.pgridsize[np.newaxis, np.newaxis, :] v = r ** 2 * dr * np.sin(th) * dth * dph return v
[docs] def integrate(self): r""" A simple integrator. :return: The integrated value. Integration is done as .. math:: \int\int\int y(r, \theta, \phi) r^2 \sin\theta dr d\theta d\phi >>> egrid_c, egrid_b, egrid_d = nptools.loggrids(1, 2, 80) >>> tgrid_c, tgrid_b, tgrid_d = nptools.lingrids(0, np.pi, 50) >>> pgrid_c, pgrid_b, pgrid_d = nptools.lingrids(-np.pi, np.pi, 90) >>> val = np.ones([80, 50, 90]) >>> vol = RegularGridVolumeContainer(val, egrid_c, tgrid_c, pgrid_c, rgridsize='log') >>> print('{:.3e}'.format(vol.integrate())) # This is theoretically the sphere of radius of r1 - that of r0, i.e. 4.1846e6 4.184e+06 """ d = self.data v = self.volume() return (d * v).sum()
[docs]def rtp2xyz(r, t, p, longitude_deg, latitude_deg): r""" RTP frame to xyz frame RTP (radius, |theta|, |phi|) frame is frequently used to represent a local vector measurement, such as velocity vector or magnetic field vector on a sphere (e.g. Earth). (XYZ is the global frame. This defines longitude and latitude.) Unit vectors for R, T, and P is defined as below: - *R*: The radial direction from the center. - |theta| : The south-ward component, tangential to the sphere surface. - |phi| : The east-ward component, tangential to the sphere surface. The conversion is .. math:: x = r \cos\lambda\cos\alpha + \theta\sin\lambda\cos\alpha - \phi\sin\alpha \\ y = r \cos\lambda\sin\alpha + \theta\sin\lambda\sin\alpha + \phi\cos\alpha \\ z = r \sin\alpha - \theta\cos\alpha Here |alpha| is the longitude, and |lambda| is the latitude. :param r: R, scalar or numpy array. :param t: T, scalar or numpy array. :param p: P, scalar or numpy array. Shapes of *r*, *t*, *p* should be the same. :param longitude_deg: Longitude in degrees. scalar or numpy array. Same shape as *r*, or scalar. :param latitude_deg: Latitude in degrees. scalar or numpy array. Same shape as *r*, or scalar. Same shape as *longitude_deg*. :return: [x, y, z] as numpy float64 or numpy array. For the spacecraft position at (1, 0, 0), then the rtp=(1, 0, 0) will be xyz=(1, 0, 0). rtp=(0, 1, 0) will be xyz=(0, 0, -1). rtp=(0, 0, 1) will be xyz=(0, 1, 0) >>> print(rtp2xyz(1, 0, 0, 0, 0)) (1.0, 0.0, 0.0) >>> print(rtp2xyz(0, 1, 0, 0, 0)) (0.0, 0.0, -1.0) >>> print(rtp2xyz(0, 0, 1, 0, 0)) (0.0, 1.0, 0.0) For the spacecraft position at (0, 1, 0)=(lon=90, lat=0), then the rtp=(1, 0, 0) will be xyz=(0, 1, 0). rtp=(0, 1, 0) will be (0, 0, -1). rtp=(0, 0, 1) will be xyz=(-1, 0, 0) >>> print(rtp2xyz(1, 0, 0, 90, 0)) # tiny residual. # doctest: +SKIP (0.0, 1.0, 0.0) >>> print(rtp2xyz(0, 1, 0, 90, 0)) (0.0, 0.0, -1.0) >>> print(rtp2xyz(0, 0, 1, 90, 0)) # tiny residual. # doctest: +SKIP (-1.0, 0.0, 0.0) For the spacecraft position at lon=0, lat=45, then the rtp=(0, 0, 1) is xyz=(0, 1, 0) >>> print(rtp2xyz(0, 0, 1, 0, 45)) (0.0, 1.0, 0.0) rtp=(1, 0, 0) is (0.707, 0, 0.707) and rtp=(0, 1, 0) is (0.707, 0, -0.707) >>> print(rtp2xyz(1, 0, 0, 0, 45)) # doctest: +NORMALIZE_WHITESPACE (0.7071067811865476, 0.0, 0.7071067811865475) >>> print(rtp2xyz(0, 1, 0, 0, 45)) # doctest: +NORMALIZE_WHITESPACE (0.7071067811865475, 0.0, -0.7071067811865476) Multidiemnsional version. >>> r = [1, 0] >>> t = [0, 1] >>> p = [0, 0] >>> lon = [0, 90] >>> lat = [0, 0] >>> x, y, z = rtp2xyz(r, t, p, lon, lat) >>> print(x) # doctest: +NORMALIZE_WHITESPACE [1. 0.] >>> print(y) # doctest: +NORMALIZE_WHITESPACE [0. 0.] >>> print(z) [ 0. -1.] """ lat = np.deg2rad(latitude_deg) lon = np.deg2rad(longitude_deg) coslat = np.cos(lat) sinlat = np.sin(lat) coslon = np.cos(lon) sinlon = np.sin(lon) [[m00, m01, m02], [m10, m11, m12], [m20, m21, m22]] = np.array([[coslat * coslon, coslat * sinlon, sinlat], [sinlat * coslon, sinlat * sinlon, -coslat], [-sinlon, coslon, 0]]).T x = m00 * r + m01 * t + m02 * p y = m10 * r + m11 * t + m12 * p z = m20 * r + m21 * t # m22 is always zero return (x, y, z)
[docs]def xyz2rtp(x, y, z, longitude_deg, latitude_deg): """ XYZ frame to RTP frame. See :func:`rtp2xyz` for details on RTP frame. :param x: X, scalar :param y: Y, scalar :param z: Z, scalar :param longitude_deg: Longitude of the "observer" :param latitude_deg: Latitude of the "observer" :return: (r, t, p) each as numpy.float or numpy.array >>> print(xyz2rtp(0, 0, -1, 90, 0)) (0.0, 1.0, 0.0) >>> print(xyz2rtp(0, 1, 0, 0, 45)) (0.0, 0.0, 1.0) >>> print(rtp2xyz(1, 0, 0, 0, 0)) (1.0, 0.0, 0.0) Multi-dimensional support. >>> x = [0, 0] >>> y = [0, 1] >>> z = [-1, 0] >>> lon = [90, 0] >>> lat = [0, 45] >>> r, t, p = xyz2rtp(x, y, z, lon, lat) >>> print(r) # doctest: +NORMALIZE_WHITESPACE [0. 0.] >>> print(t) # doctest: +NORMALIZE_WHITESPACE [1. 0.] >>> print(p) # doctest: +NORMALIZE_WHITESPACE [0. 1.] """ lat = np.deg2rad(latitude_deg) lon = np.deg2rad(longitude_deg) coslat = np.cos(lat) sinlat = np.sin(lat) coslon = np.cos(lon) sinlon = np.sin(lon) [[m00, m01, m02], [m10, m11, m12], [m20, m21, m22]] = [[coslat * coslon, coslat * sinlon, sinlat], [sinlat * coslon, sinlat * sinlon, -coslat], [-sinlon, coslon, 0]] r = m00 * x + m01 * y + m02 * z t = m10 * x + m11 * y + m12 * z p = m20 * x + m21 * y # m22 is zero return (r, t, p)