Skip to content

1. operator

Operators¤

Operators are the core of the torchfsm library. Each operator represents a specific physical process, such as convection, diffusion, or pressure calculation.

torchfsm.operator.Biharmonic ¤

Bases: LinearOperator

Biharmonic calculates the Biharmonic of a vector field. It is defined as \(\nabla^4\mathbf{u}=\left[\begin{matrix}(\sum_{i=0}^I\frac{\partial^2}{\partial i^2 })(\sum_{j=0}^I\frac{\partial^2}{\partial j^2 })u_x \\ (\sum_{i=0}^I\frac{\partial^2}{\partial i^2 })(\sum_{j=0}^I\frac{\partial^2}{\partial j^2 })u_y \\ \cdots \\ (\sum_{i=0}^I\frac{\partial^2}{\partial i^2 })(\sum_{j=0}^I\frac{\partial^2}{\partial j^2 })u_i \\ \end{matrix} \right]\) Note that this class is an operator wrapper. The real implementation of the source term is in the _BiharmonicCore class.

Source code in torchfsm/operator/generic/_biharmonic.py 
19
20
21
22
23
24
25
26
27
28
class Biharmonic(LinearOperator):
    r"""
    `Biharmonic` calculates the Biharmonic of a vector field. 
        It is defined as $\nabla^4\mathbf{u}=\left[\begin{matrix}(\sum_{i=0}^I\frac{\partial^2}{\partial i^2 })(\sum_{j=0}^I\frac{\partial^2}{\partial j^2 })u_x \\ (\sum_{i=0}^I\frac{\partial^2}{\partial i^2 })(\sum_{j=0}^I\frac{\partial^2}{\partial j^2 })u_y \\ \cdots \\ (\sum_{i=0}^I\frac{\partial^2}{\partial i^2 })(\sum_{j=0}^I\frac{\partial^2}{\partial j^2 })u_i \\ \end{matrix} \right]$
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_BiharmonicCore` class.    

    """

    def __init__(self) -> None:
        super().__init__(_BiharmonicCore())
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
solve ¤
solve(
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]
]

Solve the linear operator equation \(Ax = b\), where \(A\) is the linear operator and \(b\) is the right-hand side.

Parameters:

Name Type Description Default
b Optional[Tensor]

Right-hand side tensor in spatial domain. If None, b_fft should be provided.

 None
b_fft Optional[Tensor]

Right-hand side tensor in Fourier domain. If None, b should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. If None, the mesh registered in the operator will be used.

 None
n_channel Optional[int]

Number of channels of \(x\). If None, the number of channels registered in the operator will be used.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def solve(
    self,
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]]:

    r"""
    Solve the linear operator equation $Ax = b$, where $A$ is the linear operator and $b$ is the right-hand side.

    Args:
        b (Optional[torch.Tensor]): Right-hand side tensor in spatial domain. If None, b_fft should be provided.
        b_fft (Optional[torch.Tensor]): Right-hand side tensor in Fourier domain. If None, b should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. If None, the mesh registered in the operator will be used.
        n_channel (Optional[int]): Number of channels of $x$. If None, the number of channels registered in the operator will be used.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
    """
    if not (mesh is not None and n_channel is not None):
        assert (
            self._state_dict["f_mesh"] is not None
        ), "Mesh and n_channel should be given when calling solve"
    if not (mesh is None and n_channel is None):
        mesh = self._state_dict["f_mesh"] if mesh is None else mesh
        n_channel = (
            self._state_dict["n_channel"] if n_channel is None else n_channel
        )
        self.register_mesh(mesh, n_channel)
    if self._state_dict["invert_linear_coef"] is None:
        self._state_dict["invert_linear_coef"] = torch.where(
            self._state_dict["linear_coef"] == 0,
            1.0,
            1 / self._state_dict["linear_coef"],
        )
    if b_fft is None:
        b_fft = self._state_dict["f_mesh"].fft(b)
    value_fft = b_fft * self._state_dict["invert_linear_coef"]
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
859
860
861
862
863
864
865
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
867
868
869
870
def __neg__(self):
    return LinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_biharmonic.py 
27
28
def __init__(self) -> None:
    super().__init__(_BiharmonicCore())

torchfsm.operator.ConservativeConvection ¤

Bases: NonlinearOperator

ConservativeConvection calculates the convection of a vector field on itself. It is defined as \(\nabla \cdot \mathbf{u}\mathbf{u}=\left[\begin{matrix}\sum_{i=0}^I \frac{\partial u_i u_x }{\partial i} \\\sum_{i=0}^I \frac{\partial u_i u_y }{\partial i} \\\cdots\\\sum_{i=0}^I \frac{\partial u_i u_I }{\partial i} \\\end{matrix}\right]\). This operator only works for vector fields with the same dimension as the mesh. Note that this class is an operator wrapper. The real implementation of the source term is in the _ConservativeConvectionCore class.

Source code in torchfsm/operator/generic/_conservative_convection.py 
46
47
48
49
50
51
52
53
54
55
class ConservativeConvection(NonlinearOperator):
    r"""
    `ConservativeConvection` calculates the convection of a vector field on itself. 
        It is defined as $\nabla \cdot \mathbf{u}\mathbf{u}=\left[\begin{matrix}\sum_{i=0}^I \frac{\partial u_i u_x }{\partial i} \\\sum_{i=0}^I \frac{\partial u_i u_y }{\partial i} \\\cdots\\\sum_{i=0}^I \frac{\partial u_i u_I }{\partial i} \\\end{matrix}\right]$.
        This operator only works for vector fields with the same dimension as the mesh.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_ConservativeConvectionCore` class.
    """

    def __init__(self) -> None:
        super().__init__(_ConservativeConvectionGenerator())
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_conservative_convection.py 
54
55
def __init__(self) -> None:
    super().__init__(_ConservativeConvectionGenerator())

torchfsm.operator.Convection ¤

Bases: NonlinearOperator

Convection calculates the convection of a vector field on itself if the vector field is divergence free, i.e., \(\nabla \cdot \mathbf{u} =0\). It is defined as \(\mathbf{u} \cdot \nabla \mathbf{u}=\left[\begin{matrix}\sum_{i=0}^I u_i\frac{\partial u_x }{\partial i} \\\sum_{i=0}^I u_i\frac{\partial u_y }{\partial i} \\\cdots\\\sum_{i=0}^I u_i\frac{\partial u_I }{\partial i} \\\end{matrix}\right]\) This operator only works for vector fields with the same dimension as the mesh. Note that this class is an operator wrapper. The real implementation of the source term is in the _ConvectionCore class.

Source code in torchfsm/operator/generic/_convection.py 
67
68
69
70
71
72
73
74
75
76
class Convection(NonlinearOperator):
    r"""
    `Convection` calculates the convection of a vector field on itself if the vector field is divergence free, i.e., $\nabla \cdot \mathbf{u} =0$.
        It is defined as $\mathbf{u} \cdot \nabla  \mathbf{u}=\left[\begin{matrix}\sum_{i=0}^I u_i\frac{\partial u_x }{\partial i} \\\sum_{i=0}^I u_i\frac{\partial u_y }{\partial i} \\\cdots\\\sum_{i=0}^I u_i\frac{\partial u_I }{\partial i} \\\end{matrix}\right]$
        This operator only works for vector fields with the same dimension as the mesh.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_ConvectionCore` class.
    """

    def __init__(self) -> None:
        super().__init__(_ConvectionGenerator())
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_convection.py 
75
76
def __init__(self) -> None:
    super().__init__(_ConvectionGenerator())

torchfsm.operator.Curl ¤

Bases: NonlinearOperator

Curl operator for 2D and 3D vector fields. It is defined as: \(\nabla \times \mathbf{u} = \frac{\partial u_y}{\partial x}-\frac{\partial u_x}{\partial y}\) for 2D vector field \(\mathbf{u} = (u_x, u_y)\) and \(\nabla \times \mathbf{u} = \left[\begin{matrix} \frac{\partial u_z}{\partial y}-\frac{\partial u_y}{\partial z} \\ \frac{\partial u_x}{\partial z}-\frac{\partial u_z}{\partial x} \\ \frac{\partial u_y}{\partial x}-\frac{\partial u_x}{\partial y} \end{matrix} \right]\) for 3D vector field \(\mathbf{u} = (u_x, u_y, u_z)\). This operator only works for vector fields with the same dimension as the mesh. Note that this class is an operator wrapper. The real implementation of the source term is in the _Curl2DCore and _Curl2DCore class.

Source code in torchfsm/operator/generic/_curl.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class Curl(NonlinearOperator):

    r"""
    Curl operator for 2D and 3D vector fields. 
        It is defined as: $\nabla \times \mathbf{u} = \frac{\partial u_y}{\partial x}-\frac{\partial u_x}{\partial y}$
        for 2D vector field $\mathbf{u} = (u_x, u_y)$ and
        $\nabla \times \mathbf{u} = \left[\begin{matrix} \frac{\partial u_z}{\partial y}-\frac{\partial u_y}{\partial z} \\ \frac{\partial u_x}{\partial z}-\frac{\partial u_z}{\partial x} \\ \frac{\partial u_y}{\partial x}-\frac{\partial u_x}{\partial y} \end{matrix} \right]$
        for 3D vector field $\mathbf{u} = (u_x, u_y, u_z)$.
        This operator only works for vector fields with the same dimension as the mesh.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_Curl2DCore` and `_Curl2DCore` class.

    """

    def __init__(self) -> None:
        super().__init__(_CurlGenerator())
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_curl.py 
92
93
def __init__(self) -> None:
    super().__init__(_CurlGenerator())

torchfsm.operator.Div ¤

Bases: NonlinearOperator

Div calculates the divergence of a vector field. It is defined as\(\nabla \cdot \mathbf{u} = \sum_i \frac{\partial u_i}{\partial i}\). This operator only works for vector fields with the same dimension as the mesh. Note that this class is an operator wrapper. The actual implementation of the operator is in the _DivCore class.

Source code in torchfsm/operator/generic/_div.py 
43
44
45
46
47
48
49
50
51
52
53
class Div(NonlinearOperator):
    r"""
    `Div` calculates the divergence of a vector field.
        It is defined as$\nabla \cdot \mathbf{u} = \sum_i \frac{\partial u_i}{\partial i}$.
        This operator only works for vector fields with the same dimension as the mesh.
        Note that this class is an operator wrapper. The actual implementation of the operator is in the `_DivCore` class.

    """

    def __init__(self) -> None:
        super().__init__(_DivGenerator())
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_div.py 
52
53
def __init__(self) -> None:
    super().__init__(_DivGenerator())

torchfsm.operator.Grad ¤

Bases: LinearOperator

Grad calculates the spatial gradient of a scalar field. It is defined as \(\nabla p = \left[\begin{matrix}\frac{\partial p}{\partial x} \\\frac{\partial p}{\partial y} \\\cdots \\\frac{\partial p}{\partial i} \\\end{matrix}\right]\) Note that this class is an operator wrapper. The actual implementation of the operator is in the _GradCore class.

Source code in torchfsm/operator/generic/_grad.py 
30
31
32
33
34
35
36
37
38
class Grad(LinearOperator):
    r"""
    `Grad` calculates the spatial gradient of a scalar field.
        It is defined as $\nabla p = \left[\begin{matrix}\frac{\partial p}{\partial x} \\\frac{\partial p}{\partial y} \\\cdots \\\frac{\partial p}{\partial i} \\\end{matrix}\right]$
        Note that this class is an operator wrapper. The actual implementation of the operator is in the `_GradCore` class.
    """

    def __init__(self) -> None:
        super().__init__(_GradGenerator())
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
solve ¤
solve(
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]
]

Solve the linear operator equation \(Ax = b\), where \(A\) is the linear operator and \(b\) is the right-hand side.

Parameters:

Name Type Description Default
b Optional[Tensor]

Right-hand side tensor in spatial domain. If None, b_fft should be provided.

 None
b_fft Optional[Tensor]

Right-hand side tensor in Fourier domain. If None, b should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. If None, the mesh registered in the operator will be used.

 None
n_channel Optional[int]

Number of channels of \(x\). If None, the number of channels registered in the operator will be used.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def solve(
    self,
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]]:

    r"""
    Solve the linear operator equation $Ax = b$, where $A$ is the linear operator and $b$ is the right-hand side.

    Args:
        b (Optional[torch.Tensor]): Right-hand side tensor in spatial domain. If None, b_fft should be provided.
        b_fft (Optional[torch.Tensor]): Right-hand side tensor in Fourier domain. If None, b should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. If None, the mesh registered in the operator will be used.
        n_channel (Optional[int]): Number of channels of $x$. If None, the number of channels registered in the operator will be used.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
    """
    if not (mesh is not None and n_channel is not None):
        assert (
            self._state_dict["f_mesh"] is not None
        ), "Mesh and n_channel should be given when calling solve"
    if not (mesh is None and n_channel is None):
        mesh = self._state_dict["f_mesh"] if mesh is None else mesh
        n_channel = (
            self._state_dict["n_channel"] if n_channel is None else n_channel
        )
        self.register_mesh(mesh, n_channel)
    if self._state_dict["invert_linear_coef"] is None:
        self._state_dict["invert_linear_coef"] = torch.where(
            self._state_dict["linear_coef"] == 0,
            1.0,
            1 / self._state_dict["linear_coef"],
        )
    if b_fft is None:
        b_fft = self._state_dict["f_mesh"].fft(b)
    value_fft = b_fft * self._state_dict["invert_linear_coef"]
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
859
860
861
862
863
864
865
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
867
868
869
870
def __neg__(self):
    return LinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_grad.py 
37
38
def __init__(self) -> None:
    super().__init__(_GradGenerator())

torchfsm.operator.ExplicitSource ¤

Bases: NonlinearOperator

Explicit source term for the operator. This class is used to represent an explicit source term in the operator. Note that this class is an operator wrapper. The real implementation of the source term is in the _ExplicitSourceCore class.

Parameters:

Name Type Description Default
source Tensor

Source term in spatial domain. This is a tensor that represents the source term in the spatial domain.

 required
Source code in torchfsm/operator/_base.py 
971
972
973
974
975
976
977
978
979
980
981
982
class ExplicitSource(NonlinearOperator):

    r"""
    Explicit source term for the operator. This class is used to represent an explicit source term in the operator.
        Note that this class is an operator wrapper. The real implementation of the source term is in the _ExplicitSourceCore class. 

    Args:
        source (torch.Tensor): Source term in spatial domain. This is a tensor that represents the source term in the spatial domain.
    """

    def __init__(self, source: torch.Tensor) -> None:
        super().__init__(_ExplicitSourceCore(source))
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(source: torch.Tensor) -> None
Source code in torchfsm/operator/_base.py 
981
982
def __init__(self, source: torch.Tensor) -> None:
    super().__init__(_ExplicitSourceCore(source))

torchfsm.operator.ImplicitSource ¤

Bases: Operator

ImplicitSource allows to define a source term in the implicit form. Note that this class is an operator wrapper. The actual implementation of the operator is in the _ImplicitFuncSourceCore class and _ImplicitUnitSourceCore.

Parameters:

Name Type Description Default
source_func Callable[[Tensor], Tensor]

The f(x) function to be used as the source term. This function is used to define the source term in the implicit form. If None, the source term will be set to the unknown variable itself, i.e., f(x) = x.

 None
non_linear bool

If True, the source term is treated as a nonlinear function. If False, it is treated as a linear function. Default is True. This actually controls whether the operator wil use the dealiased version of unknown variable for the source term. If the source term is a nonlinear function, the dealiased version of the unknown variable will be used.

 True
Source code in torchfsm/operator/generic/_source.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
class ImplicitSource(Operator):

    r"""
    `ImplicitSource` allows to define a source term in the implicit form.
        Note that this class is an operator wrapper. The actual implementation of the operator is in the `_ImplicitFuncSourceCore` class and `_ImplicitUnitSourceCore`.

    Args:
        source_func (Callable[[torch.Tensor], torch.Tensor], optional): 
            The f(x) function to be used as the source term.
            This function is used to define the source term in the implicit form.
            If None, the source term will be set to the unknown variable itself, i.e., f(x) = x.

        non_linear (bool, optional): 
            If True, the source term is treated as a nonlinear function. 
            If False, it is treated as a linear function. Default is True.
            This actually controls whether the operator wil use the dealiased version of unknown variable for the source term.
            If the source term is a nonlinear function, the dealiased version of the unknown variable will be used.
    """

    def __init__(
        self,
        source_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
        non_linear: bool = True,
    ) -> None:
        if source_func is None:
            generator = lambda f_mesh, n_channel: _ImplicitUnitSourceCore()
        else:
            generator = lambda f_mesh, n_channel: _ImplicitFuncSourceCore(
                source_func, non_linear
            )
        super().__init__(generator)
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear: bool

Check if the operator is linear.

Returns:

Name Type Description
bool bool

True if the operator is linear, False otherwise.
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
775
776
777
778
779
780
781
782
783
784
785
786
787
788
def __add__(self, other):
    if isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
790
791
792
793
794
795
796
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return Operator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
798
799
def __neg__(self):
    return Operator(self.operator_generators, [-1 * coef for coef in self.coefs])
__init__ ¤
__init__(
    source_func: Optional[
        Callable[[torch.Tensor], torch.Tensor]
    ] = None,
    non_linear: bool = True,
) -> None
Source code in torchfsm/operator/generic/_source.py 
63
64
65
66
67
68
69
70
71
72
73
74
def __init__(
    self,
    source_func: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
    non_linear: bool = True,
) -> None:
    if source_func is None:
        generator = lambda f_mesh, n_channel: _ImplicitUnitSourceCore()
    else:
        generator = lambda f_mesh, n_channel: _ImplicitFuncSourceCore(
            source_func, non_linear
        )
    super().__init__(generator)

torchfsm.operator.Laplacian ¤

Bases: LinearOperator

Laplacian calculates the Laplacian of a vector field.

It is defined as \(\nabla \cdot (\nabla\mathbf{u}) = \left[\begin{matrix}\sum_i \frac{\partial^2 u_x}{\partial i^2 } \\\sum_i \frac{\partial^2 u_y}{\partial i^2 } \\\cdots \\\sum_i \frac{\partial^2 u_i}{\partial i^2 } \\\end{matrix}\right]\) Note that this class is an operator wrapper. The actual implementation of the operator is in the _LaplacianCore class.

Source code in torchfsm/operator/generic/_laplacian.py 
17
18
19
20
21
22
23
24
25
26
class Laplacian(LinearOperator):
    r"""
    `Laplacian` calculates the Laplacian of a vector field.

    It is defined as $\nabla \cdot (\nabla\mathbf{u}) = \left[\begin{matrix}\sum_i \frac{\partial^2 u_x}{\partial i^2 } \\\sum_i \frac{\partial^2 u_y}{\partial i^2 } \\\cdots \\\sum_i \frac{\partial^2 u_i}{\partial i^2 } \\\end{matrix}\right]$
    Note that this class is an operator wrapper. The actual implementation of the operator is in the `_LaplacianCore` class.
    """

    def __init__(self) -> None:
        super().__init__( _LaplacianCore())
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
solve ¤
solve(
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]
]

Solve the linear operator equation \(Ax = b\), where \(A\) is the linear operator and \(b\) is the right-hand side.

Parameters:

Name Type Description Default
b Optional[Tensor]

Right-hand side tensor in spatial domain. If None, b_fft should be provided.

 None
b_fft Optional[Tensor]

Right-hand side tensor in Fourier domain. If None, b should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. If None, the mesh registered in the operator will be used.

 None
n_channel Optional[int]

Number of channels of \(x\). If None, the number of channels registered in the operator will be used.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def solve(
    self,
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]]:

    r"""
    Solve the linear operator equation $Ax = b$, where $A$ is the linear operator and $b$ is the right-hand side.

    Args:
        b (Optional[torch.Tensor]): Right-hand side tensor in spatial domain. If None, b_fft should be provided.
        b_fft (Optional[torch.Tensor]): Right-hand side tensor in Fourier domain. If None, b should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. If None, the mesh registered in the operator will be used.
        n_channel (Optional[int]): Number of channels of $x$. If None, the number of channels registered in the operator will be used.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
    """
    if not (mesh is not None and n_channel is not None):
        assert (
            self._state_dict["f_mesh"] is not None
        ), "Mesh and n_channel should be given when calling solve"
    if not (mesh is None and n_channel is None):
        mesh = self._state_dict["f_mesh"] if mesh is None else mesh
        n_channel = (
            self._state_dict["n_channel"] if n_channel is None else n_channel
        )
        self.register_mesh(mesh, n_channel)
    if self._state_dict["invert_linear_coef"] is None:
        self._state_dict["invert_linear_coef"] = torch.where(
            self._state_dict["linear_coef"] == 0,
            1.0,
            1 / self._state_dict["linear_coef"],
        )
    if b_fft is None:
        b_fft = self._state_dict["f_mesh"].fft(b)
    value_fft = b_fft * self._state_dict["invert_linear_coef"]
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
859
860
861
862
863
864
865
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
867
868
869
870
def __neg__(self):
    return LinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_laplacian.py 
25
26
def __init__(self) -> None:
    super().__init__( _LaplacianCore())

torchfsm.operator.SpatialDerivative ¤

Bases: LinearOperator

SpatialDeritivate calculates the spatial derivative of a scalar field w.r.t to a spatial dimension. It is defined as\(\frac{\partial ^n}{\partial i} p\) where \(i = x, y, z, \cdots\) and \(n=1, 2, 3, \cdots\) Note that this class is an operator wrapper. The actual implementation of the operator is in the _SpatialDerivativeCore class.

Parameters:

Name Type Description Default
dim_index int

The index of the spatial dimension.

 required
order int

The order of the derivative.

 required
Source code in torchfsm/operator/generic/_spatial_derivative.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
class SpatialDerivative(LinearOperator):
    r"""
    `SpatialDeritivate` calculates the spatial derivative of a scalar field w.r.t to a spatial dimension.
        It is defined as$\frac{\partial ^n}{\partial i} p$ where $i = x, y, z, \cdots$ and $n=1, 2, 3, \cdots$
        Note that this class is an operator wrapper. The actual implementation of the operator is in the `_SpatialDerivativeCore` class.

    Args:
        dim_index (int): The index of the spatial dimension.
        order (int): The order of the derivative.
    """

    def __init__(self, dim_index: int, order: int) -> None:
        super().__init__(_SpatialDerivativeGenerator(dim_index, order))
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
solve ¤
solve(
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]
]

Solve the linear operator equation \(Ax = b\), where \(A\) is the linear operator and \(b\) is the right-hand side.

Parameters:

Name Type Description Default
b Optional[Tensor]

Right-hand side tensor in spatial domain. If None, b_fft should be provided.

 None
b_fft Optional[Tensor]

Right-hand side tensor in Fourier domain. If None, b should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. If None, the mesh registered in the operator will be used.

 None
n_channel Optional[int]

Number of channels of \(x\). If None, the number of channels registered in the operator will be used.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def solve(
    self,
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]]:

    r"""
    Solve the linear operator equation $Ax = b$, where $A$ is the linear operator and $b$ is the right-hand side.

    Args:
        b (Optional[torch.Tensor]): Right-hand side tensor in spatial domain. If None, b_fft should be provided.
        b_fft (Optional[torch.Tensor]): Right-hand side tensor in Fourier domain. If None, b should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. If None, the mesh registered in the operator will be used.
        n_channel (Optional[int]): Number of channels of $x$. If None, the number of channels registered in the operator will be used.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
    """
    if not (mesh is not None and n_channel is not None):
        assert (
            self._state_dict["f_mesh"] is not None
        ), "Mesh and n_channel should be given when calling solve"
    if not (mesh is None and n_channel is None):
        mesh = self._state_dict["f_mesh"] if mesh is None else mesh
        n_channel = (
            self._state_dict["n_channel"] if n_channel is None else n_channel
        )
        self.register_mesh(mesh, n_channel)
    if self._state_dict["invert_linear_coef"] is None:
        self._state_dict["invert_linear_coef"] = torch.where(
            self._state_dict["linear_coef"] == 0,
            1.0,
            1 / self._state_dict["linear_coef"],
        )
    if b_fft is None:
        b_fft = self._state_dict["f_mesh"].fft(b)
    value_fft = b_fft * self._state_dict["invert_linear_coef"]
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
859
860
861
862
863
864
865
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
867
868
869
870
def __neg__(self):
    return LinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(dim_index: int, order: int) -> None
Source code in torchfsm/operator/generic/_spatial_derivative.py 
54
55
def __init__(self, dim_index: int, order: int) -> None:
    super().__init__(_SpatialDerivativeGenerator(dim_index, order))

torchfsm.operator.KSConvection ¤

Bases: NonlinearOperator

The Kuramoto-Sivashinsky convection operator for a scalar field. It is defined as: \(\frac{1}{2}|\nabla \phi|^2=\frac{1}{2}\sum_{i=0}^{I}(\frac{\partial \phi}{\partial i})^2\) Note that this class is an operator wrapper. The real implementation of the source term is in the _KSConvectionCore class.

Parameters:

Name Type Description Default
remove_mean bool

Whether to remove the mean of the result. Default is True. Set to True will improve the stability of the simulation.

 True
Source code in torchfsm/operator/dedicated/_ks_convection.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
class KSConvection(NonlinearOperator):

    r"""
    The Kuramoto-Sivashinsky convection operator for a scalar field.
        It is defined as: $\frac{1}{2}|\nabla \phi|^2=\frac{1}{2}\sum_{i=0}^{I}(\frac{\partial \phi}{\partial i})^2$
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_KSConvectionCore` class.

    Args:
        remove_mean (bool): Whether to remove the mean of the result. Default is True. Set to True will improve the stability of the simulation.
    """

    def __init__(self, remove_mean: bool = True) -> None:
        super().__init__(_KSConvectionGenerator(remove_mean))
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(remove_mean: bool = True) -> None
Source code in torchfsm/operator/dedicated/_ks_convection.py 
66
67
def __init__(self, remove_mean: bool = True) -> None:
    super().__init__(_KSConvectionGenerator(remove_mean))

torchfsm.operator.VorticityConvection ¤

Bases: NonlinearOperator

Operator for vorticity convection in 2D. It is defined as \((\mathbf{u}\cdot\nabla) \omega\) where \(\omega\) is the vorticity and \(\mathbf{u}\) is the velocity. Note that this class is an operator wrapper. The real implementation of the source term is in the _VorticityConvectionCore class.

Source code in torchfsm/operator/dedicated/_navier_stokes.py 
62
63
64
65
66
67
68
69
70
71
class VorticityConvection(NonlinearOperator):

    r"""
    Operator for vorticity convection in 2D. 
        It is defined as $(\mathbf{u}\cdot\nabla) \omega$ where $\omega$ is the vorticity and $\mathbf{u}$ is the velocity.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_VorticityConvectionCore` class.
    """

    def __init__(self) -> None:
        super().__init__(_VorticityConvectionGenerator())
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/dedicated/_navier_stokes.py 
70
71
def __init__(self) -> None:
    super().__init__(_VorticityConvectionGenerator())

torchfsm.operator.Vorticity2Velocity ¤

Bases: LinearOperator

Operator for vorticity to velocity conversion in 2D. It is defined as \([u,v]=[-\frac{\partial \nabla^{-2}\omega}{\partial y},\frac{\partial \nabla^{-2}\omega}{\partial x}]\). Note that this class is an operator wrapper. The real implementation of the source term is in the _Vorticity2VelocityCore class.

Source code in torchfsm/operator/dedicated/_navier_stokes.py 
106
107
108
109
110
111
112
113
114
115
class Vorticity2Velocity(LinearOperator):

    r"""
    Operator for vorticity to velocity conversion in 2D.
        It is defined as $[u,v]=[-\frac{\partial \nabla^{-2}\omega}{\partial y},\frac{\partial \nabla^{-2}\omega}{\partial x}]$.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_Vorticity2VelocityCore` class.
    """

    def __init__(self):
        super().__init__(_Vorticity2VelocityGenerator())
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
solve ¤
solve(
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]
]

Solve the linear operator equation \(Ax = b\), where \(A\) is the linear operator and \(b\) is the right-hand side.

Parameters:

Name Type Description Default
b Optional[Tensor]

Right-hand side tensor in spatial domain. If None, b_fft should be provided.

 None
b_fft Optional[Tensor]

Right-hand side tensor in Fourier domain. If None, b should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. If None, the mesh registered in the operator will be used.

 None
n_channel Optional[int]

Number of channels of \(x\). If None, the number of channels registered in the operator will be used.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def solve(
    self,
    b: Optional[torch.Tensor] = None,
    b_fft: Optional[torch.Tensor] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    n_channel: Optional[int] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], SpatialTensor["B C H ..."]]:

    r"""
    Solve the linear operator equation $Ax = b$, where $A$ is the linear operator and $b$ is the right-hand side.

    Args:
        b (Optional[torch.Tensor]): Right-hand side tensor in spatial domain. If None, b_fft should be provided.
        b_fft (Optional[torch.Tensor]): Right-hand side tensor in Fourier domain. If None, b should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. If None, the mesh registered in the operator will be used.
        n_channel (Optional[int]): Number of channels of $x$. If None, the number of channels registered in the operator will be used.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Solution tensor in spatial or Fourier domain.
    """
    if not (mesh is not None and n_channel is not None):
        assert (
            self._state_dict["f_mesh"] is not None
        ), "Mesh and n_channel should be given when calling solve"
    if not (mesh is None and n_channel is None):
        mesh = self._state_dict["f_mesh"] if mesh is None else mesh
        n_channel = (
            self._state_dict["n_channel"] if n_channel is None else n_channel
        )
        self.register_mesh(mesh, n_channel)
    if self._state_dict["invert_linear_coef"] is None:
        self._state_dict["invert_linear_coef"] = torch.where(
            self._state_dict["linear_coef"] == 0,
            1.0,
            1 / self._state_dict["linear_coef"],
        )
    if b_fft is None:
        b_fft = self._state_dict["f_mesh"].fft(b)
    value_fft = b_fft * self._state_dict["invert_linear_coef"]
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
859
860
861
862
863
864
865
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
867
868
869
870
def __neg__(self):
    return LinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__()
Source code in torchfsm/operator/dedicated/_navier_stokes.py 
114
115
def __init__(self):
    super().__init__(_Vorticity2VelocityGenerator())

torchfsm.operator.Vorticity2Pressure ¤

Bases: NonlinearOperator

Operator for vorticity to pressure conversion in 2D. It is defined as \(\begin{matrix}\mathbf{u}=[u,v]=[-\frac{\partial \nabla^{-2}\omega}{\partial y},\frac{\partial \nabla^{-2}\omega}{\partial x}]\\ p= -\nabla^{-2} (\nabla \cdot (\left(\mathbf{u}\cdot\nabla\right)\mathbf{u}-f))\end{matrix}\). Note that this class is an operator wrapper. The real implementation of the source term is in the _Vorticity2PressureCore class.

Source code in torchfsm/operator/dedicated/_navier_stokes.py 
162
163
164
165
166
167
168
169
170
class Vorticity2Pressure(NonlinearOperator):
    r"""
    Operator for vorticity to pressure conversion in 2D.
        It is defined as $\begin{matrix}\mathbf{u}=[u,v]=[-\frac{\partial \nabla^{-2}\omega}{\partial y},\frac{\partial \nabla^{-2}\omega}{\partial x}]\\ p= -\nabla^{-2} (\nabla \cdot (\left(\mathbf{u}\cdot\nabla\right)\mathbf{u}-f))\end{matrix}$.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_Vorticity2PressureCore` class.
    """

    def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
        super().__init__(_Vorticity2PressureGenerator(external_force))
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes.py 
169
170
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_Vorticity2PressureGenerator(external_force))

torchfsm.operator.Velocity2Pressure ¤

Bases: NonlinearOperator

Operator for velocity to pressure conversion. It is defined as \(-\nabla^{-2} (\nabla \cdot (\left(\mathbf{u}\cdot\nabla\right)\mathbf{u}-f))\) Note that this class is an operator wrapper. The real implementation of the source term is in the _Velocity2PressureCore class.

Source code in torchfsm/operator/dedicated/_navier_stokes.py 
209
210
211
212
213
214
215
216
217
class Velocity2Pressure(NonlinearOperator):
    r"""
    Operator for velocity to pressure conversion.
        It is defined as $-\nabla^{-2} (\nabla \cdot (\left(\mathbf{u}\cdot\nabla\right)\mathbf{u}-f))$
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_Velocity2PressureCore` class.
    """

    def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
        super().__init__(_Velocity2PressureCore(external_force))
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes.py 
216
217
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_Velocity2PressureCore(external_force))

torchfsm.operator.NSPressureConvection ¤

Bases: NonlinearOperator

Operator for Navier-Stokes pressure convection. It is defined as \(-\nabla (\nabla^{-2} \nabla \cdot (\left(\mathbf{u}\cdot\nabla\right)\mathbf{u}-f))-\left(\mathbf{u}\cdot\nabla\right)\mathbf{u} + \mathbf{f}\). Note that this class is an operator wrapper. The real implementation of the source term is in the _NSPressureConvectionCore class.

Parameters:

Name Type Description Default
external_force Optional[OperatorLike]

Optional[OperatorLike], optional, default=None

 None
Source code in torchfsm/operator/dedicated/_navier_stokes.py 
257
258
259
260
261
262
263
264
265
266
267
268
class NSPressureConvection(NonlinearOperator):
    r"""
    Operator for Navier-Stokes pressure convection.
        It is defined as $-\nabla (\nabla^{-2} \nabla \cdot (\left(\mathbf{u}\cdot\nabla\right)\mathbf{u}-f))-\left(\mathbf{u}\cdot\nabla\right)\mathbf{u} + \mathbf{f}$.
        Note that this class is an operator wrapper. The real implementation of the source term is in the `_NSPressureConvectionCore` class.

    Args:
        external_force: Optional[OperatorLike], optional, default=None
    """

    def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
        super().__init__(_NSPressureConvectionCore(external_force))
_de_aliasing_rate instance-attribute ¤
_de_aliasing_rate = 2 / 3
_state_dict instance-attribute ¤
_state_dict = {
    "f_mesh": None,
    "n_channel": None,
    "linear_coef": None,
    "nonlinear_func": None,
    "operator": None,
    "integrator": None,
    "invert_linear_coef": None,
}
operator_generators instance-attribute ¤
operator_generators = default(operator_generators, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_generators))
_nonlinear_funcs instance-attribute ¤
_nonlinear_funcs = []
_value_mesh_check_func instance-attribute ¤
_value_mesh_check_func = lambda dim_value, dim_mesh: True
_integrator instance-attribute ¤
_integrator = 'auto'
_integrator_config instance-attribute ¤
_integrator_config = {}
_is_etdrk_integrator instance-attribute ¤
_is_etdrk_integrator = True
is_linear property ¤
is_linear
set_de_aliasing_rate ¤
set_de_aliasing_rate(de_aliasing_rate: float)

Set the de-aliasing rate for the nonlinear operator. Args: de_aliasing_rate (float): De-aliasing rate. Default is ⅔.

Source code in torchfsm/operator/_base.py 
275
276
277
278
279
280
281
282
283
def set_de_aliasing_rate(self, de_aliasing_rate: float):
    r"""
    Set the de-aliasing rate for the nonlinear operator.
    Args:
        de_aliasing_rate (float): De-aliasing rate. Default is 2/3.
    """

    self._de_aliasing_rate = de_aliasing_rate
    self._state_dict = None
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
176
177
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
179
180
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
182
183
184
185
186
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
188
189
190
191
192
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
194
195
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
197
198
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
200
201
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
203
204
205
206
207
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
_build_linear_coefs ¤
_build_linear_coefs(
    linear_coefs: Optional[Sequence[LinearCoef]],
)

Build the linear coefficients based on the provided linear coefficient generators.

Parameters:

Name Type Description Default
linear_coefs Optional[Sequence[LinearCoef]]

List of linear coefficient generators.

 required
Source code in torchfsm/operator/_base.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
def _build_linear_coefs(
    self, linear_coefs: Optional[Sequence[LinearCoef]]
):
    r"""
    Build the linear coefficients based on the provided linear coefficient generators.

    Args:
        linear_coefs (Optional[Sequence[LinearCoef]]): List of linear coefficient generators.

    """
    if len(linear_coefs) == 0:
        linear_coefs = None
    else:
        linear_coefs = sum(
            [
                coef * op(self._state_dict["f_mesh"], self._state_dict["n_channel"])
                for coef, op in linear_coefs
            ]
        )
    self._state_dict["linear_coef"] = linear_coefs
_build_nonlinear_funcs ¤
_build_nonlinear_funcs(
    nonlinear_funcs: Optional[Sequence[NonlinearFunc]],
)

Build the nonlinear functions based on the provided nonlinear function generators.

Parameters:

Name Type Description Default
nonlinear_funcs Optional[Sequence[NonlinearFunc]]

List of nonlinear function generators.

 required
Source code in torchfsm/operator/_base.py 
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
def _build_nonlinear_funcs(
    self, nonlinear_funcs: Optional[Sequence[NonlinearFunc]]
):
    r"""
    Build the nonlinear functions based on the provided nonlinear function generators.

    Args:
        nonlinear_funcs (Optional[Sequence[NonlinearFunc]]): List of nonlinear function generators.
    """
    if len(nonlinear_funcs) == 0:
        nonlinear_funcs_all = None
    else:
        self._state_dict["f_mesh"].set_default_freq_threshold(
            self._de_aliasing_rate
        )

        def nonlinear_funcs_all(u_fft):
            result = 0.0
            dealiased_u_fft = None
            dealiased_u = None
            u = None
            for coef, fun in nonlinear_funcs:
                if fun._dealiasing_swtich:
                    if dealiased_u_fft is None:
                        dealiased_u_fft = u_fft * self._state_dict[
                            "f_mesh"
                        ].low_pass_filter(self._de_aliasing_rate)
                        dealiased_u = (
                            self._state_dict["f_mesh"].ifft(dealiased_u_fft).real
                        )
                    result += coef * fun(
                        dealiased_u_fft,
                        self._state_dict["f_mesh"],
                        dealiased_u,
                    )
                else:
                    if u is None:
                        u = self._state_dict["f_mesh"].ifft(u_fft).real
                    result += coef * fun(
                        u_fft,
                        self._state_dict["f_mesh"],
                        u,
                    )

            return result

    self._state_dict["nonlinear_func"] = nonlinear_funcs_all
_build_operator ¤
_build_operator()

Build the operator based on the linear coefficient and nonlinear function. If both linear coefficient and nonlinear function are None, the operator is set to None.

Source code in torchfsm/operator/_base.py 
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
def _build_operator(self):
    r"""
    Build the operator based on the linear coefficient and nonlinear function.
    If both linear coefficient and nonlinear function are None, the operator is set to None.
    """
    if self._state_dict["nonlinear_func"] is None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft
    elif self._state_dict["linear_coef"] is None:
        def operator(u_fft):
            return self._state_dict["nonlinear_func"](u_fft)
    elif self._state_dict["nonlinear_func"] is not None and self._state_dict["linear_coef"] is not None:
        def operator(u_fft):
            return self._state_dict["linear_coef"] * u_fft + self._state_dict[
                "nonlinear_func"
            ](u_fft)
    else:
        raise ValueError(
            "Both linear coefficient and nonlinear function are None. Cannot build operator."
        )

    self._state_dict["operator"] = operator
_build_integrator ¤
_build_integrator(dt: float)

Build the integrator based on the provided time step and integrator type.

Parameters:

Name Type Description Default
dt float

Time step for the integrator.

 required
Source code in torchfsm/operator/_base.py 
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
def _build_integrator(
    self,
    dt: float,
):
    r"""
    Build the integrator based on the provided time step and integrator type.

    Args:
        dt (float): Time step for the integrator.
    """
    if self._integrator == "auto":
        if self.is_linear:
            solver = ETDRKIntegrator.ETDRK0
        else:
            solver = ETDRKIntegrator.ETDRK4
    else:
        solver = self._integrator
    self._is_etdrk_integrator = isinstance(solver, ETDRKIntegrator)
    if self._is_etdrk_integrator:
        if solver == ETDRKIntegrator.ETDRK0:
            assert self.is_linear, "The ETDRK0 integrator only supports linear term"
            self._state_dict["integrator"] = solver.value(
                dt,
                self._state_dict["linear_coef"],
                **self._integrator_config,
            )
        else:
            if self._state_dict["linear_coef"] is None:
                linear_coef = torch.tensor(
                    [0.0],
                    dtype=self._state_dict["f_mesh"].dtype,
                    device=self._state_dict["f_mesh"].device,
                )
            else:
                linear_coef = self._state_dict["linear_coef"]
            self._state_dict["integrator"] = solver.value(
                dt,
                linear_coef,
                self._state_dict["nonlinear_func"],
                **self._integrator_config,
            )
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(u_fft),
        )
    elif isinstance(solver, RKIntegrator):
        if self._state_dict["operator"] is None:
            self._build_operator()
        self._state_dict["integrator"] = solver.value(**self._integrator_config)
        setattr(
            self._state_dict["integrator"],
            "forward",
            lambda u_fft, dt: self._state_dict["integrator"].step(
                self._state_dict["operator"], u_fft, dt
            ),
        )
_pre_check ¤
_pre_check(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ] = None,
) -> Tuple[FourierMesh, int]

Pre-check the input tensor and mesh. If the mesh is not registered, register it.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used.

 None

Returns:

Type Description
Tuple[FourierMesh, int]

Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
Source code in torchfsm/operator/_base.py 
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
def _pre_check(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh] = None,
) -> Tuple[FourierMesh, int]:
    r"""
    Pre-check the input tensor and mesh. If the mesh is not registered, register it.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used.

    Returns:
        Tuple[FourierMesh, int]: Tuple of Fourier mesh and number of channels.
    """

    if u_fft is None and u is None:
        raise ValueError("Either u or u_fft should be given")
    if u_fft is not None and u is not None:
        assert u.shape == u_fft.shape, "The shape of u and u_fft should be the same"
    assert mesh is not None, "Mesh should be given"
    value_device = u.device if u is not None else u_fft.device
    value_dtype = u.dtype if u is not None else u_fft.dtype
    if not isinstance(mesh, FourierMesh):
        if not isinstance(mesh, MeshGrid):
            mesh = FourierMesh(mesh, device=value_device, dtype=value_dtype)
        else:
            mesh = FourierMesh(mesh)
    n_channel = u.shape[1] if u is not None else u_fft.shape[1]
    value_shape = u.shape if u is not None else u_fft.shape
    assert (
        len(value_shape) == mesh.n_dim + 2
    ), f"the value shape {value_shape} is not compatible with mesh dim {mesh.n_dim}"
    for i in range(mesh.n_dim):
        assert (
            value_shape[i + 2] == mesh.mesh_info[i][2]
        ), f"Expect to have {mesh.mesh_info[i][2]} points in dim {i} but got {value_shape[i+2]}"
    assert (
        value_device == mesh.device
    ), "The device of mesh {} and the device of value {} are not the same".format(
        mesh.device, value_device
    )
    # assert value_dtype==mesh.dtype, "The dtype of mesh {} and the dtype of value {} are not the same".format(mesh.dtype,value_dtype)
    # value fft is a complex dtype
    assert self._value_mesh_check_func(
        len(value_shape) - 2, mesh.n_dim
    ), "Value and mesh do not match the requirement"
    return mesh, n_channel
register_mesh ¤
register_mesh(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    n_channel: int,
    device=None,
    dtype=None,
)

Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

Parameters:

Name Type Description Default
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
n_channel int

Number of channels of the input tensor.

 required
device Optional[device]

Device to which the mesh should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the mesh. Default is None.

 None
Source code in torchfsm/operator/_base.py 
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
def register_mesh(
    self,
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    n_channel: int,
    device=None,
    dtype=None,
):
    r"""
    Register the mesh and number of channels for the operator. Once a mesh is registered, mesh information is not required for integration and operator call.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): Mesh information or mesh object.
        n_channel (int): Number of channels of the input tensor.
        device (Optional[torch.device]): Device to which the mesh should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the mesh. Default is None.
    """
    if isinstance(mesh, FourierMesh):
        f_mesh = mesh
        if device is not None or dtype is not None:
            f_mesh.to(device=device, dtype=dtype)
    else:
        f_mesh = FourierMesh(mesh, device=device, dtype=dtype)
    for key in self._state_dict:
        self._state_dict[key] = None
    self._state_dict.update(
        {
            "f_mesh": f_mesh,
            "n_channel": n_channel,
        }
    )
    linear_coefs = []
    nonlinear_funcs = []
    for coef, generator in zip(self.coefs, self.operator_generators):
        op = generator(f_mesh, n_channel)
        if isinstance(op, LinearCoef):
            linear_coefs.append((coef, op))
        elif isinstance(op, NonlinearFunc):
            nonlinear_funcs.append((coef, op))
        else:
            raise ValueError(f"Operator {op} is not supported")
    self._nonlinear_funcs = nonlinear_funcs
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(self._nonlinear_funcs)
regisiter_additional_check ¤
regisiter_additional_check(
    func: Callable[[int, int], bool]
)

Register an additional check function for the value and mesh compatibility.

Parameters:

Name Type Description Default
func Callable[[int, int], bool]

Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.

 required
Source code in torchfsm/operator/_base.py 
589
590
591
592
593
594
595
596
def regisiter_additional_check(self, func: Callable[[int, int], bool]):
    r"""
    Register an additional check function for the value and mesh compatibility.

    Args:
        func (Callable[[int, int], bool]): Function that takes the dimension of the value and mesh as input and returns a boolean indicating whether they are compatible.
    """
    self._value_mesh_check_func = func
add_generator ¤
add_generator(generator: GeneratorLike, coef=1)

Add a generator to the operator.

Parameters:

Name Type Description Default
generator GeneratorLike

Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.

 required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
598
599
600
601
602
603
604
605
606
607
def add_generator(self, generator: GeneratorLike, coef=1):
    r"""
    Add a generator to the operator.

    Args:
        generator (GeneratorLike): Generator to be added. It should be a callable that takes a Fourier mesh and number of channels as input and returns a linear coefficient or nonlinear function.
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_generators.append(generator)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the integrator for the operator. The integrator is used for time integration of the operator.

Parameters:

Name Type Description Default
integrator Union[Literal['auto'], ETDRKIntegrator, RKIntegrator]

Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type. If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
def set_integrator(
    self,
    integrator: Union[Literal["auto"], ETDRKIntegrator, RKIntegrator],
    **integrator_config,
):
    r"""
    Set the integrator for the operator. The integrator is used for time integration of the operator.

    Args:
        integrator (Union[Literal["auto"], ETDRKIntegrator, RKIntegrator]): Integrator to be used. If "auto", the integrator will be chosen automatically based on the operator type.
            If "auto", the integrator will be set as ETDRKIntegrator.ETDRK0 for linear operators and ETDRKIntegrator.ETDRK4 for nonlinear operators.
        **integrator_config: Additional configuration for the integrator.
    """

    if isinstance(integrator, str):
        assert (
            integrator == "auto"
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    else:
        assert isinstance(integrator, ETDRKIntegrator) or isinstance(
            integrator, RKIntegrator
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,
) -> Union[
    SpatialTensor["B C H ..."],
    SpatialTensor["B T C H ..."],
    FourierTensor["B C H ..."],
    FourierTensor["B T C H ..."],
]

Integrate the operator using the provided initial condition and time step.

Parameters:

Name Type Description Default
u_0 Optional[Tensor]

Initial condition in spatial domain. Default is None.

 None
u_0_fft Optional[Tensor]

Initial condition in Fourier domain. Default is None. At least one of u_0 or u_0_fft should be provided.

 None
dt float

Time step for the integrator. Default is 1.

 1
step int

Number of time steps to integrate. Default is 1.

 1
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before integration.

 None
progressive bool

If True, show a progress bar during integration. Default is False.

 False
trajectory_recorder Optional[_TrajRecorder]

Trajectory recorder for recording the trajectory during integration. Default is None. If None, no trajectory will be recorded. The function will only return the final frame.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], SpatialTensor['B T C H ...'], FourierTensor['B C H ...'], FourierTensor['B T C H ...']]

Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain. If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...). If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.
Source code in torchfsm/operator/_base.py 
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
def integrate(
    self,
    u_0: Optional[torch.Tensor] = None,
    u_0_fft: Optional[torch.Tensor] = None,
    dt: float = 1,
    step: int = 1,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    progressive: bool = False,
    trajectory_recorder: Optional[_TrajRecorder] = None,
    return_in_fourier: bool = False,

) -> Union[
        SpatialTensor["B C H ..."],
        SpatialTensor["B T C H ..."],
        FourierTensor["B C H ..."],
        FourierTensor["B T C H ..."],
    ]:
    r"""
    Integrate the operator using the provided initial condition and time step.  

    Args:
        u_0 (Optional[torch.Tensor]): Initial condition in spatial domain. Default is None.
        u_0_fft (Optional[torch.Tensor]): Initial condition in Fourier domain. Default is None.
            At least one of u_0 or u_0_fft should be provided.
        dt (float): Time step for the integrator. Default is 1.
        step (int): Number of time steps to integrate. Default is 1.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before integration.
        progressive (bool): If True, show a progress bar during integration. Default is False.
        trajectory_recorder (Optional[_TrajRecorder]): Trajectory recorder for recording the trajectory during integration. Default is None.
            If None, no trajectory will be recorded. The function will only return the final frame.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], SpatialTensor["B T C H ..."], FourierTensor["B C H ..."], FourierTensor["B T C H ..."]]: Integrated result in spatial or Fourier domain.
            If trajectory_recorder is provided, the result will be a trajectory tensor of shape (B, T, C, H, ...). Otherwise, the result will be a tensor of shape (B, C, H, ...).
            If return_in_fourier is True, the result will be in Fourier domain. Otherwise, it will be in spatial domain.

    """
    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u=u_0, u_fft=u_0_fft, mesh=mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u_0, u_fft=u_0_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["integrator"] is None:
        self._build_integrator(dt)
    elif self._is_etdrk_integrator:
        if self._state_dict["integrator"].dt != dt:
            self._build_integrator(dt)
    f_mesh = self._state_dict["f_mesh"]
    if u_0_fft is None:
        u_0_fft = f_mesh.fft(u_0)
    p_bar = tqdm(range(step), desc="Integrating", disable=not progressive)
    for i in p_bar:
        if trajectory_recorder is not None:
            trajectory_recorder.record(i, u_0_fft)
        u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
    if trajectory_recorder is not None:
        trajectory_recorder.record(i + 1, u_0_fft)
        trajectory_recorder.return_in_fourier = return_in_fourier
        return trajectory_recorder.trajectory
    else:
        if return_in_fourier:
            return u_0_fft
        else:
            return f_mesh.ifft(u_0_fft).real
__call__ ¤
__call__(
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[
            Sequence[tuple[float, float, int]],
            MeshGrid,
            FourierMesh,
        ]
    ] = None,
    return_in_fourier=False,
) -> Union[
    SpatialTensor["B C H ..."], FourierTensor["B C H ..."]
]

Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

Parameters:

Name Type Description Default
u Optional[SpatialTensor]

Input tensor in spatial domain. Default is None.

 None
u_fft Optional[FourierTensor]

Input tensor in Fourier domain. Default is None. At least one of u or u_fft should be provided.

 None
mesh Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]

Mesh information or mesh object. Default is None. If None, the mesh registered in the operator will be used. You can use register_mesh to register a mesh before calling the operator.

 None
return_in_fourier bool

If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

 False

Returns:

Type Description
Union[SpatialTensor['B C H ...'], FourierTensor['B C H ...']]

Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
Source code in torchfsm/operator/_base.py 
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
def __call__(
    self,
    u: Optional[SpatialTensor["B C H ..."]] = None,
    u_fft: Optional[FourierTensor["B C H ..."]] = None,
    mesh: Optional[
        Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]
    ] = None,
    return_in_fourier=False,
) -> Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]:
    r"""
    Call the operator with the provided input tensor. The operator will apply the linear coefficient and nonlinear function to the input tensor.

    Args:
        u (Optional[SpatialTensor]): Input tensor in spatial domain. Default is None.
        u_fft (Optional[FourierTensor]): Input tensor in Fourier domain. Default is None.
            At least one of u or u_fft should be provided.
        mesh (Optional[Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]]): Mesh information or mesh object. Default is None.
            If None, the mesh registered in the operator will be used. You can use `register_mesh` to register a mesh before calling the operator.
        return_in_fourier (bool): If True, return the result in Fourier domain. If False, return the result in spatial domain. Default is False.

    Returns:
        Union[SpatialTensor["B C H ..."], FourierTensor["B C H ..."]]: Result of the operator in spatial or Fourier domain.
    """    

    if self._state_dict["f_mesh"] is None or mesh is not None:
        mesh, n_channel = self._pre_check(u, u_fft, mesh)
        self.register_mesh(mesh, n_channel)
    else:
        self._pre_check(u=u, u_fft=u_fft, mesh=self._state_dict["f_mesh"])
    if self._state_dict["operator"] is None:
        self._build_operator()
    if u_fft is None:
        u_fft = self._state_dict["f_mesh"].fft(u)
    value_fft = self._state_dict["operator"](u_fft)
    if return_in_fourier:
        return value_fft
    else:
        return self._state_dict["f_mesh"].ifft(value_fft).real
to ¤
to(device=None, dtype=None)

Move the operator to the specified device and change the data type.

Parameters:

Name Type Description Default
device Optional[device]

Device to which the operator should be moved. Default is None.

 None
dtype Optional[dtype]

Data type of the operator. Default is None.

 None
Source code in torchfsm/operator/_base.py 
743
744
745
746
747
748
749
750
751
752
753
def to(self, device=None, dtype=None):
    r"""
    Move the operator to the specified device and change the data type.

    Args:
        device (Optional[torch.device]): Device to which the operator should be moved. Default is None.
        dtype (Optional[torch.dtype]): Data type of the operator. Default is None.
    """
    if self._state_dict is not None:
        self._state_dict["f_mesh"].to(device=device, dtype=dtype)
        self.register_mesh(self._state_dict["f_mesh"], self._state_dict["n_channel"])
__add__ ¤
__add__(other)
Source code in torchfsm/operator/_base.py 
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_generators + other.operator_generators,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_generators
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
929
930
931
932
933
934
935
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_generators, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
937
938
939
940
def __neg__(self):
    return NonlinearOperator(
        self.operator_generators, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes.py 
267
268
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_NSPressureConvectionCore(external_force))

Utils¤

torchfsm.operator.run_operators ¤

run_operators(
    u: SpatialTensor["B C H ..."],
    operators: Sequence[Operator],
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
) -> SpatialTensor["B C H ..."]

Run a sequence of operators on the input tensor.

Parameters:

Name Type Description Default
u SpatialTensor

Input tensor of shape (B, C, H, ...).

 required
operators Sequence[Operator]

Sequence of operators to be applied.

 required
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Resulting tensor after applying the operators.
Source code in torchfsm/operator/__init__.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def run_operators(u:SpatialTensor["B C H ..."],
                  operators:Sequence[Operator],
                  mesh: Union[Sequence[tuple[float, float, int]],MeshGrid,FourierMesh]) -> SpatialTensor["B C H ..."]:
    r"""
    Run a sequence of operators on the input tensor.

    Args:
        u (SpatialTensor): Input tensor of shape (B, C, H, ...).
        operators (Sequence[Operator]): Sequence of operators to be applied.
        mesh (Union[Sequence[tuple[float, float, int]],MeshGrid,FourierMesh]): Mesh information or mesh object.

    Returns:
        SpatialTensor: Resulting tensor after applying the operators.
    """
    if not isinstance(mesh,FourierMesh):
        mesh=FourierMesh(mesh)
    u_fft=mesh.fft(u)
    def _run_operator(operator:Operator):
        return operator(u_fft=u_fft,mesh=mesh)
    return map(_run_operator,operators)

torchfsm.operator.check_value_with_mesh ¤

check_value_with_mesh(
    u: SpatialTensor["B C H ..."],
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
)

Check if the value and mesh are compatible. If not, raise a ValueError.

Parameters:

Name Type Description Default
u SpatialTensor

Input tensor of shape (B, C, H, ...).

 required
mesh Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]

Mesh information or mesh object.

 required
Source code in torchfsm/operator/_base.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def check_value_with_mesh(
    u: SpatialTensor["B C H ..."],
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
):
    r"""
    Check if the value and mesh are compatible. If not, raise a ValueError.

    Args:
        u (SpatialTensor): Input tensor of shape (B, C, H, ...).
        mesh (Union[Sequence[tuple[float, float, int]],MeshGrid,FourierMesh]): Mesh information or mesh object.
    """
    if isinstance(mesh, FourierMesh) or isinstance(mesh, MeshGrid):
        mesh = mesh.mesh_info
    n_dim = len(mesh)
    value_shape = u.shape
    if len(value_shape) - 2 != n_dim:
        raise ValueError(
            f"the value shape {value_shape} is not compatible with mesh dim {n_dim}"
        )
    for i in range(n_dim):
        if value_shape[i + 2] != mesh[i][2]:
            raise ValueError(
                f"Expect to have {mesh[i][2]} points in dim {i} but got {value_shape[i+2]}"
            )