import xarray as xr
import xroms
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import cartopy

How to select data¶

The input/output notebook demonstrates how to load in data, but now how to select and slice it apart? Much of this is accomplished with the sel and isel methods in xarray, which are demonstrated in detail in this notebook.

Use sel to select/slice a Dataset or DataArray by dimension values; the best example of this for ROMS output is selecting certain time using a string representation of a datetime.

ds.salt.sel(ocean_time='2010-1-1 12:00')
ds.salt.sel(ocean_time=slice('2010-1-1', '2010-2-1'))

Use isel to subdivide a Dataset or DataArray by dimension indices:

ds.salt.isel(eta_rho=20, xi_rho=100)
ds.salt.isel(eta_rho=slice(20,100,10), xi_rho=slice(None,None,5))

`cf-xarray`¶

xroms includes the cf-xarray accessor, which allows you to use xarray sel and isel commands for a DataArray without needing to input the exact grid – just the axes.

With xarray alone:

ds.salt.isel(xi_rho=20, eta_rho=10, s_rho=20, ocean_time=10)

With cf-xarray accessor:

ds.salt.cf.isel(X=20, Y=10, Z=20, T=10)

and get the same thing back. Same for sel. The T, Z, Y, X names can be mixed and matched with the actual dimension names. Some of the attribute wrangling in xroms is dedicated to making sure that cf-xarray can always identify dimensions and coordinates for DataArrays.

You can always check what cf-xarray understands about a Dataset or DataArray with

ds.salt.cf.describe()

Load in data¶

More information at in input/output page

ds = xroms.datasets.fetch_ROMS_example_full_grid()
ds, xgrid = xroms.roms_dataset(ds, include_cell_volume=True)
ds.xroms.set_grid(xgrid)
ds

<xarray.Dataset> Size: 734MB
Dimensions:     (eta_rho: 191, xi_rho: 300, s_rho: 30, ocean_time: 2, s_w: 31,
                 xi_u: 299, eta_v: 190)
Coordinates: (12/21)
    lon_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    lat_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
  * s_rho       (s_rho) float64 240B -0.9833 -0.95 -0.9167 ... -0.05 -0.01667
  * s_w         (s_w) float64 248B -1.0 -0.9667 -0.9333 ... -0.03333 0.0
  * ocean_time  (ocean_time) datetime64[ns] 16B 2009-11-19T12:00:00 2009-11-1...
    lon_u       (eta_rho, xi_u) float64 457kB dask.array<chunksize=(191, 299), meta=np.ndarray>
    ...          ...
    z_w_v       (ocean_time, s_w, eta_v, xi_rho) float64 28MB dask.array<chunksize=(2, 31, 190, 300), meta=np.ndarray>
    z_w_psi     (ocean_time, s_w, eta_v, xi_u) float64 28MB dask.array<chunksize=(2, 31, 190, 299), meta=np.ndarray>
    z_rho       (ocean_time, s_rho, eta_rho, xi_rho) float64 28MB dask.array<chunksize=(2, 30, 191, 300), meta=np.ndarray>
    z_rho_u     (ocean_time, s_rho, eta_rho, xi_u) float64 27MB dask.array<chunksize=(2, 30, 191, 299), meta=np.ndarray>
    z_rho_v     (ocean_time, s_rho, eta_v, xi_rho) float64 27MB dask.array<chunksize=(2, 30, 190, 300), meta=np.ndarray>
    z_rho_psi   (ocean_time, s_rho, eta_v, xi_u) float64 27MB dask.array<chunksize=(2, 30, 190, 299), meta=np.ndarray>
Data variables: (12/45)
    angle       (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    hc          float64 8B ...
    Cs_r        (s_rho) float64 240B dask.array<chunksize=(30,), meta=np.ndarray>
    zeta        (ocean_time, eta_rho, xi_rho) float32 458kB dask.array<chunksize=(2, 191, 300), meta=np.ndarray>
    h           (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    Cs_w        (s_w) float64 248B dask.array<chunksize=(31,), meta=np.ndarray>
    ...          ...
    dV_w_u      (ocean_time, s_w, eta_rho, xi_u) float64 28MB dask.array<chunksize=(2, 31, 191, 299), meta=np.ndarray>
    dV_v        (ocean_time, s_rho, eta_v, xi_rho) float64 27MB dask.array<chunksize=(2, 30, 190, 300), meta=np.ndarray>
    dV_w_v      (ocean_time, s_w, eta_v, xi_rho) float64 28MB dask.array<chunksize=(2, 31, 190, 300), meta=np.ndarray>
    dV_psi      (ocean_time, s_rho, eta_v, xi_u) float64 27MB dask.array<chunksize=(2, 30, 190, 299), meta=np.ndarray>
    dV_w_psi    (ocean_time, s_w, eta_v, xi_u) float64 28MB dask.array<chunksize=(2, 31, 190, 299), meta=np.ndarray>
    rho0        int64 8B 1025
Attributes: (12/29)
    file:              ocean_his_0150.nc
    format:            netCDF-3 classic file
    Conventions:       CF-1.4
    type:              ROMS/TOMS history file
    title:             Texas and Louisiana Shelf case (Nesting)
    rst_file:          ocean_rst.nc
    ...                ...
    compiler_command:  /g/software/openmpi/1.4.3/intel/bin/mpif90
    compiler_flags:    -heap-arrays -fp-model precise -assume 2underscores -c...
    tiling:            016x032
    history:           ROMS/TOMS, Version 3.4, Sunday - December 4, 2011 -  7...
    ana_file:          /scratch/zhangxq/projects/txla_nesting6/Functionals/an...
    CPP_options:       TXLA, ANA_BSFLUX, ANA_BTFLUX, ASSUMED_SHAPE, BULK_FLUX...

ds.salt.cf.describe()

Coordinates:
             CF Axes: * X: ['xi_rho']
                      * Y: ['eta_rho']
                      * Z: ['s_rho']
                      * T: ['ocean_time']

      CF Coordinates:   longitude: ['lon_rho']
                        latitude: ['lat_rho']
                        vertical: ['z_rho']
                      * time: ['ocean_time']

       Cell Measures:   area, volume: n/a

      Standard Names:   depth: ['z_rho']
                        latitude: ['lat_rho']
                        longitude: ['lon_rho']
                      * ocean_s_coordinate_g1: ['s_rho']
                      * time: ['ocean_time']

              Bounds:   n/a

       Grid Mappings:   n/a

/tmp/ipykernel_3073/990108105.py:1: DeprecationWarning: 'obj.cf.describe()' will be removed in a future version. Use instead 'repr(obj.cf)' or 'obj.cf' in a Jupyter environment.
  ds.salt.cf.describe()

Select¶

Surface layer slice¶

The surface in ROMS is given by the last index in the vertical dimension. The easiest way to access this is by indexing into s_rho. While normally it is better to access coordinates through keywords to be human-readable, it’s not easy to tell what value of s_rho gives the surface. In this instance, it’s easier to just go by index.

ds.salt.isel(s_rho=-1)

ds.salt.cf.isel(Z=-1)  # with cf-xarray

You can also grab the Z level that is “nearest” to 0, the surface, which will give the same vertical level as the other options:

ds.salt.cf.sel(Z=0, method="nearest")

ds.salt.cf.sel(Z=0, method="nearest")

<xarray.DataArray 'salt' (ocean_time: 2, eta_rho: 191, xi_rho: 300)> Size: 458kB
dask.array<getitem, shape=(2, 191, 300), dtype=float32, chunksize=(2, 191, 300), chunktype=numpy.ndarray>
Coordinates:
    lon_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    lat_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    s_rho       float64 8B -0.01667
  * ocean_time  (ocean_time) datetime64[ns] 16B 2009-11-19T12:00:00 2009-11-1...
  * xi_rho      (xi_rho) int64 2kB 0 1 2 3 4 5 6 ... 293 294 295 296 297 298 299
  * eta_rho     (eta_rho) int64 2kB 0 1 2 3 4 5 6 ... 185 186 187 188 189 190
    z_rho       (ocean_time, eta_rho, xi_rho) float64 917kB dask.array<chunksize=(2, 191, 300), meta=np.ndarray>
Attributes:
    long_name:  salinity
    time:       ocean_time
    field:      salinity, scalar, series

x/y index slice¶

For a curvilinear ROMS grid, selecting by the dimensions xi_rho or eta_rho (or for whichever is the relevant grid) is not very meaningful because they are given by index. Thus the following is possible to get a slice along the index, but it cannot be used to find a slice based on the lon/lat values. For the eta and xi grids, sel is equivalent to isel.

ds.temp.sel(xi_rho=20)

ds.temp.cf.sel(X=20);  # same with cf-xarray accessor

ds.temp.cf.sel(X=20)

<xarray.DataArray 'temp' (ocean_time: 2, s_rho: 30, eta_rho: 191)> Size: 46kB
dask.array<getitem, shape=(2, 30, 191), dtype=float32, chunksize=(2, 30, 191), chunktype=numpy.ndarray>
Coordinates:
    lon_rho     (eta_rho) float64 2kB dask.array<chunksize=(191,), meta=np.ndarray>
    lat_rho     (eta_rho) float64 2kB dask.array<chunksize=(191,), meta=np.ndarray>
  * s_rho       (s_rho) float64 240B -0.9833 -0.95 -0.9167 ... -0.05 -0.01667
  * ocean_time  (ocean_time) datetime64[ns] 16B 2009-11-19T12:00:00 2009-11-1...
    xi_rho      int64 8B 20
  * eta_rho     (eta_rho) int64 2kB 0 1 2 3 4 5 6 ... 185 186 187 188 189 190
    z_rho       (ocean_time, s_rho, eta_rho) float64 92kB dask.array<chunksize=(2, 30, 191), meta=np.ndarray>
Attributes:
    long_name:  potential temperature
    units:      Celsius
    time:       ocean_time
    field:      temperature, scalar, series

Single time¶

Find the forecast model output available that is closest to now. Note that the method keyword argument is not necessary if the desired date/time is exactly a model output time. You can daisy-chain together different sel and isel calls.

date = "2009-11-19T13:00"

ds.salt.isel(s_rho=-1).sel(ocean_time=date, method='nearest')

ds.salt.cf.isel(Z=-1).cf.sel(T=date, method='nearest')  # with cf-xarray

date = "2009-11-19T13:00"
ds.salt.cf.sel(Z=0, T=date, method='nearest')

<xarray.DataArray 'salt' (eta_rho: 191, xi_rho: 300)> Size: 229kB
dask.array<getitem, shape=(191, 300), dtype=float32, chunksize=(191, 300), chunktype=numpy.ndarray>
Coordinates:
    lon_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    lat_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    s_rho       float64 8B -0.01667
    ocean_time  datetime64[ns] 8B 2009-11-19T12:00:00
  * xi_rho      (xi_rho) int64 2kB 0 1 2 3 4 5 6 ... 293 294 295 296 297 298 299
  * eta_rho     (eta_rho) int64 2kB 0 1 2 3 4 5 6 ... 185 186 187 188 189 190
    z_rho       (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
Attributes:
    long_name:  salinity
    time:       ocean_time
    field:      salinity, scalar, series

Range of time¶

time_range = slice(date, pd.Timestamp(date)+pd.Timedelta('3 hours'))

ds.salt.sel(ocean_time=time_range)

ds.salt.cf.sel(T=time_range)  # cf-xarray

time_range = slice(date, pd.Timestamp(date)+pd.Timedelta('3 hours'))
ds.salt.cf.sel(T=time_range)  # cf-xarray

<xarray.DataArray 'salt' (ocean_time: 1, s_rho: 30, eta_rho: 191, xi_rho: 300)> Size: 7MB
dask.array<getitem, shape=(1, 30, 191, 300), dtype=float32, chunksize=(1, 30, 191, 300), chunktype=numpy.ndarray>
Coordinates:
    lon_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
    lat_rho     (eta_rho, xi_rho) float64 458kB dask.array<chunksize=(191, 300), meta=np.ndarray>
  * s_rho       (s_rho) float64 240B -0.9833 -0.95 -0.9167 ... -0.05 -0.01667
  * ocean_time  (ocean_time) datetime64[ns] 8B 2009-11-19T16:00:00
  * xi_rho      (xi_rho) int64 2kB 0 1 2 3 4 5 6 ... 293 294 295 296 297 298 299
  * eta_rho     (eta_rho) int64 2kB 0 1 2 3 4 5 6 ... 185 186 187 188 189 190
    z_rho       (ocean_time, s_rho, eta_rho, xi_rho) float64 14MB dask.array<chunksize=(1, 30, 191, 300), meta=np.ndarray>
Attributes:
    long_name:  salinity
    time:       ocean_time
    field:      salinity, scalar, series

Select region¶

Select a boxed region by min/max lon and lat values.

# want model output only within the box defined by these lat/lon values
lon = np.array([-92, -91])
lat = np.array([28, 29])

# this condition defines the region of interest
box = ((lon[0] < ds["salt"].cf["longitude"]) & (ds["salt"].cf["longitude"] < lon[1])
       & (lat[0] < ds["salt"].cf["latitude"]) & (ds["salt"].cf["latitude"] < lat[1])).compute()

Plot the model output in the box at the surface

dss = ds.where(box).salt.cf.isel(Z=-1, T=0)
dss.cf.plot(x='longitude', y='latitude')

<matplotlib.collections.QuadMesh at 0x75dd53a4d4c0>

_images/31bceb7303c2904d453e243e3441b9d5cb755c35b7280ec31a655dac3ee13b4a.png

If you don’t need the rest of the model output, you can drop it by using drop=True in the where call.

dss = ds.where(box, drop=True).salt.cf.isel(Z=-1, T=0)
dss.cf.plot(x='longitude', y='latitude')

<matplotlib.collections.QuadMesh at 0x75dd5368e0c0>

_images/5e73dcff1acc7d4492053fa3a01d5d43a7269b29d01bbbd58da7bb29a5c9877f.png

Can calculate a metric within the box:

dss.mean().values

array(28.308672, dtype=float32)

Subset model output¶

Subset Dataset of model output such that subsetted domain is as if the simulation was run on that size grid. That is, the rho grid is 1 larger than the psi grid in each of xi and eta.

ds.xroms.subset(X=slice(20,40), Y=slice(50,100))  # with accessor

xroms.subset(ds, X=slice(20,40), Y=slice(50,100))

ds.xroms.subset(X=slice(20,40), Y=slice(50,100))  # with accessor

<xarray.Dataset> Size: 12MB
Dimensions:     (eta_rho: 50, xi_rho: 20, s_rho: 30, ocean_time: 2, s_w: 31,
                 xi_u: 19, eta_v: 49)
Coordinates: (12/21)
    lon_rho     (eta_rho, xi_rho) float64 8kB dask.array<chunksize=(50, 20), meta=np.ndarray>
    lat_rho     (eta_rho, xi_rho) float64 8kB dask.array<chunksize=(50, 20), meta=np.ndarray>
  * s_rho       (s_rho) float64 240B -0.9833 -0.95 -0.9167 ... -0.05 -0.01667
  * s_w         (s_w) float64 248B -1.0 -0.9667 -0.9333 ... -0.03333 0.0
  * ocean_time  (ocean_time) datetime64[ns] 16B 2009-11-19T12:00:00 2009-11-1...
    lon_u       (eta_rho, xi_u) float64 8kB dask.array<chunksize=(50, 19), meta=np.ndarray>
    ...          ...
    z_w_v       (ocean_time, s_w, eta_v, xi_rho) float64 486kB dask.array<chunksize=(2, 31, 49, 20), meta=np.ndarray>
    z_w_psi     (ocean_time, s_w, eta_v, xi_u) float64 462kB dask.array<chunksize=(2, 31, 49, 19), meta=np.ndarray>
    z_rho       (ocean_time, s_rho, eta_rho, xi_rho) float64 480kB dask.array<chunksize=(2, 30, 50, 20), meta=np.ndarray>
    z_rho_u     (ocean_time, s_rho, eta_rho, xi_u) float64 456kB dask.array<chunksize=(2, 30, 50, 19), meta=np.ndarray>
    z_rho_v     (ocean_time, s_rho, eta_v, xi_rho) float64 470kB dask.array<chunksize=(2, 30, 49, 20), meta=np.ndarray>
    z_rho_psi   (ocean_time, s_rho, eta_v, xi_u) float64 447kB dask.array<chunksize=(2, 30, 49, 19), meta=np.ndarray>
Data variables: (12/45)
    angle       (eta_rho, xi_rho) float64 8kB dask.array<chunksize=(50, 20), meta=np.ndarray>
    hc          float64 8B ...
    Cs_r        (s_rho) float64 240B dask.array<chunksize=(30,), meta=np.ndarray>
    zeta        (ocean_time, eta_rho, xi_rho) float32 8kB dask.array<chunksize=(2, 50, 20), meta=np.ndarray>
    h           (eta_rho, xi_rho) float64 8kB dask.array<chunksize=(50, 20), meta=np.ndarray>
    Cs_w        (s_w) float64 248B dask.array<chunksize=(31,), meta=np.ndarray>
    ...          ...
    dV_w_u      (ocean_time, s_w, eta_rho, xi_u) float64 471kB dask.array<chunksize=(2, 31, 50, 19), meta=np.ndarray>
    dV_v        (ocean_time, s_rho, eta_v, xi_rho) float64 470kB dask.array<chunksize=(2, 30, 49, 20), meta=np.ndarray>
    dV_w_v      (ocean_time, s_w, eta_v, xi_rho) float64 486kB dask.array<chunksize=(2, 31, 49, 20), meta=np.ndarray>
    dV_psi      (ocean_time, s_rho, eta_v, xi_u) float64 447kB dask.array<chunksize=(2, 30, 49, 19), meta=np.ndarray>
    dV_w_psi    (ocean_time, s_w, eta_v, xi_u) float64 462kB dask.array<chunksize=(2, 31, 49, 19), meta=np.ndarray>
    rho0        int64 8B 1025
Attributes: (12/29)
    file:              ocean_his_0150.nc
    format:            netCDF-3 classic file
    Conventions:       CF-1.4
    type:              ROMS/TOMS history file
    title:             Texas and Louisiana Shelf case (Nesting)
    rst_file:          ocean_rst.nc
    ...                ...
    compiler_command:  /g/software/openmpi/1.4.3/intel/bin/mpif90
    compiler_flags:    -heap-arrays -fp-model precise -assume 2underscores -c...
    tiling:            016x032
    history:           ROMS/TOMS, Version 3.4, Sunday - December 4, 2011 -  7...
    ana_file:          /scratch/zhangxq/projects/txla_nesting6/Functionals/an...
    CPP_options:       TXLA, ANA_BSFLUX, ANA_BTFLUX, ASSUMED_SHAPE, BULK_FLUX...

Find nearest in lon/lat¶

This matters for a curvilinear grid.

Can’t use sel because it will only search in one dimension for the nearest value and the dimensions are indices which are not necessarily geographic distance. Instead need to use a search for distance and use that for the where condition from the previous example. This functionality has been wrapped into xroms.sel2d (and its partner function xroms.argsel2d).

lon0, lat0 = -91, 28
saltsel = ds.salt.xroms.sel2d(lon0, lat0)

Or, if you instead want the indices of the nearest grid node returned, you can call argsel2d:

ds.salt.xroms.argsel2d(lon0, lat0)

(np.int64(16), np.int64(110))

Check this function, just to be sure:

dl = 0.05
box = (ds.lon_rho>lon0-dl) & (ds.lon_rho<lon0+dl) & (ds.lat_rho>lat0-dl) & (ds.lat_rho<lat0+dl)
dss = ds.where(box).salt.cf.isel(T=0, Z=-1)

vmin = dss.min().values
vmax = dss.max().values

dss.plot(x='lon_rho', y='lat_rho')
plt.scatter(lon0, lat0, c=saltsel.cf.isel(Z=-1, T=0), s=200, edgecolor='k', vmin=vmin, vmax=vmax)
plt.xlim(lon0-dl,lon0+dl)
plt.ylim(lat0-dl, lat0+dl)

(27.95, 28.05)

_images/f110828692e60a17bc3a6e42316160936457a7613fa4b6802069a89785796cca.png