Skip to content

Cores

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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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._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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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._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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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._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 
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
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 
1000
1001
1002
1003
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 
1005
1006
1007
1008
1009
1010
1011
1012
1013
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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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._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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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._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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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._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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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.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

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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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._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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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._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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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._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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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._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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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._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 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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._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 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
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