From 488edcb650c85d4342d6ade6f0304a1cb5ce7440 Mon Sep 17 00:00:00 2001 From: Qiang Date: Mon, 14 Oct 2024 21:06:48 +0200 Subject: [PATCH 1/2] add non-conservative form for convection operator --- exponax/nonlin_fun/_convection.py | 48 +++++++++++++++++++++++++++---- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/exponax/nonlin_fun/_convection.py b/exponax/nonlin_fun/_convection.py index 0e30f9a..ec42d7d 100644 --- a/exponax/nonlin_fun/_convection.py +++ b/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,23 +19,37 @@ 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 typical extension to higher dimensions requires u to have as many 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 @@ -56,17 +71,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"]: num_channels = u_hat.shape[0] @@ -78,12 +95,28 @@ def _multi_channel_eval( u = self.ifft(u_hat_dealiased) 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) + 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"] @@ -105,4 +138,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) From f07b32a95fb78482815b75092974608a5e2d3c28 Mon Sep 17 00:00:00 2001 From: "qiangliu.7@outlook.com" Date: Tue, 15 Oct 2024 20:15:53 +0200 Subject: [PATCH 2/2] update convection calculation to boardcast version and add conservative option to steppers --- exponax/nonlin_fun/_convection.py | 11 ++++++++--- exponax/stepper/_burgers.py | 6 ++++++ exponax/stepper/_korteweg_de_vries.py | 6 ++++++ exponax/stepper/_kuramoto_sivashinsky.py | 4 ++++ exponax/stepper/generic/_convection.py | 13 +++++++++++++ 5 files changed, 37 insertions(+), 3 deletions(-) diff --git a/exponax/nonlin_fun/_convection.py b/exponax/nonlin_fun/_convection.py index 3bcde05..693409d 100644 --- a/exponax/nonlin_fun/_convection.py +++ b/exponax/nonlin_fun/_convection.py @@ -117,7 +117,7 @@ def _multi_channel_conservative_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 = jnp.sum( self.derivative_operator[None, :] * u_outer_product_hat, @@ -136,9 +136,14 @@ def _multi_channel_nonconservative_eval( ) u_hat_dealiased = self.dealias(u_hat) u = self.ifft(u_hat_dealiased) - conv_u=sum( - [u[i:i+1]*self.ifft(self.derivative_operator[i:i+1]*u_hat) for i in range(num_channels)] + 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) diff --git a/exponax/stepper/_burgers.py b/exponax/stepper/_burgers.py index 866e983..ec4dfeb 100644 --- a/exponax/stepper/_burgers.py +++ b/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, ) diff --git a/exponax/stepper/_korteweg_de_vries.py b/exponax/stepper/_korteweg_de_vries.py index 0a54008..bbfe5c6 100644 --- a/exponax/stepper/_korteweg_de_vries.py +++ b/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, ) diff --git a/exponax/stepper/_kuramoto_sivashinsky.py b/exponax/stepper/_kuramoto_sivashinsky.py index 8596bfa..7acabbf 100644 --- a/exponax/stepper/_kuramoto_sivashinsky.py +++ b/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, ) diff --git a/exponax/stepper/generic/_convection.py b/exponax/stepper/generic/_convection.py index 00df9c2..b4241f4 100644 --- a/exponax/stepper/generic/_convection.py +++ b/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, )