4.1.6.2 pde.backends.torch.backend module

Defines base class of backends that implement computations.

class TorchBackend(config=None, *, name=None, device='config')[source]

Bases: BackendBase[Tensor]

Defines torch backend.

Initialize the torch backend.

Parameters:

  • config (Config) – Configuration data for the backend

  • name (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 = {'backend': 'inductor', 'dynamic': False, 'fullgraph': True, 'options': {'epilogue_fusion': True, 'max_autotune': True}}

defines options that affect compilation by torch

Type:

dict

copy_data = True

Flag indicating whether data needs to be copied between numpy’s representation on CPU and a native device.

Type:

bool

property device: device

The currently assigned torch device.

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

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_numpy_dtype(dtype)[source]

Determine numpy dtype suitable for the torch backend.

Parameters:

dtype (DTypeLike) – numpy dtype to convert to supported dtype

Returns:

A numpy dtype that is compatible with the torch backend

Return type:

torch.dtype

get_torch_dtype(dtype)[source]

Convert dtype to torch dtype.

Parameters:

dtype (DTypeLike) – numpy dtype to convert to corresponding torch dtype

Returns:

A proper dtype for torch

Return type:

torch.dtype

implementation = 'torch'

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_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 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_inner_prod_operator(field, *, conjugate=True)[source]

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 ndarray, 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:

Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None], torch.Tensor]

make_integrator(grid, *, dtype=<class 'numpy.float64'>)[source]

Return function that integrates discretized data over a grid.

Parameters:

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

  • dtype (DTypeLike) – The data type of the field that is being integrated

Returns:

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

Return type:

Callable[[torch.Tensor], torch.Tensor]

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

Return a torch 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, optional) – 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:

TorchDifferentialOperator

Warning

The same operator should not be assigned to different variables that are used in the same code, because torch has 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:
make_operator_no_bc(grid, operator, *, dtype=None, **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 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)[source]

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 ndarray, 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:

Callable[[torch.Tensor, torch.Tensor, torch.Tensor | None], torch.Tensor]

make_pde_rhs(eq, state)[source]

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:

TorchRHSType

make_stepper(solver, state)[source]

Return a stepper function using an explicit scheme.

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. The function call signature is (state: numpy.ndarray, t_start: float, t_end: float)

Return type:

StepperType

native_to_numpy(value)[source]

Convert native values to numpy representation.

Parameters:

value (Any)

Return type:

Any

numpy_to_native(value)[source]

Convert values from numpy to torch representation.

This method also ensures that the value is copied to the selected device.

Parameters:

value (Any)

Return type:

Tensor