Regridding of coarse to fine w/ dask. Example of using dask.map_blocks and of writing dask files as mds stores.

gendata_justregrid.py

coarsename = '/scr/jklymak/LeeCoarse/coarse3d%s01U10/_Model/input/'%runtype
print(coarsename)

def interpolateAtDepth(X,T,x0,y0,x,y,defaultval,block_id=None):
    import scipy.interpolate
    good = np.where(T[0,0,:,:]>20.)
    bad = np.where(T[0,0,:,:]<=20.)
    Tt = X.copy()[0,0,:,:]
    print(Tt)
    #print(good)
    if (len(good[0])>0) & (len(bad[0])>0):
        Tt[bad]= Tt[good].mean()
    else:
        Tt[bad]=defaultval
    if np.shape(X)[-1]>1:
        xout=np.zeros((1,1,ny,nx))   
        fit=scipy.interpolate.RectBivariateSpline(x0,y0,Tt.T)
        xout = fit(x,y).T 
    else:
        xout=np.ones((1,1,1,1))
    print(block_id)
    return xout

class writable_mds_store:
    '''
    This will write each chunk to a binary file.  Really only works if chunks are the k-index of the matrix
    '''
    def __init__(self, prefix, itern, suffix='data', dtype='>f4'):
        self.prefix = prefix
        self.itern = itern
        self.suffix = suffix
        self.dtype = dtype
        self.fname = '%s.%010d.%s' % (self.prefix, self.itern, self.suffix)
        # zero the file:
        with open(self.fname,'wb+') as f:
            pass

    def __setitem__(self, idx, data):
        import os
        tslice = idx[0]
        assert tslice.step is None
        assert (tslice.stop - tslice.start) == 1
        n = tslice.start
        nbytes = data.astype(self.dtype).nbytes
        kslice = idx[1]
        with open(self.fname,'rb+') as f:
            f.seek(nbytes*kslice.start,os.SEEK_SET)
            f.write(data.astype(self.dtype).tobytes())


# one-d field...
todo = 'Eta'
print('Regriding %s'%todo)
data = xmitgcm.open_mdsdataset(dirname=coarsename,
                               prefix={todo},iters=36000,read_grid=True,
                               geometry='cartesian',endian='<')
# x, y, nx, ny are determined elsewhere, but set the new grid...
fit=scipy.interpolate.RectBivariateSpline(data['XC'].values,data['YC'].values,data[todo].T)
xout = fit(x,y).T
print(np.shape(xout))
with open(indir+'/%sinit.bin'%todo,'wb') as f:
    xout.tofile(f)
#tm.store(tnew_store,lock=True)
#tm.close()
print('Done regriding %s'%todo)

if 1:
    for todo,defaultval in zip(['T','U','V'],[27.5,0.1,0]):
        print('Regriding %s'%todo)
        data = xmitgcm.open_mdsdataset(dirname=coarsename,
                                       prefix=[todo,'T'],iters=36000,
                                       read_grid=True,geometry='cartesian',endian='<',
                                       chunks={'Z':1,'time':1})
        print(data)
        # x, y, nx, ny are determined elsewhere, but set the new grid...
        tm = da.map_blocks(interpolateAtDepth,
                           data[todo].data,data['T'].data,
                           data['XC'].values,data['YC'].values,
                           x,y,defaultval,
                           chunks=(1,1,ny,nx),
                           dtype=data[todo].data.dtype)

        tnew_store = writable_mds_store(indir+'/%sinit'%todo,0,dtype=tm.dtype,suffix='bin')
        tm.store(tnew_store,lock=True)
        #tm.close()
        print('Done regriding %s'%todo)