From 3cf13991882db9a0f868d516a3c95498d78b7c40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Felix=20K=C3=B6hler?=
 <27728103+Ceyron@users.noreply.github.com>
Date: Tue, 3 Sep 2024 15:10:30 +0200
Subject: [PATCH] Use correct tuple typing (#27)

---
 exponax/nonlin_fun/_general_nonlinear.py         |  2 +-
 exponax/stepper/generic/_nonlinear.py            | 16 ++++++++++------
 exponax/stepper/generic/_utils.py                |  4 ++--
 .../stepper/reaction/_belousov_zhabotinsky.py    |  4 ++--
 4 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/exponax/nonlin_fun/_general_nonlinear.py b/exponax/nonlin_fun/_general_nonlinear.py
index c266572..8b09cb9 100644
--- 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
index 60315be..5c3e875 100644
--- 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,7 +156,7 @@ 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,
@@ -164,7 +164,11 @@ def __init__(
         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
index 4263913..1e90c76 100644
--- 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
index 8047c24..7fb8ca6 100644
--- 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