1. Concepts¶

Before we dive into the details of TorchFSM, let's introduce some key concepts used in the framework.

1.1 Data Types¶

TorchFSM is designed for deep learning applications, and all data is manipulated in the form of PyTorch Tensors. As a result, data shapes follow PyTorch conventions. There are two primary types of data shapes in TorchFSM:

• [B, C, H, ...]: This is the most common data shape in TorchFSM. Here, B is the batch size, C is the number of channels, and H, ... represent spatial dimensions. For example, 32 2D vector fields on a 64×64 mesh can be represented as [32, 1, 64, 64]. All inputs to the Operator class should conform to this shape.

• [B, T, C, H, ...]: This shape is used to store trajectory data, where T represents the number of time steps (solution frames) in the trajectory. For example, 32 2D vector fields on a 64×64 mesh over 100 time steps can be represented as [32, 100, 2, 64, 64]. This shape is typically used as the output when solving time-dependent PDEs.

In TorchFSM, we also introduce type annotations to help users supply valid inputs to functions. These type annotations are defined as follows:

• SpatialTensor: A tensor in physical space.
• SpatialArray: A NumPy array in physical space.
• FourierTensor: A tensor in Fourier space (i.e., a complex-valued tensor).
• FourierArray: A NumPy array in Fourier space (i.e., a complex-valued array).

These types can also be combined with shape annotations. For example, SpatialTensor["B C H ..."] indicates a tensor in physical space with the shape [B, C, H, ...].

Note

"SpatialTensor" and "FourierTensor" are not actual PyTorch tensor types. They are type annotations intended to clarify the expected data format.

1.2 Operator¶

The Operator is a central concept in TorchFSM. It represents a spatial operation applied to a variable. For example:

• The gradient operator computes the gradient of a variable.
• The divergence operator computes the divergence of a variable.
• The Laplacian operator computes the Laplacian of a variable.

There are three main types of operator classes in TorchFSM:

• LinearOperator: Represents linear terms in the context of the Fourier spectral method.
• NonlinearOperator: Represents nonlinear terms.
• Operator: A composite class that can include both linear and nonlinear components.

For more details on linear terms, see our Theory Introduction.

All Operator classes support basic linear algebra operations such as addition, subtraction, and multiplication. Thanks to the linearity of the Fourier transform, it's efficient to construct PDEs using combinations of different Operators.

For example, the Burgers' equation,

$$ \frac{\partial \mathbf{u}}{\partial t} = -\mathbf{u} \cdot \nabla \mathbf{u} + \nu \nabla^2 \mathbf{u}, $$

can be constructed as follows:

In [ ]:

from torchfsm.operator import Operator, Convection, Laplacian

def Burgers(nu:float) -> Operator:
    return nu*Laplacian()-Convection()
burgers=Burgers(0.01)

from torchfsm.operator import Operator, Convection, Laplacian def Burgers(nu:float) -> Operator: return nu*Laplacian()-Convection() burgers=Burgers(0.01)

There are three important functions for Operator classes:

• __call__: Calling the operator as a function directly will return the result of applying the operator to the input. This is useful for evaluating complex spatial derivatives.
• integrate: Integrates the operator over the time dimension. This is primarily used to solve time-dependent PDEs.
• solve: Solves the linear operation given a right-hand side. This function only works for the LinearOperator class and is typically used for solving time-independent problems, such as the Poisson equation.

In later sections, we will explain how to use these functions in detail.

All available operators in TorchFSM are summarized in the cheat sheet.

1.3 Mesh¶

In the Fourier spectral method, all boundary conditions are periodic, and internal obstacles within the domain are not supported. Therefore, TorchFSM uses a simplified Mesh to represent the spatial domain. The Mesh is another key concept in TorchFSM, encapsulating the domain size and grid resolution.

All basic operations in an Operator class require a mesh as input. The mesh can be one of the following types:

• mesh_info (Sequence[tuple[float, float, int]]): A sequence of tuples, where each tuple represents a spatial dimension. Each tuple contains three elements: the start of the domain, the end of the domain, and the number of grid points. For example, mesh_info = [(0, 1, 64), (0, 1, 64)] represents a 2D domain of size $[0,1] \times [0,1]$ with grid spacing $dx = dy = 1/64$.

• MeshGrid: An iterable class that represents a mesh grid. This class is especially useful for generating initial conditions. Its length corresponds to the number of dimensions in the mesh. The attributes x, y, and z provide mesh grids for the first three dimensions. For additional dimensions, you can access the mesh grid via indexing (e.g., mesh_grid[0] for the first dimension, which is equivalent to x). There is no limit to the number of dimensions. If the number of points in each dimension is $n_1, n_2, n_3, \dots, n_k$, the full mesh grid shape is (n_1, n_2, n_3, ..., n_k). The shapes of attributes x, y, z are (n_1), (n_2), (n_3) respectively. MeshGrid can be initialized with a mesh_info, along with optional keywords to specify the device and data type.

• FourierMesh: A class that contains FFT frequency information and basic derivative operators used in the Fourier spectral method. This class is used internally by Operator objects. All other mesh types are converted to a FourierMesh inside an Operator. If you're using multiple operators and want to save memory, you should create and pass a FourierMesh directly. Otherwise, each operator will create its own FourierMesh instance with identical content. FourierMesh can be initialized using either a mesh_info or a MeshGrid object.

All these mesh types can be used in the three main functions of the Operator class mentioned above. They can also be assigned to an operator using the register_mesh method. Once a mesh is registered, you no longer need to pass it explicitly when calling operator functions. A mesh is automatically registered the first time it is passed to any operator method.