Change the default way of normalization

exponax/normalized_stepper/convection.py

@@ -18,8 +18,8 @@ def __init__(
         num_points: int,
         *,
         dt: float = 0.1,
-        normalized_coefficients: list[float] = [0.0, 0.0, 0.01 * 0.1],
-        normalized_convection_scale: float = 0.5,
+        normalized_coefficients: list[float] = [0.0, 0.0, 0.01],
+        normalized_convection_scale: float = 1.0,
         order: int = 2,
         dealiasing_fraction: float = 2 / 3,
         n_circle_points: int = 16,
@@ -61,7 +61,7 @@ def __init__(
           derivative (default: [0.0, 0.0, 0.01 * 0.1] refers to a diffusion
           (2nd) order term)
         - `normalized_convection_scale`: convection scale for the nonlinear
-            function (default: 0.5)
+            function (default: 1.0)
         - `order`: order of exponential time differencing Runge Kutta method,
           can be 1, 2, 3, 4 (default: 2)
         - `dealiasing_fraction`: fraction of the wavenumbers being kept before
@@ -91,7 +91,7 @@ def _build_linear_operator(self, derivative_operator: Array) -> Array:
         # Now the linear operator is unscaled
         linear_operator = sum(
             jnp.sum(
-                c / self.dt * (derivative_operator) ** i,
+                c * (derivative_operator) ** i,
                 axis=0,
                 keepdims=True,
             )
@@ -106,5 +106,5 @@ def _build_nonlinear_fun(self, derivative_operator: Array):
             num_channels=self.num_channels,
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
-            convection_scale=self.normalized_convection_scale,
+            scale=self.normalized_convection_scale,
         )

exponax/normalized_stepper/gradient_norm.py

@@ -18,8 +18,10 @@ def __init__(
         num_points: int,
         *,
         dt: float = 0.1,
-        normalized_coefficients: list[float] = [0.0, 0.0, 0.01 * 0.1],
-        normalized_gradient_norm_scale: float = 0.5,
+        normalized_coefficients: list[float] = [
+            0.0, 0.0, -1.0 / (60.0**2), 0.0, -1.0 / (60.0**4)
+        ],
+        normalized_gradient_norm_scale: float = 1.0 / (60.0**2),
         order: int = 2,
         dealiasing_fraction: float = 2 / 3,
         n_circle_points: int = 16,
@@ -29,16 +31,16 @@ def __init__(
         the number of channels do **not** grow with the number of spatial
         dimensions. They are always 1.
 
+        By default: the KS equation on L=60.0
+
         **Arguments:**
         - `num_spatial_dims`: number of spatial dimensions
         - `num_points`: number of points in each spatial dimension
         - `dt`: time step (default: 0.1)
         - `normalized_coefficients`: coefficients for the linear operator,
           `normalized_coefficients[i]` is the coefficient for the `i`-th
-          derivative (default: [0.0, 0.0, 0.01 * 0.1] refers to a diffusion
-          operator)
+          derivative 
         - `normalized_gradient_norm_scale`: scale for the gradient norm
-            (default: 0.5)
         - `order`: order of the derivative operator (default: 2)
         - `dealiasing_fraction`: fraction of the wavenumbers being kept before
           applying any nonlinearity (default: 2/3)
@@ -66,7 +68,7 @@ def __init__(
     def _build_linear_operator(self, derivative_operator: Array) -> Array:
         linear_operator = sum(
             jnp.sum(
-                c / self.dt * (derivative_operator) ** i,
+                c * (derivative_operator) ** i,
                 axis=0,
                 keepdims=True,
             )

exponax/normalized_stepper/linear.py

@@ -16,19 +16,22 @@ def __init__(
         num_points: int,
         *,
         normalized_coefficients: list[float] = [0.0, -0.5, 0.01],
+        dt: float = 1.0,
     ):
         """
         By default: advection-diffusion with normalized advection of 0.5, and
         normalized diffusion of 0.01.
 
         Take care of the signs!
+
+        Normalized coefficients are alpha_i / L^i, where L is the domain extent.
         """
         self.normalized_coefficients = normalized_coefficients
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             domain_extent=1.0,  # Derivative operator is just scaled with 2 * jnp.pi
             num_points=num_points,
-            dt=1.0,
+            dt=dt,
             num_channels=1,
             order=0,
         )

exponax/normalized_stepper/utils.py

@@ -3,7 +3,6 @@
 
 def normalize_coefficients(
     domain_extent: float,
-    dt: float,
     coefficients: tuple[float],
 ) -> tuple[float]:
     """
@@ -12,19 +11,17 @@ def normalize_coefficients(
 
     **Arguments:**
     - `domain_extent`: extent of the domain
-    - `dt`: time step
     - `coefficients`: coefficients for the linear operator, `coefficients[i]` is
         the coefficient for the `i`-th derivative
     """
     normalized_coefficients = tuple(
-        c * dt / (domain_extent**i) for i, c in enumerate(coefficients)
+        c  / (domain_extent**i) for i, c in enumerate(coefficients)
     )
     return normalized_coefficients
 
 
 def denormalize_coefficients(
     domain_extent: float,
-    dt: float,
     normalized_coefficients: tuple[float],
 ) -> tuple[float]:
     """
@@ -33,13 +30,12 @@ def denormalize_coefficients(
 
     **Arguments:**
     - `domain_extent`: extent of the domain
-    - `dt`: time step
     - `normalized_coefficients`: coefficients for the linear operator,
         `normalized_coefficients[i]` is the coefficient for the `i`-th
         derivative
     """
     coefficients = tuple(
-        c_n / dt * domain_extent**i for i, c_n in enumerate(normalized_coefficients)
+        c_n * domain_extent**i for i, c_n in enumerate(normalized_coefficients)
     )
     return coefficients