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

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

# In[1]:


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


# In[2]:


# depends on the machine you are using
cluster = LocalCluster()
client = Client(cluster)


# In[3]:


client 


# Pangeo dashboard is: http://pangeo.pydata.org/user/rsignell-usgs/proxy/8787/status
# 
# So try: https://jupyter-jetstream.unidata.ucar.edu/user/rsignell/proxy/8787/status
# 
# http://pangeo.pydata.org/user/rsignell-usgs/proxy/8787/status
# 
# Hmmm, doesn't work...

# In[4]:


bucket_endpoint='https://iu.jetstream-cloud.org:8080'
fs = s3fs.S3FileSystem(anon=False, client_kwargs=dict(endpoint_url=bucket_endpoint))


# In[5]:


f_zarr = 'rsignell/nwm/test_week'


# In[6]:


d = s3fs.S3Map(f_zarr, s3=fs)


# In[7]:


ds = xr.open_zarr(d)


# In[8]:


ds


# In[9]:


ds.nbytes/1.e9


# In[10]:


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


# In[11]:


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


# In[12]:


get_ipython().run_cell_magic('time', '', "isub = 4\nmean_temp = ds['T2D'][:,::isub,::isub].mean(dim='time').compute()\n")


# In[13]:


get_ipython().run_line_magic('matplotlib', 'inline')
mean_temp.plot.imshow(figsize=(12,8))


# In[14]:


get_ipython().run_cell_magic('time', '', "t1d = ds['T2D'][:,1000,1000].persist()\nt1d.plot()\n")