From ac7093072414f6eb4fbcd19c961c253ae0235f1c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20K=C3=B6hler?=
 <27728103+Ceyron@users.noreply.github.com>
Date: Tue, 22 Oct 2024 14:38:19 +0200
Subject: [PATCH] Consistent Generic Steppers (#52)

* Use consistent attribute names

* Forward API changes to tests

* Forward changes to qualitative rollout
---
 exponax/stepper/generic/_convection.py        | 22 +++---
 exponax/stepper/generic/_gradient_norm.py     | 20 +++---
 exponax/stepper/generic/_linear.py            | 30 ++++----
 exponax/stepper/generic/_nonlinear.py         | 50 +++++++-------
 exponax/stepper/generic/_polynomial.py        | 69 ++++++++++---------
 .../stepper/generic/_vorticity_convection.py  | 10 +--
 tests/test_builtin_solvers.py                 | 12 ++--
 validation/qualitative_rollouts.py            |  4 +-
 8 files changed, 109 insertions(+), 108 deletions(-)

diff --git a/exponax/stepper/generic/_convection.py b/exponax/stepper/generic/_convection.py
index b4241f4..82b9e29 100644
--- a/exponax/stepper/generic/_convection.py
+++ b/exponax/stepper/generic/_convection.py
@@ -10,7 +10,7 @@
 
 
 class GeneralConvectionStepper(BaseStepper):
-    coefficients: tuple[float, ...]
+    linear_coefficients: tuple[float, ...]
     convection_scale: float
     dealiasing_fraction: float
     single_channel: bool
@@ -23,7 +23,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        coefficients: tuple[float, ...] = (0.0, 0.0, 0.01),
+        linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01),
         convection_scale: float = 1.0,
         single_channel: bool = False,
         conservative: bool = False,
@@ -74,7 +74,7 @@ def __init__(
             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`
+        - `linear_coefficients` (keyword-only): The list of coefficients `a_j`
             corresponding to the derivatives. The length of this tuple
             represents the highest occuring derivative. The default value `(0.0,
             0.0, 0.01)` corresponds to the Burgers equation (because of the
@@ -103,7 +103,7 @@ def __init__(
             coefficients of the exponential time differencing Runge Kutta
             method. Default: 1.0.
         """
-        self.coefficients = coefficients
+        self.linear_coefficients = linear_coefficients
         self.convection_scale = convection_scale
         self.single_channel = single_channel
         self.dealiasing_fraction = dealiasing_fraction
@@ -136,7 +136,7 @@ def _build_linear_operator(
                 axis=0,
                 keepdims=True,
             )
-            for i, c in enumerate(self.coefficients)
+            for i, c in enumerate(self.linear_coefficients)
         )
         return linear_operator
 
@@ -156,7 +156,7 @@ def _build_nonlinear_fun(
 
 
 class NormalizedConvectionStepper(GeneralConvectionStepper):
-    normalized_coefficients: tuple[float, ...]
+    normalized_linear_coefficients: tuple[float, ...]
     normalized_convection_scale: float
 
     def __init__(
@@ -164,7 +164,7 @@ def __init__(
         num_spatial_dims: int,
         num_points: int,
         *,
-        normalized_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1),
+        normalized_linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01 * 0.1),
         normalized_convection_scale: float = 1.0 * 0.1,
         single_channel: bool = False,
         conservative: bool = False,
@@ -205,7 +205,7 @@ def __init__(
             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 list of coefficients
+        - `normalized_linear_coefficients`: The list of coefficients
             `α_j` corresponding to the derivatives. The length of this tuple
             represents the highest occuring derivative. The default value `(0.0,
             0.0, 0.01)` corresponds to the Burgers equation (because of the
@@ -235,14 +235,14 @@ def __init__(
             coefficients of the exponential time differencing Runge Kutta
             method. Default: 1.0.
         """
-        self.normalized_coefficients = normalized_coefficients
+        self.normalized_linear_coefficients = normalized_linear_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,
+            linear_coefficients=normalized_linear_coefficients,
             convection_scale=normalized_convection_scale,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
@@ -364,7 +364,7 @@ def __init__(
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
-            normalized_coefficients=normalized_coefficients,
+            normalized_linear_coefficients=normalized_coefficients,
             normalized_convection_scale=normalized_convection_scale,
             single_channel=single_channel,
             order=order,
diff --git a/exponax/stepper/generic/_gradient_norm.py b/exponax/stepper/generic/_gradient_norm.py
index cd6e299..3df31a8 100644
--- a/exponax/stepper/generic/_gradient_norm.py
+++ b/exponax/stepper/generic/_gradient_norm.py
@@ -10,7 +10,7 @@
 
 
 class GeneralGradientNormStepper(BaseStepper):
-    coefficients: tuple[float, ...]
+    linear_coefficients: tuple[float, ...]
     gradient_norm_scale: float
     dealiasing_fraction: float
 
@@ -21,7 +21,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        coefficients: tuple[float, ...] = (0.0, 0.0, -1.0, 0.0, -1.0),
+        linear_coefficients: tuple[float, ...] = (0.0, 0.0, -1.0, 0.0, -1.0),
         gradient_norm_scale: float = 1.0,
         order=2,
         dealiasing_fraction: float = 2 / 3,
@@ -66,7 +66,7 @@ def __init__(
             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`
+        - `linear_coefficients` (keyword-only): The list of coefficients `a_j`
             corresponding to the derivatives. The length of this tuple
             represents the highest occuring derivative. The default value `(0.0,
             0.0, -1.0, 0.0, -1.0)` corresponds to the Kuramoto- Sivashinsky
@@ -89,7 +89,7 @@ def __init__(
             coefficients of the exponential time differencing Runge Kutta
             method. Default: 1.0.
         """
-        self.coefficients = coefficients
+        self.linear_coefficients = linear_coefficients
         self.gradient_norm_scale = gradient_norm_scale
         self.dealiasing_fraction = dealiasing_fraction
         super().__init__(
@@ -113,7 +113,7 @@ def _build_linear_operator(
                 axis=0,
                 keepdims=True,
             )
-            for i, c in enumerate(self.coefficients)
+            for i, c in enumerate(self.linear_coefficients)
         )
         return linear_operator
 
@@ -132,7 +132,7 @@ def _build_nonlinear_fun(
 
 
 class NormalizedGradientNormStepper(GeneralGradientNormStepper):
-    normalized_coefficients: tuple[float, ...]
+    normalized_linear_coefficients: tuple[float, ...]
     normalized_gradient_norm_scale: float
 
     def __init__(
@@ -140,7 +140,7 @@ def __init__(
         num_spatial_dims: int,
         num_points: int,
         *,
-        normalized_coefficients: tuple[float, ...] = (
+        normalized_linear_coefficients: tuple[float, ...] = (
             0.0,
             0.0,
             -1.0 * 0.1 / (60.0**2),
@@ -217,14 +217,14 @@ def __init__(
             coefficients of the exponential time differencing Runge Kutta
             method. Default: 1.0.
         """
-        self.normalized_coefficients = normalized_coefficients
+        self.normalized_linear_coefficients = normalized_linear_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,
+            linear_coefficients=normalized_linear_coefficients,
             gradient_norm_scale=normalized_gradient_norm_scale,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
@@ -339,7 +339,7 @@ def __init__(
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
-            normalized_coefficients=normalized_coefficients,
+            normalized_linear_coefficients=normalized_coefficients,
             normalized_gradient_norm_scale=normalized_gradient_norm_scale,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
diff --git a/exponax/stepper/generic/_linear.py b/exponax/stepper/generic/_linear.py
index 2c9fc15..8a35dca 100644
--- a/exponax/stepper/generic/_linear.py
+++ b/exponax/stepper/generic/_linear.py
@@ -11,7 +11,7 @@
 
 
 class GeneralLinearStepper(BaseStepper):
-    coefficients: tuple[float, ...]
+    linear_coefficients: tuple[float, ...]
 
     def __init__(
         self,
@@ -20,7 +20,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        coefficients: tuple[float, ...] = (0.0, -0.1, 0.01),
+        linear_coefficients: tuple[float, ...] = (0.0, -0.1, 0.01),
     ):
         """
         General timestepper for a d-dimensional (`d ∈ {1, 2, 3}`) linear
@@ -67,7 +67,7 @@ def __init__(
                 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`
+            - `linear_coefficients` (keyword-only): The list of coefficients `a_j`
                 corresponding to the derivatives. Default: `[0.0, -0.1, 0.01]`.
 
         **Notes:**
@@ -137,7 +137,7 @@ def __init__(
                 the function [`exponax.stepper.generic.normalize_coefficients`][] to
                 obtain the normalized coefficients.
         """
-        self.coefficients = coefficients
+        self.linear_coefficients = linear_coefficients
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             domain_extent=domain_extent,
@@ -157,7 +157,7 @@ def _build_linear_operator(
                 axis=0,
                 keepdims=True,
             )
-            for i, c in enumerate(self.coefficients)
+            for i, c in enumerate(self.linear_coefficients)
         )
         return linear_operator
 
@@ -172,14 +172,14 @@ def _build_nonlinear_fun(
 
 
 class NormalizedLinearStepper(GeneralLinearStepper):
-    normalized_coefficients: tuple[float, ...]
+    normalized_linear_coefficients: tuple[float, ...]
 
     def __init__(
         self,
         num_spatial_dims: int,
         num_points: int,
         *,
-        normalized_coefficients: tuple[float, ...] = (0.0, -0.5, 0.01),
+        normalized_linear_coefficients: tuple[float, ...] = (0.0, -0.5, 0.01),
     ):
         """
         Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic
@@ -218,25 +218,25 @@ def __init__(
             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
+        self.normalized_linear_coefficients = normalized_linear_coefficients
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             domain_extent=1.0,
             num_points=num_points,
             dt=1.0,
-            coefficients=normalized_coefficients,
+            linear_coefficients=normalized_linear_coefficients,
         )
 
 
 class DifficultyLinearStepper(NormalizedLinearStepper):
-    difficulties: tuple[float, ...]
+    linear_difficulties: tuple[float, ...]
 
     def __init__(
         self,
         num_spatial_dims: int = 1,
         num_points: int = 48,
         *,
-        difficulties: tuple[float, ...] = (0.0, -2.0),
+        linear_difficulties: tuple[float, ...] = (0.0, -2.0),
     ):
         """
         Timestepper for d-dimensional (`d ∈ {1, 2, 3}`) linear PDEs on periodic
@@ -275,9 +275,9 @@ def __init__(
             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
+        self.linear_difficulties = linear_difficulties
         normalized_coefficients = extract_normalized_coefficients_from_difficulty(
-            difficulties,
+            linear_difficulties,
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
         )
@@ -285,7 +285,7 @@ def __init__(
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
-            normalized_coefficients=normalized_coefficients,
+            normalized_linear_coefficients=normalized_coefficients,
         )
 
 
@@ -318,7 +318,7 @@ def __init__(
         """
         difficulties = (0.0,) * (order) + (difficulty,)
         super().__init__(
-            difficulties=difficulties,
+            linear_difficulties=difficulties,
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
         )
diff --git a/exponax/stepper/generic/_nonlinear.py b/exponax/stepper/generic/_nonlinear.py
index 5c3e875..45fd76c 100644
--- a/exponax/stepper/generic/_nonlinear.py
+++ b/exponax/stepper/generic/_nonlinear.py
@@ -10,8 +10,8 @@
 
 
 class GeneralNonlinearStepper(BaseStepper):
-    coefficients_linear: tuple[float, ...]
-    coefficients_nonlinear: tuple[float, float, float]
+    linear_coefficients: tuple[float, ...]
+    nonlinear_coefficients: tuple[float, float, float]
     dealiasing_fraction: float
 
     def __init__(
@@ -21,8 +21,8 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.01),
-        coefficients_nonlinear: tuple[float, float, float] = (0.0, -1.0, 0.0),
+        linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.01),
+        nonlinear_coefficients: tuple[float, float, float] = (0.0, -1.0, 0.0),
         order=2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -83,12 +83,12 @@ def __init__(
             dimension is the same. Hence, the total number of degrees of freedom
             is `Nᵈ`.
         - `dt`: The timestep size `Δt` between two consecutive states.
-        - `coefficients_linear`: The list of coefficients `a_j` corresponding to
+        - `linear_coefficients`: The list of coefficients `a_j` corresponding to
             the derivatives. The length of this tuple represents the highest
             occuring derivative. The default value `(0.0, 0.0, 0.01)` together
-            with the default `coefficients_nonlinear` corresponds to the Burgers
+            with the default `nonlinear_coefficients` corresponds to the Burgers
             equation.
-        - `coefficients_nonlinear`: The list of coefficients `b₀`, `b₁`, `b₂`
+        - `nonlinear_coefficients`: The list of coefficients `b₀`, `b₁`, `b₂`
             (in this order). The default value `(0.0, -1.0, 0.0)` corresponds to
             a (single-channel) convection nonlinearity scaled with `1.0`. Note
             that all nonlinear contributions are considered to be on the
@@ -107,12 +107,12 @@ def __init__(
         - `circle_radius`: The radius of the contour used to compute the
             coefficients of the exponential time differencing Runge Kutta method.
         """
-        if len(coefficients_nonlinear) != 3:
+        if len(nonlinear_coefficients) != 3:
             raise ValueError(
                 "The nonlinear coefficients list must have exactly 3 elements"
             )
-        self.coefficients_linear = coefficients_linear
-        self.coefficients_nonlinear = coefficients_nonlinear
+        self.linear_coefficients = linear_coefficients
+        self.nonlinear_coefficients = nonlinear_coefficients
         self.dealiasing_fraction = dealiasing_fraction
 
         super().__init__(
@@ -136,7 +136,7 @@ def _build_linear_operator(
                 axis=0,
                 keepdims=True,
             )
-            for i, c in enumerate(self.coefficients_linear)
+            for i, c in enumerate(self.linear_coefficients)
         )
         return linear_operator
 
@@ -149,22 +149,22 @@ def _build_nonlinear_fun(
             self.num_points,
             derivative_operator=derivative_operator,
             dealiasing_fraction=self.dealiasing_fraction,
-            scale_list=self.coefficients_nonlinear,
+            scale_list=self.nonlinear_coefficients,
             zero_mode_fix=True,  # ToDo: check this
         )
 
 
 class NormalizedNonlinearStepper(GeneralNonlinearStepper):
-    normalized_coefficients_linear: tuple[float, ...]
-    normalized_coefficients_nonlinear: tuple[float, float, float]
+    normalized_linear_coefficients: tuple[float, ...]
+    normalized_nonlinear_coefficients: tuple[float, float, 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, float, float] = (
+        normalized_linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.1 * 0.1),
+        normalized_nonlinear_coefficients: tuple[float, float, float] = (
             0.0,
             -1.0 * 0.1,
             0.0,
@@ -227,13 +227,13 @@ def __init__(
             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_linear`: The list of coefficients `αⱼ`
+        - `normalized_linear_coefficients`: The list of coefficients `αⱼ`
             corresponding to the linear derivatives. The length of this tuple
             represents the highest occuring derivative. The default value `(0.0,
             0.0, 0.1 * 0.1)` together with the default
-            `normalized_coefficients_nonlinear` corresponds to the Burgers
+            `normalized_nonlinear_coefficients` corresponds to the Burgers
             equation (in single-channel mode).
-        - `normalized_coefficients_nonlinear`: The list of coefficients `β₀`,
+        - `normalized_nonlinear_coefficients`: The list of coefficients `β₀`,
             `β₁`, and `β₂` (in this order) corresponding to the quadratic,
             (single-channel) convection, and gradient norm nonlinearity,
             respectively. The default value `(0.0, -1.0 * 0.1, 0.0)` corresponds
@@ -256,16 +256,16 @@ def __init__(
             coefficients of the exponential time differencing Runge Kutta method.
         """
 
-        self.normalized_coefficients_linear = normalized_coefficients_linear
-        self.normalized_coefficients_nonlinear = normalized_coefficients_nonlinear
+        self.normalized_linear_coefficients = normalized_linear_coefficients
+        self.normalized_nonlinear_coefficients = normalized_nonlinear_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,
-            coefficients_linear=normalized_coefficients_linear,
-            coefficients_nonlinear=normalized_coefficients_nonlinear,
+            linear_coefficients=normalized_linear_coefficients,
+            nonlinear_coefficients=normalized_nonlinear_coefficients,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
             num_circle_points=num_circle_points,
@@ -394,8 +394,8 @@ def __init__(
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
-            normalized_coefficients_linear=normalized_coefficients_linear,
-            normalized_coefficients_nonlinear=normalized_coefficients_nonlinear,
+            normalized_linear_coefficients=normalized_coefficients_linear,
+            normalized_nonlinear_coefficients=normalized_coefficients_nonlinear,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
             num_circle_points=num_circle_points,
diff --git a/exponax/stepper/generic/_polynomial.py b/exponax/stepper/generic/_polynomial.py
index 2780569..4523536 100644
--- a/exponax/stepper/generic/_polynomial.py
+++ b/exponax/stepper/generic/_polynomial.py
@@ -7,8 +7,8 @@
 
 
 class GeneralPolynomialStepper(BaseStepper):
-    coefficients: tuple[float, ...]
-    polynomial_scales: tuple[float, ...]
+    linear_coefficients: tuple[float, ...]
+    polynomial_coefficients: tuple[float, ...]
     dealiasing_fraction: float
 
     def __init__(
@@ -18,8 +18,8 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        coefficients: tuple[float, ...] = (10.0, 0.0, 1.0),
-        polynomial_scales: tuple[float, ...] = (0.0, 0.0, -10.0),
+        linear_coefficients: tuple[float, ...] = (10.0, 0.0, 1.0),
+        polynomial_coefficients: tuple[float, ...] = (0.0, 0.0, -10.0),
         order=2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -78,15 +78,15 @@ def __init__(
             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`: The list of coefficients `a_j` corresponding to the
+        - `linear_coefficients`: The list of coefficients `a_j` corresponding to the
             derivatives. The length of this tuple represents the highest
             occuring derivative. The default value `(10.0, 0.0, 0.01)` in
-            combination with the default `polynomial_scales` corresponds to the
+            combination with the default `polynomial_coefficients` corresponds to the
             Fisher-KPP equation.
-        - `polynomial_scales`: The list of scales `pₖ` corresponding to the
+        - `polynomial_coefficients`: The list of scales `pₖ` corresponding to the
             polynomial contributions. The length of this tuple represents the
             highest occuring polynomial. The default value `(0.0, 0.0, 10.0)` in
-            combination with the default `coefficients` corresponds to the
+            combination with the default `linear_coefficients` corresponds to the
             Fisher-KPP equation.
         - `order`: The order of the Exponential Time Differencing Runge
             Kutta method. Must be one of {0, 1, 2, 3, 4}. The option `0` only
@@ -105,8 +105,8 @@ def __init__(
             coefficients of the exponential time differencing Runge Kutta
             method.
         """
-        self.coefficients = coefficients
-        self.polynomial_scales = polynomial_scales
+        self.linear_coefficients = linear_coefficients
+        self.polynomial_coefficients = polynomial_coefficients
         self.dealiasing_fraction = dealiasing_fraction
 
         super().__init__(
@@ -130,7 +130,7 @@ def _build_linear_operator(
                 axis=0,
                 keepdims=True,
             )
-            for i, c in enumerate(self.coefficients)
+            for i, c in enumerate(self.linear_coefficients)
         )
         return linear_operator
 
@@ -142,25 +142,25 @@ def _build_nonlinear_fun(
             self.num_spatial_dims,
             self.num_points,
             dealiasing_fraction=self.dealiasing_fraction,
-            coefficients=self.polynomial_scales,
+            coefficients=self.polynomial_coefficients,
         )
 
 
 class NormalizedPolynomialStepper(GeneralPolynomialStepper):
-    normalized_coefficients: tuple[float, ...]
-    normalized_polynomial_scales: tuple[float, ...]
+    normalized_linear_coefficients: tuple[float, ...]
+    normalized_polynomial_coefficients: tuple[float, ...]
 
     def __init__(
         self,
         num_spatial_dims: int,
         num_points: int,
         *,
-        normalized_coefficients: tuple[float, ...] = (
+        normalized_linear_coefficients: tuple[float, ...] = (
             10.0 * 0.001 / (10.0**0),
             0.0,
             1.0 * 0.001 / (10.0**2),
         ),
-        normalized_polynomial_scales: tuple[float, ...] = (
+        normalized_polynomial_coefficients: tuple[float, ...] = (
             0.0,
             0.0,
             -10.0 * 0.001,
@@ -190,19 +190,19 @@ def __init__(
             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 list of coefficients `α_j` corresponding
-            to the derivatives. The length of this tuple represents the highest
-            occuring derivative. The default value corresponds to the Fisher-KPP
-            equation.
-        - `normalized_polynomial_scales`: The list of scales `βₖ` corresponding
-            to the polynomial contributions. The length of this tuple represents
-            the highest occuring polynomial. The default value corresponds to the
-            Fisher-KPP equation.
+        - `normalized_linear_coefficients`: The list of coefficients `α_j`
+            corresponding to the derivatives. The length of this tuple
+            represents the highest occuring derivative. The default value
+            corresponds to the Fisher-KPP equation.
+        - `normalized_polynomial_coefficients`: The list of scales `βₖ`
+            corresponding to the polynomial contributions. The length of this
+            tuple represents the highest occuring polynomial. The default value
+            corresponds to the Fisher-KPP equation.
         - `order`: The order of the Exponential Time Differencing Runge Kutta
             method. Must be one of {0, 1, 2, 3, 4}. The option `0` only solves
-            the linear part of the equation. Use higher values for higher accuracy
-            and stability. The default choice of `2` is a good compromise for
-            single precision floats.
+            the linear part of the equation. Use higher values for higher
+            accuracy and stability. The default choice of `2` is a good
+            compromise for single precision floats.
         - `dealiasing_fraction`: The fraction of the wavenumbers to keep before
             evaluating the nonlinearity. The default 2/3 corresponds to Orszag's
             2/3 rule which is sufficient if the highest occuring polynomial is
@@ -212,18 +212,19 @@ def __init__(
             integral method to compute the coefficients of the exponential time
             differencing Runge Kutta method.
         - `circle_radius`: The radius of the contour used to compute the
-            coefficients of the exponential time differencing Runge Kutta method.
+            coefficients of the exponential time differencing Runge Kutta
+            method.
         """
-        self.normalized_coefficients = normalized_coefficients
-        self.normalized_polynomial_scales = normalized_polynomial_scales
+        self.normalized_linear_coefficients = normalized_linear_coefficients
+        self.normalized_polynomial_coefficients = normalized_polynomial_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,
-            coefficients=normalized_coefficients,
-            polynomial_scales=normalized_polynomial_scales,
+            linear_coefficients=normalized_linear_coefficients,
+            polynomial_coefficients=normalized_polynomial_coefficients,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
             num_circle_points=num_circle_points,
@@ -328,8 +329,8 @@ def __init__(
         super().__init__(
             num_spatial_dims=num_spatial_dims,
             num_points=num_points,
-            normalized_coefficients=normalized_coefficients,
-            normalized_polynomial_scales=normalized_polynomial_scales,
+            normalized_linear_coefficients=normalized_coefficients,
+            normalized_polynomial_coefficients=normalized_polynomial_scales,
             order=order,
             dealiasing_fraction=dealiasing_fraction,
             num_circle_points=num_circle_points,
diff --git a/exponax/stepper/generic/_vorticity_convection.py b/exponax/stepper/generic/_vorticity_convection.py
index f75b2c4..d4d83c3 100644
--- a/exponax/stepper/generic/_vorticity_convection.py
+++ b/exponax/stepper/generic/_vorticity_convection.py
@@ -7,7 +7,7 @@
 
 class GeneralVorticityConvectionStepper(BaseStepper):
     vorticity_convection_scale: float
-    coefficients: tuple[float, ...]
+    linear_coefficients: tuple[float, ...]
     injection_mode: int
     injection_scale: float
     dealiasing_fraction: float
@@ -20,7 +20,7 @@ def __init__(
         dt: float,
         *,
         vorticity_convection_scale: float = 1.0,
-        coefficients: tuple[float, ...] = (0.0, 0.0, 0.001),
+        linear_coefficients: tuple[float, ...] = (0.0, 0.0, 0.001),
         injection_mode: int = 4,
         injection_scale: float = 0.0,
         order: int = 2,
@@ -61,7 +61,7 @@ def __init__(
             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`: The list of coefficients `a_j`
+        - `linear_coefficients`: The list of coefficients `a_j`
             corresponding to the derivatives. The length of this tuple
             represents the highest occuring derivative. The default value `(0.0,
             0.0, 0.001)` corresponds to pure regular diffusion.
@@ -89,7 +89,7 @@ def __init__(
         if num_spatial_dims != 2:
             raise ValueError(f"Expected num_spatial_dims = 2, got {num_spatial_dims}.")
         self.vorticity_convection_scale = vorticity_convection_scale
-        self.coefficients = coefficients
+        self.linear_coefficients = linear_coefficients
         self.injection_mode = injection_mode
         self.injection_scale = injection_scale
         self.dealiasing_fraction = dealiasing_fraction
@@ -114,7 +114,7 @@ def _build_linear_operator(
                 axis=0,
                 keepdims=True,
             )
-            for i, c in enumerate(self.coefficients)
+            for i, c in enumerate(self.linear_coefficients)
         )
         return linear_operator
 
diff --git a/tests/test_builtin_solvers.py b/tests/test_builtin_solvers.py
index 1c1acd4..9960920 100644
--- a/tests/test_builtin_solvers.py
+++ b/tests/test_builtin_solvers.py
@@ -107,7 +107,7 @@ def test_specific_stepper_to_general_linear_stepper(
         domain_extent,
         num_points,
         dt,
-        coefficients=general_stepper_coefficients,
+        linear_coefficients=general_stepper_coefficients,
     )
 
     specific_pred = specific_stepper(u_0)
@@ -204,7 +204,7 @@ def test_specific_stepper_to_general_convection_stepper(
         domain_extent,
         num_points,
         dt,
-        coefficients=general_stepper_coefficients,
+        linear_coefficients=general_stepper_coefficients,
         convection_scale=general_stepper_scale,
         conservative=conservative,
     )
@@ -282,7 +282,7 @@ def test_specific_to_general_gradient_norm_stepper(
         domain_extent,
         num_points,
         dt,
-        coefficients=general_stepper_coefficients,
+        linear_coefficients=general_stepper_coefficients,
         gradient_norm_scale=general_stepper_scale,
     )
 
@@ -321,12 +321,12 @@ def test_linear_normalized_stepper(coefficients):
         domain_extent,
         num_points,
         dt,
-        coefficients=coefficients,
+        linear_coefficients=coefficients,
     )
     normalized_linear_stepper = ex.stepper.generic.NormalizedLinearStepper(
         num_spatial_dims,
         num_points,
-        normalized_coefficients=ex.stepper.generic.normalize_coefficients(
+        normalized_linear_coefficients=ex.stepper.generic.normalize_coefficients(
             coefficients,
             domain_extent=domain_extent,
             dt=dt,
@@ -361,7 +361,7 @@ def test_nonlinear_normalized_stepper():
     normalized_burgers_stepper = ex.stepper.generic.NormalizedConvectionStepper(
         num_spatial_dims,
         num_points,
-        normalized_coefficients=ex.stepper.generic.normalize_coefficients(
+        normalized_linear_coefficients=ex.stepper.generic.normalize_coefficients(
             [0.0, 0.0, diffusivity],
             domain_extent=domain_extent,
             dt=dt,
diff --git a/validation/qualitative_rollouts.py b/validation/qualitative_rollouts.py
index 3e89dd5..8460548 100644
--- a/validation/qualitative_rollouts.py
+++ b/validation/qualitative_rollouts.py
@@ -66,7 +66,7 @@
             3.0,
             110,
             0.01,
-            coefficients=[0.0, 0.0, 0.1, 0.0001],
+            linear_coefficients=[0.0, 0.0, 0.1, 0.0001],
         ),
         "dispersion_diffusion",
         ex.ic.RandomTruncatedFourierSeries(1, cutoff=5),
@@ -80,7 +80,7 @@
             3.0,
             110,
             0.01,
-            coefficients=[0.0, 0.0, 0.0, 0.0001, -0.001],
+            linear_coefficients=[0.0, 0.0, 0.0, 0.0001, -0.001],
         ),
         "dispersion_hyper_diffusion",
         ex.ic.RandomTruncatedFourierSeries(1, cutoff=5),