Analyzing a Cloud-Native IVT Zarr for Atmospheric Rivers

import os
import sys
import cftime
import dask
import folium
import hvplot.xarray  
from dask.distributed import Client
import s3fs
import numpy as np
import xarray as xr
from pathlib import Path
import cartopy.crs as ccrs
import cartopy.feature as cfeature 
import matplotlib.pyplot as plt
import requests
import importlib.util
from getpass import getpass
from hsclient import HydroShare
from hsmodels.schemas.fields import BoxCoverage, PointCoverage, PeriodCoverage
from hsmodels.schemas.fields import AwardInfo
from datetime import datetime
def is_installed(package_name):
    return importlib.util.find_spec(package_name) is not None
import ARUtils
print("h5netcdf installed:", is_installed("h5netcdf"))
print("h5py installed:", is_installed("h5py"))
# Check if the required packages are installed

# use a try accept loop so we only instantiate the client
# if it doesn't already exist.
try:
    print(client.dashboard_link)
except:    
    # The client should be customized to your workstation resources.
    # This is configured for a "Large" instance on ciroh.awi.2i2c.cloud
    # client = Client()
    client = Client(n_workers=4, threads_per_worker=2, memory_limit="3GB") 
    print(client.dashboard_link)

dask.config.set({'array.slicing.split_large_chunks': True})

# sign into HydroShare 
username = "abogan" # <- Make sure to change your HS username
password = getpass('Enter your password: ')
resource_id = "53effa6f532b43a5972daa5336932224"

# create an S3 access key and secret key for the authenticated user
response = requests.post(f"https://hydroshare.org/hsapi/user/service/accounts/s3/", auth=(username, password))
response_json = response.json()
access_key = response_json['access_key']
secret_key = response_json['secret_key']

# create a filesystem-like client for S3
s3 = s3fs.S3FileSystem(key=access_key,
                       secret=secret_key,
                       endpoint_url='https://s3.hydroshare.org')

# get the s3_path for the resource_id
response = requests.get(f"https://hydroshare.org/hsapi/resource/s3/{resource_id}/", auth=(username, password))
response_json = response.json()
s3_path = f"{response_json['bucket']}/{response_json['prefix']}"

s3.ls(s3_path)

hs = HydroShare(username, password)

user_id = 1183
hs_user = hs.user(user_id)

filename='ivt.zarr'
# create interface between remote S3 access and filesystem-style access
mapper = s3.get_mapper(f"{s3_path.rstrip('/')}/{filename}")

# read metadata, variables/dimensions/chuncks, build lazy array objects
# (`consolidated=True` reads the zarr metadata from the combined metadata file)
ds = xr.open_zarr(mapper, consolidated=True) 

# you'll see this object is lazy Dask arrays, not a fully loaded numpy.ndarray
print(type(ds.IVT.data))

list(mapper.keys())[:20]

ds

# print(ds.IVT.chunks)
print(ds.IVT.chunksizes)

print(f'Variable size: {ds.nbytes/1e9:.1f} GB')
print('Number of years in one file:', round(ds.time.shape[0]/4/365)) 

def time_order_report(ds: xr.Dataset, dim: str = "time"):
    idx = ds.indexes.get(dim)
    if idx is None:
        raise KeyError(f"No index found for '{dim}'.")
    print(
        f"{dim}: monotonic↑={idx.is_monotonic_increasing}, "
        f"monotonic↓={idx.is_monotonic_decreasing}, "
        f"unique={idx.is_unique}, "
        f"length={len(idx)}"
    )
    return idx

# Example usage
idx = time_order_report(ds)
idx

# If not strictly increasing or not unique, fix it:
ds = ds.sortby("time")

# Re-check
_ = time_order_report(ds)

print(f"Start: {ds.time.min().values}")
print(f"End:   {ds.time.max().values}")

# Convert to proleptic_gregorian and drop cftime (requires modern xarray)
ds = ds.convert_calendar("proleptic_gregorian", use_cftime=False)

# If there are duplicate timestamps, drop them first (optional but helpful)
if ds.get_index("time").has_duplicates:
    ds = ds.sel(time=~ds.get_index("time").duplicated())

# isolate and load in only one timestep of the Zarr store
ivt0 = ds.IVT.isel(time=0).compute()

# use xarray to inspect metadata attributes of the loaded data
print(ivt0.name)
print(ivt0.attrs)

# compute the bounds of the full model domain (i.e., the Zarr store dataset)
bounds = [
    [float(ds.lat.min()), float(ds.lon.min())],
    [float(ds.lat.max()), float(ds.lon.max())]
]

# plot the full extent of the IVT data for one timestep
fig = plt.figure(figsize=(12,8))
ax = plt.axes(projection=ccrs.PlateCarree())

ivt0.plot(
    ax=ax,
    transform=ccrs.PlateCarree(),
    cmap='viridis',
    cbar_kwargs={"shrink": 0.6}
)

ax.coastlines()
ax.set_title(f"IVT at timestep {ds.time.isel(time=0).values}")

point_name = "Bodega Bay"
start_date="1970-01-01"
end_date="2009-12-31"
lat = 38.
lon = -123.125 

m = folium.Map(location=[lat, lon], zoom_start=10)
folium.CircleMarker([lat, lon], radius=6, color="red", fill=True, fill_opacity=0.8,
                    popup=f"Lat {lat:.4f}, Lon {lon:.4f}").add_to(m)
m

ds_point = ds.sel(time=slice(start_date, end_date)).sel(lat=lat, lon=lon+360., method='nearest').IVT
ds_point

%%time
ds_point = ds_point.load() 
ds_point.shape

# save the dataframe to CSV 
ds_point.to_dataframe().to_csv("ivt_timeseries.csv")

p = (ds_point.hvplot(label='IVT', color='black', line_dash='dashed', line_width=1)
).opts(
    width=900, height=400, show_grid=True, tools=['hover']
)

p  

# the nearest index for the lat/lon provided above
ilon = list(ds.lon.values).index(ds.sel(lon=lon+360., method='nearest').lon)
ilat = list(ds.lat.values).index(ds.sel(lat=lat, method='nearest').lat)

cell_buffer = 10
# subset these data to include cells above and below this lat/lon
ds_box = ds.sel(time=slice(start_date, end_date)).isel(lat=range(ilat-cell_buffer, ilat+cell_buffer),
                  lon=range(ilon-cell_buffer, ilon+cell_buffer)).IVT

ds_box = ds_box.to_dataset()

%%time
ds_box = ds_box.load()

fig, ax = plt.subplots()
ax.set_xticks([])
ax.set_yticks([])
ax = plt.axes(projection=ccrs.PlateCarree())

#add coastlines
ax.coastlines()

#add lat lon grids
gl = ax.gridlines(draw_labels=True, color='lightgrey', alpha=0.5, linestyle='--')
gl.xlabels_bottom = False
gl.ylabels_right = False

plot = ds_box.isel(time=0).IVT.plot(label='',
                         add_labels=False,
                         add_colorbar=False,    
                         xticks=None,
                         yticks=None)

# add custom colorbar inside the figure
cbar_ax = fig.add_axes([0.99, 0.25, 0.02, 0.5])  # [left, bottom, width, height] in figure coords
cbar = plt.colorbar(plot, cax=cbar_ax)
cbar.set_label('IVT (kg m⁻¹ s⁻¹)')

ax.plot(lon+360., lat,
        marker='o', color='red', markersize=8,
        transform=ccrs.PlateCarree())

plt.show()

ds_box

ds_box.chunks

# # Ensure time is a single chunk (so events aren’t broken across chunks)
# ds_box = ds_box.chunk({"time": -1})
# ds_box.chunks

ds_box['AR_INDEX'] = xr.apply_ufunc(
    ARUtils.identify_ar_events,  
    ds_box.IVT,  
    input_core_dims=[["time"]],
    output_core_dims=[["time"]],
    vectorize=True,
    dask='parallelized',
    dask_gufunc_kwargs={'allow_rechunk': True},
)

ds_box['AR_DURATION'] = xr.apply_ufunc(
    ARUtils.compute_ar_durations,  
    ds_box.AR_INDEX,  
    input_core_dims=[["time"]],
    output_core_dims=[["time"]],
    vectorize=True,
    dask='parallelized',
    dask_gufunc_kwargs={'allow_rechunk': True},
)

ds_box['AR_CATEGORY'] = xr.apply_ufunc(
    ARUtils.compute_ar_category,  
    ds_box.AR_INDEX,
    ds_box.AR_DURATION,
    ds_box.IVT,
    input_core_dims=[["time"],["time"],["time"]],
    output_core_dims=[["time"]],
    vectorize=True,
    dask='parallelized',
    dask_gufunc_kwargs={'allow_rechunk': True},
)

res = xr.apply_ufunc(
    ARUtils.compute_ar_probability,
    ds_box.AR_INDEX,
    ds_box.AR_CATEGORY,
    input_core_dims=[["time"],["time"]],
    output_core_dims=[['category']],
    # output_sizes={'category': 5},
    output_dtypes=[float],
    vectorize=True,
    dask='parallelized',
    dask_gufunc_kwargs={'allow_rechunk': True,
                        'output_sizes': {'category': 5}},
)

# Add the category dimension (from 1 to 5)
res = res.assign_coords(category=np.arange(1, 6))

# add the newly computed probabilty back to the original dataset
ds = ds.assign(AR_PROBABILITY=res)

ds_box['AR_DEGREE'] = xr.apply_ufunc(
    ARUtils.compute_ar_degree,
    ds_box.IVT,
    ds_box.AR_DURATION,
    ds_box.AR_INDEX,
    input_core_dims=[["time"],["time"],["time"]],
    output_core_dims=[['time']],
    output_dtypes=[float],
    vectorize=True,
    dask='parallelized',
    dask_gufunc_kwargs={'allow_rechunk': True},
)

year = ds_box["time"].dt.year
decade = np.floor(year / 10) * 10
ds_box = ds_box.assign_coords({'decade':decade})

ds_decade = ds_box.where(ds_box.AR_CATEGORY > 0) \
            .AR_CATEGORY \
            .groupby('decade') \
            .mean(dim='time').compute()

fig, ax = plt.subplots(nrows=2, ncols=2,
                       subplot_kw={'projection': ccrs.PlateCarree()},
                       figsize=(10,8),
                       layout='constrained')

axes = ax.flatten()
for i in range(0, len(ds_decade.decade)):
    axes[i].coastlines();
    axes[i].gridlines(draw_labels=True, color='lightgrey', alpha=0.5, linestyle='--');
    
    # created grid and contour plots.
    d_plot = ds_decade.isel(decade=i).plot(label='',
                   add_labels=False,
                   add_colorbar=False,                   
                   ax=axes[i])
    
    c_plot = ds_decade.isel(decade=i).plot.contour(add_labels=True,
                                                   add_colorbar=False,
                                                   colors='white',
                                                   ax=axes[i])
    axes[i].clabel(c_plot)


    st = int(ds_decade.decade[i].item())
    et = st + 9
    axes[i].set_title(f'{st} - {et}')

fig.colorbar(d_plot, ax=ax[:, 1], shrink=0.8);
fig.suptitle('Mean Decadal AR Category');

ds_decade = ds_box.where(ds_box.AR_CATEGORY > 0) \
            .AR_DEGREE \
            .groupby('decade') \
            .mean(dim='time').compute()

import matplotlib.ticker as ticker
fig, ax = plt.subplots(nrows=2, ncols=2,
                       subplot_kw={'projection': ccrs.PlateCarree()},
                       figsize=(10,8),
                       layout='constrained')

axes = ax.flatten()
for i in range(0, len(ds_decade.decade)):
    axes[i].coastlines();
    axes[i].gridlines(draw_labels=True, color='lightgrey', alpha=0.5, linestyle='--');
    
    # created grid and contour plots.
    d_plot = ds_decade.isel(decade=i).plot(label='',
                   add_labels=False,
                   add_colorbar=False,                   
                   ax=axes[i])
    
    c_plot = ds_decade.isel(decade=i).plot.contour(add_labels=True,
                                                   add_colorbar=False,
                                                   colors='white',
                                                   ax=axes[i])

    # custom formatting function for contour labels on the maps
    def fmt_func(value):
        return f'{value:.0e}'
        
        
    axes[i].clabel(c_plot, inline=True, fontsize=10, fmt=fmt_func)


    st = int(ds_decade.decade[i].item())
    et = st + 9
    axes[i].set_title(f'{st} - {et}')

fig.colorbar(d_plot, ax=ax[:, 1], shrink=0.8);
fig.suptitle('Mean Decadal AR Degree');

fig.savefig(
    "mean_decadal_ar_degree.png",
    dpi=150,
    bbox_inches="tight"
)

# Create the new, empty resource
new_resource = hs.create()

# Get the HydroShare identifier for the new resource
res_identifier = new_resource.resource_id
print(f'The HydroShare Identifier for your new resources is: {res_identifier}')

# Construct a hyperlink to access the HydroShare landing page for the new resource
print(f'Your new resource is available at: {new_resource.metadata.url}')

# Print the Creator names
print('The list of Authors includes: ')
for creator in new_resource.metadata.creators:
    if creator.name is None:
        print(creator.organization)
    else:
        print(creator.name)

# Set the Title for the resource
new_resource.metadata.title = f'Integrated Vapor Transport - {point_name}'

# Set the Abstract text for the resource
new_resource.metadata.abstract = (
    'This resource contains integrated vapor transport (IVT) data'
    'for a selected location and time period. The outputs were generated from a '
    'notebook that subsets cloud-hosted IVT data and summarizes changes in atmospheric '
    'moisture transport at the specified point.'
)

# Call the save function to save the metadata edits to HydroShare
new_resource.save()

# Print the title just added to the resource
print(f'Title: {new_resource.metadata.title}')
print(f'Abstract: {new_resource.metadata.abstract}')

# Create subject keywords for the resource using a list of strings
new_resource.metadata.subjects = ['Atmospheric Rivers', 'Integrated vapor transport', 'Hydroinformatics']

# Save the changes to the resource in HydroShare
new_resource.save()

# Print the keywords for the resource
print('The list of keywords for the resource includes:')
for keyword in new_resource.metadata.subjects:
    print(keyword)

# Import the required metadata classes for coverage objects
from hsmodels.schemas.fields import BoxCoverage, PointCoverage, PeriodCoverage
from datetime import datetime

# If you want to set the spatial coverage to a PointCoverage instead
new_resource.metadata.spatial_coverage = PointCoverage(name=point_name,
                                                       north=lat,
                                                       east=lon,
                                                       projection='WGS 84 EPSG:4326',
                                                       type='point',
                                                       units='Decimal degrees')

# Create a beginning and ending date for a time period
beginDate = datetime.strptime(start_date,'%Y-%m-%d')
endDate = datetime.strptime(end_date,'%Y-%m-%d')

# Set the temporal coverage of the resource to a PeriodCoverage object
new_resource.metadata.period_coverage = PeriodCoverage(start=beginDate,
                                                       end=endDate)

# Save the changes to the resource in HydroShare
new_resource.save()

# Print the temporal coverage information
print('Temporal Coverage:')
print(new_resource.metadata.period_coverage)

# Print the spatial coverage information
print('\nSpatial Coverage:')
print(new_resource.metadata.spatial_coverage)

# Create a new AwardInfo object
newAwardInfo = AwardInfo(funding_agency_name='National Oceanic and Atmospheric Administration (NOAA)',
                         title=('Cooperative Institute for Research to Operations in HydrologyCooperative Institute for Research to Operations in Hydrology'),
                         number='NA22NWS4320003')

# Append the new AwardInfo object to the list of funding agencies
new_resource.metadata.awards.append(newAwardInfo)

# Save the changes to the resource in HydroShare
new_resource.save()

# Print the AwardInfo
print('Funding sources added: ')
for award in new_resource.metadata.awards:
    print(f'Award Title: {award.title}')

# We can add both mean decadal AR figure and the point-based IVT 
# timeseries data to the HydroShare resource
new_resource.file_upload('mean_decadal_ar_degree.png')
new_resource.file_upload('ivt_timeseries.csv')

# Print the names of the files in the resource
print('Updated file list after adding a file: ')
for file in new_resource.files(search_aggregations=True):
    print(file.path)