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

# # Read National Water Model data from NetCDF on FUSE

# In[1]:


get_ipython().run_line_magic('matplotlib', 'inline')
import xarray as xr
import pandas as pd

from dask.distributed import Client, progress
from dask_kubernetes import KubeCluster


# In[2]:


root = '/s3/noaa-nwm-pds'


# In[3]:


dates = pd.date_range(start='2018-07-01T00:00', end='2018-07-08T00:00', freq='H')


# In[4]:


cluster = KubeCluster.from_yaml('/home/jovyan/custom-worker-template.yaml')


# In[5]:


cluster.scale(25);


# In[6]:


client = Client(cluster)


# In[7]:


client


# In[8]:


urls = ['{}/nwm.{}/forcing_short_range/nwm.t{}z.short_range.forcing.f001.conus.nc'.format(root,a.strftime('%Y%m%d'),a.strftime('%H')) for a in dates]


# In[9]:


# /s3/noaa-nwm-pds/nwm.20180615/forcing_short_range/nwm.t00z.short_range.forcing.f001.conus.nc


# In[10]:


get_ipython().run_cell_magic('time', '', "ds = xr.open_mfdataset(urls, concat_dim='time')\n")


# In[11]:


ds


# In[12]:


var = 'T2D'


# In[13]:


ds[var].nbytes/1.e9


# In[14]:


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


# In[ ]:


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


# In[ ]:


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


# In[ ]: