[wb2] add evaluation of WindVectorRMSE for (a)geostrophic wind

PiperOrigin-RevId: 574623237

Author: shoyer

scripts/evaluate.py

@@ -217,28 +217,29 @@
 def _wind_vector_rmse():
   """Defines Wind Vector RMSEs if U/V components are in variables."""
   wind_vector_rmse = []
-  if (
-      'u_component_of_wind' in VARIABLES.value
-      and 'v_component_of_wind' in VARIABLES.value
-  ):
-    wind_vector_rmse.append(
-        metrics.WindVectorRMSE(
-            u_name='u_component_of_wind',
-            v_name='v_component_of_wind',
-            vector_name='wind_vector',
-        )
-    )
-  if (
-      '10m_u_component_of_wind' in VARIABLES.value
-      and '10m_v_component_of_wind' in VARIABLES.value
-  ):
-    wind_vector_rmse.append(
-        metrics.WindVectorRMSE(
-            u_name='10m_u_component_of_wind',
-            v_name='10m_v_component_of_wind',
-            vector_name='10m_wind_vector',
-        )
-    )
+  available = set(VARIABLES.value).union(DERIVED_VARIABLES.value)
+  for u_name, v_name, vector_name in [
+      ('u_component_of_wind', 'v_component_of_wind', 'wind_vector'),
+      ('10m_u_component_of_wind', '10m_v_component_of_wind', '10m_wind_vector'),
+      (
+          'u_component_of_geostrophic_wind',
+          'v_component_of_geostrophic_wind',
+          'geostrophic_wind_vector',
+      ),
+      (
+          'u_component_of_ageostrophic_wind',
+          'v_component_of_ageostrophic_wind',
+          'ageostrophic_wind_vector',
+      ),
+  ]:
+    if u_name in available and v_name in available:
+      wind_vector_rmse.append(
+          metrics.WindVectorRMSE(
+              u_name=u_name,
+              v_name=v_name,
+              vector_name=vector_name,
+          )
+      )
   return wind_vector_rmse
 
 

scripts/evaluate_test.py

@@ -30,6 +30,11 @@ def _test(self, use_beam=True):
         'u_component_of_wind',
         'v_component_of_wind',
     ]
+    derived_variables = [
+        'wind_speed',
+        'u_component_of_ageostrophic_wind',
+        'v_component_of_ageostrophic_wind',
+    ]
     variables_2d = ['2m_temperature']
     truth = schema.mock_truth_data(
         variables_3d=variables_3d,
@@ -49,7 +54,9 @@ def _test(self, use_beam=True):
         variables_2d=variables_2d,
     )
     climatology = climatology.assign(
-        wind_speed=climatology['u_component_of_wind']
+        wind_speed=climatology['u_component_of_wind'],
+        u_component_of_ageostrophic_wind=climatology['u_component_of_wind'],
+        v_component_of_ageostrophic_wind=climatology['u_component_of_wind'],
     )
 
     truth_path = self.create_tempdir('truth').full_path
@@ -80,7 +87,7 @@ def _test(self, use_beam=True):
         eval_configs=','.join(eval_configs),
         use_beam=use_beam,
         variables=variables_3d + variables_2d,
-        derived_variables=['wind_speed'],
+        derived_variables=derived_variables,
     ):
       evaluate.main([])
 
@@ -91,12 +98,20 @@ def _test(self, use_beam=True):
       with self.subTest(config_name):
         results_path = os.path.join(output_dir, f'{config_name}.nc')
         actual = xarray.open_dataset(results_path)
+        extra_out_vars = [
+            'wind_speed',
+            'wind_vector',
+            'u_component_of_ageostrophic_wind',
+            'v_component_of_ageostrophic_wind',
+            'ageostrophic_wind_vector',
+        ]
         self.assertEqual(
-            set(actual),
-            set(variables_3d + variables_2d + ['wind_speed', 'wind_vector']),
+            set(actual), set(variables_3d + variables_2d + extra_out_vars)
         )
         self.assertEqual(actual['geopotential'].sizes, expected_sizes_3d)
         self.assertEqual(actual['2m_temperature'].sizes, expected_sizes_2d)
+        self.assertIn('wind_vector', actual)
+        self.assertIn('ageostrophic_wind_vector', actual)
 
   def test_in_memory(self):
     self._test(use_beam=False)

weatherbench2/derived_variables.py

@@ -236,7 +236,8 @@ def _geostrophic_wind(
       2 * omega * np.sin(np.deg2rad(geopotential.coords['latitude']))
   )
   # Geostrophic wind is inf on the equator. We don't clip it to ensure that the
-  # user makes an intentional choice about how handle these invalid values.
+  # user makes an intentional choice about how handle these invalid values
+  # (e.g., by evaluating over a region).
   return (
       -_d_dy(geopotential) / coriolis_parameter,
       +_d_dx(geopotential) / coriolis_parameter,
@@ -245,6 +246,7 @@ def _geostrophic_wind(
 
 @dataclasses.dataclass
 class _GeostrophicWindVariable(DerivedVariable):
+  """Base class for geostrophic wind variables."""
   geopotential_name: str = 'geopotential'
 
   @property
@@ -290,9 +292,8 @@ def compute(self, dataset: xr.Dataset) -> xr.DataArray:
 
 
 @dataclasses.dataclass
-class AgeostrophicWindSpeed(DerivedVariable):
-  """Calculate ageostrophic wind speed."""
-
+class _AgeostrophicWindVariable(DerivedVariable):
+  """Base class for ageostrophic wind variables."""
   u_name: str = 'u_component_of_wind'
   v_name: str = 'v_component_of_wind'
   geopotential_name: str = 'geopotential'
@@ -306,13 +307,35 @@ def core_dims(self) -> t.Tuple[t.Tuple[t.List[str], ...], t.List[str]]:
     lon_lat = ['longitude', 'latitude']
     return (lon_lat, lon_lat, lon_lat), lon_lat
 
+
+class AgeostrophicWindSpeed(_AgeostrophicWindVariable):
+  """Calculate ageostrophic wind speed."""
+
   def compute(self, dataset: xr.Dataset) -> xr.DataArray:
     u = dataset[self.u_name]
     v = dataset[self.v_name]
     u_geo, v_geo = _geostrophic_wind(dataset[self.geopotential_name])
     return np.sqrt((u - u_geo) ** 2 + (v - v_geo) ** 2)
 
 
+class UComponentOfAgeostrophicWind(_AgeostrophicWindVariable):
+  """East-west component of geostrophic wind."""
+
+  def compute(self, dataset: xr.Dataset) -> xr.DataArray:
+    u = dataset[self.u_name]
+    u_geo, _ = _geostrophic_wind(dataset[self.geopotential_name])
+    return u - u_geo
+
+
+class VComponentOfAgeostrophicWind(_AgeostrophicWindVariable):
+  """North-south component of geostrophic wind."""
+
+  def compute(self, dataset: xr.Dataset) -> xr.DataArray:
+    v = dataset[self.v_name]
+    _, v_geo = _geostrophic_wind(dataset[self.geopotential_name])
+    return v - v_geo
+
+
 @dataclasses.dataclass
 class LapseRate(DerivedVariable):
   """Compute lapse rate in temperature."""
@@ -704,6 +727,8 @@ def compute(self, dataset: xr.Dataset):
     'u_component_of_geostrophic_wind': UComponentOfGeostrophicWind(),
     'v_component_of_geostrophic_wind': VComponentOfGeostrophicWind(),
     'ageostrophic_wind_speed': AgeostrophicWindSpeed(),
+    'u_component_of_ageostrophic_wind': UComponentOfAgeostrophicWind(),
+    'v_component_of_ageostrophic_wind': VComponentOfAgeostrophicWind(),
     'lapse_rate': LapseRate(),
     'total_column_vapor': TotalColumnWater(
         water_species_name='specific_humidity'

weatherbench2/metrics.py

@@ -140,7 +140,7 @@ class WindVectorRMSE(Metric):
 
   Attributes:
     u_name: Name of U component.
-    v_name: Name of v component.
+    v_name: Name of V component.
     vector_name: Name of wind vector to be computed.
   """
 
@@ -793,3 +793,7 @@ def compute_chunk(
     """Energy score skill, averaged over space, for a time chunk of data."""
     _get_n_ensemble(forecast, self.ensemble_dim)  # Will raise if no ensembles.
     return _spatial_average_l2_norm(forecast - truth).mean(self.ensemble_dim)
+
+
+# TODO(shoyer): Consider adding WindVectorEnergyScore based on a pair of wind
+# components, as a sort of probabilistic variant of WindVectorRMSE.