%matplotlib inline
from dask.distributed import Client, progress, LocalCluster
from dask_kubernetes import KubeCluster
import xarray as xr
import s3fs
cluster = KubeCluster()
cluster.scale(10);
cluster
client = Client(cluster)
# 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)
# AWS s3
fs = s3fs.S3FileSystem(anon=True)
s3map = s3fs.S3Map('rsignell/nwm/test_week5c', s3=fs)
#s3map = s3fs.S3Map('rsignell/nwm/tiny3a', s3=fs)
ds = xr.open_zarr(s3map)
ds
var='T2D'
ds[var].nbytes/1.e9
uvar = ds[var].max(dim='time').persist()
progress(uvar)
isub=2
uvar[::isub,::isub].plot.imshow(figsize=(8,6));
%%time
ds1d = ds[var][:,2000,2000]
ds1d.plot()