def test_flatten_unflatten(): # first consider the 2D case where we assume a 4 X 5 grid of cells dim = 2 ncells_per_dim = IntArray(3) ncells_per_dim[0] = 4 ncells_per_dim[1] = 5 ncells_per_dim[2] = 0 # valid un-flattened cell indices cids = [[i, j] for i in range(4) for j in range(5)] for _cid in cids: cid = IntPoint(_cid[0], _cid[1], 0) flattened = nnps.py_flatten(cid, ncells_per_dim, dim) unflattened = nnps.py_unflatten(flattened, ncells_per_dim, dim) # the unflattened index should match with cid assert (cid == unflattened) # 3D dim = 3 ncells_per_dim = IntArray(3) ncells_per_dim[0] = 4 ncells_per_dim[1] = 5 ncells_per_dim[2] = 2 # valid un-flattened indices cids = [[i, j, k] for i in range(4) for j in range(5) for k in range(2)] for _cid in cids: cid = IntPoint(_cid[0], _cid[1], _cid[2]) flattened = nnps.py_flatten(cid, ncells_per_dim, dim) unflattened = nnps.py_unflatten(flattened, ncells_per_dim, dim) # the unflattened index should match with cid assert (cid == unflattened)
def test_1D_get_valid_cell_index(): dim = 1 # simulate a dummy distribution such that 10 cells are along the # 'x' direction n_cells = 10 ncells_per_dim = IntArray(3) ncells_per_dim[0] = n_cells ncells_per_dim[1] = 0 ncells_per_dim[2] = 0 # target cell cx = 1 cy = cz = 0 # as long as cy and cz are 0, the function should return the valid # flattened cell index for the cell for i in [-1, 0, 1]: index = nnps.py_get_valid_cell_index(IntPoint(cx + i, cy, cz), ncells_per_dim, dim, n_cells) assert index != -1 # index should be -1 whenever cy and cz are > 1. This is # specifically the case that was failing earlier. for j in [-1, 1]: for k in [-1, 1]: index = nnps.py_get_valid_cell_index(IntPoint(cx, cy + j, cz + k), ncells_per_dim, dim, n_cells) assert index == -1