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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-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 
70
71
72
73
74
75
76
77
78
79
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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_convection.py 
78
79
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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-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 
47
48
49
50
51
52
53
54
55
56
57
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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_div.py 
56
57
def __init__(self) -> None:
    super().__init__(_DivGenerator())

torchfsm.operator.Dispersion ¤

Bases: LinearOperator

Dispersion calculates the Laplacian of a vector field.

It is defined as \(\nabla \cdot (\nabla^2\mathbf{u}) = \left[\begin{matrix}\sum_j^I \frac{\partial}{\partial j}\sum_i^I \frac{\partial^2 u_x}{\partial i^2 } \\ \sum_j^I \frac{\partial}{\partial j}\sum_i^I \frac{\partial^2 u_y}{\partial i^2 } \\ \cdots \\ \sum_j^I \frac{\partial}{\partial j}\sum_i^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/_dispersion.py 
21
22
23
24
25
26
27
28
29
30
class Dispersion(LinearOperator):
    r"""
    `Dispersion` calculates the Laplacian of a vector field.

    It is defined as $\nabla \cdot (\nabla^2\mathbf{u}) = \left[\begin{matrix}\sum_j^I \frac{\partial}{\partial j}\sum_i^I \frac{\partial^2 u_x}{\partial i^2 } \\ \sum_j^I \frac{\partial}{\partial j}\sum_i^I \frac{\partial^2 u_y}{\partial i^2 } \\ \cdots \\ \sum_j^I \frac{\partial}{\partial j}\sum_i^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__(_DispersionCore())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_dispersion.py 
29
30
def __init__(self) -> None:
    super().__init__(_DispersionCore())

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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-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 
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
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))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(source: Tensor) -> None
Source code in torchfsm/operator/_base.py 
1025
1026
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)
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
840
841
842
843
844
845
846
847
848
849
850
851
852
853
def __add__(self, other):
    if isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
855
856
857
858
859
860
861
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return Operator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
863
864
def __neg__(self):
    return Operator(self.operator_cores, [-1 * coef for coef in self.coefs])
__init__ ¤
__init__(
    source_func: Optional[
        Callable[[Tensor], 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.HyperDiffusion ¤

Bases: LinearOperator

HyperDiffusion calculates the hyper diffusion of a vector field.

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

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

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

    def __init__(self) -> None:
        super().__init__(_HyperDiffusionCore())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/generic/_hyper_diffusion.py 
25
26
def __init__(self) -> None:
    super().__init__(_HyperDiffusionCore())

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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-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))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-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))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-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/_vorticity_convection.py 
57
58
59
60
61
62
63
64
65
66
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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_vorticity_convection.py 
65
66
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/_value_transformation.py 
49
50
51
52
53
54
55
56
57
58
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())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__()
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
57
58
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/_value_transformation.py 
105
106
107
108
109
110
111
112
113
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))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
112
113
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/_value_transformation.py 
152
153
154
155
156
157
158
159
160
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))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
159
160
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/_ns_pressure_convection.py 
50
51
52
53
54
55
56
57
58
59
60
61
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))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_ns_pressure_convection.py 
60
61
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_NSPressureConvectionCore(external_force))

torchfsm.operator.Leray ¤

Bases: NonlinearOperator

Leray calculates the Leray projection of a vector field. It is defined as \(\mathbf{u} - \nabla \nabla^{-2} \nabla \cdot \mathbf{u}\). 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 _LerayCore class.

Source code in torchfsm/operator/dedicated/_navier_stokes/_leray.py 
42
43
44
45
46
47
48
49
50
51
class Leray(NonlinearOperator):
    r"""
    `Leray` calculates the Leray projection of a vector field.
        It is defined as $\mathbf{u} - \nabla \nabla^{-2} \nabla \cdot \mathbf{u}$.
        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 `_LerayCore` class.
    """

    def __init__(self) -> None:
        super().__init__(_LerayGenerator())
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_leray.py 
50
51
def __init__(self) -> None:
    super().__init__(_LerayGenerator())

torchfsm.operator.ChannelWisedDiffusion ¤

Bases: LinearOperator

The ChannelWisedDiffusion operator applies different diffusion coefficients to each channel of a field. It is defined as: \(\nabla^2 \phi_i = \nu_i \nabla^2 \phi_i\) for each channel \(i\). Note that this class is an operator wrapper. The real implementation of the diffusion is in the _ChannelWisedDiffusionCore class.

Source code in torchfsm/operator/dedicated/_gray_scott.py 
38
39
40
41
42
43
44
45
46
47
class ChannelWisedDiffusion(LinearOperator):

    r"""
    The ChannelWisedDiffusion operator applies different diffusion coefficients to each channel of a field.
    It is defined as: $\nabla^2 \phi_i = \nu_i \nabla^2 \phi_i$ for each channel $i$.
    Note that this class is an operator wrapper. The real implementation of the diffusion is in the `_ChannelWisedDiffusionCore` class.
    """

    def __init__(self,viscosities: Sequence[Union[Tensor, float]]):
        super().__init__(_ChannelWisedDiffusionGenerator(viscosities))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
solve ¤
solve(
    b: Optional[Tensor] = None,
    b_fft: Optional[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 
216
217
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
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 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
def __add__(self, other):
    if isinstance(other, LinearOperator):
        return LinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
914
915
916
917
918
919
920
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return LinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
922
923
924
925
def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(viscosities: Sequence[Union[Tensor, float]])
Source code in torchfsm/operator/dedicated/_gray_scott.py 
46
47
def __init__(self,viscosities: Sequence[Union[Tensor, float]]):
    super().__init__(_ChannelWisedDiffusionGenerator(viscosities))

torchfsm.operator.GrayScottSource ¤

Bases: NonlinearOperator

The Gray-Scott source term operator for a two-channel field. It is defined as: \(\left[\begin{matrix}f (1 - \phi_0) - \phi_0 \phi_1^2 \\ \phi_0 \phi_1^2 - (f + k) \phi_1 \end{matrix}\right]\) Note that this class is an operator wrapper. The real implementation of the source term is in the _GrayScottSource class.

Source code in torchfsm/operator/dedicated/_gray_scott.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
class GrayScottSource(NonlinearOperator):

    r"""
    The Gray-Scott source term operator for a two-channel field.
    It is defined as: $\left[\begin{matrix}f (1 - \phi_0) - \phi_0 \phi_1^2 \\ \phi_0 \phi_1^2 - (f + k) \phi_1 \end{matrix}\right]$
    Note that this class is an operator wrapper. The real implementation of the source term is in the `_GrayScottSource` class.
    """

    def __init__(self, 
                 feed_rate: Union[Tensor, float],
                 kill_rate: Union[Tensor, float]) -> None:
        super().__init__(_GrayScottSourceGenerator(feed_rate, kill_rate))
operator_cores instance-attribute ¤
operator_cores = default(operator_cores, [])
coefs instance-attribute ¤
coefs = default(coefs, [1] * len(operator_cores))
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 
272
273
274
275
276
277
278
279
280
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 = {key:None for key in self._state_dict.keys()}
__radd__ ¤
__radd__(other)
Source code in torchfsm/operator/_base.py 
174
175
def __radd__(self, other):
    return self + other
__iadd__ ¤
__iadd__(other)
Source code in torchfsm/operator/_base.py 
177
178
def __iadd__(self, other):
    return self + other
__sub__ ¤
__sub__(other)
Source code in torchfsm/operator/_base.py 
180
181
182
183
184
def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented
__rsub__ ¤
__rsub__(other)
Source code in torchfsm/operator/_base.py 
186
187
188
189
190
def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented
__isub__ ¤
__isub__(other)
Source code in torchfsm/operator/_base.py 
192
193
def __isub__(self, other):
    return self - other
__rmul__ ¤
__rmul__(other)
Source code in torchfsm/operator/_base.py 
195
196
def __rmul__(self, other):
    return self * other
__imul__ ¤
__imul__(other)
Source code in torchfsm/operator/_base.py 
198
199
def __imul__(self, other):
    return self * other
__truediv__ ¤
__truediv__(other)
Source code in torchfsm/operator/_base.py 
201
202
203
204
205
def __truediv__(self, other):
    try:
        return self * (1 / other)
    except:
        return NotImplemented
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 
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
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, core in zip(self.coefs, self.operator_cores):
        op = (
            core(f_mesh, n_channel)
            if (
                not isinstance(core, LinearCoef)
                and not isinstance(core, NonlinearFunc)
            )
            else core
        )
        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")
    clean_up_memory()
    self._build_linear_coefs(linear_coefs)
    self._build_nonlinear_funcs(nonlinear_funcs)
register_additional_check ¤
register_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 
630
631
632
633
634
635
636
637
def register_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_core ¤
add_core(
    core: Union[LinearCoef, NonlinearFunc, GeneratorLike],
    coef=1,
)

Add a generator to the operator.

Parameters:

Name Type Description Default
core Union[LinearCoef, NonlinearFunc, GeneratorLike]

Core to be added.

required
coef float

Coefficient for the generator. Default is 1.

 1
Source code in torchfsm/operator/_base.py 
639
640
641
642
643
644
645
646
647
648
def add_core(self,core:Union[LinearCoef,NonlinearFunc,GeneratorLike],coef=1):
    r"""
    Add a generator to the operator.

    Args:
        core (Union[LinearCoef,NonlinearFunc,GeneratorLike]): Core to be added. 
        coef (float): Coefficient for the generator. Default is 1.
    """
    self.operator_cores.append(core)
    self.coefs.append(coef)
set_integrator ¤
set_integrator(
    integrator: Union[
        Literal["auto"],
        ETDRKIntegrator,
        SETDRKIntegrator,
        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, SETDRKIntegrator, 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.ETDRK2 for nonlinear operators.

 required
**integrator_config

Additional configuration for the integrator.

 {}
Source code in torchfsm/operator/_base.py 
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
def set_integrator(
    self,
    integrator: Union[
        Literal["auto"], ETDRKIntegrator, SETDRKIntegrator, 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, SETDRKIntegrator, 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.ETDRK2 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, SETDRKIntegrator or RKIntegrator"
    else:
        assert (
            isinstance(integrator, ETDRKIntegrator)
            or isinstance(integrator, SETDRKIntegrator)
            or isinstance(integrator, RKIntegrator)
        ), "The integrator should be 'auto' or an instance of ETDRKIntegrator, SETDRKIntegrator or RKIntegrator"
    self._integrator = integrator
    self._integrator_config = integrator_config
    self._state_dict["integrator"] = None
integrate ¤
integrate(
    u_0: Optional[Tensor] = None,
    u_0_fft: Optional[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,
    nan_check: 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
Nan_check bool

If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. Default is False.

 required

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 
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
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
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
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,
    nan_check: 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.
        Nan_check (bool): If True, check for NaN values in the result. If NaN values are found, raise a NanSimulationError. 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)
    clean_up_memory()
    try:
        for i in p_bar:
            if trajectory_recorder is not None:
                trajectory_recorder.record(i, u_0_fft)
                if trajectory_recorder._shutdown_flag:
                    break
            u_0_fft = self._state_dict["integrator"].forward(u_0_fft, dt)
            if nan_check:
                if torch.isnan(u_0_fft).any():
                    raise NanSimulationError(
                        f"NaN values found in the result at step {i}. Please check the input and simulation parameters."
                    )
    except TorchOutOfMemoryError as e:
        error_msg = [
            "Cuda out of memory when integrating the operator.",
            "Original error message: {}".format(str(e)),
            "Please try to use a smaller mesh or a low-order integrator.",
        ]
        raise OutOfMemoryError(os.linesep.join(error_msg))
    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 
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
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 
805
806
807
808
809
810
811
812
813
814
815
816
817
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 
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
def __add__(self, other):
    if isinstance(other, NonlinearOperator):
        return NonlinearOperator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, OperatorLike):
        return Operator(
            self.operator_cores + other.operator_cores,
            self.coefs + other.coefs,
        )
    elif isinstance(other, Tensor):
        return Operator(
            self.operator_cores
            + [lambda f_mesh, n_channel: _ExplicitSourceCore(other)],
            self.coefs + [1],
        )
    else:
        return NotImplemented
__mul__ ¤
__mul__(other)
Source code in torchfsm/operator/_base.py 
975
976
977
978
979
980
981
def __mul__(self, other):
    if isinstance(other, OperatorLike):
        return NotImplemented
    else:
        return NonlinearOperator(
            self.operator_cores, [coef * other for coef in self.coefs]
        )
__neg__ ¤
__neg__()
Source code in torchfsm/operator/_base.py 
983
984
985
986
def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )
__init__ ¤
__init__(
    feed_rate: Union[Tensor, float],
    kill_rate: Union[Tensor, float],
) -> None
Source code in torchfsm/operator/dedicated/_gray_scott.py 
102
103
104
105
def __init__(self, 
             feed_rate: Union[Tensor, float],
             kill_rate: Union[Tensor, float]) -> None:
    super().__init__(_GrayScottSourceGenerator(feed_rate, kill_rate))

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 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
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]}"
            )