quickdiags

#!/usr/bin/env python

# Model parameters to assume for this data:
tracers = ('CO2','CH4','TCO','CFF','CBB','CLA','COC')
nlev = 79   # Number of thermodynamic levels
debug = False
grav = .980616e+1  # Taken from GEM-MACH file chm_consphychm_mod.ftn90

from argparse import ArgumentParser
import subprocess
parser = ArgumentParser(description='Performs some quick diagnostics on EC-CAS model output.')
parser.add_argument('infile', help='Input file or directory to process.')
parser.add_argument('outdir', nargs='?', default='.', help='Where to put the output files.  Default is the current directory.')
parser.add_argument('--label', help='The name of the experiment.  If not specified, then a name will be determined from the input directory structure.')
parser.add_argument('--nthreads', type=int, default=int(subprocess.check_output('nproc')), help='Number of threads to use for doing the calcuations.  Default is currently %(default)s.')
args = parser.parse_args()

try:
  import fstd2nc
  import rpnpy.librmn.all as rmn
  import rpnpy.vgd.all as vgd
except ImportError:
  parser.error("You need to run the following command before using the script:\n\n. ssmuse-sh -p eccc/crd/ccmr/EC-CAS/master/fstd2nc_0.20180821.0\n")

from os.path import dirname
if args.label is None:
  dirs = dirname(args.infile).split('/')
  while dirs[-1] in ('model','pres','time_series','','.'):
    dirs.pop()
  args.label = dirs[-1]

import netCDF4

# Helper functions

# Class for handling dependencies on calculations.
class Calc (object):
  __slots__ = ('func','inputs','callbacks','__weakref__','lock')
  def __init__ (self, func, *inputs):
    import weakref
    from threading import Lock
    self.func = func
    self.inputs = list(inputs)
    self.lock = Lock()
    for inp in inputs:
      if isinstance(inp,Calc):
        if debug: print '%s will call back %s'%(inp,func.func_name)
        inp.callbacks.add(self)
    self.callbacks = weakref.WeakSet()
  def trigger (self):
    if debug: print 'trigger %s? %s'%(self.func.func_name, [not isinstance(inp,Calc) for inp in self.inputs])
    if not any(isinstance(inp,Calc) for inp in self.inputs):
      if debug: print 'triggered'
      data = self.func(*self.inputs)
      self._do_callback (data)
    else:
      if debug: print 'not triggered'
  def _do_callback (self, data):
    if debug: print 'callbacks:', list(self.callbacks)
    for callback in self.callbacks:
      ids = list(map(id,callback.inputs))
      if id(self) not in ids:
        print 'weird', self.func, callback.func, id(self), ids, map(type,callback.inputs)
        return
      i = ids.index(id(self))
      with callback.lock:
        callback.inputs[i] = data
        callback.trigger()
  def __del__ (self):
    if debug: print 'deleting %s'%(self.func.func_name)

# Decorate for turning a regular function into one that handles dependent
# calculations.
import weakref
def calc (func, lookup=weakref.WeakValueDictionary()):
  from functools import wraps
  @wraps(func)
  def f (*args):
    key = (func,tuple(args))
    if debug and key in lookup:
      print 'reusing existing %s (%s)'%(func.func_name, ','.join(map(str,args)))
    elif debug:
      print 'generating new %s (%s)'%(func.func_name, ','.join(map(str,args)))
    result = lookup.get(key,None)
    if result is None:
      result = Calc(func,*args)
    lookup[key] = result
    return result
  return f
del weakref

# Keep track of all records that are used.
all_records = set()

from threading import Lock
read_lock = Lock()

@calc
def _read (rec_id):
  with read_lock:
    if debug: print 'read'
    all_records.remove(rec_id)
    return np.asarray(fstluk(rec_id),dtype='float64')
def read (rec_id):
  all_records.add(rec_id)
  return _read (rec_id)


# Object for managing writes into a netCDF4 dataset.
class Writer (Calc):
  def __init__ (self, dataset):
    Calc.__init__(self,None)
    self.dataset = dataset
  def trigger (self):
    if not any(isinstance(inp,Calc) for inp in self.inputs):
      if debug: print 'close file'
      self.dataset.close()
  def write (self, varname, ind, data):
    if not isinstance(data,Calc):
      self.dataset[varname][ind] = data
      return
    inp = _write (self.dataset[varname], ind, data)
    self.inputs.append(inp)
    inp.callbacks.add(self)
  def __del__ (self):
    return

# Object for doing an individual write.
from threading import Lock
write_lock = Lock()
@calc
def _write (var, ind, data):
  with write_lock:
    if debug: print 'write'
    var[ind] = data


@calc
def logp (p0):
  if debug: print 'logp'
  import numpy as np
  return np.log(p0/1000.)

@calc
def log_pres (a, b, s):
  if debug: print 'log_pres'
  import numpy as np
  return a+b*s

@calc
def pres (a, b, s):
  if debug: print 'pres'
  import numpy as np
  return np.exp(a+b*s)

@calc
def layer_mass (c,q,p_above,p_below):
  if debug: print 'layer mass'
  return c*(1-q)*(p_below-p_above)

@calc
def area_integrate (x, dx, scale):
  if debug: print 'area integrate'
  import numpy as np
  # Remove repeated longitude
  x = x[:,:-1]
  dx = dx[:,:-1]
  return np.sum(x*dx*scale)

@calc
def layer_sum (*inputs):
  if debug: print 'layer sum'
  return sum(inputs)

@calc
def quotient (numerator, denominator):
  if debug: print 'quotient'
  return numerator / denominator

@calc
def stack (*inputs):
  if debug: print 'stack'
  import numpy as np
  return np.stack(inputs)

@calc
def vinterp (x, z, N):
  if debug: print 'vinterp'
  import numpy as np

  # Number of model levels and horizontal grid cells.
  nlev, nj, ni = x.shape
  ninj = ni*nj

  # Workspace for building the interpolated values.
  work = np.zeros((N+2,ninj), dtype=x.dtype)

  x = x.reshape(nlev,ninj)
  z = z.reshape(nlev,ninj)

  # Calculate rates of change
  dxdz = np.zeros((nlev+1,ninj),dtype=x.dtype)
  dxdz[1:-1] = np.diff(x,axis=0)
  dxdz[1:-1] /= np.diff(z,axis=0)
  ddxdz = np.diff(dxdz,axis=0)
  ind = np.arange(ninj)

  # Loop over one model level at a time.
  # The advanced numpy indexing doesn't work well if there are repeated
  # indices, such as repeated target levels.
  for i in range(nlev):
    c = np.ceil(z[i]).astype(int)
    c[c<0] = 0
    c[c>N] = N
    # Mark the transition between model levels.
    work[c,ind] += (c-z[i])*ddxdz[i]
    work[c+1,ind] -= (c-1-z[i])*ddxdz[i]
    # Set the lower boundary value.
    if i == 0:
      work[c,ind] += x[0]
      work[c+1,ind] -= x[0]

  c[c==z[i]] += 1
  work[c,ind] = np.float('nan')
  np.cumsum(work,axis=0,out=work)  # magic
  np.cumsum(work,axis=0,out=work)  # more magic
  work[work==0] = np.float('nan')

  return work[:N].reshape(N,nj,ni)

@calc
def zonal_mean (x):
  if debug: print 'zonal mean'
  import numpy as np
  # Remove repeated longitude
  x = x[...,:-1]
  return np.nanmean(x,axis=-1)


import numpy as np

fstd2nc.stdout.streams = ('error',)
if not debug:
  rmn.fstopt('MSGLVL','ERRORS')
b = fstd2nc.Buffer(args.infile, ignore_diag_level=True, rpnstd_metadata=True, progress=True, unique_names=False)

# Call c_fstluk directly
all_keys = np.array(b._headers['key'])
all_ni = np.array(b._headers['ni'])
all_nj = np.array(b._headers['nj'])
all_nk = np.array(b._headers['nk'])
all_shape = list(zip(all_ni,all_nj,all_nk))
all_datyp = np.array(b._headers['datyp'])
all_nbits = np.array(b._headers['nbits'])
all_file_id = np.array(b._headers['file_id'])

def fstluk (rec_id):
  import ctypes as _ct
  import numpy.ctypeslib as _npc
  from rpnpy.librmn import proto as _rp
  from rpnpy.librmn.fstd98 import dtype_fst2numpy
  import numpy as np
  file_id = all_file_id[rec_id]
  b._open(file_id)
  key = (all_keys[rec_id]<<10) + b._opened_librmn_index
  shape = all_shape[rec_id]
  dtype = dtype_fst2numpy(int(all_datyp[rec_id]),int(all_nbits[rec_id]))
  _rp.c_fstluk.argtypes = (_npc.ndpointer(dtype=dtype), _ct.c_int,
                           _ct.POINTER(_ct.c_int), _ct.POINTER(_ct.c_int),
                           _ct.POINTER(_ct.c_int))
  data = np.empty(shape, dtype=dtype, order='FORTRAN')
  (cni, cnj, cnk) = (_ct.c_int(), _ct.c_int(), _ct.c_int())
  istat = _rp.c_fstluk(data, key, _ct.byref(cni), _ct.byref(cnj),
                       _ct.byref(cnk))
  if istat < 0:
    raise FSTDError()
  return data.transpose().squeeze()

# Get vertical parameters for this data.
vgd_id = vgd.vgd_read(b._opened_funit)
am = vgd.vgd_get(vgd_id,'CA_M')[:-1]
bm = vgd.vgd_get(vgd_id,'CB_M')[:-1]
at = vgd.vgd_get(vgd_id,'CA_T')[:-2]
bt = vgd.vgd_get(vgd_id,'CB_T')[:-2]
assert len(am) == nlev+1
assert len(at) == nlev
assert np.all(am[:-1] < at)
assert np.all(am[1:] > at)

# Structure the records into time/level dimensions.
b._makevars()
vardict = {v.name:v for v in b._varlist}

# Get grid cell areas.
dx = read(vardict['DX'].record_id[0])

# Select records with full vertical extent
def get_3d_recs (recs):
  return recs[np.all(recs>=0, axis=1),:]

# Get unique timesteps from records
def get_timesteps (recs):
  if recs.ndim > 1:
    recs = recs[:,0]
  return set(b._headers[recs]['datev'])

# Limit records to those that fall on particular timesteps
def limit_timesteps (recs, target_timesteps):
  if recs.ndim > 1:
    timesteps = b._headers[recs[:,0]]['datev']
  else:
    timesteps = b._headers[recs]['datev']
  return recs[np.isin(timesteps,list(target_timesteps))]

##############################################################################
# Column and mass diagnostics

# Find timesteps available for mass diagnostics.
mass_timesteps = get_timesteps(get_3d_recs(vardict['HU'].record_id)) \
               & get_timesteps(vardict['P0'].record_id)

HU_recs = limit_timesteps(get_3d_recs(vardict['HU'].record_id), mass_timesteps)
P0_recs = limit_timesteps(vardict['P0'].record_id, mass_timesteps)

# Find the available timesteps and levels.
mass_timesteps = b._headers[HU_recs[:,0]]['datev']
mass_levels = b._headers[HU_recs[0,:]]['level']
assert np.all(vgd.vgd_get(vgd_id,'VIPT')[:-2] == b._headers[HU_recs[0,:]]['ip1'])

# Determine column / mass calculations.
total_filename = "%s/%s_totalmass.nc"%(args.outdir,args.label)
tropo_filename = "%s/%s_tropomass.nc"%(args.outdir,args.label)
colavg_filename = "%s/%s_column_avg.nc"%(args.outdir,args.label)
tropoavg_filename = "%s/%s_tropocolumn_avg.nc"%(args.outdir,args.label)
totalmass = netCDF4.Dataset(total_filename,'w')
tropomass = netCDF4.Dataset(tropo_filename,'w')
colavg = netCDF4.Dataset(colavg_filename,'w')
tropoavg = netCDF4.Dataset(tropoavg_filename,'w')
times = fstd2nc.mixins.dates.stamp2datetime(mass_timesteps)
times = np.array(times,dtype='datetime64[s]')
for dataset in totalmass, tropomass:
  dataset.createDimension('time', None)
  dataset.createVariable(varname='time', datatype='float32', dimensions=('time',))
  dataset['time'].setncatts({'units':'hours since 2015-01-01'})
  dataset['time'][:] = netCDF4.date2num(times.tolist(),units='hours since 2015-01-01')
  for tracer in tracers:
    if tracer not in vardict: continue
    dataset.createVariable(varname=tracer, datatype='float64', dimensions=('time',))
  dataset.createVariable(varname='air', datatype='float64', dimensions=('time',))
  dataset.createVariable(varname='dry_air', datatype='float64', dimensions=('time',))
for dataset in colavg, tropoavg:
  dataset.createDimension('time', None)
  dataset.createVariable(varname='time', datatype='float32', dimensions=('time',))
  dataset['time'].setncatts({'units':'hours since 2015-01-01'})
  dataset['time'][:] = netCDF4.date2num(times.tolist(),units='hours since 2015-01-01')
  dataset.createDimension('lat', all_nj[P0_recs[0]])
  dataset.createVariable(varname='lat', datatype='float32', dimensions=('lat',))
  dataset['lat'].setncatts({'units':'degrees_north'})
  dataset['lat'][:] = vardict['P0'].axes[-2].array
  dataset.createDimension('lon', all_ni[P0_recs[0]])
  dataset.createVariable(varname='lon', datatype='float32', dimensions=('lon',))
  dataset['lon'].setncatts({'units':'degrees_east'})
  dataset['lon'][:] = vardict['P0'].axes[-1].array
  for tracer in tracers:
    if tracer not in vardict: continue
    dataset.createVariable(varname=tracer, datatype='float64', dimensions=('time','lat','lon'))

totalmass = Writer(totalmass)
tropomass = Writer(tropomass)
colavg = Writer(colavg)
tropoavg = Writer(tropoavg)

# Dry-air mass
dryair_sum_save = []  # Saved for use later.
dryair_tropo_sum_save = []  # Saved for use later.
for t in range(len(mass_timesteps)):
  dryair_layers = []
  dryair_tropo_layers = []
  for k in range(len(mass_levels)):
    q = read(HU_recs[t,k])
    s = logp(read(P0_recs[t]))
    p_above = pres(am[k],bm[k],s)
    p_below = pres(am[k+1],bm[k+1],s)
    layer = layer_mass(1,q,p_above,p_below)
    dryair_layers.append(layer)
    if mass_levels[k] > 0.2:
      dryair_tropo_layers.append(layer)
  dryair_sum = layer_sum(*dryair_layers)
  dryair_sum_save.append(dryair_sum)
  mass = area_integrate(dryair_sum, dx, 1/grav*1E-12)
  totalmass.write ('dry_air', t, mass)
  dryair_tropo_sum = layer_sum(*dryair_tropo_layers)
  dryair_tropo_sum_save.append(dryair_tropo_sum)
  mass = area_integrate(dryair_tropo_sum, dx, 1/grav*1E-12)
  tropomass.write ('dry_air', t, mass)
# Total air mass
for t in range(len(mass_timesteps)):
  s = logp(read(P0_recs[t]))
  p_above = pres(am[0],bm[0],s)
  p_below = pres(am[nlev],bm[nlev],s)
  mass = layer_mass(1,0,p_above,p_below)
  mass = area_integrate(mass, dx, 1/grav*1E-12)
  totalmass.write ('air', t, mass)
  mass = layer_mass(1,0,20000,p_below)
  mass = area_integrate(mass, dx, 1/grav*1E-12)
  tropomass.write ('air', t, mass)
# Tracer mass
for tracer in tracers:
  if tracer not in vardict: continue
  recs = limit_timesteps(get_3d_recs(vardict[tracer].record_id), mass_timesteps)
  for it,t in enumerate(np.searchsorted(mass_timesteps, b._headers[recs[:,0]]['datev'])):
    layers = []
    tropo_layers = []
    for ik,k in enumerate(np.searchsorted(mass_levels, b._headers[recs[0,:]]['level'])):
      c = read(recs[it,ik])
      q = read(HU_recs[t,k])
      s = logp(read(P0_recs[t]))
      p_above = pres(am[k],bm[k],s)
      p_below = pres(am[k+1],bm[k+1],s)
      layer = layer_mass(c,q,p_above,p_below)
      layers.append(layer)
      if mass_levels[k] > 0.2:
        tropo_layers.append(layer)
    tracer_sum = layer_sum(*layers)
    mass = area_integrate(tracer_sum, dx, 1E-9/grav*1E-12)
    totalmass.write (tracer, t, mass)
    dryair_sum = dryair_sum_save[t]
    tracer_col = quotient(tracer_sum, dryair_sum)
    colavg.write (tracer, t, tracer_col)
    tracer_sum = layer_sum(*tropo_layers)
    mass = area_integrate(tracer_sum, dx, 1E-9/grav*1E-12)
    tropomass.write (tracer, t, mass)
    dryair_tropo_sum = dryair_tropo_sum_save[t]
    tracer_col = quotient(tracer_sum, dryair_tropo_sum)
    tropoavg.write (tracer, t, tracer_col)

del dryair_sum_save, dryair_tropo_sum_save

##############################################################################
# Zonal diagnostics

# Find timesteps available for zonal mean diagnostics on gph.

if 'GZ' in vardict:

  gph_timesteps = get_timesteps(get_3d_recs(vardict['GZ'].record_id))
  tracer_timesteps = set.union(*[get_timesteps(get_3d_recs(vardict[tracer].record_id)) for tracer in tracers if tracer in vardict])
  gph_timesteps = gph_timesteps & tracer_timesteps

  # Get thermodynamic GZ records.
  GZ_recs = limit_timesteps(get_3d_recs(vardict['GZ'].record_id), gph_timesteps)[:,1::2]

  # Find the available timesteps and levels.
  gph_timesteps = b._headers[GZ_recs[:,0]]['datev']
  gph_levels = b._headers[GZ_recs[0,:]]['level']
  assert np.all(vgd.vgd_get(vgd_id,'VIPT')[:-2] == b._headers[GZ_recs[0,:]]['ip1'])

  # Determine zonal mean calculations.
  filename = "%s/%s_zonalmean_gph.nc"%(args.outdir,args.label)
  zonalmean_gph = netCDF4.Dataset(filename,'w')
  times = fstd2nc.mixins.dates.stamp2datetime(gph_timesteps)
  times = np.array(times,dtype='datetime64[s]')

  zonalmean_gph.createDimension('time', None)
  zonalmean_gph.createVariable(varname='time', datatype='float32', dimensions=('time',))
  zonalmean_gph['time'].setncatts({'units':'hours since 2015-01-01'})
  zonalmean_gph['time'][:] = netCDF4.date2num(times.tolist(),units='hours since 2015-01-01')
  zonalmean_gph.createDimension('height', 68)
  zonalmean_gph.createVariable(varname='height', datatype='float32', dimensions=('height',))
  zonalmean_gph['height'].setncatts({'units':'km'})
  zonalmean_gph['height'][:] = np.arange(68)[::-1]
  zonalmean_gph.createDimension('lat', all_nj[GZ_recs[0,0]])
  zonalmean_gph.createVariable(varname='lat', datatype='float32', dimensions=('lat',))
  zonalmean_gph['lat'].setncatts({'units':'degrees_north'})
  zonalmean_gph['lat'][:] = vardict['GZ'].axes[-2].array

  for tracer in tracers:
    if tracer not in vardict: continue
    zonalmean_gph.createVariable(varname=tracer, datatype='float64', dimensions=('time','height','lat'))

  zonalmean_gph = Writer(zonalmean_gph)
  gph_indices = calc(lambda gz: 67-gz/100.)

  for tracer in tracers:
    if tracer not in vardict: continue
    recs = limit_timesteps(get_3d_recs(vardict[tracer].record_id), gph_timesteps)
    for it,t in enumerate(np.searchsorted(gph_timesteps, b._headers[recs[:,0]]['datev'])):
      x = []
      gz = []
      for ik,k in enumerate(np.searchsorted(gph_levels, b._headers[recs[0,:]]['level'])):
        x.append(read(recs[it,ik]))
        gz.append(read(GZ_recs[t,k]))
      x = stack(*x)
      gz = stack(*gz)
      z = gph_indices(gz)
      zonalmean_gph.write (tracer, t, zonal_mean(vinterp (x, z, 68)))

# Find timesteps available for zonal mean diagnostics on pressure levels.
pres_timesteps = get_timesteps(vardict['P0'].record_id)
tracer_timesteps = set.union(*[get_timesteps(get_3d_recs(vardict[tracer].record_id)) for tracer in tracers if tracer in vardict])
pres_timesteps = pres_timesteps & tracer_timesteps
P0_recs = limit_timesteps(vardict['P0'].record_id, pres_timesteps)

# Find the available timesteps and levels.
pres_timesteps = b._headers[P0_recs[:]]['datev']

# Determine zonal mean calculations.
filename = "%s/%s_zonalmean_plev.nc"%(args.outdir,args.label)
tropo_filename = "%s/%s_zonalmean_tropo.nc"%(args.outdir,args.label)
zonalmean_plev = netCDF4.Dataset(filename,'w')
zonalmean_tropo = netCDF4.Dataset(tropo_filename,'w')
times = fstd2nc.mixins.dates.stamp2datetime(pres_timesteps)
times = np.array(times,dtype='datetime64[s]')

zonalmean_plev.createDimension('time', None)
zonalmean_plev.createVariable(varname='time', datatype='float32', dimensions=('time',))
zonalmean_plev['time'].setncatts({'units':'hours since 2015-01-01'})
zonalmean_plev['time'][:] = netCDF4.date2num(times.tolist(),units='hours since 2015-01-01')
zonalmean_plev.createDimension('pres', 100)
zonalmean_plev.createVariable(varname='pres', datatype='float32', dimensions=('pres',))
zonalmean_plev['pres'].setncatts({'units':'hPa','positive':'down'})
zonalmean_plev['pres'][:] = np.logspace(-1,3,100)
zonalmean_plev.createDimension('lat', all_nj[P0_recs[0]])
zonalmean_plev.createVariable(varname='lat', datatype='float32', dimensions=('lat',))
zonalmean_plev['lat'].setncatts({'units':'degrees_north'})
zonalmean_plev['lat'][:] = vardict['P0'].axes[-2].array

zonalmean_tropo.createDimension('time', None)
zonalmean_tropo.createVariable(varname='time', datatype='float32', dimensions=('time',))
zonalmean_tropo['time'].setncatts({'units':'hours since 2015-01-01'})
zonalmean_tropo['time'][:] = netCDF4.date2num(times.tolist(),units='hours since 2015-01-01')
zonalmean_tropo.createDimension('pres', 81)
zonalmean_tropo.createVariable(varname='pres', datatype='float32', dimensions=('pres',))
zonalmean_tropo['pres'].setncatts({'units':'hPa','positive':'down'})
zonalmean_tropo['pres'][:] = np.linspace(200,1000,81)
zonalmean_tropo.createDimension('lat', all_nj[P0_recs[0]])
zonalmean_tropo.createVariable(varname='lat', datatype='float32', dimensions=('lat',))
zonalmean_tropo['lat'].setncatts({'units':'degrees_north'})
zonalmean_tropo['lat'][:] = vardict['P0'].axes[-2].array

for tracer in tracers:
  if tracer not in vardict: continue
  zonalmean_plev.createVariable(varname=tracer, datatype='float64', dimensions=('time','pres','lat'))
  zonalmean_tropo.createVariable(varname=tracer, datatype='float64', dimensions=('time','pres','lat'))

zonalmean_plev = Writer(zonalmean_plev)
zonalmean_tropo = Writer(zonalmean_tropo)
pres_indices = calc(lambda lp: 99*(lp-np.log(100)-np.log(0.1))/(np.log(1000)-np.log(0.1)))
tropo_indices = calc(lambda p: 80*(p/100-200)/(1000-200))

for tracer in tracers:
  if tracer not in vardict: continue
  recs = limit_timesteps(get_3d_recs(vardict[tracer].record_id), pres_timesteps)
  for it,t in enumerate(np.searchsorted(pres_timesteps, b._headers[recs[:,0]]['datev'])):
    x = []
    lp = []
    p = []
    for ik,k in enumerate(np.searchsorted(gph_levels, b._headers[recs[0,:]]['level'])):
      x.append(read(recs[it,ik]))
      s = logp(read(P0_recs[t]))
      lp.append(log_pres(at[k],bt[k],s))
      p.append(pres(at[k],bt[k],s))
    x = stack(*x)
    lp = stack(*lp)
    p = stack(*p)
    z = pres_indices(lp)
    zonalmean_plev.write (tracer, t, zonal_mean(vinterp (x, z, 100)))
    z = tropo_indices(p)
    zonalmean_tropo.write (tracer, t, zonal_mean(vinterp (x, z, 81)))


if debug: print 'starting I/O'

# Ignore numpy warnings about things like "mean of empty slice".
import warnings
warnings.simplefilter("ignore")

# Read DX first.
# These records seem to be found at the end of the RPN file, so all the other
# operations will be waiting a long time if this isn't pre-loaded.
dx.trigger()

pbar = fstd2nc.mixins._ProgressBar("Calculating diagnostics", suffix='%(percent)d%% [%(myeta)s]', max=len(all_records))

from multiprocessing.pool import ThreadPool
from threading import Lock
pool = ThreadPool(args.nthreads)
pbar_lock = Lock()
def update_pbar (stuff):
  with pbar_lock:
    pbar.next()

for rec_id in sorted(all_records):
  pool.apply_async(read(rec_id).trigger, callback=update_pbar)
  #read(rec_id).trigger()
  #if not debug: pbar.next()
pool.close()
pool.join()
pbar.finish()