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.Advection ¤

Bases: NonlinearOperator

Advection calculates the advection of a scalar field by a constant velocity field. If your velocity field is constant in space, please consider using LinearAdvection operator to allow you use larger simulation dt. It is defined as $\nabla \cdot (\phi\mathbf{u}) = \sum_{i=0}^I \frac{\phi\partial u_i}{\partial i}$ where $\mathbf{u}$ is the velocity field. Note that this class is an operator wrapper. The actual implementation of the operator is in the _AdvectionCore class.

Parameters:

Name	Type	Description	Default
`velocity`	`SpatialTensor['B C H ...']`	The velocity field used for advection. Please not that your velocity should be smooth enough to avoid aliasing error in Fourier space.	required

Source code in torchfsm/operator/generic/_advection.py

class Advection(NonlinearOperator):
    r"""
    `Advection` calculates the advection of a scalar field by a constant velocity field.
        If your velocity field is constant in space, please consider using `LinearAdvection` operator to allow you use larger simulation dt.
        It is defined as $\nabla \cdot (\phi\mathbf{u}) = \sum_{i=0}^I \frac{\phi\partial u_i}{\partial i}$
        where $\mathbf{u}$ is the velocity field.
        Note that this class is an operator wrapper. The actual implementation of the operator is in the `_AdvectionCore` class.

    Args:
        velocity (SpatialTensor["B C H ..."]): The velocity field used for advection. Please not that your velocity should be smooth enough to avoid aliasing error in Fourier space.
    """

    def __init__(self,
                 velocity:SpatialTensor["B C H ..."]
                 ) -> None:
        super().__init__(_AdvectionCore(velocity))

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

def __neg__(self):
    return NonlinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )

init ¤

__init__(velocity: SpatialTensor['B C H ...']) -> None

Source code in torchfsm/operator/generic/_advection.py

def __init__(self,
             velocity:SpatialTensor["B C H ..."]
             ) -> None:
    super().__init__(_AdvectionCore(velocity))

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_BiharmonicCore())

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self,viscosities: Sequence[Union[Tensor, float]]):
    super().__init__(_ChannelWisedDiffusionGenerator(viscosities))

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_CurlGenerator())

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_DispersionCore())

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_DivGenerator())

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, source: torch.Tensor) -> None:
    super().__init__(_ExplicitSourceCore(source))

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_GradGenerator())

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, 
             feed_rate: Union[Tensor, float],
             kill_rate: Union[Tensor, float]) -> None:
    super().__init__(_GrayScottSourceGenerator(feed_rate, kill_rate))

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_HyperDiffusionCore())

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

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.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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, remove_mean: bool = True) -> None:
    super().__init__(_KSConvectionGenerator(remove_mean))

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__( _LaplacianCore())

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_LerayGenerator())

torchfsm.operator.LinearAdvection ¤

Bases: LinearOperator

LinearAdvection calculates the advection of a scalar field by a constant velocity field. It is defined as $ \mathbf{u} \cdot \nabla \phi = \sum_{i=0}^I u_i \frac{\partial \phi}{\partial i} $ where $\mathbf{u} = [u_x, u_y, \cdots, u_I]$ is the constant velocity field. If your velocity is not constant in space, please consider using Advection operator instead. Note that this class is an operator wrapper. The actual implementation of the operator is in the _AdvectionCore class.

Parameters:

Name	Type	Description	Default
`velocity`	`Union[float, Sequence[float]]`	The constant velocity field. If a float is provided, it is assumed that the velocity is the same in all dimensions. If a sequence is provided, its length must match the number of dimensions of the mesh. Default is 1.0.	`1.0`

Source code in torchfsm/operator/generic/_linear_advection.py

class LinearAdvection(LinearOperator):
    r"""
    `LinearAdvection` calculates the advection of a scalar field by a constant velocity field.
        It is defined as $ \mathbf{u} \cdot \nabla \phi = \sum_{i=0}^I u_i \frac{\partial \phi}{\partial i} $
        where $\mathbf{u} = [u_x, u_y, \cdots, u_I]$ is the constant velocity field.
        If your velocity is not constant in space, please consider using `Advection` operator instead.
        Note that this class is an operator wrapper. The actual implementation of the operator is in the `_AdvectionCore` class.

    Args:
        velocity (Union[float, Sequence[float]]): The constant velocity field.
            If a float is provided, it is assumed that the velocity is the same in all dimensions.
            If a sequence is provided, its length must match the number of dimensions of the mesh.
            Default is 1.0.
    """

    def __init__(self, velocity: Union[float, Sequence[float]] = 1.0) -> None:
        super().__init__(_LinearAdvectionCore(velocity))

    def set_advection_velocity(self, velocity: Union[float, Sequence[float]]) -> None:
        """Set the advection velocity.

        Args:
            velocity (Union[float, Sequence[float]]): The constant velocity field.
                If a float is provided, it is assumed that the velocity is the same in all dimensions.
                If a sequence is provided, its length must match the number of dimensions of the mesh.
        """
        self._core.set_advection_velocity(velocity)

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

def __neg__(self):
    return LinearOperator(
        self.operator_cores, [-1 * coef for coef in self.coefs]
    )

init ¤

__init__(
    velocity: Union[float, Sequence[float]] = 1.0,
) -> None

Source code in torchfsm/operator/generic/_linear_advection.py

def __init__(self, velocity: Union[float, Sequence[float]] = 1.0) -> None:
    super().__init__(_LinearAdvectionCore(velocity))

set_advection_velocity ¤

set_advection_velocity(
    velocity: Union[float, Sequence[float]],
) -> None

Set the advection velocity.

Parameters:

Name	Type	Description	Default
`velocity`	`Union[float, Sequence[float]]`	The constant velocity field. If a float is provided, it is assumed that the velocity is the same in all dimensions. If a sequence is provided, its length must match the number of dimensions of the mesh.	required

Source code in torchfsm/operator/generic/_linear_advection.py

def set_advection_velocity(self, velocity: Union[float, Sequence[float]]) -> None:
    """Set the advection velocity.

    Args:
        velocity (Union[float, Sequence[float]]): The constant velocity field.
            If a float is provided, it is assumed that the velocity is the same in all dimensions.
            If a sequence is provided, its length must match the number of dimensions of the mesh.
    """
    self._core.set_advection_velocity(velocity)

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_NSPressureConvectionCore(external_force))

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, dim_index: int, order: int) -> None:
    super().__init__(_SpatialDerivativeGenerator(dim_index, order))

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_Velocity2PressureCore(external_force))

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(_Vorticity2PressureGenerator(external_force))

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

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self):
    super().__init__(_Vorticity2VelocityGenerator())

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

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

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

def __radd__(self, other):
    return self + other

iadd ¤

__iadd__(other)

Source code in torchfsm/operator/_base.py

def __iadd__(self, other):
    return self + other

sub ¤

__sub__(other)

Source code in torchfsm/operator/_base.py

def __sub__(self, other):
    try:
        return self + (-1 * other)
    except Exception:
        return NotImplemented

rsub ¤

__rsub__(other)

Source code in torchfsm/operator/_base.py

def __rsub__(self, other):
    try:
        return other + (-1 * self)
    except Exception:
        return NotImplemented

isub ¤

__isub__(other)

Source code in torchfsm/operator/_base.py

def __isub__(self, other):
    return self - other

rmul ¤

__rmul__(other)

Source code in torchfsm/operator/_base.py

def __rmul__(self, other):
    return self * other

imul ¤

__imul__(other)

Source code in torchfsm/operator/_base.py

def __imul__(self, other):
    return self * other

truediv ¤

__truediv__(other)

Source code in torchfsm/operator/_base.py

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

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

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

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

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

set_default_nonlinear_integrator ¤

set_default_nonlinear_integrator(
    integrator: Union[
        ETDRKIntegrator, SETDRKIntegrator, RKIntegrator
    ],
    **integrator_config
)

Set the default nonlinear integrator for the operator.

Parameters:

Name	Type	Description	Default
`integrator`	`Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]`	Integrator to be used.	required
`**integrator_config`		Additional configuration for the integrator.	`{}`

Source code in torchfsm/operator/_base.py

def set_default_nonlinear_integrator(
    self,
    integrator: Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator],
    **integrator_config,
    ):
    r"""
    Set the default nonlinear integrator for the operator.

    Args:
        integrator (Union[ETDRKIntegrator, SETDRKIntegrator, RKIntegrator]): Integrator to be used.
        **integrator_config: Additional configuration for the integrator.

    """
    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._default_nonlinear_integrator = integrator
    self._default_nonlinear_integrator_config = integrator_config
    if self._integrator == "auto":
        self._state_dict["integrator"] = None
    else:
        warn("Not automatically setting the integrator. The default nonlinear integrator will only be used if the integrator is set to 'auto'.")

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.	`False`

Returns:

Type Description

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

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

Source code in torchfsm/operator/_base.py

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

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

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

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

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

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

def __init__(self) -> None:
    super().__init__(_VorticityConvectionGenerator())

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

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

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]}"
            )

1. operator

Operators¤

torchfsm.operator.Advection ¤

operator_cores instance-attribute ¤

coefs instance-attribute ¤

is_linear property ¤

set_de_aliasing_rate ¤

__radd__ ¤

__iadd__ ¤

__sub__ ¤

__rsub__ ¤

__isub__ ¤

__rmul__ ¤

__imul__ ¤

__truediv__ ¤

register_mesh ¤

register_additional_check ¤

add_core ¤

set_integrator ¤

set_default_nonlinear_integrator ¤

integrate ¤

__call__ ¤

to ¤

__add__ ¤

__mul__ ¤

__neg__ ¤

__init__ ¤

torchfsm.operator.Biharmonic ¤

operator_cores instance-attribute ¤

coefs instance-attribute ¤

is_linear property ¤

register_mesh ¤

solve ¤

__radd__ ¤

__iadd__ ¤

__sub__ ¤

__rsub__ ¤

__isub__ ¤

__rmul__ ¤

__imul__ ¤

__truediv__ ¤

register_additional_check ¤

add_core ¤

set_integrator ¤

set_default_nonlinear_integrator ¤

integrate ¤

__call__ ¤

to ¤

__add__ ¤

__mul__ ¤

__neg__ ¤

__init__ ¤

torchfsm.operator.ChannelWisedDiffusion ¤

operator_cores instance-attribute ¤

coefs instance-attribute ¤

is_linear property ¤

register_mesh ¤

solve ¤

__radd__ ¤

__iadd__ ¤

__sub__ ¤

__rsub__ ¤

__isub__ ¤

__rmul__ ¤

__imul__ ¤

__truediv__ ¤

register_additional_check ¤

add_core ¤

set_integrator ¤

set_default_nonlinear_integrator ¤

integrate ¤

__call__ ¤

to ¤

__add__ ¤

__mul__ ¤

__neg__ ¤

__init__ ¤

torchfsm.operator.ConservativeConvection ¤

operator_cores instance-attribute ¤

coefs instance-attribute ¤

operator_cores `instance-attribute` ¤

coefs `instance-attribute` ¤

is_linear `property` ¤

radd ¤

iadd ¤

sub ¤

rsub ¤

isub ¤

rmul ¤

imul ¤

truediv ¤

call ¤

add ¤

mul ¤

neg ¤

init ¤

operator_cores `instance-attribute` ¤

coefs `instance-attribute` ¤

is_linear `property` ¤

radd ¤

iadd ¤

sub ¤

rsub ¤

isub ¤

rmul ¤

imul ¤

truediv ¤

call ¤

add ¤

mul ¤

neg ¤

init ¤

operator_cores `instance-attribute` ¤

coefs `instance-attribute` ¤

is_linear `property` ¤

radd ¤

iadd ¤

sub ¤

rsub ¤

isub ¤

rmul ¤

imul ¤

truediv ¤

call ¤

add ¤

mul ¤

neg ¤

init ¤

operator_cores `instance-attribute` ¤

coefs `instance-attribute` ¤

is_linear `property` ¤

radd ¤

iadd ¤

sub ¤

rsub ¤

isub ¤

rmul ¤

imul ¤

truediv ¤

call ¤

add ¤

mul ¤

neg ¤

init ¤

operator_cores `instance-attribute` ¤

coefs `instance-attribute` ¤

is_linear `property` ¤

radd ¤

iadd ¤

sub ¤

rsub ¤

isub ¤

rmul ¤

imul ¤

truediv ¤

call ¤

add ¤

mul ¤

neg ¤

init ¤