<jupyter_start><jupyter_text>Sample DiagnosticExample of a diagnostic showing how the COSIMA cookbook works. An objective of the COSIMA cookbook is to catalogue useful diagnostics for ocean and ice models.  Certaintools and patterns are used extensively in these examples.  While the diagnostic itself should be portableto another framework, there are some conventions used throughout that require explanation.  Background of a diagnostic notebook Each diagnostic is written up as a Jupyter notebook with the extension `.ipynb`.  The first cell in the notebookmust be a Markdown cell with a header and a one-line description.  This cell is used by sphinx-nbgallery to collectall of the diagnostics together the http://cosima-cookbook.readthedocs.io site. In this first section, a brief background on the theory of the diagnostic is presented. For this worked example, we will be calculating the eddy kinetic energy (see the Kinetic Energy notebook).  While in real diagnostic notebook, commentary about the diagnostic is presented, here we provide commentary about the technical aspects of how these diagnostics have been implemented. A note on names- The project's long name is "COSIMA Cookbook".- The GitHub project name is "cosima-cookbook".- The Python package is called "cosima_cookbook". This conventions appears to be consistent with other Python based projects.  Python import statements Early in the notebook, there is a code cell that imports all of the needed Python packages.  Internal to cosima_cookbook, other packages may also be important. But, if they are needed in this notebook, they must beimported explicitly into the namespace.  Example of a import cell block<jupyter_code>%matplotlib inline

import cosima_cookbook as cc
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import xarray as xr

from tqdm import tqdm_notebook<jupyter_output><empty_output><jupyter_text>For static plots, the plotting package matplotlib is used. The inline statement tells Jupyter to place those plots within the notebook file.It is common to import standard packages with abbreviated package names such as plt, np, pd, and xr. XArray is fornamed arrays and can be thought of as layer that sits above the netCDF4 package.The package tqdm is for progress bars. These diagnostics are usually very memory and/or computationally expensive.  We leverage the `dask` library http://dask.pydata.org and its related package `distributed` https://distributed.readthedocs.io.<jupyter_code>import dask
import distributed<jupyter_output><empty_output><jupyter_text>By default, we create a collection of workers -- one work for each core. The memory is set as 70% of the total memory of the node. Beyond that, distributed will start caching results locally.<jupyter_code>client = distributed.Client()
client<jupyter_output><empty_output><jupyter_text>You see above that there is a URL for the Dashboard. This is a very useful tool for inspecting the progress of a daskcalculation. If you are running the VDI over VNC you should be able to click on the link to make the dashboard open in another tab.  If you are running this notebook over a SSH tunnel, you will also have to tunnel the port for the dashboard to your local machine. Here's some code which generates the need string.  Run that command on your local machine. Then the link above should work.<jupyter_code>import os

params = {'host': os.environ['HOSTNAME'],
          'user': os.environ['USER'],
          'port': client.scheduler_info()['services']['bokeh']}

tunnel_cmd = "ssh {host}.nci.org.au -l {user} -L {port}:127.0.0.1:{port}".format(**params)
print(tunnel_cmd)<jupyter_output>ssh vdi-n11.nci.org.au -l jm0634 -L 8787:127.0.0.1:8787<jupyter_text>Organization of the model data By default, all of the model output is assumed to stored in the directory given by<jupyter_code>import cosima_cookbook.netcdf_index
cosima_cookbook.netcdf_index.directoriesToSearch<jupyter_output><empty_output><jupyter_text>This global variable may be changed if needed.<jupyter_code>datadir = cosima_cookbook.netcdf_index.directoriesToSearch[0]
!ls {datadir}<jupyter_output>access-om2  access-om2-025  bathymetry	mom01v5  mom025  mom-sis  README  woa13<jupyter_text>The data directory contains several model __configurations__ (mom01v5 or mom025)<jupyter_code>!ls {datadir}/mom01v5<jupyter_output>climatology  KDS75	       KDS75_saltfluxes  KDS75_WOA13
figures      KDS75_PI	       KDS75_UP		 MOM01_Diagnostics.ipynb
GFDL50	     KDS75_salt10days  KDS75_wind<jupyter_text>Each configuration contains a number of __experiments__ (KDS75 or KDS75_wind)<jupyter_code>!ls {datadir}/mom01v5/KDS75_salt10days<jupyter_output>output266  output273  output280  output287  output294  output301   restart300
output267  output274  output281  output288  output295  output302   restart306
output268  output275  output282  output289  output296  output303
output269  output276  output283  output290  output297  output304
output270  output277  output284  output291  output298  output305
output271  output278  output285  output292  output299  output306
output272  output279  output286  output293  output300  restart280<jupyter_text>Which are each made up of a set of several __runs__ (e.g. output266)<jupyter_code>!ls {datadir}/mom01v5/KDS75_salt10days/output266<jupyter_output>config.yaml  field_table   logfile.000000.out  ocean_grid.nc   ocean_scalar.nc
data_table   ice_month.nc  mom.err	       ocean_month.nc  overturning.nc
diag_table   input.nml	   mom.out	       ocean.nc        time_stamp.out<jupyter_text>The actual model out in stored in NetCDF4 files (denoted by the extention .nc).<jupyter_code>fn = os.path.join(datadir, 'mom01v5/KDS75_salt10days/output266/ocean.nc')
xr.open_dataset(fn)<jupyter_output>/g/data1/v45/jm0634/miniconda3/envs/cosima/lib/python3.6/site-packages/xarray/conventions.py:393: RuntimeWarning: Unable to decode time axis into full numpy.datetime64 objects, continuing using dummy netCDF4.datetime objects instead, reason: dates out of range
  result = decode_cf_datetime(example_value, units, calendar)
/g/data1/v45/jm0634/miniconda3/envs/cosima/lib/python3.6/site-packages/xarray/conventions.py:412: RuntimeWarning: Unable to decode time axis into full numpy.datetime64 objects, continuing using dummy netCDF4.datetime objects instead, reason: dates out of range
  calendar=self.calendar)<jupyter_text>There are many, many such NetCDF4 files.<jupyter_code>!find {datadir} -name '*.nc' | wc<jupyter_output><empty_output><jupyter_text>__Runs__ may change which variables they saved and at what temporal resolution over the course of an __experiment__.  Rather than trying to enumerate the variables _a priori_, we can a data discovery approach using a glob.<jupyter_code>directoriesToSearch = ['/g/data3/hh5/tmp/cosima/', 
                      ]
import netCDF4
import dataset
import re
import os
import fnmatch<jupyter_output><empty_output><jupyter_text>Build index of all NetCDF files found in directories to search.<jupyter_code>%%time
m = re.compile('.*\.nc$')

ncfiles = []
for directoryToSearch in directoriesToSearch:
    for root, dirs, filenames in os.walk(directoryToSearch):
        for filename in filenames:
            if m.match(filename) is not None:
                ncfiles.append(os.path.join(root, filename))

print(len(ncfiles))<jupyter_output>17598
CPU times: user 1.69 s, sys: 1.38 s, total: 3.07 s
Wall time: 18.2 s<jupyter_text>We can persist this index by storing it in a sqlite database placed in a centrally available location.<jupyter_code>cosima_cookbook_dir = '/g/data1/v45/cosima-cookbook'
if not os.path.exists(cosima_cookbook_dir):
    os.mkdir(cosima_cookbook_dir)

database_file = '/{}/cosima-cookbook.db'.format(cosima_cookbook_dir)
# os.remove(database_file)<jupyter_output><empty_output><jupyter_text>The use of the `dataset` module hides the details of working with SQL directly.<jupyter_code>db = dataset.connect('sqlite://' + database_file)<jupyter_output><empty_output><jupyter_text>In this database is a single table listing all variables in NetCDF4 seen previously.<jupyter_code>files_already_seen = set([_['ncfile'] for _ in db['ncfiles'].distinct('ncfile')])
print(len(files_already_seen))<jupyter_output>17129<jupyter_text>NetCDF files found on disk not seen before:<jupyter_code>files_to_add = set(ncfiles) - set(files_already_seen)
print(len(files_to_add))<jupyter_output>469<jupyter_text>For these new files, we can determine their configuration, experiment, and run. Using NetCDF4 to get list of all variables in each file.<jupyter_code># output* directories
# match the parent and grandparent directory to configuration/experiment
m = re.compile('(.*)/(.*)/(.*)/(output\d+)/.*\.nc')

def index_variables(ncfile):

    matched = m.match(ncfile)
    if matched is None:
        return []
    
    if not os.path.exists(ncfile):
        return []
    
    try: 
        with netCDF4.Dataset(ncfile) as ds:
            ncvars = [ {'ncfile': ncfile,
                   'rootdir': matched.group(1),
                   'configuration': matched.group(2),
                   'experiment' : matched.group(3),
                   'run' : matched.group(4),
                   'basename' : os.path.basename(ncfile),
                   'variable' : v.name
                   } for v in ds.variables.values()]
    except:
        return []
    
    return ncvars<jupyter_output><empty_output><jupyter_text>Parallel approach<jupyter_code>import dask.bag
from distributed.diagnostics.progressbar import progress
%%time

bag = dask.bag.from_sequence(files_to_add)
bag = bag.map(index_variables).flatten()
ncvars = bag.compute()
print(len(ncvars))
db['ncfiles'].insert_many(ncvars)
import pandas as pd
files = pd.DataFrame(list(db['ncfiles'].all()))
files<jupyter_output><empty_output><jupyter_text>Files seen before but now not found on disk Notice that, unlike the find command above, the glob has only identified .nc files that with in the__configuration__/__experiment__/__run__ directory structure.We want to produce an index over all of these NetCDF files. Once we do that, we can build our diagnostics by firstquerying that index.<jupyter_code>rows = [cosima_cookbook.index_ncfile(fn[0]) for fn in ncfiles[:15]]<jupyter_output><empty_output><jupyter_text>index_ncfile() returns a list of dictionaries describing each variable in an the NetCDF file.<jupyter_code>rows[0][:3]<jupyter_output><empty_output><jupyter_text>The list comprehension gives back a list of lists.  To continue, we first flatten this list of lists:<jupyter_code>rows = [item for sublist in rows for item in sublist]<jupyter_output><empty_output><jupyter_text>Finally, we can convert this dictionary<jupyter_code>df = pd.DataFrame(rows)
df<jupyter_output><empty_output><jupyter_text>So this gives us a nice index for all variables over all NetCDF files.  However, since we have tens of thousands of ncfiles, generating this index can be slow.  To improve performance, we can use a dask bag.<jupyter_code>bag = dask.bag.from_sequence([_[0] for _ in ncfiles],npartitions=1000)
rows = bag.map(cosima_cookbook.index_ncfile)
futures = client.compute(rows)
futures
import dataset
progress(futures)<jupyter_output><empty_output><jupyter_text>To actually get the computation to occur, we can conver the bag into a list.  This takes a few minutes. As before, we convert the list of lists to a single list<jupyter_code>rows = [item for sublist in rows for item in sublist]<jupyter_output><empty_output><jupyter_text>Finally, we can put this all into a pandas DataFrame<jupyter_code>df = pd.DataFrame.from_records(rows)
df<jupyter_output><empty_output><jupyter_text>The above steps are implemented in<jupyter_code>df = cosima_cookbook.build_index()<jupyter_output><empty_output><jupyter_text>Here are all of the unique experiments found in the data<jupyter_code>expts = df.experiment.unique()
expts
expts[7]
list(df)<jupyter_output><empty_output><jupyter_text>Calculation of EKE Let's choose a specific experiment:<jupyter_code>pd.DataFrame(list(db['ncfiles'].distinct('configuration', 'experiment')))
db['ncfiles'].columns
expt = 'KDS75'
expt<jupyter_output><empty_output><jupyter_text>To calculate the eddy kinetic energy, we are going to consider only  portions of simulations which have 5-day average velocities saved, which means directories with `ocean__*.nc` files.<jupyter_code>res = db.query('SELECT ncfile FROM ncfiles \
                WHERE experiment = "KDS75" \
                AND basename LIKE "%ocean__%" \
                AND variable = "u" \
                ORDER BY ncfile \
               ')
ncfiles = [row['ncfile'] for row in res]
ncfiles<jupyter_output><empty_output><jupyter_text>Using our index of ncfiles, we search for such nc files.<jupyter_code>u = cosima_cookbook.get_nc_variable(expt, 'ocean__', 'u', n=6)
@cosima_cookbook.memory.cache
def calc_eke(expt, box_index):
    yi, xi = box_index
   
    box = {'yu_ocean': slice(300*yi, 300*(yi+1)),
           'xu_ocean': slice(400*xi, 400*(xi+1))}
    
    op = lambda p: p.isel(**box)

    u = get_nc_variable(expt, 'ocean__', 'u', op=op, n=72)
    v = get_nc_variable(expt, 'ocean__', 'v', op=op, n=72)
    
    u_avg = u.mean('time')
    v_avg = v.mean('time')

    MKE = 0.5 * (u_avg**2 + v_avg**2)
    MKE = MKE.sum(dim='st_ocean')
    MKE = MKE.to_dataset(name='MKE')
    
    u_ = u - u_avg
    v_ = v - v_avg
    
    EKE = 0.5 * (u_**2 + v_**2)

    EKE = EKE.sum(dim='st_ocean')
    EKE = EKE.to_dataset(name='EKE')

    dsx = xr.merge([MKE, EKE])
    dsx.load()
        
    return dsx
from itertools import product

# ceil(x/y) = (x+y+1)//y
yi = range(2700//(300)) 
xi = range(3600//(400))

box_indexes = list(product(*[yi, xi]))<jupyter_output><empty_output><jupyter_text>Maps<jupyter_code>from mpl_toolkits.basemap import Basemap, shiftgrid
expt = expts[1]

print(expt)
# Plot in basemap 

plt.figure(figsize=(15,6))
lev = np.arange(0, 1.0, 0.05)
map = Basemap(projection='mbtfpq',
              lon_0 = -100, resolution='l')
map.drawcoastlines(linewidth=0.25)
map.fillcontinents(color='gray',lake_color='gray')
map.drawparallels(np.arange(-60.,61.,30.),
                  labels=[True,False,False,False])
map.drawmeridians(np.arange(-180.,181.,90.),
                  labels=[False,False,False,True])

for box_index in tqdm_notebook(box_indexes):
    #print(box_index)
    dsx = calc_eke(expt, box_index)
    
    x=dsx.xu_ocean[:]
    y=dsx.yu_ocean[:]
    lon, lat = np.meshgrid(x, y)
    
    X, Y = map(lon,lat) 
    
    map.contourf(X, Y, dsx.MKE,
                 cmap=plt.cm.hot,
                 levels=lev,extend='both')

cb = plt.colorbar(orientation='vertical',shrink = 0.7)

cb.ax.set_xlabel('m^2 s$^{-2}$')
plt.title('Mean KE {}'.format(expt))<jupyter_output>mom01v5/KDS75
{'fraction_non_alphanumeric': 0.09302168021680217, 'fraction_numerical': 0.028184281842818428, 'mean_word_length': 4.591287878787878, 'pattern_counts': {'":': 0, '<': 119, '<?xml version=': 0, '>': 119, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '32', 'dir': 'jupyter-structured-clean-dedup', 'id': '658788', 'license': 'Apache-2.0', 'n_tokens_mistral': 5133, 'n_tokens_neox': 4791, 'n_words': 1410, 'path': 'DocumentedExamples/sample_diagnostic.ipynb', 'repo_name': 'ruth-moorman/cosima-recipes', 'size': '14760'}