# An Example of Using Intake-ESM#

This past week, NCAR CISL updated the Casper Node to use PBS instead of Slurm for scheduling jobs. This led a post in which an example of spinning up dask clusters on the new configuration. This was also an opportunity to dig into dask, and try applying it to a sample task, specifically looking at ecosystem variables in the CESM-LE dataset, using notebooks included in Matt Long’s krill-cesm-le repository, modified by Kristen Krumhardt.

%matplotlib inline
import os
import shutil

from glob import glob

import cftime

import numpy as np
import xarray as xr

import matplotlib.pyplot as plt
import matplotlib.colors as colors

import cartopy.crs as ccrs

import intake
import pop_tools
import esmlab
import util

import warnings
warnings.filterwarnings('ignore')


Here, we spin up our dask cluster. At first, running this notebook resulted in a killed worker error. After further expection, we noticed that additional resources would be needed to read in the notebook since the data are so large (on the order of ~1-2 TB). Increasing the individual worker to a higher amount (ex. 256 GB) solved the issue. Scale up to as many workers as you think are neccessary for the calculation (this may take some trial and error).

You should base your chunks on your analysis - so in this case, we are interested in the top 100 m in the ocean, so chunking by z_t, the vertical coordinate, will help reduce the memory.

Good general advice for working with dask is to check the dask graph (accessed via the url when you call the client). When there are “orange” regions within your task graph, this means your workers are getting close to maxing out on memory.

import dask

# Import a client

cluster = PBSCluster(
cores=2, # The number of cores you want
memory='256 GB', # Amount of memory
processes=1, # How many processes
queue='casper', # The type of queue to utilize (/glade/u/apps/dav/opt/usr/bin/execcasper)
local_directory='\$TMPDIR', # Use your local directory
resource_spec='select=1:ncpus=2:mem=256GB', # Specify resources
project='project_id', # Input your project ID here
walltime='02:00:00', # Amount of wall time
interface='ib0', # Interface to use
)
# Scale up
cluster.scale(4)

# Change your url to the dask dashboard so you can see it

client = Client(cluster)


We will use intake-esm, which is a data catalog tool. It enables querying a database for the files we want, then loading those directly as an xarray.Dataset.

First step is to set the “collection” for the CESM-LE, which depends on a json file conforming to the ESM Catalog Specification.

catalog_file = '/glade/u/home/kristenk/TTE_CESM-LE/krill-cesm-le/notebooks/data/glade-cesm1-le.json'
variables = ['TEMP', 'diatC']

experiments = ['20C', 'RCP85']
stream = 'pop.h'

col = intake.open_esm_datastore(catalog_file)


Now we search the collection for the ensemble members (unique member_id’s) that have a chlorophyll field (diatChl). This is necessary because the ocean biogeochemistry was corrupted in some members and the data were deleted.

In this cell, member_id is a list of the ensemble members we want to operate on.

col_sub = col.search(experiment=['20C'],
stream='pop.h',
variable=['diatChl'])

member_id = list(col_sub.df.member_id.unique())
print(member_id)


### Apply the search for data#

Specify a list of variables and perform a search. Under the hood, the search functionality uses pandas data frames. We can view that frame here using the .df syntax.

col_sub = col.search(
experiment=experiments,
stream=stream,
variable=variables,
member_id=member_id,
)

print(col_sub)



Use .to_dataset_dict to read in the datasets as a dictionary, specifying the chunk size. This is also a point where it is recommended you take a look at the chunksize for the dask array, ensuring that it is a reasonable number (~100-200 mb). Test it out on one of the variables (ex. TEMP).

Since we are only interested in the top 100 m, we can chunk by z_t by 10 which should reduce the memory usage.

dsets = col_sub.to_dataset_dict(cdf_kwargs={'chunks': {'time':5, 'z_t':10}, 'decode_times': False})


## Apply an operation#

Apply a calculation. In this case, we are calculating the average temperature within the top 100 m of the ocean.

def compute_TEMP_100m(ds):
"""compute top 100m mean temperature"""

ds['TEMP_100m_mean'] = ds.TEMP.isel(z_t=slice(0,10)).mean(dim='z_t')
ds.TEMP_100m_mean.attrs = ds.TEMP.attrs
ds.TEMP_100m_mean.attrs['long_name'] = 'Mean temperature over top 100m'

return ds.drop(['TEMP'])


Now apply this function to the dataset(s)

# compute top 100m temperature
dsets2 = {key: compute_TEMP_100m(ds) for key, ds in dsets2.items()}
print('computed top 100m temp')


### Combine the datasets#

Concatenate the datasets in time, i.e. 20C + RCP8.5 experiments.

ordered_dsets_keys = ['ocn,20C,pop.h', 'ocn,RCP85,pop.h']
#ordered_dsets_keys = ['ocn.20C.pop.h', 'ocn.RCP85.pop.h']
ds = xr.concat(
[dsets2[exp] for exp in ordered_dsets_keys],
dim='time',
data_vars='minimal',
)
time_encoding = dsets2[ordered_dsets_keys[0]].time.encoding


### Calculate an annual mean#

Create an annual mean for temperature, using a utility function which deals with the times correctly.

ds_ann = util.ann_mean(ds, time_bnds_varname='time_bound', time_centered=True)


## Write the dataset to disk#

Here, we are just looking to write out a few of the variables.


# Create a list of variables to write out
keep_vars = ['time_bound', 'TAREA', 'time', 'dz', 'KMT', 'member_id', 'TLAT', 'TLONG', 'TEMP_100m_mean']

# Subset for these variables
ds_out = ds_ann[keep_vars])

# Run the computations
ds_out.compute()

# Write out to scratch since this is likely a large file