Notebook

Explore the COAWST US East Coast and Gulf of Mexico Forecast Archive Dataset¶

This is a cloud-optimized version of the NetCDF files accessed from the USGS ScienceBase item Collection of COAWST model forecast for the US East Coast and Gulf of Mexico. The original daily forecast files were converted into weekly NetCDF files with 168 points in the time dimension to facilitate time series access.

In [1]:

import os
import fsspec
import xarray as xr
import hvplot.xarray
import intake
import cf_xarray
import numpy as np
import panel as pn
from matplotlib import path
import xoak

Open Dataset¶

The details of data loading are stored in an intake catalog, which simplifies use. Metadata and coordinate data are loaded, but not the actual data variables, which are loaded only as needed by subsequent analysis and visualization.

In [2]:

intake_catalog_url = 's3://usgs-coawst/coawst_us_east_and_gulf.yml'
cat = intake.open_catalog(intake_catalog_url, storage_options={"anon": True})
list(cat)

Out[2]:

['COAWST_USEAST_Archive']

In [3]:

dataset = 'COAWST_USEAST_Archive' 

This is a big dataset, so it takes up to 30s to open the dataset (which involves reading all the metadata and index coordinate variable data). Here we load the data into xarray using .to_dask() so that if we have a Dask cluster, we can speed up data processing by loading and processing chunks of data in parallel.

In [4]:

%%time
ds = cat[dataset].to_dask()

CPU times: user 12.1 s, sys: 576 ms, total: 12.7 s
Wall time: 15 s

Let's look at that metadata. We can explore the different attributes and variables by clicking on the variables and icons below.

In [5]:

ds

Out[5]:

<xarray.Dataset>
Dimensions:                 (ocean_time: 112392, s_w: 17, eta_rho: 336,
                             xi_rho: 896, tracer: 8, s_rho: 16, NST: 6,
                             boundary: 4, Nbed: 1, eta_u: 336, xi_u: 895,
                             eta_v: 335, xi_v: 896, eta_psi: 335, xi_psi: 895)
Coordinates:
    lat_psi                 (eta_psi, xi_psi) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lat_rho                 (eta_rho, xi_rho) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lat_u                   (eta_u, xi_u) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lat_v                   (eta_v, xi_v) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lon_psi                 (eta_psi, xi_psi) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lon_rho                 (eta_rho, xi_rho) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lon_u                   (eta_u, xi_u) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lon_v                   (eta_v, xi_v) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
  * ocean_time              (ocean_time) datetime64[ns] 2009-08-21T01:00:00 ....
  * s_rho                   (s_rho) float64 -0.9688 -0.9062 ... -0.03125
  * s_w                     (s_w) float64 -1.0 -0.9375 -0.875 ... -0.0625 0.0
Dimensions without coordinates: eta_rho, xi_rho, tracer, NST, boundary, Nbed,
                                eta_u, xi_u, eta_v, xi_v, eta_psi, xi_psi
Data variables: (12/188)
    AKs                     (ocean_time, s_w, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 1, 168, 224), meta=np.ndarray>
    AKt                     (ocean_time, s_w, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 1, 168, 224), meta=np.ndarray>
    AKv                     (ocean_time, s_w, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 1, 168, 224), meta=np.ndarray>
    Akk_bak                 (ocean_time) float64 dask.array<chunksize=(1,), meta=np.ndarray>
    Akp_bak                 (ocean_time) float64 dask.array<chunksize=(1,), meta=np.ndarray>
    Akt_bak                 (ocean_time, tracer) float64 dask.array<chunksize=(1, 8), meta=np.ndarray>
    ...                      ...
    wetdry_mask_psi         (ocean_time, eta_psi, xi_psi) float32 dask.array<chunksize=(168, 168, 224), meta=np.ndarray>
    wetdry_mask_rho         (ocean_time, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 168, 224), meta=np.ndarray>
    wetdry_mask_u           (ocean_time, eta_u, xi_u) float32 dask.array<chunksize=(168, 168, 224), meta=np.ndarray>
    wetdry_mask_v           (ocean_time, eta_v, xi_v) float32 dask.array<chunksize=(168, 168, 224), meta=np.ndarray>
    xl                      (ocean_time) float64 dask.array<chunksize=(1,), meta=np.ndarray>
    zeta                    (ocean_time, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 168, 224), meta=np.ndarray>
Attributes: (12/37)
    CPP_options:       COAWST, ANA_BPFLUX, ANA_BSFLUX, ANA_BTFLUX, ANA_FSOBC,...
    Conventions:       CF-1.4, SGRID-0.3
    NLM_LBC:           \nEDGE:         WEST   SOUTH  EAST   NORTH  \nzeta:   ...
    ana_file:          ROMS/Functionals/ana_btflux.h, ROMS/Functionals/ana_fs...
    bry_file_01:       ./forcings3/USE_coawst_bdy.nc
    clm_file_01:       ./forcings3/USE_coawst_clm.nc
    ...                ...
    svn_url:            
    tide_file:         ../forcings/tide_forc_USeast_grd16_osu_rev2.nc
    tiling:            006x004
    title:             COAWST ROMS SWAN
    type:              ROMS/TOMS history file
    var_info:          ROMS/External/varinfo.dat

We can also explore a specific variable of interest:

In [6]:

var = 'Hwave'
da = ds[var]
da

Out[6]:

<xarray.DataArray 'Hwave' (ocean_time: 112392, eta_rho: 336, xi_rho: 896)>
dask.array<open_dataset-cde4b49188bf20ba82fd25d09eac1357Hwave, shape=(112392, 336, 896), dtype=float32, chunksize=(168, 168, 224), chunktype=numpy.ndarray>
Coordinates:
    lat_rho     (eta_rho, xi_rho) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
    lon_rho     (eta_rho, xi_rho) float64 dask.array<chunksize=(168, 224), meta=np.ndarray>
  * ocean_time  (ocean_time) datetime64[ns] 2009-08-21T01:00:00 ... 2022-06-17
Dimensions without coordinates: eta_rho, xi_rho
Attributes:
    field:      Hwave, scalar, series
    grid:       grid
    location:   face
    long_name:  wind-induced significant wave height
    time:       ocean_time
    units:      meter

Use the CF conventions to identify the coordinate variables for longitude, latitude and time

In [7]:

x = da.cf['longitude']
y = da.cf['latitude']
t = da.cf['time']
print(x.name, y.name, t.name)

lon_rho lat_rho ocean_time

Example: Load the entire spatial domain for a variable at a specific time step¶

Loading the entire spatial domain at a time step only requires reading 8 chunks of data, so it loads in a few seconds. A dask cluster doesn't help much in this case as it's already fast.

In [8]:

%%time
da2d = da.cf.sel(T='2012-10-29 12:00', method='nearest').load()

CPU times: user 1.53 s, sys: 360 ms, total: 1.89 s
Wall time: 4.3 s

In [9]:

da2d.hvplot.quadmesh(x=x.name, y=y.name, rasterize=True, geo=True, tiles='OSM', cmap='viridis')

Out[9]:

Example: Load a time series for a variable at a specific lon,lat location for a specified time range.¶

To identify a point, we will start with its lat/lon coordinates. If lon and lat were 1D coordinates, we could use lon,lat values to select using xarray, but instead we need to extract using indices, which we need to find. For this we use the xoak package:

In [10]:

lat,lon = 42.0, -69.0

In [11]:

da.xoak.set_index([y.name, x.name], 'sklearn_geo_balltree')

In [12]:

ds_point = xr.Dataset({"lon": ("point", [lon]), "lat": ("point", [lat])})

Before we read the data, let's see how many chunks we will be reading:

In [13]:

da.xoak.sel(lat_rho=ds_point.lat, lon_rho=ds_point.lon).cf.sel(T='2012-10')

Out[13]:

<xarray.DataArray 'Hwave' (ocean_time: 744, point: 1)>
dask.array<getitem, shape=(744, 1), dtype=float32, chunksize=(168, 1), chunktype=numpy.ndarray>
Coordinates:
    lat_rho     (point) float64 41.98
    lon_rho     (point) float64 -68.98
  * ocean_time  (ocean_time) datetime64[ns] 2012-10-01 ... 2012-10-31T23:00:00
Dimensions without coordinates: point
Attributes:
    field:      Hwave, scalar, series
    grid:       grid
    location:   face
    long_name:  wind-induced significant wave height
    time:       ocean_time
    units:      meter

To load one month means reading 6 chunks of data, so still don't need a cluster:

In [14]:

%%time
da1d = da.xoak.sel(lat_rho=ds_point.lat, lon_rho=ds_point.lon).cf.sel(T='2012-10').load()

CPU times: user 805 ms, sys: 183 ms, total: 988 ms
Wall time: 2.11 s

In [15]:

da1d.hvplot(x=t.name, grid=True)

Out[15]:

How many chunks of data will we read to load the entire time series of record at a point?

In [16]:

da.xoak.sel(lat_rho=ds_point.lat, lon_rho=ds_point.lon)

Out[16]:

<xarray.DataArray 'Hwave' (ocean_time: 112392, point: 1)>
dask.array<transpose, shape=(112392, 1), dtype=float32, chunksize=(168, 1), chunktype=numpy.ndarray>
Coordinates:
    lat_rho     (point) float64 41.98
    lon_rho     (point) float64 -68.98
  * ocean_time  (ocean_time) datetime64[ns] 2009-08-21T01:00:00 ... 2022-06-17
Dimensions without coordinates: point
Attributes:
    field:      Hwave, scalar, series
    grid:       grid
    location:   face
    long_name:  wind-induced significant wave height
    time:       ocean_time
    units:      meter

Since we now need to read 669 chunks of data, we should use a Dask cluster if we have access to one

Parallelize with Dask¶

We opened the dataset so that we can take advantage of parallel compute environments using dask. We're going to start a cluster now so that future steps can take advantage of this ability.

This is an optional step, but speeds up data loading and processing significantly, especially when accessing data from the cloud.

There are many options to set up Dask clusters.

Use LocalCluster¶

LocalCluster is available in any computing environment. It uses the number of CPUs of the computer running the notebook to create a cluster.

from dask.distributed import LocalCluster, Client cluster = LocalCluster() client = Client(cluster)

Use Coiled¶

Coiled provides access to remote clusters that can be used from anywhere. It requires a Coiled account.

import coiled cluster = coiled.Cluster( region="us-west-2", compute_purchase_option="spot_with_fallback", arm=True, scheduler_port=443, n_workers=10, account='dask' ) client = cluster.get_client() cluster.scale(20) client

Use a Dask Gateway Cluster¶

Dask Gateway is available on Nebari and Daskhub,

In [17]:

from dask_gateway import Gateway

gateway = Gateway()  # instantiate Dask gateway 

# Cluster options on Nebari
options = gateway.cluster_options()
options.conda_environment='global/global-pangeo'
options.profile = 'Small Worker'   # single CPU worker

# pass environment vars to workers. This could include AWS environment vars 
# needed to access requester-pays and private buckets
#os.environ['LD_PRELOAD'] = ''
#options.environment_vars = dict(os.environ)

# Create a Dask Gateway cluster
cluster = gateway.new_cluster(options)

# Get the Dask client for the Dask Gateway cluster
client = cluster.get_client()

# Scale the cluster
#cluster.adapt(minimum=4, maximum=30)
cluster.scale(30)

In [18]:

client

Out[18]:

Client

Client-c3060cba-8768-11ee-8e0e-820f45ca7c85

Connection method: Cluster object	Cluster type: dask_gateway.GatewayCluster
Dashboard: https://nebari.esipfed.org/gateway/clusters/dev.d679e15da53b467c8c0e7a0b4adb7d6d/status

Cluster Info

GatewayCluster

Name: dev.d679e15da53b467c8c0e7a0b4adb7d6d
Dashboard: https://nebari.esipfed.org/gateway/clusters/dev.d679e15da53b467c8c0e7a0b4adb7d6d/status

Load the entire time series:

In [19]:

%%time
ds_selection = da.xoak.sel(lat_rho=ds_point.lat, lon_rho=ds_point.lon).load()       

CPU times: user 345 ms, sys: 23.4 ms, total: 368 ms
Wall time: 1min 18s

In [20]:

ds_selection.hvplot(x=t.name, grid=True) 

Out[20]:

Example: Compute the time mean for a variable over the entire domain for a specific time period¶

In [21]:

%%time
da_mean = da.cf.sel(T=slice('2016-01-01 00:00','2017-01-01 00:00')).mean(dim=t.name).compute()

CPU times: user 94.3 ms, sys: 612 µs, total: 94.9 ms
Wall time: 10.8 s

In [22]:

da_mean.hvplot.quadmesh(x=x.name, y=y.name, rasterize=True, geo=True,  tiles='OSM', cmap='viridis')

Out[22]:

Example: Subset a time and space region and export to NetCDF¶

In [23]:

def bbox2ij(lon,lat,bbox=[-160., -155., 18., 23.]):
    """Return indices for i,j that will completely cover the specified bounding box.     
    i0,i1,j0,j1 = bbox2ij(lon,lat,bbox)
    lon,lat = 2D arrays that are the target of the subset
    bbox = list containing the bounding box: [lon_min, lon_max, lat_min, lat_max]

    Example
    -------  
    >>> i0,i1,j0,j1 = bbox2ij(lon_rho,lat_rho,[-71, -63., 39., 46])
    >>> h_subset = nc.variables['h'][j0:j1,i0:i1]       
    """
    bbox=np.array(bbox)
    mypath=np.array([bbox[[0,1,1,0]],bbox[[2,2,3,3]]]).T
    p = path.Path(mypath)
    points = np.vstack((lon.ravel(),lat.ravel())).T   
    n,m = np.shape(lon)
    inside = p.contains_points(points).reshape((n,m))
    ii,jj = np.meshgrid(range(m),range(n))
    return min(ii[inside]),max(ii[inside]),min(jj[inside]),max(jj[inside])

In [24]:

bbox = [-76.63290610753754, -73.55671530588432, 37.57888442021855, 41.225532965406224]   # DRB

In [25]:

i0,i1,j0,j1 = bbox2ij(x.values, y.values, bbox=bbox)
print(i0,i1,j0,j1)

518 603 253 328

In [26]:

ds_drb = ds[['temp', 'salt', 'Hwave']].isel(eta_rho=slice(j0,j1), xi_rho=slice(i0,i1))

In [27]:

ds_drb

Out[27]:

<xarray.Dataset>
Dimensions:     (ocean_time: 112392, s_rho: 16, eta_rho: 75, xi_rho: 85)
Coordinates:
    lat_rho     (eta_rho, xi_rho) float64 36.5 36.52 36.54 ... 42.18 42.2 42.22
    lon_rho     (eta_rho, xi_rho) float64 -76.16 -76.11 -76.07 ... -74.11 -74.06
  * ocean_time  (ocean_time) datetime64[ns] 2009-08-21T01:00:00 ... 2022-06-17
  * s_rho       (s_rho) float64 -0.9688 -0.9062 -0.8438 ... -0.09375 -0.03125
Dimensions without coordinates: eta_rho, xi_rho
Data variables:
    temp        (ocean_time, s_rho, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 1, 75, 85), meta=np.ndarray>
    salt        (ocean_time, s_rho, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 1, 75, 85), meta=np.ndarray>
    Hwave       (ocean_time, eta_rho, xi_rho) float32 dask.array<chunksize=(168, 75, 85), meta=np.ndarray>
Attributes: (12/37)
    CPP_options:       COAWST, ANA_BPFLUX, ANA_BSFLUX, ANA_BTFLUX, ANA_FSOBC,...
    Conventions:       CF-1.4, SGRID-0.3
    NLM_LBC:           \nEDGE:         WEST   SOUTH  EAST   NORTH  \nzeta:   ...
    ana_file:          ROMS/Functionals/ana_btflux.h, ROMS/Functionals/ana_fs...
    bry_file_01:       ./forcings3/USE_coawst_bdy.nc
    clm_file_01:       ./forcings3/USE_coawst_clm.nc
    ...                ...
    svn_url:            
    tide_file:         ../forcings/tide_forc_USeast_grd16_osu_rev2.nc
    tiling:            006x004
    title:             COAWST ROMS SWAN
    type:              ROMS/TOMS history file
    var_info:          ROMS/External/varinfo.dat

In [28]:

ds_drb_timeslice = ds_drb.cf.sel(T=slice('2022-04-01 00:00','2022-04-08 00:00'))

In [29]:

ds_drb_timeslice = ds_drb_timeslice.chunk({'eta_rho':-1, 'xi_rho':-1})  # chunk to full spatial subset domain
print(f'Uncompressed dataset size: {ds_drb_timeslice.nbytes/1e6} MB')

Uncompressed dataset size: 142.31698 MB

In [30]:

var = 'salt'
da_drb = ds_drb_timeslice[var].load()

In [31]:

viz = da_drb.hvplot.quadmesh(x=x.name, y=y.name, geo=True,
                    cmap='turbo', rasterize=True, tiles='OSM', title=var)
viz = pn.panel(viz, widgets={'ocean_time': pn.widgets.Select} )
pn.Column(viz).servable('DRB Explorer')

Out[31]:

Close the Dask client since we can't write NetCDF in parallel

In [32]:

client.close()

Specify the encoding to enable compression in the NetCDF file

In [33]:

%%time
encoding={}
for var in ds_drb_timeslice.variables:
    encoding[var] = dict(zlib=True, complevel=4, 
                         fletcher32=False, shuffle=True,
                         _FillValue=None)

ds_drb_timeslice.to_netcdf('drb.nc', encoding=encoding, mode='w')

CPU times: user 16.4 s, sys: 1.02 s, total: 17.4 s
Wall time: 21.8 s

Stop cluster¶

In [34]:

cluster.shutdown()