4.1.6.5 pde.backends.torch.utils module

Defines utilities for the torch backend.

class TorchGaussianNoise(data_shape, dtype, scale=1, generator=None)[source]

Bases: TorchOperatorBase

Operator that returns uncorrelated Gaussian random field.

Parameters:

  • data_shape (tuple of ints) – Shape of the output array

  • dtype – Dtype of the

  • scale (float or array) – Scaling of each entry in the field

  • generator (torch.Generator or None) – Random number generator

forward()[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class TorchOperatorBase(*, dtype)[source]

Bases: Module

Base class for operators implemented in torch.

Initialize the torch operator.

Parameters:

dtype (type[Any] | dtype[Any] | _HasDType[dtype[Any]] | _HasNumPyDType[dtype[Any]] | tuple[Any, Any] | list[Any] | _DTypeDict | str) – The data type of the field using the numpy convention

register_array(name, arr)[source]

Register an array as a buffer in the torch module.

Parameters:

  • name (str) – The name under which the buffer is registered

  • arr (numpy.ndarray or torch.Tensor) – The array to register. If a numpy array is provided, it will be converted to a torch tensor with the appropriate dtype.

Return type:

None

torch_heaviside(x1, x2=None)[source]

Return the Heaviside step function using torch.

This does not use torch.heaviside() since this is not implemented for the MPS device.

Parameters:

  • x1 (torch.Tensor) – Input values at which the Heaviside function is evaluated.

  • x2 (torch.Tensor, optional) – Value used where x1 == 0. If omitted, 0.5 is used.

Returns:

Tensor containing the Heaviside values of x1.

Return type:

torch.Tensor

torch_hypot(x1, x2)[source]

Return the Euclidean norm sqrt(x1**2 + x2**2) using torch.

This wraps torch.hypot() and ensures that both inputs are converted to tensors before evaluation.

Parameters:
Returns:

Tensor containing the element-wise hypotenuse of x1 and x2.

Return type:

torch.Tensor