bin_density.py

# output from the MITgcm using xmitgcm

import xarray as xr
import numpy as np
import xgcm

with xr.open_dataset(fname, decode_times=False) as ds:
    grid = xgcm.Grid(ds, coords={'Z':{'center':'Z', 'outer':'Zp1'}}, periodic=False)
    ds = preprocess(ds)

    pdlevels = ds['pden'][0, :, -3].values[0::3]

    T2 = grid.transform((ds.TRAC02) * ds.drF , 'Z', pdlevels, target_data=ds.pden, method='conservative')
    pd = T2.to_dataset(name='TRAC02')
    for td in ['TRAC01', 'TRAC03']:
        pd[td] = grid.transform((ds[td]) * ds.drF , 'Z', pdlevels, target_data=ds.pden, method='conservative')

    pdZ = grid.transform(ds.Z.values[np.newaxis,:, np.newaxis] + ds.TRAC01 * 0, 'Z', pdlevels, target_data=ds.pden)
    pdDz = -pdZ.diff(dim='pden')
#. to get the mean concentration:
x = pd['TRAC01'] / pdDz.values