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_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, *out.impl[0],
*vec_null_arg_impl[0], *vec_null_arg_impl[1],
*scl_null_arg_impl[0], *scl_null_arg_impl[1],
*scl_null_arg_impl[0], *scl_null_arg_impl[1],
*scl_null_arg_impl[0], *scl_null_arg_impl[1])
# assert
data = out.get()
assert data.shape == grid
focus = (-halo, -halo)
for i in range(halo, halo + grid[0]):
for j in (-1, ) if n_dims == 1 else range(halo, halo + grid[1]):
value = indexers[n_dims].at[0](focus, data, i, j)
assert value == (n_dims if loop else 1) * cell_id(i, j)
assert scl_null_arg_impl[0][0][meta_halo_valid]
assert vec_null_arg_impl[0][0][meta_halo_valid]
assert not out.impl[0][0][meta_halo_valid]
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)))
else:
raise NotImplementedError()
out = VectorField(data, halo,
[ConstantBoundaryCondition(np.nan)] * n_dims)
# act
sut(_cell_id_vector, _cell_id_vector, *out.impl[0],
*scl_null_arg_impl[0], *scl_null_arg_impl[1],
*vec_null_arg_impl[0], *vec_null_arg_impl[1],
*scl_null_arg_impl[0], *scl_null_arg_impl[1])
# assert
halos = ((halo - 1, halo), (halo, halo - 1))
for d in range(n_dims):
data = out.get_component(d)
focus = tuple(-halos[d][i] for i in range(n_dims))
for i in range(halos[d][0], halos[d][0] + data.shape[0]):
for j in (-1, ) if n_dims == 1 else range(
halos[d][1], halos[d][1] + data.shape[1]):
value = indexers[n_dims].at[0](focus, data, i, j)
assert value == cell_id(i, j)
assert scl_null_arg_impl[0][0][meta_halo_valid]
assert vec_null_arg_impl[0][0][meta_halo_valid]
assert not out.impl[0][0][meta_halo_valid]