Move dx from being a grid input to being a property.

Currently `dx` is always `1/nx`.

Drive-by: fix formatting check.
PiperOrigin-RevId: 778535641

Author: Nush395

torax/_src/array_typing.py

@@ -14,6 +14,7 @@
 # ============================================================================
 """Common types for using jaxtyping in TORAX."""
 from typing import TypeAlias
+
 import jax
 import jaxtyping as jt
 import numpy as np

torax/_src/config/tests/build_runtime_params_test.py

@@ -33,7 +33,7 @@ class RuntimeParamsSliceTest(parameterized.TestCase):
 
   def setUp(self):
     super().setUp()
-    self._torax_mesh = torax_pydantic.Grid1D(nx=4, dx=0.25)
+    self._torax_mesh = torax_pydantic.Grid1D(nx=4,)
 
   def test_time_dependent_provider_is_time_dependent(self):
     """Tests that the runtime_params slice provider is time dependent."""

torax/_src/geometry/circular_geometry.py

@@ -45,9 +45,8 @@ def build_circular_geometry(
   """
   # circular geometry assumption of r/a_minor = rho_norm, the normalized
   # toroidal flux coordinate.
-  drho_norm = 1.0 / n_rho
   # Define mesh (Slab Uniform 1D with Jacobian = 1)
-  mesh = torax_pydantic.Grid1D(nx=n_rho, dx=drho_norm)
+  mesh = torax_pydantic.Grid1D(nx=n_rho,)
   # toroidal flux coordinate (rho) at boundary (last closed flux surface)
   rho_b = np.asarray(a_minor)
 

torax/_src/geometry/standard_geometry.py

@@ -1057,9 +1057,8 @@ def build_standard_geometry(
   j_total = np.concatenate([np.array([j_total_face_axis]), j_total_face_bulk])
 
   # fill geometry structure
-  drho_norm = float(rho_norm_intermediate[-1]) / intermediate.n_rho
   # normalized grid
-  mesh = torax_pydantic.Grid1D(nx=intermediate.n_rho, dx=drho_norm)
+  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

torax/_src/sources/tests/ion_cyclotron_source_test.py

@@ -112,7 +112,7 @@ def test_build_dynamic_params(self):
     self.assertIsInstance(source, self._source_config_class)
     torax_pydantic.set_grid(
         source,
-        torax_pydantic.Grid1D(nx=4, dx=0.25),
+        torax_pydantic.Grid1D(nx=4,),
     )
     dynamic_params = source.build_dynamic_params(t=0.0)
     self.assertIsInstance(

torax/_src/sources/tests/mavrin_impurity_radiation_heat_sink_test.py

@@ -43,7 +43,7 @@ def test_correct_dynamic_params_built(self):
     })
     # Set the grid to allows the dynamic params to be built without making the
     # full config.
-    torax_pydantic.set_grid(sources, torax_pydantic.Grid1D(nx=4, dx=0.25))
+    torax_pydantic.set_grid(sources, torax_pydantic.Grid1D(nx=4,))
     runtime_params = getattr(sources, self._source_name).build_dynamic_params(
         t=0.0
     )

torax/_src/sources/tests/pydantic_model_test.py

@@ -162,7 +162,7 @@ def test_adding_a_source_with_prescribed_values(self):
             ),
         },
     })
-    mesh = torax_pydantic.Grid1D(nx=4, dx=0.25)
+    mesh = torax_pydantic.Grid1D(nx=4,)
     torax_pydantic.set_grid(sources, mesh)
     source = sources.generic_current
     self.assertLen(source.prescribed_values, 1)

torax/_src/sources/tests/source_profiles_test.py

@@ -65,7 +65,7 @@ def test_summed_T_i_profiles_dont_change_when_jitting(self):
 
   def test_merging_source_profiles(self):
     """Tests that the implicit and explicit source profiles merge correctly."""
-    torax_mesh = torax_pydantic.Grid1D(nx=10, dx=0.1)
+    torax_mesh = torax_pydantic.Grid1D(nx=10,)
     sources = sources_pydantic_model.Sources.from_dict(
         default_sources.get_default_source_config()
     )

torax/_src/sources/tests/test_lib.py

@@ -59,7 +59,7 @@ def test_build_dynamic_params(self):
     self.assertIsInstance(source, self._source_config_class)
     torax_pydantic.set_grid(
         source,
-        torax_pydantic.Grid1D(nx=4, dx=0.25),
+        torax_pydantic.Grid1D(nx=4,),
     )
     dynamic_params = source.build_dynamic_params(t=0.0)
     self.assertIsInstance(

torax/_src/torax_pydantic/interpolated_param_2d.py

@@ -39,11 +39,13 @@ class Grid1D(model_base.BaseModelFrozen):
 
   Attributes:
     nx: Number of cells.
-    dx: Distance between cell centers.
   """
 
   nx: typing_extensions.Annotated[pydantic.conint(ge=4), model_base.JAX_STATIC]
-  dx: typing_extensions.Annotated[pydantic.PositiveFloat, model_base.JAX_STATIC]
+
+  @functools.cached_property
+  def dx(self) -> float:
+    return 1 / self.nx
 
   @property
   def face_centers(self) -> np.ndarray:
@@ -441,7 +443,6 @@ def _update_rule(submodel):
     # the same NumPy arrays.
     new_grid = Grid1D.model_construct(
         nx=grid.nx,
-        dx=grid.dx,
         face_centers=grid.face_centers,
         cell_centers=grid.cell_centers,
     )