Improve type annotations

- `utilFcns.write_weather_vars_to_ds`: annotate number of dimensions
- types: Add generic version of the numpy array helper types so you can annotate an array of booleans as `Array1D[np.bool]`
- Add type annotations to `models.hrrr.get_bounds_indices`

Author: nate-kean

test/test_util.py

@@ -33,7 +33,7 @@
     transform_bbox,
     unproject,
     writeArrayToRaster,
-    writeWeatherVarsXarray,
+    write_weather_vars_to_ds,
 )
 from test import TEST_DIR, pushd
 
@@ -917,7 +917,7 @@ def test_UTM_to_WGS84_empty_input():
 
 
 # Test writeWeatherVarsXarray
-def test_writeWeatherVarsXarray(tmp_path):
+def test_write_weather_vars_to_ds(tmp_path):
     """Test writing weather variables to an xarray dataset and NetCDF file."""
     # Mock inputs
     lat = np.random.rand(91, 144) * 180 - 90  # Random latitudes between -90 and 90
@@ -938,7 +938,7 @@ def test_writeWeatherVarsXarray(tmp_path):
     out_path = tmp_path / "test_output.nc"
 
     # Call the function
-    writeWeatherVarsXarray(lat, lon, h, q, p, t, datetime_value, crs, out_path)
+    write_weather_vars_to_ds(lat, lon, h, q, p, t, datetime_value, crs, out_path)
 
     # Check that the file was created
     assert out_path.exists()

test/test_weather_model.py

@@ -403,7 +403,7 @@ def test_ztd(model: MockWeatherModel) -> None:
 
 def test_get_bounds_indices() -> None:
     """Test bounds indices."""
-    snwe = [-10, 10, -10, 10]
+    snwe = (-10, 10, -10, 10)
     ll = np.arange(-20, 20)
     lats, lons = np.meshgrid(ll, ll)
     xmin, xmax, ymin, ymax = get_bounds_indices(snwe, lats, lons)
@@ -415,7 +415,7 @@ def test_get_bounds_indices() -> None:
 
 def test_get_bounds_indices_2() -> None:
     """Test bounds indices."""
-    snwe = [-10, 10, 170, -170]
+    snwe = (-10, 10, 170, -170)
     l = np.arange(-20, 20)
     l2 = (((np.arange(160, 200) + 180) % 360) - 180)
     lats, lons = np.meshgrid(l, l2)
@@ -425,7 +425,7 @@ def test_get_bounds_indices_2() -> None:
 
 def test_get_bounds_indices_2b() -> None:
     """Test bounds indices."""
-    snwe = [-10, 10, 170, 190]
+    snwe = (-10, 10, 170, 190)
     l = np.arange(-20, 20)
     l2 = np.arange(160, 200)
     lats, lons = np.meshgrid(l, l2)
@@ -438,7 +438,7 @@ def test_get_bounds_indices_2b() -> None:
 
 def test_get_bounds_indices_3() -> None:
     """Test bounds indices"""
-    snwe = [-10, 10, -10, 10]
+    snwe = (-10, 10, -10, 10)
     l = np.arange(-20, 20)
     l2 = (((np.arange(160, 200) + 180) % 360) - 180)
     lats, lons = np.meshgrid(l, l2)
@@ -448,7 +448,7 @@ def test_get_bounds_indices_3() -> None:
 
 def test_get_bounds_indices_4() -> None:
     """Test bounds_indices."""
-    snwe = [55, 60, 175, 185]
+    snwe = (55, 60, 175, 185)
     l = np.arange(55, 60, 1)
     l2 = np.arange(175, 185, 1)
     lats, lons = np.meshgrid(l, l2)

tools/RAiDER/models/gmao.py

@@ -12,7 +12,7 @@
     LEVELS_137_HEIGHTS,
 )
 from RAiDER.models.weatherModel import TIME_RES, WeatherModel
-from RAiDER.utilFcns import requests_retry_session, round_date, writeWeatherVarsXarray
+from RAiDER.utilFcns import requests_retry_session, round_date, write_weather_vars_to_ds
 
 
 class GMAO(WeatherModel):
@@ -142,7 +142,7 @@ def _fetch(self, out: Path) -> None:
 
         try:
             # Note that lat/lon gets written twice for GMAO because they are the same as y/x
-            writeWeatherVarsXarray(lat, lon, h, q, p, t, self._time, self._proj, out)
+            write_weather_vars_to_ds(lat, lon, h, q, p, t, self._time, self._proj, out)
         except:
             logger.exception('Unable to save weathermodel to file:')
             raise

tools/RAiDER/models/hrrr.py

@@ -1,14 +1,15 @@
 import datetime as dt
-import os
 from pathlib import Path
+from typing import List, Tuple, Union, cast
 
+from RAiDER.types import BB, Array2D, FloatArray2D
 import geopandas as gpd
+import herbie
+import herbie.accessors
 import numpy as np
 import xarray as xr
-from herbie import Herbie
 from pyproj import CRS, Transformer
 from shapely.geometry import Polygon, box
-from typing import Optional, Union, List, Tuple
 
 from RAiDER.logger import logger
 from RAiDER.models.customExceptions import NoWeatherModelData
@@ -29,13 +30,13 @@
 
 def check_hrrr_dataset_availability(datetime: dt.datetime, model='hrrr') -> bool:
     """Note a file could still be missing within the models valid range."""
-    herbie = Herbie(
+    h = herbie.Herbie(
         datetime,
         model=model,
         product='nat',
         fxx=0,
     )
-    return herbie.grib_source is not None
+    return h.grib_source is not None
 
 
 def download_hrrr_file(ll_bounds, DATE, out: Path, model='hrrr', product='nat', fxx=0, verbose=False) -> None:
@@ -54,7 +55,7 @@ def download_hrrr_file(ll_bounds, DATE, out: Path, model='hrrr', product='nat',
     Returns:
         None, writes data to a netcdf file
     """
-    herbie = Herbie(
+    h = herbie.Herbie(
         DATE.strftime('%Y-%m-%d %H:%M'),
         model=model,
         product=product,
@@ -66,22 +67,27 @@ def download_hrrr_file(ll_bounds, DATE, out: Path, model='hrrr', product='nat',
 
     # Iterate through the list of datasets
     try:
-        ds_list = herbie.xarray(':(SPFH|PRES|TMP|HGT):', verbose=verbose)
+        # cast: Herbie.xarray can return one Dataset or a list
+        ds_list = cast(
+            Union[xr.Dataset, List[xr.Dataset]],
+            h.xarray(':(SPFH|PRES|TMP|HGT):', verbose=verbose),
+        )
+        assert isinstance(ds_list, list)
     except ValueError as e:
         logger.error(e)
         raise
 
     # Note order coord names are request for `test_HRRR_ztd` matters
     # when both coord names are retreived by Herbie is ds_list possibly in
     # Different orders on different machines; `hybrid` is what is expected for the test.
-    ds_list_filt_0 = [ds for ds in ds_list if 'hybrid' in ds._coord_names]
-    ds_list_filt_1 = [ds for ds in ds_list if 'isobaricInhPa' in ds._coord_names]
-    if ds_list_filt_0:
+    ds_list_filt_0 = [ds for ds in ds_list if 'hybrid' in ds.coords]
+    ds_list_filt_1 = [ds for ds in ds_list if 'isobaricInhPa' in ds.coords]
+    if len(ds_list_filt_0) > 0:
         ds_out = ds_list_filt_0[0]
         coord = 'hybrid'
     #  I do not think that this coord name will result in successful processing nominally as variables are
     #  gh,gribfile_projection for test_HRRR_ztd
-    elif ds_list_filt_1:
+    elif len(ds_list_filt_1) > 0:
         ds_out = ds_list_filt_1[0]
         coord = 'isobaricInhPa'
     else:
@@ -91,22 +97,25 @@ def download_hrrr_file(ll_bounds, DATE, out: Path, model='hrrr', product='nat',
     try:
         x_min, x_max, y_min, y_max = get_bounds_indices(
             ll_bounds,
-            ds_out.latitude.to_numpy(),
-            ds_out.longitude.to_numpy(),
+            ds_out['latitude'].to_numpy(),
+            ds_out['longitude'].to_numpy(),
         )
     except NoWeatherModelData as e:
         logger.error(e)
-        logger.error('lat/lon bounds: %s', ll_bounds)
-        logger.error('Weather model is {}'.format(model))
+        logger.error(f'lat/lon bounds: {ll_bounds}')
+        logger.error(f'Weather model is {model}')
         raise
 
     # bookkeepping
     ds_out = ds_out.rename({'gh': 'z', coord: 'levels'})
 
+    # Herbie injects brainworms into all its xarray Datasets
+    crs = cast(herbie.accessors.HerbieAccessor, ds_out.herbie).crs
+
     # projection information
     ds_out['proj'] = 0
-    for k, v in CRS.from_user_input(ds_out.herbie.crs).to_cf().items():
-        ds_out.proj.attrs[k] = v
+    for k, v in CRS.from_user_input(crs).to_cf().items():
+        ds_out['proj'].attrs[k] = v
     for var in ds_out.data_vars:
         ds_out[var].attrs['grid_mapping'] = 'proj'
 
@@ -134,40 +143,37 @@ def download_hrrr_file(ll_bounds, DATE, out: Path, model='hrrr', product='nat',
     ds_sub.to_netcdf(out)
 
 
-def get_bounds_indices(SNWE, lats, lons):
+def get_bounds_indices(snwe: BB.SNWE, lats: FloatArray2D, lons: FloatArray2D) -> BB.WSEN:
     """Convert SNWE lat/lon bounds to index bounds."""
     # Unpack the bounds and find the relevent indices
-    S, N, W, E = SNWE
+    s, n, w, e = snwe
 
     # Need to account for crossing the international date line
-    if W < E:
-        m1 = (S <= lats) & (N >= lats) & (W <= lons) & (E >= lons)
-    else:
+    if w >= e:
         raise ValueError(
-            'Longitude is either flipped or you are crossing the international date line;'
-            + 'if the latter please give me longitudes from 0-360'
+            'Longitude is either flipped or you are crossing the international date line; '
+            'if the latter please give me longitudes from 0-360'
         )
 
-    if np.sum(m1) == 0:
+    m1: Array2D[np.bool] = (s <= lats) & (n >= lats) & (w <= lons) & (e >= lons)
+    if np.sum(m1) == 0:  # All False
         lons = np.mod(lons, 360)
-        W, E = np.mod([W, E], 360)
-        m1 = (S <= lats) & (N >= lats) & (W <= lons) & (E >= lons)
-        if np.sum(m1) == 0:
+        w, e = np.mod([w, e], 360)
+        # try again
+        m1 = (s <= lats) & (n >= lats) & (w <= lons) & (e >= lons)
+        if np.sum(m1) == 0:  # All False
             raise NoWeatherModelData('Area of Interest has no overlap with the HRRR model available extent')
 
     # Y extent
-    shp = lats.shape
-    m1_y = np.argwhere(np.sum(m1, axis=1) != 0)
-    y_min = max(m1_y[0][0], 0)
-    y_max = min(m1_y[-1][0], shp[0])
-    m1_y = None
+    m1_y: Array2D[np.integer] = np.argwhere(np.sum(m1, axis=1) != 0)
+    y_min: int = max(m1_y[0][0], 0)
+    y_max: int = min(m1_y[-1][0], lats.shape[0])
+    del m1_y  # big
 
     # X extent
-    m1_x = np.argwhere(np.sum(m1, axis=0) != 0)
-    x_min = max(m1_x[0][0], 0)
-    x_max = min(m1_x[-1][0], shp[1])
-    m1_x = None
-    m1 = None
+    m1_x: Array2D[np.integer] = np.argwhere(np.sum(m1, axis=0) != 0)
+    x_min: int = max(m1_x[0][0], 0)
+    x_max: int = min(m1_x[-1][0], lats.shape[1])
 
     return x_min, x_max, y_min, y_max
 

tools/RAiDER/models/merra2.py

@@ -11,7 +11,7 @@
     LEVELS_137_HEIGHTS,
 )
 from RAiDER.models.weatherModel import WeatherModel
-from RAiDER.utilFcns import writeWeatherVarsXarray
+from RAiDER.utilFcns import write_weather_vars_to_ds
 
 
 # Path to Netrc file, can be controlled by env var
@@ -137,7 +137,7 @@ def _fetch(self, out: Path) -> None:
         ].data.squeeze()
 
         try:
-            writeWeatherVarsXarray(lat, lon, h, q, p, t, time, self._proj, out_path=out)
+            write_weather_vars_to_ds(lat, lon, h, q, p, t, time, self._proj, out_path=out)
         except Exception as e:
             logger.debug(e)
             logger.exception('MERRA-2: Unable to save weather model query to file')

tools/RAiDER/models/ncmr.py

@@ -18,7 +18,7 @@
 from RAiDER.utilFcns import (
     read_NCMR_loginInfo,
     show_progress,
-    writeWeatherVarsXarray,
+    write_weather_vars_to_ds,
 )
 
 
@@ -193,7 +193,7 @@ def _download_ncmr_file(self, out: Path, date_time, bounding_box) -> None:
         ########################################################################################################################
 
         try:
-            writeWeatherVarsXarray(lats, lons, hgt, q, p, t, self._time, self._proj, out_path=out)
+            write_weather_vars_to_ds(lats, lons, hgt, q, p, t, self._time, self._proj, out_path=out)
         except:
             logger.exception('Unable to save weathermodel to file')
 

tools/RAiDER/types/__init__.py

@@ -17,16 +17,19 @@
 # Helpers for type annotating the dimensions of a numpy array.
 # When you use these, your type checker will alert you when an annotated array
 # is used in a way inconsistent with its dimensions.
-Array1D = np.ndarray[tuple[int], np.dtype[_ScalarT_co]]
-Array2D = np.ndarray[tuple[int, int], np.dtype[_ScalarT_co]]
-Array3D = np.ndarray[tuple[int, int, int], np.dtype[_ScalarT_co]]
+Array1D = np.ndarray[tuple[int], np.dtype[np._ScalarT_co]]
+Array2D = np.ndarray[tuple[int, int], np.dtype[np._ScalarT_co]]
+Array3D = np.ndarray[tuple[int, int, int], np.dtype[np._ScalarT_co]]
 # ... (repeat the pattern as needed for higher dimensions)
 
 # Any number of dimensions -- when ndim is not able to be known ahead of time
-ArrayND = np.ndarray[tuple[int, ...], np.dtype[_ScalarT_co]]
+ArrayND = np.ndarray[tuple[int, ...], np.dtype[np._ScalarT_co]]
 
 
 FloatArray1D = Array1D[np.floating]
 FloatArray2D = Array2D[np.floating]
 FloatArray3D = Array3D[np.floating]
+# ... (repeat the pattern as needed for higher dimensions)
+
+# Any number of dimensions -- when ndim is not able to be known ahead of time
 FloatArrayND = ArrayND[np.floating]