"""This module implements differential operators on Cartesian grids.
.. autosummary::
:nosignatures:
CartesianLaplacian
CartesianGradient
CartesianGradientSquared
CartesianDivergence
CartesianVectorGradient
CartesianVectorLaplacian
CartesianTensorDivergence
.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
from ....grids import CartesianGrid, GridBase
from ..backend import TorchBackend
from .common import TorchDifferentialOperator
if TYPE_CHECKING:
import numpy as np
from torch import Tensor
from ....grids.boundaries import BoundariesList
[docs]
@TorchBackend.register_operator(CartesianGrid, "laplace", rank_in=0, rank_out=0)
class CartesianLaplacian(TorchDifferentialOperator):
"""Cartesian Laplace using torch."""
rank_in = 0
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
dtype: np.dtype,
):
"""Initialize the Cartesian Laplacian operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
self.scale = self.grid.discretization**-2
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
data_full = self.get_full_data(arr, args=args)
if self.grid.num_axes == 1:
scale = self.scale[0]
return (data_full[:-2] - 2 * data_full[1:-1] + data_full[2:]) * scale # type: ignore
if self.grid.num_axes == 2:
lap_x = (
data_full[:-2, 1:-1] - 2 * data_full[1:-1, 1:-1] + data_full[2:, 1:-1]
) * self.scale[0]
lap_y = (
data_full[1:-1, :-2] - 2 * data_full[1:-1, 1:-1] + data_full[1:-1, 2:]
) * self.scale[1]
return lap_x + lap_y # type: ignore
if self.grid.num_axes == 3:
data_mid2 = 2 * data_full[1:-1, 1:-1, 1:-1]
lap_x = (
data_full[:-2, 1:-1, 1:-1] - data_mid2 + data_full[2:, 1:-1, 1:-1]
) * self.scale[0]
lap_y = (
data_full[1:-1, :-2, 1:-1] - data_mid2 + data_full[1:-1, 2:, 1:-1]
) * self.scale[1]
lap_z = (
data_full[1:-1, 1:-1, :-2] - data_mid2 + data_full[1:-1, 1:-1, 2:]
) * self.scale[2]
return lap_x + lap_y + lap_z # type: ignore
raise NotImplementedError
[docs]
@TorchBackend.register_operator(CartesianGrid, "gradient", rank_in=0, rank_out=1)
class CartesianGradient(TorchDifferentialOperator):
"""Cartesian gradient operator using torch."""
rank_in = 0
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
dtype: np.dtype,
):
"""Initialize the Cartesian gradient operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
self.scale = 0.5 / self.grid.discretization
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
data_full = self.get_full_data(arr, args=args)
if self.grid.num_axes == 1:
# one-dimensional grids support various implementations of finite difference
x = (data_full[2:] - data_full[:-2]) * self.scale[0]
return x[None, :] # type: ignore
if self.grid.num_axes == 2:
# two-dimensional grids only support central differences
x = (data_full[2:, 1:-1] - data_full[:-2, 1:-1]) * self.scale[0]
y = (data_full[1:-1, 2:] - data_full[1:-1, :-2]) * self.scale[1]
return torch.stack((x, y))
if self.grid.num_axes == 3:
# three-dimensional grids only support central differences
x = (data_full[2:, 1:-1, 1:-1] - data_full[:-2, 1:-1, 1:-1]) * self.scale[0]
y = (data_full[1:-1, 2:, 1:-1] - data_full[1:-1, :-2, 1:-1]) * self.scale[1]
z = (data_full[1:-1, 1:-1, 2:] - data_full[1:-1, 1:-1, :-2]) * self.scale[2]
return torch.stack((x, y, z))
raise NotImplementedError
[docs]
@TorchBackend.register_operator(
CartesianGrid, "gradient_squared", rank_in=0, rank_out=0
)
class CartesianGradientSquared(TorchDifferentialOperator):
"""Cartesian gradient squared operator using torch."""
rank_in = 0
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
central: bool = True,
dtype: np.dtype,
):
"""Initialize the Cartesian gradient squared operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
central (bool):
Whether to use central differences. If `False`, forward and backward
differences are used.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
self.central = central
if self.central:
self.grad_central = CartesianGradient(grid, bcs, dtype=dtype)
else:
self.scale = 0.5 / self.grid.discretization**2
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
if self.central:
# simple squared sum of central differences
return torch.sum(self.grad_central(arr) ** 2, dim=0)
# use forward and backward differences
data_full = self.get_full_data(arr, args=args)
if self.grid.num_axes == 1:
# use forward and backward differences
diff_l = (data_full[2:] - data_full[1:-1]) ** 2
diff_r = (data_full[1:-1] - data_full[:-2]) ** 2
return (diff_l + diff_r) * self.scale[0] # type: ignore
if self.grid.num_axes == 2:
# two-dimensional grids only support central differences
x = (
(data_full[2:, 1:-1] - data_full[1:-1, 1:-1]) ** 2
+ (data_full[1:-1, 1:-1] - data_full[:-2, 1:-1]) ** 2
) * self.scale[0]
y = (
(data_full[1:-1, 2:] - data_full[1:-1, 1:-1]) ** 2
+ (data_full[1:-1, 1:-1] - data_full[1:-1, :-2]) ** 2
) * self.scale[1]
return x + y # type: ignore
if self.grid.num_axes == 3:
# three-dimensional grids only support central differences
x = (
(data_full[2:, 1:-1, 1:-1] - data_full[1:-1, 1:-1, 1:-1]) ** 2
+ (data_full[1:-1, 1:-1, 1:-1] - data_full[:-2, 1:-1, 1:-1]) ** 2
) * self.scale[0]
y = (
(data_full[1:-1, 2:, 1:-1] - data_full[1:-1, 1:-1, 1:-1]) ** 2
+ (data_full[1:-1, 1:-1, 1:-1] - data_full[1:-1, :-2, 1:-1]) ** 2
) * self.scale[1]
z = (
(data_full[1:-1, 1:-1, 2:] - data_full[1:-1, 1:-1, 1:-1]) ** 2
+ (data_full[1:-1, 1:-1, 1:-1] - data_full[1:-1, 1:-1, :-2]) ** 2
) * self.scale[2]
return x + y + z # type: ignore
raise NotImplementedError
[docs]
@TorchBackend.register_operator(CartesianGrid, "divergence", rank_in=1, rank_out=0)
class CartesianDivergence(TorchDifferentialOperator):
"""Cartesian divergence operator using torch."""
rank_in = 1
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
dtype: np.dtype,
):
"""Initialize the Cartesian divergence operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
self.scale = 0.5 / self.grid.discretization
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
data_full = self.get_full_data(arr, args=args)
if self.grid.num_axes == 1:
return (data_full[0, 2:] - data_full[0, :-2]) * self.scale[0] # type: ignore
if self.grid.num_axes == 2:
# two-dimensional grids only support central differences
x = (data_full[0, 2:, 1:-1] - data_full[0, :-2, 1:-1]) * self.scale[0]
y = (data_full[1, 1:-1, 2:] - data_full[1, 1:-1, :-2]) * self.scale[1]
return x + y # type: ignore
if self.grid.num_axes == 3:
# three-dimensional grids only support central differences
x = data_full[0, 2:, 1:-1, 1:-1] - data_full[0, :-2, 1:-1, 1:-1]
y = data_full[1, 1:-1, 2:, 1:-1] - data_full[1, 1:-1, :-2, 1:-1]
z = data_full[2, 1:-1, 1:-1, 2:] - data_full[2, 1:-1, 1:-1, :-2]
return x * self.scale[0] + y * self.scale[1] + z * self.scale[2] # type: ignore
raise NotImplementedError
[docs]
@TorchBackend.register_operator(CartesianGrid, "vector_gradient", rank_in=1, rank_out=2)
class CartesianVectorGradient(TorchDifferentialOperator):
"""Cartesian vector gradient operator using torch."""
rank_in = 1
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
dtype: np.dtype,
):
"""Initialize the Cartesian vector gradient operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
# define gradient operator that does not define the boundary conditions
self.result_shape = (grid.dim, grid.dim, *grid.shape)
self.grad = CartesianGradient(grid, bcs=None, dtype=dtype)
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
# apply boundary conditions and get full data array
data_full = self.get_full_data(arr, args=args)
result = torch.empty(self.result_shape, dtype=arr.dtype, device=arr.device)
# apply differential operators for all dimensions
for i in range(self.grid.num_axes):
result[i] = self.grad(data_full[i], args=args)
return result
[docs]
@TorchBackend.register_operator(CartesianGrid, "vector_laplace", rank_in=1, rank_out=1)
class CartesianVectorLaplacian(TorchDifferentialOperator):
"""Cartesian vector Laplacian operator using torch."""
rank_in = 1
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
dtype: np.dtype,
):
"""Initialize the Cartesian vector Laplacian operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
# define Laplace operator that does not define the boundary conditions
self.result_shape = (grid.dim, *grid.shape)
self.lap = CartesianLaplacian(grid, bcs=None, dtype=dtype)
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
# apply boundary conditions and get full data array
data_full = self.get_full_data(arr, args=args)
result = torch.empty(self.result_shape, dtype=arr.dtype, device=arr.device)
# apply differential operators for all dimensions
for i in range(self.grid.num_axes):
result[i] = self.lap(data_full[i], args=args)
return result
[docs]
@TorchBackend.register_operator(
CartesianGrid, "tensor_divergence", rank_in=2, rank_out=1
)
class CartesianTensorDivergence(TorchDifferentialOperator):
"""Cartesian tensor divergence operator using torch."""
rank_in = 2
def __init__(
self,
grid: GridBase,
bcs: BoundariesList | None,
*,
dtype: np.dtype,
):
"""Initialize the Cartesian tensor divergence operator.
Args:
grid (:class:`~pde.grids.base.GridBase`):
The grid on which the operator acts
bcs (:class:`~pde.grids.boundaries.axes.BoundariesList` or None):
The boundary conditions applied to the field. If `None`, no boundary
conditions are enforced.
dtype:
The data type of the field
"""
super().__init__(grid, bcs, dtype=dtype)
# define divergence operator that does not define the boundary conditions
self.result_shape = (grid.dim, *grid.shape)
self.div = CartesianDivergence(grid, bcs=None, dtype=dtype)
[docs]
def forward(self, arr: Tensor, args=None) -> Tensor:
"""Fill internal data array, apply operator, and return valid data."""
# apply boundary conditions and get full data array
data_full = self.get_full_data(arr, args=args)
result = torch.empty(self.result_shape, dtype=arr.dtype, device=arr.device)
# apply differential operators for all dimensions
for i in range(self.grid.num_axes):
result[i] = self.div(data_full[i], args=args)
return result
# apply differential operators for all dimensions
return torch.stack(
tuple(self.div(data_full[i], args=args) for i in range(self.grid.num_axes))
)
__all__ = [
"CartesianDivergence",
"CartesianGradient",
"CartesianGradientSquared",
"CartesianLaplacian",
"CartesianTensorDivergence",
"CartesianVectorGradient",
"CartesianVectorLaplacian",
]