suspendtools/helper.py

def load_mat(f,var,squeeze=True,cast_string=True,cast_integer=False):
    try:
        return load_matbin(f,var,squeeze=squeeze)
    except NotImplementedError:
        return load_matv73(f,var,squeeze=squeeze,cast_string=cast_string,cast_integer=cast_integer)

def load_matbin(f,var,squeeze=True):
    '''Loads variables from a matlab binary mat file (not v7.3).'''
    from scipy.io import loadmat
    import numpy
    data = loadmat(f,variable_names=var,squeeze_me=squeeze,struct_as_record=True)
    def parse(val):
        if not isinstance(val,numpy.ndarray):
            return val
        if val.dtype.names is not None:
            d = {}
            for key in val.dtype.names:
                d[key] = parse(val[key].item())
            return d
        else:
            return val
    if isinstance(var,(tuple,list)):
        return tuple(parse(data[x]) for x in var)
    else:
        return parse(data[var])

def load_matv73(f,var,squeeze=True,cast_string=True,cast_integer=False):
    '''Loads variables from a matlab hdf5 mat file (v7.3).'''
    import numpy
    import h5py
    if isinstance(f,str):
        f = h5py.File(f,'r')
    if isinstance(var,(list,tuple)):
        return tuple(load_matv73(f,x,cast_string=cast_string,cast_integer=cast_integer) for x in var)
    def parse(val):
        # String data
        if cast_string and val.dtype==numpy.dtype('uint16'):
            return ''.join([chr(x) for x in val.squeeze()])
        # Numeric data
        val = val.transpose()
        if cast_integer:
            val_int = numpy.rint(val)
            if numpy.all(numpy.abs(val-val_int)<1e-8):
                val = val_int.astype('int')
        if squeeze:
            val = val.squeeze()
            if val.shape==():
                val = val.item()
        return val
    t = f[var]
    if isinstance(t,h5py.Dataset):
        val = t[:]
        if val.dtype==numpy.dtype('object'):
            # Cell arrays are stored as references
            for idx,x in numpy.ndenumerate(val):
                val[idx] = parse(f[val[idx]][:])
            return val
        else:
            return parse(val)
    else:
        d = {}
        for key in t:
            d[key] = load_matv73(t,key,
                        squeeze=squeeze,
                        cast_string=cast_string,
                        cast_integer=cast_integer)
        return d

def h5write(file,arr,name='data',truncate=True):
    import h5py
    is_open = isinstance(file,(h5py.File,h5py.Group))
    flag = 'w' if truncate else 'a'
    f = file if is_open else h5py.File(file,flag)
    f.create_dataset(name,data=arr)
    if not is_open: f.close()
    return

def h5read(file,name='data'):
    import h5py
    is_open = isinstance(file,(h5py.File,h5py.Group))
    f = file if is_open else h5py.File(file,'r')
    data = f[name][:]
    if not is_open: f.close()
    return data

def convert_coordinate_to_absolute(coord,bounds):
    assert len(coord) in (3,6)
    if len(coord)==3:
        return tuple(bounds[2*ii]+(bounds[2*ii+1]-bounds[2*ii])*coord[ii] for ii in range(3))
    else:
        return tuple(bounds[2*(ii//2)]+(bounds[2*(ii//2)+1]-bounds[2*(ii//2)])*coord[ii] for ii in range(6))

def rank_from_position(ip,jp,kp,nyp,nzp):
    return ip*nyp*nzp+jp*nzp+kp

def position_from_rank(rank,nyp,nzp):
    ip = rank//(nyp*nzp)
    jp = (rank//nzp)%nyp
    kp = rank%nzp
    return (ip,jp,kp)

def key_grid(key):
    '''Get the key for the grid: scalar fields all share a common grid,
    so this routine returns 's' if key is 'sX' with X being a (potentially
    multi-digit) number. Otherwise the key itself is returned.'''
    if key[0]=='s' and key[1:].isnumeric():
        return 's'
    else:
        return key

def eval_procslice(proc,origin,spacing,key_grid,procslice):
    # Get number of procs
    nxp = len(proc[key_grid][0])
    nyp = len(proc[key_grid][2])
    nzp = len(proc[key_grid][4])
    nxyzp = (nxp,nyp,nzp)
    # Get processor indices for slicing
    if procslice is not None:
        ipb = procslice[0] if procslice[0] is not None else 0
        ipe = procslice[1] if procslice[1] is not None else nxp-1
        jpb = procslice[2] if procslice[2] is not None else 0
        jpe = procslice[3] if procslice[3] is not None else nyp-1
        kpb = procslice[4] if procslice[4] is not None else 0
        kpe = procslice[5] if procslice[5] is not None else nzp-1
    else:
        ipb = 0
        ipe = nxp-1
        jpb = 0
        jpe = nyp-1
        kpb = 0
        kpe = nzp-1
    idxproc = (ipb,ipe,jpb,jpe,kpb,kpe)
    for ii_ in range(6):
        assert idxproc[ii_]>=0, \
            "procslice[{:d}] must be greater than or equal to 0, but is {:d}.".format(
                ii_,idxproc[ii_])
        assert idxproc[ii_]<nxyzp[ii_//2], \
            "procslice[{:d}] must be smaller than {:d}, but is {:d}.".format(
                ii_,nxyzp[ii_//2],idxproc[ii_])
    for ii_ in range(3):
        assert idxproc[2*ii_]<=idxproc[2*ii_+1], \
            "procslice[{:d}] must be smaller than or equal to procslice[{:d}].".format(
                2*ii_,2*ii_+1)
    ib = proc[key_grid][0][ipb]
    ie = proc[key_grid][1][ipe]
    jb = proc[key_grid][2][jpb]
    je = proc[key_grid][3][jpe]
    kb = proc[key_grid][4][kpb]
    ke = proc[key_grid][5][kpe]
    dim = (ie-ib+1,je-jb+1,ke-kb+1)
    ori = (origin[key_grid][0]+(ib-1)*spacing[0],
           origin[key_grid][1]+(jb-1)*spacing[1],
           origin[key_grid][2]+(kb-1)*spacing[2])
    return idxproc,ori,dim