Change the default scaling for advection-like nonlinarities

exponax/nonlinear_functions/convection.py

@@ -11,7 +11,7 @@
 
 
 class ConvectionNonlinearFun(BaseNonlinearFun):
-    convection_scale: float
+    scale: float
     zero_mode_fix: bool
 
     def __init__(
@@ -22,10 +22,13 @@ def __init__(
         *,
         derivative_operator: Complex[Array, "D ... (N//2)+1"],
         dealiasing_fraction: float,
-        convection_scale: float = 0.5,
+        scale: float = 1.0,
         zero_mode_fix: bool = False,
     ):
-        self.convection_scale = convection_scale
+        """
+        Uses by default a scaling of 0.5 to take into account the conservative evaluation
+        """
+        self.scale = scale
         self.zero_mode_fix = zero_mode_fix
         super().__init__(
             num_spatial_dims,
@@ -65,4 +68,4 @@ def evaluate(
             axis=1,
         )
         # Requires minus to move term to the rhs
-        return -self.convection_scale * u_divergence_on_outer_product_hat
+        return - self.scale * 0.5 * u_divergence_on_outer_product_hat

exponax/nonlinear_functions/gradient_norm.py

@@ -23,8 +23,11 @@ def __init__(
         derivative_operator: Complex[Array, "D ... (N//2)+1"],
         dealiasing_fraction: float,
         zero_mode_fix: bool = True,
-        scale: float = 0.5,
+        scale: float = 1.0,
     ):
+        """
+        Uses by default a scaling of 0.5 to take into account the conservative evaluation
+        """
         super().__init__(
             num_spatial_dims,
             num_points,
@@ -70,4 +73,4 @@ def evaluate(
         #     )
 
         # Requires minus to move term to the rhs
-        return -self.scale * u_gradient_norm_squared_hat
+        return - self.scale * 0.5 * u_gradient_norm_squared_hat

exponax/sample_stepper/burgers.py

@@ -21,12 +21,16 @@ def __init__(
         dt: float,
         *,
         diffusivity: float = 0.1,
-        convection_scale: float = 0.5,
+        convection_scale: float = 1.0,
         order=2,
         dealiasing_fraction: float = 2 / 3,
         n_circle_points: int = 16,
         circle_radius: float = 1.0,
     ):
+        """
+        Convection is always scaled by 0.5, use `convection_scale` to multiply
+        an additional factor.
+        """
         self.diffusivity = diffusivity
         self.convection_scale = convection_scale
         self.dealiasing_fraction = dealiasing_fraction
@@ -59,5 +63,5 @@ def _build_nonlinear_fun(
             num_channels=self.num_channels,
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
-            convection_scale=self.convection_scale,
+            scale=self.convection_scale,
         )

exponax/sample_stepper/convection.py

@@ -20,7 +20,7 @@ def __init__(
         dt: float,
         *,
         coefficients: list[float] = [0.0, 0.0, 0.01],
-        convection_scale: float = 0.5,
+        convection_scale: float = 1.0,
         order=2,
         dealiasing_fraction: float = 2 / 3,
         n_circle_points: int = 16,
@@ -70,6 +70,6 @@ def _build_nonlinear_fun(
             num_channels=self.num_channels,
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
-            convection_scale=self.convection_scale,
+            scale=self.convection_scale,
             zero_mode_fix=False,  # Todo: check this
         )

exponax/sample_stepper/gradient_norm.py

@@ -20,7 +20,7 @@ def __init__(
         dt: float,
         *,
         coefficients: list[float] = [0.0, 0.0, -1.0, 0.0, -1.0],
-        gradient_norm_scale: float = 0.5,
+        gradient_norm_scale: float = 1.0,
         order=2,
         dealiasing_fraction: float = 2 / 3,
         n_circle_points: int = 16,

exponax/sample_stepper/korteveg_de_vries.py

@@ -24,7 +24,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        convection_scale: float = -6 / 2,
+        convection_scale: float = -6.0,
         pure_dispersivity: Union[Float[Array, "D"], float] = 1.0,
         advect_over_diffuse_dispersivity: Union[Float[Array, "D"], float] = 0.0,
         diffusivity: float = 0.0,
@@ -82,5 +82,5 @@ def _build_nonlinear_fun(
             num_channels=self.num_channels,
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
-            convection_scale=self.convection_scale,
+            scale=self.convection_scale,
         )

exponax/sample_stepper/kuramoto_sivashinsky.py

@@ -12,6 +12,7 @@
 
 
 class KuramotoSivashinsky(BaseStepper):
+    gradient_norm_scale: float
     second_order_diffusivity: float
     fourth_order_diffusivity: float
     dealiasing_fraction: float
@@ -23,6 +24,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
+        gradient_norm_scale: float = 1.0,
         second_order_diffusivity: float = 1.0,
         fourth_order_diffusivity: float = 1.0,
         dealiasing_fraction: float = 2 / 3,
@@ -36,6 +38,7 @@ def __init__(
         The advantage is that the number of channels is always 1 no matter the
         number of spatial dimensions.
         """
+        self.gradient_norm_scale = gradient_norm_scale
         self.second_order_diffusivity = second_order_diffusivity
         self.fourth_order_diffusivity = fourth_order_diffusivity
         self.dealiasing_fraction = dealiasing_fraction
@@ -72,11 +75,12 @@ def _build_nonlinear_fun(
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
             zero_mode_fix=True,
-            scale=0.5,
+            scale=self.gradient_norm_scale,
         )
 
 
 class KuramotoSivashinskyConservative(BaseStepper):
+    convection_scale: float
     second_order_diffusivity: float
     fourth_order_diffusivity: float
     dealiasing_fraction: float
@@ -88,6 +92,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
+        convection_scale: float = 1.0,
         second_order_diffusivity: float = 1.0,
         fourth_order_diffusivity: float = 1.0,
         dealiasing_fraction: float = 2 / 3,
@@ -100,6 +105,7 @@ def __init__(
         Burgers equation). This also means that the number of channels grow with
         the number of spatial dimensions.
         """
+        self.convection_scale = convection_scale
         self.second_order_diffusivity = second_order_diffusivity
         self.fourth_order_diffusivity = fourth_order_diffusivity
         self.dealiasing_fraction = dealiasing_fraction
@@ -136,5 +142,5 @@ def _build_nonlinear_fun(
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
             zero_mode_fix=True,
-            convection_scale=0.5,
+            scale=self.convection_scale,
         )

exponax/sample_stepper/nikolaevskiy.py

@@ -12,6 +12,7 @@
 
 
 class Nikolaevskiy(BaseStepper):
+    gradient_norm_scale: float
     second_order_diffusivity: float
     fourth_order_diffusivity: float
     sixth_order_diffusivity: float
@@ -24,6 +25,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
+        gradient_norm_scale: float = 1.0,
         second_order_diffusivity: float = 0.1,
         fourth_order_diffusivity: float = 1.0,
         sixth_order_diffusivity: float = 1.0,
@@ -32,6 +34,7 @@ def __init__(
         n_circle_points: int = 16,
         circle_radius: float = 1.0,
     ):
+        self.gradient_norm_scale = gradient_norm_scale
         self.second_order_diffusivity = second_order_diffusivity
         self.fourth_order_diffusivity = fourth_order_diffusivity
         self.sixth_order_diffusivity = sixth_order_diffusivity
@@ -72,11 +75,12 @@ def _build_nonlinear_fun(
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
             zero_mode_fix=True,
-            scale=0.5,
+            scale=self.gradient_norm_scale,
         )
 
 
 class NikolaevskiyConservative(BaseStepper):
+    convection_scale: float
     second_order_diffusivity: float
     fourth_order_diffusivity: float
     sixth_order_diffusivity: float
@@ -89,6 +93,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
+        convection_scale: float = 1.0,
         second_order_diffusivity: float = 0.1,
         fourth_order_diffusivity: float = 1.0,
         sixth_order_diffusivity: float = 1.0,
@@ -97,6 +102,7 @@ def __init__(
         n_circle_points: int = 16,
         circle_radius: float = 1.0,
     ):
+        self.convection_scale = convection_scale
         self.second_order_diffusivity = second_order_diffusivity
         self.fourth_order_diffusivity = fourth_order_diffusivity
         self.sixth_order_diffusivity = sixth_order_diffusivity
@@ -137,5 +143,5 @@ def _build_nonlinear_fun(
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
             zero_mode_fix=True,
-            convection_scale=0.5,
+            scale=self.convection_scale,
         )