"""
Class for computing and storing the coordinates on the grid and other related
quantities.
"""
import numpy as np
[docs]
class Grid:
r"""
Computes the grid on which the Boltzmann equation is solved.
Grid is 3d, and consists of the physical coordinates:
- :math:`\xi`, position perpendicular to the wall,
- :math:`p_z`, momentum perpendicular to the wall,
- :math:`p_\Vert`, momentum magnitude parallel to the wall.
In addition there are the corresponding compactified coordinates on the
interval [-1, 1],
.. math::
\chi \equiv \frac{\xi}{\sqrt{\xi^2 + L_{\xi}^2}}, \qquad
\rho_{z} \equiv \tanh\left(\frac{p_z}{2 T_0}\right), \qquad
\rho_{\Vert} \equiv 1 - 2 e^{-p_\Vert/T_0}.
All coordinates are in the wall frame.
Attributes
----------
chiValues : array_like
Grid of the :math:`\chi` direction.
rzValues : array_like
Grid of the :math:`\rho_z` direction.
rpValues : array_like
Grid of the :math:`\rho_\Vert` direction.
xiValues : array_like
Grid of the :math:`\xi` direction.
pzValues : array_like
Grid of the :math:`p_z` direction.
ppValues : array_like
Grid of the :math:`p_\Vert` direction.
"""
[docs]
def __init__(
self,
M: int, # pylint: disable=invalid-name
N: int, # pylint: disable=invalid-name
positionFalloff: float,
momentumFalloffT: float,
spacing: str = "Spectral",
):
r"""
Initialises Grid object.
Compactified coordinates are chosen according to
.. math::
\chi = -\cos\left(\frac{\pi \alpha}{M}\right), \qquad
\rho_{z} = -\cos\left(\frac{\pi \beta}{N}\right), \qquad
\rho_{\Vert} = -\cos\left(\frac{\pi \gamma}{N-1}\right),
with integers :math:`\alpha, \beta, \gamma` taken over
.. math::
\alpha = 0, 1, \dots, M, \qquad
\beta = 0, 1, \dots, N, \qquad
\gamma = 0, 1, \dots, N-1.
These are the Gauss-Lobatto collocation points, here with all
boundary points included.
The boundary points :math:`\chi=\pm 1`, :math:`\rho_z=\pm 1` and
:math:`\rho_{\Vert}=1` correspond to points at infinity. The
deviation from equilibrium is assumed to equal zero at infinity, so
these points are dropped when solving the Boltzmann equations. The
resulting grid is
.. math::
\alpha = 1, 2, \dots, M-1, \qquad
\beta = 1, 2, \dots, N-1, \qquad
\gamma = 0, 1, \dots, N-2.
Parameters
----------
M : int
Number of basis functions in the :math:`\xi` (and :math:`\chi`)
direction.
N : int
Number of basis functions in the :math:`p_z` and :math:`p_\Vert`
(and :math:`\rho_z` and :math:`\rho_\Vert`) directions.
positionFalloff : float
Length scale determining transform in :math:`\xi` direction. Should be
expressed in physical units (the units used in EffectivePotential).
momentumFalloffT : float
Temperature scale determining transform in momentum directions.
Should be close to the plasma temperature.
spacing : {'Spectral', 'Uniform'}
Choose 'Spectral' for the Gauss-Lobatto collocation points, as
required for WallGo's spectral representation, or 'Uniform' for
a uniform grid. Default is 'Spectral'.
"""
self.M = M # pylint: disable=invalid-name
# This number has to be odd
self.N = N # pylint: disable=invalid-name
self.positionFalloff = positionFalloff
assert spacing in [
"Spectral",
"Uniform",
], f"Unknown spacing {spacing}, not 'Spectral' or 'Uniform'"
self.spacing = spacing
self.momentumFalloffT = momentumFalloffT
# Computing the grids in the chi, rz and rp directions
if self.spacing == "Spectral":
# See equation (34) in arXiv:2204.13120.
# Additional signs are so that each coordinate starts from -1.
self.chiValues = -np.cos(np.arange(1, self.M) * np.pi / self.M)
self.rzValues = -np.cos(np.arange(1, self.N) * np.pi / self.N)
self.rpValues = -np.cos(np.arange(0, self.N - 1) * np.pi / (self.N - 1))
elif self.spacing == "Uniform":
dchi = 2 / self.M
drz = 2 / self.N
self.chiValues = np.linspace(
-1 + dchi,
1,
num=self.M - 1,
endpoint=False,
)
self.rzValues = np.linspace(
-1.0 + drz,
1.0,
num=self.N - 1,
endpoint=False,
)
self.rpValues = np.linspace(-1, 1, num=self.N - 1, endpoint=False)
self._cacheCoordinates()
def _cacheCoordinates(self) -> None:
"""
Compute physical coordinates and store them internally.
"""
(self.xiValues, self.pzValues, self.ppValues) = self.decompactify(
self.chiValues, self.rzValues, self.rpValues
)
(self.dxidchi, self.dpzdrz, self.dppdrp) = self.compactificationDerivatives(
self.chiValues, self.rzValues, self.rpValues
)
[docs]
def changeMomentumFalloffScale(self, newScale: float) -> None:
"""
Change the momentum falloff scale.
Parameters
----------
newScale : float
New momentum falloff scale.
"""
self.momentumFalloffT = newScale
self._cacheCoordinates()
[docs]
def changePositionFalloffScale(self, newScale: float) -> None:
"""
Change the position falloff scale.
Parameters
----------
newScale : float
New position falloff scale.
"""
self.positionFalloff = newScale
self._cacheCoordinates()
[docs]
def getCompactCoordinates(
self,
endpoints: bool=False,
direction: str | None=None,
) -> tuple[np.ndarray, ...] | np.ndarray:
r"""
Return compact coordinates of grid.
Parameters
----------
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
direction : string or None, optional
Specifies which coordinates to return. Can either be 'z', 'pz',
'pp' or None. If None, returns a tuple containing the 3 directions.
Default is None.
Returns
----------
chiValues : array_like
Grid of the :math:`\chi` direction.
rzValues : array_like
Grid of the :math:`\rho_z` direction.
rpValues : array_like
Grid of the :math:`\rho_\Vert` direction.
"""
if endpoints:
chi = np.array([-1] + list(self.chiValues) + [1])
rz = np.array([-1]+list(self.rzValues)+[1]) # pylint: disable=invalid-name
rp = np.array(list(self.rpValues) + [1]) # pylint: disable=invalid-name
else:
chi, rz, rp = ( # pylint: disable=invalid-name
self.chiValues, self.rzValues, self.rpValues)
if direction == "z":
return chi
if direction == "pz":
return rz
if direction == "pp":
return rp
return chi, rz, rp
[docs]
def getCoordinates(self, endpoints: bool=False) -> tuple[np.ndarray, ...]:
r"""
Return coordinates of grid, not compactified.
Parameters
----------
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
Returns
----------
xiValues : array_like
Grid of the :math:`\xi` direction.
pzValues : array_like
Grid of the :math:`p_z` direction.
ppValues : array_like
Grid of the :math:`p_\Vert` direction.
"""
if endpoints:
xi = np.array( # pylint: disable=invalid-name
[-np.inf] + list(self.xiValues) + [np.inf])
pz = np.array( # pylint: disable=invalid-name
[-np.inf] + list(self.pzValues) + [np.inf])
pp = np.array( # pylint: disable=invalid-name
list(self.ppValues) + [np.inf])
return xi, pz, pp
return self.xiValues, self.pzValues, self.ppValues
[docs]
def getCompactificationDerivatives(
self,
endpoints: bool=False,
) -> tuple[np.ndarray, ...]:
r"""
Return derivatives of compactified coordinates of grid, with respect to
uncompactified derivatives.
Parameters
----------
endpoints : Bool, optional
If True, include endpoints of grid. Default is False.
Returns
----------
dchiValues : array_like
Grid of the :math:`\partial_\xi\chi` direction.
drzValues : array_like
Grid of the :math:`\partial_{p_z}\rho_z` direction.
drpValues : array_like
Grid of the :math:`\partial_{p_\Vert}\rho_\Vert` direction.
"""
if endpoints:
dxidchi = np.array([np.inf] + list(self.dxidchi) + [np.inf])
dpzdrz = np.array([np.inf] + list(self.dpzdrz) + [np.inf])
dppdrp = np.array(list(self.dppdrp) + [np.inf])
return dxidchi, dpzdrz, dppdrp
return self.dxidchi, self.dpzdrz, self.dppdrp
[docs]
def compactify(
self,
z: np.ndarray, # pylint: disable=invalid-name
pz: np.ndarray, # pylint: disable=invalid-name
pp: np.ndarray, # pylint: disable=invalid-name
) -> tuple[np.ndarray, ...]:
"""
Transforms coordinates to [-1, 1] interval
Parameters
----------
z : array-like
Physical z (or xi) coordinate.
pz : array-like
Physical pz coordinate.
pp : array-like
Physical p_par coordinate.
Returns
-------
z_compact : array-like
Compact z coordinate (chi).
pz_compact : array-like
Compact pz coordinate (rho_z).
pp_compact : array-like
Compact p_par coordinate (rho_par).
"""
zCompact = z / np.sqrt(self.positionFalloff**2 + z**2)
pzCompact = np.tanh(pz / 2 / self.momentumFalloffT)
ppCompact = 1 - 2 * np.exp(-pp / self.momentumFalloffT)
return zCompact, pzCompact, ppCompact
[docs]
def decompactify(
self,
zCompact: np.ndarray,
pzCompact: np.ndarray,
ppCompact: np.ndarray,
) -> tuple[np.ndarray, ...]:
"""
Transforms coordinates to [-1, 1] interval
Parameters
----------
z_compact : array-like
Compact z coordinate (chi).
pz_compact : array-like
Compact pz coordinate (rho_z).
pp_compact : array-like
Compact p_par coordinate (rho_par).
Returns
-------
z : array-like
Physical z (or xi) coordinate.
pz : array-like
Physical pz coordinate.
pp : array-like
Physical p_par coordinate.
"""
z = ( # pylint: disable=invalid-name
self.positionFalloff * zCompact / np.sqrt(1 - zCompact**2))
pz = ( # pylint: disable=invalid-name
2 * self.momentumFalloffT * np.arctanh(pzCompact))
pp = ( # pylint: disable=invalid-name
-self.momentumFalloffT * np.log((1 - ppCompact) / 2))
return z, pz, pp
[docs]
def compactificationDerivatives(
self,
zCompact: np.ndarray,
pzCompact: np.ndarray,
ppCompact: np.ndarray,
) -> tuple[np.ndarray, ...]:
r"""
Derivative :math:`d(X)/d(X_\text{compact})` of coordinate transforms to [-1, 1] interval.
Parameters
----------
z_compact : array-like
Compact z coordinate (chi).
pz_compact : array-like
Compact pz coordinate (rho_z).
pp_compact : array-like
Compact p_par coordinate (rho_par).
Returns
-------
dzdzCompact : array-like
Derivative d(z)/d(chi).
dpzdpzCompact : array-like
PDerivative d(p_z)/d(rho_z).
dppdppCompact : array-like
Derivative d(p_par)/d(rho_par).
"""
dzdzCompact = self.positionFalloff / (1 - zCompact**2) ** 1.5
dpzdpzCompact = 2 * self.momentumFalloffT / (1 - pzCompact**2)
dppdppCompact = self.momentumFalloffT / (1 - ppCompact)
return dzdzCompact, dpzdpzCompact, dppdppCompact