def test_apply_scalar(n_threads, halo, grid, loop): n_dims = len(grid) if n_dims == 1 and n_threads > 1: return # arrange traversals = Traversals(grid, halo, jit_flags, n_threads) sut = traversals.apply_scalar(loop=loop) scl_null_arg_impl = ScalarField.make_null(n_dims).impl vec_null_arg_impl = VectorField.make_null(n_dims).impl out = ScalarField(np.zeros(grid), halo, [ConstantBoundaryCondition(np.nan)] * n_dims) # act sut(_cell_id_scalar, _cell_id_scalar if loop else None, _cell_id_scalar if loop else None, *out.impl[IMPL_META_AND_DATA], *vec_null_arg_impl[IMPL_META_AND_DATA], *vec_null_arg_impl[IMPL_BC], *scl_null_arg_impl[IMPL_META_AND_DATA], *scl_null_arg_impl[IMPL_BC], *scl_null_arg_impl[IMPL_META_AND_DATA], *scl_null_arg_impl[IMPL_BC], *scl_null_arg_impl[IMPL_META_AND_DATA], *scl_null_arg_impl[IMPL_BC]) # assert data = out.get() assert data.shape == grid focus = (-halo, -halo, -halo) for i in range(halo, halo + grid[OUTER]) if n_dims > 1 else (INVALID_INDEX, ): for j in range(halo, halo + grid[MID3D]) if n_dims > 2 else (INVALID_INDEX, ): for k in range(halo, halo + grid[INNER]): if n_dims == 1: ijk = (k, INVALID_INDEX, INVALID_INDEX) elif n_dims == 2: ijk = (i, k, INVALID_INDEX) else: raise NotImplementedError() value = indexers[n_dims].at[INNER if n_dims == 1 else OUTER](focus, data, *ijk) assert (n_dims if loop else 1) * cell_id(i, j, k) == value assert scl_null_arg_impl[IMPL_META_AND_DATA][META_AND_DATA_META][ META_HALO_VALID] assert vec_null_arg_impl[IMPL_META_AND_DATA][META_AND_DATA_META][ META_HALO_VALID] assert not out.impl[IMPL_META_AND_DATA][META_AND_DATA_META][ META_HALO_VALID]
def test_vector_1d(self, data, halo, side): # arrange field = VectorField((data, ), halo, (PeriodicBoundaryCondition(), )) meta_and_data, fill_halos = field.impl traversals = Traversals(grid=(data.shape[0] - 1, ), halo=halo, jit_flags={}, n_threads=1) sut = traversals._fill_halos_vector thread_id = 0 # act sut(thread_id, *meta_and_data, *fill_halos) # assert if halo == 1: return if side == LEFT: np.testing.assert_array_equal(field.data[0][:(halo - 1)], data[-(halo - 1):]) elif side == RIGHT: np.testing.assert_array_equal(field.data[0][-(halo - 1):], data[:(halo - 1)]) else: raise ValueError()
def test_vector_2d(self, halo, n_threads): # arrange grid = (4, 2) data = (np.array([ [1, 6], [2, 7], [3, 8], [4, 9], [5, 10], ], dtype=float), np.array([ [1, 5, 9], [2, 6, 10], [3, 7, 11], [4, 8, 12], ], dtype=float)) bc = (PeriodicBoundaryCondition(), PolarBoundaryCondition(grid=grid, longitude_idx=OUTER, latitude_idx=INNER)) field = VectorField(data, halo, bc) meta_and_data, fill_halos = field.impl traversals = Traversals(grid=grid, halo=halo, jit_flags={}, n_threads=n_threads) sut = traversals._fill_halos_vector # act for thread_id in numba.prange(n_threads): sut(thread_id, *meta_and_data, *fill_halos)
def test_scalar_2d(self, halo, n_threads): # arrange data = np.array([[1, 6], [2, 7], [3, 8], [4, 9]], dtype=float) bc = (PeriodicBoundaryCondition(), PolarBoundaryCondition(grid=data.shape, longitude_idx=OUTER, latitude_idx=INNER)) field = ScalarField(data, halo, bc) meta_and_data, fill_halos = field.impl traversals = Traversals(grid=data.shape, halo=halo, jit_flags={}, n_threads=n_threads) sut = traversals._fill_halos_scalar # act for thread_id in numba.prange(n_threads): sut(thread_id, *meta_and_data, *fill_halos) # assert np.testing.assert_array_equal( field.data[halo:-halo, :halo], np.roll(field.get()[:, :halo], data.shape[OUTER] // 2, axis=OUTER)) np.testing.assert_array_equal( field.data[halo:-halo, -halo:], np.roll(field.get()[:, -halo:], data.shape[OUTER] // 2, axis=OUTER))
def test_make_upwind(self): # Arrange psi_data = np.array((0, 1, 0)) flux_data = np.array((0, 0, 1, 0)) options = Options() halo = options.n_halo traversals = Traversals(grid=psi_data.shape, halo=halo, jit_flags={}, n_threads=1) upwind = make_upwind(options=options, non_unit_g_factor=False, traversals=traversals) bc = [PeriodicBoundaryCondition()] psi = ScalarField(psi_data, halo, bc) psi_impl = psi.impl flux_impl = VectorField((flux_data, ), halo, bc).impl null_impl = ScalarField.make_null(len(psi_data.shape)).impl # Act upwind(psi_impl[0], *flux_impl, *null_impl) # Assert np.testing.assert_array_equal(psi.get(), np.roll(psi_data, 1))
def test_scalar_1d(self, data, halo, side): # arrange field = ScalarField(data, halo, (PeriodicBoundaryCondition(), )) meta_and_data, fill_halos = field.impl traversals = Traversals(grid=data.shape, halo=halo, jit_flags={}, n_threads=1) sut = traversals._fill_halos_scalar thread_id = 0 # act sut(thread_id, *meta_and_data, *fill_halos) # assert if side == LEFT: np.testing.assert_array_equal(field.data[:halo], data[-halo:]) elif side == RIGHT: np.testing.assert_array_equal(field.data[-halo:], data[:halo]) else: raise ValueError()
def test_apply_vector(n_threads, halo, grid): n_dims = len(grid) if n_dims == 1 and n_threads > 1: return # arrange traversals = Traversals(grid, halo, jit_flags, n_threads) sut = traversals.apply_vector() scl_null_arg_impl = ScalarField.make_null(n_dims).impl vec_null_arg_impl = VectorField.make_null(n_dims).impl if n_dims == 1: data = (np.zeros(grid[0] + 1), ) elif n_dims == 2: data = (np.zeros( (grid[0] + 1, grid[1])), np.zeros((grid[0], grid[1] + 1))) elif n_dims == 3: pass # TODO else: raise NotImplementedError() out = VectorField(data, halo, [ConstantBoundaryCondition(np.nan)] * n_dims) # act sut(*[_cell_id_vector] * MAX_DIM_NUM, *out.impl[IMPL_META_AND_DATA], *scl_null_arg_impl[IMPL_META_AND_DATA], *scl_null_arg_impl[IMPL_BC], *vec_null_arg_impl[IMPL_META_AND_DATA], *vec_null_arg_impl[IMPL_BC], *scl_null_arg_impl[IMPL_META_AND_DATA], *scl_null_arg_impl[IMPL_BC]) # assert halos = ((halo - 1, halo, halo), (halo, halo - 1, halo), (halo, halo, halo - 1)) if n_dims == 1: dims = (INNER, ) elif n_dims == 2: dims = (OUTER, INNER) else: raise NotImplementedError() for d in dims: print("DIM", d) data = out.get_component(d) focus = tuple(-halos[d][i] for i in range(MAX_DIM_NUM)) print("focus", focus) for i in range( halos[d][OUTER], halos[d][OUTER] + data.shape[OUTER]) if n_dims > 1 else (INVALID_INDEX, ): for j in range(halos[d][MID3D], halos[d][MID3D] + data.shape[MID3D]) if n_dims > 2 else ( INVALID_INDEX, ): for k in range(halos[d][INNER], halos[d][INNER] + data.shape[INNER]): if n_dims == 1: ijk = (k, INVALID_INDEX, INVALID_INDEX) elif n_dims == 2: ijk = (i, k, INVALID_INDEX) else: raise NotImplementedError() print("check at", i, j, k) value = indexers[n_dims].at[INNER if n_dims == 1 else OUTER](focus, data, *ijk) assert cell_id(i, j, k) == value assert scl_null_arg_impl[IMPL_META_AND_DATA][META_AND_DATA_META][ META_HALO_VALID] assert vec_null_arg_impl[IMPL_META_AND_DATA][META_AND_DATA_META][ META_HALO_VALID] assert not out.impl[IMPL_META_AND_DATA][META_AND_DATA_META][ META_HALO_VALID]