add non-conservative form for convection operator

exponax/nonlin_fun/_convection.py

@@ -8,6 +8,7 @@ class ConvectionNonlinearFun(BaseNonlinearFun):
     derivative_operator: Complex[Array, "D ... (N//2)+1"]
     scale: float
     single_channel: bool
+    conservative: bool
 
     def __init__(
         self,
@@ -18,13 +19,14 @@ def __init__(
         dealiasing_fraction: float = 2 / 3,
         scale: float = 1.0,
         single_channel: bool = False,
+        conservative: bool = False,
     ):
         """
         Performs a pseudo-spectral evaluation of the nonlinear convection, e.g.,
         found in the Burgers equation. In 1d and state space, this reads
 
         ```
-            𝒩(u) = - b₁ 1/2 (u²)ₓ
+            𝒩(u) = b₁ u (u)ₓ
         ```
 
         with a scale `b₁`. The minus arises because `Exponax` follows the
@@ -37,10 +39,23 @@ def __init__(
         channels as spatial dimensions and then gives
 
         ```
-            𝒩(u) = - b₁ 1/2 ∇ ⋅ (u ⊗ u)
+            𝒩(u) = b₁ u ⋅ ∇ u
         ```
-
         with `∇ ⋅` the divergence operator and the outer product `u ⊗ u`.
+
+        Meanwhile, if you use a conservative form, the convection term is given by
+
+        ```
+            𝒩(u) = b₁ u (u)ₓ
+        ```
+        for 1D and
+
+        ```
+            𝒩(u) = b₁ ∇ ⋅ (u ⊗ u)
+        ```
+
+        for 2D and 3D.
+
         Another option is a "single-channel" hack requiring only one channel no
         matter the spatial dimensions. This reads
 
@@ -49,7 +64,6 @@ def __init__(
         ```
 
         **Arguments:**
-
         - `num_spatial_dims`: The number of spatial dimensions `d`.
         - `num_points`: The number of points `N` used to discretize the
             domain. This **includes** the left boundary point and **excludes**
@@ -63,17 +77,19 @@ def __init__(
         - `scale`: The scale `b₁` of the convection term. Defaults to `1.0`.
         - `single_channel`: Whether to use the single-channel hack. Defaults
             to `False`.
+        - `conservative`: Whether to use the conservative form. Defaults to `False`.
         """
         self.derivative_operator = derivative_operator
         self.scale = scale
         self.single_channel = single_channel
+        self.conservative=conservative
         super().__init__(
             num_spatial_dims,
             num_points,
             dealiasing_fraction=dealiasing_fraction,
         )
 
-    def _multi_channel_eval(
+    def _multi_channel_conservative_eval(
         self, u_hat: Complex[Array, "C ... (N//2)+1"]
     ) -> Complex[Array, "C ... (N//2)+1"]:
         """
@@ -101,14 +117,35 @@ def _multi_channel_eval(
             )
         u_hat_dealiased = self.dealias(u_hat)
         u = self.ifft(u_hat_dealiased)
-        u_outer_product = u[:, None] * u[None, :]
+        u_outer_product = u[None, :] * u[:, None]
         u_outer_product_hat = self.fft(u_outer_product)
-        convection = 0.5 * jnp.sum(
+        convection = jnp.sum(
             self.derivative_operator[None, :] * u_outer_product_hat,
             axis=1,
         )
         # Requires minus to move term to the rhs
         return -self.scale * convection
+
+    def _multi_channel_nonconservative_eval(
+        self, u_hat: Complex[Array, "C ... (N//2)+1"]
+    ) -> Complex[Array, "C ... (N//2)+1"]:
+        num_channels = u_hat.shape[0]
+        if num_channels != self.num_spatial_dims:
+            raise ValueError(
+                "Number of channels in u_hat should match number of spatial dimensions"
+            )
+        u_hat_dealiased = self.dealias(u_hat)
+        u = self.ifft(u_hat_dealiased)
+        nabla_u = self.ifft(self.derivative_operator[None, :] * u_hat[:, None])
+        conv_u = jnp.sum(
+            u[None, :] * nabla_u,
+            axis=1,
+        )
+        #conv_u=sum(
+        #    [u[i:i+1]*self.ifft(self.derivative_operator[i:i+1]*u_hat) for i in range(num_channels)]
+        #)
+        # Requires minus to move term to the rhs
+        return -self.scale * self.fft(conv_u)
 
     def _single_channel_eval(
         self, u_hat: Complex[Array, "C ... (N//2)+1"]
@@ -148,4 +185,7 @@ def __call__(
         if self.single_channel:
             return self._single_channel_eval(u_hat)
         else:
-            return self._multi_channel_eval(u_hat)
+            if self.conservative:
+                return self._multi_channel_conservative_eval(u_hat)
+            else:
+                return self._multi_channel_nonconservative_eval(u_hat)

exponax/stepper/_burgers.py

@@ -10,6 +10,7 @@ class Burgers(BaseStepper):
     convection_scale: float
     dealiasing_fraction: float
     single_channel: bool
+    conservative: bool
 
     def __init__(
         self,
@@ -21,6 +22,7 @@ def __init__(
         diffusivity: float = 0.1,
         convection_scale: float = 1.0,
         single_channel: bool = False,
+        conservative: bool = False,
         order=2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -75,6 +77,8 @@ def __init__(
             dimensions. In this case the the convection is `b₁ (∇ ⋅ 1)(u²)`. In
             this case, the state always has a single channel, no matter the
             spatial dimension. Default: False.
+        - `conservative`: Whether to use the conservative form of the convection
+            term. Default: False.
         - `order`: The order of the Exponential Time Differencing Runge
             Kutta method. Must be one of {0, 1, 2, 3, 4}. The option `0` only
             solves the linear part of the equation. Use higher values for higher
@@ -112,6 +116,7 @@ def __init__(
         self.diffusivity = diffusivity
         self.convection_scale = convection_scale
         self.single_channel = single_channel
+        self.conservative = conservative
         self.dealiasing_fraction = dealiasing_fraction
 
         if single_channel:
@@ -149,4 +154,5 @@ def _build_nonlinear_fun(
             dealiasing_fraction=self.dealiasing_fraction,
             scale=self.convection_scale,
             single_channel=self.single_channel,
+            conservative=False,
         )

exponax/stepper/_korteweg_de_vries.py

@@ -19,6 +19,7 @@ class KortewegDeVries(BaseStepper):
     advect_over_diffuse: bool
     diffuse_over_diffuse: bool
     single_channel: bool
+    conservaitve: bool
 
     def __init__(
         self,
@@ -34,6 +35,7 @@ def __init__(
         advect_over_diffuse: bool = False,
         diffuse_over_diffuse: bool = False,
         single_channel: bool = False,
+        conservative: bool = False,
         order: int = 2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -108,6 +110,8 @@ def __init__(
             dimensions. In this case the the convection is `b₁ (∇ ⋅ 1)(u²)`. In
             this case, the state always has a single channel, no matter the
             spatial dimension. Default: False.
+        - `conservative`: Whether to use the conservative form of the convection
+            term. Default: False.
         - `order`: The order of the Exponential Time Differencing Runge
             Kutta method. Must be one of {0, 1, 2, 3, 4}. The option `0` only
             solves the linear part of the equation. Use higher values for higher
@@ -145,6 +149,7 @@ def __init__(
         self.advect_over_diffuse = advect_over_diffuse
         self.diffuse_over_diffuse = diffuse_over_diffuse
         self.single_channel = single_channel
+        self.conservative = conservative
         self.dealiasing_fraction = dealiasing_fraction
 
         if single_channel:
@@ -210,4 +215,5 @@ def _build_nonlinear_fun(
             dealiasing_fraction=self.dealiasing_fraction,
             scale=self.convection_scale,
             single_channel=self.single_channel,
+            conservative=self.conservative,
         )

exponax/stepper/_kuramoto_sivashinsky.py

@@ -176,6 +176,7 @@ class KuramotoSivashinskyConservative(BaseStepper):
     second_order_diffusivity: float
     fourth_order_diffusivity: float
     single_channel: bool
+    conservative: bool
     dealiasing_fraction: float
 
     def __init__(
@@ -189,6 +190,7 @@ def __init__(
         second_order_diffusivity: float = 1.0,
         fourth_order_diffusivity: float = 1.0,
         single_channel: bool = False,
+        conservative: bool = False,
         dealiasing_fraction: float = 2 / 3,
         order: int = 2,
         num_circle_points: int = 16,
@@ -203,6 +205,7 @@ def __init__(
         self.second_order_diffusivity = second_order_diffusivity
         self.fourth_order_diffusivity = fourth_order_diffusivity
         self.single_channel = single_channel
+        self.conservative = conservative
         self.dealiasing_fraction = dealiasing_fraction
 
         if num_spatial_dims > 1:
@@ -252,4 +255,5 @@ def _build_nonlinear_fun(
             dealiasing_fraction=self.dealiasing_fraction,
             scale=self.convection_scale,
             single_channel=self.single_channel,
+            conservative=self.conservative,
         )

exponax/stepper/generic/_convection.py

@@ -14,6 +14,7 @@ class GeneralConvectionStepper(BaseStepper):
     convection_scale: float
     dealiasing_fraction: float
     single_channel: bool
+    conservative: bool
 
     def __init__(
         self,
@@ -25,6 +26,7 @@ def __init__(
         coefficients: tuple[float, ...] = (0.0, 0.0, 0.01),
         convection_scale: float = 1.0,
         single_channel: bool = False,
+        conservative: bool = False,
         order=2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -83,6 +85,8 @@ def __init__(
             dimensions. In this case the the convection is `b₁ (∇ ⋅ 1)(u²)`. In
             this case, the state always has a single channel, no matter the
             spatial dimension. Default: False.
+        - `conservative`: Whether to use the conservative form of the convection
+            term. Default: False.
         - `order`: The order of the Exponential Time Differencing Runge
             Kutta method. Must be one of {0, 1, 2, 3, 4}. The option `0` only
             solves the linear part of the equation. Use higher values for higher
@@ -103,6 +107,7 @@ def __init__(
         self.convection_scale = convection_scale
         self.single_channel = single_channel
         self.dealiasing_fraction = dealiasing_fraction
+        self.conservative = conservative
 
         if single_channel:
             num_channels = 1
@@ -146,6 +151,7 @@ def _build_nonlinear_fun(
             dealiasing_fraction=self.dealiasing_fraction,
             scale=self.convection_scale,
             single_channel=self.single_channel,
+            conservative=self.conservative,
         )
 
 
@@ -161,6 +167,7 @@ def __init__(
         normalized_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1),
         normalized_convection_scale: float = 1.0 * 0.1,
         single_channel: bool = False,
+        conservative: bool = False,
         order: int = 2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -210,6 +217,8 @@ def __init__(
             dimensions. In this case the the convection is `β (∇ ⋅ 1)(u²)`. In
             this case, the state always has a single channel, no matter the
             spatial dimension. Default: False.
+        - `conservative`: Whether to use the conservative form of the convection
+            term. Default: False.
         - `order`: The order of the Exponential Time Differencing Runge
             Kutta method. Must be one of {0, 1, 2, 3, 4}. The option `0` only
             solves the linear part of the equation. Use higher values for higher
@@ -240,6 +249,7 @@ def __init__(
             num_circle_points=num_circle_points,
             circle_radius=circle_radius,
             single_channel=single_channel,
+            conservative=conservative,
         )
 
 
@@ -255,6 +265,7 @@ def __init__(
         linear_difficulties: tuple[float, ...] = (0.0, 0.0, 4.5),
         convection_difficulty: float = 5.0,
         single_channel: bool = False,
+        conservative: bool = False,
         maximum_absolute: float = 1.0,
         order: int = 2,
         dealiasing_fraction: float = 2 / 3,
@@ -315,6 +326,7 @@ def __init__(
             dimensions. In this case the the convection is `δ (∇ ⋅ 1)(u²)`. In
             this case, the state always has a single channel, no matter the
             spatial dimension. Default: False.
+        - `conservative`: Whether to use the conservative form of the convection
         - `maximum_absolute`: The maximum absolute value of the state. This is
             used to extract the normalized dynamics from the convection
             difficulty.
@@ -359,4 +371,5 @@ def __init__(
             dealiasing_fraction=dealiasing_fraction,
             num_circle_points=num_circle_points,
             circle_radius=circle_radius,
+            conservative=conservative,
         )