Skip to content

Convection¤

exponax.normalized.DifficultyConvectionStepper ¤

Bases: NormalizedConvectionStepper

Source code in exponax/normalized/_convection.py
 73
 74
 75
 76
 77
 78
 79
 80
 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
class DifficultyConvectionStepper(NormalizedConvectionStepper):
    linear_difficulties: tuple[float, ...]
    convection_difficulty: float

    def __init__(
        self,
        num_spatial_dims: int = 1,
        num_points: int = 48,
        *,
        linear_difficulties: tuple[float, ...] = (0.0, 0.0, 4.5),
        convection_difficulty: float = 5.0,
        maximum_absolute: float = 1.0,
        order: int = 2,
        dealiasing_fraction: float = 2 / 3,
        num_circle_points: int = 16,
        circle_radius: float = 1.0,
    ):
        """
        By default: Behaves like a Burgers

        """
        self.linear_difficulties = linear_difficulties
        self.convection_difficulty = convection_difficulty
        normalized_coefficients = extract_normalized_coefficients_from_difficulty(
            linear_difficulties,
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
        )
        normalized_convection_scale = (
            extract_normalized_convection_scale_from_difficulty(
                convection_difficulty,
                num_spatial_dims=num_spatial_dims,
                num_points=num_points,
                maximum_absolute=maximum_absolute,
            )
        )
        super().__init__(
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            normalized_coefficients=normalized_coefficients,
            normalized_convection_scale=normalized_convection_scale,
            order=order,
            dealiasing_fraction=dealiasing_fraction,
            num_circle_points=num_circle_points,
            circle_radius=circle_radius,
        )
__init__ ¤
__init__(
    num_spatial_dims: int = 1,
    num_points: int = 48,
    *,
    linear_difficulties: tuple[float, ...] = (
        0.0,
        0.0,
        4.5,
    ),
    convection_difficulty: float = 5.0,
    maximum_absolute: float = 1.0,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0
)

By default: Behaves like a Burgers

Source code in exponax/normalized/_convection.py
 77
 78
 79
 80
 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
def __init__(
    self,
    num_spatial_dims: int = 1,
    num_points: int = 48,
    *,
    linear_difficulties: tuple[float, ...] = (0.0, 0.0, 4.5),
    convection_difficulty: float = 5.0,
    maximum_absolute: float = 1.0,
    order: int = 2,
    dealiasing_fraction: float = 2 / 3,
    num_circle_points: int = 16,
    circle_radius: float = 1.0,
):
    """
    By default: Behaves like a Burgers

    """
    self.linear_difficulties = linear_difficulties
    self.convection_difficulty = convection_difficulty
    normalized_coefficients = extract_normalized_coefficients_from_difficulty(
        linear_difficulties,
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
    )
    normalized_convection_scale = (
        extract_normalized_convection_scale_from_difficulty(
            convection_difficulty,
            num_spatial_dims=num_spatial_dims,
            num_points=num_points,
            maximum_absolute=maximum_absolute,
        )
    )
    super().__init__(
        num_spatial_dims=num_spatial_dims,
        num_points=num_points,
        normalized_coefficients=normalized_coefficients,
        normalized_convection_scale=normalized_convection_scale,
        order=order,
        dealiasing_fraction=dealiasing_fraction,
        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)