Skip to content

Vorticity Convection¤

exponax.normalized.NormalizedVorticityConvection ¤

Bases: BaseStepper

Source code in exponax/normalized/_vorticity_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
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class NormalizedVorticityConvection(BaseStepper):
    normalized_vorticity_convection_scale: float
    normalized_coefficients: tuple[float, ...]
    injection_mode: int
    normalized_injection_scale: float
    dealiasing_fraction: float

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        normalized_vorticity_convection_scale: float = 0.01 * 1.0 / (1.0**0),
        normalized_coefficients: tuple[float, ...] = (
            0.01 * 0.0 / (1.0**0),
            0.0,
            0.01 * 0.001 / (1.0**2),
        ),
        injection_mode: int = 4,
        normalized_injection_scale: float = 0.01 * 0.0 / (1.0**0),
        order: int = 2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        if num_spatial_dims != 2:
            raise ValueError(f"Expected num_spatial_dims = 2, got {num_spatial_dims}.")
        self.normalized_vorticity_convection_scale = (
            normalized_vorticity_convection_scale
        )
        self.normalized_coefficients = normalized_coefficients
        self.injection_mode = injection_mode
        self.normalized_injection_scale = normalized_injection_scale
        self.dealiasing_fraction = dealiasing_fraction
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            domain_extent=1.0,
            num_points=num_points,
            dt=1.0,
            num_channels=1,
            order=order,
            num_circle_points=num_circle_points,
            circle_radius=circle_radius,
        )

    def _build_linear_operator(
        self,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
    ) -> Complex[Array, "1 ... (N//2)+1"]:
        linear_operator = sum(
            jnp.sum(
                c * (derivative_operator) ** i,
                axis=0,
                keepdims=True,
            )
            for i, c in enumerate(self.normalized_coefficients)
        )
        return linear_operator

    def _build_nonlinear_fun(
        self,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
    ) -> VorticityConvection2d:
        if self.normalized_injection_scale == 0.0:
            return VorticityConvection2d(
                self.num_spatial_dims,
                self.num_points,
                convection_scale=self.normalized_vorticity_convection_scale,
                derivative_operator=derivative_operator,
                dealiasing_fraction=self.dealiasing_fraction,
            )
        else:
            return VorticityConvection2dKolmogorov(
                self.num_spatial_dims,
                self.num_points,
                convection_scale=self.normalized_vorticity_convection_scale,
                derivative_operator=derivative_operator,
                dealiasing_fraction=self.dealiasing_fraction,
                injection_mode=self.injection_mode,
                injection_scale=self.normalized_injection_scale,
            )
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_vorticity_convection_scale: float = 0.01
    * 1.0
    / 1.0**0,
    normalized_coefficients: tuple[float, ...] = (
        0.01 * 0.0 / 1.0**0,
        0.0,
        0.01 * 0.001 / 1.0**2,
    ),
    injection_mode: int = 4,
    normalized_injection_scale: float = 0.01 * 0.0 / 1.0**0,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)
Source code in exponax/normalized/_vorticity_convection.py 
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
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    normalized_vorticity_convection_scale: float = 0.01 * 1.0 / (1.0**0),
    normalized_coefficients: tuple[float, ...] = (
        0.01 * 0.0 / (1.0**0),
        0.0,
        0.01 * 0.001 / (1.0**2),
    ),
    injection_mode: int = 4,
    normalized_injection_scale: float = 0.01 * 0.0 / (1.0**0),
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    if num_spatial_dims != 2:
        raise ValueError(f"Expected num_spatial_dims = 2, got {num_spatial_dims}.")
    self.normalized_vorticity_convection_scale = (
        normalized_vorticity_convection_scale
    )
    self.normalized_coefficients = normalized_coefficients
    self.injection_mode = injection_mode
    self.normalized_injection_scale = normalized_injection_scale
    self.dealiasing_fraction = dealiasing_fraction
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        domain_extent=1.0,
        num_points=num_points,
        dt=1.0,
        num_channels=1,
        order=order,
        num_circle_points=num_circle_points,
        circle_radius=circle_radius,
    )
__call__ ¤
__call__(
    u: Float[Array, "C ... N"]
) -> Float[Array, "C ... N"]

Performs a check

Source code in exponax/_base_stepper.py 
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
def __call__(
    self,
    u: Float[Array, "C ... N"],
) -> Float[Array, "C ... N"]:
    """
    Performs a check
    """
    expected_shape = (self.num_channels,) + spatial_shape(
        self.num_spatial_dims, self.num_points
    )
    if u.shape != expected_shape:
        raise ValueError(
            f"Expected shape {expected_shape}, got {u.shape}. For batched operation use `jax.vmap` on this function."
        )
    return self.step(u)