Remove unused attributes from Geometry.

PiperOrigin-RevId: 720583158
google-deepmind · Jan 28, 2025 · 0631499 · 0631499
1 parent 9f8c120
commit 0631499
Show file tree

Hide file tree

Showing 4 changed files with 2 additions and 12 deletions.
diff --git a/torax/geometry/circular_geometry.py b/torax/geometry/circular_geometry.py
@@ -257,13 +257,10 @@ def build_circular_geometry(
       # Set the circular geometry-specific params.
       elongation=elongation,
       elongation_face=elongation_face,
-      volume_hires=volume_hires,
-      area_hires=area_hires,
       spr_hires=spr_hires,
       rho_hires_norm=rho_hires_norm,
       rho_hires=rho_hires,
       elongation_hires=elongation_hires,
-      vpr_hires=vpr_hires,
       # always initialize Phibdot as zero. It will be replaced once both geo_t
       # and geo_t_plus_dt are provided, and set to be the same for geo_t and
       # geo_t_plus_dt for each given time interval.

diff --git a/torax/geometry/geometry.py b/torax/geometry/geometry.py
@@ -165,12 +165,9 @@ class Geometry:
   Rin_face: chex.Array
   Rout: chex.Array
   Rout_face: chex.Array
-  volume_hires: chex.Array
-  area_hires: chex.Array
   spr_hires: chex.Array
   rho_hires_norm: chex.Array
   rho_hires: chex.Array
-  vpr_hires: chex.Array
   Phibdot: chex.Array
   _z_magnetic_axis: chex.Array | None
 

diff --git a/torax/geometry/standard_geometry.py b/torax/geometry/standard_geometry.py
@@ -931,7 +931,6 @@ def build_standard_geometry(
   # V' for volume integrations on face grid
   vpr_face = rhon_interpolation_func(rho_face_norm, vpr)
   # V' for volume integrations on cell grid
-  vpr_hires = rhon_interpolation_func(rho_hires_norm, vpr)
   vpr = rhon_interpolation_func(rho_norm, vpr)
 
   # S' for area integrals on face grid
@@ -999,11 +998,9 @@ def build_standard_geometry(
   g2g3_over_rhon = rhon_interpolation_func(rho_norm, g2g3_over_rhon)
 
   volume_face = rhon_interpolation_func(rho_face_norm, volume_intermediate)
-  volume_hires = rhon_interpolation_func(rho_hires_norm, volume_intermediate)
   volume = rhon_interpolation_func(rho_norm, volume_intermediate)
 
   area_face = rhon_interpolation_func(rho_face_norm, area_intermediate)
-  area_hires = rhon_interpolation_func(rho_hires_norm, area_intermediate)
   area = rhon_interpolation_func(rho_norm, area_intermediate)
 
   return StandardGeometry(
@@ -1052,12 +1049,9 @@ def build_standard_geometry(
       delta_lower_face=delta_lower_face,
       elongation=elongation,
       elongation_face=elongation_face,
-      volume_hires=volume_hires,
-      area_hires=area_hires,
       spr_hires=spr_hires,
       rho_hires_norm=rho_hires_norm,
       rho_hires=rho_hires,
-      vpr_hires=vpr_hires,
       # always initialize Phibdot as zero. It will be replaced once both geo_t
       # and geo_t_plus_dt are provided, and set to be the same for geo_t and
       # geo_t_plus_dt for each given time interval.

diff --git a/torax/output.py b/torax/output.py
@@ -378,6 +378,8 @@ def _save_geometry(
     xr_dict = {}
     stacked_geometry = geometry_lib.stack_geometries(self.geometry)
     for field_name, data in stacked_geometry.items():
+      if "hires" in field_name:
+        continue
       data_array = self._pack_into_data_array(field_name, data,)
       if data_array is not None:
         xr_dict[field_name] = data_array