Stricter typing for geometry.

PiperOrigin-RevId: 800410735

Author: sbodenstein

torax/_src/geometry/geometry.py

@@ -19,14 +19,13 @@
 import chex
 import jax
 import jax.numpy as jnp
+import jaxtyping as jt
 import numpy as np
-from torax._src import array_typing
+from torax._src import array_typing as at
 from torax._src.torax_pydantic import torax_pydantic
 
 
-def face_to_cell(
-    face: array_typing.FloatVectorFace,
-) -> array_typing.FloatVectorCell:
+def face_to_cell(face: at.FloatVectorFace) -> at.FloatVectorCell:
   """Infers cell values corresponding to a vector of face values.
 
   Simply a linear interpolation between face values.
@@ -59,6 +58,7 @@ class GeometryType(enum.IntEnum):
 # pylint: disable=invalid-name
 
 
+@at.jaxtyped
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True)
 class Geometry:
@@ -170,47 +170,47 @@ class Geometry:
       [:math:`\mathrm{m}`].
   """
 
-  geometry_type: GeometryType
+  geometry_type: GeometryType = dataclasses.field(metadata=dict(static=True))
   torax_mesh: torax_pydantic.Grid1D
-  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
+  Phi: jt.Float[at.Array, '*stack rhon']
+  Phi_face: jt.Float[at.Array, '*stack rhon+1']
+  R_major: jt.Float[at.Array, '*stack']
+  a_minor: jt.Float[at.Array, '*stack']
+  B_0: jt.Float[at.Array, '*stack']
+  volume: jt.Float[at.Array, '*stack rhon']
+  volume_face: jt.Float[at.Array, '*stack rhon+1']
+  area: jt.Float[at.Array, '*stack rhon']
+  area_face: jt.Float[at.Array, '*stack rhon+1']
+  vpr: jt.Float[at.Array, '*stack rhon']
+  vpr_face: jt.Float[at.Array, '*stack rhon+1']
+  spr: jt.Float[at.Array, '*stack rhon']
+  spr_face: jt.Float[at.Array, '*stack rhon+1']
+  delta_face: jt.Float[at.Array, '*stack rhon+1']
+  elongation: jt.Float[at.Array, '*stack rhon']
+  elongation_face: jt.Float[at.Array, '*stack rhon+1']
+  g0: jt.Float[at.Array, '*stack rhon']
+  g0_face: jt.Float[at.Array, '*stack rhon+1']
+  g1: jt.Float[at.Array, '*stack rhon']
+  g1_face: jt.Float[at.Array, '*stack rhon+1']
+  g2: jt.Float[at.Array, '*stack rhon']
+  g2_face: jt.Float[at.Array, '*stack rhon+1']
+  g3: jt.Float[at.Array, '*stack rhon']
+  g3_face: jt.Float[at.Array, '*stack rhon+1']
+  g2g3_over_rhon: jt.Float[at.Array, '*stack rhon']
+  g2g3_over_rhon_face: jt.Float[at.Array, '*stack rhon+1']
+  g2g3_over_rhon_hires: jt.Float[at.Array, '*stack channel']
+  F: jt.Float[at.Array, '*stack rhon']
+  F_face: jt.Float[at.Array, '*stack rhon+1']
+  F_hires: jt.Float[at.Array, '*stack channel']
+  R_in: jt.Float[at.Array, '*stack rhon']
+  R_in_face: jt.Float[at.Array, '*stack rhon+1']
+  R_out: jt.Float[at.Array, '*stack rhon']
+  R_out_face: jt.Float[at.Array, '*stack rhon+1']
+  spr_hires: jt.Float[at.Array, '*stack channel']
+  rho_hires_norm: jt.Float[at.Array, '*stack channel']
+  rho_hires: jt.Float[at.Array, '*stack channel']
+  Phi_b_dot: jt.Float[at.Array | float, '*stack']
+  _z_magnetic_axis: jt.Float[at.Array, '*stack'] | None
 
   def __eq__(self, other: 'Geometry') -> bool:
     try:
@@ -220,7 +220,7 @@ def __eq__(self, other: 'Geometry') -> bool:
     return True
 
   @property
-  def q_correction_factor(self) -> chex.Numeric:
+  def q_correction_factor(self) -> jt.Float[jax.Array, '']:
     """Ad-hoc fix for non-physical circular geometry model.
 
     Set such that q(r=a) = 3 for standard ITER parameters.
@@ -232,27 +232,27 @@ def q_correction_factor(self) -> chex.Numeric:
     )
 
   @property
-  def rho_norm(self) -> array_typing.Array:
+  def rho_norm(self) -> at.FloatVectorCell:
     r"""Normalized toroidal flux coordinate on cell grid [dimensionless]."""
     return self.torax_mesh.cell_centers
 
   @property
-  def rho_face_norm(self) -> array_typing.Array:
+  def rho_face_norm(self) -> at.FloatVectorFace:
     r"""Normalized toroidal flux coordinate on face grid [dimensionless]."""
     return self.torax_mesh.face_centers
 
   @property
-  def drho_norm(self) -> array_typing.Array:
+  def drho_norm(self) -> at.FloatVectorCell:
     r"""Grid size for rho_norm [dimensionless]."""
     return jnp.array(self.torax_mesh.dx)
 
   @property
-  def rho_face(self) -> array_typing.Array:
+  def rho_face(self) -> at.FloatVectorFace:
     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) -> array_typing.Array:
+  def rho(self) -> jt.Float[jax.Array, 'rhon']:
     r"""Toroidal flux coordinate on cell grid :math:`\mathrm{m}`.
 
     The toroidal flux coordinate is defined as
@@ -263,54 +263,54 @@ def rho(self) -> array_typing.Array:
     return self.rho_norm * jnp.expand_dims(self.rho_b, axis=-1)
 
   @property
-  def r_mid(self) -> array_typing.Array:
+  def r_mid(self) -> jt.Float[at.Array, '*stack rhon']:
     """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) -> array_typing.Array:
+  def r_mid_face(self) -> jt.Float[at.Array, '*stack rhon+1']:
     """Midplane radius of the plasma on the face grid [m]."""
     return (self.R_out_face - self.R_in_face) / 2
 
   @property
-  def epsilon(self) -> array_typing.Array:
+  def epsilon(self) -> jt.Float[at.Array, '*stack rhon']:
     """Local midplane inverse aspect ratio [dimensionless]."""
     return (self.R_out - self.R_in) / (self.R_out + self.R_in)
 
   @property
-  def epsilon_face(self) -> array_typing.Array:
+  def epsilon_face(self) -> jt.Float[at.Array, '*stack rhon+1']:
     """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) -> array_typing.Array:
+  def drho(self) -> jt.Float[at.Array, 'rhon']:
     """Grid size for rho [m]."""
     return self.drho_norm * self.rho_b
 
   @property
-  def rho_b(self) -> array_typing.FloatScalar:
+  def rho_b(self) -> jt.Float[jax.Array, '*stack']:
     """Toroidal flux coordinate [m] at boundary (LCFS)."""
     return jnp.sqrt(self.Phi_b / np.pi / self.B_0)
 
   @property
-  def Phi_b(self) -> array_typing.FloatScalar:
+  def Phi_b(self) -> jt.Float[at.Array, '*stack']:
     r"""Toroidal flux at boundary (LCFS) :math:`\mathrm{Wb}`."""
     return self.Phi_face[..., -1]
 
   @property
-  def g1_over_vpr(self) -> array_typing.Array:
+  def g1_over_vpr(self) -> jt.Float[at.Array, '*stack rhon']:
     r"""g1/vpr [:math:`\mathrm{m}`]."""
     return self.g1 / self.vpr
 
   @property
-  def g1_over_vpr2(self) -> array_typing.Array:
+  def g1_over_vpr2(self) -> jt.Float[at.Array, '*stack rhon']:
     r"""g1/vpr**2 [:math:`\mathrm{m}^{-2}`]."""
     return self.g1 / self.vpr**2
 
   @property
-  def g0_over_vpr_face(self) -> jax.Array:
+  def g0_over_vpr_face(self) -> jt.Float[at.Array, '*stack rhon_1']:
     """g0_face/vpr_face [:math:`m^{-1}`], equal to 1/rho_b on-axis."""
     # Calculate the bulk of the array (excluding the first element)
     # to avoid division by zero.
@@ -322,7 +322,7 @@ def g0_over_vpr_face(self) -> jax.Array:
     )
 
   @property
-  def g1_over_vpr_face(self) -> jax.Array:
+  def g1_over_vpr_face(self) -> jt.Float[at.Array, '*stack rhon_1']:
     r"""g1_face/vpr_face [:math:`\mathrm{m}`]. Zero on-axis."""
     bulk = self.g1_face[..., 1:] / self.vpr_face[..., 1:]
     first_element = jnp.zeros_like(self.rho_b)
@@ -331,7 +331,7 @@ def g1_over_vpr_face(self) -> jax.Array:
     )
 
   @property
-  def g1_over_vpr2_face(self) -> jax.Array:
+  def g1_over_vpr2_face(self) -> jt.Float[at.Array, '*stack rhon_1']:
     """g1_face/vpr_face**2 [:math:`m^{-2}`], equal to 1/rho_b**2 on-axis."""
     bulk = self.g1_face[..., 1:] / self.vpr_face[..., 1:] ** 2
     first_element = jnp.ones_like(self.rho_b) / self.rho_b**2
@@ -340,20 +340,20 @@ def g1_over_vpr2_face(self) -> jax.Array:
     )
 
   @property
-  def gm9(self) -> jax.Array:
+  def gm9(self) -> jt.Float[at.Array, '*stack rhon']:
     r"""<1/R> on cell grid [:math:`\mathrm{m}^{-1}`]."""
     return 2 * jnp.pi * self.spr / self.vpr
 
   @property
-  def gm9_face(self) -> jax.Array:
+  def gm9_face(self) -> jt.Float[at.Array, '*stack rhon_1']:
     r"""<1/R> on face grid [:math:`\mathrm{m}^{-1}`]."""
     bulk = 2 * jnp.pi * self.spr_face[..., 1:] / self.vpr_face[..., 1:]
     first_element = 1 / self.R_major
     return jnp.concatenate(
         [jnp.expand_dims(first_element, axis=-1), bulk], axis=-1
     )
 
-  def z_magnetic_axis(self) -> chex.Numeric:
+  def z_magnetic_axis(self) -> jt.Float[jax.Array, '*stack']:
     """z position of magnetic axis [m]."""
     z_magnetic_axis = self._z_magnetic_axis
     if z_magnetic_axis is not None:
@@ -391,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, (array_typing.Array, array_typing.FloatScalar)):
+    if isinstance(field_value, (at.Array, at.FloatScalar)):
       field_values = [getattr(geo, field_name) for geo in geometries]
       stacked_data[field_name] = np.stack(field_values)
     else:
@@ -401,9 +401,10 @@ def stack_geometries(geometries: Sequence[Geometry]) -> Geometry:
   return first_geo.__class__(**stacked_data)
 
 
+@at.jaxtyped
 def update_geometries_with_Phibdot(
     *,
-    dt: chex.Numeric,
+    dt: at.FloatScalar,
     geo_t: Geometry,
     geo_t_plus_dt: Geometry,
 ) -> tuple[Geometry, Geometry]:

torax/_src/geometry/tests/geometry_test.py

@@ -151,8 +151,8 @@ def test_stack_geometries_error_handling_different_geometry_types(self):
   def test_update_phibdot(self):
     """Test update_phibdot for circular geometries."""
     geo = geometry_pydantic_model.CircularConfig().build_geometry()
-    geo0 = dataclasses.replace(geo, Phi_face=np.array([1.0]))
-    geo1 = dataclasses.replace(geo, Phi_face=np.array([2.0]))
+    geo0 = dataclasses.replace(geo, Phi_face=np.ones_like(geo.Phi_face) * 1.0)
+    geo1 = dataclasses.replace(geo, Phi_face=np.ones_like(geo.Phi_face) * 2.0)
     geo0_updated, geo1_updated = geometry.update_geometries_with_Phibdot(
         dt=0.1, geo_t=geo0, geo_t_plus_dt=geo1
     )
@@ -166,7 +166,9 @@ def test_geometry_eq(self):
       self.assertEqual(geo1, geo2)
 
     with self.subTest('different_geometries_are_not_equal'):
-      geo3 = dataclasses.replace(geo1, Phi_face=np.array([2.0]))
+      geo3 = dataclasses.replace(
+          geo1, Phi_face=np.ones_like(geo1.Phi_face) * 2.0
+      )
       self.assertNotEqual(geo1, geo3)