""" 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