# Legacy module pending type-checking rework; excluded for now.
# type: ignore
"""Nabla utilities manager."""
from __future__ import annotations
import warnings
from typing import Any
import numpy as np
from pyPLUTO.loadmixin import LoadMixin
from pyPLUTO.loadstate import LoadState
class NablaManager(LoadMixin):
"""Manager for gradient, divergence and curl utilities."""
def __init__(self, state: LoadState) -> None:
"""Initialize the nabla manager with the given load state."""
self.state = state
def _is_number(self, value: Any) -> bool:
"""Check if value is a number.
Parameters
----------
- value (not optional): Any
The value to check.
Returns
-------
- bool
Examples
--------
- Example #1: Check if 1 is a number
>>> _is_number(1)
True
- Example #2: Check if '1' is a number
>>> _is_number("1")
False
"""
return isinstance(value, (int, float))
def _islice_imin_imax(self, xvalue: float, xgrid: np.ndarray) -> list[int]:
"""Return i, imin, imax for xgrid.
The computation is done such that xgrid[i] <= xvalue <=
xgrid[i+1] with a stencil of 3 cells.
If the grid is too small, imin and imax are set to 0 and N,
respectively.
Parameters
----------
- xvalue (not optional): float
The value to check.
- xgrid (not optional): numpy.ndarray
The grid to check.
Returns
-------
- list[int]
Examples
--------
- Example #1: Compute the indices for the stencil of 3 cells
>>> _islice_imin_imax(0.5, np.linspace(0, 1, 11))
"""
# Compute the grid length and the index of the closest grid point
N = len(xgrid)
i = np.argmin(abs(xgrid - xvalue))
# Compute the limits of the stencil
i_min = max(0, i - 1)
i_max = min(N, i + 2)
# If the grid is too small, set imin and imax to 0 and N
if N < 3:
i_min, i_max = 0, N
# If imin or imax are outside the grid, set them to 0 and/or N
else:
i_min = N - 3 if i_max == N else i_min
i_max = 3 if i_min == 0 else i_max
# Return i, imin, imax
return [int(i - i_min), int(i_min), int(i_max)]
def _get_slice_indices(
self,
slice_val: float,
grid: np.ndarray,
grid_size: int,
) -> tuple[int, slice]:
"""Get the slice indices from the slice value.
Parameters
----------
- slice_val: float
The value of the slice.
- grid: np.ndarray
The grid of the data.
- grid_size: int
The size of the grid.
Returns
-------
- tuple
Examples
--------
- Example #1: Get the slice indices
>>> _get_slice_indices(0.5, np.linspace(0, 1, 11), 11)
"""
# If the slice value is a number return the index and the slice
if self._is_number(slice_val):
# Get the slice indices
idx, idx_min, idx_max = self._islice_imin_imax(slice_val, grid)
# Return the indices and the slice
return idx, slice(idx_min, idx_max)
# If the slice value is a list return the indices and the slice
# Return the indices and the slice
return slice(0, grid_size), slice(0, grid_size)
def _warning_cylindrical(self):
"""Emit a DeprecationWarning for the unsupported CYLINDRICAL geometry."""
warning_message = """"""
# CYLINDRICAL geometry has been deprecated since PLUTO v4.4.
# POLAR may be used instead by excluding the phi-direction
# (simply set INCLUDE_JDIR to NO in definitions.h).
# """
warnings.warn(warning_message, DeprecationWarning, stacklevel=3)
[docs]
def gradient(
self,
var: np.ndarray,
x1slice: float | None = None,
x2slice: float | None = None,
x3slice: float | None = None,
edge_order: int = 2,
) -> np.ndarray:
"""Compute the gradient of a specified field 'var'.
The gradient is computed in all available directions using second-order
accurate central differences via the NumPy gradient() function.
The first index of the resulting array
represents the N gradient components. If 'x1slice', 'x2slice', or
'x3slice' are specified, the gradient is only computed at the
corresponding x1, x2, or x3 values. N corresponds to the number of
employed dimensions unless a slice is taken.
Parameters
----------
- var (not optional): np.ndarray
The field whose gradient is calculated (e.g., 'rho', 'vx1'). Must
have the same shape as self.rho.
- x1slice: float | None, default None
If not None, specifies the constant value for the x1 axis.
- x2slice: float | None, default None
If not None, specifies the constant value for the x2 axis.
- x3slice: float | None, default None
If not None, specifies the constant value for the x3 axis.
- edge_order: int | None, default 2
The order of accuracy of derivatives at the domain boundaries.
Returns
-------
- np.ndarray
Examples
--------
- Example #1: Compute the gradient of the density field
>>> import pyPLUTO as pp
>>> D = pp.Load(0)
>>> D.gradient(D.rho)
"""
# If the geometry is not defined raise an error
if self.geom == "UNKNOWN":
raise ValueError("Unknown geometry nabla cannot be computed!")
# if self.geom == 'CYLINDRICAL':
# self._warning_cylindrical()
# Unpack the slice values and grids into tuples
slices = [
(x1slice, self.x1, self.nx1),
(x2slice, self.x2, self.nx2),
(x3slice, self.x3, self.nx3),
]
# Process each slice and store the results
indices, ranges = zip(
*[self._get_slice_indices(s, g, size) for s, g, size in slices],
strict=False,
)
# Unpack the results back into individual variables
i, j, k = indices
irange, jrange, krange = ranges
if self.dim == 1:
grad = np.gradient(
var[irange],
self.x1[irange],
edge_order=edge_order,
)
return np.asarray(grad)[i]
if self.dim == 2:
grad = np.gradient(
var[irange, jrange],
self.x1[irange],
self.x2[jrange],
edge_order=edge_order,
)
if self.geom in ["SPHERICAL", "POLAR"]:
rr, _ = np.meshgrid(
self.x1[irange],
self.x2[jrange],
indexing="ij",
)
grad[1] /= rr
return np.asarray(grad)[:, i, j]
if self.dim == 3:
if self.nx2 == 1:
grad = np.gradient(
var[irange, 0, krange],
self.x1[irange],
self.x3[krange],
edge_order=edge_order,
)
if self.geom == "SPHERICAL":
rr, _ = np.meshgrid(
self.x1[irange],
self.x3[krange],
indexing="ij",
)
grad[1] /= rr * np.sin(self.x2[0])
return np.asarray(grad)[:, i, k]
grad = np.gradient(
var[irange, jrange, krange],
self.x1[irange],
self.x2[jrange],
self.x3[krange],
edge_order=edge_order,
)
if self.geom != "CARTESIAN":
xx1, xx2, _ = np.meshgrid(
self.x1[irange],
self.x2[jrange],
self.x3[krange],
indexing="ij",
)
grad[1] /= xx1
if self.geom == "SPHERICAL":
grad[2] /= xx1 * np.sin(xx2)
return np.asarray(grad)[:, i, j, k]
[docs]
def divergence(
self,
v1: np.ndarray | None = None,
v2: np.ndarray | None = None,
v3: np.ndarray | None = None,
x1slice: float | None = None,
x2slice: float | None = None,
x3slice: float | None = None,
edge_order: int = 2,
) -> np.ndarray:
"""Calculate the divergence of a vector field.
The divergence is specified by its
components v1, v2, and v3 using second-order accurate central
differences via the NumPy gradient() function. If 'x1slice',
'x2slice', or 'x3slice' are specified, the divergence is only
computed at the corresponding x1, x2, or x3 values.
Parameters
----------
- v1: np.ndarray | None
Field corresponding to the x1 vector component. Must have the same
shape as self.rho. Can only be None is a given direction is not
used.
- v2: np.ndarray | None
Field corresponding to the x2 vector component. Must have the same
shape as self.rho. Can only be None is a given direction is not
used.
- v3: np.ndarray | None
Field corresponding to the x3 vector component. Must have the same
shape as self.rho. Can only be None is a given direction is not
used.
- x1slice: float | None
If not None, specifies the constant value for the x1 axis.
- x2slice: float | None
If not None, specifies the constant value for the x2 axis.
- x3slice: float | None
If not None, specifies the constant value for the x3 axis.
- edge_order: int, default 2
The order of accuracy of derivatives at the domain boundaries.
Returns
-------
- np.ndarray
Examples
--------
- Example #1: Calculate the divergence of a vector field
>>> import pyPLUTO as pp
>>> D = pp.Load()
>>> D.divergence(D.vx1, D.vx2, D.vx3)
"""
if self.geom == "CYLINDRICAL":
self._warning_cylindrical()
if self._is_number(x1slice):
i, imin, imax = self._islice_imin_imax(x1slice, self.x1)
irange = slice(imin, imax)
else:
i = slice(0, self.nx1)
irange = i
if self._is_number(x2slice):
j, jmin, jmax = self._islice_imin_imax(x2slice, self.x2)
jrange = slice(jmin, jmax)
else:
j = slice(0, self.nx2)
jrange = j
if self._is_number(x3slice):
k, kmin, kmax = self._islice_imin_imax(x3slice, self.x3)
krange = slice(kmin, kmax)
else:
k = slice(0, self.nx3)
krange = k
if self.dim == 1:
var1 = np.copy(v1[irange])
rr = self.x1[irange]
if self.geom in ["POLAR", "CYLINDRICAL"]:
var1 *= rr
elif self.geom == "SPHERICAL":
var1 *= rr**2
div1 = np.gradient(var1, rr, edge_order=edge_order)
if self.geom in ["POLAR", "CYLINDRICAL"]:
div1 /= rr
elif self.geom == "SPHERICAL":
div1 /= rr**2
return div1[i]
if self.dim == 2:
var1 = np.copy(v1[irange, jrange])
var2 = np.copy(v2[irange, jrange])
if self.geom != "CARTESIAN":
rr, tt = np.meshgrid(
self.x1[irange],
self.x2[jrange],
indexing="ij",
)
if self.geom in ["POLAR", "CYLINDRICAL"]:
var1 *= rr
elif self.geom == "SPHERICAL":
var1 *= rr**2
var2 *= np.sin(tt)
div1 = np.gradient(
var1,
self.x1[irange],
axis=0,
edge_order=edge_order,
)
div2 = np.gradient(
var2,
self.x2[jrange],
axis=1,
edge_order=edge_order,
)
if self.geom in ["POLAR", "CYLINDRICAL"]:
div1 /= rr
if self.geom == "POLAR":
div2 /= rr
elif self.geom == "SPHERICAL":
div1 /= rr**2
div2 /= rr * np.sin(tt)
return div1[i, j] + div2[i, j]
if self.dim == 3:
if self.nx2 == 1:
var1 = np.copy(v1[irange, 0, krange])
var3 = np.copy(v3[irange, 0, krange])
if self.geom != "CARTESIAN":
rr, _ = np.meshgrid(
self.x1[irange],
self.x3[krange],
indexing="ij",
)
if self.geom == "POLAR":
var1 *= rr
elif self.geom == "SPHERICAL":
var1 *= rr**2
div1 = np.gradient(
var1,
self.x1[irange],
axis=0,
edge_order=edge_order,
)
div3 = np.gradient(
var3,
self.x3[krange],
axis=1,
edge_order=edge_order,
)
if self.geom == "POLAR":
div1 /= rr
elif self.geom == "SPHERICAL":
div1 /= rr**2
div3 /= rr * np.sin(self.x2[0])
return div1[i, k] + div3[i, k]
var1 = np.copy(v1[irange, jrange, krange])
var2 = np.copy(v2[irange, jrange, krange])
var3 = np.copy(v3[irange, jrange, krange])
if self.geom != "CARTESIAN":
rr, tt, _ = np.meshgrid(
self.x1[irange],
self.x2[jrange],
self.x3[krange],
indexing="ij",
)
if self.geom == "POLAR":
var1 *= rr
elif self.geom == "SPHERICAL":
var1 *= rr**2
var2 *= np.sin(tt)
div1 = np.gradient(
var1,
self.x1[irange],
axis=0,
edge_order=edge_order,
)
div2 = np.gradient(
var2,
self.x2[jrange],
axis=1,
edge_order=edge_order,
)
div3 = np.gradient(
var3,
self.x3[krange],
axis=2,
edge_order=edge_order,
)
if self.geom == "POLAR":
div1 /= rr
div2 /= rr
elif self.geom == "SPHERICAL":
div1 /= rr**2
div2 /= rr * np.sin(tt)
div3 /= rr * np.sin(tt)
return div1[i, j, k] + div2[i, j, k] + div3[i, j, k]
[docs]
def curl(
self,
v1: np.ndarray | None = None,
v2: np.ndarray | None = None,
v3: np.ndarray | None = None,
x1slice: float | None = None,
x2slice: float | None = None,
x3slice: float | None = None,
edge_order: int = 2,
) -> np.ndarray:
"""Calculate the curl of a specified vector field (v1, v2, v3).
Uses second-order accurate central differences via the NumPy
gradient() function. Unlike in divergence(), all three vector
components must be specified. The resulting array has its first
index representing the 3 curl components, while the remaining
indices correspond to the grid location. If 'x1slice', 'x2slice', or
'x3slice' are specified, the curl is only computed at the
corresponding x1, x2, or x3 values.
Parameters
----------
- v1: np.ndarray | None
Field corresponding to the x1 vector component.
Must have the same shape as self.rho.
- v2: np.ndarray | None
Field corresponding to the x2 vector component.
Must have the same shape as self.rho.
- v3: np.ndarray | None
Field corresponding to the x3 vector component.
Must have the same shape as self.rho.
- x1slice: float | None
If not None, specifies the constant value for the x1 axis.
- x2slice: float | None
If not None, specifies the constant value for the x2 axis.
- x3slice: float | None
If not None, specifies the constant value for the x3 axis.
- edge_order: int, default 2
The order of accuracy of derivatives at the domain boundaries.
Returns
-------
- np.ndarray
Examples
--------
- Example #1: Calculate the curl of a vector field
>>> import pyPLUTO as pp
>>> D = pp.Load()
>>> D.curl(D.vx1, D.vx2, D.vx3)
"""
if self.geom == "CYLINDRICAL":
self._warning_cylindrical()
if self._is_number(x1slice):
i, imin, imax = self._islice_imin_imax(x1slice, self.x1)
irange = slice(imin, imax)
else:
i = slice(0, self.nx1)
irange = i
if self._is_number(x2slice):
j, jmin, jmax = self._islice_imin_imax(x2slice, self.x2)
jrange = slice(jmin, jmax)
else:
j = slice(0, self.nx2)
jrange = j
if self._is_number(x3slice):
k, kmin, kmax = self._islice_imin_imax(x3slice, self.x3)
krange = slice(kmin, kmax)
else:
k = slice(0, self.nx3)
krange = k
if self.dim == 1:
raise ValueError("curl() requires DIMENSION >= 2")
if self.dim == 2:
var1 = np.copy(v1[irange, jrange])
var2 = np.copy(v2[irange, jrange])
var3 = np.copy(v3[irange, jrange])
if self.geom != "CARTESIAN":
rr, tt = np.meshgrid(
self.x1[irange],
self.x2[jrange],
indexing="ij",
)
if self.geom == "POLAR":
var2 *= rr
elif self.geom in "CYLINDRICAL":
var3 *= rr
elif self.geom == "SPHERICAL":
var2 *= rr
var3 *= rr * np.sin(tt)
dv1_dx2 = np.gradient(var1, self.x2, axis=1, edge_order=edge_order)
dv2_dx1 = np.gradient(var2, self.x1, axis=0, edge_order=edge_order)
dv3_dx1, dv3_dx2 = np.gradient(
var3,
self.x1,
self.x2,
edge_order=edge_order,
)
curl1 = dv3_dx2
curl2 = -dv3_dx1
curl3 = dv2_dx1 - dv1_dx2
if self.geom == "CYLINDRICAL":
curl1 *= -1
curl2 *= -1
curl3 *= -1
if self.geom == "POLAR":
curl1 /= rr
curl3 /= rr
elif self.geom == "CYLINDRICAL":
curl1 /= rr
curl2 /= rr
elif self.geom == "SPHERICAL":
curl1 /= rr**2 * np.sin(tt)
curl2 /= rr * np.sin(tt)
curl3 /= rr
return np.asarray([curl1[i, j], curl2[i, j], curl3[i, j]])
if self.dim == 3:
if self.nx2 == 1:
var1 = np.copy(v1[irange, 0, krange])
var2 = np.copy(v2[irange, 0, krange])
var3 = np.copy(v3[irange, 0, krange])
if self.geom != "CARTESIAN":
rr, _ = np.meshgrid(
self.x1[irange],
self.x3[krange],
indexing="ij",
)
if self.geom == "POLAR":
var2 *= rr
elif self.geom == "SPHERICAL":
var2 *= rr
var3 *= rr * np.sin(self.x2[0])
dv1_dx3 = np.gradient(
var1,
self.x3,
axis=1,
edge_order=edge_order,
)
dv2_dx1, dv2_dx3 = np.gradient(
var2,
self.x1,
self.x3,
edge_order=edge_order,
)
dv3_dx1 = np.gradient(
var3,
self.x1,
axis=0,
edge_order=edge_order,
)
curl1 = -dv2_dx3
curl2 = dv1_dx3 - dv3_dx1
curl3 = dv2_dx1
if self.geom == "POLAR":
curl1 /= rr
curl3 /= rr
elif self.geom == "SPHERICAL":
curl1 /= rr**2 * np.sin(self.x2[0])
curl2 /= rr * np.sin(self.x2[0])
curl3 /= rr
return np.asarray([curl1[i, k], curl2[i, k], curl3[i, k]])
var1 = np.copy(v1[irange, jrange, krange])
var2 = np.copy(v2[irange, jrange, krange])
var3 = np.copy(v3[irange, jrange, krange])
if self.geom != "CARTESIAN":
rr, tt, _ = np.meshgrid(
self.x1[irange],
self.x2[jrange],
self.x3[krange],
indexing="ij",
)
if self.geom == "POLAR":
var2 *= rr
elif self.geom == "SPHERICAL":
var2 *= rr
var3 *= rr * np.sin(tt)
dv1_dx2, dv1_dx3 = np.gradient(
var1,
self.x2,
self.x3,
axis=[1, 2],
edge_order=edge_order,
)
dv2_dx1, dv2_dx3 = np.gradient(
var2,
self.x1,
self.x3,
axis=[0, 2],
edge_order=edge_order,
)
dv3_dx1, dv3_dx2 = np.gradient(
var3,
self.x1,
self.x2,
axis=[0, 1],
edge_order=edge_order,
)
curl1 = dv3_dx2 - dv2_dx3
curl2 = dv1_dx3 - dv3_dx1
curl3 = dv2_dx1 - dv1_dx2
if self.geom == "POLAR":
curl1 /= rr
curl3 /= rr
elif self.geom == "SPHERICAL":
curl1 /= rr**2 * np.sin(tt)
curl2 /= rr * np.sin(tt)
curl3 /= rr
return np.asarray(
[curl1[i, j, k], curl2[i, j, k], curl3[i, j, k]],
)