Source code for pde.grids.operators.spherical_sym

r"""
This module implements differential operators on spherical grids 

.. autosummary::
   :nosignatures:

   make_laplace
   make_gradient
   make_divergence
   make_vector_gradient
   make_tensor_divergence
   make_poisson_solver

.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""

from __future__ import annotations

from typing import Literal

import numpy as np

from ...tools.docstrings import fill_in_docstring
from ...tools.numba import jit
from ...tools.typing import OperatorType
from ..boundaries import Boundaries
from ..spherical import SphericalSymGrid
from .common import make_general_poisson_solver


[docs] @SphericalSymGrid.register_operator("laplace", rank_in=0, rank_out=0) @fill_in_docstring def make_laplace(grid: SphericalSymGrid, *, conservative: bool = True) -> OperatorType: """make a discretized laplace operator for a spherical grid {DESCR_SPHERICAL_GRID} Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined conservative (bool): Flag indicating whether the laplace operator should be conservative (which results in slightly slower computations). Conservative operators ensure mass conservation. Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] dr = grid.discretization[0] rs = grid.axes_coords[0] r_min, r_max = grid.axes_bounds[0] if conservative: # create a conservative spherical laplace operator rl = rs - dr / 2 # inner radii of spherical shells rh = rs + dr / 2 # outer radii assert np.isclose(rl[0], r_min) and np.isclose(rh[-1], r_max) volumes = (rh**3 - rl**3) / 3 # volume of the spherical shells factor_l = rl**2 / (dr * volumes) factor_h = rh**2 / (dr * volumes) @jit def laplace(arr: np.ndarray, out: np.ndarray) -> None: """apply laplace operator to array `arr`""" for i in range(1, dim_r + 1): # iterate inner radial points term_h = factor_h[i - 1] * (arr[i + 1] - arr[i]) term_l = factor_l[i - 1] * (arr[i] - arr[i - 1]) out[i - 1] = term_h - term_l else: # create an operator that is not conservative dr2 = 1 / dr**2 @jit def laplace(arr: np.ndarray, out: np.ndarray) -> None: """apply laplace operator to array `arr`""" for i in range(1, dim_r + 1): # iterate inner radial points diff_2 = (arr[i + 1] - 2 * arr[i] + arr[i - 1]) * dr2 diff_1 = (arr[i + 1] - arr[i - 1]) / (rs[i - 1] * dr) out[i - 1] = diff_2 + diff_1 return laplace # type: ignore
[docs] @SphericalSymGrid.register_operator("gradient", rank_in=0, rank_out=1) @fill_in_docstring def make_gradient( grid: SphericalSymGrid, *, method: Literal["central", "forward", "backward"] = "central", ) -> OperatorType: """make a discretized gradient operator for a spherical grid {DESCR_SPHERICAL_GRID} Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined method (str): The method for calculating the derivative. Possible values are 'central', 'forward', and 'backward'. Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] if method == "central": scale_r = 0.5 / grid.discretization[0] elif method in {"forward", "backward"}: scale_r = 1 / grid.discretization[0] else: raise ValueError(f"Unknown derivative type `{method}`") @jit def gradient(arr: np.ndarray, out: np.ndarray) -> None: """apply gradient operator to array `arr`""" for i in range(1, dim_r + 1): # iterate inner radial points if method == "central": out[0, i - 1] = (arr[i + 1] - arr[i - 1]) * scale_r elif method == "forward": out[0, i - 1] = (arr[i + 1] - arr[i]) * scale_r elif method == "backward": out[0, i - 1] = (arr[i] - arr[i - 1]) * scale_r out[1, i - 1] = out[2, i - 1] = 0 # no angular dependence by definition return gradient # type: ignore
[docs] @SphericalSymGrid.register_operator("gradient_squared", rank_in=0, rank_out=0) @fill_in_docstring def make_gradient_squared( grid: SphericalSymGrid, *, central: bool = True ) -> OperatorType: """make a discretized gradient squared operator for a spherical grid {DESCR_SPHERICAL_GRID} Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined central (bool): Whether a central difference approximation is used for the gradient operator. If this is False, the squared gradient is calculated as the mean of the squared values of the forward and backward derivatives. Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] dr = grid.discretization[0] if central: # use central differences scale = 0.25 / dr**2 @jit def gradient_squared(arr: np.ndarray, out: np.ndarray) -> None: """apply squared gradient operator to array `arr`""" for i in range(1, dim_r + 1): # iterate inner radial points out[i - 1] = (arr[i + 1] - arr[i - 1]) ** 2 * scale else: # use forward and backward differences scale = 0.5 / dr**2 @jit def gradient_squared(arr: np.ndarray, out: np.ndarray) -> None: """apply squared gradient operator to array `arr`""" for i in range(1, dim_r + 1): # iterate inner radial points term = (arr[i + 1] - arr[i]) ** 2 + (arr[i] - arr[i - 1]) ** 2 out[i - 1] = term * scale return gradient_squared # type: ignore
[docs] @SphericalSymGrid.register_operator("divergence", rank_in=1, rank_out=0) @fill_in_docstring def make_divergence( grid: SphericalSymGrid, *, safe: bool = True, conservative: bool = True, method: Literal["central", "forward", "backward"] = "central", ) -> OperatorType: """make a discretized divergence operator for a spherical grid {DESCR_SPHERICAL_GRID} Warning: This operator ignores the θ-component of the field when calculating the divergence. This is because the resulting scalar field could not be expressed on a :class:`~pde.grids.spherical_sym.SphericalSymGrid`. Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined safe (bool): Add extra checks for the validity of the input conservative (bool): Flag indicating whether the operator should be conservative (which results in slightly slower computations). Conservative operators ensure mass conservation. method (str): The method for calculating the derivative. Possible values are 'central', 'forward', and 'backward'. Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] dr = grid.discretization[0] rs = grid.axes_coords[0] if conservative: # implement conservative version of the divergence operator rl = rs - dr / 2 # inner radii of spherical shells rh = rs + dr / 2 # outer radii volumes = (rh**3 - rl**3) / 3 # volume of the spherical shells factor_l = rl**2 / (2 * volumes) factor_h = rh**2 / (2 * volumes) @jit def divergence(arr: np.ndarray, out: np.ndarray) -> None: """apply divergence operator to array `arr`""" if safe: # the θ-component of the vector field are required to be zero. If this # was not the case the scale field resulting from the divergence would # contain components that cannot be expressed in spherically symmetric # coordinates assert np.all(arr[1, 1:-1] == 0) arr_r = arr[0, :] for i in range(1, dim_r + 1): # iterate radial points if method == "central": term_h = factor_h[i - 1] * (arr_r[i] + arr_r[i + 1]) term_l = factor_l[i - 1] * (arr_r[i - 1] + arr_r[i]) elif method == "forward": term_h = 2 * factor_h[i - 1] * arr_r[i + 1] term_l = 2 * factor_l[i - 1] * arr_r[i] elif method == "backward": term_h = 2 * factor_h[i - 1] * arr_r[i] term_l = 2 * factor_l[i - 1] * arr_r[i - 1] out[i - 1] = term_h - term_l else: # implement naive divergence operator factors = 2 / rs # factors that need to be multiplied below @jit def divergence(arr: np.ndarray, out: np.ndarray) -> None: """apply divergence operator to array `arr`""" if safe: # the θ-component of the vector field are required to be zero. If this # was not the case the scale field resulting from the divergence would # contain components that cannot be expressed in spherically symmetric # coordinates assert np.all(arr[1, 1:-1] == 0) arr_r = arr[0, :] for i in range(1, dim_r + 1): # iterate radial points if method == "central": diff_r = (arr_r[i + 1] - arr_r[i - 1]) / (2 * dr) elif method == "forward": diff_r = (arr_r[i + 1] - arr_r[i]) / dr elif method == "backward": diff_r = (arr_r[i] - arr_r[i - 1]) / dr out[i - 1] = diff_r + factors[i - 1] * arr_r[i] return divergence # type: ignore
[docs] @SphericalSymGrid.register_operator("vector_gradient", rank_in=1, rank_out=2) @fill_in_docstring def make_vector_gradient( grid: SphericalSymGrid, *, method: Literal["central", "forward", "backward"] = "central", safe: bool = True, ) -> OperatorType: """make a discretized vector gradient operator for a spherical grid Warning: This operator ignores the two angular components of the field when calculating the gradient. This is because the resulting field could not be expressed on a :class:`~pde.grids.spherical_sym.SphericalSymGrid`. {DESCR_SPHERICAL_GRID} Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined method (str): The method for calculating the derivative. Possible values are 'central', 'forward', and 'backward'. safe (bool): Add extra checks for the validity of the input Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] rs = grid.axes_coords[0] if method == "central": scale_r = 0.5 / grid.discretization[0] elif method in {"forward", "backward"}: scale_r = 1 / grid.discretization[0] else: raise ValueError(f"Unknown derivative type `{method}`") @jit def vector_gradient(arr: np.ndarray, out: np.ndarray) -> None: """apply vector gradient operator to array `arr`""" if safe: # the θ- and φ-components are required to be zero. If this was not the case # the tensor field resulting from the gradient would contain components that # cannot be expressed in spherically symmetric coordinates assert np.all(arr[1:, 1:-1] == 0) # assign aliases arr_r = arr[0, :] out_rr, out_rθ, out_rφ = out[0, 0, :], out[0, 1, :], out[0, 2, :] out_θr, out_θθ, out_θφ = out[1, 0, :], out[1, 1, :], out[1, 2, :] out_φr, out_φθ, out_φφ = out[2, 0, :], out[2, 1, :], out[2, 2, :] # set all components to zero that are not affected out_rθ[:] = 0 out_rφ[:] = 0 out_θr[:] = 0 out_θφ[:] = 0 out_φr[:] = 0 out_φθ[:] = 0 # inner radial boundary condition for i in range(1, dim_r + 1): # iterate radial points if method == "central": out_rr[i - 1] = (arr_r[i + 1] - arr_r[i - 1]) * scale_r elif method == "forward": out_rr[i - 1] = (arr_r[i + 1] - arr_r[i]) * scale_r elif method == "backward": out_rr[i - 1] = (arr_r[i] - arr_r[i - 1]) * scale_r out_θθ[i - 1] = arr_r[i] / rs[i - 1] out_φφ[i - 1] = arr_r[i] / rs[i - 1] return vector_gradient # type: ignore
[docs] @SphericalSymGrid.register_operator("tensor_divergence", rank_in=2, rank_out=1) @fill_in_docstring def make_tensor_divergence( grid: SphericalSymGrid, *, safe: bool = True, conservative: bool = False ) -> OperatorType: """make a discretized tensor divergence operator for a spherical grid {DESCR_SPHERICAL_GRID} Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined safe (bool): Add extra checks for the validity of the input conservative (bool): Flag indicating whether the operator should be conservative (which results in slightly slower computations). Conservative operators ensure mass conservation. Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] rs = grid.axes_coords[0] dr = grid.discretization[0] if conservative: # conservative implementation of the tensor divergence rl = rs - dr / 2 # inner radii of spherical shells rh = rs + dr / 2 # outer radii volumes = (rh**3 - rl**3) / 3 # volume of the spherical shells factor_l = rl**2 / (2 * volumes) factor_h = rh**2 / (2 * volumes) area_factor = (rh**2 - rl**2) / volumes @jit def tensor_divergence(arr: np.ndarray, out: np.ndarray) -> None: """apply tensor divergence operator to array `arr`""" # assign aliases arr_rr, arr_rθ, arr_rφ = arr[0, 0, :], arr[0, 1, :], arr[0, 2, :] arr_θr, arr_θθ, arr_θφ = arr[1, 0, :], arr[1, 1, :], arr[1, 2, :] arr_φr, arr_φθ, arr_φφ = arr[2, 0, :], arr[2, 1, :], arr[2, 2, :] out_r, out_θ, out_φ = out[0, :], out[1, :], out[2, :] # check inputs if safe: # the following conditions need to be met. Otherwise, the vector resulting # from the divergence might contain components that cannot be expressed in # spherically symmetric coordinates assert np.all(arr_φr[1:-1] == 0) assert np.all(arr_rφ[1:-1] == 0) assert np.all(arr_rθ[1:-1] == 0) assert np.all(arr_θr[1:-1] == 0) assert np.all(arr_θθ[1:-1] == arr_φφ[1:-1]) assert np.all(arr_φθ[1:-1] == -arr_θφ[1:-1]) # iterate over inner points for i in range(1, dim_r + 1): term_r_h = factor_h[i - 1] * (arr_rr[i] + arr_rr[i + 1]) term_r_l = factor_l[i - 1] * (arr_rr[i - 1] + arr_rr[i]) out_r[i - 1] = term_r_h - term_r_l - area_factor[i - 1] * arr_φφ[i] out_θ[i - 1] = 0 out_φ[i - 1] = 0 else: # naive implementation of the tensor divergence scale_r = 1 / (2 * dr) @jit def tensor_divergence(arr: np.ndarray, out: np.ndarray) -> None: """apply tensor divergence operator to array `arr`""" # assign aliases arr_rr, arr_rθ, arr_rφ = arr[0, 0, :], arr[0, 1, :], arr[0, 2, :] arr_θr, arr_θθ, arr_θφ = arr[1, 0, :], arr[1, 1, :], arr[1, 2, :] arr_φr, arr_φθ, arr_φφ = arr[2, 0, :], arr[2, 1, :], arr[2, 2, :] out_r, out_θ, out_φ = out[0, :], out[1, :], out[2, :] # check inputs if safe: # the following conditions need to be met. Otherwise, the vector resulting # from the divergence might contain components that cannot be expressed in # spherically symmetric coordinates assert np.all(arr_rθ[1:-1] == 0) assert np.all(arr_θθ[1:-1] == arr_φφ[1:-1]) assert np.all(arr_φθ[1:-1] == -arr_θφ[1:-1]) # iterate over inner points for i in range(1, dim_r + 1): deriv_r = (arr_rr[i + 1] - arr_rr[i - 1]) * scale_r out_r[i - 1] = deriv_r + 2 * (arr_rr[i] - arr_φφ[i]) / rs[i - 1] deriv_r = (arr_θr[i + 1] - arr_θr[i - 1]) * scale_r out_θ[i - 1] = deriv_r + 2 * arr_θr[i] / rs[i - 1] deriv_r = (arr_φr[i + 1] - arr_φr[i - 1]) * scale_r out_φ[i - 1] = deriv_r + (2 * arr_φr[i] + arr_rφ[i]) / rs[i - 1] return tensor_divergence # type: ignore
[docs] @SphericalSymGrid.register_operator("tensor_double_divergence", rank_in=2, rank_out=0) @fill_in_docstring def make_tensor_double_divergence( grid: SphericalSymGrid, *, safe: bool = True, conservative: bool = True ) -> OperatorType: """make a discretized tensor double divergence operator for a spherical grid {DESCR_SPHERICAL_GRID} Args: grid (:class:`~pde.grids.spherical.SphericalSymGrid`): The polar grid for which this operator will be defined safe (bool): Add extra checks for the validity of the input conservative (bool): Flag indicating whether the operator should be conservative (which results in slightly slower computations). Conservative operators ensure mass conservation. Returns: A function that can be applied to an array of values """ assert isinstance(grid, SphericalSymGrid) # calculate preliminary quantities dim_r = grid.shape[0] rs = grid.axes_coords[0] dr = grid.discretization[0] if conservative: # create a conservative double divergence laplace operator rl = rs - dr / 2 # inner radii of spherical shells rh = rs + dr / 2 # outer radii r_min, r_max = grid.axes_bounds[0] assert np.isclose(rl[0], r_min) and np.isclose(rh[-1], r_max) volumes = (rh**3 - rl**3) / 3 # volume of the spherical shells factor_l = rl / volumes factor_h = rh / volumes factor2_l = rl**2 / (dr * volumes) factor2_h = rh**2 / (dr * volumes) @jit def tensor_double_divergence(arr: np.ndarray, out: np.ndarray) -> None: """apply double divergence operator to tensor array `arr`""" # assign aliases arr_rr, arr_rθ, ______ = arr[0, 0, :], arr[0, 1, :], arr[0, 2, :] arr_θr, arr_θθ, ______ = arr[1, 0, :], arr[1, 1, :], arr[1, 2, :] ______, ______, arr_φφ = arr[2, 0, :], arr[2, 1, :], arr[2, 2, :] # check inputs if safe: # the following conditions need to be met. Otherwise, the vector # resulting from the divergence might contain components that cannot be # expressed in spherically symmetric coordinates assert np.all(arr_rθ[1:-1] == -arr_θr[1:-1]) assert np.all(arr_θθ[1:-1] == arr_φφ[1:-1]) # iterate over inner points for i in range(1, dim_r + 1): # radial part arr_rr_h = arr_rr[i] + arr_rr[i + 1] arr_rr_l = arr_rr[i - 1] + arr_rr[i] arr_rr_dr_h = arr_rr[i + 1] - arr_rr[i] arr_rr_dr_l = arr_rr[i] - arr_rr[i - 1] div2_rr_h = factor_h[i - 1] * arr_rr_h + factor2_h[i - 1] * arr_rr_dr_h div2_rr_l = factor_l[i - 1] * arr_rr_l + factor2_l[i - 1] * arr_rr_dr_l div2_rr = div2_rr_h - div2_rr_l # angular part arr_φφ_h = arr_φφ[i] + arr_φφ[i + 1] arr_φφ_l = arr_φφ[i - 1] + arr_φφ[i] div2_φφ = factor_h[i - 1] * arr_φφ_h - factor_l[i - 1] * arr_φφ_l out[i - 1] = div2_rr - div2_φφ else: # naive, non-conservative implementation of the double divergence dr2 = 1 / dr**2 scale_r = 1 / (2 * dr) @jit def tensor_double_divergence(arr: np.ndarray, out: np.ndarray) -> None: """apply double divergence operator to tensor array `arr`""" # assign aliases arr_rr, arr_rθ, ______ = arr[0, 0, :], arr[0, 1, :], arr[0, 2, :] arr_θr, arr_θθ, ______ = arr[1, 0, :], arr[1, 1, :], arr[1, 2, :] ______, ______, arr_φφ = arr[2, 0, :], arr[2, 1, :], arr[2, 2, :] # check inputs if safe: # the following conditions need to be met. Otherwise, the vector # resulting from the divergence might contain components that cannot be # expressed in spherically symmetric coordinates assert np.all(arr_rθ[1:-1] == -arr_θr[1:-1]) assert np.all(arr_θθ[1:-1] == arr_φφ[1:-1]) # iterate over inner points for i in range(1, dim_r + 1): # first derivatives arr_rr_dr = (arr_rr[i + 1] - arr_rr[i - 1]) * scale_r arr_φφ_dr = (arr_φφ[i + 1] - arr_φφ[i - 1]) * scale_r # second derivative term1 = (arr_rr[i + 1] - arr_rr[i - 1]) / (rs[i - 1] * dr) term2 = (arr_rr[i + 1] - 2 * arr_rr[i] + arr_rr[i - 1]) * dr2 lap_rr = term1 + term2 enum = (arr_rr[i] - arr_φφ[i]) / rs[i - 1] + arr_rr_dr - arr_φφ_dr out[i - 1] = lap_rr + 2 * enum / rs[i - 1] return tensor_double_divergence # type: ignore
@fill_in_docstring def _get_laplace_matrix(bcs: Boundaries) -> tuple[np.ndarray, np.ndarray]: """get sparse matrix for laplace operator on a polar grid Args: bcs (:class:`~pde.grids.boundaries.axes.Boundaries`): {ARG_BOUNDARIES_INSTANCE} Returns: tuple: A sparse matrix and a sparse vector that can be used to evaluate the discretized laplacian """ from scipy import sparse assert isinstance(bcs.grid, SphericalSymGrid) bcs.check_value_rank(0) # calculate preliminary quantities dim_r = bcs.grid.shape[0] dr = bcs.grid.discretization[0] rs = bcs.grid.axes_coords[0] r_min, r_max = bcs.grid.axes_bounds[0] # create a conservative spherical laplace operator rl = r_min + dr * np.arange(dim_r) # inner radii of spherical shells rh = rl + dr # outer radii assert np.isclose(rh[-1], r_max) volumes = (rh**3 - rl**3) / 3 # volume of the spherical shells factor_l = (rs - 0.5 * dr) ** 2 / (dr * volumes) factor_h = (rs + 0.5 * dr) ** 2 / (dr * volumes) matrix = sparse.dok_matrix((dim_r, dim_r)) vector = sparse.dok_matrix((dim_r, 1)) for i in range(dim_r): matrix[i, i] += -factor_l[i] - factor_h[i] if i == 0: if r_min == 0: matrix[i, i + 1] = factor_l[i] else: const, entries = bcs[0].get_sparse_matrix_data((-1,)) vector[i] += const * factor_l[i] for k, v in entries.items(): matrix[i, k] += v * factor_l[i] else: matrix[i, i - 1] = factor_l[i] if i == dim_r - 1: const, entries = bcs[0].get_sparse_matrix_data((dim_r,)) vector[i] += const * factor_h[i] for k, v in entries.items(): matrix[i, k] += v * factor_h[i] else: matrix[i, i + 1] = factor_h[i] return matrix, vector
[docs] @SphericalSymGrid.register_operator("poisson_solver", rank_in=0, rank_out=0) @fill_in_docstring def make_poisson_solver( bcs: Boundaries, *, method: Literal["auto", "scipy"] = "auto" ) -> OperatorType: """make a operator that solves Poisson's equation {DESCR_POLAR_GRID} Args: bcs (:class:`~pde.grids.boundaries.axes.Boundaries`): {ARG_BOUNDARIES_INSTANCE} method (str): The chosen method for implementing the operator Returns: A function that can be applied to an array of values """ matrix, vector = _get_laplace_matrix(bcs) return make_general_poisson_solver(matrix, vector, method)