diff --git a/docs/api/stepper/physical/poisson.md b/docs/api/stepper/additional/poisson.md similarity index 100% rename from docs/api/stepper/physical/poisson.md rename to docs/api/stepper/additional/poisson.md diff --git a/docs/api/stepper/difficulty/convection.md b/docs/api/stepper/generic/difficulty/difficulty_convection.md similarity index 61% rename from docs/api/stepper/difficulty/convection.md rename to docs/api/stepper/generic/difficulty/difficulty_convection.md index eafd182..8dec1e6 100644 --- a/docs/api/stepper/difficulty/convection.md +++ b/docs/api/stepper/generic/difficulty/difficulty_convection.md @@ -1,6 +1,6 @@ # Convection -::: exponax.normalized.DifficultyConvectionStepper +::: exponax.stepper.generic.DifficultyConvectionStepper options: members: - __init__ diff --git a/docs/api/stepper/difficulty/gradient_norm.md b/docs/api/stepper/generic/difficulty/difficulty_gradient_norm.md similarity index 60% rename from docs/api/stepper/difficulty/gradient_norm.md rename to docs/api/stepper/generic/difficulty/difficulty_gradient_norm.md index cc2fbae..7399145 100644 --- a/docs/api/stepper/difficulty/gradient_norm.md +++ b/docs/api/stepper/generic/difficulty/difficulty_gradient_norm.md @@ -1,6 +1,6 @@ # Gradient Norm -::: exponax.normalized.DifficultyGradientNormStepper +::: exponax.stepper.generic.DifficultyGradientNormStepper options: members: - __init__ diff --git a/docs/api/stepper/difficulty/linear.md b/docs/api/stepper/generic/difficulty/difficulty_linear.md similarity index 61% rename from docs/api/stepper/difficulty/linear.md rename to docs/api/stepper/generic/difficulty/difficulty_linear.md index 100621c..371bd27 100644 --- a/docs/api/stepper/difficulty/linear.md +++ b/docs/api/stepper/generic/difficulty/difficulty_linear.md @@ -1,6 +1,6 @@ # Linear -::: exponax.normalized.DifficultyLinearStepper +::: exponax.stepper.generic.DifficultyLinearStepper options: members: - __init__ diff --git a/docs/api/stepper/difficulty/nonlinear.md b/docs/api/stepper/generic/difficulty/difficulty_nonlinear.md similarity index 60% rename from docs/api/stepper/difficulty/nonlinear.md rename to docs/api/stepper/generic/difficulty/difficulty_nonlinear.md index 717189b..5962c39 100644 --- a/docs/api/stepper/difficulty/nonlinear.md +++ b/docs/api/stepper/generic/difficulty/difficulty_nonlinear.md @@ -1,6 +1,6 @@ # Nonlinear -::: exponax.normalized.DifficultyGeneralNonlinearStepper +::: exponax.stepper.generic.DifficultyNonlinearStepper options: members: - __init__ diff --git a/docs/api/stepper/difficulty/polynomial.md b/docs/api/stepper/generic/difficulty/difficulty_polynomial.md similarity index 60% rename from docs/api/stepper/difficulty/polynomial.md rename to docs/api/stepper/generic/difficulty/difficulty_polynomial.md index 882c62f..2dbb4b0 100644 --- a/docs/api/stepper/difficulty/polynomial.md +++ b/docs/api/stepper/generic/difficulty/difficulty_polynomial.md @@ -1,6 +1,6 @@ # Polynomial -::: exponax.normalized.DifficultyPolynomialStepper +::: exponax.stepper.generic.DifficultyPolynomialStepper options: members: - __init__ diff --git a/docs/api/stepper/normalized/convection.md b/docs/api/stepper/generic/normalized/normalized_convection.md similarity index 60% rename from docs/api/stepper/normalized/convection.md rename to docs/api/stepper/generic/normalized/normalized_convection.md index 7f3982d..6222489 100644 --- a/docs/api/stepper/normalized/convection.md +++ b/docs/api/stepper/generic/normalized/normalized_convection.md @@ -1,6 +1,6 @@ # Convection -::: exponax.normalized.NormalizedConvectionStepper +::: exponax.stepper.generic.NormalizedConvectionStepper options: members: - __init__ diff --git a/docs/api/stepper/normalized/gradient_norm.md b/docs/api/stepper/generic/normalized/normalized_gradient_norm.md similarity index 60% rename from docs/api/stepper/normalized/gradient_norm.md rename to docs/api/stepper/generic/normalized/normalized_gradient_norm.md index 64689e0..7715874 100644 --- a/docs/api/stepper/normalized/gradient_norm.md +++ b/docs/api/stepper/generic/normalized/normalized_gradient_norm.md @@ -1,6 +1,6 @@ # Gradient NormA -::: exponax.normalized.NormalizedGradientNormStepper +::: exponax.stepper.generic.NormalizedGradientNormStepper options: members: - __init__ diff --git a/docs/api/stepper/normalized/linear.md b/docs/api/stepper/generic/normalized/normalized_linear.md similarity index 59% rename from docs/api/stepper/normalized/linear.md rename to docs/api/stepper/generic/normalized/normalized_linear.md index 931d8a2..8bf37f4 100644 --- a/docs/api/stepper/normalized/linear.md +++ b/docs/api/stepper/generic/normalized/normalized_linear.md @@ -1,5 +1,5 @@ -::: exponax.normalized.NormalizedLinearStepper +::: exponax.stepper.generic.NormalizedLinearStepper options: members: - __init__ diff --git a/docs/api/stepper/normalized/nonlinear.md b/docs/api/stepper/generic/normalized/normalized_nonlinear.md similarity index 59% rename from docs/api/stepper/normalized/nonlinear.md rename to docs/api/stepper/generic/normalized/normalized_nonlinear.md index 2171583..379866e 100644 --- a/docs/api/stepper/normalized/nonlinear.md +++ b/docs/api/stepper/generic/normalized/normalized_nonlinear.md @@ -1,6 +1,6 @@ # Nonlinear -::: exponax.normalized.NormalizedGeneralNonlinearStepper +::: exponax.stepper.generic.NormalizedNonlinearStepper options: members: - __init__ diff --git a/docs/api/stepper/normalized/polynomial.md b/docs/api/stepper/generic/normalized/normalized_polynomial.md similarity index 60% rename from docs/api/stepper/normalized/polynomial.md rename to docs/api/stepper/generic/normalized/normalized_polynomial.md index 294c9f2..b6bd71d 100644 --- a/docs/api/stepper/normalized/polynomial.md +++ b/docs/api/stepper/generic/normalized/normalized_polynomial.md @@ -1,6 +1,6 @@ # Polynomial -::: exponax.normalized.NormalizedPolynomialStepper +::: exponax.stepper.generic.NormalizedPolynomialStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/general/general_convection.md b/docs/api/stepper/generic/physical/general_convection.md similarity index 66% rename from docs/api/stepper/physical/general/general_convection.md rename to docs/api/stepper/generic/physical/general_convection.md index 64b8920..851876f 100644 --- a/docs/api/stepper/physical/general/general_convection.md +++ b/docs/api/stepper/generic/physical/general_convection.md @@ -1,6 +1,6 @@ # General Convection Stepper -::: exponax.stepper.GeneralConvectionStepper +::: exponax.stepper.generic.GeneralConvectionStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/general/general_gradient_norm.md b/docs/api/stepper/generic/physical/general_gradient_norm.md similarity index 66% rename from docs/api/stepper/physical/general/general_gradient_norm.md rename to docs/api/stepper/generic/physical/general_gradient_norm.md index 0be3dba..aa37fd9 100644 --- a/docs/api/stepper/physical/general/general_gradient_norm.md +++ b/docs/api/stepper/generic/physical/general_gradient_norm.md @@ -1,6 +1,6 @@ # General Gradient Norm Stepper -::: exponax.stepper.GeneralGradientNormStepper +::: exponax.stepper.generic.GeneralGradientNormStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/general/general_linear.md b/docs/api/stepper/generic/physical/general_linear.md similarity index 67% rename from docs/api/stepper/physical/general/general_linear.md rename to docs/api/stepper/generic/physical/general_linear.md index e867c51..f99caf4 100644 --- a/docs/api/stepper/physical/general/general_linear.md +++ b/docs/api/stepper/generic/physical/general_linear.md @@ -1,6 +1,6 @@ # General Linear Stepper -::: exponax.stepper.GeneralLinearStepper +::: exponax.stepper.generic.GeneralLinearStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/general/general_nonlinear.md b/docs/api/stepper/generic/physical/general_nonlinear.md similarity index 66% rename from docs/api/stepper/physical/general/general_nonlinear.md rename to docs/api/stepper/generic/physical/general_nonlinear.md index 8714da1..36a5595 100644 --- a/docs/api/stepper/physical/general/general_nonlinear.md +++ b/docs/api/stepper/generic/physical/general_nonlinear.md @@ -1,6 +1,6 @@ # General Nonlinear Stepper -::: exponax.stepper.GeneralNonlinearStepper +::: exponax.stepper.generic.GeneralNonlinearStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/general/general_polynomial.md b/docs/api/stepper/generic/physical/general_polynomial.md similarity index 66% rename from docs/api/stepper/physical/general/general_polynomial.md rename to docs/api/stepper/generic/physical/general_polynomial.md index 9de794c..7716d63 100644 --- a/docs/api/stepper/physical/general/general_polynomial.md +++ b/docs/api/stepper/generic/physical/general_polynomial.md @@ -1,6 +1,6 @@ # General Polynomial Stepper -::: exponax.stepper.GeneralPolynomialStepper +::: exponax.stepper.generic.GeneralPolynomialStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/general/general_vorticity_convection.md b/docs/api/stepper/generic/physical/general_vorticity_convection.md similarity index 64% rename from docs/api/stepper/physical/general/general_vorticity_convection.md rename to docs/api/stepper/generic/physical/general_vorticity_convection.md index 637a524..357847b 100644 --- a/docs/api/stepper/physical/general/general_vorticity_convection.md +++ b/docs/api/stepper/generic/physical/general_vorticity_convection.md @@ -1,6 +1,6 @@ # General Vorticity Convection Stepper -::: exponax.stepper.GeneralVorticityConvectionStepper +::: exponax.stepper.generic.GeneralVorticityConvectionStepper options: members: - __init__ diff --git a/docs/api/stepper/physical/linear/advection.md b/docs/api/stepper/linear/advection.md similarity index 100% rename from docs/api/stepper/physical/linear/advection.md rename to docs/api/stepper/linear/advection.md diff --git a/docs/api/stepper/physical/linear/advection_diffusion.md b/docs/api/stepper/linear/advection_diffusion.md similarity index 100% rename from docs/api/stepper/physical/linear/advection_diffusion.md rename to docs/api/stepper/linear/advection_diffusion.md diff --git a/docs/api/stepper/physical/linear/diffusion.md b/docs/api/stepper/linear/diffusion.md similarity index 100% rename from docs/api/stepper/physical/linear/diffusion.md rename to docs/api/stepper/linear/diffusion.md diff --git a/docs/api/stepper/physical/linear/dispersion.md b/docs/api/stepper/linear/dispersion.md similarity index 100% rename from docs/api/stepper/physical/linear/dispersion.md rename to docs/api/stepper/linear/dispersion.md diff --git a/docs/api/stepper/physical/linear/hyper_diffusion.md b/docs/api/stepper/linear/hyper_diffusion.md similarity index 100% rename from docs/api/stepper/physical/linear/hyper_diffusion.md rename to docs/api/stepper/linear/hyper_diffusion.md diff --git a/docs/api/stepper/physical/burgers.md b/docs/api/stepper/nonlinear/burgers.md similarity index 100% rename from docs/api/stepper/physical/burgers.md rename to docs/api/stepper/nonlinear/burgers.md diff --git a/docs/api/stepper/physical/kdv.md b/docs/api/stepper/nonlinear/kdv.md similarity index 100% rename from docs/api/stepper/physical/kdv.md rename to docs/api/stepper/nonlinear/kdv.md diff --git a/docs/api/stepper/physical/ks.md b/docs/api/stepper/nonlinear/ks.md similarity index 100% rename from docs/api/stepper/physical/ks.md rename to docs/api/stepper/nonlinear/ks.md diff --git a/docs/api/stepper/physical/ks_cons.md b/docs/api/stepper/nonlinear/ks_cons.md similarity index 100% rename from docs/api/stepper/physical/ks_cons.md rename to docs/api/stepper/nonlinear/ks_cons.md diff --git a/docs/api/stepper/physical/navier_stokes.md b/docs/api/stepper/nonlinear/navier_stokes.md similarity index 100% rename from docs/api/stepper/physical/navier_stokes.md rename to docs/api/stepper/nonlinear/navier_stokes.md diff --git a/docs/api/stepper/normalized/vorticity_convection.md b/docs/api/stepper/normalized/vorticity_convection.md deleted file mode 100644 index fbd9790..0000000 --- a/docs/api/stepper/normalized/vorticity_convection.md +++ /dev/null @@ -1,7 +0,0 @@ -# Vorticity Convection - -::: exponax.normalized.NormalizedVorticityConvection - options: - members: - - __init__ - - __call__ \ No newline at end of file diff --git a/docs/api/stepper/overview.md b/docs/api/stepper/overview.md new file mode 100644 index 0000000..2b894dc --- /dev/null +++ b/docs/api/stepper/overview.md @@ -0,0 +1,23 @@ +# Overview + +`Exponax` comes with many pre-built dynamics. They are all based on the idea of +solving semi-linear PDEs + +$$ +\partial_t u = \mathcal{L} u + \mathcal{N}(u) +$$ + +where $\mathcal{L}$ is a linear operator and $\mathcal{N}$ is a non-linear +differential operator. + +TODO: add classification of dynamics based on: + +- linear or nonlinear or reaction-diffusion +- specific (i.e. concrete equations like advection, Burgers etc.) or generic + based on providing a collection of linear and nonlinear coefficients +- if generic: whether it uses a physical, normalized or difficulty-based + interface +- in higher dimensions: whether it behaves isotropic or anisotropic (with + different coefficients per direction) + +Then the resulting behavior can also be classified into ... \ No newline at end of file diff --git a/docs/api/stepper/physical/reaction_diffusion/allen_cahn.md b/docs/api/stepper/reaction/allen_cahn.md similarity index 69% rename from docs/api/stepper/physical/reaction_diffusion/allen_cahn.md rename to docs/api/stepper/reaction/allen_cahn.md index e74b2d2..ade5156 100644 --- a/docs/api/stepper/physical/reaction_diffusion/allen_cahn.md +++ b/docs/api/stepper/reaction/allen_cahn.md @@ -1,6 +1,6 @@ # Allen-Cahn -::: exponax.reaction.AllenCahn +::: exponax.stepper.reaction.AllenCahn options: members: - __init__ diff --git a/docs/api/stepper/physical/reaction_diffusion/cahn_hilliard.md b/docs/api/stepper/reaction/cahn_hilliard.md similarity index 68% rename from docs/api/stepper/physical/reaction_diffusion/cahn_hilliard.md rename to docs/api/stepper/reaction/cahn_hilliard.md index b0d441f..728efab 100644 --- a/docs/api/stepper/physical/reaction_diffusion/cahn_hilliard.md +++ b/docs/api/stepper/reaction/cahn_hilliard.md @@ -1,6 +1,6 @@ # Cahn-Hilliad -::: exponax.reaction.CahnHilliard +::: exponax.stepper.reaction.CahnHilliard options: members: - __init__ diff --git a/docs/api/stepper/physical/reaction_diffusion/fisher_kpp.md b/docs/api/stepper/reaction/fisher_kpp.md similarity index 69% rename from docs/api/stepper/physical/reaction_diffusion/fisher_kpp.md rename to docs/api/stepper/reaction/fisher_kpp.md index d4b7a66..0da53f7 100644 --- a/docs/api/stepper/physical/reaction_diffusion/fisher_kpp.md +++ b/docs/api/stepper/reaction/fisher_kpp.md @@ -1,6 +1,6 @@ # Fisher-KPP -::: exponax.reaction.FisherKPP +::: exponax.stepper.reaction.FisherKPP options: members: - __init__ diff --git a/docs/api/stepper/physical/reaction_diffusion/gray_scott.md b/docs/api/stepper/reaction/gray_scott.md similarity index 69% rename from docs/api/stepper/physical/reaction_diffusion/gray_scott.md rename to docs/api/stepper/reaction/gray_scott.md index c3dce33..30357b9 100644 --- a/docs/api/stepper/physical/reaction_diffusion/gray_scott.md +++ b/docs/api/stepper/reaction/gray_scott.md @@ -1,6 +1,6 @@ # Gray-Scott -::: exponax.reaction.GrayScott +::: exponax.stepper.reaction.GrayScott options: members: - __init__ diff --git a/docs/api/stepper/physical/reaction_diffusion/swift_hohenberg.md b/docs/api/stepper/reaction/swift_hohenberg.md similarity index 68% rename from docs/api/stepper/physical/reaction_diffusion/swift_hohenberg.md rename to docs/api/stepper/reaction/swift_hohenberg.md index b38ab8c..89594c0 100644 --- a/docs/api/stepper/physical/reaction_diffusion/swift_hohenberg.md +++ b/docs/api/stepper/reaction/swift_hohenberg.md @@ -1,6 +1,6 @@ # Swift-Hohenberg -::: exponax.reaction.SwiftHohenberg +::: exponax.stepper.reaction.SwiftHohenberg options: members: - __init__ diff --git a/docs/api/utilities/normalized_and_difficulty.md b/docs/api/utilities/normalized_and_difficulty.md index 64b9579..ba7b4d5 100644 --- a/docs/api/utilities/normalized_and_difficulty.md +++ b/docs/api/utilities/normalized_and_difficulty.md @@ -1,63 +1,63 @@ # Conversion Utilties for Normalized and Difficulty Steppers -:::exponax.normalized.normalize_coefficients +:::exponax.stepper.generic.normalize_coefficients --- -:::exponax.normalized.denormalize_coefficients +:::exponax.stepper.generic.denormalize_coefficients --- -:::exponax.normalized.normalize_convection_scale +:::exponax.stepper.generic.normalize_convection_scale --- -:::exponax.normalized.denormalize_convection_scale +:::exponax.stepper.generic.denormalize_convection_scale --- -:::exponax.normalized.normalize_gradient_norm_scale +:::exponax.stepper.generic.normalize_gradient_norm_scale --- -:::exponax.normalized.denormalize_gradient_norm_scale +:::exponax.stepper.generic.denormalize_gradient_norm_scale --- -:::exponax.normalized.normalize_polynomial_scales +:::exponax.stepper.generic.normalize_polynomial_scales --- -:::exponax.normalized.denormalize_polynomial_scales +:::exponax.stepper.generic.denormalize_polynomial_scales --- -:::exponax.normalized.reduce_normalized_coefficients_to_difficulty +:::exponax.stepper.generic.reduce_normalized_coefficients_to_difficulty --- -:::exponax.normalized.extract_normalized_coefficients_from_difficulty +:::exponax.stepper.generic.extract_normalized_coefficients_from_difficulty --- -:::exponax.normalized.reduce_normalized_convection_scale_to_difficulty +:::exponax.stepper.generic.reduce_normalized_convection_scale_to_difficulty --- -:::exponax.normalized.extract_normalized_convection_scale_from_difficulty +:::exponax.stepper.generic.extract_normalized_convection_scale_from_difficulty --- -:::exponax.normalized.reduce_normalized_gradient_norm_scale_to_difficulty +:::exponax.stepper.generic.reduce_normalized_gradient_norm_scale_to_difficulty --- -:::exponax.normalized.extract_normalized_gradient_norm_scale_from_difficulty +:::exponax.stepper.generic.extract_normalized_gradient_norm_scale_from_difficulty \ No newline at end of file +:::exponax.stepper.generic.extract_normalized_nonlinear_scales_from_difficulty --> \ No newline at end of file diff --git a/exponax/__init__.py b/exponax/__init__.py index 804849f..2366a33 100644 --- a/exponax/__init__.py +++ b/exponax/__init__.py @@ -1,7 +1,7 @@ from . import _metrics as metrics from . import _poisson as poisson from . import _spectral as spectral -from . import etdrk, ic, nonlin_fun, normalized, reaction, stepper, viz +from . import etdrk, ic, nonlin_fun, stepper, viz from ._base_stepper import BaseStepper from ._forced_stepper import ForcedStepper from ._repeated_stepper import RepeatedStepper @@ -36,8 +36,6 @@ "etdrk", "ic", "nonlin_fun", - "normalized", - "reaction", "stepper", "viz", "spectral", diff --git a/exponax/normalized/_convection.py b/exponax/normalized/_convection.py deleted file mode 100644 index ba0dd79..0000000 --- a/exponax/normalized/_convection.py +++ /dev/null @@ -1,118 +0,0 @@ -import jax.numpy as jnp -from jaxtyping import Array - -from .._base_stepper import BaseStepper -from ..nonlin_fun import ConvectionNonlinearFun -from ._utils import ( - extract_normalized_coefficients_from_difficulty, - extract_normalized_convection_scale_from_difficulty, -) - - -class NormalizedConvectionStepper(BaseStepper): - normalized_coefficients: tuple[float, ...] - normalized_convection_scale: float - dealiasing_fraction: float - - def __init__( - self, - num_spatial_dims: int, - num_points: int, - *, - normalized_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1), - normalized_convection_scale: float = 1.0 * 0.1, - order: int = 2, - dealiasing_fraction: float = 2 / 3, - num_circle_points: int = 16, - circle_radius: float = 1.0, - ): - """ - By default: Behaves like a Burgers with - - ``` Burgers( - D=D, L=1, N=N, dt=0.1, diffusivity=0.01, - ) - ``` - """ - self.normalized_coefficients = normalized_coefficients - self.normalized_convection_scale = normalized_convection_scale - self.dealiasing_fraction = dealiasing_fraction - 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, - num_channels=num_spatial_dims, - order=order, - num_circle_points=num_circle_points, - circle_radius=circle_radius, - ) - - def _build_linear_operator(self, derivative_operator: Array) -> Array: - # Now the linear operator is unscaled - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.normalized_coefficients) - ) - return linear_operator - - def _build_nonlinear_fun(self, derivative_operator: Array): - return ConvectionNonlinearFun( - self.num_spatial_dims, - self.num_points, - derivative_operator=derivative_operator, - dealiasing_fraction=self.dealiasing_fraction, - scale=self.normalized_convection_scale, - ) - - -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, - ) diff --git a/exponax/normalized/_gradient_norm.py b/exponax/normalized/_gradient_norm.py deleted file mode 100644 index 57cce23..0000000 --- a/exponax/normalized/_gradient_norm.py +++ /dev/null @@ -1,124 +0,0 @@ -import jax.numpy as jnp -from jaxtyping import Array - -from .._base_stepper import BaseStepper -from ..nonlin_fun import GradientNormNonlinearFun -from ._utils import ( - extract_normalized_coefficients_from_difficulty, - extract_normalized_gradient_norm_scale_from_difficulty, -) - - -class NormalizedGradientNormStepper(BaseStepper): - normalized_coefficients: tuple[float, ...] - normalized_gradient_norm_scale: float - dealiasing_fraction: float - - def __init__( - self, - num_spatial_dims: int, - num_points: int, - *, - normalized_coefficients: tuple[float, ...] = ( - 0.0, - 0.0, - -1.0 * 0.1 / (60.0**2), - 0.0, - -1.0 * 0.1 / (60.0**4), - ), - normalized_gradient_norm_scale: float = 1.0 * 0.1 / (60.0**2), - order: int = 2, - dealiasing_fraction: float = 2 / 3, - num_circle_points: int = 16, - circle_radius: float = 1.0, - ): - """ - 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 - - """ - self.normalized_coefficients = normalized_coefficients - self.normalized_gradient_norm_scale = normalized_gradient_norm_scale - self.dealiasing_fraction = dealiasing_fraction - 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, - num_channels=1, - order=order, - num_circle_points=num_circle_points, - circle_radius=circle_radius, - ) - - def _build_linear_operator(self, derivative_operator: Array) -> Array: - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.normalized_coefficients) - ) - return linear_operator - - def _build_nonlinear_fun(self, derivative_operator: Array): - return GradientNormNonlinearFun( - num_spatial_dims=self.num_spatial_dims, - num_points=self.num_points, - derivative_operator=derivative_operator, - dealiasing_fraction=self.dealiasing_fraction, - scale=self.normalized_gradient_norm_scale, - zero_mode_fix=True, - ) - - -class DifficultyGradientNormStepper(NormalizedGradientNormStepper): - linear_difficulties: tuple[float, ...] - gradient_norm_difficulty: float - - def __init__( - self, - num_spatial_dims: int = 1, - num_points: int = 48, - *, - linear_difficulties: tuple[float, ...] = (0.0, 0.0, -0.128, 0.0, -0.32768), - gradient_norm_difficulty: float = 0.064, - 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: KS equation - """ - self.linear_difficulties = linear_difficulties - self.gradient_norm_difficulty = gradient_norm_difficulty - - normalized_coefficients = extract_normalized_coefficients_from_difficulty( - linear_difficulties, - num_spatial_dims=num_spatial_dims, - num_points=num_points, - ) - normalized_gradient_norm_scale = ( - extract_normalized_gradient_norm_scale_from_difficulty( - gradient_norm_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_gradient_norm_scale=normalized_gradient_norm_scale, - order=order, - dealiasing_fraction=dealiasing_fraction, - num_circle_points=num_circle_points, - circle_radius=circle_radius, - ) diff --git a/exponax/normalized/_linear.py b/exponax/normalized/_linear.py deleted file mode 100644 index 043b618..0000000 --- a/exponax/normalized/_linear.py +++ /dev/null @@ -1,177 +0,0 @@ -import jax.numpy as jnp -from jaxtyping import Array - -from .._base_stepper import BaseStepper -from ..nonlin_fun import ZeroNonlinearFun -from ._utils import extract_normalized_coefficients_from_difficulty - - -class NormalizedLinearStepper(BaseStepper): - normalized_coefficients: tuple[float, ...] - - def __init__( - self, - num_spatial_dims: int, - num_points: int, - *, - normalized_coefficients: tuple[float, ...] = (0.0, -0.5, 0.01), - ): - """ - Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic - boundary conditions with normalized dynamics. - - If the PDE in physical description is - - uₜ = ∑ᵢ aᵢ (∂ₓ)ⁱ u - - with `aᵢ` the coefficients on `domain_extent` `L` with time step `dt`, - the normalized coefficients are - - αᵢ = (aᵢ Δt)/(Lⁱ) - - Important: note that the `domain_extent` is raised to the order of - linear derivative `i`. - - One can also think of normalized dynamics, as a PDE on `domain_extent` - `1.0` with time step `dt=1.0`. - - Take care of the signs! - - In the defaulf configuration of this timestepper, the PDE is an - advection-diffusion equation with normalized advection of 0.5 and - normalized diffusion of 0.01. - - **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** - the right boundary point. In higher dimensions; the number of points - in each dimension is the same. Hence, the total number of degrees of - freedom is `Nᵈ`. - - `normalized_coefficients`: The coefficients of the normalized - dynamics. This must a tuple of floats. The length of the tuple - defines the highest occuring linear derivative in the PDE. - """ - 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, - num_channels=1, - order=0, - ) - - def _build_linear_operator(self, derivative_operator: Array) -> Array: - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.normalized_coefficients) - ) - return linear_operator - - def _build_nonlinear_fun(self, derivative_operator: Array): - return ZeroNonlinearFun( - num_spatial_dims=self.num_spatial_dims, - num_points=self.num_points, - ) - - -class DifficultyLinearStepper(NormalizedLinearStepper): - difficulties: tuple[float, ...] - - def __init__( - self, - num_spatial_dims: int = 1, - num_points: int = 48, - *, - difficulties: tuple[float, ...] = (0.0, -2.0), - ): - """ - Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic - boundary conditions with normalized dynamics in a difficulty-based - interface. - - Different to `NormalizedLinearStepper`, the dynamics are defined by - difficulties. The difficulties are a different combination of normalized - dynamics, `num_spatial_dims`, and `num_points`. - - γᵢ = αᵢ Nⁱ 2ⁱ⁻¹ d - - with `d` the number of spatial dimensions, `N` the number of points, and - `αᵢ` the normalized coefficient. - - This interface is more natural because the difficulties for all orders - (given by `i`) are around 1.0. Additionally, they relate to stability - condition of explicit Finite Difference schemes for the particular - equations. For example, for advection (`i=1`), the absolute of the - difficulty is the Courant-Friedrichs-Lewy (CFL) number. - - In the default configuration of this timestepper, the PDE is an - advection equation with CFL number 2 solved in 1d with 48 resolution - points to discretize the domain. - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. Default is - 1. - - `num_points`: The number of points `N` used to discretize the domain. - This **includes** the left boundary point and **excludes** the right - boundary point. In higher dimensions; the number of points in each - dimension is the same. Hence, the total number of degrees of freedom - is `Nᵈ`. Default is 48. - - `difficulties`: The difficulties of the normalized dynamics. This must - be a tuple of floats. The length of the tuple defines the highest - occuring linear derivative in the PDE. Default is `(0.0, -2.0)`. - """ - self.difficulties = difficulties - normalized_coefficients = extract_normalized_coefficients_from_difficulty( - difficulties, - num_spatial_dims=num_spatial_dims, - num_points=num_points, - ) - - super().__init__( - num_spatial_dims=num_spatial_dims, - num_points=num_points, - normalized_coefficients=normalized_coefficients, - ) - - -class DiffultyLinearStepperSimple(DifficultyLinearStepper): - def __init__( - self, - num_spatial_dims: int = 1, - num_points: int = 48, - *, - difficulty: float = -2.0, - order: int = 1, - ): - """ - A simple interface for `DifficultyLinearStepper` with only one - difficulty and a given order. - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. Default is - 1. - - `num_points`: The number of points `N` used to discretize the domain. - This **includes** the left boundary point and **excludes** the right - boundary point. In higher dimensions; the number of points in each - dimension is the same. Hence, the total number of degrees of freedom - is `Nᵈ`. Default is 48. - - `difficulty`: The difficulty of the normalized dynamics. This must be - a float. Default is -2.0. - - `order`: The order of the derivative associated with the provided - difficulty. The default of 1 is the advection equation. - """ - difficulties = (0.0,) * (order) + (difficulty,) - super().__init__( - difficulties=difficulties, - num_spatial_dims=num_spatial_dims, - num_points=num_points, - ) diff --git a/exponax/normalized/_vorticity_convection.py b/exponax/normalized/_vorticity_convection.py deleted file mode 100644 index 43a99cb..0000000 --- a/exponax/normalized/_vorticity_convection.py +++ /dev/null @@ -1,88 +0,0 @@ -import jax.numpy as jnp -from jaxtyping import Array, Complex - -from .._base_stepper import BaseStepper -from ..nonlin_fun import VorticityConvection2d, VorticityConvection2dKolmogorov - - -class NormalizedVorticityConvection(BaseStepper): - normalized_vorticity_convection_scale: float - normalized_coefficients: tuple[float, ...] - injection_mode: int - normalized_injection_scale: float - dealiasing_fraction: float - - def __init__( - self, - num_spatial_dims: int, - num_points: int, - *, - normalized_vorticity_convection_scale: float = 0.01 * 1.0 / (1.0**0), - normalized_coefficients: tuple[float, ...] = ( - 0.01 * 0.0 / (1.0**0), - 0.0, - 0.01 * 0.001 / (1.0**2), - ), - injection_mode: int = 4, - normalized_injection_scale: float = 0.01 * 0.0 / (1.0**0), - order: int = 2, - dealiasing_fraction: float = 2 / 3, - num_circle_points: int = 16, - circle_radius: float = 1.0, - ): - if num_spatial_dims != 2: - raise ValueError(f"Expected num_spatial_dims = 2, got {num_spatial_dims}.") - self.normalized_vorticity_convection_scale = ( - normalized_vorticity_convection_scale - ) - self.normalized_coefficients = normalized_coefficients - self.injection_mode = injection_mode - self.normalized_injection_scale = normalized_injection_scale - self.dealiasing_fraction = dealiasing_fraction - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=1.0, - num_points=num_points, - dt=1.0, - num_channels=1, - order=order, - num_circle_points=num_circle_points, - circle_radius=circle_radius, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.normalized_coefficients) - ) - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> VorticityConvection2d: - if self.normalized_injection_scale == 0.0: - return VorticityConvection2d( - self.num_spatial_dims, - self.num_points, - convection_scale=self.normalized_vorticity_convection_scale, - derivative_operator=derivative_operator, - dealiasing_fraction=self.dealiasing_fraction, - ) - else: - return VorticityConvection2dKolmogorov( - self.num_spatial_dims, - self.num_points, - convection_scale=self.normalized_vorticity_convection_scale, - derivative_operator=derivative_operator, - dealiasing_fraction=self.dealiasing_fraction, - injection_mode=self.injection_mode, - injection_scale=self.normalized_injection_scale, - ) diff --git a/exponax/stepper/__init__.py b/exponax/stepper/__init__.py index 5ae2c3e..1cf4f7f 100644 --- a/exponax/stepper/__init__.py +++ b/exponax/stepper/__init__.py @@ -65,23 +65,17 @@ modes for the linear terms. """ +from . import generic as generic +from . import reaction as reaction +from ._advection import Advection +from ._advection_diffusion import AdvectionDiffusion from ._burgers import Burgers -from ._convection import GeneralConvectionStepper -from ._general_nonlinear import GeneralNonlinearStepper -from ._gradient_norm import GeneralGradientNormStepper +from ._diffusion import Diffusion +from ._dispersion import Dispersion +from ._hyper_diffusion import HyperDiffusion from ._korteweg_de_vries import KortewegDeVries from ._kuramoto_sivashinsky import KuramotoSivashinsky, KuramotoSivashinskyConservative -from ._linear import ( - Advection, - AdvectionDiffusion, - Diffusion, - Dispersion, - GeneralLinearStepper, - HyperDiffusion, -) from ._navier_stokes import KolmogorovFlowVorticity, NavierStokesVorticity -from ._polynomial import GeneralPolynomialStepper -from ._vorticity_convection import GeneralVorticityConvectionStepper __all__ = [ "Advection", @@ -93,12 +87,8 @@ "KortewegDeVries", "KuramotoSivashinsky", "KuramotoSivashinskyConservative", - "GeneralPolynomialStepper", - "GeneralNonlinearStepper", - "GeneralLinearStepper", - "GeneralConvectionStepper", - "GeneralGradientNormStepper", - "GeneralVorticityConvectionStepper", "NavierStokesVorticity", "KolmogorovFlowVorticity", + "reaction", + "generic", ] diff --git a/exponax/stepper/_advection.py b/exponax/stepper/_advection.py new file mode 100644 index 0000000..34fc3c8 --- /dev/null +++ b/exponax/stepper/_advection.py @@ -0,0 +1,104 @@ +from typing import TypeVar, Union + +import jax.numpy as jnp +from jaxtyping import Array, Complex, Float + +from .._base_stepper import BaseStepper +from .._spectral import build_gradient_inner_product_operator +from ..nonlin_fun import ZeroNonlinearFun + +D = TypeVar("D") + + +class Advection(BaseStepper): + velocity: Float[Array, "D"] + + def __init__( + self, + num_spatial_dims: int, + domain_extent: float, + num_points: int, + dt: float, + *, + velocity: Union[Float[Array, "D"], float] = 1.0, + ): + """ + Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) advection equation + on periodic boundary conditions. + + In 1d, the advection equation is given by + + ``` + uₜ + c uₓ = 0 + ``` + + with `c ∈ ℝ` being the velocity/advection speed. + + In higher dimensions, the advection equation can written as the inner + product between velocity vector and gradient + + ``` + uₜ + c ⋅ ∇u = 0 + ``` + + with `c ∈ ℝᵈ` being the velocity/advection vector. + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. + - `domain_extent`: The size of the domain `L`; in higher dimensions + the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. + - `num_points`: The number of points `N` used to discretize the + domain. This **includes** the left boundary point and **excludes** + the right boundary point. In higher dimensions; the number of points + in each dimension is the same. Hence, the total number of degrees of + freedom is `Nᵈ`. + - `dt`: The timestep size `Δt` between two consecutive states. + - `velocity` (keyword-only): The advection speed `c`. In higher + dimensions, this can be a scalar (=float) or a vector of length `d`. + If a scalar is given, the advection speed is assumed to be the same + in all spatial dimensions. Default: `1.0`. + + **Notes:** + + - The stepper is unconditionally stable, not matter the choice of + any argument because the equation is solved analytically in Fourier + space. **However**, note that initial conditions with modes higher + than the Nyquist freuency (`(N//2)+1` with `N` being the + `num_points`) lead to spurious oscillations. + - Ultimately, only the factor `c Δt / L` affects the characteristic + of the dynamics. See also + [`exponax.stepper.generic.NormalizedLinearStepper`][] with + `normalized_coefficients = [0, alpha_1]` with `alpha_1 = - velocity + * dt / domain_extent`. + """ + # TODO: better checks on the desired type of velocity + if isinstance(velocity, float): + velocity = jnp.ones(num_spatial_dims) * velocity + self.velocity = velocity + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=domain_extent, + num_points=num_points, + dt=dt, + num_channels=1, + order=0, + ) + + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: + # Requires minus to move term to the rhs + return -build_gradient_inner_product_operator( + derivative_operator, self.velocity, order=1 + ) + + def _build_nonlinear_fun( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> ZeroNonlinearFun: + return ZeroNonlinearFun( + self.num_spatial_dims, + self.num_points, + ) diff --git a/exponax/stepper/_advection_diffusion.py b/exponax/stepper/_advection_diffusion.py new file mode 100644 index 0000000..4a00b75 --- /dev/null +++ b/exponax/stepper/_advection_diffusion.py @@ -0,0 +1,141 @@ +from typing import TypeVar, Union + +import jax.numpy as jnp +from jaxtyping import Array, Complex, Float + +from .._base_stepper import BaseStepper +from .._spectral import build_gradient_inner_product_operator +from ..nonlin_fun import ZeroNonlinearFun + +D = TypeVar("D") + + +class AdvectionDiffusion(BaseStepper): + velocity: Float[Array, "D"] + diffusivity: Float[Array, "D D"] + + def __init__( + self, + num_spatial_dims: int, + domain_extent: float, + num_points: int, + dt: float, + *, + velocity: Union[Float[Array, "D"], float] = 1.0, + diffusivity: Union[ + Float[Array, "D D"], + Float[Array, "D"], + float, + ] = 0.01, + ): + """ + Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) advection-diffusion + equation on periodic boundary conditions. + + In 1d, the advection-diffusion equation is given by + + ``` + uₜ + c uₓ = ν uₓₓ + ``` + + with `c ∈ ℝ` being the velocity/advection speed and `ν ∈ ℝ` being the + diffusivity. + + In higher dimensions, the advection-diffusion equation can be written as + + ``` + uₜ + c ⋅ ∇u = ν Δu + ``` + + with `c ∈ ℝᵈ` being the velocity/advection vector. + + See also [`exponax.stepper.Diffusion`][] for additional details on + anisotropic diffusion. + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. + - `domain_extent`: The size of the domain `L`; in higher dimensions + the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. + - `num_points`: The number of points `N` used to discretize the + domain. This **includes** the left boundary point and **excludes** + the right boundary point. In higher dimensions; the number of points + in each dimension is the same. Hence, the total number of degrees of + freedom is `Nᵈ`. + - `dt`: The timestep size `Δt` between two consecutive states. + - `velocity` (keyword-only): The advection speed `c`. In higher + dimensions, this can be a scalar (=float) or a vector of length `d`. + If a scalar is given, the advection speed is assumed to be the same + in all spatial dimensions. Default: `1.0`. + - `diffusivity` (keyword-only): The diffusivity `ν`. In higher + dimensions, this can be a scalar (=float), a vector of length `d`, + or a matrix of shape `d ˣ d`. If a scalar is given, the diffusivity + is assumed to be the same in all spatial dimensions. If a vector (of + length `d`) is given, the diffusivity varies across dimensions (=> + diagonal diffusion). For a matrix, there is fully anisotropic + diffusion. In this case, `A` must be symmetric positive definite + (SPD). Default: `0.01`. + + **Notes:** + + - The stepper is unconditionally stable, not matter the choice of + any argument because the equation is solved analytically in Fourier + space. **However**, note that initial conditions with modes higher + than the Nyquist freuency (`(N//2)+1` with `N` being the + `num_points`) lead to spurious oscillations. + - Ultimately, only the factors `c Δt / L` and `ν Δt / L²` affect the + characteristic of the dynamics. See also + [`exponax.stepper.generic.NormalizedLinearStepper`][] with + `normalized_coefficients = [0, alpha_1, alpha_2]` with `alpha_1 = - + velocity * dt / domain_extent` and `alpha_2 = diffusivity * dt / + domain_extent**2`. + """ + # TODO: more sophisticated checks here + if isinstance(velocity, float): + velocity = jnp.ones(num_spatial_dims) * velocity + self.velocity = velocity + if isinstance(diffusivity, float): + diffusivity = jnp.diag(jnp.ones(num_spatial_dims)) * diffusivity + elif len(diffusivity.shape) == 1: + diffusivity = jnp.diag(diffusivity) + self.diffusivity = diffusivity + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=domain_extent, + num_points=num_points, + dt=dt, + num_channels=1, + order=0, + ) + + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: + laplace_outer_producct = ( + derivative_operator[:, None] * derivative_operator[None, :] + ) + diffusion_operator = jnp.einsum( + "ij,ij...->...", + self.diffusivity, + laplace_outer_producct, + ) + # Add the necessary singleton channel axis + diffusion_operator = diffusion_operator[None, ...] + + advection_operator = -build_gradient_inner_product_operator( + derivative_operator, self.velocity, order=1 + ) + + linear_operator = advection_operator + diffusion_operator + + return linear_operator + + def _build_nonlinear_fun( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> ZeroNonlinearFun: + return ZeroNonlinearFun( + self.num_spatial_dims, + self.num_points, + ) diff --git a/exponax/stepper/_diffusion.py b/exponax/stepper/_diffusion.py new file mode 100644 index 0000000..ccc1f85 --- /dev/null +++ b/exponax/stepper/_diffusion.py @@ -0,0 +1,126 @@ +from typing import TypeVar, Union + +import jax.numpy as jnp +from jaxtyping import Array, Complex, Float + +from .._base_stepper import BaseStepper +from ..nonlin_fun import ZeroNonlinearFun + +D = TypeVar("D") + + +class Diffusion(BaseStepper): + diffusivity: Float[Array, "D D"] + + def __init__( + self, + num_spatial_dims: int, + domain_extent: float, + num_points: int, + dt: float, + *, + diffusivity: Union[ + Float[Array, "D D"], + Float[Array, "D"], + float, + ] = 0.01, + ): + """ + Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) diffusion equation + on periodic boundary conditions. + + In 1d, the diffusion equation is given by + + ``` + uₜ = ν uₓₓ + ``` + + with `ν ∈ ℝ` being the diffusivity. + + In higher dimensions, the diffusion equation can written using the + Laplacian operator. + + ``` + uₜ = ν Δu + ``` + + More generally speaking, there can be anistropic diffusivity given by a + `A ∈ ℝᵈ ˣ ᵈ` sandwiched between the gradient and divergence operators. + + ``` + uₜ = ∇ ⋅ (A ∇u) + ``` + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. + - `domain_extent`: The size of the domain `L`; in higher dimensions + the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. + - `num_points`: The number of points `N` used to discretize the + domain. This **includes** the left boundary point and **excludes** + the right boundary point. In higher dimensions; the number of points + in each dimension is the same. Hence, the total number of degrees of + freedom is `Nᵈ`. + - `dt`: The timestep size `Δt` between two consecutive states. + - `diffusivity` (keyword-only): The diffusivity `ν`. In higher + dimensions, this can be a scalar (=float), a vector of length `d`, + or a matrix of shape `d ˣ d`. If a scalar is given, the diffusivity + is assumed to be the same in all spatial dimensions. If a vector (of + length `d`) is given, the diffusivity varies across dimensions (=> + diagonal diffusion). For a matrix, there is fully anisotropic + diffusion. In this case, `A` must be symmetric positive definite + (SPD). Default: `0.01`. + + **Notes:** + + - The stepper is unconditionally stable, not matter the choice of + any argument because the equation is solved analytically in Fourier + space. + - A `ν > 0` leads to stable and decaying solutions (i.e., energy is + removed from the system). A `ν < 0` leads to unstable and growing + solutions (i.e., energy is added to the system). + - Ultimately, only the factor `ν Δt / L²` affects the characteristic + of the dynamics. See also + [`exponax.stepper.generic.NormalizedLinearStepper`][] with + `normalized_coefficients = [0, 0, alpha_2]` with `alpha_2 = + diffusivity * dt / domain_extent**2`. + """ + # ToDo: more sophisticated checks here + if isinstance(diffusivity, float): + diffusivity = jnp.diag(jnp.ones(num_spatial_dims)) * diffusivity + elif len(diffusivity.shape) == 1: + diffusivity = jnp.diag(diffusivity) + self.diffusivity = diffusivity + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=domain_extent, + num_points=num_points, + dt=dt, + num_channels=1, + order=0, + ) + + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: + laplace_outer_producct = ( + derivative_operator[:, None] * derivative_operator[None, :] + ) + linear_operator = jnp.einsum( + "ij,ij...->...", + self.diffusivity, + laplace_outer_producct, + ) + # Add the necessary singleton channel axis + linear_operator = linear_operator[None, ...] + return linear_operator + + def _build_nonlinear_fun( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> ZeroNonlinearFun: + return ZeroNonlinearFun( + self.num_spatial_dims, + self.num_points, + ) diff --git a/exponax/stepper/_dispersion.py b/exponax/stepper/_dispersion.py new file mode 100644 index 0000000..b7d8347 --- /dev/null +++ b/exponax/stepper/_dispersion.py @@ -0,0 +1,126 @@ +from typing import TypeVar, Union + +import jax.numpy as jnp +from jaxtyping import Array, Complex, Float + +from .._base_stepper import BaseStepper +from .._spectral import build_gradient_inner_product_operator, build_laplace_operator +from ..nonlin_fun import ZeroNonlinearFun + +D = TypeVar("D") + + +class Dispersion(BaseStepper): + dispersivity: Float[Array, "D"] + advect_on_diffusion: bool + + def __init__( + self, + num_spatial_dims: int, + domain_extent: float, + num_points: int, + dt: float, + *, + dispersivity: Union[Float[Array, "D"], float] = 1.0, + advect_on_diffusion: bool = False, + ): + """ + Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) dispersion equation + on periodic boundary conditions. Essentially, a dispersion equation is + an advection equation with different velocities (=advection speeds) for + different wavenumbers/modes. Higher wavenumbers/modes are advected + faster. + + In 1d, the dispersion equation is given by + + ``` + uₜ = 𝒸 uₓₓₓ + ``` + + with `𝒸 ∈ ℝ` being the dispersivity. + + In higher dimensions, the dispersion equation can be written as + + ``` + uₜ = 𝒸 ⋅ (∇⊙∇⊙(∇u)) + ``` + + or + + ``` + uₜ = 𝒸 ⋅ ∇(Δu) + ``` + + with `𝒸 ∈ ℝᵈ` being the dispersivity vector + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. + - `domain_extent`: The size of the domain `L`; in higher dimensions + the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. + - `num_points`: The number of points `N` used to discretize the + domain. This **includes** the left boundary point and **excludes** + the right boundary point. In higher dimensions; the number of points + in each dimension is the same. Hence, the total number of degrees of + freedom is `Nᵈ`. + - `dt`: The timestep size `Δt` between two consecutive states. + - `dispersivity` (keyword-only): The dispersivity `𝒸`. In higher + dimensions, this can be a scalar (=float) or a vector of length `d`. + If a scalar is given, the dispersivity is assumed to be the same in + all spatial dimensions. Default: `1.0`. + - `advect_on_diffusion` (keyword-only): If `True`, the second form + of the dispersion equation in higher dimensions is used. As a + consequence, there will be mixing in the spatial derivatives. + Default: `False`. + + **Notes:** + + - The stepper is unconditionally stable, not matter the choice of + any argument because the equation is solved analytically in Fourier + space. **However**, note that initial conditions with modes higher + than the Nyquist freuency (`(N//2)+1` with `N` being the + `num_points`) lead to spurious oscillations. + - Ultimately, only the factor `𝒸 Δt / L³` affects the + characteristic of the dynamics. See also + [`exponax.stepper.generic.NormalizedLinearStepper`][] with + `normalized_coefficients = [0, 0, 0, alpha_3]` with `alpha_3 = + dispersivity * dt / domain_extent**3`. + """ + if isinstance(dispersivity, float): + dispersivity = jnp.ones(num_spatial_dims) * dispersivity + self.dispersivity = dispersivity + self.advect_on_diffusion = advect_on_diffusion + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=domain_extent, + num_points=num_points, + dt=dt, + num_channels=1, + order=0, + ) + + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: + if self.advect_on_diffusion: + laplace_operator = build_laplace_operator(derivative_operator) + advection_operator = build_gradient_inner_product_operator( + derivative_operator, self.dispersivity, order=1 + ) + linear_operator = advection_operator * laplace_operator + else: + linear_operator = build_gradient_inner_product_operator( + derivative_operator, self.dispersivity, order=3 + ) + + return linear_operator + + def _build_nonlinear_fun( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> ZeroNonlinearFun: + return ZeroNonlinearFun( + self.num_spatial_dims, + self.num_points, + ) diff --git a/exponax/stepper/_general_nonlinear.py b/exponax/stepper/_general_nonlinear.py deleted file mode 100644 index e128175..0000000 --- a/exponax/stepper/_general_nonlinear.py +++ /dev/null @@ -1,74 +0,0 @@ -import jax.numpy as jnp -from jaxtyping import Array, Complex - -from .._base_stepper import BaseStepper -from ..nonlin_fun import GeneralNonlinearFun - - -class GeneralNonlinearStepper(BaseStepper): - coefficients_linear: tuple[float, ...] - coefficients_nonlinear: tuple[float, 3] - dealiasing_fraction: float - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.01), - coefficients_nonlinear: tuple[float, 3] = (0.0, -1.0, 0.0), - order=2, - dealiasing_fraction: float = 2 / 3, - num_circle_points: int = 16, - circle_radius: float = 1.0, - ): - """ - By default Burgers equation - """ - if len(coefficients_nonlinear) != 3: - raise ValueError( - "The nonlinear coefficients list must have exactly 3 elements" - ) - self.coefficients_linear = coefficients_linear - self.coefficients_nonlinear = coefficients_nonlinear - self.dealiasing_fraction = dealiasing_fraction - - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=order, - num_circle_points=num_circle_points, - circle_radius=circle_radius, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.coefficients_linear) - ) - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> GeneralNonlinearFun: - return GeneralNonlinearFun( - self.num_spatial_dims, - self.num_points, - derivative_operator=derivative_operator, - dealiasing_fraction=self.dealiasing_fraction, - scale_list=self.coefficients_nonlinear, - zero_mode_fix=True, # ToDo: check this - ) diff --git a/exponax/stepper/_hyper_diffusion.py b/exponax/stepper/_hyper_diffusion.py new file mode 100644 index 0000000..73bac08 --- /dev/null +++ b/exponax/stepper/_hyper_diffusion.py @@ -0,0 +1,119 @@ +from jaxtyping import Array, Complex + +from .._base_stepper import BaseStepper +from .._spectral import build_laplace_operator +from ..nonlin_fun import ZeroNonlinearFun + + +class HyperDiffusion(BaseStepper): + hyper_diffusivity: float + diffuse_on_diffuse: bool + + def __init__( + self, + num_spatial_dims: int, + domain_extent: float, + num_points: int, + dt: float, + *, + hyper_diffusivity: float = 0.0001, + diffuse_on_diffuse: bool = False, + ): + """ + Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) hyper-diffusion + equation on periodic boundary conditions. A hyper-diffusion equation + acts like a diffusion equation but higher wavenumbers/modes are damped + even faster. + + In 1d, the hyper-diffusion equation is given by + + ``` + uₜ = - μ uₓₓₓₓ + ``` + + with `μ ∈ ℝ` being the hyper-diffusivity. + + Note the minus sign because by default, a fourth-order derivative + dampens with a negative coefficient. To match the concept of + second-order diffusion, a negation is introduced. + + In higher dimensions, the hyper-diffusion equation can be written as + + ``` + uₜ = − μ ((∇⊙∇) ⋅ (∇⊙∇)) u + ``` + + or + + ``` + uₜ = - μ Δ(Δu) + ``` + + The latter introduces spatial mixing. + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. + - `domain_extent`: The size of the domain `L`; in higher dimensions + the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. + - `num_points`: The number of points `N` used to discretize the + domain. This **includes** the left boundary point and **excludes** + the right boundary point. In higher dimensions; the number of points + in each dimension is the same. Hence, the total number of degrees of + freedom is `Nᵈ`. + - `dt`: The timestep size `Δt` between two consecutive states. + - `hyper_diffusivity` (keyword-only): The hyper-diffusivity `ν`. + This stepper only supports scalar (=isotropic) hyper-diffusivity. + Default: 0.0001. + - `diffuse_on_diffuse` (keyword-only): If `True`, the second form + of the hyper-diffusion equation in higher dimensions is used. As a + consequence, there will be mixing in the spatial derivatives. + Default: `False`. + + **Notes:** + + - The stepper is unconditionally stable, not matter the choice of + any argument because the equation is solved analytically in Fourier + space. + - Ultimately, only the factor `μ Δt / L⁴` affects the characteristic + of the dynamics. See also + [`exponax.stepper.generic.NormalizedLinearStepper`][] with + `normalized_coefficients = [0, 0, 0, 0, alpha_4]` with `alpha_4 = - + hyper_diffusivity * dt / domain_extent**4`. + """ + self.hyper_diffusivity = hyper_diffusivity + self.diffuse_on_diffuse = diffuse_on_diffuse + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=domain_extent, + num_points=num_points, + dt=dt, + num_channels=1, + order=0, + ) + + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: + # Use minus sign to have diffusion work in "correct direction" by default + if self.diffuse_on_diffuse: + laplace_operator = build_laplace_operator(derivative_operator) + linear_operator = ( + -self.hyper_diffusivity * laplace_operator * laplace_operator + ) + else: + linear_operator = -self.hyper_diffusivity * build_laplace_operator( + derivative_operator, order=4 + ) + + return linear_operator + + def _build_nonlinear_fun( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> ZeroNonlinearFun: + return ZeroNonlinearFun( + self.num_spatial_dims, + self.num_points, + ) diff --git a/exponax/stepper/_linear.py b/exponax/stepper/_linear.py deleted file mode 100644 index 05e6a79..0000000 --- a/exponax/stepper/_linear.py +++ /dev/null @@ -1,743 +0,0 @@ -from typing import TypeVar, Union - -import jax.numpy as jnp -from jaxtyping import Array, Complex, Float - -from .._base_stepper import BaseStepper -from .._spectral import build_gradient_inner_product_operator, build_laplace_operator -from ..nonlin_fun import ZeroNonlinearFun - -D = TypeVar("D") - - -class Advection(BaseStepper): - velocity: Float[Array, "D"] - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - velocity: Union[Float[Array, "D"], float] = 1.0, - ): - """ - Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) advection equation - on periodic boundary conditions. - - In 1d, the advection equation is given by - - ``` - uₜ + c uₓ = 0 - ``` - - with `c ∈ ℝ` being the velocity/advection speed. - - In higher dimensions, the advection equation can written as the inner - product between velocity vector and gradient - - ``` - uₜ + c ⋅ ∇u = 0 - ``` - - with `c ∈ ℝᵈ` being the velocity/advection vector. - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. - - `domain_extent`: The size of the domain `L`; in higher dimensions - the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. - - `num_points`: The number of points `N` used to discretize the - domain. This **includes** the left boundary point and **excludes** - the right boundary point. In higher dimensions; the number of points - in each dimension is the same. Hence, the total number of degrees of - freedom is `Nᵈ`. - - `dt`: The timestep size `Δt` between two consecutive states. - - `velocity` (keyword-only): The advection speed `c`. In higher - dimensions, this can be a scalar (=float) or a vector of length `d`. - If a scalar is given, the advection speed is assumed to be the same - in all spatial dimensions. Default: `1.0`. - - **Notes:** - - - The stepper is unconditionally stable, not matter the choice of - any argument because the equation is solved analytically in Fourier - space. **However**, note that initial conditions with modes higher - than the Nyquist freuency (`(N//2)+1` with `N` being the - `num_points`) lead to spurious oscillations. - - Ultimately, only the factor `c Δt / L` affects the characteristic - of the dynamics. See also - [`exponax.normalized.NormalizedLinearStepper`][] with - `normalized_coefficients = [0, alpha_1]` with `alpha_1 = - velocity - * dt / domain_extent`. - """ - # TODO: better checks on the desired type of velocity - if isinstance(velocity, float): - velocity = jnp.ones(num_spatial_dims) * velocity - self.velocity = velocity - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=0, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - # Requires minus to move term to the rhs - return -build_gradient_inner_product_operator( - derivative_operator, self.velocity, order=1 - ) - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> ZeroNonlinearFun: - return ZeroNonlinearFun( - self.num_spatial_dims, - self.num_points, - ) - - -class Diffusion(BaseStepper): - diffusivity: Float[Array, "D D"] - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - diffusivity: Union[ - Float[Array, "D D"], - Float[Array, "D"], - float, - ] = 0.01, - ): - """ - Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) diffusion equation - on periodic boundary conditions. - - In 1d, the diffusion equation is given by - - ``` - uₜ = ν uₓₓ - ``` - - with `ν ∈ ℝ` being the diffusivity. - - In higher dimensions, the diffusion equation can written using the - Laplacian operator. - - ``` - uₜ = ν Δu - ``` - - More generally speaking, there can be anistropic diffusivity given by a - `A ∈ ℝᵈ ˣ ᵈ` sandwiched between the gradient and divergence operators. - - ``` - uₜ = ∇ ⋅ (A ∇u) - ``` - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. - - `domain_extent`: The size of the domain `L`; in higher dimensions - the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. - - `num_points`: The number of points `N` used to discretize the - domain. This **includes** the left boundary point and **excludes** - the right boundary point. In higher dimensions; the number of points - in each dimension is the same. Hence, the total number of degrees of - freedom is `Nᵈ`. - - `dt`: The timestep size `Δt` between two consecutive states. - - `diffusivity` (keyword-only): The diffusivity `ν`. In higher - dimensions, this can be a scalar (=float), a vector of length `d`, - or a matrix of shape `d ˣ d`. If a scalar is given, the diffusivity - is assumed to be the same in all spatial dimensions. If a vector (of - length `d`) is given, the diffusivity varies across dimensions (=> - diagonal diffusion). For a matrix, there is fully anisotropic - diffusion. In this case, `A` must be symmetric positive definite - (SPD). Default: `0.01`. - - **Notes:** - - - The stepper is unconditionally stable, not matter the choice of - any argument because the equation is solved analytically in Fourier - space. - - A `ν > 0` leads to stable and decaying solutions (i.e., energy is - removed from the system). A `ν < 0` leads to unstable and growing - solutions (i.e., energy is added to the system). - - Ultimately, only the factor `ν Δt / L²` affects the characteristic - of the dynamics. See also - [`exponax.normalized.NormalizedLinearStepper`][] with - `normalized_coefficients = [0, 0, alpha_2]` with `alpha_2 = - diffusivity * dt / domain_extent**2`. - """ - # ToDo: more sophisticated checks here - if isinstance(diffusivity, float): - diffusivity = jnp.diag(jnp.ones(num_spatial_dims)) * diffusivity - elif len(diffusivity.shape) == 1: - diffusivity = jnp.diag(diffusivity) - self.diffusivity = diffusivity - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=0, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - laplace_outer_producct = ( - derivative_operator[:, None] * derivative_operator[None, :] - ) - linear_operator = jnp.einsum( - "ij,ij...->...", - self.diffusivity, - laplace_outer_producct, - ) - # Add the necessary singleton channel axis - linear_operator = linear_operator[None, ...] - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> ZeroNonlinearFun: - return ZeroNonlinearFun( - self.num_spatial_dims, - self.num_points, - ) - - -class AdvectionDiffusion(BaseStepper): - velocity: Float[Array, "D"] - diffusivity: Float[Array, "D D"] - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - velocity: Union[Float[Array, "D"], float] = 1.0, - diffusivity: Union[ - Float[Array, "D D"], - Float[Array, "D"], - float, - ] = 0.01, - ): - """ - Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) advection-diffusion - equation on periodic boundary conditions. - - In 1d, the advection-diffusion equation is given by - - ``` - uₜ + c uₓ = ν uₓₓ - ``` - - with `c ∈ ℝ` being the velocity/advection speed and `ν ∈ ℝ` being the - diffusivity. - - In higher dimensions, the advection-diffusion equation can be written as - - ``` - uₜ + c ⋅ ∇u = ν Δu - ``` - - with `c ∈ ℝᵈ` being the velocity/advection vector. - - See also [`exponax.stepper.Diffusion`][] for additional details on - anisotropic diffusion. - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. - - `domain_extent`: The size of the domain `L`; in higher dimensions - the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. - - `num_points`: The number of points `N` used to discretize the - domain. This **includes** the left boundary point and **excludes** - the right boundary point. In higher dimensions; the number of points - in each dimension is the same. Hence, the total number of degrees of - freedom is `Nᵈ`. - - `dt`: The timestep size `Δt` between two consecutive states. - - `velocity` (keyword-only): The advection speed `c`. In higher - dimensions, this can be a scalar (=float) or a vector of length `d`. - If a scalar is given, the advection speed is assumed to be the same - in all spatial dimensions. Default: `1.0`. - - `diffusivity` (keyword-only): The diffusivity `ν`. In higher - dimensions, this can be a scalar (=float), a vector of length `d`, - or a matrix of shape `d ˣ d`. If a scalar is given, the diffusivity - is assumed to be the same in all spatial dimensions. If a vector (of - length `d`) is given, the diffusivity varies across dimensions (=> - diagonal diffusion). For a matrix, there is fully anisotropic - diffusion. In this case, `A` must be symmetric positive definite - (SPD). Default: `0.01`. - - **Notes:** - - - The stepper is unconditionally stable, not matter the choice of - any argument because the equation is solved analytically in Fourier - space. **However**, note that initial conditions with modes higher - than the Nyquist freuency (`(N//2)+1` with `N` being the - `num_points`) lead to spurious oscillations. - - Ultimately, only the factors `c Δt / L` and `ν Δt / L²` affect the - characteristic of the dynamics. See also - [`exponax.normalized.NormalizedLinearStepper`][] with - `normalized_coefficients = [0, alpha_1, alpha_2]` with `alpha_1 = - - velocity * dt / domain_extent` and `alpha_2 = diffusivity * dt / - domain_extent**2`. - """ - # TODO: more sophisticated checks here - if isinstance(velocity, float): - velocity = jnp.ones(num_spatial_dims) * velocity - self.velocity = velocity - if isinstance(diffusivity, float): - diffusivity = jnp.diag(jnp.ones(num_spatial_dims)) * diffusivity - elif len(diffusivity.shape) == 1: - diffusivity = jnp.diag(diffusivity) - self.diffusivity = diffusivity - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=0, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - laplace_outer_producct = ( - derivative_operator[:, None] * derivative_operator[None, :] - ) - diffusion_operator = jnp.einsum( - "ij,ij...->...", - self.diffusivity, - laplace_outer_producct, - ) - # Add the necessary singleton channel axis - diffusion_operator = diffusion_operator[None, ...] - - advection_operator = -build_gradient_inner_product_operator( - derivative_operator, self.velocity, order=1 - ) - - linear_operator = advection_operator + diffusion_operator - - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> ZeroNonlinearFun: - return ZeroNonlinearFun( - self.num_spatial_dims, - self.num_points, - ) - - -class Dispersion(BaseStepper): - dispersivity: Float[Array, "D"] - advect_on_diffusion: bool - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - dispersivity: Union[Float[Array, "D"], float] = 1.0, - advect_on_diffusion: bool = False, - ): - """ - Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) dispersion equation - on periodic boundary conditions. Essentially, a dispersion equation is - an advection equation with different velocities (=advection speeds) for - different wavenumbers/modes. Higher wavenumbers/modes are advected - faster. - - In 1d, the dispersion equation is given by - - ``` - uₜ = 𝒸 uₓₓₓ - ``` - - with `𝒸 ∈ ℝ` being the dispersivity. - - In higher dimensions, the dispersion equation can be written as - - ``` - uₜ = 𝒸 ⋅ (∇⊙∇⊙(∇u)) - ``` - - or - - ``` - uₜ = 𝒸 ⋅ ∇(Δu) - ``` - - with `𝒸 ∈ ℝᵈ` being the dispersivity vector - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. - - `domain_extent`: The size of the domain `L`; in higher dimensions - the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. - - `num_points`: The number of points `N` used to discretize the - domain. This **includes** the left boundary point and **excludes** - the right boundary point. In higher dimensions; the number of points - in each dimension is the same. Hence, the total number of degrees of - freedom is `Nᵈ`. - - `dt`: The timestep size `Δt` between two consecutive states. - - `dispersivity` (keyword-only): The dispersivity `𝒸`. In higher - dimensions, this can be a scalar (=float) or a vector of length `d`. - If a scalar is given, the dispersivity is assumed to be the same in - all spatial dimensions. Default: `1.0`. - - `advect_on_diffusion` (keyword-only): If `True`, the second form - of the dispersion equation in higher dimensions is used. As a - consequence, there will be mixing in the spatial derivatives. - Default: `False`. - - **Notes:** - - - The stepper is unconditionally stable, not matter the choice of - any argument because the equation is solved analytically in Fourier - space. **However**, note that initial conditions with modes higher - than the Nyquist freuency (`(N//2)+1` with `N` being the - `num_points`) lead to spurious oscillations. - - Ultimately, only the factor `𝒸 Δt / L³` affects the - characteristic of the dynamics. See also - [`exponax.normalized.NormalizedLinearStepper`][] with - `normalized_coefficients = [0, 0, 0, alpha_3]` with `alpha_3 = - dispersivity * dt / domain_extent**3`. - """ - if isinstance(dispersivity, float): - dispersivity = jnp.ones(num_spatial_dims) * dispersivity - self.dispersivity = dispersivity - self.advect_on_diffusion = advect_on_diffusion - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=0, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - if self.advect_on_diffusion: - laplace_operator = build_laplace_operator(derivative_operator) - advection_operator = build_gradient_inner_product_operator( - derivative_operator, self.dispersivity, order=1 - ) - linear_operator = advection_operator * laplace_operator - else: - linear_operator = build_gradient_inner_product_operator( - derivative_operator, self.dispersivity, order=3 - ) - - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> ZeroNonlinearFun: - return ZeroNonlinearFun( - self.num_spatial_dims, - self.num_points, - ) - - -class HyperDiffusion(BaseStepper): - hyper_diffusivity: float - diffuse_on_diffuse: bool - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - hyper_diffusivity: float = 0.0001, - diffuse_on_diffuse: bool = False, - ): - """ - Timestepper for the d-dimensional (`d ∈ {1, 2, 3}`) hyper-diffusion - equation on periodic boundary conditions. A hyper-diffusion equation - acts like a diffusion equation but higher wavenumbers/modes are damped - even faster. - - In 1d, the hyper-diffusion equation is given by - - ``` - uₜ = - μ uₓₓₓₓ - ``` - - with `μ ∈ ℝ` being the hyper-diffusivity. - - Note the minus sign because by default, a fourth-order derivative - dampens with a negative coefficient. To match the concept of - second-order diffusion, a negation is introduced. - - In higher dimensions, the hyper-diffusion equation can be written as - - ``` - uₜ = − μ ((∇⊙∇) ⋅ (∇⊙∇)) u - ``` - - or - - ``` - uₜ = - μ Δ(Δu) - ``` - - The latter introduces spatial mixing. - - **Arguments:** - - - `num_spatial_dims`: The number of spatial dimensions `d`. - - `domain_extent`: The size of the domain `L`; in higher dimensions - the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. - - `num_points`: The number of points `N` used to discretize the - domain. This **includes** the left boundary point and **excludes** - the right boundary point. In higher dimensions; the number of points - in each dimension is the same. Hence, the total number of degrees of - freedom is `Nᵈ`. - - `dt`: The timestep size `Δt` between two consecutive states. - - `hyper_diffusivity` (keyword-only): The hyper-diffusivity `ν`. - This stepper only supports scalar (=isotropic) hyper-diffusivity. - Default: 0.0001. - - `diffuse_on_diffuse` (keyword-only): If `True`, the second form - of the hyper-diffusion equation in higher dimensions is used. As a - consequence, there will be mixing in the spatial derivatives. - Default: `False`. - - **Notes:** - - - The stepper is unconditionally stable, not matter the choice of - any argument because the equation is solved analytically in Fourier - space. - - Ultimately, only the factor `μ Δt / L⁴` affects the characteristic - of the dynamics. See also - [`exponax.normalized.NormalizedLinearStepper`][] with - `normalized_coefficients = [0, 0, 0, 0, alpha_4]` with `alpha_4 = - - hyper_diffusivity * dt / domain_extent**4`. - """ - self.hyper_diffusivity = hyper_diffusivity - self.diffuse_on_diffuse = diffuse_on_diffuse - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=0, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - # Use minus sign to have diffusion work in "correct direction" by default - if self.diffuse_on_diffuse: - laplace_operator = build_laplace_operator(derivative_operator) - linear_operator = ( - -self.hyper_diffusivity * laplace_operator * laplace_operator - ) - else: - linear_operator = -self.hyper_diffusivity * build_laplace_operator( - derivative_operator, order=4 - ) - - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> ZeroNonlinearFun: - return ZeroNonlinearFun( - self.num_spatial_dims, - self.num_points, - ) - - -class GeneralLinearStepper(BaseStepper): - coefficients: tuple[float, ...] - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - coefficients: tuple[float, ...] = (0.0, -0.1, 0.01), - ): - """ - General timestepper for a d-dimensional (`d ∈ {1, 2, 3}`) linear - equation with an arbitrary combination of derivative terms and - respective coefficients. To simplify the interface, only isotropicity is - supported. For example, the advection speed is the same in all spatial - dimensions, or diffusion is equally strong in all spatial dimensions. - - In 1d the equation is given by - - ``` - uₜ = sum_j a_j uₓˢ - ``` - - with `uₓˢ` denoting the s-th derivative of `u` with respect to `x`. The - coefficient corresponding to this derivative is `a_j`. - - The isotropic version in higher dimensions can expressed as - - ``` - uₜ = sum_j a_j (1⋅∇ʲ)u - ``` - - with `1⋅∇ʲ` denoting the j-th repeated elementwise product of the nabla - operator with itself in an inner product with the one vector. For - example, `1⋅∇¹` is the collection of first derivatives, `1⋅∇²` is the - collection of second derivatives (=Laplace operator), etc. - - The interface to this general stepper is the list of coefficients - containing the `a_j`. Its length determines the highes occuring order of - derivative. Note that this list starts at zero. If only one specific - linear term is wanted, have all prior coefficients set to zero. - - The default configuration is an advection-diffusion equation with `a_0 = - 0`, `a_1 = -0.1`, and `a_2 = 0.01`. - - **Arguments:** - - `num_spatial_dims`: The number of spatial dimensions `d`. - - `domain_extent`: The size of the domain `L`; in higher dimensions - the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. - - `num_points`: The number of points `N` used to discretize the - domain. This **includes** the left boundary point and - **excludes** the right boundary point. In higher dimensions; the - number of points in each dimension is the same. Hence, the total - number of degrees of freedom is `Nᵈ`. - - `dt`: The timestep size `Δt` between two consecutive states. - - `coefficients` (keyword-only): The list of coefficients `a_j` - corresponding to the derivatives. Default: `[0.0, -0.1, 0.01]`. - - **Notes:** - - There is a repeating pattern in the effect of orders of - derivatives: - - Even derivatives (i.e., 0, 2, 4, 6, ...) scale the - solution. Order 0 scales all wavenumbers/modes equally (if - its coefficient is negative, this is also called a drag). - Order 2 scales higher wavenumbers/modes stronger with the - dependence on the effect on the wavenumber being - quadratically. Order 4 also scales but stronger than order - 4. Its dependency on the wavenumber is quartically. This - pattern continues for higher orders. - - Odd derivatives (i.e, 1, 3, 5, 7, ...) rotate the solution in - Fourier space. In state space, this is observed as - advection. Order 1 rotates all wavenumbers/modes equally. In - state space, this is observed in that the initial condition - just moves over the domain. Order 3 rotates higher - wavenumbers/modes stronger with the dependence on the - wavenumber being quadratic. If certain wavenumbers are - rotated at a different speed, there is still advection in - state space but certain patterns in the initial condition - are advected at different speeds. As such, it is observed - that the shape of the initial condition dissolves. The - effect continues for higher orders with the dependency on - the wavenumber becoming continuously stronger. - - Take care of the signs of coefficients. In contrast to the - indivial linear steppers ([`exponax.stepper.Advection`][], - [`exponax.stepper.Diffusion`][], etc.), the signs are not - automatically taken care of to produce meaningful coefficients. - For the general linear stepper all linear derivatives are on the - right-hand side of the equation. This has the following effect - based on the order of derivative (this a consequence of squaring - the imaginary unit returning -1): - - Zeroth-Order: A negative coeffcient is a drag and removes - energy from the system. A positive coefficient adds energy - to the system. - - First-Order: A negative coefficient rotates the solution - clockwise. A positive coefficient rotates the solution - counter-clockwise. Hence, negative coefficients advect - solutions to the right, positive coefficients advect - solutions to the left. - - Second-Order: A positive coefficient diffuses the solution - (i.e., removes energy from the system). A negative - coefficient adds energy to the system. - - Third-Order: A negative coefficient rotates the solution - counter-clockwise. A positive coefficient rotates the - solution clockwise. Hence, negative coefficients advect - solutions to the left, positive coefficients advect - solutions to the right. - - Fourth-Order: A negative coefficient diffuses the solution - (i.e., removes energy from the system). A positive - coefficient adds energy to the system. - - ... - - The stepper is unconditionally stable, no matter the choice of - any argument because the equation is solved analytically in - Fourier space. **However**, note if you have odd-order - derivative terms (e.g., advection or dispersion) and your - initial condition is **not** bandlimited (i.e., it contains - modes beyond the Nyquist frequency of `(N//2)+1`) there is a - chance spurious oscillations can occur. - - Ultimately, only the factors `a_j Δt / Lʲ` affect the - characteristic of the dynamics. See also - [`exponax.normalized.NormalizedLinearStepper`][] with - `normalized_coefficients = [0, alpha_1, alpha_2, ...]` with - `alpha_j = coefficients[j] * dt / domain_extent**j`. You can use - the function [`exponax.normalized.normalize_coefficients`][] to - obtain the normalized coefficients. - """ - self.coefficients = coefficients - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=0, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.coefficients) - ) - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> ZeroNonlinearFun: - return ZeroNonlinearFun( - self.num_spatial_dims, - self.num_points, - ) diff --git a/exponax/stepper/_polynomial.py b/exponax/stepper/_polynomial.py deleted file mode 100644 index ae8b9d2..0000000 --- a/exponax/stepper/_polynomial.py +++ /dev/null @@ -1,72 +0,0 @@ -import jax.numpy as jnp -from jaxtyping import Array, Complex - -from .._base_stepper import BaseStepper -from ..nonlin_fun import PolynomialNonlinearFun - - -class GeneralPolynomialStepper(BaseStepper): - coefficients: tuple[float, ...] - polynomial_scales: tuple[float, ...] - dealiasing_fraction: float - - def __init__( - self, - num_spatial_dims: int, - domain_extent: float, - num_points: int, - dt: float, - *, - coefficients: tuple[float, ...] = (10.0, 0.0, 0.01), - polynomial_scales: tuple[float, ...] = (0.0, 0.0, 10.0), - order=2, - dealiasing_fraction: float = 2 / 3, - num_circle_points: int = 16, - circle_radius: float = 1.0, - ): - """ - By default: Fisher-KPP with a small diffusion and 10.0 reactivity - - Note that the first two entries in the polynomial_scales list are often zero. - - The effect of polynomial_scale[1] is similar to the effect of coefficients[0] - """ - self.coefficients = coefficients - self.polynomial_scales = polynomial_scales - self.dealiasing_fraction = dealiasing_fraction - - super().__init__( - num_spatial_dims=num_spatial_dims, - domain_extent=domain_extent, - num_points=num_points, - dt=dt, - num_channels=1, - order=order, - num_circle_points=num_circle_points, - circle_radius=circle_radius, - ) - - def _build_linear_operator( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> Complex[Array, "1 ... (N//2)+1"]: - linear_operator = sum( - jnp.sum( - c * (derivative_operator) ** i, - axis=0, - keepdims=True, - ) - for i, c in enumerate(self.coefficients) - ) - return linear_operator - - def _build_nonlinear_fun( - self, - derivative_operator: Complex[Array, "D ... (N//2)+1"], - ) -> PolynomialNonlinearFun: - return PolynomialNonlinearFun( - self.num_spatial_dims, - self.num_points, - dealiasing_fraction=self.dealiasing_fraction, - coefficients=self.polynomial_scales, - ) diff --git a/exponax/normalized/__init__.py b/exponax/stepper/generic/__init__.py similarity index 64% rename from exponax/normalized/__init__.py rename to exponax/stepper/generic/__init__.py index 8d177c9..fcd94a0 100644 --- a/exponax/normalized/__init__.py +++ b/exponax/stepper/generic/__init__.py @@ -1,24 +1,29 @@ -""" -Submodule for timesteppers with normalized dynamics. Those steppers operate on -unit dynamics, i.e., `domain_extent = 1.0` and `dt = 1.0`. As such, the -constitutive coefficients are normalized. One has to supply fewer arguments to -uniquely describe a dyanamics. - -Additionally, there are Difficulty steppers that interface the same concept -slightly differently. -""" -from ._convection import DifficultyConvectionStepper, NormalizedConvectionStepper -from ._general_nonlinear import ( - DifficultyGeneralNonlinearStepper, - NormalizedGeneralNonlinearStepper, +from ._convection import ( + DifficultyConvectionStepper, + GeneralConvectionStepper, + NormalizedConvectionStepper, +) +from ._gradient_norm import ( + DifficultyGradientNormStepper, + GeneralGradientNormStepper, + NormalizedGradientNormStepper, ) -from ._gradient_norm import DifficultyGradientNormStepper, NormalizedGradientNormStepper from ._linear import ( DifficultyLinearStepper, DiffultyLinearStepperSimple, + GeneralLinearStepper, NormalizedLinearStepper, ) -from ._polynomial import DifficultyPolynomialStepper, NormalizedPolynomialStepper +from ._nonlinear import ( + DifficultyNonlinearStepper, + GeneralNonlinearStepper, + NormalizedNonlinearStepper, +) +from ._polynomial import ( + DifficultyPolynomialStepper, + GeneralPolynomialStepper, + NormalizedPolynomialStepper, +) from ._utils import ( denormalize_coefficients, denormalize_convection_scale, @@ -35,21 +40,26 @@ reduce_normalized_convection_scale_to_difficulty, reduce_normalized_gradient_norm_scale_to_difficulty, ) -from ._vorticity_convection import NormalizedVorticityConvection +from ._vorticity_convection import GeneralVorticityConvectionStepper __all__ = [ + "GeneralLinearStepper", + "NormalizedLinearStepper", "DifficultyLinearStepper", - "DiffultyLinearStepperSimple", - "DifficultyConvectionStepper", - "DifficultyGradientNormStepper", - "DifficultyPolynomialStepper", - "DifficultyGeneralNonlinearStepper", + "GeneralConvectionStepper", "NormalizedConvectionStepper", - "NormalizedGeneralNonlinearStepper", + "DifficultyConvectionStepper", + "GeneralGradientNormStepper", "NormalizedGradientNormStepper", - "NormalizedLinearStepper", + "DifficultyGradientNormStepper", + "GeneralVorticityConvectionStepper", + "GeneralPolynomialStepper", "NormalizedPolynomialStepper", - "NormalizedVorticityConvection", + "DifficultyPolynomialStepper", + "GeneralNonlinearStepper", + "NormalizedNonlinearStepper", + "DifficultyNonlinearStepper", + "DiffultyLinearStepperSimple", "denormalize_coefficients", "denormalize_convection_scale", "denormalize_gradient_norm_scale", diff --git a/exponax/stepper/_convection.py b/exponax/stepper/generic/_convection.py similarity index 63% rename from exponax/stepper/_convection.py rename to exponax/stepper/generic/_convection.py index b54159e..c5b5aca 100644 --- a/exponax/stepper/_convection.py +++ b/exponax/stepper/generic/_convection.py @@ -1,8 +1,12 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from ..nonlin_fun import ConvectionNonlinearFun +from ..._base_stepper import BaseStepper +from ...nonlin_fun import ConvectionNonlinearFun +from ._utils import ( + extract_normalized_coefficients_from_difficulty, + extract_normalized_convection_scale_from_difficulty, +) class GeneralConvectionStepper(BaseStepper): @@ -142,3 +146,95 @@ def _build_nonlinear_fun( scale=self.convection_scale, single_channel=self.single_channel, ) + + +class NormalizedConvectionStepper(GeneralConvectionStepper): + normalized_coefficients: tuple[float, ...] + normalized_convection_scale: float + + def __init__( + self, + num_spatial_dims: int, + num_points: int, + *, + normalized_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1), + normalized_convection_scale: float = 1.0 * 0.1, + single_channel: bool = False, + order: int = 2, + dealiasing_fraction: float = 2 / 3, + num_circle_points: int = 16, + circle_radius: float = 1.0, + ): + """ + By default: Behaves like a Burgers with + + ``` Burgers( + D=D, L=1, N=N, dt=0.1, diffusivity=0.01, + ) + ``` + """ + self.normalized_coefficients = normalized_coefficients + self.normalized_convection_scale = normalized_convection_scale + 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, + coefficients=normalized_coefficients, + convection_scale=normalized_convection_scale, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + single_channel=single_channel, + ) + + +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, + single_channel: bool = False, + 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, + single_channel=single_channel, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + ) diff --git a/exponax/stepper/_gradient_norm.py b/exponax/stepper/generic/_gradient_norm.py similarity index 60% rename from exponax/stepper/_gradient_norm.py rename to exponax/stepper/generic/_gradient_norm.py index b463276..fe03f7d 100644 --- a/exponax/stepper/_gradient_norm.py +++ b/exponax/stepper/generic/_gradient_norm.py @@ -1,8 +1,12 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from ..nonlin_fun import GradientNormNonlinearFun +from ..._base_stepper import BaseStepper +from ...nonlin_fun import GradientNormNonlinearFun +from ._utils import ( + extract_normalized_coefficients_from_difficulty, + extract_normalized_gradient_norm_scale_from_difficulty, +) class GeneralGradientNormStepper(BaseStepper): @@ -124,3 +128,97 @@ def _build_nonlinear_fun( scale=self.gradient_norm_scale, zero_mode_fix=True, # Todo: check this ) + + +class NormalizedGradientNormStepper(GeneralGradientNormStepper): + normalized_coefficients: tuple[float, ...] + normalized_gradient_norm_scale: float + + def __init__( + self, + num_spatial_dims: int, + num_points: int, + *, + normalized_coefficients: tuple[float, ...] = ( + 0.0, + 0.0, + -1.0 * 0.1 / (60.0**2), + 0.0, + -1.0 * 0.1 / (60.0**4), + ), + normalized_gradient_norm_scale: float = 1.0 * 0.1 / (60.0**2), + order: int = 2, + dealiasing_fraction: float = 2 / 3, + num_circle_points: int = 16, + circle_radius: float = 1.0, + ): + """ + 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 + + """ + self.normalized_coefficients = normalized_coefficients + self.normalized_gradient_norm_scale = normalized_gradient_norm_scale + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=1.0, + num_points=num_points, + dt=1.0, + coefficients=normalized_coefficients, + gradient_norm_scale=normalized_gradient_norm_scale, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + ) + + +class DifficultyGradientNormStepper(NormalizedGradientNormStepper): + linear_difficulties: tuple[float, ...] + gradient_norm_difficulty: float + + def __init__( + self, + num_spatial_dims: int = 1, + num_points: int = 48, + *, + linear_difficulties: tuple[float, ...] = (0.0, 0.0, -0.128, 0.0, -0.32768), + gradient_norm_difficulty: float = 0.064, + 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: KS equation + """ + self.linear_difficulties = linear_difficulties + self.gradient_norm_difficulty = gradient_norm_difficulty + + normalized_coefficients = extract_normalized_coefficients_from_difficulty( + linear_difficulties, + num_spatial_dims=num_spatial_dims, + num_points=num_points, + ) + normalized_gradient_norm_scale = ( + extract_normalized_gradient_norm_scale_from_difficulty( + gradient_norm_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_gradient_norm_scale=normalized_gradient_norm_scale, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + ) diff --git a/exponax/stepper/generic/_linear.py b/exponax/stepper/generic/_linear.py new file mode 100644 index 0000000..2c9fc15 --- /dev/null +++ b/exponax/stepper/generic/_linear.py @@ -0,0 +1,324 @@ +from typing import TypeVar + +import jax.numpy as jnp +from jaxtyping import Array, Complex + +from ..._base_stepper import BaseStepper +from ...nonlin_fun import ZeroNonlinearFun +from ._utils import extract_normalized_coefficients_from_difficulty + +D = TypeVar("D") + + +class GeneralLinearStepper(BaseStepper): + coefficients: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int, + domain_extent: float, + num_points: int, + dt: float, + *, + coefficients: tuple[float, ...] = (0.0, -0.1, 0.01), + ): + """ + General timestepper for a d-dimensional (`d ∈ {1, 2, 3}`) linear + equation with an arbitrary combination of derivative terms and + respective coefficients. To simplify the interface, only isotropicity is + supported. For example, the advection speed is the same in all spatial + dimensions, or diffusion is equally strong in all spatial dimensions. + + In 1d the equation is given by + + ``` + uₜ = sum_j a_j uₓˢ + ``` + + with `uₓˢ` denoting the s-th derivative of `u` with respect to `x`. The + coefficient corresponding to this derivative is `a_j`. + + The isotropic version in higher dimensions can expressed as + + ``` + uₜ = sum_j a_j (1⋅∇ʲ)u + ``` + + with `1⋅∇ʲ` denoting the j-th repeated elementwise product of the nabla + operator with itself in an inner product with the one vector. For + example, `1⋅∇¹` is the collection of first derivatives, `1⋅∇²` is the + collection of second derivatives (=Laplace operator), etc. + + The interface to this general stepper is the list of coefficients + containing the `a_j`. Its length determines the highes occuring order of + derivative. Note that this list starts at zero. If only one specific + linear term is wanted, have all prior coefficients set to zero. + + The default configuration is an advection-diffusion equation with `a_0 = + 0`, `a_1 = -0.1`, and `a_2 = 0.01`. + + **Arguments:** + - `num_spatial_dims`: The number of spatial dimensions `d`. + - `domain_extent`: The size of the domain `L`; in higher dimensions + the domain is assumed to be a scaled hypercube `Ω = (0, L)ᵈ`. + - `num_points`: The number of points `N` used to discretize the + domain. This **includes** the left boundary point and + **excludes** the right boundary point. In higher dimensions; the + number of points in each dimension is the same. Hence, the total + number of degrees of freedom is `Nᵈ`. + - `dt`: The timestep size `Δt` between two consecutive states. + - `coefficients` (keyword-only): The list of coefficients `a_j` + corresponding to the derivatives. Default: `[0.0, -0.1, 0.01]`. + + **Notes:** + - There is a repeating pattern in the effect of orders of + derivatives: + - Even derivatives (i.e., 0, 2, 4, 6, ...) scale the + solution. Order 0 scales all wavenumbers/modes equally (if + its coefficient is negative, this is also called a drag). + Order 2 scales higher wavenumbers/modes stronger with the + dependence on the effect on the wavenumber being + quadratically. Order 4 also scales but stronger than order + 4. Its dependency on the wavenumber is quartically. This + pattern continues for higher orders. + - Odd derivatives (i.e, 1, 3, 5, 7, ...) rotate the solution in + Fourier space. In state space, this is observed as + advection. Order 1 rotates all wavenumbers/modes equally. In + state space, this is observed in that the initial condition + just moves over the domain. Order 3 rotates higher + wavenumbers/modes stronger with the dependence on the + wavenumber being quadratic. If certain wavenumbers are + rotated at a different speed, there is still advection in + state space but certain patterns in the initial condition + are advected at different speeds. As such, it is observed + that the shape of the initial condition dissolves. The + effect continues for higher orders with the dependency on + the wavenumber becoming continuously stronger. + - Take care of the signs of coefficients. In contrast to the + indivial linear steppers ([`exponax.stepper.Advection`][], + [`exponax.stepper.Diffusion`][], etc.), the signs are not + automatically taken care of to produce meaningful coefficients. + For the general linear stepper all linear derivatives are on the + right-hand side of the equation. This has the following effect + based on the order of derivative (this a consequence of squaring + the imaginary unit returning -1): + - Zeroth-Order: A negative coeffcient is a drag and removes + energy from the system. A positive coefficient adds energy + to the system. + - First-Order: A negative coefficient rotates the solution + clockwise. A positive coefficient rotates the solution + counter-clockwise. Hence, negative coefficients advect + solutions to the right, positive coefficients advect + solutions to the left. + - Second-Order: A positive coefficient diffuses the solution + (i.e., removes energy from the system). A negative + coefficient adds energy to the system. + - Third-Order: A negative coefficient rotates the solution + counter-clockwise. A positive coefficient rotates the + solution clockwise. Hence, negative coefficients advect + solutions to the left, positive coefficients advect + solutions to the right. + - Fourth-Order: A negative coefficient diffuses the solution + (i.e., removes energy from the system). A positive + coefficient adds energy to the system. + - ... + - The stepper is unconditionally stable, no matter the choice of + any argument because the equation is solved analytically in + Fourier space. **However**, note if you have odd-order + derivative terms (e.g., advection or dispersion) and your + initial condition is **not** bandlimited (i.e., it contains + modes beyond the Nyquist frequency of `(N//2)+1`) there is a + chance spurious oscillations can occur. + - Ultimately, only the factors `a_j Δt / Lʲ` affect the + characteristic of the dynamics. See also + [`exponax.stepper.generic.NormalizedLinearStepper`][] with + `normalized_coefficients = [0, alpha_1, alpha_2, ...]` with + `alpha_j = coefficients[j] * dt / domain_extent**j`. You can use + the function [`exponax.stepper.generic.normalize_coefficients`][] to + obtain the normalized coefficients. + """ + self.coefficients = coefficients + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=domain_extent, + num_points=num_points, + dt=dt, + num_channels=1, + order=0, + ) + + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: + linear_operator = sum( + jnp.sum( + c * (derivative_operator) ** i, + axis=0, + keepdims=True, + ) + for i, c in enumerate(self.coefficients) + ) + return linear_operator + + def _build_nonlinear_fun( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> ZeroNonlinearFun: + return ZeroNonlinearFun( + self.num_spatial_dims, + self.num_points, + ) + + +class NormalizedLinearStepper(GeneralLinearStepper): + normalized_coefficients: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int, + num_points: int, + *, + normalized_coefficients: tuple[float, ...] = (0.0, -0.5, 0.01), + ): + """ + Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic + boundary conditions with normalized dynamics. + + If the PDE in physical description is + + uₜ = ∑ᵢ aᵢ (∂ₓ)ⁱ u + + with `aᵢ` the coefficients on `domain_extent` `L` with time step `dt`, + the normalized coefficients are + + αᵢ = (aᵢ Δt)/(Lⁱ) + + Important: note that the `domain_extent` is raised to the order of + linear derivative `i`. + + One can also think of normalized dynamics, as a PDE on `domain_extent` + `1.0` with time step `dt=1.0`. + + Take care of the signs! + + In the defaulf configuration of this timestepper, the PDE is an + advection-diffusion equation with normalized advection of 0.5 and + normalized diffusion of 0.01. + + **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** + the right boundary point. In higher dimensions; the number of points + in each dimension is the same. Hence, the total number of degrees of + freedom is `Nᵈ`. + - `normalized_coefficients`: The coefficients of the normalized + dynamics. This must a tuple of floats. The length of the tuple + defines the highest occuring linear derivative in the PDE. + """ + self.normalized_coefficients = normalized_coefficients + super().__init__( + num_spatial_dims=num_spatial_dims, + domain_extent=1.0, + num_points=num_points, + dt=1.0, + coefficients=normalized_coefficients, + ) + + +class DifficultyLinearStepper(NormalizedLinearStepper): + difficulties: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int = 1, + num_points: int = 48, + *, + difficulties: tuple[float, ...] = (0.0, -2.0), + ): + """ + Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic + boundary conditions with normalized dynamics in a difficulty-based + interface. + + Different to `NormalizedLinearStepper`, the dynamics are defined by + difficulties. The difficulties are a different combination of normalized + dynamics, `num_spatial_dims`, and `num_points`. + + γᵢ = αᵢ Nⁱ 2ⁱ⁻¹ d + + with `d` the number of spatial dimensions, `N` the number of points, and + `αᵢ` the normalized coefficient. + + This interface is more natural because the difficulties for all orders + (given by `i`) are around 1.0. Additionally, they relate to stability + condition of explicit Finite Difference schemes for the particular + equations. For example, for advection (`i=1`), the absolute of the + difficulty is the Courant-Friedrichs-Lewy (CFL) number. + + In the default configuration of this timestepper, the PDE is an + advection equation with CFL number 2 solved in 1d with 48 resolution + points to discretize the domain. + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. Default is + 1. + - `num_points`: The number of points `N` used to discretize the domain. + This **includes** the left boundary point and **excludes** the right + boundary point. In higher dimensions; the number of points in each + dimension is the same. Hence, the total number of degrees of freedom + is `Nᵈ`. Default is 48. + - `difficulties`: The difficulties of the normalized dynamics. This must + be a tuple of floats. The length of the tuple defines the highest + occuring linear derivative in the PDE. Default is `(0.0, -2.0)`. + """ + self.difficulties = difficulties + normalized_coefficients = extract_normalized_coefficients_from_difficulty( + difficulties, + num_spatial_dims=num_spatial_dims, + num_points=num_points, + ) + + super().__init__( + num_spatial_dims=num_spatial_dims, + num_points=num_points, + normalized_coefficients=normalized_coefficients, + ) + + +class DiffultyLinearStepperSimple(DifficultyLinearStepper): + def __init__( + self, + num_spatial_dims: int = 1, + num_points: int = 48, + *, + difficulty: float = -2.0, + order: int = 1, + ): + """ + A simple interface for `DifficultyLinearStepper` with only one + difficulty and a given order. + + **Arguments:** + + - `num_spatial_dims`: The number of spatial dimensions `d`. Default is + 1. + - `num_points`: The number of points `N` used to discretize the domain. + This **includes** the left boundary point and **excludes** the right + boundary point. In higher dimensions; the number of points in each + dimension is the same. Hence, the total number of degrees of freedom + is `Nᵈ`. Default is 48. + - `difficulty`: The difficulty of the normalized dynamics. This must be + a float. Default is -2.0. + - `order`: The order of the derivative associated with the provided + difficulty. The default of 1 is the advection equation. + """ + difficulties = (0.0,) * (order) + (difficulty,) + super().__init__( + difficulties=difficulties, + num_spatial_dims=num_spatial_dims, + num_points=num_points, + ) diff --git a/exponax/normalized/_general_nonlinear.py b/exponax/stepper/generic/_nonlinear.py similarity index 69% rename from exponax/normalized/_general_nonlinear.py rename to exponax/stepper/generic/_nonlinear.py index 202abac..67c43ec 100644 --- a/exponax/normalized/_general_nonlinear.py +++ b/exponax/stepper/generic/_nonlinear.py @@ -1,62 +1,66 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from ..nonlin_fun import GeneralNonlinearFun +from ..._base_stepper import BaseStepper +from ...nonlin_fun import GeneralNonlinearFun from ._utils import ( extract_normalized_coefficients_from_difficulty, extract_normalized_nonlinear_scales_from_difficulty, ) -class NormalizedGeneralNonlinearStepper(BaseStepper): - normalized_coefficients_linear: tuple[float, ...] - normalized_coefficients_nonlinear: tuple[float, ...] +class GeneralNonlinearStepper(BaseStepper): + coefficients_linear: tuple[float, ...] + coefficients_nonlinear: tuple[float, 3] dealiasing_fraction: float def __init__( self, num_spatial_dims: int, + domain_extent: float, num_points: int, + dt: float, *, - normalized_coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.1 * 0.1), - normalized_coefficients_nonlinear: tuple[float, ...] = (0.0, -1.0 * 0.1, 0.0), + coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.01), + coefficients_nonlinear: tuple[float, 3] = (0.0, -1.0, 0.0), order=2, dealiasing_fraction: float = 2 / 3, num_circle_points: int = 16, circle_radius: float = 1.0, ): """ - By default Burgers. + By default Burgers equation """ - - if len(normalized_coefficients_nonlinear) != 3: + if len(coefficients_nonlinear) != 3: raise ValueError( "The nonlinear coefficients list must have exactly 3 elements" ) - self.normalized_coefficients_linear = normalized_coefficients_linear - self.normalized_coefficients_nonlinear = normalized_coefficients_nonlinear + self.coefficients_linear = coefficients_linear + self.coefficients_nonlinear = coefficients_nonlinear self.dealiasing_fraction = dealiasing_fraction super().__init__( num_spatial_dims=num_spatial_dims, - domain_extent=1.0, # Derivative operator is just scaled with 2 * jnp.pi + domain_extent=domain_extent, num_points=num_points, - dt=1.0, + dt=dt, num_channels=1, order=order, num_circle_points=num_circle_points, circle_radius=circle_radius, ) - def _build_linear_operator(self, derivative_operator: Array) -> Array: + def _build_linear_operator( + self, + derivative_operator: Complex[Array, "D ... (N//2)+1"], + ) -> Complex[Array, "1 ... (N//2)+1"]: linear_operator = sum( jnp.sum( c * (derivative_operator) ** i, axis=0, keepdims=True, ) - for i, c in enumerate(self.normalized_coefficients_linear) + for i, c in enumerate(self.coefficients_linear) ) return linear_operator @@ -69,12 +73,49 @@ def _build_nonlinear_fun( self.num_points, derivative_operator=derivative_operator, dealiasing_fraction=self.dealiasing_fraction, - scale_list=self.normalized_coefficients_nonlinear, + scale_list=self.coefficients_nonlinear, zero_mode_fix=True, # ToDo: check this ) -class DifficultyGeneralNonlinearStepper(NormalizedGeneralNonlinearStepper): +class NormalizedNonlinearStepper(GeneralNonlinearStepper): + normalized_coefficients_linear: tuple[float, ...] + normalized_coefficients_nonlinear: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int, + num_points: int, + *, + normalized_coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.1 * 0.1), + normalized_coefficients_nonlinear: tuple[float, ...] = (0.0, -1.0 * 0.1, 0.0), + order=2, + dealiasing_fraction: float = 2 / 3, + num_circle_points: int = 16, + circle_radius: float = 1.0, + ): + """ + By default Burgers. + """ + + self.normalized_coefficients_linear = normalized_coefficients_linear + self.normalized_coefficients_nonlinear = normalized_coefficients_nonlinear + + 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, + coefficients_linear=normalized_coefficients_linear, + coefficients_nonlinear=normalized_coefficients_nonlinear, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, + ) + + +class DifficultyNonlinearStepper(NormalizedNonlinearStepper): linear_difficulties: tuple[float, ...] nonlinear_difficulties: tuple[float, ...] diff --git a/exponax/normalized/_polynomial.py b/exponax/stepper/generic/_polynomial.py similarity index 69% rename from exponax/normalized/_polynomial.py rename to exponax/stepper/generic/_polynomial.py index 0ff1402..bb95451 100644 --- a/exponax/normalized/_polynomial.py +++ b/exponax/stepper/generic/_polynomial.py @@ -1,48 +1,46 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from ..nonlin_fun import PolynomialNonlinearFun +from ..._base_stepper import BaseStepper +from ...nonlin_fun import PolynomialNonlinearFun from ._utils import extract_normalized_coefficients_from_difficulty -class NormalizedPolynomialStepper(BaseStepper): - normalized_coefficients: tuple[float, ...] - normalized_polynomial_scales: tuple[float, ...] +class GeneralPolynomialStepper(BaseStepper): + coefficients: tuple[float, ...] + polynomial_scales: tuple[float, ...] dealiasing_fraction: float def __init__( self, num_spatial_dims: int, + domain_extent: float, num_points: int, + dt: float, *, - normalized_coefficients: tuple[float, ...] = ( - 10.0 * 0.001 / (10.0**0), - 0.0, - 1.0 * 0.001 / (10.0**2), - ), - normalized_polynomial_scales: tuple[float, ...] = ( - 0.0, - 0.0, - -10.0 * 0.001, - ), - order: int = 2, + coefficients: tuple[float, ...] = (10.0, 0.0, 0.01), + polynomial_scales: tuple[float, ...] = (0.0, 0.0, 10.0), + order=2, dealiasing_fraction: float = 2 / 3, num_circle_points: int = 16, circle_radius: float = 1.0, ): """ - By default: Fisher-KPP + By default: Fisher-KPP with a small diffusion and 10.0 reactivity + + Note that the first two entries in the polynomial_scales list are often zero. + + The effect of polynomial_scale[1] is similar to the effect of coefficients[0] """ - self.normalized_coefficients = normalized_coefficients - self.normalized_polynomial_scales = normalized_polynomial_scales + self.coefficients = coefficients + self.polynomial_scales = polynomial_scales self.dealiasing_fraction = dealiasing_fraction super().__init__( num_spatial_dims=num_spatial_dims, - domain_extent=1.0, # Derivative operator is just scaled with 2 * jnp.pi + domain_extent=domain_extent, num_points=num_points, - dt=1.0, + dt=dt, num_channels=1, order=order, num_circle_points=num_circle_points, @@ -59,7 +57,7 @@ def _build_linear_operator( axis=0, keepdims=True, ) - for i, c in enumerate(self.normalized_coefficients) + for i, c in enumerate(self.coefficients) ) return linear_operator @@ -70,8 +68,52 @@ def _build_nonlinear_fun( return PolynomialNonlinearFun( self.num_spatial_dims, self.num_points, - coefficients=self.normalized_polynomial_scales, dealiasing_fraction=self.dealiasing_fraction, + coefficients=self.polynomial_scales, + ) + + +class NormalizedPolynomialStepper(GeneralPolynomialStepper): + normalized_coefficients: tuple[float, ...] + normalized_polynomial_scales: tuple[float, ...] + + def __init__( + self, + num_spatial_dims: int, + num_points: int, + *, + normalized_coefficients: tuple[float, ...] = ( + 10.0 * 0.001 / (10.0**0), + 0.0, + 1.0 * 0.001 / (10.0**2), + ), + normalized_polynomial_scales: tuple[float, ...] = ( + 0.0, + 0.0, + -10.0 * 0.001, + ), + order: int = 2, + dealiasing_fraction: float = 2 / 3, + num_circle_points: int = 16, + circle_radius: float = 1.0, + ): + """ + By default: Fisher-KPP + """ + self.normalized_coefficients = normalized_coefficients + self.normalized_polynomial_scales = normalized_polynomial_scales + + 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, + coefficients=normalized_coefficients, + polynomial_scales=normalized_polynomial_scales, + order=order, + dealiasing_fraction=dealiasing_fraction, + num_circle_points=num_circle_points, + circle_radius=circle_radius, ) diff --git a/exponax/normalized/_utils.py b/exponax/stepper/generic/_utils.py similarity index 100% rename from exponax/normalized/_utils.py rename to exponax/stepper/generic/_utils.py diff --git a/exponax/stepper/_vorticity_convection.py b/exponax/stepper/generic/_vorticity_convection.py similarity index 95% rename from exponax/stepper/_vorticity_convection.py rename to exponax/stepper/generic/_vorticity_convection.py index f318085..8f64e81 100644 --- a/exponax/stepper/_vorticity_convection.py +++ b/exponax/stepper/generic/_vorticity_convection.py @@ -1,8 +1,8 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from ..nonlin_fun import VorticityConvection2d, VorticityConvection2dKolmogorov +from ..._base_stepper import BaseStepper +from ...nonlin_fun import VorticityConvection2d, VorticityConvection2dKolmogorov class GeneralVorticityConvectionStepper(BaseStepper): diff --git a/exponax/reaction/__init__.py b/exponax/stepper/reaction/__init__.py similarity index 100% rename from exponax/reaction/__init__.py rename to exponax/stepper/reaction/__init__.py diff --git a/exponax/reaction/_allen_cahn.py b/exponax/stepper/reaction/_allen_cahn.py similarity index 97% rename from exponax/reaction/_allen_cahn.py rename to exponax/stepper/reaction/_allen_cahn.py index 27d4f29..fa4a5f2 100644 --- a/exponax/reaction/_allen_cahn.py +++ b/exponax/stepper/reaction/_allen_cahn.py @@ -1,8 +1,8 @@ from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from .._spectral import build_laplace_operator -from ..nonlin_fun import PolynomialNonlinearFun +from ..._base_stepper import BaseStepper +from ..._spectral import build_laplace_operator +from ...nonlin_fun import PolynomialNonlinearFun class AllenCahn(BaseStepper): diff --git a/exponax/reaction/_belousov_zhabotinsky.py b/exponax/stepper/reaction/_belousov_zhabotinsky.py similarity index 95% rename from exponax/reaction/_belousov_zhabotinsky.py rename to exponax/stepper/reaction/_belousov_zhabotinsky.py index 8a35e2e..8047c24 100644 --- a/exponax/reaction/_belousov_zhabotinsky.py +++ b/exponax/stepper/reaction/_belousov_zhabotinsky.py @@ -4,9 +4,9 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from .._spectral import build_laplace_operator -from ..nonlin_fun import BaseNonlinearFun +from ..._base_stepper import BaseStepper +from ..._spectral import build_laplace_operator +from ...nonlin_fun import BaseNonlinearFun class BelousovZhabotinskyNonlinearFun(BaseNonlinearFun): diff --git a/exponax/reaction/_cahn_hilliard.py b/exponax/stepper/reaction/_cahn_hilliard.py similarity index 97% rename from exponax/reaction/_cahn_hilliard.py rename to exponax/stepper/reaction/_cahn_hilliard.py index 20580e8..27c9170 100644 --- a/exponax/reaction/_cahn_hilliard.py +++ b/exponax/stepper/reaction/_cahn_hilliard.py @@ -1,8 +1,8 @@ from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from .._spectral import build_laplace_operator -from ..nonlin_fun import BaseNonlinearFun +from ..._base_stepper import BaseStepper +from ..._spectral import build_laplace_operator +from ...nonlin_fun import BaseNonlinearFun class CahnHilliardNonlinearFun(BaseNonlinearFun): diff --git a/exponax/reaction/_fisher_kpp.py b/exponax/stepper/reaction/_fisher_kpp.py similarity index 97% rename from exponax/reaction/_fisher_kpp.py rename to exponax/stepper/reaction/_fisher_kpp.py index a4c92e6..2070947 100644 --- a/exponax/reaction/_fisher_kpp.py +++ b/exponax/stepper/reaction/_fisher_kpp.py @@ -1,8 +1,8 @@ from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from .._spectral import build_laplace_operator -from ..nonlin_fun import PolynomialNonlinearFun +from ..._base_stepper import BaseStepper +from ..._spectral import build_laplace_operator +from ...nonlin_fun import PolynomialNonlinearFun class FisherKPP(BaseStepper): diff --git a/exponax/reaction/_gray_scott.py b/exponax/stepper/reaction/_gray_scott.py similarity index 97% rename from exponax/reaction/_gray_scott.py rename to exponax/stepper/reaction/_gray_scott.py index c71274a..169edb8 100644 --- a/exponax/reaction/_gray_scott.py +++ b/exponax/stepper/reaction/_gray_scott.py @@ -1,9 +1,9 @@ import jax.numpy as jnp from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from .._spectral import build_laplace_operator -from ..nonlin_fun import BaseNonlinearFun +from ..._base_stepper import BaseStepper +from ..._spectral import build_laplace_operator +from ...nonlin_fun import BaseNonlinearFun class GrayScottNonlinearFun(BaseNonlinearFun): diff --git a/exponax/reaction/_swift_hohenberg.py b/exponax/stepper/reaction/_swift_hohenberg.py similarity index 97% rename from exponax/reaction/_swift_hohenberg.py rename to exponax/stepper/reaction/_swift_hohenberg.py index 1ffcccf..92b3e5b 100644 --- a/exponax/reaction/_swift_hohenberg.py +++ b/exponax/stepper/reaction/_swift_hohenberg.py @@ -1,8 +1,8 @@ from jaxtyping import Array, Complex -from .._base_stepper import BaseStepper -from .._spectral import build_laplace_operator -from ..nonlin_fun import PolynomialNonlinearFun +from ..._base_stepper import BaseStepper +from ..._spectral import build_laplace_operator +from ...nonlin_fun import PolynomialNonlinearFun class SwiftHohenberg(BaseStepper): diff --git a/mkdocs.yml b/mkdocs.yml index 256b8ce..d8be3c4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -99,45 +99,47 @@ nav: - Nice Features: 'examples/additional_features.ipynb' - Performance Hints: 'examples/performance_hints.ipynb' - Stepper: - - Physical: - - Linear: - - Advection: 'api/stepper/physical/linear/advection.md' - - Diffusion: 'api/stepper/physical/linear/diffusion.md' - - Advection-Diffusion: 'api/stepper/physical/linear/advection_diffusion.md' - - Dispersion: 'api/stepper/physical/linear/dispersion.md' - - Hyper-Diffusion: 'api/stepper/physical/linear/hyper_diffusion.md' - - Burgers: 'api/stepper/physical/burgers.md' - - Korteweg-de Vries: 'api/stepper/physical/kdv.md' - - Kuramoto-Sivashinsky: 'api/stepper/physical/ks.md' - - Kuramoto-Sivashinsky (Conservative): 'api/stepper/physical/ks_cons.md' - - Navier-Stokes: 'api/stepper/physical/navier_stokes.md' - - Reaction-Diffusion: - - Fisher-KPP: 'api/stepper/physical/reaction_diffusion/fisher_kpp.md' - - Allen-Cahn: 'api/stepper/physical/reaction_diffusion/allen_cahn.md' - - Cahn-Hilliard: 'api/stepper/physical/reaction_diffusion/cahn_hilliard.md' - - Swift-Hohenberg: 'api/stepper/physical/reaction_diffusion/swift_hohenberg.md' - - Gray-Scott: 'api/stepper/physical/reaction_diffusion/gray_scott.md' - - General: - - General Linear: 'api/stepper/physical/general/general_linear.md' - - General Convection: 'api/stepper/physical/general/general_convection.md' - - General Gradient Norm: 'api/stepper/physical/general/general_gradient_norm.md' - - General Polynomial: 'api/stepper/physical/general/general_polynomial.md' - - General Nonlinear: 'api/stepper/physical/general/general_nonlinear.md' - - General Vorticity Convection: 'api/stepper/physical/general/general_vorticity_convection.md' - - Poisson: 'api/stepper/physical/poisson.md' - - Normalized: - - Linear: 'api/stepper/normalized/linear.md' - - Convection: 'api/stepper/normalized/convection.md' - - Gradient Norm: 'api/stepper/normalized/gradient_norm.md' - - Polynomial: 'api/stepper/normalized/polynomial.md' - - Nonlinear: 'api/stepper/normalized/nonlinear.md' - - Vorticity Convection: 'api/stepper/normalized/vorticity_convection.md' - - Difficulty: - - Linear: 'api/stepper/difficulty/linear.md' - - Convection: 'api/stepper/difficulty/convection.md' - - Gradient Norm: 'api/stepper/difficulty/gradient_norm.md' - - Polynomial: 'api/stepper/difficulty/polynomial.md' - - Nonlinear: 'api/stepper/difficulty/nonlinear.md' + - Overview: 'api/stepper/overview.md' + - Linear: + - Advection: 'api/stepper/linear/advection.md' + - Diffusion: 'api/stepper/linear/diffusion.md' + - Advection-Diffusion: 'api/stepper/linear/advection_diffusion.md' + - Dispersion: 'api/stepper/linear/dispersion.md' + - Hyper-Diffusion: 'api/stepper/linear/hyper_diffusion.md' + - Nonlinear: + - Burgers: 'api/stepper/nonlinear/burgers.md' + - Korteweg-de Vries: 'api/stepper/nonlinear/kdv.md' + - Kuramoto-Sivashinsky: 'api/stepper/nonlinear/ks.md' + - Kuramoto-Sivashinsky (Conservative): 'api/stepper/nonlinear/ks_cons.md' + - Navier-Stokes: 'api/stepper/nonlinear/navier_stokes.md' + - Reaction: + - Fisher-KPP: 'api/stepper/reaction/fisher_kpp.md' + - Allen-Cahn: 'api/stepper/reaction/allen_cahn.md' + - Cahn-Hilliard: 'api/stepper/reaction/cahn_hilliard.md' + - Swift-Hohenberg: 'api/stepper/reaction/swift_hohenberg.md' + - Gray-Scott: 'api/stepper/reaction/gray_scott.md' + - Generic: + - Physical: + - General Linear: 'api/stepper/generic/physical/general_linear.md' + - General Convection: 'api/stepper/generic/physical/general_convection.md' + - General Gradient Norm: 'api/stepper/generic/physical/general_gradient_norm.md' + - General Polynomial: 'api/stepper/generic/physical/general_polynomial.md' + - General Nonlinear: 'api/stepper/generic/physical/general_nonlinear.md' + - General Vorticity Convection: 'api/stepper/generic/physical/general_vorticity_convection.md' + - Normalized: + - Linear: 'api/stepper/generic/normalized/normalized_linear.md' + - Convection: 'api/stepper/generic/normalized/normalized_convection.md' + - Gradient Norm: 'api/stepper/generic/normalized/normalized_gradient_norm.md' + - Polynomial: 'api/stepper/generic/normalized/normalized_polynomial.md' + - Nonlinear: 'api/stepper/generic/normalized/normalized_nonlinear.md' + - Difficulty: + - Linear: 'api/stepper/generic/difficulty/difficulty_linear.md' + - Convection: 'api/stepper/generic/difficulty/difficulty_convection.md' + - Gradient Norm: 'api/stepper/generic/difficulty/difficulty_gradient_norm.md' + - Polynomial: 'api/stepper/generic/difficulty/difficulty_polynomial.md' + - Nonlinear: 'api/stepper/generic/difficulty/difficulty_nonlinear.md' + - Additional: + - Poisson: 'api/stepper/additional/poisson.md' - Utilities: - Nonlinear Functions: - Zero: 'api/utilities/nonlin_fun/zero.md' diff --git a/tests/test_builtin_solvers.py b/tests/test_builtin_solvers.py index c9b36e5..3241ec5 100644 --- a/tests/test_builtin_solvers.py +++ b/tests/test_builtin_solvers.py @@ -21,22 +21,22 @@ def test_instantiate(): ex.stepper.KuramotoSivashinsky, ex.stepper.KuramotoSivashinskyConservative, ex.stepper.KortewegDeVries, - ex.stepper.GeneralConvectionStepper, - ex.stepper.GeneralGradientNormStepper, - ex.stepper.GeneralLinearStepper, - ex.stepper.GeneralNonlinearStepper, - ex.stepper.GeneralPolynomialStepper, + ex.stepper.generic.GeneralConvectionStepper, + ex.stepper.generic.GeneralGradientNormStepper, + ex.stepper.generic.GeneralLinearStepper, + ex.stepper.generic.GeneralNonlinearStepper, + ex.stepper.generic.GeneralPolynomialStepper, ]: simulator(num_spatial_dims, domain_extent, num_points, dt) for num_spatial_dims in [1, 2, 3]: for simulator in [ - ex.reaction.FisherKPP, - ex.reaction.AllenCahn, - ex.reaction.CahnHilliard, - ex.reaction.SwiftHohenberg, - # ex.reaction.BelousovZhabotinsky, - ex.reaction.GrayScott, + ex.stepper.reaction.FisherKPP, + ex.stepper.reaction.AllenCahn, + ex.stepper.reaction.CahnHilliard, + ex.stepper.reaction.SwiftHohenberg, + # ex.stepper.reaction.BelousovZhabotinsky, + ex.stepper.reaction.GrayScott, ]: simulator(num_spatial_dims, domain_extent, num_points, dt) @@ -51,17 +51,14 @@ def test_instantiate(): for num_spatial_dims in [1, 2, 3]: for normalized_simulator in [ - ex.normalized.NormalizedLinearStepper, - ex.normalized.NormalizedConvectionStepper, - ex.normalized.NormalizedGradientNormStepper, - ex.normalized.NormalizedPolynomialStepper, - ex.normalized.NormalizedGeneralNonlinearStepper, + ex.stepper.generic.NormalizedLinearStepper, + ex.stepper.generic.NormalizedConvectionStepper, + ex.stepper.generic.NormalizedGradientNormStepper, + ex.stepper.generic.NormalizedPolynomialStepper, + ex.stepper.generic.NormalizedNonlinearStepper, ]: normalized_simulator(num_spatial_dims, num_points) - for simulator in [ex.normalized.NormalizedVorticityConvection]: - simulator(2, num_points) - @pytest.mark.parametrize( "specific_stepper,general_stepper_coefficients", @@ -105,7 +102,7 @@ def test_specific_stepper_to_general_linear_stepper( cutoff=5, )(num_points, key=jax.random.PRNGKey(0)) - general_stepper = ex.stepper.GeneralLinearStepper( + general_stepper = ex.stepper.generic.GeneralLinearStepper( num_spatial_dims, domain_extent, num_points, @@ -193,7 +190,7 @@ def test_specific_stepper_to_general_convection_stepper( cutoff=5, )(num_points, key=jax.random.PRNGKey(0)) - general_stepper = ex.stepper.GeneralConvectionStepper( + general_stepper = ex.stepper.generic.GeneralConvectionStepper( num_spatial_dims, domain_extent, num_points, @@ -270,7 +267,7 @@ def test_specific_to_general_gradient_norm_stepper( cutoff=5, )(num_points, key=jax.random.PRNGKey(0)) - general_stepper = ex.stepper.GeneralGradientNormStepper( + general_stepper = ex.stepper.generic.GeneralGradientNormStepper( num_spatial_dims, domain_extent, num_points, @@ -309,17 +306,17 @@ def test_linear_normalized_stepper(coefficients): cutoff=5, )(num_points, key=jax.random.PRNGKey(0)) - regular_linear_stepper = ex.stepper.GeneralLinearStepper( + regular_linear_stepper = ex.stepper.generic.GeneralLinearStepper( num_spatial_dims, domain_extent, num_points, dt, coefficients=coefficients, ) - normalized_linear_stepper = ex.normalized.NormalizedLinearStepper( + normalized_linear_stepper = ex.stepper.generic.NormalizedLinearStepper( num_spatial_dims, num_points, - normalized_coefficients=ex.normalized.normalize_coefficients( + normalized_coefficients=ex.stepper.generic.normalize_coefficients( coefficients, domain_extent=domain_extent, dt=dt, @@ -351,15 +348,15 @@ def test_nonlinear_normalized_stepper(): diffusivity=diffusivity, convection_scale=convection_scale, ) - normalized_burgers_stepper = ex.normalized.NormalizedConvectionStepper( + normalized_burgers_stepper = ex.stepper.generic.NormalizedConvectionStepper( num_spatial_dims, num_points, - normalized_coefficients=ex.normalized.normalize_coefficients( + normalized_coefficients=ex.stepper.generic.normalize_coefficients( [0.0, 0.0, diffusivity], domain_extent=domain_extent, dt=dt, ), - normalized_convection_scale=ex.normalized.normalize_convection_scale( + normalized_convection_scale=ex.stepper.generic.normalize_convection_scale( convection_scale, domain_extent=domain_extent, dt=dt,