Source code for waveformtools.grids

""" Classes to hold grid information


Classes
-------
spherical_grid: grid info.
                Stores information on the 2d grid in spherical polar coordinates.
gl_grid : grid info
          Stores a Gauss-Legendre type grid on a spherical surface.
"""

import numpy as np

# from numba import jit, njit

# import numba as nb

# from numba.experimental import jitclass


[docs]class spherical_grid: """A class to store the coordinate grid on a sphere. Attributes ---------- ntheta: int The number of angular points in the :math:`\\theta` direction, including ghost zones. nphi: int The number of angular points in the :math:`\\phi` direction, including ghost zones. nghosts: int The number of ghost zones at the end of each direction. meshgrid: tuple of 2d array The 2d array containing the meshgrid of (:math:`\\theta, \\phi`) angular points. theta_1d: 1d array The 1d array of angular points along the :math:`\\theta` axis. phi_1d: 1d array The 1d array of angular points along the :math:`\\phi` axis. dtheta: float The angular step size in the :math:`\\theta` direction. dphi: float The angular step size inthe :math:`\\phi` direction. npix_act: int The total number of gridpoints on the sphere, excluding the ghost points. meshgrid: Get the 2d angular grid. Methods ------- theta_1d: Get the :math:`\\theta` axis. phi_1d: Get the :math:`\\phi` axis. """ def __init__(self, nphi=80, ntheta=41, nphimax=124, nthetamax=66, nghosts=2): # Number of gridpoints along phi direction including ghost points. self.nphi = nphi # Number of gridpoints along theta direction including ghost points. self.ntheta = ntheta # Total length of phi array used by ETK. self.nphimax = nphimax # Total length of theta array used by ETK. self.nthetamax = nthetamax # Number of ghost points in theta/phi direction. self.nghosts = nghosts @property def npix(self): """Return the total number of pixels, including the ghost zones present at one iteration.""" return (self.ntheta) * (self.nphi) @property def npix_act(self): """Return the actual number of pixels, excluding the ghost zones present at one iteration""" return (self.ntheta - 2 * self.nghosts) * (self.nphi - 2 * self.nghosts) @property def npix_max(self): """Return the (max) total number of pixels, including the ghost and buffer zones at one iteration.""" return (self.nthetamax) * (self.nphimax) @property def ntheta_act(self): """Return the actual number of valid pixels, excluding the ghost and buffer zones, along the theta axis at one iteration.""" return self.ntheta - 2 * self.nghosts @property def nphi_act(self): """Return the actual number of valid pixels, excluding the ghost and buffer zones, along the phi axis at one iteration.""" return self.nphi - 2 * self.nghosts @property def dtheta(self): """Return the coodinate spacing :math:`d\\theta.`""" return np.pi / (self.ntheta - 2 * self.nghosts) @property def dphi(self): """Return the coordinate spacing :math:`d\\phi`.""" return 2 * np.pi / (self.nphi - 2 * self.nghosts) @property def nbuffer(self): """Return the number of buffer zones (excluding ghosts)""" return self.nthetamax - self.ntheta
[docs] def theta_1d(self, theta_index=None): """Returns the coordinate value theta given the coordinate index. The coordinate index ranges from (0, ntheta). The actual indices without the ghost and extra zones is (nghosts, ntheta-nghosts). Parameters ----------- theta_index: int/ 1d array The theta coordinate index or axis. Returns ------- theta_1d: float The coordinate(s) :math:`\\theta` on the sphere. """ if not theta_index: theta_index = np.arange(self.nghosts, self.ntheta - self.nghosts) return (theta_index - self.nghosts + 0.5) * np.pi / (self.ntheta - 2 * self.nghosts)
[docs] def phi_1d(self, phi_index=None): """Returns the coordinate value theta given the coordinate index. The coordinate index lies in (0, nphi). The actual indices without the ghost and extra zones is (nghosts, nphi-nghosts). Parameters ----------- phi_1d: int / 1d array The phi coordinate index or axis. Returns ------- phi_1d: float or 1d array The coordinate(s) :math:`\\phi` on the sphere. """ if not phi_index: phi_index = np.arange(self.nghosts, self.nphi - self.nghosts) return (phi_index - self.nghosts) * 2 * np.pi / (self.nphi - 2 * self.nghosts)
@property def meshgrid(self): """The (:math:`\\theta, \\phi)`: coordinate meshes. Excludes the ghost zones. Returns ------- theta: 2d array The :math:`\\theta` coordinate matrix for vectorization. phi: 2d array The :math:`\\phi` coordinate matrix for vectorization. """ theta, phi = np.meshgrid(self.theta_1d(), self.phi_1d()) # theta = np.zeros((self.ntheta_act, self.nphi_act)) # phi = np.zeros((self.ntheta_act, self.nphi_act)) # for theta_index, theta_val in enumerate(self.theta_1d()): # for phi_index, phi_val in enumerate(self.phi_1d()): # theta[theta_index, phi_index] = theta_val # phi[theta_index, phi_index] = phi_val return np.transpose(theta), np.transpose(phi)
[docs]class gl_grid: """A class to store the coordinate grid on a sphere. Attributes ---------- ntheta: int The number of angular points in the :math:`\\theta` direction, including ghost zones. nphi: int The number of angular points in the :math:`\\phi` direction, including ghost zones. nghosts: int The number of ghost zones at the end of each direction. meshgrid: tuple of 2d array The 2d array containing the meshgrid of (:math:`\\theta, \\phi`) angular points. theta_1d: 1d array The 1d array of angular points along the :math:`\\theta` axis. phi_1d: 1d array The 1d array of angular points along the :math:`\\phi` axis. dtheta: float The angular step size in the :math:`\\theta` direction. dphi: float The angular step size inthe :math:`\\phi` direction. npix_act: int The total number of gridpoints on the sphere, excluding the ghost points. meshgrid: Get the 2d angular grid. Methods ------- theta_1d: Get the :math:`\\theta` axis. phi_1d: Get the :math:`\\phi` axis. Notes ----- The total number of points on the sphere is assumed to be :math:`(L+1)^2` :math:`N_\theta = L+1` :math:`N_\phi = 2(L+1)` This integrates out spherical harmonics of degree L exactly, given a regular function on the sphere. """ def __init__(self, nphi=None, ntheta=None, L=47, nghosts=2): # Number of gridpoints along phi direction including ghost points. self._nphi = nphi # Number of gridpoints along theta direction including ghost points. self._ntheta = ntheta # Total length of phi array used by ETK. # self._nphimax = nphimax # Total length of theta array used by ETK. # self._nthetamax = nthetamax # Number of ghost points in theta/phi direction. self._nghosts = nghosts self._L = L self._theta_1d = None self._phi_1d = None self._meshgrid = None if self._ntheta is None: if self._L is None: raise ValueError("Please specify L or angular points!") else: self._ntheta_act = L + 1 self._nphi_act = 2 * self._ntheta_act self._ntheta = self._ntheta_act + 2 * self._nghosts self._nphi = self._nphi_act + 2 * self._nghosts elif self._L is None: if self.nthetha is None: raise ValueError("Please specify L or angular points!") else: self.L = self.ntheta - 1 assert nphi % 2 == 0 from scipy.special import roots_legendre cpoints, self._weights, self._sum_of_weights = roots_legendre(L + 1, mu=True) # xpoints = (np.pi-np.arccos(cpoints)) xpoints = np.arccos(cpoints[::-1]) self._theta_1d = xpoints dphi = 2 * np.pi / self._nphi_act self._phi_1d = np.linspace(0, 2 * np.pi - dphi, self._nphi_act) theta_grid, phi_grid = np.meshgrid(self._theta_1d, self._phi_1d) self._meshgrid = np.transpose(theta_grid), np.transpose(phi_grid) dtheta_axis = np.diff(self._theta_1d) dtheta_axis = np.append(dtheta_axis, np.pi - self._theta_1d[-1]) dtheta_axis = np.insert(dtheta_axis, 0, self._theta_1d[0]) self._dtheta_1d = dtheta_axis self._dphi = dphi # self._phi_1d[1] @property def nphi(self): """Return the (max) total number of pixels, including the ghost and buffer zones at one iteration.""" return self._nphi @property def ntheta(self): """Return the actual number of valid pixels, excluding the ghost and buffer zones, along the theta axis at one iteration.""" return self._ntheta @property def nghosts(self): """Return the number of ghost zones at each end in each direction""" return self._nghosts @property def ntheta_act(self): """Return the actual number of valid pixels, excluding the ghost and buffer zones, along the phi axis at one iteration.""" return self.ntheta - 2 * self.nghosts @property def nphi_act(self): """Return the actual number of valid pixels, excluding the ghost and buffer zones, along the phi axis at one iteration.""" return self.nphi - 2 * self.nghosts @property def npix_act(self): """Return the actual number of pixels, excluding the ghost zones present at one iteration""" return (self.ntheta_act) * (self.nphi_act) @property def dtheta_1d(self): """Return the non-uniform angular stepping in :math:`\theta` direction""" return self._dtheta_1d @property def dphi(self): """Return the uniform angular stepping in :math:`\phi` direction""" return self._dphi @property def L(self): """Return the maximum mode number that this grid can integrate out.""" return self._L @property def npix(self): """Return the total number of pixels, including the ghost zones present at one iteration.""" return (self.ntheta) * (self.nphi) @property def weights(self): """Return the integration weights along theta.""" return self._weights @property def weights_grid(self): """Return the integration weights on the coordinate mesh grid""" return np.outer(self.weights, np.ones(self.nphi_act)) @property def shape(self): """Return the shape of the grid.""" return self.weights_grid.shape @property def theta_1d(self): """Returns the coordinate values theta given the coordinate index. The coordinate index ranges from (0, ntheta). The actual indices without the ghost and extra zones is (nghosts, ntheta-nghosts). Parameters ----------- theta_index: int/ 1d array The theta coordinate index or axis. Returns ------- theta_1d: float The coordinate(s) :math:`\\theta` on the sphere. """ return self._theta_1d @property def phi_1d(self): """Returns the coordinate values phi given the coordinate index. The coordinate index lies in (0, nphi). The actual indices without the ghost and extra zones is (nghosts, nphi-nghosts). Parameters ----------- phi_1d: int / 1d array The phi coordinate index or axis. Returns ------- phi_1d: float or 1d array The coordinate(s) :math:`\\phi` on the sphere. """ return self._phi_1d @property def meshgrid(self): """The (:math:`\\theta, \\phi)`: coordinate meshes. Excludes the ghost zones. Returns ------- theta: 2d array The :math:`\\theta` coordinate matrix for vectorization. phi: 2d array The :math:`\\phi` coordinate matrix for vectorization. """ return self._meshgrid