Use correct tuple typing (#27)

Ceyron · Sep 3, 2024 · 3cf1399 · 3cf1399
1 parent bea4079
commit 3cf1399
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 11 deletions.
diff --git a/exponax/nonlin_fun/_general_nonlinear.py b/exponax/nonlin_fun/_general_nonlinear.py
@@ -18,7 +18,7 @@ def __init__(
         *,
         derivative_operator: Complex[Array, "D ... (N//2)+1"],
         dealiasing_fraction: float,
-        scale_list: tuple[float, ...] = [0.0, -1.0, 0.0],
+        scale_list: tuple[float, float, float] = (0.0, -1.0, 0.0),
         zero_mode_fix: bool = True,
     ):
         """

diff --git a/exponax/stepper/generic/_nonlinear.py b/exponax/stepper/generic/_nonlinear.py
@@ -11,7 +11,7 @@
 
 class GeneralNonlinearStepper(BaseStepper):
     coefficients_linear: tuple[float, ...]
-    coefficients_nonlinear: tuple[float, 3]
+    coefficients_nonlinear: tuple[float, float, float]
     dealiasing_fraction: float
 
     def __init__(
@@ -22,7 +22,7 @@ def __init__(
         dt: float,
         *,
         coefficients_linear: tuple[float, ...] = (0.0, 0.0, 0.01),
-        coefficients_nonlinear: tuple[float, 3] = (0.0, -1.0, 0.0),
+        coefficients_nonlinear: tuple[float, float, float] = (0.0, -1.0, 0.0),
         order=2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -156,15 +156,19 @@ def _build_nonlinear_fun(
 
 class NormalizedNonlinearStepper(GeneralNonlinearStepper):
     normalized_coefficients_linear: tuple[float, ...]
-    normalized_coefficients_nonlinear: tuple[float, ...]
+    normalized_coefficients_nonlinear: 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, ...] = (0.0, -1.0 * 0.1, 0.0),
+        normalized_coefficients_nonlinear: tuple[float, float, float] = (
+            0.0,
+            -1.0 * 0.1,
+            0.0,
+        ),
         order=2,
         dealiasing_fraction: float = 2 / 3,
         num_circle_points: int = 16,
@@ -271,7 +275,7 @@ def __init__(
 
 class DifficultyNonlinearStepper(NormalizedNonlinearStepper):
     linear_difficulties: tuple[float, ...]
-    nonlinear_difficulties: tuple[float, ...]
+    nonlinear_difficulties: tuple[float, float, float]
 
     def __init__(
         self,
@@ -283,7 +287,7 @@ def __init__(
             0.0,
             0.1 * 0.1 / 1.0 * 48**2 * 2,
         ),
-        nonlinear_difficulties: tuple[float, ...] = (
+        nonlinear_difficulties: tuple[float, float, float] = (
             0.0,
             -1.0 * 0.1 / 1.0 * 48,
             0.0,

diff --git a/exponax/stepper/generic/_utils.py b/exponax/stepper/generic/_utils.py
@@ -448,7 +448,7 @@ def extract_normalized_gradient_norm_scale_from_difficulty(
 
 
 def reduce_normalized_nonlinear_scales_to_difficulty(
-    normalized_nonlinear_scales: tuple[float],
+    normalized_nonlinear_scales: tuple[float, float, float],
     *,
     num_spatial_dims: int,
     num_points: int,
@@ -501,7 +501,7 @@ def reduce_normalized_nonlinear_scales_to_difficulty(
 
 
 def extract_normalized_nonlinear_scales_from_difficulty(
-    nonlinear_difficulties: tuple[float],
+    nonlinear_difficulties: tuple[float, float, float],
     *,
     num_spatial_dims: int,
     num_points: int,

diff --git a/exponax/stepper/reaction/_belousov_zhabotinsky.py b/exponax/stepper/reaction/_belousov_zhabotinsky.py
@@ -47,7 +47,7 @@ def __call__(
 
 
 class BelousovZhabotinsky(BaseStepper):
-    diffusivities: tuple[float, 3]
+    diffusivities: tuple[float, float, float]
     dealiasing_fraction: float
 
     def __init__(
@@ -57,7 +57,7 @@ def __init__(
         num_points: int,
         dt: float,
         *,
-        diffusivities: tuple[float, 3] = (1e-5, 2e-5, 1e-5),
+        diffusivities: tuple[float, float, float] = (1e-5, 2e-5, 1e-5),
         order: int = 2,
         dealiasing_fraction: float = 1
         / 2,  # Needs lower value due to cubic nonlinearity