4.1 pde.backends package

Defines backends, which implement efficient numerical simulations.

Backends are classes that provide logic to carry out numerical calculations. In particular, each of these classes implements various operators. In principle, backend classes can be defined independently, but we use some inheritance to share logic.

Moreover, we provide a BackendRegistry, which allows selecting backends by their identifier, so users do not usually need to construct backend classes. Users should access the registry via the function get_backend() to load a backend in their code.

BackendRegistry	Class handling all backends and their configurations.
JaxBackend	Defines jax backend.
NumbaBackend	Defines numba backend.
NumbaMPIBackend	Defines MPI-compatible numba backend.
NumpyBackend	Defines numpy backend, from which all other backends inherit.
ScipyBackend	Defines scipy backend.
TorchBackend	Defines torch backend.

Inheritance structure of the classes:

class BackendBase(config, *, name=None)

Bases: Generic[TNativeArray]

Basic backend from which all other backends inherit.

The generic type parameter TNativeArray determines the type of the native data representation of the backend.

Initialize the backend.

Parameters:

• config (Config) – Configuration data for the backend

• name (str) – The name of the backend

compile_function(func, **kwargs)

General method that compiles a user function.

Parameters:

• func (callable) – The function that needs to be compiled for this backend

• **kwargs – Additional arguments affecting how the function is compiled

Return type:

TFunc

config: Config

Configuration options of this backend.

Type:

dict

config_inheritance: list[str] = []

Additional backends that are queried for configuration parameters.

Type:

list

copy_data: bool = False

Data must be copied from CPU numpy representation to a native device.

Type:

bool

classmethod from_args(config, args='', *, name=None)

Initialize backend with extra arguments.

Parameters:

• config (Config) – Configuration data for the backend

• args (str) – Additional arguments that determine how the backend is initialized

• name (str) – The name of the backend

get_operator_info(grid, operator)

Return an operator for a particular grid.

Parameters:

• grid (GridBase) – Grid for which the operator is needed

• operator (str) – Identifier for the operator. Some examples are ‘laplace’, ‘gradient’, or ‘divergence’. The registered operators for this grid can be obtained from the operators attribute.

Returns:

information for the operator

Return type:

OperatorInfo

get_registered_operators(grid_id)

Returns all operators defined for a grid.

Parameters:

grid_id (GridBase or its type) – Grid or grid class for which the operators need to be returned

Return type:

set[str]

implementation: str = 'undefined'

The name of the python module that is used to implement this backend. This information can be used to distinguish the general implementation of backends.

Type:

str

property info: dict[str, Any]

relevant information about the backend

Type:

dict

make_data_setter(grid, rank, bcs=None)

Create a function to set the valid part of a full data array.

Parameters:

• grid (GridBase) – Grid for which the data setter is defined

• rank (int) – Rank of the data represented on the grid

• bcs (BoundariesBase, optional) – If supplied, the returned function also enforces boundary conditions by setting the ghost cells to the correct values

Returns:

Takes two numpy arrays, setting the valid data in the first one, using the second array. The arrays need to be allocated already and they need to have the correct dimensions, which are not checked. If bcs are given, a third argument is allowed, which sets arguments for the BCs.

Return type:

callable

make_expression_function(expression, *, single_arg=False, user_funcs=None)

Return a function evaluating an expression.

Parameters:

• expression (ExpressionBase) – The expression that is converted to a function

• single_arg (bool) – Determines whether the returned function accepts all variables in a single argument as an array or whether all variables need to be supplied separately.

• user_funcs (dict) – Additional functions that can be used in the expression.

Returns:

the function

Return type:

function

make_full_data_setter(bcs)

Create a function to set the valid part of a full data array.

Parameters:

bcs (BoundariesBase, optional) – If supplied, the returned function also enforces boundary conditions by setting the ghost cells to the correct values

Returns:

Takes two numpy arrays, setting the valid data in the first one, using the second array. The arrays need to be allocated already and they need to have the correct dimensions, which are not checked. If bcs are given, a third argument is allowed, which sets arguments for the BCs.

Return type:

callable

make_gaussian_noise(field, *, rng)

Create a function generating Gaussian white noise.

Parameters:

• field (FieldBase) – An example for the field from which the grid and other information can be extracted

• rng (Generator) – Random number generator (default: default_rng()).

Return type:

Callable[[], TNativeArray]

make_ghost_cell_setter(bcs)

Return function that sets the ghost cells on a full array.

Parameters:

bcs (BoundariesBase) – Defines the boundary conditions for a particular grid, for which the setter should be defined.

Returns:

Callable with signature (data_full: NumericArray, args=None), which sets the ghost cells of the full data, potentially using additional information in args (e.g., the time t during solving a PDE)

Return type:

GhostCellSetter

make_inner_prod_operator(field, *, conjugate=True)

Return operator calculating the dot product between two fields.

This supports both products between two vectors as well as products between a vector and a tensor.

Parameters:

• field (DataFieldBase) – Field for which the inner product is defined

• conjugate (bool) – Whether to use the complex conjugate for the second operand

Returns:

Function that takes two instance of native data arrays, which contain the discretized data of the two operands. An optional third argument can specify the output array to which the result is written.

Return type:

BinaryOperatorImplType

make_integrator(grid)

Return function that integrates discretized data over a grid.

Note that this function takes and returns data in the native representation of the backend. If this function is used in a multiprocessing run (using MPI), the integrals are performed on all subgrids and then accumulated. Each process then receives the same value representing the global integral.

Parameters:

grid (GridBase) – Grid for which the operator is needed

Returns:

A function that takes a numpy array and returns the integral with the correct weights given by the cell volumes.

Return type:

Callable[[TNativeArray], TNativeArray]

make_interpolator(field, *, fill=None, with_ghost_cells=False)

Returns a function that can be used to interpolate values.

Parameters:

• field (DataFieldBase) – Field for which the interpolator is defined

• fill (Number, optional) – Determines how values out of bounds are handled. If None, a ValueError is raised when out-of-bounds points are requested. Otherwise, the given value is returned.

• with_ghost_cells (bool) – Flag indicating that the interpolator should work on the full data array that includes values for the ghost points. If this is the case, the boundaries are not checked and the coordinates are used as is.

Returns:

A function which returns interpolated values when called with arbitrary positions within the space of the grid.

Return type:

Callable[[FloatingArray, NumericArray], NumberOrArray]

make_mpi_synchronizer(operator='MAX', mpi_run=False)

Return function that synchronizes values between multiple MPI processes.

Warning

The default implementation does not synchronize anything. This is simply a hook, which can be used by backends that support MPI

Parameters:

• operator (str or int) – Flag determining how the value from multiple nodes is combined. Possible values include “MAX”, “MIN”, and “SUM”.

• mpi_run (bool) – Whether MPI is actually used. If False, the method returns a no-op.

Returns:

Function that can be used to synchronize values across nodes

Return type:

Callable[[float], float]

make_operator(grid, operator, *, bcs, dtype=None, **kwargs)

Return a compiled function applying an operator with boundary conditions.

Parameters:

• grid (GridBase) – Grid for which the operator is needed

• operator (str) – Identifier for the operator. Some examples are ‘laplace’, ‘gradient’, or ‘divergence’. The registered operators for this grid can be obtained from the operators attribute.

• bcs (BoundariesBase) – The boundary conditions used before the operator is applied

• dtype (numpy dtype) – The data type of the field.

• **kwargs – Specifies extra arguments influencing how the operator is created.

Return type:

OperatorType

The returned function takes the discretized data on the grid as an input and returns the data to which the operator operator has been applied. The function only takes the valid grid points and allocates memory for the ghost points internally to apply the boundary conditions specified as bc. Note that the function supports an optional argument out, which if given should provide space for the valid output array without the ghost cells. The result of the operator is then written into this output array.

The function also accepts an optional parameter args, which is forwarded to set_ghost_cells. This allows setting boundary conditions based on external parameters, like time.

Returns:

the function that applies the operator. This function has the signature (arr: NumericArray, out: NumericArray = None, args=None).

Return type:

callable

Parameters:

• grid (GridBase)

• operator (str | OperatorInfo)

• bcs (BoundariesBase)

• dtype (DTypeLike | None)

make_operator_no_bc(grid, operator, *, dtype=None, **kwargs)

Return a compiled function applying an operator without boundary conditions.

A function that takes the discretized full data as an input and an array of valid data points to which the result of applying the operator is written.

Note

The resulting function does not check whether the ghost cells of the input array have been supplied with sensible values. It is the responsibility of the user to set the values of the ghost cells beforehand. Use this function only if you absolutely know what you’re doing. In all other cases, make_operator() is probably the better choice.

Parameters:

• grid (GridBase) – Grid for which the operator is needed

• operator (str) – Identifier for the operator. Some examples are ‘laplace’, ‘gradient’, or ‘divergence’. The registered operators for this grid can be obtained from the operators attribute.

• dtype (numpy dtype) – The data type of the field.

• **kwargs – Specifies extra arguments influencing how the operator is created.

Returns:

the function that applies the operator. This function has the signature (arr: NumericArray, out: NumericArray), so they out array need to be supplied explicitly.

Return type:

callable

make_outer_prod_operator(field)

Return operator calculating the outer product between two fields.

This supports typically only supports products between two vector fields.

Parameters:

field (DataFieldBase) – Field for which the outer product is defined

Returns:

Function that takes two instance of native data arrays, which contain the discretized data of the two operands. An optional third argument can specify the output array to which the result is written.

Return type:

BinaryOperatorImplType

make_pde_rhs(eq, state)

Return a function for evaluating the right hand side of the PDE.

Parameters:

• eq (PDEBase) – The object describing the differential equation

• state (FieldBase) – An example for the state from which information can be extracted

Returns:

Function returning deterministic part of the right hand side of the PDE

Return type:

Callable[[TNativeArray, float], TNativeArray]

make_stepper(solver, state)

Create a field-based stepping function for a given solver.

Parameters:

• solver (SolverBase) – The solver instance, which determines how the stepper is constructed

• state (FieldBase) – An example for the state from which the grid and other information can be extracted

Returns:

Function that can be called to advance the state from time t_start to time t_end.

Return type:

StepperType

make_valid_data_setter(grid, rank)

Create a function to set the valid part of a full data array.

Parameters:

• grid (GridBase) – Grid for which the data setter is defined

• rank (int) – Rank of the data represented on the grid

Returns:

Takes two numpy arrays, setting the valid data in the first one, using the second array. The arrays need to be allocated already and they need to have the correct dimensions, which are not checked. If bcs are given, a third argument is allowed, which sets arguments for the BCs.

Return type:

callable

native_to_numpy(value: TNativeArray) → ndarray[Any, dtype[number]]

native_to_numpy(value: TValue) → TValue

Convert native values to numpy representation.

Parameters:

value (Any) – The value to convert to numpy representation

Return type:

Any

numpy_to_native(value: ndarray[Any, dtype[number]]) → TNativeArray

numpy_to_native(value: TValue) → TValue

Convert values from numpy to native representation.

Parameters:

value (Any) – The value to convert from numpy representation

Return type:

Any

classmethod register_operator(grid_cls, name, factory_func=None, *, rank_in=0, rank_out=0)

Register an operator for a particular grid.

Example

The method can either be used directly:

BackendClass.register_operator(grid_class, "operator", make_operator)

or as a decorator for the factory function:

@BackendClass.register_operator(grid_class, "operator")
def make_operator(grid: GridBase): ...

Parameters:

• grid_cls (GridBase) – Grid class for which the operator is defined

• name (str) – The name of the operator to register

• factory_func (callable) – A function with signature (grid: GridBase, **kwargs), which takes a grid object and optional keyword arguments and returns an implementation of the given operator. This implementation is a function that takes a ndarray of discretized values as arguments and returns the resulting discretized data in a ndarray after applying the operator.

• rank_in (int) – The rank of the input field for the operator

• rank_out (int) – The rank of the field that is returned by the operator

supports_mpi: bool = False

This backend supports parallel simulations with MPI.

Type:

bool

get_backend(backend)

Return backend specified by string of instance.

Parameters:

backend (str or BackendBase) – Backend specified by name given as a string. If the string contains a colon, the first part determines the backend, whereas the second part can be used to convey additional information. For example, torch:cuda may load a torch backend and use a cuda device. As a special case, we also allow full backend objects, which are simply returned. This is a simple way to allow providing full backend objects in places where we otherwise would expect a backend name.

Return type:

BackendBase

registered_backends()

Returns all registered backends.

Return type:

list[str]

Subpackages:

• 4.1.1 pde.backends.jax package
  • 4.1.1.1 pde.backends.jax.operators package
    • 4.1.1.1.1 pde.backends.jax.operators.cartesian module
    • 4.1.1.1.2 pde.backends.jax.operators.cylindrical_sym module
    • 4.1.1.1.3 pde.backends.jax.operators.polar_sym module
    • 4.1.1.1.4 pde.backends.jax.operators.spherical_sym module
  • 4.1.1.2 pde.backends.jax.backend module
  • 4.1.1.3 pde.backends.jax.config module
  • 4.1.1.4 pde.backends.jax.typing module
• 4.1.2 pde.backends.numba package
  • 4.1.2.1 pde.backends.numba.operators package
    • 4.1.2.1.1 pde.backends.numba.operators.cartesian module
    • 4.1.2.1.2 pde.backends.numba.operators.common module
    • 4.1.2.1.3 pde.backends.numba.operators.cylindrical_sym module
    • 4.1.2.1.4 pde.backends.numba.operators.polar_sym module
    • 4.1.2.1.5 pde.backends.numba.operators.spherical_sym module
  • 4.1.2.2 pde.backends.numba.backend module
  • 4.1.2.3 pde.backends.numba.config module
  • 4.1.2.4 pde.backends.numba.grids module
  • 4.1.2.5 pde.backends.numba.overloads module
  • 4.1.2.6 pde.backends.numba.utils module
• 4.1.3 pde.backends.numba_mpi package
  • 4.1.3.1 pde.backends.numba_mpi.backend module
  • 4.1.3.2 pde.backends.numba_mpi.overloads module
• 4.1.4 pde.backends.numpy package
  • 4.1.4.1 pde.backends.numpy.backend module
• 4.1.5 pde.backends.scipy package
  • 4.1.5.1 pde.backends.scipy.operators package
    • 4.1.5.1.1 pde.backends.scipy.operators.cartesian module
    • 4.1.5.1.2 pde.backends.scipy.operators.common module
    • 4.1.5.1.3 pde.backends.scipy.operators.cylindrical_sym module
    • 4.1.5.1.4 pde.backends.scipy.operators.polar_sym module
    • 4.1.5.1.5 pde.backends.scipy.operators.spherical_sym module
  • 4.1.5.2 pde.backends.scipy.backend module
• 4.1.6 pde.backends.torch package
  • 4.1.6.1 pde.backends.torch.operators package
    • 4.1.6.1.1 pde.backends.torch.operators.cartesian module
    • 4.1.6.1.2 pde.backends.torch.operators.common module
    • 4.1.6.1.3 pde.backends.torch.operators.cylindrical_sym module
    • 4.1.6.1.4 pde.backends.torch.operators.polar_sym module
    • 4.1.6.1.5 pde.backends.torch.operators.spherical_sym module
  • 4.1.6.2 pde.backends.torch.backend module
  • 4.1.6.3 pde.backends.torch.config module
  • 4.1.6.4 pde.backends.torch.typing module
  • 4.1.6.5 pde.backends.torch.utils module

• 4.1.7 pde.backends.base module
• 4.1.8 pde.backends.registry module