Replace chex.Array with array_typing.Array.

chex.Array includes things like np.bool_ and np.number that are not what functions accepting either NumPy or JAX arrays intend.

PiperOrigin-RevId: 802070254

Author: sbodenstein

torax/_src/array_typing.py

@@ -32,6 +32,8 @@
 FloatVector: TypeAlias = jt.Float[Array, "_"]
 BoolVector: TypeAlias = jt.Bool[Array, "_"]
 FloatVectorCell: TypeAlias = jt.Float[Array, "rhon"]
+FloatVectorCellPlusBoundaries: TypeAlias = jt.Float[Array, "rhon+2"]
+FloatMatrixCell: TypeAlias = jt.Float[Array, "rhon rhon"]
 FloatVectorFace: TypeAlias = jt.Float[Array, "rhon+1"]
 
 

torax/_src/fvm/diffusion_terms.py

@@ -19,13 +19,14 @@
 
 import chex
 from jax import numpy as jnp
+from torax._src import array_typing
 from torax._src import math_utils
 from torax._src.fvm import cell_variable
 
 
 def make_diffusion_terms(
-    d_face: chex.Array, var: cell_variable.CellVariable
-) -> tuple[chex.Array, chex.Array]:
+    d_face: array_typing.FloatVectorFace, var: cell_variable.CellVariable
+) -> tuple[array_typing.FloatMatrixCell, array_typing.FloatVectorCell]:
   """Makes the terms of the matrix equation derived from the diffusion term.
 
   The diffusion term is of the form

torax/_src/fvm/newton_raphson_solve_block.py

@@ -19,7 +19,7 @@
 
 import functools
 from typing import Final
-import chex
+from torax._src import array_typing
 from torax._src import jax_utils
 from torax._src import physics_models as physics_models_lib
 from torax._src import state as state_module
@@ -52,7 +52,7 @@
     ],
 )
 def newton_raphson_solve_block(
-    dt: chex.Array,
+    dt: array_typing.FloatScalar,
     dynamic_runtime_params_slice_t: runtime_params_slice.RuntimeParams,
     dynamic_runtime_params_slice_t_plus_dt: runtime_params_slice.RuntimeParams,
     geo_t: geometry.Geometry,

torax/_src/geometry/geometry.py

@@ -16,15 +16,17 @@
 from collections.abc import Sequence
 import dataclasses
 import enum
-
 import chex
 import jax
 import jax.numpy as jnp
 import numpy as np
+from torax._src import array_typing
 from torax._src.torax_pydantic import torax_pydantic
 
 
-def face_to_cell(face: chex.Array) -> chex.Array:
+def face_to_cell(
+    face: array_typing.FloatVectorFace,
+) -> array_typing.FloatVectorCell:
   """Infers cell values corresponding to a vector of face values.
 
   Simply a linear interpolation between face values.
@@ -170,45 +172,45 @@ class Geometry:
 
   geometry_type: GeometryType
   torax_mesh: torax_pydantic.Grid1D
-  Phi: chex.Array
-  Phi_face: chex.Array
-  R_major: chex.Array
-  a_minor: chex.Array
-  B_0: chex.Array
-  volume: chex.Array
-  volume_face: chex.Array
-  area: chex.Array
-  area_face: chex.Array
-  vpr: chex.Array
-  vpr_face: chex.Array
-  spr: chex.Array
-  spr_face: chex.Array
-  delta_face: chex.Array
-  elongation: chex.Array
-  elongation_face: chex.Array
-  g0: chex.Array
-  g0_face: chex.Array
-  g1: chex.Array
-  g1_face: chex.Array
-  g2: chex.Array
-  g2_face: chex.Array
-  g3: chex.Array
-  g3_face: chex.Array
-  g2g3_over_rhon: chex.Array
-  g2g3_over_rhon_face: chex.Array
-  g2g3_over_rhon_hires: chex.Array
-  F: chex.Array
-  F_face: chex.Array
-  F_hires: chex.Array
-  R_in: chex.Array
-  R_in_face: chex.Array
-  R_out: chex.Array
-  R_out_face: chex.Array
-  spr_hires: chex.Array
-  rho_hires_norm: chex.Array
-  rho_hires: chex.Array
-  Phi_b_dot: chex.Array
-  _z_magnetic_axis: chex.Array | None
+  Phi: array_typing.Array
+  Phi_face: array_typing.Array
+  R_major: array_typing.FloatScalar
+  a_minor: array_typing.FloatScalar
+  B_0: array_typing.FloatScalar
+  volume: array_typing.Array
+  volume_face: array_typing.Array
+  area: array_typing.Array
+  area_face: array_typing.Array
+  vpr: array_typing.Array
+  vpr_face: array_typing.Array
+  spr: array_typing.Array
+  spr_face: array_typing.Array
+  delta_face: array_typing.Array
+  elongation: array_typing.Array
+  elongation_face: array_typing.Array
+  g0: array_typing.Array
+  g0_face: array_typing.Array
+  g1: array_typing.Array
+  g1_face: array_typing.Array
+  g2: array_typing.Array
+  g2_face: array_typing.Array
+  g3: array_typing.Array
+  g3_face: array_typing.Array
+  g2g3_over_rhon: array_typing.Array
+  g2g3_over_rhon_face: array_typing.Array
+  g2g3_over_rhon_hires: array_typing.Array
+  F: array_typing.Array
+  F_face: array_typing.Array
+  F_hires: array_typing.Array
+  R_in: array_typing.Array
+  R_in_face: array_typing.Array
+  R_out: array_typing.Array
+  R_out_face: array_typing.Array
+  spr_hires: array_typing.Array
+  rho_hires_norm: array_typing.Array
+  rho_hires: array_typing.Array
+  Phi_b_dot: array_typing.FloatScalar
+  _z_magnetic_axis: array_typing.FloatScalar | None
 
   def __eq__(self, other: 'Geometry') -> bool:
     try:
@@ -230,27 +232,27 @@ def q_correction_factor(self) -> chex.Numeric:
     )
 
   @property
-  def rho_norm(self) -> chex.Array:
+  def rho_norm(self) -> array_typing.Array:
     r"""Normalized toroidal flux coordinate on cell grid [dimensionless]."""
     return self.torax_mesh.cell_centers
 
   @property
-  def rho_face_norm(self) -> chex.Array:
+  def rho_face_norm(self) -> array_typing.Array:
     r"""Normalized toroidal flux coordinate on face grid [dimensionless]."""
     return self.torax_mesh.face_centers
 
   @property
-  def drho_norm(self) -> chex.Array:
+  def drho_norm(self) -> array_typing.Array:
     r"""Grid size for rho_norm [dimensionless]."""
     return jnp.array(self.torax_mesh.dx)
 
   @property
-  def rho_face(self) -> chex.Array:
+  def rho_face(self) -> array_typing.Array:
     r"""Toroidal flux coordinate on face grid :math:`\mathrm{m}`."""
     return self.rho_face_norm * jnp.expand_dims(self.rho_b, axis=-1)
 
   @property
-  def rho(self) -> chex.Array:
+  def rho(self) -> array_typing.Array:
     r"""Toroidal flux coordinate on cell grid :math:`\mathrm{m}`.
 
     The toroidal flux coordinate is defined as
@@ -261,49 +263,49 @@ def rho(self) -> chex.Array:
     return self.rho_norm * jnp.expand_dims(self.rho_b, axis=-1)
 
   @property
-  def r_mid(self) -> chex.Array:
+  def r_mid(self) -> array_typing.Array:
     """Midplane radius of the plasma [m], defined as (Rout-Rin)/2."""
     return (self.R_out - self.R_in) / 2
 
   @property
-  def r_mid_face(self) -> chex.Array:
+  def r_mid_face(self) -> array_typing.Array:
     """Midplane radius of the plasma on the face grid [m]."""
     return (self.R_out_face - self.R_in_face) / 2
 
   @property
-  def epsilon(self) -> chex.Array:
+  def epsilon(self) -> array_typing.Array:
     """Local midplane inverse aspect ratio [dimensionless]."""
     return (self.R_out - self.R_in) / (self.R_out + self.R_in)
 
   @property
-  def epsilon_face(self) -> chex.Array:
+  def epsilon_face(self) -> array_typing.Array:
     """Local midplane inverse aspect ratio on the face grid [dimensionless]."""
     return (self.R_out_face - self.R_in_face) / (
         self.R_out_face + self.R_in_face
     )
 
   @property
-  def drho(self) -> chex.Array:
+  def drho(self) -> array_typing.Array:
     """Grid size for rho [m]."""
     return self.drho_norm * self.rho_b
 
   @property
-  def rho_b(self) -> chex.Array:
+  def rho_b(self) -> array_typing.FloatScalar:
     """Toroidal flux coordinate [m] at boundary (LCFS)."""
     return jnp.sqrt(self.Phi_b / np.pi / self.B_0)
 
   @property
-  def Phi_b(self) -> chex.Array:
+  def Phi_b(self) -> array_typing.FloatScalar:
     r"""Toroidal flux at boundary (LCFS) :math:`\mathrm{Wb}`."""
     return self.Phi_face[..., -1]
 
   @property
-  def g1_over_vpr(self) -> chex.Array:
+  def g1_over_vpr(self) -> array_typing.Array:
     r"""g1/vpr [:math:`\mathrm{m}`]."""
     return self.g1 / self.vpr
 
   @property
-  def g1_over_vpr2(self) -> chex.Array:
+  def g1_over_vpr2(self) -> array_typing.Array:
     r"""g1/vpr**2 [:math:`\mathrm{m}^{-2}`]."""
     return self.g1 / self.vpr**2
 
@@ -389,7 +391,7 @@ def stack_geometries(geometries: Sequence[Geometry]) -> Geometry:
     field_name = field.name
     field_value = getattr(first_geo, field_name)
     # Stack stackable fields. Save first geo's value for non-stackable fields.
-    if isinstance(field_value, chex.Array) or isinstance(field_value, float):
+    if isinstance(field_value, (array_typing.Array, array_typing.FloatScalar)):
       field_values = [getattr(geo, field_name) for geo in geometries]
       stacked_data[field_name] = np.stack(field_values)
     else:

torax/_src/geometry/standard_geometry.py

@@ -21,13 +21,13 @@
 from collections.abc import Mapping
 import dataclasses
 import logging
-
 import chex
 import contourpy
 from imas import ids_toplevel
 import jax
 import numpy as np
 import scipy
+from torax._src import array_typing
 from torax._src import constants
 from torax._src import interpolated_param
 from torax._src.geometry import geometry
@@ -74,14 +74,14 @@ class StandardGeometry(geometry.Geometry):
   """
 
   Ip_from_parameters: bool = dataclasses.field(metadata=dict(static=True))
-  Ip_profile_face: chex.Array
-  psi: chex.Array
-  psi_from_Ip: chex.Array
-  psi_from_Ip_face: chex.Array
-  j_total: chex.Array
-  j_total_face: chex.Array
-  delta_upper_face: chex.Array
-  delta_lower_face: chex.Array
+  Ip_profile_face: array_typing.FloatVectorFace
+  psi: array_typing.FloatVectorCell
+  psi_from_Ip: array_typing.FloatVectorCell
+  psi_from_Ip_face: array_typing.FloatVectorFace
+  j_total: array_typing.Array
+  j_total_face: array_typing.FloatVectorFace
+  delta_upper_face: array_typing.FloatVectorFace
+  delta_lower_face: array_typing.FloatVectorFace
 
 
 @jax.tree_util.register_dataclass
@@ -169,28 +169,28 @@ class StandardGeometryIntermediates:
 
   geometry_type: geometry.GeometryType
   Ip_from_parameters: bool
-  R_major: chex.Numeric
-  a_minor: chex.Numeric
-  B_0: chex.Numeric
-  psi: chex.Array
-  Ip_profile: chex.Array
-  Phi: chex.Array
-  R_in: chex.Array
-  R_out: chex.Array
-  F: chex.Array
-  int_dl_over_Bp: chex.Array
-  flux_surf_avg_1_over_R: chex.Array
-  flux_surf_avg_1_over_R2: chex.Array
-  flux_surf_avg_Bp2: chex.Array
-  flux_surf_avg_RBp: chex.Array
-  flux_surf_avg_R2Bp2: chex.Array
-  delta_upper_face: chex.Array
-  delta_lower_face: chex.Array
-  elongation: chex.Array
-  vpr: chex.Array
+  R_major: array_typing.FloatScalar
+  a_minor: array_typing.FloatScalar
+  B_0: array_typing.FloatScalar
+  psi: array_typing.Array
+  Ip_profile: array_typing.Array
+  Phi: array_typing.Array
+  R_in: array_typing.Array
+  R_out: array_typing.Array
+  F: array_typing.Array
+  int_dl_over_Bp: array_typing.Array
+  flux_surf_avg_1_over_R: array_typing.Array
+  flux_surf_avg_1_over_R2: array_typing.Array
+  flux_surf_avg_Bp2: array_typing.Array
+  flux_surf_avg_RBp: array_typing.Array
+  flux_surf_avg_R2Bp2: array_typing.Array
+  delta_upper_face: array_typing.Array
+  delta_lower_face: array_typing.Array
+  elongation: array_typing.Array
+  vpr: array_typing.Array
   n_rho: int
   hires_factor: int
-  z_magnetic_axis: chex.Numeric | None
+  z_magnetic_axis: array_typing.FloatScalar | None
 
   def __post_init__(self):
     """Extrapolates edge values and smooths near-axis values.
@@ -1079,7 +1079,7 @@ def build_standard_geometry(
 
   # fill geometry structure
   # normalized grid
-  mesh = torax_pydantic.Grid1D(nx=intermediate.n_rho,)
+  mesh = torax_pydantic.Grid1D(nx=intermediate.n_rho)
   rho_b = rho_intermediate[-1]  # radius denormalization constant
   # helper variables for mesh cells and faces
   rho_face_norm = mesh.face_centers