Changes after review

env/environment.yml

@@ -18,3 +18,4 @@ dependencies:
   - xarray
   - dask
   - cdsapi
+  - scikit-learn

jobs/tools/ICON_to_point.py

@@ -4,15 +4,9 @@
 from scipy import argmin
 import argparse
 
-# TODO: re-use the interpolator & sampler when the function is called with a list of files.
-#       I've tested some options, but I fail to achieve better than linear scaling of the
-#       performance, whereas the second (etc.) calls to the function can entirely skip
-#       recomputing the interpolation weights etc. Probably just a limitation of this method
-#       and a limitation of passing around and merging/concatenating xarray data.
 
-
-def get_horizontal_distances(longitude, latitude, ICON_grid_path, k=5):
-    '''
+def get_horizontal_distances(longitude, latitude, icon_grid_path, k=5):
+    """
     Get horizontal distances between points and their k nearest
     neighbours on the ICON grid using a quick BallTree algorithm
 
@@ -24,7 +18,7 @@ def get_horizontal_distances(longitude, latitude, ICON_grid_path, k=5):
     latitude : list or 1D np.array 
         e.g., [52] or np.array([52,53,54])
 
-    ICON_grid_path : str
+    icon_grid_path : str
         Contains the path to the ICON grid
 
     k : int, default is 5
@@ -39,15 +33,15 @@ def get_horizontal_distances(longitude, latitude, ICON_grid_path, k=5):
     indices: 2D np.array
         Contains the indices to the ICON grid cells of the corresponding
         nearest neighbours
-    '''
+    """
     # Get ICON grid specifics
-    ICON_GRID = xr.open_dataset(ICON_grid_path)
-    clon = ICON_GRID.clon.values
-    clat = ICON_GRID.clat.values
+    icon_grid = xr.open_dataset(icon_grid_path)
+    clon = icon_grid.clon.values
+    clat = icon_grid.clat.values
 
     # Generate BallTree
-    ICON_lat_lon = np.column_stack([clat, clon])
-    tree = BallTree(ICON_lat_lon, metric='haversine')
+    icon_lat_lon = np.column_stack([clat, clon])
+    tree = BallTree(icon_lat_lon, metric='haversine')
 
     # Query BallTree
     target_lat_lon = np.column_stack(
@@ -76,12 +70,11 @@ def get_horizontal_distances(longitude, latitude, ICON_grid_path, k=5):
 
     return distances, indices
 
-
-def get_nearestvertical_distances(model_topography, model_levels,
-                                  station_ground_elevation,
-                                  station_altitude_over_ground,
-                                  interpolation_strategy):
-    '''
+def get_nearest_vertical_distances(model_topography, model_levels,
+                                   base_height_msl,
+                                   inlet_height_agl,
+                                   interpolation_strategy):
+    """
     Get the 2 nearest distances between ICON grid points and specified
     station altitudes
 
@@ -93,14 +86,12 @@ def get_nearestvertical_distances(model_topography, model_levels,
     model_levels : 2D np.array
         Dimensions [ICON_heights, number_of_samples]
 
-    station_ground_elevation : list or 1D np.array
+    base_height_msl : list or 1D np.array
         e.g., [20,] or np.array([72,180,40])
-        This is the elevation of the *ground over mean sea level*
 
-    station_altitude_over_ground : list or 1D np.array
+    inlet_height_agl : list or 1D np.array
         e.g., [15,] or np.array([15, 21, 42])
-        This is the altitude of the *station above the ground*
-    
+
     interpolation_strategy : list of strings
         e.g., ['ground',] or ['ground','mountain','ground']
         Can be 'ground' or 'mountain', or 'middle' (the latter is between the ground and mountain approach)
@@ -116,19 +107,19 @@ def get_nearestvertical_distances(model_topography, model_levels,
     vertical_indices: 3D np.array
         Contains the indices to the ICON height levels of the corresponding 2
         nearest neighbour levels
-    '''
+    """
     # Get the target sampling altitude with a list comprehension
     target_altitude = [
         model_topography.isel({
             "station": i
         }).values +
-        station_altitude_over_ground[i] if strategy == 'ground' else
-        np.repeat(station_ground_elevation[i], model_topography.shape[1]) +
-        station_altitude_over_ground[i] if strategy == 'mountain' else
-        np.repeat(station_ground_elevation[i], model_topography.shape[1]) / 2 +
+        inlet_height_agl[i] if strategy == 'ground' else
+        np.repeat(base_height_msl[i], model_topography.shape[1]) +
+        inlet_height_agl[i] if strategy == 'mountain' else
+        np.repeat(base_height_msl[i], model_topography.shape[1]) / 2 +
         model_topography.isel({
             "station": i
-        }).values / 2 + station_altitude_over_ground[i]
+        }).values / 2 + inlet_height_agl[i]
         # if strategy=='middle'
         for (i, strategy) in enumerate(interpolation_strategy)
     ]
@@ -149,17 +140,16 @@ def get_nearestvertical_distances(model_topography, model_levels,
 
     return np.abs(vertical_distances).T, vertical_indices.T
 
-
-def ICON_to_point(longitude,
+def icon_to_point(longitude,
                   latitude,
-                  station_ground_elevation,
-                  station_altitude_over_ground,
-                  ICON_field_path,
-                  ICON_grid_path,
+                  inlet_height_agl,
+                  base_height_msl,
+                  icon_field_path,
+                  icon_grid_path,
                   interpolation_strategy,
                   k=5,
                   field_name=None):
-    '''
+    """
     Function to interpolate ICON fields to point locations
     
     Parameters
@@ -170,18 +160,22 @@ def ICON_to_point(longitude,
     latitude : list or 1D np.array 
         e.g., [52,] or np.array([52,53,54])
 
-    station_ground_elevation : list or 1D np.array
+    inlet_height_agl : list or 1D np.array
         e.g., [20,] or np.array([72,180,40])
-        This is the elevation of the *ground over mean sea level*
+        This is the height of the *base station over mean sea level*
+            (e.g., for Cabau: base_height_msl=0, 
+                              inlet_height_agl=27)
 
-    station_altitude_over_ground : list or 1D np.array
+    base_height_msl : list or 1D np.array
         e.g., [15,] or np.array([15, 21, 42])
         This is the altitude of the *station above the ground*
+            (e.g., for Jungfraujoch: base_height_msl=3850, 
+                                     inlet_height_agl=5)
 
-    ICON_field_path : str
+    icon_field_path : str
         Contains the path to the unstructured ICON output
 
-    ICON_grid_path : str
+    icon_grid_path : str
         Contains the path to the ICON grid
     
     interpolation_strategy : list of strings
@@ -204,37 +198,37 @@ def ICON_to_point(longitude,
         An Xarray dataset organised by 'station', containing the original
         input specifications, and the vertically and horizontally interpolated
         values
-    '''
+    """
 
     # Load dataset
-    ICON_field = xr.open_dataset(ICON_field_path)
+    icon_field = xr.open_dataset(icon_field_path)
     # Get dimension names
-    icon_heights = ICON_field.z_mc.dims[
+    icon_heights = icon_field.z_mc.dims[
         0]  # Dimension name (something like "heights_5")
-    icon_cells = ICON_field.z_mc.dims[
+    icon_cells = icon_field.z_mc.dims[
         1]  # Dimension name (something like "ncells")
-    ICON_field[icon_cells] = ICON_field[
+    icon_field[icon_cells] = icon_field[
         icon_cells]  # Explicitly assign 'ncells'
 
     # --- Horizontal grid selection & interpolation weights
     # Get k nearest horizontal distances (for use in inverse distance weighing)
-    horizontal_distances, ICON_grid_indices = get_horizontal_distances(
-        longitude, latitude, ICON_grid_path, k=k)
+    horizontal_distances, icon_grid_indices = get_horizontal_distances(
+        longitude, latitude, icon_grid_path, k=k)
 
     horizontal_interp = 1 / horizontal_distances / (
         1 / horizontal_distances).sum(axis=1, keepdims=True)
     weights_horizontal = xr.DataArray(horizontal_interp,
                                       dims=["station", icon_cells])
-    ind_X = xr.DataArray(ICON_grid_indices, dims=["station", icon_cells])
-    ICON_subset = ICON_field.isel({icon_cells: ind_X})
+    ind_X = xr.DataArray(icon_grid_indices, dims=["station", icon_cells])
+    icon_subset = icon_field.isel({icon_cells: ind_X})
 
     # --- Vertical level selection & interpolation weights
     # Get 2 nearest vertical distances (for use in linear interpolation)
-    model_topography = ICON_subset.z_ifc[-1]
-    model_levels = ICON_subset.z_mc
-    vertical_distances, ICON_level_indices = get_nearestvertical_distances(
-        model_topography, model_levels, station_ground_elevation,
-        station_altitude_over_ground, interpolation_strategy)
+    model_topography = icon_subset.z_ifc[-1]
+    model_levels = icon_subset.z_mc
+    vertical_distances, icon_level_indices = get_nearest_vertical_distances(
+        model_topography, model_levels, inlet_height_agl,
+        base_height_msl, interpolation_strategy)
 
     vertical_interp = vertical_distances[:, :, ::-1] / (vertical_distances.sum(
         axis=-1, keepdims=True))
@@ -244,32 +238,32 @@ def ICON_to_point(longitude,
 
     weights_vertical = xr.DataArray(vertical_interp,
                                     dims=["ncells", "station", icon_heights])
-    ind_Z = xr.DataArray(ICON_level_indices,
+    ind_Z = xr.DataArray(icon_level_indices,
                          dims=["ncells", "station", icon_heights])
 
     # --- Generate output
     # Subset the ICON field if we want only a few fields of output
     if field_name is not None:
-        ICON_subset = ICON_subset[field_name]
+        icon_subset = icon_subset[field_name]
     # Include the input station parameters in the output
     ds = xr.Dataset({
         'longitude': (['station'], longitude),
         'latitude': (['station'], latitude),
-        'station_ground_elevation': (['station'], station_ground_elevation),
-        'station_altitude_over_ground':
-        (['station'], station_altitude_over_ground),
+        'inlet_height_agl': (['station'], inlet_height_agl),
+        'base_height_msl':
+        (['station'], base_height_msl),
         'interpolation_strategy': (['station'], interpolation_strategy)
     })
     # Perform the interpolations
-    ICON_subset = ICON_subset.isel({icon_heights: ind_Z})
-    ICON_out = ICON_subset.weighted(weights_vertical.fillna(0)).sum(
+    icon_subset = icon_subset.isel({icon_heights: ind_Z})
+    icon_out = icon_subset.weighted(weights_vertical.fillna(0)).sum(
         dim=icon_heights,
         skipna=True).weighted(weights_horizontal).sum(dim=icon_cells)
-    ICON_out = ICON_out.where(
+    icon_out = icon_out.where(
         ~(weights_vertical.sum(dim=[icon_cells, icon_heights],
                                skipna=False)).isnull()
     )  # Remove out of bounds values where weights_vertical has NaNs
-    return xr.merge([ICON_out, ds])
+    return xr.merge([icon_out, ds])
 
 
 if __name__ == '__main__':
@@ -302,12 +296,12 @@ def ICON_to_point(longitude,
         'Station altitude over surface [absolute height asl: elevation+altitude]'
     )
     parser.add_argument('-fields',
-                        dest='ICON_field',
+                        dest='icon_field',
                         default=None,
                         type=str,
                         help='The ICON output fields')
     parser.add_argument('-grid',
-                        dest='ICON_grid',
+                        dest='icon_grid',
                         default=None,
                         type=str,
                         help='The ICON grid dynamic grid file')
@@ -340,20 +334,20 @@ def ICON_to_point(longitude,
     args = parser.parse_args()
 
     # Example run (note: most inputs should be lists, and the performance is optimized for these lists!)
-    output = ICON_to_point(longitude=[
+    output = icon_to_point(longitude=[
         args.longitude,
     ],
                            latitude=[
                                args.latitude,
                            ],
-                           station_ground_elevation=[
+                           inlet_height_agl=[
                                args.elevation,
                            ],
-                           station_altitude_over_ground=[
+                           base_height_msl=[
                                args.altitude,
                            ],
-                           ICON_field_path=args.ICON_field,
-                           ICON_grid_path=args.ICON_grid,
+                           icon_field_path=args.icon_field,
+                           icon_grid_path=args.icon_grid,
                            interpolation_strategy=[
                                args.strategy,
                            ],