precommit

Author: kthyng

docs/whats_new.rst

@@ -1,6 +1,11 @@
 :mod:`What's New`
 -----------------
 
+v1.2.0 (September 13, 2023)
+===========================
+* Improvements to interpolation
+
+
 v1.1.4 (January 27, 2023)
 =========================
 * fixed docs to run fully

extract_model/__init__.py

@@ -9,8 +9,8 @@
 import extract_model.accessor  # noqa: F401
 
 from .extract_model import sel2d, sel2dcf, select, selZ  # noqa: F401
-from .utils import filter, order, sub_bbox, sub_grid, guess_model_type  # noqa: F401
 from .preprocessing import preprocess
+from .utils import filter, guess_model_type, order, sub_bbox, sub_grid  # noqa: F401
 
 
 try:

extract_model/extract_model.py

@@ -1,14 +1,20 @@
 """Preprocessing-related functions for model output."""
 
 
+from typing import Optional
+
 import numpy as np
 import xarray as xr
-from typing import Optional
-from .utils import order, guess_model_type
+
 from extract_model.model_type import ModelType
 
+from .utils import guess_model_type, order
 
-def preprocess_roms(ds, grid=None,):
+
+def preprocess_roms(
+    ds,
+    grid=None,
+):
     """Preprocess ROMS model output for use with cf-xarray.
 
     Also fixes any other known issues with model output.
@@ -24,7 +30,7 @@ def preprocess_roms(ds, grid=None,):
     -------
     Same Dataset but with some metadata added and/or altered.
     """
-    
+
     rename = {}
     if "eta_u" in ds.dims:
         rename["eta_u"] = "eta_rho"
@@ -36,7 +42,6 @@ def preprocess_roms(ds, grid=None,):
         rename["eta_psi"] = "eta_v"
     ds = ds.rename(rename)
 
-
     # add axes attributes for dimensions
     dims = [dim for dim in ds.dims if dim.startswith("s_")]
     for dim in dims:
@@ -125,44 +130,48 @@ def preprocess_roms(ds, grid=None,):
         # }
 
         ds.coords["z_rho"] = order(ds["z_rho"])
-        ds.coords["z_rho_u"] = grid.interp(ds.z_rho.chunk({ds.z_rho.cf["X"].name: -1}), "X")
+        ds.coords["z_rho_u"] = grid.interp(
+            ds.z_rho.chunk({ds.z_rho.cf["X"].name: -1}), "X"
+        )
         ds.coords["z_rho_u"].attrs = {
             "long_name": "depth of U-points on vertical RHO grid",
             "time": "ocean_time",
             "field": "z_rho_u, scalar, series",
             "units": "m",
         }
 
-        ds.coords["z_rho_v"] = grid.interp(ds.z_rho.chunk({ds.z_rho.cf["Y"].name: -1}), "Y")
+        ds.coords["z_rho_v"] = grid.interp(
+            ds.z_rho.chunk({ds.z_rho.cf["Y"].name: -1}), "Y"
+        )
         ds.coords["z_rho_v"].attrs = {
             "long_name": "depth of V-points on vertical RHO grid",
             "time": "ocean_time",
             "field": "z_rho_v, scalar, series",
             "units": "m",
         }
 
-        ds.coords["z_rho_psi"] = grid.interp(ds.z_rho_u.chunk({ds.z_rho_u.cf["Y"].name: -1}), "Y")
+        ds.coords["z_rho_psi"] = grid.interp(
+            ds.z_rho_u.chunk({ds.z_rho_u.cf["Y"].name: -1}), "Y"
+        )
         ds.coords["z_rho_psi"].attrs = {
             "long_name": "depth of PSI-points on vertical RHO grid",
             "time": "ocean_time",
             "field": "z_rho_psi, scalar, series",
             "units": "m",
         }
-        
-        # will use this to update coordinate encoding
-        name_dict.update({"filler1": "z_rho_u", "filler2": "z_rho_v", "filler3": "z_rho_psi"})#, "None": "z_w_u", "None": "z_w_v", "None": "z_w_psi"})
-        
 
+        # will use this to update coordinate encoding
+        name_dict.update(
+            {"filler1": "z_rho_u", "filler2": "z_rho_v", "filler3": "z_rho_psi"}
+        )  # , "None": "z_w_u", "None": "z_w_v", "None": "z_w_psi"})
 
     # fix attrs
     # for zname in ["z_rho", "z_w"]:
     for zname in [var for var in ds.coords if "z_rho" in var or "z_w" in var]:
         if zname in ds:
             ds[
                 zname
-            ].attrs = (
-                {}
-            )  # coord inherits from one of the vars going into calculation
+            ].attrs = {}  # coord inherits from one of the vars going into calculation
             ds[zname].attrs["positive"] = "up"
             ds[zname].attrs["units"] = "m"
             ds[zname] = order(ds[zname])
@@ -224,6 +233,7 @@ def preprocess_roms(ds, grid=None,):
 def preprocess_roms_grid(ds):
     # use xgcm
     from xgcm import Grid
+
     coords = {
         "X": {"center": "xi_rho", "inner": "xi_u"},
         "Y": {"center": "eta_rho", "inner": "eta_v"},

extract_model/preprocessing.py

@@ -9,7 +9,9 @@
 import dask
 import numpy as np
 import xarray as xr
+
 from sklearn.neighbors import BallTree
+
 from extract_model.grids.triangular_mesh import UnstructuredGridSubset
 from extract_model.model_type import ModelType
 
@@ -482,10 +484,15 @@ def order(da):
     )
 
 
-def tree_query(lon_coords: xr.DataArray, lat_coords: xr.DataArray, 
-               lons_to_find: np.array, lats_to_find: np.array, k: int = 3) -> Tuple[np.array]:
+def tree_query(
+    lon_coords: xr.DataArray,
+    lat_coords: xr.DataArray,
+    lons_to_find: np.array,
+    lats_to_find: np.array,
+    k: int = 3,
+) -> Tuple[np.array]:
     """Set up and query BallTree for k nearest points
-    
+
     Uses haversine for the metric because we are dealing with lon/lat coordinates.
 
     Parameters
@@ -505,38 +512,40 @@ def tree_query(lon_coords: xr.DataArray, lat_coords: xr.DataArray,
     -------
     Tuple[np.array]
         distances, (iys, ixs) 2D indices for coordinates
-    
+
     Notes
     -----
     Reference: https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.BallTree.html
     """
-    
+
     # create tree
     coords = [lon_coords, lat_coords]
     X = np.stack([np.ravel(c) for c in coords]).T
-    tree = BallTree(np.deg2rad(X), metric='haversine')
+    tree = BallTree(np.deg2rad(X), metric="haversine")
 
     # set up coordinates we want to search for
     coords_to_find = [lons_to_find, lats_to_find]
     X_to_find = np.stack([np.ravel(c) for c in coords_to_find]).T
-    
+
     # query tree
     distances, inds = tree.query(np.deg2rad(X_to_find), k=k)
-    
+
     # convert flat indies to 2D indices
     iys, ixs = np.unravel_index(inds, lon_coords.shape)
-    
+
     return distances, (iys, ixs)
 
 
-def calc_barycentric(x: np.array, y: np.array, xs: np.array, ys: np.array) -> xr.DataArray:
+def calc_barycentric(
+    x: np.array, y: np.array, xs: np.array, ys: np.array
+) -> xr.DataArray:
     """Calculate barycentric weights for npts
-    
+
     Parameters
     ----------
-    x 
+    x
         npts x 1 vector of x locations, can be in lon or projection coordinates.
-    y 
+    y
         npts x 1 vector of y locations, can be in lat or projection coordinates.
     xs
         npts x 3 array of triangle x vertices with which to calculate the barycentric weights for each of npts
@@ -550,19 +559,32 @@ def calc_barycentric(x: np.array, y: np.array, xs: np.array, ys: np.array) -> xr
     """
     # barycentric weights
     # npts x 1 (vectors)
-    L1 = ( (ys[:,1] - ys[:,2])*(x[:] - xs[:,2]) + (xs[:,2] - xs[:,1])*(y[:] - ys[:,2]) )/ \
-        ( (ys[:,1] - ys[:,2])*(xs[:,0] - xs[:,2]) + (xs[:,2] - xs[:,1])*(ys[:,0] - ys[:,2]) )
-    L2 = ( (ys[:,2] - ys[:,0])*(x[:] - xs[:,2]) + (xs[:,0] - xs[:,2])*(y[:] - ys[:,2]) )/ \
-        ( (ys[:,1] - ys[:,2])*(xs[:,0] - xs[:,2]) + (xs[:,2] - xs[:,1])*(ys[:,0] - ys[:,2]))
+    L1 = (
+        (ys[:, 1] - ys[:, 2]) * (x[:] - xs[:, 2])
+        + (xs[:, 2] - xs[:, 1]) * (y[:] - ys[:, 2])
+    ) / (
+        (ys[:, 1] - ys[:, 2]) * (xs[:, 0] - xs[:, 2])
+        + (xs[:, 2] - xs[:, 1]) * (ys[:, 0] - ys[:, 2])
+    )
+    L2 = (
+        (ys[:, 2] - ys[:, 0]) * (x[:] - xs[:, 2])
+        + (xs[:, 0] - xs[:, 2]) * (y[:] - ys[:, 2])
+    ) / (
+        (ys[:, 1] - ys[:, 2]) * (xs[:, 0] - xs[:, 2])
+        + (xs[:, 2] - xs[:, 1]) * (ys[:, 0] - ys[:, 2])
+    )
     L3 = 1 - L1 - L2
 
-    lam = xr.DataArray(dims=("npts","triangle"), data=np.vstack((L1, L2, L3)).T)
-    
+    lam = xr.DataArray(dims=("npts", "triangle"), data=np.vstack((L1, L2, L3)).T)
+
     return lam
 
 
 def interp_with_barycentric(da, ixs, iys, lam):
-    vector = da.cf.isel(X=xr.DataArray(ixs, dims=("npts","triangle")), Y=xr.DataArray(iys, dims=("npts","triangle")))
+    vector = da.cf.isel(
+        X=xr.DataArray(ixs, dims=("npts", "triangle")),
+        Y=xr.DataArray(iys, dims=("npts", "triangle")),
+    )
     with xr.set_options(keep_attrs=True):
         da = xr.dot(vector, lam, dims=("triangle"))
 
@@ -576,7 +598,7 @@ def interp_with_barycentric(da, ixs, iys, lam):
 
             # add vertical coords into da
             da = da.assign_coords({zkey: da_vert})
-            
+
     # add "X" axis to npts
     da["npts"] = ("npts", da.npts.values, {"axis": "X"})
 

extract_model/utils.py

@@ -7,8 +7,7 @@
 import pytest
 import xarray as xr
 
-from extract_model import utils
-from extract_model import preprocessing
+from extract_model import preprocessing, utils
 from extract_model.grids.triangular_mesh import UnstructuredGridSubset
 
 

tests/grids/test_triangular_mesh.py

@@ -46,7 +46,9 @@ def test_2dsel():
     assert np.allclose(da_sel2d_check.squeeze(), da_check)
 
     da_test, kwargs_out = da.em.sel2dcf(
-        longitude=lon_comp, latitude=lat_comp, return_info=True, #distances_name="distance"
+        longitude=lon_comp,
+        latitude=lat_comp,
+        return_info=True,  # distances_name="distance"
     )
     assert np.allclose(da_sel2d[varname], da_test[varname])
     assert np.allclose(kwargs_out_sel2d_acc_check["distances"], kwargs_out["distances"])

tests/test_accessor.py

@@ -65,7 +65,9 @@ def test_hor_interp_no_xesmf():
     XESMF_AVAILABLE = em.extract_model.XESMF_AVAILABLE
     em.extract_model.XESMF_AVAILABLE = False
     with pytest.raises(ModuleNotFoundError):
-        em.select(da, longitude=longitude, latitude=latitude, T=0.5, horizontal_interp=True)
+        em.select(
+            da, longitude=longitude, latitude=latitude, T=0.5, horizontal_interp=True
+        )
     em.extract_model.XESMF_AVAILABLE = XESMF_AVAILABLE
 
 
@@ -95,7 +97,7 @@ def test_sel2d(model):
         inputs = {
             da.cf["longitude"].name: lon_comp,
             da.cf["latitude"].name: lat_comp,
-            "return_info": True
+            "return_info": True,
         }
         da_sel2d, kwargs_out = em.sel2d(da, **inputs)
         da_check = da.cf.isel(X=i, Y=j)
@@ -369,9 +371,7 @@ def test_sel2d_simple_2D():
     assert ds_out.lon == 3
     assert ds_out.lat == 7
 
-    ds_outcf = em.sel2dcf(
-        ds, longitude=0, latitude=4, mask=mask
-    )
+    ds_outcf = em.sel2dcf(ds, longitude=0, latitude=4, mask=mask)
     assert ds_out.coords == ds_outcf.coords
 
     # if distance_name=None, no distance returned