Skip to content

5. field

torchfsm.field.diffused_noise ¤

diffused_noise(
    mesh: Union[
        Sequence[tuple[float, float, int]],
        MeshGrid,
        FourierMesh,
    ],
    diffusion_coef: float = 1.0,
    zero_centered: bool = True,
    unit_variance: bool = True,
    unit_magnitude: bool = True,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    n_batch: int = 1,
    n_channel: int = 1,
) -> SpatialTensor["B C H ..."]

Generate a diffused noise field. The noise is generated by integrating a Laplacian operator with a random initial condition. The diffusion coefficient controls the amount of diffusion applied to the noise.

Parameters:

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

The mesh to generate the noise on. If a sequence is provided, it should be in the form of [(x_min, x_max, n_points), ...].

required
diffusion_coef float

The diffusion coefficient. Default is 1.0.

1.0
zero_centered bool

If True, the noise will be zero-centered. Default is True.

True
unit_variance bool

If True, the noise will have unit variance. Default is True.

True
unit_magnitude bool

If True, the noise will have unit magnitude. Default is True.

True
device Optional[device]

The device to generate the noise on. Default is None.

None
dtype Optional[dtype]

The data type of the generated noise. Default is None.

None
n_batch int

The number of batches. Default is 1.

1
n_channel int

The number of channels. Default is 1.

1

Returns:

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

SpatialTensor["B C H ...]: The generated diffused noise field.

Source code in torchfsm/field.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
45
46
47
48
49
50
51
52
53
54
55
def diffused_noise(
    mesh: Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh],
    diffusion_coef: float = 1.0,
    zero_centered: bool = True,
    unit_variance: bool = True,
    unit_magnitude: bool = True,
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
    n_batch: int = 1,
    n_channel: int = 1,
) -> SpatialTensor["B C H ..."]:
    r"""
    Generate a diffused noise field.
        The noise is generated by integrating a Laplacian operator with a random initial condition.
        The diffusion coefficient controls the amount of diffusion applied to the noise.

    Args:
        mesh (Union[Sequence[tuple[float, float, int]], MeshGrid, FourierMesh]): The mesh to generate the noise on.
            If a sequence is provided, it should be in the form of [(x_min, x_max, n_points), ...].
        diffusion_coef (float): The diffusion coefficient. Default is 1.0.
        zero_centered (bool): If True, the noise will be zero-centered. Default is True.
        unit_variance (bool): If True, the noise will have unit variance. Default is True.
        unit_magnitude (bool): If True, the noise will have unit magnitude. Default is True.
        device (Optional[torch.device]): The device to generate the noise on. Default is None.
        dtype (Optional[torch.dtype]): The data type of the generated noise. Default is None.
        n_batch (int): The number of batches. Default is 1.
        n_channel (int): The number of channels. Default is 1.

    Returns:
        SpatialTensor["B C H ...]: The generated diffused noise field.
    """
    if device is None and (isinstance(mesh, FourierMesh) or isinstance(mesh, MeshGrid)):
        device = mesh.device
    if dtype is None and (isinstance(mesh, FourierMesh) or isinstance(mesh, MeshGrid)):
        dtype = mesh.dtype
    u_0 = torch.randn(
        *mesh_shape(mesh, n_batch=n_batch, n_channel=n_channel),
        device=device,
        dtype=dtype
    )
    diffusion = diffusion_coef * Laplacian()
    u_0 = diffusion.integrate(u_0, dt=1, step=1, mesh=mesh)
    if zero_centered:
        u_0 = u_0 - u_0.mean()
    if unit_variance:
        u_0 = u_0 / u_0.std()
    if unit_magnitude:
        u_0 = u_0 / u_0.abs().max()
    return u_0

torchfsm.field.kolm_force ¤

kolm_force(
    x: torch.Tensor,
    drag_coef: float = -0.1,
    k: float = 4.0,
    length_scale: float = 1.0,
) -> Operator

Generate cosine force field for 2d kolmogorov flow in vorticity form. It is defined as \(a \omega - k cos (k l x)\)

Parameters:

Name Type Description Default
x Tensor

The input tensor.

required
drag_coef float

The drag coefficient \(a\). Default is -0.1.

-0.1
k float

The wave number. Default is 4.0.

4.0
length_scale float

The length scale \(l\). Default is 1.0.

1.0
Source code in torchfsm/field.py
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def kolm_force(
    x: torch.Tensor,
    drag_coef: float = -0.1,
    k: float = 4.0,
    length_scale: float = 1.0,
) -> Operator:
    r"""
    Generate cosine force field for 2d kolmogorov flow in vorticity form.
        It is defined as $a \omega - k cos (k l x)$

    Args:
        x (torch.Tensor): The input tensor.
        drag_coef (float): The drag coefficient $a$. Default is -0.1.
        k (float): The wave number. Default is 4.0.
        length_scale (float): The length scale $l$. Default is 1.0.
    """


    return drag_coef * ImplicitSource() - ExplicitSource(
        k * torch.cos(k * length_scale * x)
    )

torchfsm.field.wave_1d ¤

wave_1d(
    x: SpatialTensor["B C H ..."],
    min_k: int = 1,
    max_k: int = 5,
    min_amplitude: float = 0.5,
    max_amplitude: float = 1.0,
    n_polynomial: int = 5,
    zero_mean: bool = False,
    mean_shift_coef=0.3,
    batched: bool = False,
) -> SpatialTensor["B C H ..."]

Generate a 1D wave field with multiple harmonics.

Parameters:

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

The input tensor.

required
min_k int

The minimum wave number. Default is 1.

1
max_k int

The maximum wave number. Default is 5.

5
min_amplitude float

The minimum amplitude. Default is 0.5.

0.5
max_amplitude float

The maximum amplitude. Default is 1.0.

1.0
n_polynomial int

The number of polynomial terms. Default is 5.

5
zero_mean bool

If True, the mean of the wave will be zero. Default is False.

False
mean_shift_coef float

The coefficient for mean shift. Default is 0.3.

0.3
batched bool

If True, the input tensor is batched. Default is False.

False

Returns:

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

SpatialTensor["B C H ..."]: The generated wave field.

Source code in torchfsm/field.py
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def wave_1d(
    x: SpatialTensor["B C H ..."],
    min_k: int = 1,
    max_k: int = 5,
    min_amplitude: float = 0.5,
    max_amplitude: float = 1.0,
    n_polynomial: int = 5,
    zero_mean: bool = False,
    mean_shift_coef=0.3,
    batched: bool = False,
) -> SpatialTensor["B C H ..."]:
    r"""
    Generate a 1D wave field with multiple harmonics.

    Args:
        x (SpatialTensor["B C H ..."]): The input tensor.
        min_k (int): The minimum wave number. Default is 1.
        max_k (int): The maximum wave number. Default is 5.
        min_amplitude (float): The minimum amplitude. Default is 0.5.
        max_amplitude (float): The maximum amplitude. Default is 1.0.
        n_polynomial (int): The number of polynomial terms. Default is 5.
        zero_mean (bool): If True, the mean of the wave will be zero. Default is False.
        mean_shift_coef (float): The coefficient for mean shift. Default is 0.3.
        batched (bool): If True, the input tensor is batched. Default is False.

    Returns:
        SpatialTensor["B C H ..."]: The generated wave field.
    """
    x_new = x / x.max() * torch.pi * 2
    y = torch.zeros_like(x)
    if not batched:
        x_new = x_new.unsqueeze(0)
        y = y.unsqueeze(0)
    batch = x_new.shape[0]
    shape = [batch, n_polynomial] + [1] * (x_new.dim() - 2)
    k = torch.randint(min_k, max_k + 1, shape, device=x.device, dtype=x.dtype)
    amplitude = (
        torch.rand(*shape, device=x.device, dtype=x.dtype)
        * (max_amplitude - min_amplitude)
        + min_amplitude
    )
    shift = torch.rand(*shape, device=x.device, dtype=x.dtype) * torch.pi * 2
    for i in range(n_polynomial):
        y += amplitude[:, i : i + 1, ...] * torch.sin(
            k[:, i : i + 1, ...] * (x_new + shift[:, i : i + 1, ...])
        )
    if not zero_mean:
        value_shift = torch.rand(
            [batch] + [1] * (x_new.dim() - 1), device=x.device, dtype=x.dtype
        )
        value_shift = (value_shift - 0.5) * 2 * (
            max_amplitude - min_amplitude
        ) * mean_shift_coef + min_amplitude
        y += value_shift
    if not batched:
        y = y.squeeze(0)
    return y