Skip to content

Commit

Permalink
new Station sampleMethod (#377)
Browse files Browse the repository at this point in the history 
* allow multiple points

* fix bug

* new samp method

* formatting

* format

* persist+compute data extraction

* unpersist

* format

* comment

* restore

* fix bug

* fix bug

* rename args to coincide with python `slice`

* reformat

* isort

* include slicing in test

* format

* format

* only data vars

* fix unref dataarray

* format

* fix format

* formatting

* set `iface=None` as default argument

* squeeze dimensions of len=0

* explicit rotated facets

* format

* format

* format

* final fix

* fix bug

* more fixes

* format

* pin down scipy

---------

Co-authored-by: Miguel Jimenez <mjimen17@jhu.edu>
  • Loading branch information
@Mikejmnez
Miguel Jimenez and Mikejmnez authored Jul 6, 2023
1 parent cea1845 commit c9ba7bc
Show file tree
Hide file tree
Showing 7 changed files with 652 additions and 69 deletions.
1 change: 1 addition & 0 deletions ci/environment.yml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ dependencies:
- xarray
- xoak
- cartopy
- scipy < 1.11
- intake-xarray
- geopy
- xesmf > 0.6.3
Expand Down
252 changes: 250 additions & 2 deletions oceanspy/llc_rearrange.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -892,8 +892,6 @@ def flip_v(_ds, co_list=metrics, dims=True, _len=3):
_ds[_varName] = -_ds[_varName]
elif _varName == "SN":
_ds[_varName] = -_ds[_varName]
# elif _varName == "CS":
# _ds[_varName] = -_ds[_varName]
return _ds


Expand Down Expand Up @@ -1232,6 +1230,256 @@ def llc_local_to_lat_lon(ds, co_list=metrics):
return _ds


def arctic_eval(_ds, _ix, _iy, _dim_name="mooring"):
"""
Evaluates a dataset
at arctic face. _ix and _iy are vectors of index points, associated with
different locations around the face=6.
"""

_ds = mates(_ds)

y = DataArray(
_np.arange(1),
dims=("y"),
attrs={"long_name": "j-index of cell center", "units": "none"},
)
x = DataArray(
_np.arange(1),
dims=("x"),
attrs={"long_name": "i-index of cell center", "units": "none"},
)
yp1 = DataArray(
_np.arange(2),
dims=("yp1"),
attrs={"long_name": "j-index of cell corner", "units": "none"},
)
xp1 = DataArray(
_np.arange(2),
dims=("xp1"),
attrs={"long_name": "i-index of cell corner", "units": "none"},
)

_XC = _ds["XC"].isel(face=6)
XR5 = _np.min(_XC.isel(Y=0, X=0).values), _np.max(_XC.isel(Y=0, X=-1).values)
XR2 = _np.min(_XC.isel(X=0, Y=-1).values), _np.max(_XC.isel(X=0, Y=0).values)
XR7 = _np.min(_XC.isel(X=-1, Y=0).values), _np.max(_XC.isel(X=-1, Y=-1).values)
XR10 = _np.min(_XC.isel(X=-1, Y=-1).values), _np.max(_XC.isel(X=0, Y=-1).values)
DS = []
co_list = [var for var in _ds.coords]

for i in range(len(_ix)):
_ix0, _iy0 = _np.array([_ix[i]]), _np.array([_iy[i]])
p = _XC.isel(X=_ix[i], Y=_iy[i]).values
if p > XR2[0] and p < XR2[-1]:
new_dim = DataArray(
_np.arange(len(_ix0)),
dims=(_dim_name),
attrs={"long_name": "index of " + _dim_name, "units": "none"},
)

# Transform indexes in DataArray
iY = DataArray(
_np.reshape(_iy0, (len(new_dim), len(y))),
coords={_dim_name: new_dim, "y": y},
dims=(_dim_name, "y"),
)
iX = DataArray(
_np.reshape(_ix0, (len(new_dim), len(x))),
coords={_dim_name: new_dim, "x": x},
dims=(_dim_name, "x"),
)

iYp1 = DataArray(
_np.stack((_iy0, _iy0 + 1)[::-1], 1),
coords={_dim_name: new_dim, "yp1": yp1},
dims=(_dim_name, "yp1"),
)
iXp1 = DataArray(
_np.stack((_ix0, _ix0 + 1), 1),
coords={_dim_name: new_dim, "xp1": xp1},
dims=(_dim_name, "xp1"),
)
args = {
"X": iX,
"Y": iY,
"Xp1": iXp1,
"Yp1": iYp1,
"face": 6,
}
rename = {"yp1": "Xp1", "xp1": "Yp1", "x": "Y", "y": "X"}
new_ds = _ds.isel(**args).drop_vars(["Xp1", "Yp1", "X", "Y"])
new_ds = new_ds.rename_dims(rename).rename_vars(rename)
new_ds = rotate_vars(new_ds)

for _varName in new_ds.variables:
if "mate" in new_ds[_varName].attrs:
_dims = new_ds[_varName].dims
if _varName not in co_list and "Xp1" in _dims:
new_ds[_varName] = -new_ds[_varName]
if _varName == "SN":
new_ds[_varName] = -new_ds[_varName]

elif p > XR5[0] and p < XR5[-1]:
new_dim = DataArray(
_np.arange(len(_ix0)),
dims=(_dim_name),
attrs={"long_name": "index of " + _dim_name, "units": "none"},
)

# Transform indexes in DataArray
iY = DataArray(
_np.reshape(_iy0, (len(new_dim), len(y))),
coords={_dim_name: new_dim, "y": y},
dims=(_dim_name, "y"),
)
iX = DataArray(
_np.reshape(_ix0, (len(new_dim), len(x))),
coords={_dim_name: new_dim, "x": x},
dims=(_dim_name, "x"),
)

iYp1 = DataArray(
_np.stack((_iy0, _iy0 + 1), 1),
coords={_dim_name: new_dim, "yp1": yp1},
dims=(_dim_name, "yp1"),
)
iXp1 = DataArray(
_np.stack((_ix0, _ix0 + 1), 1),
coords={_dim_name: new_dim, "xp1": xp1},
dims=(_dim_name, "xp1"),
)
args = {
"X": iX,
"Y": iY,
"Xp1": iXp1,
"Yp1": iYp1,
"face": 6,
}
rename = {"yp1": "Yp1", "xp1": "Xp1", "x": "X", "y": "Y"}

new_ds = _ds.isel(**args).drop_vars(["Xp1", "Yp1", "X", "Y"])
new_ds = new_ds.rename_dims(rename).rename_vars(rename)

elif p > XR7[0] or p < XR7[-1]:
new_dim = DataArray(
_np.arange(len(_ix0)),
dims=(_dim_name),
attrs={"long_name": "index of " + _dim_name, "units": "none"},
)

# Transform indexes in DataArray
iY = DataArray(
_np.reshape(_iy0, (len(new_dim), len(y))),
coords={_dim_name: new_dim, "y": y},
dims=(_dim_name, "y"),
)
iX = DataArray(
_np.reshape(_ix0, (len(new_dim), len(x))),
coords={_dim_name: new_dim, "x": x},
dims=(_dim_name, "x"),
)

iYp1 = DataArray(
_np.stack((_iy0, _iy0 + 1), 1),
coords={_dim_name: new_dim, "yp1": yp1},
dims=(_dim_name, "yp1"),
)

iXp1 = DataArray(
_np.stack((_ix0, _ix0 + 1)[::-1], 1),
coords={_dim_name: new_dim, "xp1": xp1},
dims=(_dim_name, "xp1"),
)
args = {
"X": iX,
"Y": iY,
"Xp1": iXp1,
"Yp1": iYp1,
"face": 6,
}
rename = {"yp1": "Xp1", "xp1": "Yp1", "x": "Y", "y": "X"}
new_ds = _ds.isel(**args).drop_vars(["Xp1", "Yp1", "X", "Y"])
new_ds = new_ds.rename_dims(rename).rename_vars(rename)
new_ds = rotate_vars(new_ds)

for _varName in new_ds.variables:
if "mate" in new_ds[_varName].attrs:
_dims = new_ds[_varName].dims
if _varName not in co_list and "Yp1" in _dims:
new_ds[_varName] = -new_ds[_varName]
if _varName == "CS":
new_ds[_varName] = -new_ds[_varName]

elif p > XR10[0] and p < XR10[-1]:
new_dim = DataArray(
_np.arange(len(_ix0)),
dims=(_dim_name),
attrs={"long_name": "index of " + _dim_name, "units": "none"},
)

# Transform indexes in DataArray
iY = DataArray(
_np.reshape(_iy0, (len(new_dim), len(y))),
coords={_dim_name: new_dim, "y": y},
dims=(_dim_name, "y"),
)
iX = DataArray(
_np.reshape(_ix0, (len(new_dim), len(x))),
coords={_dim_name: new_dim, "x": x},
dims=(_dim_name, "x"),
)

iYp1 = DataArray(
_np.stack((_iy0, _iy0 + 1)[::-1], 1),
coords={_dim_name: new_dim, "yp1": yp1},
dims=(_dim_name, "yp1"),
)

iXp1 = DataArray(
_np.stack((_ix0, _ix0 + 1)[::-1], 1),
coords={_dim_name: new_dim, "xp1": xp1},
dims=(_dim_name, "xp1"),
)
args = {
"X": iX,
"Y": iY,
"Xp1": iXp1,
"Yp1": iYp1,
"face": 6,
}
rename = {"yp1": "Yp1", "xp1": "Xp1", "x": "X", "y": "Y"}

new_ds = _ds.isel(**args).drop_vars(["Xp1", "Yp1", "X", "Y"])
new_ds = new_ds.rename_dims(rename).rename_vars(rename)

for _varName in new_ds.variables:
if "mate" in new_ds[_varName].attrs:
_dims = new_ds[_varName].dims
if _varName not in co_list and ("Yp1" in _dims or "Xp1" in _dims):
new_ds[_varName] = -new_ds[_varName]
if _varName == "SN":
new_ds[_varName] = -new_ds[_varName]
if _varName == "CS":
new_ds[_varName] = -new_ds[_varName]

DS.append(new_ds)
dsf = DS[0].reset_coords()
if len(DS) > 1:
for i in range(1, len(DS)):
nmds = reset_dim(DS[i], i, dim="station")
dsf = dsf.combine_first(nmds.reset_coords())
return dsf.set_coords(co_list)


def reset_dim(_ds, N, dim="mooring"):
"""resets the dimension mooring by shifting it by a value set by N"""
_ds["n" + dim] = N + _ds[dim]
_ds = _ds.swap_dims({dim: "n" + dim}).drop_vars(dim).rename({"n" + dim: dim})

return _ds


class Dims:
"""Creates a shortcut for dimension`s names associated with an arbitrary
variable."""
Expand Down
27 changes: 27 additions & 0 deletions oceanspy/plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
from .compute import _add_missing_variables
from .compute import integral as _integral
from .compute import weighted_mean as _weighted_mean
from .llc_rearrange import Dims

# Additional dependencies (private)
try:
Expand Down Expand Up @@ -651,6 +652,20 @@ def horizontal_section(
col_wrap = kwargs.pop("col_wrap", None)
subplot_kws = kwargs.pop("subplot_kws", None)
transform = kwargs.pop("transform", None)
xstep, ystep = kwargs.pop("xstep", None), kwargs.pop("ystep", None)

DIMS = [dim for dim in da.dims if dim[0] in ["X", "Y"]]
dims_var = Dims(DIMS[::-1])

if xstep is not None and ystep is not None:
xslice = slice(0, len(da[dims_var.X]), xstep)
yslice = slice(0, len(da[dims_var.Y]), ystep)
else:
xslice = slice(0, len(da[dims_var.X]))
yslice = slice(0, len(da[dims_var.Y]))

sargs = {dims_var.X: xslice, dims_var.Y: yslice}
da = da.isel(**sargs)

# Projection
if ax is None:
Expand Down Expand Up @@ -937,6 +952,18 @@ def vertical_section(
ver_name = [dim for dim in od.grid_coords["Z"] if dim in da.dims][0]
da = da.squeeze()

# slicing along section
step = kwargs.pop("step", None)
DIMS = [dim for dim in da.dims]
dims_var = Dims(DIMS[::-1])
if step is not None and dims_var.X in ["mooring", "station"]:
xslice = slice(0, len(da[dims_var.X]), step)
else:
xslice = slice(0, len(da[dims_var.X]))

sargs = {dims_var.X: xslice}
da = da.isel(**sargs)

# CONTOURNAME
if contourName is not None:
# Apply mean and sum
Expand Down
Loading

0 comments on commit c9ba7bc

Please sign in to comment.