Cache grid properties but not value properties of CellVariables.

The grid will be fixed for a whole simulation, so the cell properties can be cached. The left and right face values are not necessarily fixed for the whole simulation, so should not be cached.

Also fixed the fact that right_face_value was a property and left_face_value was a method.

PiperOrigin-RevId: 889163086

Author: theo-brown

torax/_src/fvm/cell_variable.py

@@ -139,12 +139,12 @@ def cell_centers(self) -> jt.Float[chex.Array, 'cell']:
     """Locations of the cell centers."""
     return (self.face_centers[..., 1:] + self.face_centers[..., :-1]) / 2.0
 
-  @property
+  @functools.cached_property
   def cell_widths(self) -> jt.Float[chex.Array, 'cell']:
     """Size of each cell."""
     return jnp.diff(self.face_centers)
 
-  @property
+  @functools.cached_property
   def cell_spacings(self) -> jt.Float[chex.Array, 'cell-1']:
     """Spacing between each cell."""
     return jnp.diff(self.cell_centers)
@@ -272,6 +272,7 @@ def constrained_grad(
     right = jnp.expand_dims(right_grad, axis=0)
     return jnp.concatenate([left, inner_grad, right])
 
+  @property
   def left_face_value(self) -> jt.Float[chex.Array, '']:
     """Calculates the value of the leftmost face."""
     if self.left_face_constraint is not None:
@@ -284,7 +285,7 @@ def left_face_value(self) -> jt.Float[chex.Array, '']:
       value = self.value[..., 0:1]
     return value
 
-  @functools.cached_property
+  @property
   def right_face_value(self) -> jt.Float[chex.Array, '']:
     """Calculates the value of the rightmost face."""
     if self.right_face_constraint is not None:
@@ -311,7 +312,7 @@ def face_value(self) -> jt.Float[chex.Array, 'face']:
     )
 
     return jnp.concatenate(
-        [self.left_face_value(), inner, self.right_face_value], axis=-1
+        [self.left_face_value, inner, self.right_face_value], axis=-1
     )
 
   def grad(self) -> jt.Float[chex.Array, 'cell']:
@@ -338,9 +339,6 @@ def __str__(self) -> str:
 
   def cell_plus_boundaries(self) -> jt.Float[chex.Array, 'cell+2']:
     """Returns the value of this variable plus left and right boundaries."""
-    right_value = self.right_face_value
-    left_value = self.left_face_value()
     return jnp.concatenate(
-        [left_value, self.value, right_value],
-        axis=-1,
+        [self.left_face_value, self.value, self.right_face_value], axis=-1
     )

torax/_src/imas_tools/output/core_profiles.py

@@ -246,11 +246,11 @@ def _fill_profiles_1d_grid(
       [[0.0], geometry_slice.rho, [geometry_slice.rho_b]]
   )
   ids.profiles_1d[i].grid.psi = cp_state.psi.cell_plus_boundaries()
-  ids.profiles_1d[i].grid.psi_magnetic_axis = cp_state.psi.left_face_value()[0]
+  ids.profiles_1d[i].grid.psi_magnetic_axis = cp_state.psi.left_face_value[0]
   ids.profiles_1d[i].grid.psi_boundary = cp_state.psi.right_face_value[0]
   ids.profiles_1d[i].grid.rho_pol_norm = np.sqrt(
-      (cp_state.psi.cell_plus_boundaries() - cp_state.psi.left_face_value()[0])
-      / (cp_state.psi.right_face_value[0] - cp_state.psi.left_face_value()[0])
+      (cp_state.psi.cell_plus_boundaries() - cp_state.psi.left_face_value[0])
+      / (cp_state.psi.right_face_value[0] - cp_state.psi.left_face_value[0])
   )
   ids.profiles_1d[i].grid.volume = output.extend_cell_grid_to_boundaries(
       [geometry_slice.volume], np.array([geometry_slice.volume_face])