from abc import ABC, abstractmethod
from typing import Union, List, Optional, Tuple, Sequence
import numpy as np
from pararealml.constraint import Constraint, \
apply_constraints_along_last_axis
from pararealml.mesh import CoordinateSystem, Mesh
Slicer = List[Union[int, slice]]
BoundaryConstraintPair = Tuple[
Optional[Constraint],
Optional[Constraint]
]
[docs]class NumericalDifferentiator(ABC):
"""
A base class for numerical differentiators.
"""
def __init__(self, tol: float = 1e-3):
"""
:param tol: the stopping criterion for the Jacobi algorithm when
computing the anti-Laplacian; once the second norm of the
difference of the estimate and the updated estimate drops below
this threshold, the equation is considered to be solved
"""
if tol < 0.:
raise ValueError('tolerance must be non-negative')
self._tol = tol
@abstractmethod
def _derivative(
self,
y: np.ndarray,
d_x: float,
x_axis: int,
derivative_boundary_constraints:
Union[Sequence[Optional[BoundaryConstraintPair]], np.ndarray]
) -> np.ndarray:
"""
Computes the derivative of y with respect to the spatial dimension
defined by x_axis at every point of the mesh.
:param y: the values of y at every point of the mesh
:param d_x: the step size of the mesh along the specified axis
:param x_axis: the spatial dimension that y is to be differentiated
with respect to
:param derivative_boundary_constraints: a sequence of boundary
constraint pairs with an element for each component of y that
allows for applying constraints to the calculated first
derivatives at the boundaries normal to the axis that y is
differentiated with respect to
:return: the derivative of y with respect to the spatial dimension
defined by x_axis
"""
@abstractmethod
def _second_derivative(
self,
y: np.ndarray,
d_x1: float,
d_x2: float,
x_axis1: int,
x_axis2: int,
derivative_boundary_constraints:
Union[Sequence[Optional[BoundaryConstraintPair]], np.ndarray]
) -> np.ndarray:
"""
Computes the second derivative of y with respect to the spatial
dimensions defined by x_axis1 and x_axis2 at every point of the mesh.
:param y: the values of y at every point of the mesh
:param d_x1: the step size of the mesh along the axis defined by
x_axis1
:param d_x2: the step size of the mesh along the axis defined by
x_axis2
:param x_axis1: the first spatial dimension that y is to be
differentiated with respect to
:param x_axis2: the second spatial dimension that y is to be
differentiated with respect to
:param derivative_boundary_constraints: a sequence of boundary
constraint pairs with an element for each component of y
that allows for applying constraints to the calculated first
derivatives at the boundaries normal to the first axis of
differentiation before computing the second derivatives
:return: the second derivative of y with respect to the spatial
dimensions defined by x_axis1 and x_axis2
"""
@abstractmethod
def _next_anti_laplacian_estimate(
self,
y_hat: np.ndarray,
laplacian: np.ndarray,
mesh: Mesh,
derivative_boundary_constraints: Optional[np.ndarray]
) -> np.ndarray:
"""
Given an estimate of the anti-Laplacian, it computes an improved
estimate.
:param y_hat: the current estimated values of the solution at every
point of the mesh
:param laplacian: the Laplacian for which y is to be determined
:param mesh: the mesh representing the discretized spatial domain
:param derivative_boundary_constraints: an optional 2D array
(x dimension, y dimension) of boundary constraint pairs that
specify constraints on the first derivatives of the solution
:return: an improved estimate of y_hat
"""
[docs] def gradient(
self,
y: np.ndarray,
mesh: Mesh,
x_axis: int,
derivative_boundary_constraints: Optional[np.ndarray] = None
) -> np.ndarray:
"""
Computes the column of the Jacobian matrix of y corresponding to the
spatial dimension specified by x_axis at every point of the mesh.
:param y: the values of y at every point of the mesh
:param mesh: the mesh representing the discretized spatial domain
:param x_axis: the spatial dimension that y is to be differentiated
with respect to
:param derivative_boundary_constraints: a 2D array (x dimension,
y dimension) of boundary constraint pairs that allow for applying
constraints to the calculated first derivatives
:return: the element(s) of the gradient of y corresponding to the
specified axis
"""
self._verify_input_shape_matches_mesh(y, mesh)
if not (0 <= x_axis < mesh.dimensions):
raise ValueError(
f'x-axis ({x_axis}) must be non-negative and less than number '
f'of x dimensions ({mesh.dimensions})')
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
y.shape[-1])
derivative = self._derivative(
y,
mesh.d_x[x_axis],
x_axis,
derivative_boundary_constraints[x_axis])
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
return derivative
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
if x_axis == 0:
return derivative
elif x_axis == 1:
return derivative / (r * np.sin(phi))
else:
return derivative / r
else:
if x_axis == 1:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
return derivative / r
else:
return derivative
[docs] def hessian(
self,
y: np.ndarray,
mesh: Mesh,
x_axis1: int,
x_axis2: int,
derivative_boundary_constraints: Optional[np.ndarray] = None
) -> np.ndarray:
"""
Computes the column of the Hessian tensor of y corresponding to the
spatial dimensions defined by x_axis1 and x_axis2 at every point of
the mesh.
:param y: the values of y at every point of the mesh
:param mesh: the mesh representing the discretized spatial domain
:param x_axis1: the first spatial dimension that y is to be
differentiated with respect to
:param x_axis2: the second spatial dimension that y is to be
differentiated with respect to
:param derivative_boundary_constraints: a 2D array (x dimension,
y dimension) of boundary constraint pairs that allow for applying
constraints to the calculated first derivatives before using them
to compute the second derivatives
:return: the element(s) of the Hessian of y corresponding to the
specified axes
"""
self._verify_input_shape_matches_mesh(y, mesh)
if not (0 <= x_axis1 < mesh.dimensions) \
or not (0 <= x_axis2 < mesh.dimensions):
raise ValueError(
f'both first x-axis ({x_axis1}) and second x-axis ({x_axis2}) '
'must be non-negative and less than number of x dimensions '
f'({mesh.dimensions})')
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
y.shape[-1])
second_derivative = self._second_derivative(
y,
mesh.d_x[x_axis1],
mesh.d_x[x_axis2],
x_axis1,
x_axis2,
derivative_boundary_constraints[x_axis1])
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
return second_derivative
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
if x_axis1 == 0 and x_axis2 == 0:
return second_derivative
elif x_axis1 == 1 and x_axis2 == 1:
sin_phi = np.sin(phi)
cos_phi = np.cos(phi)
d_y_over_d_r = self._derivative(
y,
mesh.d_x[0],
0,
derivative_boundary_constraints[0])
d_y_over_d_phi = self._derivative(
y,
mesh.d_x[2],
2,
derivative_boundary_constraints[2])
return (
d_y_over_d_r +
(
second_derivative / sin_phi +
cos_phi * d_y_over_d_phi
) / (r * sin_phi)
) / r
elif x_axis1 == 2 and x_axis2 == 2:
d_y_over_d_r = self._derivative(
y,
mesh.d_x[0],
0,
derivative_boundary_constraints[0])
return (second_derivative / r + d_y_over_d_r) / r
elif (x_axis1 == 0 and x_axis2 == 1) \
or (x_axis1 == 1 and x_axis2 == 0):
d_y_over_d_theta = self._derivative(
y,
mesh.d_x[1],
1,
derivative_boundary_constraints[1])
return (
second_derivative - d_y_over_d_theta / r
) / (r * np.sin(phi))
elif (x_axis1 == 0 and x_axis2 == 2) \
or (x_axis1 == 2 and x_axis2 == 0):
d_y_over_d_phi = self._derivative(
y,
mesh.d_x[2],
2,
derivative_boundary_constraints[2])
return (second_derivative - d_y_over_d_phi / r) / r
else:
sin_phi = np.sin(phi)
cos_phi = np.cos(phi)
d_y_over_d_theta = self._derivative(
y,
mesh.d_x[1],
1,
derivative_boundary_constraints[1])
return (
sin_phi * second_derivative - cos_phi * d_y_over_d_theta
) / (r * sin_phi) ** 2
else:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
if (x_axis1 == 0 or x_axis1 == 2) \
and (x_axis2 == 0 or x_axis2 == 2):
return second_derivative
elif x_axis1 == 1 and x_axis2 == 1:
d_y_over_d_r = self._derivative(
y,
mesh.d_x[0],
0,
derivative_boundary_constraints[0])
return (second_derivative / r + d_y_over_d_r) / r
elif (x_axis1 == 1 and x_axis2 == 0) \
or (x_axis1 == 0 and x_axis2 == 1):
d_y_over_d_theta = self._derivative(
y,
mesh.d_x[1],
1,
derivative_boundary_constraints[1])
return (second_derivative - d_y_over_d_theta / r) / r
else:
return second_derivative / r
[docs] def divergence(
self,
y: np.ndarray,
mesh: Mesh,
derivative_boundary_constraints: Optional[np.ndarray] = None
) -> np.ndarray:
"""
Computes the divergence of y with respect to x at every point of the
mesh.
:param y: the values of y at every point of the mesh
:param mesh: the mesh representing the discretized spatial domain
:param derivative_boundary_constraints: a 2D array (x dimension,
y dimension) of boundary constraint pairs that allow for applying
constraints to the calculated first derivatives before using them
to compute the divergence
:return: the divergence of y
"""
self._verify_input_is_a_vector_field(y, mesh)
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
y.shape[-1])
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
div = np.zeros(y.shape[:-1] + (1,))
for i in range(y.shape[-1]):
div += self._derivative(
y[..., i:i + 1],
mesh.d_x[i],
i,
derivative_boundary_constraints[i, i:i + 1])
return div
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
sin_phi = np.sin(phi)
cos_phi = np.cos(phi)
y_r = y[..., :1]
y_theta = y[..., 1:2]
y_phi = y[..., 2:]
d_y_r_over_d_r = self._derivative(
y_r,
mesh.d_x[0],
0,
derivative_boundary_constraints[0, :1])
d_y_theta_over_d_theta = self._derivative(
y_theta,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 1:2])
d_y_phi_over_d_phi = self._derivative(
y_phi,
mesh.d_x[2],
2,
derivative_boundary_constraints[2, 2:])
return d_y_r_over_d_r + (
d_y_phi_over_d_phi + 2. * y_r +
(d_y_theta_over_d_theta + cos_phi * y_phi) / sin_phi
) / r
else:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
y_r = y[..., :1]
y_theta = y[..., 1:2]
d_y_r_over_d_r = self._derivative(
y_r,
mesh.d_x[0],
0,
derivative_boundary_constraints[0, :1])
d_y_theta_over_d_theta = self._derivative(
y_theta,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 1:2])
div = d_y_r_over_d_r + (y_r + d_y_theta_over_d_theta) / r
if mesh.coordinate_system_type == CoordinateSystem.POLAR:
return div
else:
y_z = y[..., 2:]
d_y_z_over_d_z = self._derivative(
y_z,
mesh.d_x[2],
2,
derivative_boundary_constraints[2, 2:])
return div + d_y_z_over_d_z
[docs] def curl(
self,
y: np.ndarray,
mesh: Mesh,
curl_ind: int = 0,
derivative_boundary_constraints: Optional[np.ndarray] = None
) -> np.ndarray:
"""
Computes the curl_ind-th component of the curl of y at every point of
the mesh.
:param y: the values of y at every point of the mesh
:param mesh: the mesh representing the discretized spatial domain
:param curl_ind: the index of the component of the curl of y to
compute; if y is a two dimensional vector field, it must be 0
:param derivative_boundary_constraints: a 2D array (x dimension,
y dimension) of boundary constraint pairs that allow for applying
constraints to the calculated first derivatives before using them
to compute the curl
:return: the curl of y
"""
self._verify_input_is_a_vector_field(y, mesh)
if not (2 <= mesh.dimensions <= 3):
raise ValueError(
f'number of x dimensions ({mesh.dimensions}) must be 2 or 3')
if mesh.dimensions == 2 and curl_ind != 0:
raise ValueError(f'curl index ({curl_ind}) must be 0 for 2D curl')
if not (0 <= curl_ind < mesh.dimensions):
raise ValueError(
f'curl index ({curl_ind}) must be non-negative and less than '
f'number of x dimensions ({mesh.dimensions})')
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
y.shape[-1])
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
if mesh.dimensions == 2 or curl_ind == 2:
return self._derivative(
y[..., 1:2],
mesh.d_x[0],
0,
derivative_boundary_constraints[0, 1:2]
) - self._derivative(
y[..., :1],
mesh.d_x[1],
1,
derivative_boundary_constraints[1, :1]
)
elif curl_ind == 0:
return self._derivative(
y[..., 2:],
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 2:]
) - self._derivative(
y[..., 1:2],
mesh.d_x[2],
2,
derivative_boundary_constraints[2, 1:2]
)
else:
return self._derivative(
y[..., :1],
mesh.d_x[2],
2,
derivative_boundary_constraints[2, :1]
) - self._derivative(
y[..., 2:],
mesh.d_x[0],
0,
derivative_boundary_constraints[0, 2:]
)
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
if curl_ind == 0:
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
sin_phi = np.sin(phi)
cos_phi = np.cos(phi)
y_theta = y[..., 1:2]
y_phi = y[..., 2:]
d_y_theta_over_d_phi = self._derivative(
y_theta,
mesh.d_x[2],
2,
derivative_boundary_constraints[2, 1:2])
d_y_phi_over_d_theta = self._derivative(
y_phi,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 2:])
return (
d_y_theta_over_d_phi +
(cos_phi * y_theta - d_y_phi_over_d_theta) / sin_phi
) / r
elif curl_ind == 1:
y_r = y[..., :1]
y_phi = y[..., 2:]
d_y_r_over_d_phi = self._derivative(
y_r,
mesh.d_x[2],
2,
derivative_boundary_constraints[2, :1])
d_y_phi_over_d_r = self._derivative(
y_phi,
mesh.d_x[0],
0,
derivative_boundary_constraints[0, 2:])
return d_y_phi_over_d_r + (y_phi - d_y_r_over_d_phi) / r
else:
sin_phi = np.sin(
mesh.vertex_coordinate_grids[2][..., np.newaxis])
y_r = y[..., :1]
y_theta = y[..., 1:2]
d_y_r_over_d_theta = self._derivative(
y_r,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, :1])
d_y_theta_over_d_r = self._derivative(
y_theta,
mesh.d_x[0],
0,
derivative_boundary_constraints[0, 1:2])
return -d_y_theta_over_d_r + \
(d_y_r_over_d_theta / sin_phi - y_theta) / r
else:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
if mesh.coordinate_system_type == CoordinateSystem.POLAR \
or curl_ind == 2:
y_r = y[..., :1]
y_theta = y[..., 1:2]
d_y_r_over_d_theta = self._derivative(
y_r,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, :1])
d_y_theta_over_d_r = self._derivative(
y_theta,
mesh.d_x[0],
0,
derivative_boundary_constraints[0, 1:2])
return d_y_theta_over_d_r + (y_theta - d_y_r_over_d_theta) / r
elif curl_ind == 0:
d_y_z_over_d_theta = self._derivative(
y[..., 2:],
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 2:])
d_y_theta_over_d_z = self._derivative(
y[..., 1:2],
mesh.d_x[2],
2,
derivative_boundary_constraints[2, 1:2])
return d_y_z_over_d_theta / r - d_y_theta_over_d_z
else:
d_y_r_over_d_z = self._derivative(
y[..., :1],
mesh.d_x[2],
2,
derivative_boundary_constraints[2, :1])
d_y_z_over_d_r = self._derivative(
y[..., 2:],
mesh.d_x[0],
0,
derivative_boundary_constraints[0, 2:])
return d_y_r_over_d_z - d_y_z_over_d_r
[docs] def laplacian(
self,
y: np.ndarray,
mesh: Mesh,
derivative_boundary_constraints:
Optional[np.ndarray] = None) -> np.ndarray:
"""
Computes the Laplacian of y at every point of the mesh.
If the last rank of y has a dimension greater than one, the
element-wise scalar Laplacian is computed instead of the vector
Laplacian (which only makes a difference in curvilinear coordinate
systems).
:param y: the values of y at every point of the mesh
:param mesh: the mesh representing the discretized spatial domain
:param derivative_boundary_constraints: a 2D array (x dimension,
y dimension) of boundary constraint pairs that allow for applying
constraints to the calculated first derivatives before using them
to compute the second derivatives and the Laplacian
:return: the Laplacian of y
"""
self._verify_input_shape_matches_mesh(y, mesh)
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
y.shape[-1])
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
laplacian = np.zeros_like(y)
for axis in range(y.ndim - 1):
laplacian += self._second_derivative(
y,
mesh.d_x[axis],
mesh.d_x[axis],
axis,
axis,
derivative_boundary_constraints[axis, :])
return laplacian
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
sin_phi = np.sin(phi)
cos_phi = np.cos(phi)
d_y_over_d_r = self._derivative(
y, mesh.d_x[0], 0, derivative_boundary_constraints[0])
d_y_over_d_phi = self._derivative(
y, mesh.d_x[2], 2, derivative_boundary_constraints[2])
d_sqr_y_over_d_r_sqr = self._second_derivative(
y,
mesh.d_x[0],
mesh.d_x[0],
0,
0,
derivative_boundary_constraints[0])
d_sqr_y_over_d_theta_sqr = self._second_derivative(
y,
mesh.d_x[1],
mesh.d_x[1],
1,
1,
derivative_boundary_constraints[1])
d_sqr_y_over_d_phi_sqr = self._second_derivative(
y,
mesh.d_x[2],
mesh.d_x[2],
2,
2,
derivative_boundary_constraints[2])
return d_sqr_y_over_d_r_sqr + (
2 * d_y_over_d_r + (
d_sqr_y_over_d_phi_sqr + (
cos_phi * d_y_over_d_phi +
d_sqr_y_over_d_theta_sqr / sin_phi
) / sin_phi
) / r
) / r
else:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
d_y_over_d_r = self._derivative(
y, mesh.d_x[0], 0, derivative_boundary_constraints[0])
d_sqr_y_over_d_r_sqr = self._second_derivative(
y,
mesh.d_x[0],
mesh.d_x[0],
0,
0,
derivative_boundary_constraints[0])
d_sqr_y_over_d_theta_sqr = self._second_derivative(
y,
mesh.d_x[1],
mesh.d_x[1],
1,
1,
derivative_boundary_constraints[1])
laplacian = d_sqr_y_over_d_r_sqr + \
(d_sqr_y_over_d_theta_sqr / r + d_y_over_d_r) / r
if mesh.coordinate_system_type == CoordinateSystem.POLAR:
return laplacian
else:
d_sqr_y_over_d_z_sqr = self._second_derivative(
y,
mesh.d_x[2],
mesh.d_x[2],
2,
2,
derivative_boundary_constraints[2])
return laplacian + d_sqr_y_over_d_z_sqr
[docs] def vector_laplacian(
self,
y: np.ndarray,
mesh: Mesh,
vector_laplacian_ind: int,
derivative_boundary_constraints: Optional[np.ndarray] = None
) -> np.ndarray:
"""
Computes the vector_laplacian_ind-th component of the vector Laplacian
of y at every point of the mesh.
:param y: the values of y at every point of the mesh
:param mesh: the mesh representing the discretized spatial domain
:param vector_laplacian_ind: the index of the component of the vector
Laplacian of y to compute
:param derivative_boundary_constraints: a 2D array (x dimension,
y dimension) of boundary constraint pairs that allow for applying
constraints to the calculated first derivatives before using them
to compute the vector Laplacian
:return: the vector Laplacian of y
"""
self._verify_input_is_a_vector_field(y, mesh)
if not (0 <= vector_laplacian_ind < mesh.dimensions):
raise ValueError(
f'vector Laplacian index ({vector_laplacian_ind}) must be '
'non-negative and less than number of x dimensions '
f'({mesh.dimensions})')
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
y.shape[-1])
laplacian = self.laplacian(
y[..., vector_laplacian_ind:vector_laplacian_ind + 1],
mesh,
derivative_boundary_constraints[
:, vector_laplacian_ind:vector_laplacian_ind + 1
])
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
return laplacian
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
y_r = y[..., :1]
y_theta = y[..., 1:2]
y_phi = y[..., 2:]
sin_phi = np.sin(phi)
cos_phi = np.cos(phi)
if vector_laplacian_ind == 1:
d_y_theta_over_d_theta = self._derivative(
y_theta,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 1:2])
d_y_phi_over_d_phi = self._derivative(
y_phi,
mesh.d_x[2],
2,
derivative_boundary_constraints[2, 2:])
return laplacian - 2. * (
y_r + d_y_phi_over_d_phi +
(cos_phi * y_phi + d_y_theta_over_d_theta) / sin_phi
) / r ** 2
elif vector_laplacian_ind == 2:
d_y_r_over_d_theta = self._derivative(
y_r,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, :1])
d_y_phi_over_d_theta = self._derivative(
y_phi,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 2:])
return laplacian + 2. * (
d_y_r_over_d_theta +
(cos_phi * d_y_phi_over_d_theta - y_theta / 2.) / sin_phi
) / (sin_phi * r ** 2)
else:
d_y_r_over_d_phi = self._derivative(
y_r,
mesh.d_x[2],
2,
derivative_boundary_constraints[2, :1])
d_y_theta_over_d_theta = self._derivative(
y_theta,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 1:2])
return laplacian + 2. * (
d_y_r_over_d_phi -
(
y_phi / 2. + cos_phi * d_y_theta_over_d_theta
) / sin_phi ** 2
) / r ** 2
else:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
if vector_laplacian_ind == 0:
y_r = y[..., :1]
y_theta = y[..., 1:2]
d_y_theta_over_d_theta = self._derivative(
y_theta,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, 1:2])
return laplacian - (y_r + 2. * d_y_theta_over_d_theta) / r ** 2
elif vector_laplacian_ind == 1:
y_theta = y[..., 1:2]
y_r = y[..., :1]
d_y_r_over_d_theta = self._derivative(
y_r,
mesh.d_x[1],
1,
derivative_boundary_constraints[1, :1])
return laplacian - (y_theta - 2. * d_y_r_over_d_theta) / r ** 2
else:
return laplacian
[docs] def anti_laplacian(
self,
laplacian: np.ndarray,
mesh: Mesh,
y_constraints: Union[Sequence[Optional[Constraint]], np.ndarray],
derivative_boundary_constraints: Optional[np.ndarray] = None,
y_init: Optional[np.ndarray] = None) -> np.ndarray:
"""
Computes the inverse of the element-wise scalar Laplacian using the
Jacobi method.
:param laplacian: the right-hand side of the equation
:param mesh: the mesh representing the discretized spatial domain
:param y_constraints: a sequence of constraints on the values of the
solution containing a constraint for each element of y; each
constraint must constrain the boundary values of corresponding
element of y for the system to be solvable
:param derivative_boundary_constraints: an optional 2D array
(x dimension, y dimension) of boundary constraint pairs that
specify constraints on the first derivatives of the solution
:param y_init: an optional initial estimate of the solution; if it is
None, a random array is used
:return: the array representing the solution to Poisson's equation at
every point of the mesh
"""
self._verify_input_shape_matches_mesh(laplacian, mesh, 'Laplacian')
derivative_boundary_constraints = \
self._verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints,
mesh.dimensions,
laplacian.shape[-1])
if y_init is None:
y = np.random.random(laplacian.shape)
else:
if y_init.shape != laplacian.shape:
raise ValueError
y = y_init
apply_constraints_along_last_axis(y_constraints, y)
diff = np.inf
while diff > self._tol:
y_old = y
y = self._next_anti_laplacian_estimate(
y_old,
laplacian,
mesh,
derivative_boundary_constraints)
apply_constraints_along_last_axis(y_constraints, y)
diff = float(np.linalg.norm(y - y_old))
return y
@staticmethod
def _verify_input_shape_matches_mesh(
input_array: np.ndarray, mesh: Mesh, input_name: str = 'y'):
"""
Throws an error if the shape of the input array up to the last axis
does not match the shape of the vertices of the mesh.
:param input_array: the input array
:param mesh: the mesh to compare against
:param input_name: the name of the input array to include in the error
message
"""
if input_array.shape[:-1] != mesh.vertices_shape:
raise ValueError(
f'{input_name} shape up to second to last axis '
f'{input_array.shape[:-1]} must match mesh vertices shape '
f'{mesh.vertices_shape}')
@staticmethod
def _verify_input_is_a_vector_field(
input_array: np.ndarray, mesh: Mesh):
"""
Throws an error if the shape of the input array is not that of a vector
field evaluated over the vertices of the provided mesh.
:param input_array: the input array
:param mesh: the mesh to compare against
"""
NumericalDifferentiator._verify_input_shape_matches_mesh(
input_array, mesh)
if input_array.shape[-1] != mesh.dimensions:
raise ValueError(
f'y value vector length ({input_array.shape[-1]}) '
f'must match number of x dimensions ({mesh.dimensions})')
@staticmethod
def _verify_and_get_derivative_boundary_constraints(
derivative_boundary_constraints: Optional[np.ndarray],
x_axes: int,
y_elements: int) -> np.ndarray:
"""
If the provided derivative boundary constraints are not None, it just
asserts that the input array's shape matches expectations and returns
it. Otherwise it creates an array of empty objects of the correct
shape and returns that.
:param derivative_boundary_constraints: an array of derivative boundary
constraint pairs
:param x_axes: the number of spatial dimensions
:param y_elements: the number of elements y has
:return: an array of derivative boundary constraint pairs or empty
objects depending on whether the input array is None
"""
if derivative_boundary_constraints is None:
return np.empty((x_axes, y_elements), dtype=object)
if derivative_boundary_constraints.shape != (x_axes, y_elements):
raise ValueError(
'expected derivative boundary constraints shape to be '
f'{(x_axes, y_elements)} but got '
f'{derivative_boundary_constraints.shape}')
return derivative_boundary_constraints
[docs]class ThreePointCentralDifferenceMethod(NumericalDifferentiator):
"""
A numerical differentiator using a three-point (second order) central
difference approximation.
"""
def __init__(self, tol: float = 1e-3):
"""
:param tol: the stopping criterion for the Jacobi algorithm when
computing the anti-Laplacian
"""
super(ThreePointCentralDifferenceMethod, self).__init__(tol)
def _derivative(
self,
y: np.ndarray,
d_x: float,
x_axis: int,
derivative_boundary_constraints:
Union[Sequence[Optional[BoundaryConstraintPair]], np.ndarray]
) -> np.ndarray:
if y.shape[x_axis] <= 2:
raise ValueError(
f'y must contain at least 3 points along x-axis ({x_axis})')
slicer: Slicer = [slice(None)] * y.ndim
halo = np.zeros(y.shape[:x_axis] + (1,) + y.shape[x_axis + 1:])
y_extended = np.concatenate([halo, y, halo], axis=x_axis)
slicer[x_axis] = slice(0, -2)
y_prev = y_extended[tuple(slicer)]
slicer[x_axis] = slice(2, None)
y_next = y_extended[tuple(slicer)]
derivative = (y_next - y_prev) / (2. * d_x)
slicer[x_axis] = slice(None)
for y_ind, constraint_pair in \
enumerate(derivative_boundary_constraints):
if constraint_pair is None:
continue
slicer[-1] = slice(y_ind, y_ind + 1)
lower_boundary_constraint = constraint_pair[0]
if lower_boundary_constraint is not None:
slicer[x_axis] = slice(0, 1)
lower_boundary_constraint.apply(derivative[tuple(slicer)])
upper_boundary_constraint = constraint_pair[1]
if upper_boundary_constraint is not None:
slicer[x_axis] = slice(-1, None)
upper_boundary_constraint.apply(derivative[tuple(slicer)])
return derivative
def _second_derivative(
self,
y: np.ndarray,
d_x1: float,
d_x2: float,
x_axis1: int,
x_axis2: int,
derivative_boundary_constraints:
Union[Sequence[Optional[BoundaryConstraintPair]], np.ndarray]
) -> np.ndarray:
if x_axis1 != x_axis2:
first_derivative = self._derivative(
y,
d_x1,
x_axis1,
derivative_boundary_constraints)
return self._derivative(
first_derivative,
d_x2,
x_axis2,
[None] * y.shape[-1])
if y.shape[x_axis1] <= 2:
raise ValueError(
f'y must contain at least 3 points along x-axis ({x_axis1})')
slicer: Slicer = [slice(None)] * y.ndim
y_extended = self._add_halos_along_axis(
y, x_axis1, d_x1, slicer, derivative_boundary_constraints)
slicer[x_axis1] = slice(0, -2)
y_prev = y_extended[tuple(slicer)]
slicer[x_axis1] = slice(1, -1)
y_curr = y_extended[tuple(slicer)]
slicer[x_axis1] = slice(2, None)
y_next = y_extended[tuple(slicer)]
return (y_next - 2. * y_curr + y_prev) / (d_x1 * d_x2)
def _next_anti_laplacian_estimate(
self,
y_hat: np.ndarray,
laplacian: np.ndarray,
mesh: Mesh,
derivative_boundary_constraints: Optional[np.ndarray]
) -> np.ndarray:
if not np.all(np.array(y_hat.shape[:-1]) > 2):
raise ValueError(
'y must contain at least 3 points along all x axes')
slicer: Slicer = [slice(None)] * y_hat.ndim
anti_laplacian = np.zeros_like(y_hat)
all_d_x_sqr = np.square(mesh.d_x)
r = r_sqr = phi = sin_phi = r_sqr_sin_phi_sqr = None
if mesh.coordinate_system_type != CoordinateSystem.CARTESIAN:
r = mesh.vertex_coordinate_grids[0][..., np.newaxis]
r_sqr = r ** 2
if mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
phi = mesh.vertex_coordinate_grids[2][..., np.newaxis]
sin_phi = np.sin(phi)
r_sqr_sin_phi_sqr = r_sqr * sin_phi ** 2
for axis, d_x in enumerate(mesh.d_x):
d_x_sqr = all_d_x_sqr[axis]
y_hat_extended = self._add_halos_along_axis(
y_hat,
axis,
d_x,
slicer,
derivative_boundary_constraints[axis])
slicer[axis] = slice(0, -2)
y_hat_prev = y_hat_extended[tuple(slicer)]
slicer[axis] = slice(2, None)
y_hat_next = y_hat_extended[tuple(slicer)]
anti_laplacian_update = (y_hat_prev + y_hat_next) / d_x_sqr
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
anti_laplacian += anti_laplacian_update
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
if axis == 0:
anti_laplacian += anti_laplacian_update + \
(y_hat_next - y_hat_prev) / (d_x * r)
elif axis == 1:
anti_laplacian += anti_laplacian_update / r_sqr_sin_phi_sqr
else:
anti_laplacian += (
anti_laplacian_update +
np.cos(phi) * (
y_hat_next - y_hat_prev
) / (2. * d_x * sin_phi)
) / r_sqr
else:
if axis == 0:
anti_laplacian += anti_laplacian_update + \
(y_hat_next - y_hat_prev) / (2. * d_x * r)
elif axis == 1:
anti_laplacian += anti_laplacian_update / r_sqr
else:
anti_laplacian += anti_laplacian_update
slicer[axis] = slice(None)
anti_laplacian -= laplacian
if mesh.coordinate_system_type == CoordinateSystem.CARTESIAN:
return anti_laplacian / (2. / all_d_x_sqr).sum()
elif mesh.coordinate_system_type == CoordinateSystem.SPHERICAL:
return anti_laplacian / (
2. / all_d_x_sqr[0] +
2. / (all_d_x_sqr[1] * r_sqr_sin_phi_sqr) +
2. / (all_d_x_sqr[2] * r_sqr)
)
else:
step_size_coefficient = 2. / all_d_x_sqr[0] + \
2. / (all_d_x_sqr[1] * r_sqr)
if mesh.coordinate_system_type == CoordinateSystem.POLAR:
return anti_laplacian / step_size_coefficient
else:
step_size_coefficient += 2. / all_d_x_sqr[2]
return anti_laplacian / step_size_coefficient
@staticmethod
def _add_halos_along_axis(
y: np.ndarray,
x_axis: int,
d_x: float,
slicer: Slicer,
derivative_boundary_constraints:
Union[Sequence[Optional[BoundaryConstraintPair]], np.ndarray]
) -> np.ndarray:
"""
Adds halo vertices to y along the specified axis based on the provided
derivative boundary constraints and the values at the vertices adjacent
to the boundary vertices.
:param y: the input array to extend with halos
:param x_axis: the axis along which the halos are to be added
:param d_x: the spatial step size
:param slicer: the slicer array to use for indexing the vertices
adjacent to the boundaries
:param derivative_boundary_constraints: the derivative boundary
constraints
:return: y extended with halo vertices along the specified axis
"""
slicer[x_axis] = slice(1, 2)
y_lower_boundary_adjacent = y[tuple(slicer)]
slicer[x_axis] = slice(-2, -1)
y_upper_boundary_adjacent = y[tuple(slicer)]
y_lower_halo = np.zeros_like(y_lower_boundary_adjacent)
y_upper_halo = np.zeros_like(y_upper_boundary_adjacent)
for y_ind, boundary_constraint_pair in \
enumerate(derivative_boundary_constraints):
if boundary_constraint_pair is None:
continue
lower_boundary_constraint = boundary_constraint_pair[0]
if lower_boundary_constraint is not None:
lower_boundary_constraint.multiply_and_add(
y_lower_boundary_adjacent[..., y_ind:y_ind + 1],
-2. * d_x,
y_lower_halo[..., y_ind:y_ind + 1])
upper_boundary_constraint = boundary_constraint_pair[1]
if upper_boundary_constraint is not None:
upper_boundary_constraint.multiply_and_add(
y_upper_boundary_adjacent[..., y_ind:y_ind + 1],
2. * d_x,
y_upper_halo[..., y_ind:y_ind + 1])
return np.concatenate([y_lower_halo, y, y_upper_halo], axis=x_axis)