data-test-2-planes.py

# to run unit tests:
#
# 2-planes-flow-mesh.py data-test-2-planes 16 8 1
# 2-planes-invmap.py data-test-2-planes 16 8
# 2-planes-test.py data-test-2-planes 16 8

import numpy as np

shift=8
nx,ny=256,128
nxc,nyc=128,64
nxtc,nytc=64,32
dtype=np.uint8

def load_data(basepath, top, bot):
    tmp = np.random.randint(0, np.iinfo(dtype).max//2, size=(nx+shift,ny+shift), dtype=dtype)
    tmp[nx//2:shift+nx//2, ny//2:shift+ny//2] = np.iinfo(dtype).max
    ttop = tmp[:nx, :ny]
    tbot = tmp[shift:shift+nx, shift:shift+ny]
    return ttop, tbot

def crop_data(ttop, tbot):
    ix = (nx-nxc)//2
    iy = (ny-nyc)//2
    return ttop[ix:ix+nxc, iy:iy+nyc], tbot[ix:ix+nxc, iy:iy+nyc], (iy,ix), (nyc,nxc)

def tight_crop_data(ttop, tbot):
    ix = (nx-nxtc)//2
    iy = (ny-nytc)//2
    return ttop[ix:ix+nxtc, iy:iy+nytc], tbot[ix:ix+nxtc, iy:iy+nytc], (iy,ix), (nytc,nxtc)

def save_flow(flow, basepath, params):
    np.save('1.flow'+params+'.npy', flow)

def save_mesh(mesh, basepath, params):
    return np.save('2.mesh'+params+'.npy', mesh)

def load_flow(basepath, params):
    flow = np.load('1.flow'+params+'.npy')

def load_mesh(basepath, params):
    return np.load('2.mesh'+params+'.npy')

def save_invmap(invmap, basepath, params):
    np.save('3.invmap'+params+'.npy', invmap)

def load_invmap(basepath, params):
    invmap = np.load('3.invmap'+params+'.npy')
    return invmap

def write_warp(shape, chunk_size, basepath, params, planes):
    return ts.open({
        'driver': 'zarr',
        'kvstore': {"driver":"file", "path":'4.warped.'+params+'.zarr')},
        'metadata': {
            "compressor":{"id":"zstd","level":3},
            "shape":shape,
            "chunks":[2,16,16],
            "fill_value":0,
            'dtype': '|u1',
            'dimension_separator': '/',
        },
        'create': True,
        'delete_existing': True,
        }).result().write(planes).result()