#!/usr/bin/env python
# coding: utf-8

# # Read National Water Model (NWM) model data from Zarr

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')

from dask.distributed import Client, progress, LocalCluster
from dask_kubernetes import KubeCluster
import xarray as xr
import s3fs


# In[2]:


cluster = KubeCluster.from_yaml('/home/jovyan/worker-template.yaml')
cluster.scale(10);


# In[3]:


cluster


# In[4]:


client = Client(cluster)
client


# In[ ]:


# jetstream s3
# url='https://iu.jetstream-cloud.org:8080'
# fs = s3fs.S3FileSystem(client_kwargs=dict(endpoint_url=url), anon=True)
# s3map = s3fs.S3Map('rsignell/nwm/test_week', s3=fs)


# In[5]:


# AWS s3
fs = s3fs.S3FileSystem(anon=True)
s3map = s3fs.S3Map('rsignell/nwm/test_week5c', s3=fs)
#s3map = s3fs.S3Map('rsignell/nwm/tiny3a', s3=fs)


# In[6]:


ds = xr.open_zarr(s3map)


# In[7]:


ds


# In[8]:


var='T2D'
ds[var].nbytes/1.e9


# In[9]:


ds[var].mean(dim='time')


# In[10]:


mean_var = ds[var].mean(dim='time').persist()
progress(mean_var)


# In[11]:


isub=2
mean_var[::isub,::isub].plot.imshow(figsize=(8,6));


# In[12]:


get_ipython().run_cell_magic('time', '', 'ds1d = ds[var][:,2000,2000]\nds1d.plot()\n')


# In[ ]:





# In[ ]: