Skip to content

Cores

// ...existing code...

torchfsm.operator.generic._advection._AdvectionCore ¤

Bases: NonlinearFunc

Implementation of the Advection operator.

Source code in torchfsm/operator/generic/_advection.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
class _AdvectionCore(NonlinearFunc):
    r"""
    Implementation of the Advection operator.
    """

    def __init__(self,
                 velocity:SpatialTensor["B C H ..."]):
        super().__init__(False)
        self.velocity = velocity
        self.div=_DivCore()

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        assert self.velocity.shape[1] == f_mesh.n_dim, \
            f"advection operator only works for scalar field with velocity dimension {f_mesh.n_dim}, "\
            f"but got velocity with dimension {self.velocity.shape[1]}"
        if u is None:
            u = f_mesh.ifft(u_fft).real
        return self.div(f_mesh.fft(self.velocity*u),f_mesh)
velocity instance-attribute ¤
velocity = velocity
div instance-attribute ¤
div = _DivCore()
__init__ ¤
__init__(velocity: SpatialTensor['B C H ...'])
Source code in torchfsm/operator/generic/_advection.py 
13
14
15
16
17
def __init__(self,
             velocity:SpatialTensor["B C H ..."]):
    super().__init__(False)
    self.velocity = velocity
    self.div=_DivCore()
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_advection.py 
19
20
21
22
23
24
25
26
27
28
29
30
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    assert self.velocity.shape[1] == f_mesh.n_dim, \
        f"advection operator only works for scalar field with velocity dimension {f_mesh.n_dim}, "\
        f"but got velocity with dimension {self.velocity.shape[1]}"
    if u is None:
        u = f_mesh.ifft(u_fft).real
    return self.div(f_mesh.fft(self.velocity*u),f_mesh)
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._biharmonic._BiharmonicCore ¤

Bases: LinearCoef

Implementation of the Biharmonic operator.

Source code in torchfsm/operator/generic/_biharmonic.py 
 8
 9
10
11
12
13
14
15
16
class _BiharmonicCore(LinearCoef):
    r"""
    Implementation of the Biharmonic operator.
    """

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        return torch.cat([f_mesh.laplacian() * f_mesh.laplacian()] * n_channel, dim=1)
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_biharmonic.py 
13
14
15
16
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    return torch.cat([f_mesh.laplacian() * f_mesh.laplacian()] * n_channel, dim=1)
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.dedicated._gray_scott._ChannelWisedDiffusionCore ¤

Bases: LinearCoef

Implementation of the ChannelWisedDiffusion operator.

Source code in torchfsm/operator/dedicated/_gray_scott.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class _ChannelWisedDiffusionCore(LinearCoef):
    r"""
    Implementation of the ChannelWisedDiffusion operator.
    """

    def __init__(self,
                 viscosities: Sequence[Union[Tensor, float]]) -> None:
        super().__init__()
        self.viscosities= viscosities


    def __call__(self, 
                 f_mesh: FourierMesh, 
                 n_channel: int) -> FourierTensor["B C H ..."]:
        return torch.cat([viscosity*f_mesh.laplacian() for viscosity in self.viscosities],dim=1)
viscosities instance-attribute ¤
viscosities = viscosities
__init__ ¤
__init__(
    viscosities: Sequence[Union[Tensor, float]],
) -> None
Source code in torchfsm/operator/dedicated/_gray_scott.py 
13
14
15
16
def __init__(self,
             viscosities: Sequence[Union[Tensor, float]]) -> None:
    super().__init__()
    self.viscosities= viscosities
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_gray_scott.py 
19
20
21
22
def __call__(self, 
             f_mesh: FourierMesh, 
             n_channel: int) -> FourierTensor["B C H ..."]:
    return torch.cat([viscosity*f_mesh.laplacian() for viscosity in self.viscosities],dim=1)
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.generic._conservative_convection._ConservativeConvectionCore ¤

Bases: NonlinearFunc

Implementation of the Conservative Convection operator.

Source code in torchfsm/operator/generic/_conservative_convection.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class _ConservativeConvectionCore(NonlinearFunc):

    r"""
    Implementation of the Conservative Convection operator.
    """

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

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        if u is None:
            u = f_mesh.ifft(u_fft).real
        uu = u.unsqueeze(2) * u.unsqueeze(1)
        uu_fft = f_mesh.fft(uu)
        return (f_mesh.nabla_vector(1).unsqueeze(2) * uu_fft).sum(1)
__init__ ¤
__init__()
Source code in torchfsm/operator/generic/_conservative_convection.py 
14
15
def __init__(self):
    super().__init__()
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_conservative_convection.py 
17
18
19
20
21
22
23
24
25
26
27
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    if u is None:
        u = f_mesh.ifft(u_fft).real
    uu = u.unsqueeze(2) * u.unsqueeze(1)
    uu_fft = f_mesh.fft(uu)
    return (f_mesh.nabla_vector(1).unsqueeze(2) * uu_fft).sum(1)
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._convection._ConvectionCore ¤

Bases: NonlinearFunc

Implementation of the Convection operator.

Source code in torchfsm/operator/generic/_convection.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class _ConvectionCore(NonlinearFunc):

    r"""
    Implementation of the Convection operator.
    """

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

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        return f_mesh.fft(self.spatial_value(u_fft, f_mesh, u))

    def spatial_value(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> SpatialTensor["B C H ..."]:
        r"""
        Return the result of the nonlinear function in spatial domain.

        Args:
            u_fft (FourierTensor): Fourier-transformed input tensor.
            f_mesh (FourierMesh): Fourier mesh object.
            u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

        Returns:
            SpatialTensor: Result of the nonlinear function in spatial domain.
        """
        if u is None:
            u = f_mesh.ifft(u_fft).real
        re=0.0
        for i in range(u.shape[1]):
            re += u[:,i:i+1,...]*f_mesh.ifft(f_mesh.grad(i,1)*u_fft).real
        return re
__init__ ¤
__init__()
Source code in torchfsm/operator/generic/_convection.py 
15
16
def __init__(self):
    super().__init__()
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_convection.py 
18
19
20
21
22
23
24
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    return f_mesh.fft(self.spatial_value(u_fft, f_mesh, u))
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/generic/_convection.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """
    if u is None:
        u = f_mesh.ifft(u_fft).real
    re=0.0
    for i in range(u.shape[1]):
        re += u[:,i:i+1,...]*f_mesh.ifft(f_mesh.grad(i,1)*u_fft).real
    return re

torchfsm.operator.generic._curl._Curl2DCore ¤

Bases: NonlinearFunc

Implementation of the Curl operator for 2D vector fields.

Source code in torchfsm/operator/generic/_curl.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class _Curl2DCore(NonlinearFunc):

    r"""
    Implementation of the Curl operator for 2D vector fields.
    """

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

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        return (
            f_mesh.grad(0, 1) * u_fft[:, 1:2, ...]
            - f_mesh.grad(1, 1) * u_fft[:, 0:1, ...]
        )
__init__ ¤
__init__()
Source code in torchfsm/operator/generic/_curl.py 
15
16
def __init__(self):
    super().__init__(False)
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_curl.py 
18
19
20
21
22
23
24
25
26
27
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    return (
        f_mesh.grad(0, 1) * u_fft[:, 1:2, ...]
        - f_mesh.grad(1, 1) * u_fft[:, 0:1, ...]
    )
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._curl._Curl3DCore ¤

Bases: NonlinearFunc

Implementation of the Curl operator for 3D vector fields.

Source code in torchfsm/operator/generic/_curl.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
class _Curl3DCore(NonlinearFunc):

    r"""
    Implementation of the Curl operator for 3D vector fields.
    """

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

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        return torch.cat(
            [
                f_mesh.grad(1, 1) * u_fft[:, 2:3, ...]
                - f_mesh.grad(2, 1) * u_fft[:, 1:2, ...],
                f_mesh.grad(2, 1) * u_fft[:, 0:1, ...]
                - f_mesh.grad(0, 1) * u_fft[:, 2:3, ...],
                f_mesh.grad(0, 1) * u_fft[:, 1:2, ...]
                - f_mesh.grad(1, 1) * u_fft[:, 0:1, ...],
            ],
            dim=1,
        )
__init__ ¤
__init__()
Source code in torchfsm/operator/generic/_curl.py 
36
37
def __init__(self):
    super().__init__(False)
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_curl.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    return torch.cat(
        [
            f_mesh.grad(1, 1) * u_fft[:, 2:3, ...]
            - f_mesh.grad(2, 1) * u_fft[:, 1:2, ...],
            f_mesh.grad(2, 1) * u_fft[:, 0:1, ...]
            - f_mesh.grad(0, 1) * u_fft[:, 2:3, ...],
            f_mesh.grad(0, 1) * u_fft[:, 1:2, ...]
            - f_mesh.grad(1, 1) * u_fft[:, 0:1, ...],
        ],
        dim=1,
    )
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._dispersion._DispersionCore ¤

Bases: LinearCoef

Implementation of the Dispersion operator.

Source code in torchfsm/operator/generic/_dispersion.py 
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class _DispersionCore(LinearCoef):
    r"""
    Implementation of the Dispersion operator.
    """

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        return torch.sum(
            f_mesh.nabla_vector(1) * torch.cat([f_mesh.laplacian()] * n_channel, dim=1),
            dim=1,
            keepdim=True,
        )
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_dispersion.py 
11
12
13
14
15
16
17
18
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    return torch.sum(
        f_mesh.nabla_vector(1) * torch.cat([f_mesh.laplacian()] * n_channel, dim=1),
        dim=1,
        keepdim=True,
    )
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.generic._div._DivCore ¤

Bases: NonlinearFunc

Implementation of the Divergence operator.

Source code in torchfsm/operator/generic/_div.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class _DivCore(NonlinearFunc):
    r"""
    Implementation of the Divergence operator.
    """

    def __init__(self):
        super().__init__(False)
        """
        Although the divergence is implemented as a nonlinear function, 
        it is actually a linear operation where dealising is not required.
        We implement it as a nonlinear function since its linear feature is a dot product operation, which is not supported by the spectral method.
        """

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        return torch.sum(f_mesh.nabla_vector(1) * u_fft, dim=1, keepdim=True)
__init__ ¤
__init__()
Source code in torchfsm/operator/generic/_div.py 
14
15
16
17
18
19
20
def __init__(self):
    super().__init__(False)
    """
    Although the divergence is implemented as a nonlinear function, 
    it is actually a linear operation where dealising is not required.
    We implement it as a nonlinear function since its linear feature is a dot product operation, which is not supported by the spectral method.
    """
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_div.py 
22
23
24
25
26
27
28
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    return torch.sum(f_mesh.nabla_vector(1) * u_fft, dim=1, keepdim=True)
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator._base._ExplicitSourceCore ¤

Bases: NonlinearFunc

Implementation of the explicit source term for the operator.

Parameters:

Name Type Description Default
source SpatialTensor['B C H ...']

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 
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
class _ExplicitSourceCore(NonlinearFunc):
    r"""
    Implementation of the explicit source term for the operator.

    Args:
        source (SpatialTensor["B C H ..."]): Source term in spatial domain. This is a tensor that represents the source term in the spatial domain.
    """

    def __init__(self, source: SpatialTensor["B C H ..."]) -> None:
        super().__init__(dealiasing_swtich=False)
        fft_dim = [i + 2 for i in range(source.dim() - 2)]
        self.source = torch.fft.fftn(source, dim=fft_dim)

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        if self.source.device != f_mesh.device:
            self.source = self.source.to(f_mesh.device)
        return self.source
source instance-attribute ¤
source = fftn(source, dim=fft_dim)
__init__ ¤
__init__(source: SpatialTensor['B C H ...']) -> None
Source code in torchfsm/operator/_base.py 
1030
1031
1032
1033
def __init__(self, source: SpatialTensor["B C H ..."]) -> None:
    super().__init__(dealiasing_swtich=False)
    fft_dim = [i + 2 for i in range(source.dim() - 2)]
    self.source = torch.fft.fftn(source, dim=fft_dim)
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/_base.py 
1035
1036
1037
1038
1039
1040
1041
1042
1043
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    if self.source.device != f_mesh.device:
        self.source = self.source.to(f_mesh.device)
    return self.source
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.dedicated._gray_scott._GrayScottSourceCore ¤

Bases: NonlinearFunc

Implementation of the Gray-Scott source term.

Source code in torchfsm/operator/dedicated/_gray_scott.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
class _GrayScottSourceCore(NonlinearFunc):

    r"""
    Implementation of the Gray-Scott source term.
    """

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

    def __call__(
            self,
            u_fft: FourierTensor["B C H ..."],
            f_mesh: FourierMesh,
            u: Optional[SpatialTensor["B C H ..."]] = None,
        ) -> FourierTensor["B C H ..."]:
        u0u12 = u[:, 0, ...] * u[:, 1, ...]* u[:, 1, ...]
        return  f_mesh.fft(torch.stack(
            [
                self.feed_rate *(1 - u[:, 0, ...]) - u0u12,
                u0u12-(self.kill_rate + self.feed_rate) * u[:, 1, ...]
            ],
            dim=1
        ))
feed_rate instance-attribute ¤
feed_rate = feed_rate
kill_rate instance-attribute ¤
kill_rate = kill_rate
__init__ ¤
__init__(
    feed_rate: Union[Tensor, float],
    kill_rate: Union[Tensor, float],
) -> None
Source code in torchfsm/operator/dedicated/_gray_scott.py 
55
56
57
58
59
60
def __init__(self, 
             feed_rate: Union[Tensor, float],
             kill_rate: Union[Tensor, float]) -> None:
    super().__init__()
    self.feed_rate = feed_rate
    self.kill_rate = kill_rate
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_gray_scott.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
    u0u12 = u[:, 0, ...] * u[:, 1, ...]* u[:, 1, ...]
    return  f_mesh.fft(torch.stack(
        [
            self.feed_rate *(1 - u[:, 0, ...]) - u0u12,
            u0u12-(self.kill_rate + self.feed_rate) * u[:, 1, ...]
        ],
        dim=1
    ))
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._grad._GradCore ¤

Bases: LinearCoef

Implementation of the Grad operator.

Source code in torchfsm/operator/generic/_grad.py 
 6
 7
 8
 9
10
11
12
13
14
15
class _GradCore(LinearCoef):

    r"""
    Implementation of the Grad operator.
    """

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        return f_mesh.nabla_vector(1)
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_grad.py 
12
13
14
15
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    return f_mesh.nabla_vector(1)
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.generic._hyper_diffusion._HyperDiffusionCore ¤

Bases: LinearCoef

Implementation of the Hyper Diffusion operator.

Source code in torchfsm/operator/generic/_hyper_diffusion.py 
 6
 7
 8
 9
10
11
12
13
14
class _HyperDiffusionCore(LinearCoef):
    r"""
    Implementation of the Hyper Diffusion operator.
    """

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        return torch.cat([f_mesh.laplacian()] * n_channel,dim=1)**2
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_hyper_diffusion.py 
11
12
13
14
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    return torch.cat([f_mesh.laplacian()] * n_channel,dim=1)**2
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.generic._source._ImplicitFuncSourceCore ¤

Bases: NonlinearFunc

Implementation of the ImplicitSource operator with function form.

Source code in torchfsm/operator/generic/_source.py 
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class _ImplicitFuncSourceCore(NonlinearFunc):
    r"""
    Implementation of the ImplicitSource operator with function form.
    """

    def __init__(
        self,
        source_func: Callable[[torch.Tensor], torch.Tensor],
        non_linear: bool = True,
    ) -> None:
        super().__init__(non_linear)
        self.source_func = source_func

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        if u is None:
            u = f_mesh.ifft(u_fft).real
        return f_mesh.fft(self.source_func(u))
source_func instance-attribute ¤
source_func = source_func
__init__ ¤
__init__(
    source_func: Callable[[Tensor], Tensor],
    non_linear: bool = True,
) -> None
Source code in torchfsm/operator/generic/_source.py 
25
26
27
28
29
30
31
def __init__(
    self,
    source_func: Callable[[torch.Tensor], torch.Tensor],
    non_linear: bool = True,
) -> None:
    super().__init__(non_linear)
    self.source_func = source_func
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_source.py 
33
34
35
36
37
38
39
40
41
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    if u is None:
        u = f_mesh.ifft(u_fft).real
    return f_mesh.fft(self.source_func(u))
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._source._ImplicitUnitSourceCore ¤

Bases: LinearCoef

Implementation of the ImplicitSource operator with unit form.

Source code in torchfsm/operator/generic/_source.py 
 9
10
11
12
13
14
15
16
17
class _ImplicitUnitSourceCore(LinearCoef):
    r"""
    Implementation of the ImplicitSource operator with unit form.
    """

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        return torch.ones_like(f_mesh.bf_x)
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_source.py 
14
15
16
17
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    return torch.ones_like(f_mesh.bf_x)
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.dedicated._ks_convection._KSConvectionCore ¤

Bases: NonlinearFunc

Implementation of the KSConvection operator.

Source code in torchfsm/operator/dedicated/_ks_convection.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class _KSConvectionCore(NonlinearFunc):

    r"""
    Implementation of the KSConvection operator.
    """

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

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        n_channel: int,
        u: Optional[SpatialTensor["B C H ..."]]=None,
    ) -> FourierTensor["B C H ..."]:
        return f_mesh.fft(self.spatial_value(u_fft, f_mesh, u))

    def spatial_value(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]]=None,
    ) -> SpatialTensor["B C H ..."]:
        grad_u = f_mesh.ifft(f_mesh.nabla_vector(1) * u_fft).real
        if self.remove_mean:
            re = 0.5 * torch.sum(grad_u**2, dim=1, keepdim=True)
            return re - re.mean()
        else:
            return 0.5 * torch.sum(grad_u**2, dim=1, keepdim=True)
remove_mean instance-attribute ¤
remove_mean = remove_mean
__init__ ¤
__init__(remove_mean: bool) -> None
Source code in torchfsm/operator/dedicated/_ks_convection.py 
14
15
16
def __init__(self, remove_mean: bool) -> None:
    super().__init__()
    self.remove_mean = remove_mean
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    n_channel: int,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_ks_convection.py 
18
19
20
21
22
23
24
25
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    n_channel: int,
    u: Optional[SpatialTensor["B C H ..."]]=None,
) -> FourierTensor["B C H ..."]:
    return f_mesh.fft(self.spatial_value(u_fft, f_mesh, u))
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_ks_convection.py 
27
28
29
30
31
32
33
34
35
36
37
38
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]]=None,
) -> SpatialTensor["B C H ..."]:
    grad_u = f_mesh.ifft(f_mesh.nabla_vector(1) * u_fft).real
    if self.remove_mean:
        re = 0.5 * torch.sum(grad_u**2, dim=1, keepdim=True)
        return re - re.mean()
    else:
        return 0.5 * torch.sum(grad_u**2, dim=1, keepdim=True)

torchfsm.operator.generic._laplacian._LaplacianCore ¤

Bases: LinearCoef

Implementation of the Laplacian operator.

Source code in torchfsm/operator/generic/_laplacian.py 
 7
 8
 9
10
11
12
13
14
15
class _LaplacianCore(LinearCoef):
    r"""
    Implementation of the Laplacian operator.
    """

    def __call__(self, 
                 f_mesh: FourierMesh, 
                 n_channel: int) -> FourierTensor["B C H ..."]:
        return torch.cat([f_mesh.laplacian()]*n_channel,dim=1)
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_laplacian.py 
12
13
14
15
def __call__(self, 
             f_mesh: FourierMesh, 
             n_channel: int) -> FourierTensor["B C H ..."]:
    return torch.cat([f_mesh.laplacian()]*n_channel,dim=1)
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.dedicated._navier_stokes._leray._LerayCore ¤

Bases: NonlinearFunc

Implementation of the Leray projection operator.

Source code in torchfsm/operator/dedicated/_navier_stokes/_leray.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class _LerayCore(NonlinearFunc):

    r"""
    Implementation of the Leray projection operator.
    """

    def __init__(self):
        super().__init__(False)
        """
        Although the Leray projection is implemented as a nonlinear function, 
        it is actually a linear operation where dealising is not required.
        We implement it as a nonlinear function since its linear feature is a dot product operation, which is not supported by the spectral method.
        """

    def __call__(self, u_fft, f_mesh, u=None):
        return u_fft - f_mesh.nabla_vector(1) * f_mesh.invert_laplacian() * torch.sum(
            f_mesh.nabla_vector(1) * u_fft, dim=1, keepdim=True
        )
__init__ ¤
__init__()
Source code in torchfsm/operator/dedicated/_navier_stokes/_leray.py 
14
15
16
17
18
19
20
def __init__(self):
    super().__init__(False)
    """
    Although the Leray projection is implemented as a nonlinear function, 
    it is actually a linear operation where dealising is not required.
    We implement it as a nonlinear function since its linear feature is a dot product operation, which is not supported by the spectral method.
    """
__call__ ¤
__call__(u_fft, f_mesh, u=None)
Source code in torchfsm/operator/dedicated/_navier_stokes/_leray.py 
22
23
24
25
def __call__(self, u_fft, f_mesh, u=None):
    return u_fft - f_mesh.nabla_vector(1) * f_mesh.invert_laplacian() * torch.sum(
        f_mesh.nabla_vector(1) * u_fft, dim=1, keepdim=True
    )
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._linear_advection._LinearAdvectionCore ¤

Bases: LinearCoef

Implementation of the LinearAdvection operator.

Source code in torchfsm/operator/generic/_linear_advection.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
class _LinearAdvectionCore(LinearCoef):
    r"""
    Implementation of the LinearAdvection operator.
    """

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

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        if isinstance(self.velocity, float):
            velocity = [self.velocity] * f_mesh.n_dim
        else:
            assert len(self.velocity) == f_mesh.n_dim, (
                "Length of velocity list must match the number of dimensions of the mesh."
                + " If you have same operator for a different mesh before, please reset the velocity using `set_advection_velocity` method."
            )
            velocity = self.velocity
        if n_channel != 1:
            raise ValueError("The Advection operator only supports scalar field.")
        return sum(
            [
                v * f_mesh.nabla_vector(1)[:, i : i + 1, ...]
                for i, v in enumerate(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.velocity = velocity
velocity instance-attribute ¤
velocity = velocity
__init__ ¤
__init__(
    velocity: Union[float, Sequence[float]] = 1.0,
) -> None
Source code in torchfsm/operator/generic/_linear_advection.py 
12
13
14
def __init__(self, velocity: Union[float, Sequence[float]] = 1.0) -> None:
    super().__init__()
    self.velocity = velocity
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_linear_advection.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    if isinstance(self.velocity, float):
        velocity = [self.velocity] * f_mesh.n_dim
    else:
        assert len(self.velocity) == f_mesh.n_dim, (
            "Length of velocity list must match the number of dimensions of the mesh."
            + " If you have same operator for a different mesh before, please reset the velocity using `set_advection_velocity` method."
        )
        velocity = self.velocity
    if n_channel != 1:
        raise ValueError("The Advection operator only supports scalar field.")
    return sum(
        [
            v * f_mesh.nabla_vector(1)[:, i : i + 1, ...]
            for i, v in enumerate(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 
36
37
38
39
40
41
42
43
44
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.velocity = velocity
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.dedicated._navier_stokes._ns_pressure_convection._NSPressureConvectionCore ¤

Bases: NonlinearFunc

Implementation of the Navier-Stokes pressure convection operator.

Source code in torchfsm/operator/dedicated/_navier_stokes/_ns_pressure_convection.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class _NSPressureConvectionCore(NonlinearFunc):
    r"""
    Implementation of the Navier-Stokes pressure convection operator.
    """

    def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
        super().__init__(dealiasing_swtich=external_force is None)
        self.external_force = external_force
        self._convection = _ConvectionCore()

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> torch.Tensor:
        if self.external_force is not None:  # u_fft is original version
            force = self.external_force(
                u_fft=u_fft, mesh=f_mesh, return_in_fourier=True
            )
            u_fft *= f_mesh.low_pass_filter()
            u = f_mesh.ifft(u_fft).real
        else:  # u_fft is dealiased version
            if u is None:
                u = f_mesh.ifft(u_fft).real
        convection = self._convection(u_fft, f_mesh,u)
        if self.external_force is not None:
            convection -= force
        p = f_mesh.invert_laplacian() * torch.sum(
            f_mesh.nabla_vector(1) * convection, dim=1, keepdim=True
        )  # - p = nabla.(u.nabla_u)/laplacian
        if self.external_force is not None:
            return f_mesh.nabla_vector(1) * p - convection + force # -nabla(p) - nabla.(u.nabla_u) + f
        return f_mesh.nabla_vector(1) * p - convection  # -nabla(p) - nabla.(u.nabla_u)
external_force instance-attribute ¤
external_force = external_force
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_ns_pressure_convection.py 
19
20
21
22
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    super().__init__(dealiasing_swtich=external_force is None)
    self.external_force = external_force
    self._convection = _ConvectionCore()
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> torch.Tensor
Source code in torchfsm/operator/dedicated/_navier_stokes/_ns_pressure_convection.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> torch.Tensor:
    if self.external_force is not None:  # u_fft is original version
        force = self.external_force(
            u_fft=u_fft, mesh=f_mesh, return_in_fourier=True
        )
        u_fft *= f_mesh.low_pass_filter()
        u = f_mesh.ifft(u_fft).real
    else:  # u_fft is dealiased version
        if u is None:
            u = f_mesh.ifft(u_fft).real
    convection = self._convection(u_fft, f_mesh,u)
    if self.external_force is not None:
        convection -= force
    p = f_mesh.invert_laplacian() * torch.sum(
        f_mesh.nabla_vector(1) * convection, dim=1, keepdim=True
    )  # - p = nabla.(u.nabla_u)/laplacian
    if self.external_force is not None:
        return f_mesh.nabla_vector(1) * p - convection + force # -nabla(p) - nabla.(u.nabla_u) + f
    return f_mesh.nabla_vector(1) * p - convection  # -nabla(p) - nabla.(u.nabla_u)
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.generic._spatial_derivative._SpatialDerivativeCore ¤

Bases: LinearCoef

Implementation of the SpatialDerivative operator.

Source code in torchfsm/operator/generic/_spatial_derivative.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class _SpatialDerivativeCore(LinearCoef):
    r"""
    Implementation of the SpatialDerivative operator.
    """

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

    def __call__(
        self, f_mesh: FourierMesh, n_channel: int
    ) -> FourierTensor["B C H ..."]:
        return f_mesh.grad(self.dim_index, self.order)
dim_index instance-attribute ¤
dim_index = dim_index
order instance-attribute ¤
order = order
__init__ ¤
__init__(dim_index, order) -> None
Source code in torchfsm/operator/generic/_spatial_derivative.py 
12
13
14
15
def __init__(self, dim_index, order) -> None:
    super().__init__()
    self.dim_index = dim_index
    self.order = order
__call__ ¤
__call__(
    f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/generic/_spatial_derivative.py 
17
18
19
20
def __call__(
    self, f_mesh: FourierMesh, n_channel: int
) -> FourierTensor["B C H ..."]:
    return f_mesh.grad(self.dim_index, self.order)
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.dedicated._navier_stokes._value_transformation._Velocity2PressureCore ¤

Bases: NonlinearFunc

Implementation of the Velocity2Pressure operator.

Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
class _Velocity2PressureCore(NonlinearFunc):
    r"""
    Implementation of the Velocity2Pressure operator.
    """

    def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
        # if exeternal_force is not None, the external force may need the original u_fft, thus we do not dealiasing at the beginning
        super().__init__(dealiasing_swtich=external_force is None)
        self.external_force = external_force
        self._convection = _ConvectionCore()

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]] = None,
    ) -> FourierTensor["B C H ..."]:
        if self.external_force is not None:  # u_fft is original version
            force = self.external_force(
                u_fft=u_fft, mesh=f_mesh, return_in_fourier=True
            )
            u_fft *= f_mesh.low_pass_filter()
            u = f_mesh.ifft(u_fft).real
        else:  # u_fft is dealiased version
            if u is None:
                u = f_mesh.ifft(u_fft).real
        convection = self._convection(u_fft, f_mesh, u)
        if self.external_force is not None:
            convection -= force
        p = torch.sum(
            f_mesh.nabla_vector(1) * convection, dim=1, keepdim=True
        )  # nabla.(u.nabla_u)
        return -1 * p * f_mesh.invert_laplacian()  #  p = - nabla.(u.nabla_u)/laplacian
external_force instance-attribute ¤
external_force = external_force
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
122
123
124
125
126
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    # if exeternal_force is not None, the external force may need the original u_fft, thus we do not dealiasing at the beginning
    super().__init__(dealiasing_swtich=external_force is None)
    self.external_force = external_force
    self._convection = _ConvectionCore()
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    if self.external_force is not None:  # u_fft is original version
        force = self.external_force(
            u_fft=u_fft, mesh=f_mesh, return_in_fourier=True
        )
        u_fft *= f_mesh.low_pass_filter()
        u = f_mesh.ifft(u_fft).real
    else:  # u_fft is dealiased version
        if u is None:
            u = f_mesh.ifft(u_fft).real
    convection = self._convection(u_fft, f_mesh, u)
    if self.external_force is not None:
        convection -= force
    p = torch.sum(
        f_mesh.nabla_vector(1) * convection, dim=1, keepdim=True
    )  # nabla.(u.nabla_u)
    return -1 * p * f_mesh.invert_laplacian()  #  p = - nabla.(u.nabla_u)/laplacian
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.dedicated._navier_stokes._value_transformation._Vorticity2PressureCore ¤

Bases: NonlinearFunc

Implementation of the Vorticity2Pressure operator.

Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
class _Vorticity2PressureCore(NonlinearFunc):
    r"""
    Implementation of the Vorticity2Pressure operator.
    """
    def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
        # if exeternal_force is not None, the external force may need the original u_fft, thus we do not dealiasing at the beginning
        super().__init__(dealiasing_swtich=external_force is None)
        self.external_force = external_force
        self._vorticity2velocity = _Vorticity2VelocityCore()
        self._convection = _ConvectionCore()

    def __call__(self, u_fft, f_mesh, u) -> FourierTensor["B C H ..."]:
        velocity_fft = u_fft * self._vorticity2velocity(f_mesh, 1)
        if self.external_force is not None:
            velocity_fft *= f_mesh.low_pass_filter()
        convection = self._convection(
            velocity_fft, f_mesh, f_mesh.ifft(velocity_fft).real
        )
        if self.external_force is not None:
            convection -= self.external_force(
                u_fft=u_fft, mesh=f_mesh, return_in_fourier=True
            )
        p = torch.sum(
            f_mesh.nabla_vector(1) * convection, dim=1, keepdim=True
        )  # nabla.(u.nabla_u)
        return -1 * p * f_mesh.invert_laplacian()  #  p = - nabla.(u.nabla_u)/laplacian
external_force instance-attribute ¤
external_force = external_force
__init__ ¤
__init__(
    external_force: Optional[OperatorLike] = None,
) -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
66
67
68
69
70
71
def __init__(self, external_force: Optional[OperatorLike] = None) -> None:
    # if exeternal_force is not None, the external force may need the original u_fft, thus we do not dealiasing at the beginning
    super().__init__(dealiasing_swtich=external_force is None)
    self.external_force = external_force
    self._vorticity2velocity = _Vorticity2VelocityCore()
    self._convection = _ConvectionCore()
__call__ ¤
__call__(u_fft, f_mesh, u) -> FourierTensor['B C H ...']
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
def __call__(self, u_fft, f_mesh, u) -> FourierTensor["B C H ..."]:
    velocity_fft = u_fft * self._vorticity2velocity(f_mesh, 1)
    if self.external_force is not None:
        velocity_fft *= f_mesh.low_pass_filter()
    convection = self._convection(
        velocity_fft, f_mesh, f_mesh.ifft(velocity_fft).real
    )
    if self.external_force is not None:
        convection -= self.external_force(
            u_fft=u_fft, mesh=f_mesh, return_in_fourier=True
        )
    p = torch.sum(
        f_mesh.nabla_vector(1) * convection, dim=1, keepdim=True
    )  # nabla.(u.nabla_u)
    return -1 * p * f_mesh.invert_laplacian()  #  p = - nabla.(u.nabla_u)/laplacian
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]

Return the result of the nonlinear function in spatial domain.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
SpatialTensor SpatialTensor['B C H ...']

Result of the nonlinear function in spatial domain.
Source code in torchfsm/operator/_base.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]:
    r"""
    Return the result of the nonlinear function in spatial domain.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        SpatialTensor: Result of the nonlinear function in spatial domain.
    """

    return f_mesh.ifft(self(u_fft, f_mesh, u)).real

torchfsm.operator.dedicated._navier_stokes._value_transformation._Vorticity2VelocityCore ¤

Bases: LinearCoef

Implementation of the Vorticity2Velocity operator.

Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class _Vorticity2VelocityCore(LinearCoef):
    r"""
    Implementation of the Vorticity2Velocity operator.
    """

    def __call__(self, f_mesh, n_channel) -> FourierTensor["B C H ..."]:
        return (
            -1
            * f_mesh.invert_laplacian()
            * torch.cat(
                [
                    f_mesh.grad(1, 1).repeat([1, 1, f_mesh.mesh_info[0][-1], 1]),
                    -f_mesh.grad(0, 1).repeat([1, 1, 1, f_mesh.mesh_info[1][-1]]),
                ],
                dim=1,
            )
        )
__call__ ¤
__call__(f_mesh, n_channel) -> FourierTensor['B C H ...']
Source code in torchfsm/operator/dedicated/_navier_stokes/_value_transformation.py 
23
24
25
26
27
28
29
30
31
32
33
34
def __call__(self, f_mesh, n_channel) -> FourierTensor["B C H ..."]:
    return (
        -1
        * f_mesh.invert_laplacian()
        * torch.cat(
            [
                f_mesh.grad(1, 1).repeat([1, 1, f_mesh.mesh_info[0][-1], 1]),
                -f_mesh.grad(0, 1).repeat([1, 1, 1, f_mesh.mesh_info[1][-1]]),
            ],
            dim=1,
        )
    )
nonlinear_like ¤
nonlinear_like(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]

Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

Parameters:

Name Type Description Default
u_fft FourierTensor

Fourier-transformed input tensor.

 required
f_mesh FourierMesh

Fourier mesh object.

 required
u Optional[SpatialTensor]

Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

 None

Returns:

Name Type Description
FourierTensor FourierTensor['B C H ...']

Nonlinear-like tensor.
Source code in torchfsm/operator/_base.py 
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
def nonlinear_like(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]:
    r"""
    Calculate the result out based on the linear coefficient. It is designed to have same pattern as the nonlinear function.

    Args:
        u_fft (FourierTensor): Fourier-transformed input tensor.
        f_mesh (FourierMesh): Fourier mesh object.
        u (Optional[SpatialTensor]): Corresponding tensor of u_fft in spatial domain. This option aims to avoid repeating the inverse FFT operation in operators.

    Returns:
        FourierTensor: Nonlinear-like tensor.
    """
    return self(f_mesh, u_fft.shape[1]) * u_fft

torchfsm.operator.dedicated._navier_stokes._vorticity_convection._VorticityConvectionCore ¤

Bases: NonlinearFunc

Implementation of the VorticityConvection operator.

Source code in torchfsm/operator/dedicated/_navier_stokes/_vorticity_convection.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
class _VorticityConvectionCore(NonlinearFunc):

    r"""
    Implementation of the VorticityConvection operator.
    """

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

    def __call__(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]]=None,
    ) -> FourierTensor["B C H ..."]:
        return f_mesh.fft(self.spatial_value(u_fft, f_mesh, u))

    def spatial_value(
        self,
        u_fft: FourierTensor["B C H ..."],
        f_mesh: FourierMesh,
        u: Optional[SpatialTensor["B C H ..."]]=None,
    ) -> SpatialTensor["B C H ..."]:
        psi = -u_fft * f_mesh.invert_laplacian()
        ux = f_mesh.ifft(f_mesh.grad(1, 1) * psi).real
        uy = f_mesh.ifft(-f_mesh.grad(0, 1) * psi).real
        grad_x_w = f_mesh.ifft(f_mesh.grad(0, 1) * u_fft).real
        grad_y_w = f_mesh.ifft(f_mesh.grad(1, 1) * u_fft).real
        return ux * grad_x_w + uy * grad_y_w
__init__ ¤
__init__() -> None
Source code in torchfsm/operator/dedicated/_navier_stokes/_vorticity_convection.py 
19
20
def __init__(self) -> None:
    super().__init__()
__call__ ¤
__call__(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> FourierTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_navier_stokes/_vorticity_convection.py 
22
23
24
25
26
27
28
def __call__(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]]=None,
) -> FourierTensor["B C H ..."]:
    return f_mesh.fft(self.spatial_value(u_fft, f_mesh, u))
spatial_value ¤
spatial_value(
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]] = None,
) -> SpatialTensor["B C H ..."]
Source code in torchfsm/operator/dedicated/_navier_stokes/_vorticity_convection.py 
30
31
32
33
34
35
36
37
38
39
40
41
def spatial_value(
    self,
    u_fft: FourierTensor["B C H ..."],
    f_mesh: FourierMesh,
    u: Optional[SpatialTensor["B C H ..."]]=None,
) -> SpatialTensor["B C H ..."]:
    psi = -u_fft * f_mesh.invert_laplacian()
    ux = f_mesh.ifft(f_mesh.grad(1, 1) * psi).real
    uy = f_mesh.ifft(-f_mesh.grad(0, 1) * psi).real
    grad_x_w = f_mesh.ifft(f_mesh.grad(0, 1) * u_fft).real
    grad_y_w = f_mesh.ifft(f_mesh.grad(1, 1) * u_fft).real
    return ux * grad_x_w + uy * grad_y_w

// ...existing code...