Skip to content

General Nonlinear Function¤

exponax.nonlin_fun.GeneralNonlinearFun ¤

Bases: BaseNonlinearFun

Source code in exponax/nonlin_fun/_general_nonlinear.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
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class GeneralNonlinearFun(BaseNonlinearFun):
    square_nonlinear_fun: PolynomialNonlinearFun
    convection_nonlinear_fun: ConvectionNonlinearFun
    gradient_norm_nonlinear_fun: GradientNormNonlinearFun

    def __init__(
        self,
        num_spatial_dims: int,
        num_points: int,
        *,
        derivative_operator: Complex[Array, "D ... (N//2)+1"],
        dealiasing_fraction: float,
        scale_list: tuple[float, ...] = [0.0, -1.0, 0.0],
        zero_mode_fix: bool = True,
    ):
        """
        Uses an additional scaling of 0.5 on the latter two components only

        By default: Burgers equation
        """
        if len(scale_list) != 3:
            raise ValueError("The scale list must have exactly 3 elements")

        self.square_nonlinear_fun = PolynomialNonlinearFun(
            num_spatial_dims,
            num_points,
            dealiasing_fraction=dealiasing_fraction,
            coefficients=[0.0, 0.0, scale_list[0]],
        )
        self.convection_nonlinear_fun = ConvectionNonlinearFun(
            num_spatial_dims,
            num_points,
            derivative_operator=derivative_operator,
            dealiasing_fraction=dealiasing_fraction,
            # Minus required because it internally has another minus
            scale=-scale_list[1],
            single_channel=True,
        )
        self.gradient_norm_nonlinear_fun = GradientNormNonlinearFun(
            num_spatial_dims,
            num_points,
            derivative_operator=derivative_operator,
            dealiasing_fraction=dealiasing_fraction,
            # Minus required because it internally has another minus
            scale=-scale_list[2],
            zero_mode_fix=zero_mode_fix,
        )

        super().__init__(
            num_spatial_dims,
            num_points,
            dealiasing_fraction=dealiasing_fraction,
        )

    def __call__(
        self,
        u_hat: Complex[Array, "C ... (N//2)+1"],
    ) -> Complex[Array, "C ... (N//2)+1"]:
        return (
            self.square_nonlinear_fun(u_hat)
            + self.convection_nonlinear_fun(u_hat)
            + self.gradient_norm_nonlinear_fun(u_hat)
        )
__init__ ¤
__init__(
    num_spatial_dims: int,
    num_points: int,
    *,
    derivative_operator: Complex[Array, "D ... (N//2)+1"],
    dealiasing_fraction: float,
    scale_list: tuple[float, ...] = [0.0, -1.0, 0.0],
    zero_mode_fix: bool = True
)

Uses an additional scaling of 0.5 on the latter two components only

By default: Burgers equation

Source code in exponax/nonlin_fun/_general_nonlinear.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
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def __init__(
    self,
    num_spatial_dims: int,
    num_points: int,
    *,
    derivative_operator: Complex[Array, "D ... (N//2)+1"],
    dealiasing_fraction: float,
    scale_list: tuple[float, ...] = [0.0, -1.0, 0.0],
    zero_mode_fix: bool = True,
):
    """
    Uses an additional scaling of 0.5 on the latter two components only

    By default: Burgers equation
    """
    if len(scale_list) != 3:
        raise ValueError("The scale list must have exactly 3 elements")

    self.square_nonlinear_fun = PolynomialNonlinearFun(
        num_spatial_dims,
        num_points,
        dealiasing_fraction=dealiasing_fraction,
        coefficients=[0.0, 0.0, scale_list[0]],
    )
    self.convection_nonlinear_fun = ConvectionNonlinearFun(
        num_spatial_dims,
        num_points,
        derivative_operator=derivative_operator,
        dealiasing_fraction=dealiasing_fraction,
        # Minus required because it internally has another minus
        scale=-scale_list[1],
        single_channel=True,
    )
    self.gradient_norm_nonlinear_fun = GradientNormNonlinearFun(
        num_spatial_dims,
        num_points,
        derivative_operator=derivative_operator,
        dealiasing_fraction=dealiasing_fraction,
        # Minus required because it internally has another minus
        scale=-scale_list[2],
        zero_mode_fix=zero_mode_fix,
    )

    super().__init__(
        num_spatial_dims,
        num_points,
        dealiasing_fraction=dealiasing_fraction,
    )
__call__ ¤
__call__(
    u_hat: Complex[Array, "C ... (N//2)+1"]
) -> Complex[Array, "C ... (N//2)+1"]
Source code in exponax/nonlin_fun/_general_nonlinear.py 
63
64
65
66
67
68
69
70
71
def __call__(
    self,
    u_hat: Complex[Array, "C ... (N//2)+1"],
) -> Complex[Array, "C ... (N//2)+1"]:
    return (
        self.square_nonlinear_fun(u_hat)
        + self.convection_nonlinear_fun(u_hat)
        + self.gradient_norm_nonlinear_fun(u_hat)
    )