decompose.py 文件源码

def split_array(gle, gre, shape, psize):
    """ Split array into px*py*pz subarrays. """
    n_d = np.array(shape, dtype=np.int64)
    dds = (gre-gle)/shape
    left_edges = []
    right_edges = []
    shapes = []
    slices = []
    for i in range(psize[0]):
        for j in range(psize[1]):
            for k in range(psize[2]):
                piece = np.array((i, j, k), dtype=np.int64)
                lei = n_d * piece // psize
                rei = n_d * (piece + np.ones(3, dtype=np.int64)) // psize
                lle = gle + lei*dds
                lre = gle + rei*dds
                left_edges.append(lle)
                right_edges.append(lre)
                shapes.append(rei-lei)
                slices.append(np.s_[lei[0]:rei[0], lei[1]:
                                    rei[1], lei[2]:rei[2]])

    return left_edges, right_edges, shapes, slices