firedrake.ml.pytorch package

Submodules

firedrake.ml.pytorch.fem_operator module

class firedrake.ml.pytorch.fem_operator.FiredrakeTorchOperator

Bases: Function

PyTorch custom operator representing a set of Firedrake operations expressed as a reduced functional F.

FiredrakeTorchOperator is a wrapper around torch.autograd.Function that executes forward and backward passes by directly calling the reduced functional F.

Parameters:

• metadata (dict) – Dictionary used to stash Firedrake objects.

• *x_P (torch.Tensor) – PyTorch tensors representing the inputs to the Firedrake operator F.

Returns:

PyTorch tensor representing the output of the Firedrake operator F.

Return type:

torch.Tensor

static backward(ctx, grad_output)

Backward pass of the PyTorch custom operator.

static forward(ctx, metadata, *x_P)

Forward pass of the PyTorch custom operator.

firedrake.ml.pytorch.fem_operator.fem_operator(F)

Cast a Firedrake reduced functional to a PyTorch operator.

The resulting FiredrakeTorchOperator will take PyTorch tensors as inputs and return PyTorch tensors as outputs.

Parameters:

F (pyadjoint.ReducedFunctional) – The reduced functional to wrap.

Returns:

A PyTorch custom operator that wraps the reduced functional F.

Return type:

firedrake.ml.pytorch.fem_operator.FiredrakeTorchOperator

firedrake.ml.pytorch.fem_operator.from_torch(x, V=None)

Convert a PyTorch tensor x into a Firedrake object.

Parameters:

• x (torch.Tensor) – PyTorch tensor to convert.

• V (firedrake.functionspaceimpl.WithGeometry or None) – Function space of the corresponding Function or None when x is to be mapped to a Constant.

Returns:

Firedrake object representing the PyTorch tensor x.

Return type:

firedrake.function.Function or firedrake.constant.Constant

firedrake.ml.pytorch.fem_operator.to_torch(x, gather=False, batched=True, **kwargs)

Convert a Firedrake object x into a PyTorch tensor.

Parameters:

• x (firedrake.function.Function or firedrake.constant.Constant) – Firedrake object to convert.

• gather (bool) – If True, gather data from all processes

• batched (bool) – If True, add a batch dimension to the tensor

• kwargs (dict) –

  Additional arguments to be passed to the torch.Tensor constructor such as:

  • device: device on which the tensor is allocated (default: “cpu”)

  • dtype: the desired data type of returned tensor (default: type of x.dat.data)

  • requires_grad: if the tensor should be annotated (default: False)

Returns:

PyTorch tensor representing the Firedrake object x.

Return type:

torch.Tensor

firedrake.ml.pytorch.fem_operator.torch_operator(F)

Cast a Firedrake reduced functional to a PyTorch operator.

The resulting FiredrakeTorchOperator will take PyTorch tensors as inputs and return PyTorch tensors as outputs.

Parameters:

F (pyadjoint.ReducedFunctional) – The reduced functional to wrap.

Returns:

A PyTorch custom operator that wraps the reduced functional F.

Return type:

firedrake.ml.pytorch.fem_operator.FiredrakeTorchOperator

firedrake.ml.pytorch.ml_operator module

class firedrake.ml.pytorch.ml_operator.PytorchOperator(*operands, function_space, derivatives=None, argument_slots=(), operator_data)

Bases: MLOperator

External operator class representing machine learning models implemented in PyTorch.

The PytorchOperator allows users to embed machine learning models implemented in PyTorch into PDE systems implemented in Firedrake. The actual evaluation of the PytorchOperator is delegated to the specified PyTorch model. Similarly, differentiation through the PytorchOperator class is achieved via the torch.autograd module, which provides automatic differentiation capabilities that can be applied on the PyTorch model associated with the PytorchOperator object.

Parameters:

• *operands (ufl.core.expr.Expr or ufl.form.BaseForm) – Operands of the PytorchOperator.

• function_space (firedrake.functionspaceimpl.WithGeometryBase) – The function space the ML operator is mapping to.

• derivatives (tuple) – Tuple specifiying the derivative multiindex.

• *argument_slots (ufl.coefficient.BaseCoefficient or ufl.argument.BaseArgument) – Tuple containing the arguments of the linear form associated with the ML operator, i.e. the arguments with respect to which the ML operator is linear. Those arguments can be ufl.argument.BaseArgument objects, as a result of differentiation, or both ufl.coefficient.BaseCoefficient and ufl.argument.BaseArgument object, as a result of taking the action on a given function.

• operator_data (dict) – Dictionary to stash external data specific to the ML operator. This dictionary must at least contain the following: (i) 'model': The machine learning model implemented in PyTorch. (ii) 'inputs_format': The format of the inputs to the ML model: 0 for models acting globally on the inputs, 1 when acting locally/pointwise on the inputs. Other strategies can also be considered by subclassing the PytorchOperator class.

property model_output

property torch_grad_enabled

ufl_operands: tuple[FormArgument, ...]

firedrake.ml.pytorch.ml_operator.ml_operator(model, function_space, inputs_format=0)

Helper function for instantiating the PytorchOperator class.

This function facilitates having a two-stage instantiation which dissociates between class arguments that are fixed, such as the function space or the ML model, and the operands of the operator, which may change, e.g. when the operator is used in a time-loop.

Example

# Stage 1: Partially initialise the operator.
N = ml_operator(model, function_space=V)
# Stage 2: Define the operands and use the operator in a UFL expression.
F = (inner(grad(u), grad(v)) + inner(N(u), v) - inner(f, v)) * dx

Parameters:

• model (collections.abc.Callable) – The PyTorch model to embed in Firedrake.

• function_space (firedrake.functionspaceimpl.WithGeometryBase) – The function space into which the machine learning model is mapping.

• inputs_format (int) – The format of the input data of the ML model: 0 for models acting globally on the inputs, 1 when acting locally/pointwise on the inputs. Other strategies can also be considered by subclassing the PytorchOperator class.

Returns:

The partially initialised PytorchOperator class.

Return type:

collections.abc.Callable

firedrake.ml.pytorch.ml_operator.neuralnet(model, function_space, inputs_format=0)

Helper function for instantiating the PytorchOperator class.

This function facilitates having a two-stage instantiation which dissociates between class arguments that are fixed, such as the function space or the ML model, and the operands of the operator, which may change, e.g. when the operator is used in a time-loop.

Example

# Stage 1: Partially initialise the operator.
N = ml_operator(model, function_space=V)
# Stage 2: Define the operands and use the operator in a UFL expression.
F = (inner(grad(u), grad(v)) + inner(N(u), v) - inner(f, v)) * dx

Parameters:

• model (collections.abc.Callable) – The PyTorch model to embed in Firedrake.

• function_space (firedrake.functionspaceimpl.WithGeometryBase) – The function space into which the machine learning model is mapping.

• inputs_format (int) – The format of the input data of the ML model: 0 for models acting globally on the inputs, 1 when acting locally/pointwise on the inputs. Other strategies can also be considered by subclassing the PytorchOperator class.

Returns:

The partially initialised PytorchOperator class.

Return type:

collections.abc.Callable

Module contents