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