Commit 761a070b authored by Ezra Eisbrenner's avatar Ezra Eisbrenner

Fix and extend wrapper

parent 17358d31
Pipeline #12172 passed with stage
in 24 seconds
import windeval
from windeval import wrapper
def test_ekman(X):
windeval.ekman(X, drag_coefficient="large_and_pond_1981", extend_ranges=True)
wrapper.ekman(X, drag_coefficient="large_and_pond_1981", extend_ranges=True)
def test_sverdrup(X):
windeval.sverdrup(X, drag_coefficient="large_and_yeager_2004")
wrapper.sverdrup(X, drag_coefficient="large_and_yeager_2004")
"""Wrapper."""
import xarray as xr
from typing import Optional
from . import processing as pr
from functools import singledispatch
from typing import Optional, Any, Dict, List
from . import processing
@singledispatch
def calculate(ds, *args, **kwargs):
raise NotImplementedError("Data type not supported.")
@calculate.register
def _(
ds: xr.Dataset, var: str, diag: str, *args: Any, **kwargs: Dict[str, Any]
) -> xr.DataArray:
return calculate(ds[var], diag, *args, **kwargs)
@calculate.register # type: ignore
def _(
da: xr.DataArray, diag: str, *args: Any, **kwargs: Dict[str, Any]
) -> xr.DataArray:
f = getattr(processing, diag, None)
if f is not None:
y = f(da, *args, **kwargs)
else:
y = getattr(da, diag)(*args, **kwargs)
return y
@calculate.register # type: ignore
def _(
d: dict, keys: List[str], diag: str, *args: Any, **kwargs: Dict[str, Any]
) -> xr.DataArray:
f = getattr(processing, diag)
return f(*[d[k] for k in keys], *args, **kwargs)
def ekman(
......@@ -40,8 +76,8 @@ def ekman(
continue
d[x] = locals()[x]
c = "northward" if component == "eastward" else "eastward"
getattr(pr, "surface_downward_" + c + "_stress")(X, **d)
getattr(pr, component + "_ekman_transport")(X)
getattr(processing, "surface_downward_" + c + "_stress")(X, **d)
getattr(processing, component + "_ekman_transport")(X)
return X
......@@ -77,8 +113,8 @@ def sverdrup(
if locals()[x] is None:
continue
d[x] = locals()[x]
pr.surface_downward_eastward_stress(X, **d)
pr.surface_downward_northward_stress(X, **d)
pr.sverdrup_transport(X)
processing.surface_downward_eastward_stress(X, **d)
processing.surface_downward_northward_stress(X, **d)
processing.sverdrup_transport(X)
return X
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment