FFT Functionals#

physicsnemo.nn.functional.rfft(
input: Float[Tensor, '...'],
n: int | None = None,
dim: int = -1,
norm: str | None = None,
) → Tensor#

ONNX-compatible 1D real FFT.

Parameters:

  • input (torch.Tensor) – Real input tensor.

  • n (int, optional) – Signal length along the FFT dimension.

  • dim (int, optional) – Dimension along which to take the FFT.

  • norm (str, optional) – Normalization mode ("forward", "backward", or "ortho").

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

physicsnemo.nn.functional.rfft2(
input: Float[Tensor, '...'],
s: tuple[int, int] | None = None,
dim: tuple[int, int] = (-2, -1),
norm: str | None = None,
) → Tensor#

ONNX-compatible 2D real FFT.

Parameters:

  • input (torch.Tensor) – Real input tensor.

  • s (tuple[int, int], optional) – Signal size in the transformed dimensions.

  • dim (tuple[int, int], optional) – Dimensions along which to take the FFT.

  • norm (str, optional) – Normalization mode ("forward", "backward", or "ortho").

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

physicsnemo.nn.functional.irfft(
input: Complex[Tensor, '...'] | Float[Tensor, '... 2'],
n: int | None = None,
dim: int = -1,
norm: str | None = None,
) → Tensor#

ONNX-compatible inverse 1D real FFT.

Parameters:

  • input (torch.Tensor) – Complex input tensor in the frequency domain.

  • n (int, optional) – Signal length along the inverse FFT dimension.

  • dim (int, optional) – Dimension along which to take the inverse FFT.

  • norm (str, optional) – Normalization mode ("forward", "backward", or "ortho").

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

physicsnemo.nn.functional.irfft2(
input: Complex[Tensor, '...'] | Float[Tensor, '... 2'],
s: tuple[int, int] | None = None,
dim: tuple[int, int] = (-2, -1),
norm: str | None = None,
) → Tensor#

ONNX-compatible inverse 2D real FFT.

Parameters:

  • input (torch.Tensor) – Complex input tensor in the frequency domain.

  • s (tuple[int, int], optional) – Signal size in the transformed dimensions.

  • dim (tuple[int, int], optional) – Dimensions along which to take the inverse FFT.

  • norm (str, optional) – Normalization mode ("forward", "backward", or "ortho").

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

Utilities

physicsnemo.nn.functional.view_as_complex(
input: Float[Tensor, '... 2'],
) → Tensor#

ONNX-compatible view of real-valued tensors as complex tensors.

Parameters:

  • input (torch.Tensor) – Real tensor with a final dimension of size 2 storing real/imag parts.

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

physicsnemo.nn.functional.real(
input: Complex[Tensor, '...'] | Float[Tensor, '... 2'],
) → Tensor#

ONNX-compatible view of the real component from complex tensors.

Parameters:

  • input (torch.Tensor) – Complex tensor.

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.

physicsnemo.nn.functional.imag(
input: Complex[Tensor, '...'] | Float[Tensor, '... 2'],
) → Tensor#

ONNX-compatible view of the imaginary component from complex tensors.

Parameters:

  • input (torch.Tensor) – Complex tensor.

  • implementation ({"torch"} or None) – Implementation to use. When None, dispatch selects the available implementation.