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

# # National Water Model analysis using Zarr and Dask distributed

# In[1]:


import s3fs
import xarray as xr


# In[2]:


from dask.distributed import Client
client = Client("dask-scheduler:8786")
client


# Dashboard: https://js-157-11.jetstream-cloud.org/user/rsignell-usgs/proxy/8787/status

# In[3]:


url='https://iu.jetstream-cloud.org:8080'


# In[4]:


fs = s3fs.S3FileSystem(client_kwargs=dict(endpoint_url=url), anon=True)


# In[5]:


s3map = s3fs.S3Map('rsignell/nwm/test_week', s3=fs) 


# In[6]:


ds = xr.open_zarr(s3map)


# In[7]:


ds


# In[8]:


var='T2D'


# In[9]:


print('start:',ds[var][0].time.values)
print(' stop:',ds[var][-1].time.values)


# In[10]:


ds.nbytes/1.e9


# In[11]:


ds[var].nbytes/1.e9


# In[ ]:


get_ipython().run_cell_magic('time', '', "mean_var = ds[var][:,:,:].mean(dim='time').compute()\n")


# In[ ]:


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


# Try plotting first time step

# In[ ]:


isub=50
ds[var][0,::isub,::isub].plot.imshow(figsize=(8,6))


# In[ ]: