diff --git a/torax/geometry/circular_geometry.py b/torax/geometry/circular_geometry.py
index 8b4df414..85081340 100644
--- a/torax/geometry/circular_geometry.py
+++ b/torax/geometry/circular_geometry.py
@@ -220,7 +220,6 @@ def build_circular_geometry(
   return CircularAnalyticalGeometry(
       # Set the standard geometry params.
       geometry_type=geometry.GeometryType.CIRCULAR.value,
-      drho_norm=np.asarray(drho_norm),
       torax_mesh=mesh,
       Phi=Phi,
       Phi_face=Phi_face,
diff --git a/torax/geometry/geometry.py b/torax/geometry/geometry.py
index 2ed46869..41d19019 100644
--- a/torax/geometry/geometry.py
+++ b/torax/geometry/geometry.py
@@ -17,7 +17,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping, Sequence
+from collections.abc import Sequence
 import dataclasses
 import enum
 
@@ -25,7 +25,6 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from torax import array_typing
 from torax import jax_utils
 
 
@@ -125,12 +124,14 @@ class Geometry:
 
   Most users should default to using the StandardGeometry class, whether the
   source of their geometry comes from CHEASE, MEQ, etc.
+
+  Properties work for both 1D radial arrays and 2D stacked arrays where the
+  leading dimension is time.
   """
 
   # TODO(b/356356966): extend documentation to define what each attribute is.
   geometry_type: int
   torax_mesh: Grid1D
-  drho_norm: array_typing.ArrayFloat
   Phi: chex.Array
   Phi_face: chex.Array
   Rmaj: chex.Array
@@ -182,13 +183,17 @@ def rho_norm(self) -> chex.Array:
   def rho_face_norm(self) -> chex.Array:
     return self.torax_mesh.face_centers
 
+  @property
+  def drho_norm(self) -> chex.Array:
+    return jnp.array(self.torax_mesh.dx)
+
   @property
   def rho_face(self) -> chex.Array:
-    return self.rho_face_norm * self.rho_b
+    return self.rho_face_norm * jnp.expand_dims(self.rho_b, axis=-1)
 
   @property
   def rho(self) -> chex.Array:
-    return self.rho_norm * self.rho_b
+    return self.rho_norm * jnp.expand_dims(self.rho_b, axis=-1)
 
   @property
   def rmid(self) -> chex.Array:
@@ -210,7 +215,7 @@ def rho_b(self) -> chex.Array:
   @property
   def Phib(self) -> chex.Array:
     """Toroidal flux at boundary (LCFS)."""
-    return self.Phi_face[-1]
+    return self.Phi_face[..., -1]
 
   @property
   def g1_over_vpr(self) -> chex.Array:
@@ -222,24 +227,33 @@ def g1_over_vpr2(self) -> chex.Array:
 
   @property
   def g0_over_vpr_face(self) -> jax.Array:
-    return jnp.concatenate((
-        jnp.ones(1) / self.rho_b,  # correct value is 1/rho_b on-axis
-        self.g0_face[1:] / self.vpr_face[1:],  # avoid div by zero on-axis
-    ))
+    # Calculate the bulk of the array (excluding the first element)
+    # to avoid division by zero.
+    bulk = self.g0_face[..., 1:] / self.vpr_face[..., 1:]
+    # Correct value on-axis is 1/rho_b
+    first_element = jnp.ones_like(self.rho_b) / self.rho_b
+    # Concatenate to handle both 1D (no leading dim) and 2D cases
+    return jnp.concatenate(
+        [jnp.expand_dims(first_element, axis=-1), bulk], axis=-1
+    )
 
   @property
   def g1_over_vpr_face(self) -> jax.Array:
-    return jnp.concatenate((
-        jnp.zeros(1),  # correct value is zero on-axis
-        self.g1_face[1:] / self.vpr_face[1:],  # avoid div by zero on-axis
-    ))
+    bulk = self.g1_face[..., 1:] / self.vpr_face[..., 1:]
+    # Correct value on-axis is 0
+    first_element = jnp.zeros_like(self.rho_b)
+    return jnp.concatenate(
+        [jnp.expand_dims(first_element, axis=-1), bulk], axis=-1
+    )
 
   @property
   def g1_over_vpr2_face(self) -> jax.Array:
-    return jnp.concatenate((
-        jnp.ones(1) / self.rho_b**2,  # correct value is 1/rho_b**2 on-axis
-        self.g1_face[1:] / self.vpr_face[1:] ** 2,  # avoid div by zero on-axis
-    ))
+    bulk = self.g1_face[..., 1:] / self.vpr_face[..., 1:]**2
+    # Correct value on-axis is 1/rho_b**2
+    first_element = jnp.ones_like(self.rho_b) / self.rho_b**2
+    return jnp.concatenate(
+        [jnp.expand_dims(first_element, axis=-1), bulk], axis=-1
+    )
 
   def z_magnetic_axis(self) -> chex.Numeric:
     z_magnetic_axis = self._z_magnetic_axis
@@ -251,43 +265,40 @@ def z_magnetic_axis(self) -> chex.Numeric:
       )
 
 
-def stack_geometries(
-    geometries: Sequence[Geometry],
-) -> Mapping[str, chex.Array]:
+def stack_geometries(geometries: Sequence[Geometry]) -> Geometry:
   """Batch together a sequence of geometries.
 
-  The batched geometries are returned as a dictionary of arrays. Any fields
-  that are not arrays are not included in the dictionary and all properties
-  are excluded as well.
-
   Args:
-    geometries: A sequence of geometries to stack. The geometries must have the
-      same mesh, geometry type, and drho_norm.
+    geometries: A sequence of geometries to stack. The geometries must have
+      the same mesh, geometry type.
 
   Returns:
-    A dictionary of arrays, where each array has the same shape as the
-    corresponding attribute in the input geometries, but with an additional
-    leading axis (e.g. for the time dimension).
+    A Geometry object, where each array attribute has an additional
+    leading axis (e.g. for the time dimension) compared to each Geometry in
+    the input sequence.
   """
   if not geometries:
     raise ValueError('No geometries provided.')
-  # Stack the geometries.
-  torax_mesh = geometries[0].torax_mesh
-  geometry_type = geometries[0].geometry_type
-  drho_norm = geometries[0].drho_norm
+  # Ensure that all geometries have same mesh and are of same type.
+  first_geo = geometries[0]
+  torax_mesh = first_geo.torax_mesh
+  geometry_type = first_geo.geometry_type
   for geometry in geometries[1:]:
     if geometry.torax_mesh != torax_mesh:
       raise ValueError('All geometries must have the same mesh.')
     if geometry.geometry_type != geometry_type:
       raise ValueError('All geometries must have the same geometry type.')
-    if geometry.drho_norm != drho_norm:
-      raise ValueError('All geometries must have the same drho_norm.')
-  array_fields = [
-      attr
-      for attr, val in dataclasses.asdict(geometries[0]).items()
-      if isinstance(val, chex.Array)
-  ]
-  return {
-      attr: jnp.stack([getattr(geo, attr) for geo in geometries])
-      for attr in array_fields
-  }
+
+  stacked_data = {}
+  for field in dataclasses.fields(first_geo):
+    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):
+      field_values = [getattr(geo, field_name) for geo in geometries]
+      stacked_data[field_name] = jnp.stack(field_values)
+    else:
+      stacked_data[field_name] = field_value
+  # Create a new object with the stacked data with the same class (i.e.
+  # could be child classes of Geometry)
+  return first_geo.__class__(**stacked_data)
diff --git a/torax/geometry/standard_geometry.py b/torax/geometry/standard_geometry.py
index 19ce7e44..39b3536b 100644
--- a/torax/geometry/standard_geometry.py
+++ b/torax/geometry/standard_geometry.py
@@ -261,9 +261,9 @@ def from_chease(
     return cls(
         geometry_type=geometry.GeometryType.CHEASE,
         Ip_from_parameters=Ip_from_parameters,
-        Rmaj=Rmaj,
-        Rmin=Rmin,
-        B=B0,
+        Rmaj=np.array(Rmaj),
+        Rmin=np.array(Rmin),
+        B=np.array(B0),
         psi=psi,
         Ip_profile=Ip_chease,
         Phi=Phi,
@@ -1008,7 +1008,6 @@ def build_standard_geometry(
 
   return StandardGeometry(
       geometry_type=intermediate.geometry_type.value,
-      drho_norm=np.asarray(drho_norm),
       torax_mesh=mesh,
       Phi=Phi,
       Phi_face=Phi_face,
diff --git a/torax/geometry/tests/geometry_test.py b/torax/geometry/tests/geometry_test.py
index dea7fc58..074c1dc3 100644
--- a/torax/geometry/tests/geometry_test.py
+++ b/torax/geometry/tests/geometry_test.py
@@ -66,18 +66,70 @@ def test_none_z_magnetic_axis_raises_an_error(self):
         _ = foo()
 
   def test_stack_geometries(self):
-    """Test that geometries can be stacked."""
-    geo_0 = circular_geometry.build_circular_geometry(Rmaj=1, B0=5.3)
-    geo_1 = circular_geometry.build_circular_geometry(Rmaj=2, B0=3.7)
-    stacked_geometries = geometry.stack_geometries([geo_0, geo_1])
-    self.assertNotIn('geometry_type', stacked_geometries.keys())
-    self.assertNotIn('torax_mesh', stacked_geometries.keys())
-
-    for key, value in stacked_geometries.items():
-      self.assertEqual(value.shape, (2,) + getattr(geo_0, key).shape)
-
-    np.testing.assert_allclose(stacked_geometries['Rmaj'], np.array([1, 2]))
-    np.testing.assert_allclose(stacked_geometries['B0'], np.array([5.3, 3.7]))
+    """Test for stack_geometries."""
+    # Create a few different geometries
+    geo0 = circular_geometry.build_circular_geometry(Rmaj=1.0, B0=2.0, n_rho=10)
+    geo1 = circular_geometry.build_circular_geometry(Rmaj=1.5, B0=2.5, n_rho=10)
+    geo2 = circular_geometry.build_circular_geometry(Rmaj=2.0, B0=3.0, n_rho=10)
+
+    # Stack them
+    stacked_geo = geometry.stack_geometries([geo0, geo1, geo2])
+
+    # Check that the stacked geometry has the correct type and mesh
+    self.assertEqual(stacked_geo.geometry_type, geo0.geometry_type)
+    self.assertEqual(stacked_geo.torax_mesh, geo0.torax_mesh)
+
+    # Check some specific stacked values
+    np.testing.assert_allclose(stacked_geo.Rmaj, np.array([1.0, 1.5, 2.0]))
+    np.testing.assert_allclose(stacked_geo.B0, np.array([2.0, 2.5, 3.0]))
+    np.testing.assert_allclose(
+        stacked_geo.Phi_face[:, -1],
+        np.array([geo0.Phi_face[-1], geo1.Phi_face[-1], geo2.Phi_face[-1]]),
+    )
+
+    # Check stacking of derived properties
+    np.testing.assert_allclose(
+        stacked_geo.rho_b, np.array([geo0.rho_b, geo1.rho_b, geo2.rho_b])
+    )
+
+    # Check a property that depends on a stacked property (rho depends on rho_b)
+    np.testing.assert_allclose(
+        stacked_geo.rho,
+        np.array([
+            stacked_geo.rho_norm * geo0.rho_b,
+            stacked_geo.rho_norm * geo1.rho_b,
+            stacked_geo.rho_norm * geo2.rho_b,
+        ]),
+    )
+
+    # Check properties with special handling for on-axis values.
+    np.testing.assert_allclose(
+        stacked_geo.g0_over_vpr_face[:, 0], 1 / stacked_geo.rho_b
+    )
+    np.testing.assert_allclose(
+        stacked_geo.g1_over_vpr2_face[:, 0], 1 / stacked_geo.rho_b**2
+    )
+
+    # Test stacking with an empty list (should raise ValueError)
+    with self.assertRaisesRegex(ValueError, 'No geometries provided.'):
+      geometry.stack_geometries([])
+
+    # Test stacking with geometries of different mesh sizes
+    # (should raise ValueError)
+    geo_diff_mesh = circular_geometry.build_circular_geometry(
+        Rmaj=1.0, B0=2.0, n_rho=20
+    )  # Different n_rho
+    with self.assertRaisesRegex(
+        ValueError, 'All geometries must have the same mesh.'
+    ):
+      geometry.stack_geometries([geo0, geo_diff_mesh])
+
+    # Test with geometries that has a different geometry type
+    geo_diff_mesh = dataclasses.replace(geo1, geometry_type=3)
+    with self.assertRaisesRegex(
+        ValueError, 'All geometries must have the same geometry type'
+    ):
+      geometry.stack_geometries([geo0, geo_diff_mesh])
 
 
 def _pint_face_to_cell(n_rho, face):
diff --git a/torax/geometry/tests/standard_geometry_test.py b/torax/geometry/tests/standard_geometry_test.py
index 8bc23647..503755cf 100644
--- a/torax/geometry/tests/standard_geometry_test.py
+++ b/torax/geometry/tests/standard_geometry_test.py
@@ -19,7 +19,6 @@
 import jax
 import numpy as np
 from torax.config import build_sim
-from torax.geometry import circular_geometry
 from torax.geometry import geometry
 from torax.geometry import geometry_loader
 from torax.geometry import standard_geometry
@@ -151,20 +150,6 @@ def test_validate_fbt_data_incorrect_L_pQ_shape(self):
     with self.assertRaisesRegex(ValueError, 'Incorrect shape'):
       standard_geometry._validate_fbt_data(LY, L)
 
-  def test_stack_geometries(self):
-    """Test that geometries can be stacked."""
-    geo_0 = circular_geometry.build_circular_geometry(Rmaj=1, B0=5.3)
-    geo_1 = circular_geometry.build_circular_geometry(Rmaj=2, B0=3.7)
-    stacked_geometries = geometry.stack_geometries([geo_0, geo_1])
-    self.assertNotIn('geometry_type', stacked_geometries.keys())
-    self.assertNotIn('torax_mesh', stacked_geometries.keys())
-
-    for key, value in stacked_geometries.items():
-      self.assertEqual(value.shape, (2,) + getattr(geo_0, key).shape)
-
-    np.testing.assert_allclose(stacked_geometries['Rmaj'], np.array([1, 2]))
-    np.testing.assert_allclose(stacked_geometries['B0'], np.array([5.3, 3.7]))
-
 
 def _get_example_L_LY_data(len_psinorm: int, len_times: int):
   LY = {
diff --git a/torax/output.py b/torax/output.py
index ed29bf01..9f8e88b9 100644
--- a/torax/output.py
+++ b/torax/output.py
@@ -16,6 +16,7 @@
 from __future__ import annotations
 
 import dataclasses
+import inspect
 
 from absl import logging
 import chex
@@ -199,7 +200,8 @@ def __init__(
     post_processed_output = [
         state.post_processed_outputs for state in sim_outputs.sim_history
     ]
-    self.geometry = [state.geometry for state in sim_outputs.sim_history]
+    geometries = [state.geometry for state in sim_outputs.sim_history]
+    self.geometry = geometry_lib.stack_geometries(geometries)
     stack = lambda *ys: jnp.stack(ys)
     self.core_profiles: state.CoreProfiles = jax.tree_util.tree_map(
         stack, *core_profiles
@@ -238,6 +240,7 @@ def _pack_into_data_array(
         len(self.rho_norm),
     )
     is_scalar = lambda x: x.ndim == 1 and x.shape == (len(self.times),)
+    is_constant = lambda x: x.ndim == 0
 
     match data:
       case data if is_face_var(data):
@@ -246,6 +249,8 @@ def _pack_into_data_array(
         dims = [TIME, RHO_CELL]
       case data if is_scalar(data):
         dims = [TIME]
+      case data if is_constant(data):
+        dims = []
       case _:
         logging.warning(
             "Unsupported data shape for %s: %s. Skipping persisting.",
@@ -298,7 +303,10 @@ def _get_core_profiles(
     xr_dict[NREF] = self.core_profiles.nref
 
     xr_dict = {
-        name: self._pack_into_data_array(name, data,)
+        name: self._pack_into_data_array(
+            name,
+            data,
+        )
         for name, data in xr_dict.items()
     }
 
@@ -316,7 +324,10 @@ def _save_core_transport(
     xr_dict[V_FACE_EL] = self.core_transport.v_face_el
 
     xr_dict = {
-        name: self._pack_into_data_array(name, data,)
+        name: self._pack_into_data_array(
+            name,
+            data,
+        )
         for name, data in xr_dict.items()
     }
 
@@ -358,23 +369,40 @@ def _save_post_processed_outputs(
     for field_name, data in dataclasses.asdict(
         self.post_processed_outputs
     ).items():
-      xr_dict[field_name] = self._pack_into_data_array(field_name, data,)
-
+      xr_dict[field_name] = self._pack_into_data_array(field_name, data)
     return xr_dict
 
   def _save_geometry(
       self,
   ) -> dict[str, xr.DataArray]:
-    """Saves the geometry to a dict. We skip over all hires quantities."""
+    """Save geometry to a dict. We skip over hires and non-array quantities."""
     xr_dict = {}
-    stacked_geometry = geometry_lib.stack_geometries(self.geometry)
-    for field_name, data in stacked_geometry.items():
-      if "hires" in field_name:
+
+    # Get the variables from dataclass fields.
+    for field_name, data in dataclasses.asdict(self.geometry).items():
+      if "hires" in field_name or not isinstance(data, jax.Array):
         continue
-      data_array = self._pack_into_data_array(field_name, data,)
+      data_array = self._pack_into_data_array(
+          field_name,
+          data,
+      )
       if data_array is not None:
         xr_dict[field_name] = data_array
 
+    # Get variables from property methods
+
+    # Exclude coordinates from geometry since they are at the top DataTree level
+    excluded_names = {RHO_CELL, RHO_FACE, RHO_CELL_NORM, RHO_FACE_NORM}
+    for name, value in inspect.getmembers(type(self.geometry)):
+      if name in excluded_names:
+        continue
+      if isinstance(value, property):
+        property_data = value.fget(self.geometry)
+        if isinstance(property_data, jax.Array):
+          data_array = self._pack_into_data_array(name, property_data)
+          if data_array is not None:
+            xr_dict[name] = data_array
+
     return xr_dict
 
   def simulation_output_to_xr(
@@ -433,7 +461,10 @@ def simulation_output_to_xr(
     # Update dict with flattened StateHistory dataclass containers
     core_profiles_ds = xr.Dataset(self._get_core_profiles(), coords=coords)
     core_transport_ds = xr.Dataset(self._save_core_transport(), coords=coords)
-    core_sources_ds = xr.Dataset(self._save_core_sources(), coords=coords,)
+    core_sources_ds = xr.Dataset(
+        self._save_core_sources(),
+        coords=coords,
+    )
     post_processed_outputs_ds = xr.Dataset(
         self._save_post_processed_outputs(), coords=coords
     )
diff --git a/torax/tests/test_data/test_all_transport_crank_nicolson.nc b/torax/tests/test_data/test_all_transport_crank_nicolson.nc
index 42e87aee..6e6a5243 100644
Binary files a/torax/tests/test_data/test_all_transport_crank_nicolson.nc and b/torax/tests/test_data/test_all_transport_crank_nicolson.nc differ
diff --git a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc
index b19a38e0..96ac432e 100644
Binary files a/torax/tests/test_data/test_all_transport_fusion_qlknn.nc and b/torax/tests/test_data/test_all_transport_fusion_qlknn.nc differ
diff --git a/torax/tests/test_data/test_bohmgyrobohm_all.nc b/torax/tests/test_data/test_bohmgyrobohm_all.nc
index 2af91467..ac12a1c1 100644
Binary files a/torax/tests/test_data/test_bohmgyrobohm_all.nc and b/torax/tests/test_data/test_bohmgyrobohm_all.nc differ
diff --git a/torax/tests/test_data/test_bootstrap.nc b/torax/tests/test_data/test_bootstrap.nc
index 6cec50dd..4e2c7d9d 100644
Binary files a/torax/tests/test_data/test_bootstrap.nc and b/torax/tests/test_data/test_bootstrap.nc differ
diff --git a/torax/tests/test_data/test_bremsstrahlung.nc b/torax/tests/test_data/test_bremsstrahlung.nc
index 0dae532d..7ce10eba 100644
Binary files a/torax/tests/test_data/test_bremsstrahlung.nc and b/torax/tests/test_data/test_bremsstrahlung.nc differ
diff --git a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc
index 360edb7a..2c0ae6f7 100644
Binary files a/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc and b/torax/tests/test_data/test_bremsstrahlung_time_dependent_Zimp.nc differ
diff --git a/torax/tests/test_data/test_cgmheat.nc b/torax/tests/test_data/test_cgmheat.nc
index d0bfb752..6156e9b3 100644
Binary files a/torax/tests/test_data/test_cgmheat.nc and b/torax/tests/test_data/test_cgmheat.nc differ
diff --git a/torax/tests/test_data/test_changing_config_after.nc b/torax/tests/test_data/test_changing_config_after.nc
index b9aed945..8ab6838a 100644
Binary files a/torax/tests/test_data/test_changing_config_after.nc and b/torax/tests/test_data/test_changing_config_after.nc differ
diff --git a/torax/tests/test_data/test_changing_config_before.nc b/torax/tests/test_data/test_changing_config_before.nc
index 0f7b90bd..6c657834 100644
Binary files a/torax/tests/test_data/test_changing_config_before.nc and b/torax/tests/test_data/test_changing_config_before.nc differ
diff --git a/torax/tests/test_data/test_chease.nc b/torax/tests/test_data/test_chease.nc
index c3065722..f30f5c2a 100644
Binary files a/torax/tests/test_data/test_chease.nc and b/torax/tests/test_data/test_chease.nc differ
diff --git a/torax/tests/test_data/test_eqdsk.nc b/torax/tests/test_data/test_eqdsk.nc
index 4aad6862..04f0d9c2 100644
Binary files a/torax/tests/test_data/test_eqdsk.nc and b/torax/tests/test_data/test_eqdsk.nc differ
diff --git a/torax/tests/test_data/test_explicit.nc b/torax/tests/test_data/test_explicit.nc
index c1adf679..d35bb4ac 100644
Binary files a/torax/tests/test_data/test_explicit.nc and b/torax/tests/test_data/test_explicit.nc differ
diff --git a/torax/tests/test_data/test_fixed_dt.nc b/torax/tests/test_data/test_fixed_dt.nc
index 258ac60f..dbe7722a 100644
Binary files a/torax/tests/test_data/test_fixed_dt.nc and b/torax/tests/test_data/test_fixed_dt.nc differ
diff --git a/torax/tests/test_data/test_fusion_power.nc b/torax/tests/test_data/test_fusion_power.nc
index a4aade5d..ff619f5f 100644
Binary files a/torax/tests/test_data/test_fusion_power.nc and b/torax/tests/test_data/test_fusion_power.nc differ
diff --git a/torax/tests/test_data/test_implicit.nc b/torax/tests/test_data/test_implicit.nc
index 95d4cd47..893fd7ad 100644
Binary files a/torax/tests/test_data/test_implicit.nc and b/torax/tests/test_data/test_implicit.nc differ
diff --git a/torax/tests/test_data/test_implicit_short_optimizer.nc b/torax/tests/test_data/test_implicit_short_optimizer.nc
index 0bdbafbe..22ab0d0d 100644
Binary files a/torax/tests/test_data/test_implicit_short_optimizer.nc and b/torax/tests/test_data/test_implicit_short_optimizer.nc differ
diff --git a/torax/tests/test_data/test_iterbaseline_mockup.nc b/torax/tests/test_data/test_iterbaseline_mockup.nc
index da69a598..d5d27806 100644
Binary files a/torax/tests/test_data/test_iterbaseline_mockup.nc and b/torax/tests/test_data/test_iterbaseline_mockup.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_mockup.nc b/torax/tests/test_data/test_iterhybrid_mockup.nc
index 3a95b8d8..8937dcd5 100644
Binary files a/torax/tests/test_data/test_iterhybrid_mockup.nc and b/torax/tests/test_data/test_iterhybrid_mockup.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_newton.nc b/torax/tests/test_data/test_iterhybrid_newton.nc
index 87d1501e..59f79b67 100644
Binary files a/torax/tests/test_data/test_iterhybrid_newton.nc and b/torax/tests/test_data/test_iterhybrid_newton.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc
index 1920272f..355965c5 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc
index d6c20cbd..612ebcaa 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_clip_inputs.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc
index e771de5e..000d082f 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_constant_fraction_impurity_radiation.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc
index ed4eca37..62f6f5b5 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_cyclotron.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc
index d1a9b0ec..caf2d60f 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_ec_linliu.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc
index ca1bbc1e..e0609697 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_eqdsk.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc
index b5c21d0a..38cc4e09 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_mavrin_impurity_radiation.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc
index a2652e17..b7c5eb1d 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_set_pped_tpedratio_nped.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc
index 1a5a0505..fa36e496 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_timedependent_isotopes.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc
index 3d82638a..ab628b3a 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_tungsten.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc
index 00f3bd1b..4cabd11a 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zeffprofile.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc
index bc872f3c..7649fa29 100644
Binary files a/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc and b/torax/tests/test_data/test_iterhybrid_predictor_corrector_zi2.nc differ
diff --git a/torax/tests/test_data/test_iterhybrid_rampup.nc b/torax/tests/test_data/test_iterhybrid_rampup.nc
index 9a669051..796f4d27 100644
Binary files a/torax/tests/test_data/test_iterhybrid_rampup.nc and b/torax/tests/test_data/test_iterhybrid_rampup.nc differ
diff --git a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc
index 730eb5c2..10ca16d8 100644
Binary files a/torax/tests/test_data/test_ne_qlknn_deff_veff.nc and b/torax/tests/test_data/test_ne_qlknn_deff_veff.nc differ
diff --git a/torax/tests/test_data/test_ne_qlknn_defromchie.nc b/torax/tests/test_data/test_ne_qlknn_defromchie.nc
index 369cf220..19985802 100644
Binary files a/torax/tests/test_data/test_ne_qlknn_defromchie.nc and b/torax/tests/test_data/test_ne_qlknn_defromchie.nc differ
diff --git a/torax/tests/test_data/test_ohmic_power.nc b/torax/tests/test_data/test_ohmic_power.nc
index 5c0b6dd7..aed708c6 100644
Binary files a/torax/tests/test_data/test_ohmic_power.nc and b/torax/tests/test_data/test_ohmic_power.nc differ
diff --git a/torax/tests/test_data/test_particle_sources_cgm.nc b/torax/tests/test_data/test_particle_sources_cgm.nc
index 3b92b6d8..f1ea1ef8 100644
Binary files a/torax/tests/test_data/test_particle_sources_cgm.nc and b/torax/tests/test_data/test_particle_sources_cgm.nc differ
diff --git a/torax/tests/test_data/test_particle_sources_constant.nc b/torax/tests/test_data/test_particle_sources_constant.nc
index 531e3a3b..84b3677e 100644
Binary files a/torax/tests/test_data/test_particle_sources_constant.nc and b/torax/tests/test_data/test_particle_sources_constant.nc differ
diff --git a/torax/tests/test_data/test_pc_method_ne.nc b/torax/tests/test_data/test_pc_method_ne.nc
index adb88e52..5da56e75 100644
Binary files a/torax/tests/test_data/test_pc_method_ne.nc and b/torax/tests/test_data/test_pc_method_ne.nc differ
diff --git a/torax/tests/test_data/test_pedestal.nc b/torax/tests/test_data/test_pedestal.nc
index 1017148f..2a55121b 100644
Binary files a/torax/tests/test_data/test_pedestal.nc and b/torax/tests/test_data/test_pedestal.nc differ
diff --git a/torax/tests/test_data/test_prescribed_generic_current_source.nc b/torax/tests/test_data/test_prescribed_generic_current_source.nc
index 06f75038..96b7dc91 100644
Binary files a/torax/tests/test_data/test_prescribed_generic_current_source.nc and b/torax/tests/test_data/test_prescribed_generic_current_source.nc differ
diff --git a/torax/tests/test_data/test_prescribed_timedependent_ne.nc b/torax/tests/test_data/test_prescribed_timedependent_ne.nc
index 142ac7bd..6d2d4272 100644
Binary files a/torax/tests/test_data/test_prescribed_timedependent_ne.nc and b/torax/tests/test_data/test_prescribed_timedependent_ne.nc differ
diff --git a/torax/tests/test_data/test_psi_and_heat.nc b/torax/tests/test_data/test_psi_and_heat.nc
index 9e0d48c5..00a1cdf5 100644
Binary files a/torax/tests/test_data/test_psi_and_heat.nc and b/torax/tests/test_data/test_psi_and_heat.nc differ
diff --git a/torax/tests/test_data/test_psi_heat_dens.nc b/torax/tests/test_data/test_psi_heat_dens.nc
index 239b0bf5..0a742281 100644
Binary files a/torax/tests/test_data/test_psi_heat_dens.nc and b/torax/tests/test_data/test_psi_heat_dens.nc differ
diff --git a/torax/tests/test_data/test_psichease_ip_chease.nc b/torax/tests/test_data/test_psichease_ip_chease.nc
index 9bdc99f1..9875045c 100644
Binary files a/torax/tests/test_data/test_psichease_ip_chease.nc and b/torax/tests/test_data/test_psichease_ip_chease.nc differ
diff --git a/torax/tests/test_data/test_psichease_ip_parameters.nc b/torax/tests/test_data/test_psichease_ip_parameters.nc
index 31f48251..82b217e8 100644
Binary files a/torax/tests/test_data/test_psichease_ip_parameters.nc and b/torax/tests/test_data/test_psichease_ip_parameters.nc differ
diff --git a/torax/tests/test_data/test_psichease_prescribed_johm.nc b/torax/tests/test_data/test_psichease_prescribed_johm.nc
index 721da4a0..131af4b3 100644
Binary files a/torax/tests/test_data/test_psichease_prescribed_johm.nc and b/torax/tests/test_data/test_psichease_prescribed_johm.nc differ
diff --git a/torax/tests/test_data/test_psichease_prescribed_jtot.nc b/torax/tests/test_data/test_psichease_prescribed_jtot.nc
index 02ba595e..1ab508cf 100644
Binary files a/torax/tests/test_data/test_psichease_prescribed_jtot.nc and b/torax/tests/test_data/test_psichease_prescribed_jtot.nc differ
diff --git a/torax/tests/test_data/test_psiequation.nc b/torax/tests/test_data/test_psiequation.nc
index ddfe2233..dae01895 100644
Binary files a/torax/tests/test_data/test_psiequation.nc and b/torax/tests/test_data/test_psiequation.nc differ
diff --git a/torax/tests/test_data/test_qei.nc b/torax/tests/test_data/test_qei.nc
index f97553d9..55910ad2 100644
Binary files a/torax/tests/test_data/test_qei.nc and b/torax/tests/test_data/test_qei.nc differ
diff --git a/torax/tests/test_data/test_qei_chease_highdens.nc b/torax/tests/test_data/test_qei_chease_highdens.nc
index 9e24a99e..d4f5e67f 100644
Binary files a/torax/tests/test_data/test_qei_chease_highdens.nc and b/torax/tests/test_data/test_qei_chease_highdens.nc differ
diff --git a/torax/tests/test_data/test_qlknnheat.nc b/torax/tests/test_data/test_qlknnheat.nc
index 9c0695a0..152029f0 100644
Binary files a/torax/tests/test_data/test_qlknnheat.nc and b/torax/tests/test_data/test_qlknnheat.nc differ
diff --git a/torax/tests/test_data/test_semiimplicit_convection.nc b/torax/tests/test_data/test_semiimplicit_convection.nc
index 8ae838f4..e2bcbff9 100644
Binary files a/torax/tests/test_data/test_semiimplicit_convection.nc and b/torax/tests/test_data/test_semiimplicit_convection.nc differ
diff --git a/torax/tests/test_data/test_time_dependent_circular_geo.nc b/torax/tests/test_data/test_time_dependent_circular_geo.nc
index de964c68..1ac3b3c1 100644
Binary files a/torax/tests/test_data/test_time_dependent_circular_geo.nc and b/torax/tests/test_data/test_time_dependent_circular_geo.nc differ
diff --git a/torax/tests/test_data/test_timedependence.nc b/torax/tests/test_data/test_timedependence.nc
index c9c5d310..9eeb35ba 100644
Binary files a/torax/tests/test_data/test_timedependence.nc and b/torax/tests/test_data/test_timedependence.nc differ