"""A simple diffusion equation.
.. codeauthor:: David Zwicker <david.zwicker@ds.mpg.de>
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from ..fields import ScalarField
from ..grids.boundaries import set_default_bc
from ..tools.docstrings import fill_in_docstring
from .base import SDEBase, expr_prod
if TYPE_CHECKING:
from collections.abc import Callable
import numpy as np
from ..backends import BackendBase
from ..grids.boundaries.axes import BoundariesData
from ..tools.typing import TNativeArray
[docs]
class DiffusionPDE(SDEBase):
r"""A simple diffusion equation.
The mathematical definition is
.. math::
\partial_t c = D \nabla^2 c
where :math:`c` is a scalar field and :math:`D` denotes the diffusivity.
"""
explicit_time_dependence = False
default_bc = "auto_periodic_neumann"
"""Default boundary condition used when no specific conditions are chosen."""
@fill_in_docstring
def __init__(
self,
diffusivity: float = 1,
*,
bc: BoundariesData | None = None,
noise: float = 0,
rng: np.random.Generator | None = None,
):
"""
Args:
diffusivity (float):
The diffusivity of the described species
bc:
The boundary conditions applied to the field.
{ARG_BOUNDARIES}
noise (float):
Variance of the (additive) noise term
rng (:class:`~numpy.random.Generator`):
Random number generator (default: :func:`~numpy.random.default_rng()`)
used for stochastic simulations. Note that this random number generator
is only used for numpy functions, while compiled numba code uses the
random number generator of numba. Moreover, in simulations using
multiprocessing, setting the same generator in all processes might yield
unintended correlations in the simulation results.
"""
super().__init__(noise=noise, rng=rng)
self.diffusivity = diffusivity
self.bc = set_default_bc(bc, self.default_bc)
@property
def expression(self) -> str:
"""str: the expression of the right hand side of this PDE"""
return expr_prod(self.diffusivity, "∇²(c)")
[docs]
def evolution_rate( # type: ignore
self,
state: ScalarField,
t: float = 0,
) -> ScalarField:
"""Evaluate the right hand side of the PDE.
Args:
state (:class:`~pde.fields.ScalarField`):
The scalar field describing the concentration distribution
t (float):
The current time point
Returns:
:class:`~pde.fields.ScalarField`:
Scalar field describing the evolution rate of the PDE
"""
if not isinstance(state, ScalarField):
msg = "`state` must be ScalarField"
raise TypeError(msg)
laplace = state.laplace(bc=self.bc, label="evolution rate", args={"t": t})
return self.diffusivity * laplace # type: ignore
[docs]
def make_evolution_rate(
self, state: ScalarField, backend: BackendBase
) -> Callable[[TNativeArray, float], TNativeArray]:
"""Create a compiled function evaluating the right hand side of the PDE.
Args:
state (:class:`~pde.fields.ScalarField`):
An example for the state defining the grid and data types
backend (str or :class:`~pde.backends.base.BackendBase`):
The backend used for numerical operations
Returns:
A function with signature `(state_data, t)`, which can be called with an
instance of the state data and time to obtain the associated evolution rate.
"""
diffusivity_value = self.diffusivity
laplace = state.grid.make_operator(
operator="laplace", bc=self.bc, backend=backend, dtype=state.dtype
)
def pde_rhs(state_data, t=0):
"""Evaluate right hand side of PDE."""
return diffusivity_value * laplace(state_data, args={"t": t})
return pde_rhs