4.1.5.2 pde.backends.torch.backend module
Defines base class of backends that implement computations.
- class TorchBackend(name, registry, *, device='config')[source]
Bases:
NumpyBackendDefines
torchbackend.Initialize the torch backend.
- Parameters:
registry (
BackendRegistry) – The registry to which this backend is addedname (str) – The name of the backend
device (str) – The torch device to use. Special values are “config” (read from configuration) and “auto” (use CUDA if available, otherwise CPU)
- compile_function(func, **compile_options)[source]
General method that compiles a user function.
- Parameters:
func (callable) – The function that needs to be compiled for this backend
**compile_options – Additional keyword arguments will be forwarded to
torch.compile()
- Return type:
TFunc
- compile_options = {'dynamic': False, 'fullgraph': True, 'options': {'epilogue_fusion': True, 'max_autotune': True}}
- make_expression_function(expression, *, single_arg=False, user_funcs=None)[source]
Return a function evaluating an expression.
- Parameters:
expression (
ExpressionBase) – The expression that is converted to a functionsingle_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_integrator(grid)[source]
Return function that integrates discretized data over a grid.
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 integrator is defined- Returns:
A function that takes a numpy array and returns the integral with the correct weights given by the cell volumes.
- Return type:
Callable[[NumericArray], NumberOrArray]
- make_operator(grid, operator, *, bcs, native=False, **kwargs)[source]
Return a torch function applying an operator with boundary conditions.
- Parameters:
grid (
GridBase) – Grid for which the operator is neededoperator (str) – Identifier for the operator. Some examples are ‘laplace’, ‘gradient’, or ‘divergence’. The registered operators for this grid can be obtained from the
operatorsattribute.bcs (
BoundariesBase, optional) – The boundary conditions used before the operator is appliednative (bool) – If True, the returned functions expects the native data representation of the backend. Otherwise, the input and output are expected to be
ndarray.**kwargs – Specifies extra arguments influencing how the operator is created.
- Return type:
Warning
The same operator should not be assigned to different variables that are used in the same code, because
torchhas problems compiling the resulting code. This particularly precludes caching the operators, since they then might be reused, e.g., if boundary conditions agree between different operators.- 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)
native (bool)
- make_operator_no_bc(grid, operator, *, native=False, **kwargs)[source]
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 neededoperator (str) – Identifier for the operator. Some examples are ‘laplace’, ‘gradient’, or ‘divergence’. The registered operators for this grid can be obtained from the
operatorsattribute.native (bool) – If True, the returned functions expects the native data representation of the backend. Otherwise, the input and output are expected to be
ndarray.**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