# Copyright 2014-2025 The ODL contributors
#
# This file is part of ODL.
#
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
# obtain one at https://mozilla.org/MPL/2.0/.
"""Operators defined for tensor fields."""
import numpy as np
from math import prod
from odl.core.discr.discr_space import DiscretizedSpace
from odl.core.operator.tensor_ops import PointwiseTensorFieldOperator
from odl.core.space import ProductSpace
from odl.core.util import indent, signature_string, writable_array
from odl.core.array_API_support import asarray, get_array_and_backend
import warnings
__all__ = ('PartialDerivative', 'Gradient', 'Divergence', 'Laplacian')
_SUPPORTED_DIFF_METHODS = ('central', 'forward', 'backward')
_SUPPORTED_PAD_MODES = ('constant',
'symmetric', 'symmetric_adjoint',
'periodic',
'order0', 'order0_adjoint',
'order1', 'order1_adjoint',
'order2', 'order2_adjoint')
_ADJ_METHOD = {'central': 'central',
'forward': 'backward',
'backward': 'forward'}
_ADJ_PADDING = {'constant': 'constant',
'symmetric': 'symmetric_adjoint',
'symmetric_adjoint': 'symmetric',
'periodic': 'periodic',
'order0': 'order0_adjoint',
'order0_adjoint': 'order0',
'order1': 'order1_adjoint',
'order1_adjoint': 'order1',
'order2': 'order2_adjoint',
'order2_adjoint': 'order2'}
class PartialDerivative(PointwiseTensorFieldOperator):
"""Calculate the discrete partial derivative along a given axis.
Calls helper function `finite_diff` to calculate finite difference.
Preserves the shape of the underlying grid.
"""
[docs]
def __init__(self, domain, axis, range=None, method='forward',
pad_mode='constant', pad_const=0):
"""Initialize a new instance.
Parameters
----------
domain : `DiscretizedSpace`
Space of elements on which the operator can act.
axis : int
Axis along which the partial derivative is evaluated.
range : `DiscretizedSpace`, optional
Space of elements to which the operator maps, must have the
same shape as ``domain``.
For the default ``None``, the range is the same as ``domain``.
method : {'forward', 'backward', 'central'}, optional
Finite difference method which is used in the interior of the
domain of ``f``.
pad_mode : string, optional
The padding mode to use outside the domain.
``'constant'``: Fill with ``pad_const``.
``'symmetric'``: Reflect at the boundaries, not doubling the
outmost values.
``'periodic'``: Fill in values from the other side, keeping
the order.
``'order0'``: Extend constantly with the outmost values
(ensures continuity).
``'order1'``: Extend with constant slope (ensures continuity of
the first derivative). This requires at least 2 values along
each axis where padding is applied.
``'order2'``: Extend with second order accuracy (ensures continuity
of the second derivative). This requires at least 3 values along
the ``axis`` where padding is applied.
pad_const : float, optional
For ``pad_mode == 'constant'``, ``f`` assumes
``pad_const`` for indices outside the domain of ``f``
Examples
--------
>>> f = np.array([[ 0., 1., 2., 3., 4.],
... [ 0., 2., 4., 6., 8.]])
>>> discr = odl.uniform_discr([0, 0], [2, 1], f.shape)
>>> par_deriv = PartialDerivative(discr, axis=0, pad_mode='order1')
>>> par_deriv(f)
uniform_discr([ 0., 0.], [ 2., 1.], (2, 5)).element(
[[ 0., 1., 2., 3., 4.],
[ 0., 1., 2., 3., 4.]]
)
"""
if not isinstance(domain, DiscretizedSpace):
raise TypeError(f"`domain` {domain} is not a DiscretizedSpace instance")
if range is None:
range = domain
# Method is affine if nonzero padding is given.
linear = not (pad_mode == 'constant' and pad_const != 0)
super().__init__(domain, range, base_space=domain, linear=linear)
self.axis = int(axis)
self.dx = self.domain.cell_sides[axis]
self.method, method_in = str(method).lower(), method
if method not in _SUPPORTED_DIFF_METHODS:
raise ValueError(f"`method` {method_in} not understood")
self.pad_mode, pad_mode_in = str(pad_mode).lower(), pad_mode
if pad_mode not in _SUPPORTED_PAD_MODES:
raise ValueError(f"`pad_mode` {pad_mode_in} not understood")
self.pad_const = self.domain.field.element(pad_const)
[docs]
def _call(self, x, out=None):
"""Calculate partial derivative of ``x``."""
if out is None:
out = self.range.element()
# TODO: this pipes CUDA arrays through NumPy. Write native operator.
with writable_array(out) as out_arr:
finite_diff(x.asarray(), axis=self.axis, dx=self.dx,
method=self.method, pad_mode=self.pad_mode,
pad_const=self.pad_const, out=out_arr)
return out
[docs]
def derivative(self, point=None):
"""Return the derivative operator.
The partial derivative is usually linear, but in case the 'constant'
``pad_mode`` is used with nonzero ``pad_const``, the
derivative is given by the derivative with 0 ``pad_const``.
Parameters
----------
point : `domain` `element-like`, optional
The point to take the derivative in. Does not change the result
since the operator is affine.
"""
if self.pad_mode == 'constant' and self.pad_const != 0:
return PartialDerivative(self.domain, self.axis, self.range,
self.method, self.pad_mode, 0)
else:
return self
@property
def adjoint(self):
"""Return the adjoint operator."""
if not self.is_linear:
raise ValueError(
f"operator with nonzero pad_const ({self.pad_const}) is not linear and has no adjoint"
)
return -PartialDerivative(self.range, self.axis, self.domain,
_ADJ_METHOD[self.method],
_ADJ_PADDING[self.pad_mode],
self.pad_const)
def __repr__(self):
"""Return ``repr(self)``."""
posargs = [self.domain]
optargs = [('axis', self.axis, None),
('range', self.range, self.domain),
('method', self.method, 'forward'),
('pad_mode', self.pad_mode, 'constant'),
('pad_const', self.pad_const, 0)]
inner_str = signature_string(posargs, optargs,
sep=',\n', mod=['!r', ''])
return f"{self.__class__.__name__}(\n{indent(inner_str)}\n)"
def __str__(self):
"""Return ``str(self)``."""
dom_ran_str = "\n-->\n".join([repr(self.domain), repr(self.range)])
return f"{self.__class__.__name__}:\n{indent(dom_ran_str)}"
class Gradient(PointwiseTensorFieldOperator):
"""Spatial gradient operator for `DiscretizedSpace` spaces.
Calls helper function `finite_diff` to calculate each component of the
resulting product space element. For the adjoint of the `Gradient`
operator, zero padding is assumed to match the negative `Divergence`
operator
"""
[docs]
def __init__(self, domain=None, range=None, method='forward',
pad_mode='constant', pad_const=0):
"""Initialize a new instance.
Zero padding is assumed for the adjoint of the `Gradient`
operator to match negative `Divergence` operator.
Parameters
----------
domain : `DiscretizedSpace`, optional
Space of elements which the operator acts on.
This is required if ``range`` is not given.
range : power space of `DiscretizedSpace`, optional
Space of elements to which the operator maps.
This is required if ``domain`` is not given.
method : {'forward', 'backward', 'central'}, optional
Finite difference method to be used.
pad_mode : string, optional
The padding mode to use outside the domain.
``'constant'``: Fill with ``pad_const``.
``'symmetric'``: Reflect at the boundaries, not doubling the
outmost values.
``'periodic'``: Fill in values from the other side, keeping
the order.
``'order0'``: Extend constantly with the outmost values
(ensures continuity).
``'order1'``: Extend with constant slope (ensures continuity of
the first derivative). This requires at least 2 values along
each axis where padding is applied.
``'order2'``: Extend with second order accuracy (ensures continuity
of the second derivative). This requires at least 3 values along
each axis.
pad_const : float, optional
For ``pad_mode == 'constant'``, ``f`` assumes
``pad_const`` for indices outside the domain of ``f``
Examples
--------
Creating a Gradient operator:
>>> dom = odl.uniform_discr([0, 0], [1, 1], (10, 20))
>>> ran = odl.ProductSpace(dom, dom.ndim) # 2-dimensional
>>> grad_op = Gradient(dom)
>>> grad_op.range == ran
True
>>> grad_op2 = Gradient(range=ran)
>>> grad_op2.domain == dom
True
>>> grad_op3 = Gradient(domain=dom, range=ran)
>>> grad_op3.domain == dom
True
>>> grad_op3.range == ran
True
Calling the operator:
>>> data = np.array([[ 0., 1., 2., 3., 4.],
... [ 0., 2., 4., 6., 8.]])
>>> discr = odl.uniform_discr([0, 0], [2, 5], data.shape)
>>> f = discr.element(data)
>>> grad = Gradient(discr)
>>> grad_f = grad(f)
>>> grad_f[0]
uniform_discr([ 0., 0.], [ 2., 5.], (2, 5)).element(
[[ 0., 1., 2., 3., 4.],
[ 0., -2., -4., -6., -8.]]
)
>>> grad_f[1]
uniform_discr([ 0., 0.], [ 2., 5.], (2, 5)).element(
[[ 1., 1., 1., 1., -4.],
[ 2., 2., 2., 2., -8.]]
)
Verify adjoint:
>>> g = grad.range.element((data, data ** 2))
>>> adj_g = grad.adjoint(g)
>>> adj_g
uniform_discr([ 0., 0.], [ 2., 5.], (2, 5)).element(
[[ -0., -2., -5., -8., -11.],
[ -0., -5., -14., -23., -32.]]
)
>>> g.inner(grad_f) / f.inner(adj_g)
1.0
"""
if domain is None and range is None:
raise ValueError("either `domain` or `range` must be specified")
if domain is None:
try:
domain = range[0]
except TypeError:
pass
if range is None:
range = ProductSpace(domain, domain.ndim)
# Check range first since `domain` may end up to be `None` in
# the case filtered out here (see above)
if not isinstance(range, ProductSpace):
raise TypeError(f"`range` {range} is not a `ProductSpace` instance")
if not range.is_power_space:
raise ValueError(f"`range` {range} is not a power space")
if not isinstance(domain, DiscretizedSpace):
raise TypeError(f"`domain` {domain} is not a `DiscretizedSpace` instance")
if len(range) != domain.ndim:
raise ValueError(
f"`range` must be a power space of length n = {domain.ndim},with `n == domain.ndim`, got n = {len(range)} instead"
)
linear = not (pad_mode == 'constant' and pad_const != 0)
super().__init__(domain, range, base_space=domain, linear=linear)
self.method, method_in = str(method).lower(), method
if method not in _SUPPORTED_DIFF_METHODS:
raise ValueError(f"`method` {method_in} not understood")
self.pad_mode, pad_mode_in = str(pad_mode).lower(), pad_mode
if pad_mode not in _SUPPORTED_PAD_MODES:
raise ValueError(f"`pad_mode` {pad_mode_in} not understood")
self.pad_const = domain.field.element(pad_const)
[docs]
def _call(self, x, out=None):
"""Calculate the spatial gradient of ``x``."""
x_arr = x.asarray()
ndim = self.domain.ndim
dx = self.domain.cell_sides
if out is None:
return self.range.element([
finite_diff(x_arr, axis=axis, dx=dx[axis], method=self.method,
pad_mode=self.pad_mode,
pad_const=self.pad_const,
)
for axis in range(ndim)])
else:
for axis in range(ndim):
with writable_array(out[axis]) as out_arr:
finite_diff(x_arr, axis=axis, dx=dx[axis], method=self.method,
pad_mode=self.pad_mode,
pad_const=self.pad_const,
out=out_arr)
return out
[docs]
def derivative(self, point=None):
"""Return the derivative operator.
The gradient is usually linear, but in case the 'constant'
``pad_mode`` is used with nonzero ``pad_const``, the
derivative is given by the Gradient with ``pad_const=0``.
Parameters
----------
point : `domain` element, optional
The point to take the derivative in. Does not change the result
since the operator is affine.
"""
if self.pad_mode == 'constant' and self.pad_const != 0:
return Gradient(self.domain, self.range, self.method,
pad_mode=self.pad_mode,
pad_const=0)
else:
return self
@property
def adjoint(self):
"""Adjoint of this operator.
The adjoint is given by the negative `Divergence` with corrections for
the method and padding.
The Divergence is constructed from a ``space`` as a product space
operator ``space^n --> space``, hence we need to provide the domain of
this operator.
"""
if not self.is_linear:
raise ValueError(
f"operator with nonzero pad_const ({self.pad_const}) is not linear and has no adjoint"
)
return - Divergence(domain=self.range, range=self.domain,
method=_ADJ_METHOD[self.method],
pad_mode=_ADJ_PADDING[self.pad_mode],
pad_const=self.pad_const)
def __repr__(self):
"""Return ``repr(self)``."""
posargs = [self.domain]
optargs = [('range', self.range, self.domain ** self.domain.ndim),
('method', self.method, 'forward'),
('pad_mode', self.pad_mode, 'constant'),
('pad_const', self.pad_const, 0)]
inner_str = signature_string(posargs, optargs,
sep=[',\n', ', ', ',\n'],
mod=['!r', ''])
return f"{self.__class__.__name__}(\n{indent(inner_str)}\n)"
def __str__(self):
"""Return ``str(self)``."""
dom_ran_str = "\n-->\n".join([repr(self.domain), repr(self.range)])
return f"{self.__class__.__name__}:\n{indent(dom_ran_str)}"
class Divergence(PointwiseTensorFieldOperator):
"""Divergence operator for `DiscretizedSpace` spaces.
Calls helper function `finite_diff` for each component of the input
product space vector. For the adjoint of the `Divergence` operator to
match the negative `Gradient` operator implicit zero is assumed.
"""
[docs]
def __init__(self, domain=None, range=None, method='forward',
pad_mode='constant', pad_const=0):
"""Initialize a new instance.
Zero padding is assumed for the adjoint of the `Divergence`
operator to match the negative `Gradient` operator.
Parameters
----------
domain : power space of `DiscretizedSpace`, optional
Space of elements which the operator acts on.
This is required if ``range`` is not given.
range : `DiscretizedSpace`, optional
Space of elements to which the operator maps.
This is required if ``domain`` is not given.
method : {'forward', 'backward', 'central'}, optional
Finite difference method to be used
pad_mode : string, optional
The padding mode to use outside the domain.
``'constant'``: Fill with ``pad_const``.
``'symmetric'``: Reflect at the boundaries, not doubling the
``'periodic'``: Fill in values from the other side, keeping
the order.
``'order0'``: Extend constantly with the outmost values
(ensures continuity).
``'order1'``: Extend with constant slope (ensures continuity of
the first derivative). This requires at least 2 values along
each axis.
``'order2'``: Extend with second order accuracy (ensures continuity
of the second derivative). This requires at least 3 values along
each axis.
pad_const : float, optional
For ``pad_mode == 'constant'``, ``f`` assumes
``pad_const`` for indices outside the domain of ``f``
Examples
--------
Initialize a Divergence opeator:
>>> ran = odl.uniform_discr([0, 0], [3, 5], (3, 5))
>>> dom = odl.ProductSpace(ran, ran.ndim) # 2-dimensional
>>> div = Divergence(dom)
>>> div.range == ran
True
>>> div2 = Divergence(range=ran)
>>> div2.domain == dom
True
>>> div3 = Divergence(domain=dom, range=ran)
>>> div3.domain == dom
True
>>> div3.range == ran
True
Call the operator:
>>> data = np.array([[0., 1., 2., 3., 4.],
... [1., 2., 3., 4., 5.],
... [2., 3., 4., 5., 6.]])
>>> f = div.domain.element([data, data])
>>> div_f = div(f)
>>> print(div_f)
[[ 2., 2., 2., 2., -3.],
[ 2., 2., 2., 2., -4.],
[ -1., -2., -3., -4., -12.]]
Verify adjoint:
>>> g = div.range.element(data ** 2)
>>> adj_div_g = div.adjoint(g)
>>> g.inner(div_f) / f.inner(adj_div_g)
1.0
"""
if domain is None and range is None:
raise ValueError("either `domain` or `range` must be specified")
if domain is None:
domain = ProductSpace(range, range.ndim)
if range is None:
try:
range = domain[0]
except TypeError:
pass
# Check `domain` first since `range` may end up to be `None` in
# the case filtered out here (see above)
if not isinstance(domain, ProductSpace):
raise TypeError(f"`domain` {domain} is not a `ProductSpace` instance")
if not domain.is_power_space:
raise ValueError(f"`domain` {domain} is not a power space")
if not isinstance(range, DiscretizedSpace):
raise TypeError(f"`range` {range} is not a `DiscretizedSpace` " "instance")
if len(domain) != range.ndim:
raise ValueError(
f"`domain` must be a power space of length n = {range.ndim},with `n == range.ndim`, got n = {len(domain)} instead"
)
linear = not (pad_mode == 'constant' and pad_const != 0)
super().__init__(domain, range, base_space=range, linear=linear)
self.method, method_in = str(method).lower(), method
if method not in _SUPPORTED_DIFF_METHODS:
raise ValueError(f"`method` {method_in} not understood")
self.pad_mode, pad_mode_in = str(pad_mode).lower(), pad_mode
if pad_mode not in _SUPPORTED_PAD_MODES:
raise ValueError(f"`pad_mode` {pad_mode_in} not understood")
self.pad_const = range.field.element(pad_const)
[docs]
def _call(self, x, out=None):
"""Calculate the divergence of ``x``."""
ndim = self.range.ndim
dx = self.range.cell_sides
backend = self.range.array_backend
def directional_derivative(axis, out=None):
return finite_diff( x[axis], axis=axis, dx=dx[axis]
, method=self.method, pad_mode=self.pad_mode
, pad_const=self.pad_const
, out=out )
if out is None:
result = directional_derivative(0)
for axis in range(1,len(x)):
result += directional_derivative(axis)
return self.range.element(result)
else:
assert(backend.impl != 'pytorch')
tmp = self.range.element().asarray()
with writable_array(out) as out_arr:
for axis in range(ndim):
directional_derivative(axis, out=tmp)
if axis == 0:
out_arr[:] = tmp
else:
out_arr += tmp
return out
[docs]
def derivative(self, point=None):
"""Return the derivative operator.
The Divergence is usually linear, but in case the 'constant'
``pad_mode`` is used with nonzero ``pad_const``, the
derivative is given by the Divergence with ``pad_const=0``.
Parameters
----------
point : `domain` element, optional
The point to take the derivative in. Does not change the result
since the operator is affine.
"""
if self.pad_mode == 'constant' and self.pad_const != 0:
return Divergence(self.domain, self.range, self.method,
pad_mode=self.pad_mode, pad_const=0)
else:
return self
@property
def adjoint(self):
"""Adjoint of this operator.
The adjoint is given by the negative `Gradient` with corrections for
the method and padding.
"""
if not self.is_linear:
raise ValueError(
f"operator with nonzero pad_const ({self.pad_const}) is not linear and has no adjoint"
)
return - Gradient(self.range, self.domain,
method=_ADJ_METHOD[self.method],
pad_mode=_ADJ_PADDING[self.pad_mode])
def __repr__(self):
"""Return ``repr(self)``."""
posargs = [self.domain]
optargs = [('range', self.range, self.domain[0]),
('method', self.method, 'forward'),
('pad_mode', self.pad_mode, 'constant'),
('pad_const', self.pad_const, 0)]
inner_str = signature_string(posargs, optargs,
sep=[',\n', ', ', ',\n'],
mod=['!r', ''])
return f"{self.__class__.__name__}(\n{indent(inner_str)}\n)"
def __str__(self):
"""Return ``str(self)``."""
dom_ran_str = "\n-->\n".join([repr(self.domain), repr(self.range)])
return f"{self.__class__.__name__}:\n{indent(dom_ran_str)}"
class Laplacian(PointwiseTensorFieldOperator):
"""Spatial Laplacian operator for `DiscretizedSpace` spaces.
Calls helper function `finite_diff` to calculate each component of the
resulting product space vector.
Outside the domain zero padding is assumed.
"""
[docs]
def __init__(self, domain, range=None, pad_mode='constant', pad_const=0):
"""Initialize a new instance.
Parameters
----------
domain : `DiscretizedSpace`
Space of elements which the operator is acting on.
pad_mode : string, optional
The padding mode to use outside the domain.
``'constant'``: Fill with ``pad_const``.
``'symmetric'``: Reflect at the boundaries, not doubling the
outmost values.
``'periodic'``: Fill in values from the other side, keeping
the order.
``'order0'``: Extend constantly with the outmost values
(ensures continuity).
``'order1'``: Extend with constant slope (ensures continuity of
the first derivative). This requires at least 2 values along
each axis.
``'order2'``: Extend with second order accuracy (ensures continuity
of the second derivative). This requires at least 3 values along
each axis where padding is applied.
pad_const : float, optional
For ``pad_mode == 'constant'``, ``f`` assumes
``pad_const`` for indices outside the domain of ``f``
Examples
--------
>>> data = np.array([[ 0., 0., 0.],
... [ 0., 1., 0.],
... [ 0., 0., 0.]])
>>> space = odl.uniform_discr([0, 0], [3, 3], [3, 3])
>>> f = space.element(data)
>>> lap = Laplacian(space)
>>> lap(f)
uniform_discr([ 0., 0.], [ 3., 3.], (3, 3)).element(
[[ 0., 1., 0.],
[ 1., -4., 1.],
[ 0., 1., 0.]]
)
"""
if not isinstance(domain, DiscretizedSpace):
raise TypeError(f"`domain` {domain} is not a DiscretizedSpace instance")
if range is None:
range = domain
super().__init__(domain, range, base_space=domain, linear=True)
self.pad_mode, pad_mode_in = str(pad_mode).lower(), pad_mode
if pad_mode not in _SUPPORTED_PAD_MODES:
raise ValueError(f"`pad_mode` {pad_mode_in} not understood")
if pad_mode in ('order1', 'order1_adjoint',
'order2', 'order2_adjoint'):
# TODO: Add these pad modes
raise ValueError(f"`pad_mode` {pad_mode_in} not implemented for Laplacian.")
self.pad_const = self.domain.field.element(pad_const)
[docs]
def _call(self, x, out=None):
"""Calculate the spatial Laplacian of ``x``."""
if out is None:
out = self.range.zero()
else:
out.set_zero()
out_arr = out.asarray()
x_arr, backend = get_array_and_backend(x)
tmp = backend.array_namespace.empty(
shape=out.shape, dtype=out.dtype,device=x.device
)
ndim = self.domain.ndim
dx = self.domain.cell_sides
with writable_array(out) as out_arr:
for axis in range(ndim):
# TODO: this can be optimized
finite_diff(x_arr, axis=axis, dx=dx[axis] ** 2,
method='forward',
pad_mode=self.pad_mode,
pad_const=self.pad_const, out=tmp)
out_arr += tmp
finite_diff(x_arr, axis=axis, dx=dx[axis] ** 2,
method='backward',
pad_mode=self.pad_mode,
pad_const=self.pad_const, out=tmp)
out_arr -= tmp
return out
[docs]
def derivative(self, point=None):
"""Return the derivative operator.
The Laplacian is usually linear, but in case the 'constant'
``pad_mode`` is used with nonzero ``pad_const``, the
derivative is given by the derivative with 0 ``pad_const``.
Parameters
----------
point : ``domain`` element, optional
The point to take the derivative in. Does not change the result
since the operator is affine.
"""
if self.pad_mode == 'constant' and self.pad_const != 0:
return Laplacian(self.domain, self.range,
pad_mode=self.pad_mode, pad_const=0)
else:
return self
@property
def adjoint(self):
"""Return the adjoint operator.
The laplacian is self-adjoint, so this returns ``self``.
"""
return Laplacian(self.range, self.domain,
pad_mode=self.pad_mode, pad_const=0)
def __repr__(self):
"""Return ``repr(self)``."""
posargs = [self.domain]
optargs = [('range', self.range, self.domain ** self.domain.ndim),
('pad_mode', self.pad_mode, 'constant'),
('pad_const', self.pad_const, 0)]
inner_str = signature_string(posargs, optargs,
sep=[',\n', ', ', ',\n'],
mod=['!r', ''])
return f"{self.__class__.__name__}(\n{indent(inner_str)}\n)"
def __str__(self):
"""Return ``str(self)``."""
dom_ran_str = "\n-->\n".join([repr(self.domain), repr(self.range)])
return f"{self.__class__.__name__}:\n{indent(dom_ran_str)}"
def _finite_diff_pure(f_arr, axis, dx=1.0, method='forward',
pad_mode='constant', pad_const=0):
""" Purely-functional version of `finite_diff`, using only primitives from
the Python Array API. """
ndim = f_arr.ndim
if f_arr.shape[axis] < 2:
raise ValueError(
f"in axis {axis}: at least two elements required, got {f_arr.shape[axis]}"
)
if axis < 0:
axis += ndim
if not (0 <= axis < ndim):
raise IndexError(f"`axis` {axis} outside the valid range 0 ... {ndim - 1}")
dx, dx_in = float(dx), dx
if dx <= 0 or not np.isfinite(dx):
raise ValueError(f"`dx` must be positive, got {dx_in}")
method, method_in = str(method).lower(), method
if method not in _SUPPORTED_DIFF_METHODS:
raise ValueError(f"`method` {method_in} was not understood")
if pad_mode not in _SUPPORTED_PAD_MODES:
raise ValueError(f"`pad_mode` {pad_mode} not understood")
f_arr, backend = get_array_and_backend(f_arr)
namespace = backend.array_namespace
device = f_arr.device
orig_shape = f_arr.shape
if orig_shape[axis] < 2 and pad_mode == "order1":
raise ValueError(
f"size of array to small to use 'order1', needs at least 2 elements along axis {axis}."
)
if orig_shape[axis] < 3 and pad_mode == "order2":
raise ValueError(
f"size of array to small to use 'order2', needs at least 3 elements along axis {axis}."
)
# Reshape (in O(1)), so the axis of interest is the middle, all previous
# axes are flattened into the batch dimension, and all subsequent axes flattened
# into the final dimension. This allows a batched summing and indexing in only
# one of the dimension, regardless of which axis is actually worked on.
f_arr = f_arr.reshape([ prod(orig_shape[:axis])
, orig_shape[axis]
, prod(orig_shape[axis+1:])
])
# Interior of the domain of f
if method == 'central':
interior = (f_arr[:,2:] - f_arr[:,:-2]) / 2
elif method == 'forward':
interior = f_arr[:,2:] - f_arr[:,1:-1]
elif method == 'backward':
interior = f_arr[:,1:-1] - f_arr[:,:-2]
corrections = {}
# Boundaries
if pad_mode == 'constant':
# Assume constant value c for indices outside the domain of ``f``
# With padding the method used on endpoints is the same as in the
# interior of the domain of f
if method == 'central':
bnd_left = (f_arr[:,1] - pad_const) / 2.0
bnd_right = (pad_const - f_arr[:,-2]) / 2.0
elif method == 'forward':
bnd_left = f_arr[:,1] - f_arr[:,0]
bnd_right = pad_const - f_arr[:,-1]
elif method == 'backward':
bnd_left = f_arr[:,0] - pad_const
bnd_right = f_arr[:,-1] - f_arr[:,-2]
elif pad_mode == 'symmetric':
# Values of f for indices outside the domain of f are replicates of
# the edge values
# With padding the method used on endpoints is the same as in the
# interior of the domain of f
if method == 'central':
bnd_left = (f_arr[:,1] - f_arr[:,0]) / 2.0
bnd_right = (f_arr[:,-1] - f_arr[:,-2]) / 2.0
elif method == 'forward':
bnd_left = f_arr[:,1] - f_arr[:,0]
bnd_right = namespace.zeros_like(f_arr[:,-1])
elif method == 'backward':
bnd_left = namespace.zeros_like(f_arr[:,0])
bnd_right = f_arr[:,-1] - f_arr[:,-2]
elif pad_mode == 'symmetric_adjoint':
# The adjoint case of symmetric
if method == 'central':
bnd_left = (f_arr[:,1] + f_arr[:,0]) / 2.0
bnd_right = (-f_arr[:,-1] - f_arr[:,-2]) / 2.0
elif method == 'forward':
bnd_left = f_arr[:,1]
bnd_right = -f_arr[:,-1]
elif method == 'backward':
bnd_left = f_arr[:,0]
bnd_right = -f_arr[:,-2]
elif pad_mode == 'periodic':
# Values of f for indices outside the domain of f are replicates of
# the edge values on the other side
if method == 'central':
bnd_left = (f_arr[:,1] - f_arr[:,-1]) / 2.0
bnd_right = (f_arr[:,0] - f_arr[:,-2]) / 2.0
elif method == 'forward':
bnd_left = f_arr[:,1] - f_arr[:,0]
bnd_right = f_arr[:,0] - f_arr[:,-1]
elif method == 'backward':
bnd_left = f_arr[:,0] - f_arr[:,-1]
bnd_right = f_arr[:,-1] - f_arr[:,-2]
elif pad_mode == 'order0':
# Values of f for indices outside the domain of f are replicates of
# the edge value.
if method == 'central':
bnd_left = (f_arr[:,1] - f_arr[:,0]) / 2.0
bnd_right = (f_arr[:,-1] - f_arr[:,-2]) / 2.0
elif method == 'forward':
bnd_left = f_arr[:,1] - f_arr[:,0]
bnd_right = namespace.zeros_like(f_arr[:,-1])
elif method == 'backward':
bnd_left = namespace.zeros_like(f_arr[:,0])
bnd_right = f_arr[:,-1] - f_arr[:,-2]
elif pad_mode == 'order0_adjoint':
# Values of f for indices outside the domain of f are replicates of
# the edge value.
if method == 'central':
bnd_left = (f_arr[:,0] + f_arr[:,1]) / 2.0
bnd_right = -(f_arr[:,-1] + f_arr[:,-2]) / 2.0
elif method == 'forward':
bnd_left = f_arr[:,1]
bnd_right = -f_arr[:,-1]
elif method == 'backward':
bnd_left = f_arr[:,0]
bnd_right = -f_arr[:,-2]
elif pad_mode == 'order1':
# Values of f for indices outside the domain of f are linearly
# extrapolated from the inside.
# independent of ``method``
bnd_left = f_arr[:,1] - f_arr[:,0]
bnd_right = f_arr[:,-1] - f_arr[:,-2]
elif pad_mode == 'order1_adjoint':
# Values of f for indices outside the domain of f are linearly
# extrapolated from the inside.
if method == 'central':
bnd_left = f_arr[:,0] + f_arr[:,1] / 2.0
bnd_right = -f_arr[:,-1] - f_arr[:,-2] / 2.0
corrections[1] = -f_arr[:,0] / 2.0
corrections[-2] = f_arr[:,-1] / 2.0
elif method == 'forward':
bnd_left = f_arr[:,0] + f_arr[:,1]
bnd_right = -f_arr[:,-1]
corrections[1] = -f_arr[:,0]
elif method == 'backward':
bnd_left = f_arr[:,0]
bnd_right = -f_arr[:,-1] - f_arr[:,-2]
corrections[-2] = f_arr[:,-1]
elif pad_mode == 'order2':
# 2nd order edges
bnd_left = -(3.0 * f_arr[:,0] - 4.0 * f_arr[:,1] + f_arr[:,2]) / 2.0
bnd_right = (3.0 * f_arr[:,-1] - 4.0 * f_arr[:,-2] + f_arr[:,-3]) / 2.0
elif pad_mode == 'order2_adjoint':
# Values of f for indices outside the domain of f are quadratically
# extrapolated from the inside.
if method == 'central':
bnd_left = 1.5 * f_arr[:,0] + 0.5 * f_arr[:,1]
bnd_right = -1.5 * f_arr[:,-1] - 0.5 * f_arr[:,-2]
corrections[1] = -1.5 * f_arr[:,0]
corrections[2] = +0.5 * f_arr[:,0]
corrections[-3] = -0.5 * f_arr[:,-1]
corrections[-2] = +1.5 * f_arr[:,-1]
elif method == 'forward':
bnd_left = 1.5 * f_arr[:,0] + 1.0 * f_arr[:,1]
bnd_right = -1.5 * f_arr[:,-1]
corrections[1] = -2.0 * f_arr[:,0]
corrections[2] = 0.5 * f_arr[:,0]
corrections[-3] = -0.5 * f_arr[:,-1]
corrections[-2] = 1.0 * f_arr[:,-1]
elif method == 'backward':
bnd_left = 1.5 * f_arr[:,0]
bnd_right = -1.0 * f_arr[:,-2] - 1.5 * f_arr[:,-1]
corrections[1] = -1.0 * f_arr[:,0]
corrections[2] = 0.5 * f_arr[:,0]
corrections[-3] = -0.5 * f_arr[:,-1]
corrections[-2] = 2.0 * f_arr[:,-1]
else:
raise NotImplementedError("Unknown pad_mode")
if corrections:
uncorrected_start = max(corrections.keys()) + 1
if uncorrected_start < 1:
uncorrected_start = 1
uncorrected_end = min(corrections.keys())
if uncorrected_end >= 0:
uncorrected_end = -1
dummy_cols = namespace.zeros((f_arr.shape[0], 0, f_arr.shape[2]), device=device)
correct_columns_l = namespace.stack(
[interior[:,i-1] + corrections[i]
for i in range(1, uncorrected_start)],
axis = 1) if uncorrected_start > 1 else dummy_cols
correct_columns_r = namespace.stack(
[interior[:,i+1] + corrections[i]
for i in range(uncorrected_end, -1)],
axis = 1) if uncorrected_end < -1 else dummy_cols
unaffected_width = f_arr.shape[1] + uncorrected_end - uncorrected_start
if uncorrected_start <= 1:
relevant_interior = interior[:, :uncorrected_end+1]
elif uncorrected_end >= -1:
relevant_interior = interior[:, uncorrected_start-1:]
elif unaffected_width >= 0:
relevant_interior = interior[:, uncorrected_start-1 : uncorrected_end+1]
else:
warnings.warn("Very short array, so the corrections from the left and right boundary"
+" overlap. This should generally not happen in practice (it indicates"
+" that the resolution is so low that the whole premise of approximating"
+" derivatives by finite differences fails)")
overlap = -unaffected_width
relevant_interior = (
correct_columns_l[:,-overlap:]
+ correct_columns_r[:,:overlap]
- interior[:, uncorrected_end+1 : uncorrected_start-1])
correct_columns_l = correct_columns_l[:,:-overlap]
correct_columns_r = correct_columns_r[:,overlap:]
result = namespace.concat( [ bnd_left[:,None],
correct_columns_l,
relevant_interior,
correct_columns_r,
bnd_right[:,None] ],
axis = 1
) / dx
else:
result = namespace.concat([bnd_left[:,None], interior, bnd_right[:,None]], axis = 1) / dx
return result.reshape(orig_shape)
[docs]
def finite_diff(f, axis, dx=1.0, method='forward', out=None,
pad_mode='constant', pad_const=0):
"""Calculate the partial derivative of ``f`` along a given ``axis``.
In the interior of the domain of f, the partial derivative is computed
using first-order accurate forward or backward difference or
second-order accurate central differences.
With padding the same method and thus accuracy is used on endpoints as
in the interior i.e. forward and backward differences use first-order
accuracy on edges while central differences use second-order accuracy at
edges.
Without padding one-sided forward or backward differences are used at
the boundaries. The accuracy at the endpoints can then also be
triggered by the edge order.
The returned array has the same shape as the input array ``f``.
Per default forward difference with dx=1 and no padding is used.
Parameters
----------
f : `array-like`
An N-dimensional array.
axis : int
The axis along which the partial derivative is evaluated.
dx : float, optional
Scalar specifying the distance between sampling points along ``axis``.
method : {'central', 'forward', 'backward'}, optional
Finite difference method which is used in the interior of the domain
of ``f``.
out : `numpy.ndarray`, optional
An N-dimensional array to which the output is written. Has to have
the same shape as the input array ``f``.
pad_mode : string, optional
The padding mode to use outside the domain.
``'constant'``: Fill with ``pad_const``.
``'symmetric'``: Reflect at the boundaries, not doubling the
outmost values.
``'periodic'``: Fill in values from the other side, keeping
the order.
``'order0'``: Extend constantly with the outmost values
(ensures continuity).
``'order1'``: Extend with constant slope (ensures continuity of
the first derivative). This requires at least 2 values along
each axis where padding is applied.
``'order2'``: Extend with second order accuracy (ensures continuity
of the second derivative). This requires at least 3 values along
each axis where padding is applied.
pad_const : float, optional
For ``pad_mode == 'constant'``, ``f`` assumes ``pad_const`` for
indices outside the domain of ``f``
Returns
-------
out : `numpy.ndarray`
N-dimensional array of the same shape as ``f``. If ``out`` was
provided, the returned object is a reference to it.
Examples
--------
>>> f = np.array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
>>> finite_diff(f, axis=0)
array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., -9.])
Without arguments the above defaults to:
>>> finite_diff(f, axis=0, dx=1.0, method='forward', pad_mode='constant')
array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., -9.])
Parameters can be changed one by one:
>>> finite_diff(f, axis=0, dx=0.5)
array([ 2., 2., 2., 2., 2., 2., 2., 2., 2., -18.])
>>> finite_diff(f, axis=0, pad_mode='order1')
array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.])
Central differences and different edge orders:
>>> finite_diff(0.5 * f ** 2, axis=0, method='central', pad_mode='order1')
array([ 0.5, 1. , 2. , 3. , 4. , 5. , 6. , 7. , 8. , 8.5])
>>> finite_diff(0.5 * f ** 2, axis=0, method='central', pad_mode='order2')
array([-0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])
In-place evaluation:
>>> out = f.copy()
>>> out is finite_diff(f, axis=0, out=out)
True
"""
_, backend = get_array_and_backend(f)
if out is None:
return _finite_diff_pure(
f, axis, dx=dx, method=method, pad_mode=pad_mode, pad_const=pad_const)
else:
if out.shape != f.shape:
raise ValueError(f"expected output shape {f.shape}, got {out.shape}")
out[:] = _finite_diff_pure(
f, axis, dx=dx, method=method, pad_mode=pad_mode, pad_const=pad_const)
return out
if __name__ == '__main__':
from odl.core.util.testutils import run_doctests
run_doctests()