def test_scalar_2d(halo, n_threads):
# arrange
data = np.array([[1, 6], [2, 7], [3, 8], [4, 9]], dtype=float)
boundary_condition = (Periodic(),
Polar(grid=data.shape,
longitude_idx=OUTER,
latitude_idx=INNER))
field = ScalarField(data, halo, boundary_condition)
traversals = Traversals(grid=data.shape,
halo=halo,
jit_flags=JIT_FLAGS,
n_threads=n_threads)
field.assemble(traversals)
meta_and_data, fill_halos = field.impl
sut = traversals._code['fill_halos_scalar'] # pylint:disable=protected-access
# act
# pylint: disable-next=not-an-iterable
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_apply_scalar(n_threads: int, halo: int, grid: tuple, loop: bool):
if len(grid) == 1 and n_threads > 1:
return
cmn = make_commons(grid, halo, n_threads)
# arrange
sut = cmn.traversals.apply_scalar(loop=loop)
out = ScalarField(np.zeros(grid), halo,
tuple([Constant(np.nan)] * cmn.n_dims))
out.assemble(cmn.traversals)
# 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],
*cmn.vec_null_arg_impl[IMPL_META_AND_DATA],
*cmn.vec_null_arg_impl[IMPL_BC],
*cmn.scl_null_arg_impl[IMPL_META_AND_DATA],
*cmn.scl_null_arg_impl[IMPL_BC],
*cmn.scl_null_arg_impl[IMPL_META_AND_DATA],
*cmn.scl_null_arg_impl[IMPL_BC],
*cmn.scl_null_arg_impl[IMPL_META_AND_DATA],
*cmn.scl_null_arg_impl[IMPL_BC],
*cmn.scl_null_arg_impl[IMPL_META_AND_DATA],
*cmn.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 cmn.n_dims > 1 else (INVALID_INDEX, ):
for j in range(
halo, halo +
grid[MID3D]) if cmn.n_dims > 2 else (INVALID_INDEX, ):
for k in range(halo, halo + grid[INNER]):
if cmn.n_dims == 1:
ijk = (k, INVALID_INDEX, INVALID_INDEX)
elif cmn.n_dims == 2:
ijk = (i, k, INVALID_INDEX)
else:
ijk = (i, j, k)
value = cmn.traversals.indexers[
cmn.n_dims].ats[INNER if cmn.n_dims == 1 else OUTER](
focus, data, *ijk)
assert (cmn.n_dims if loop else 1) * cell_id(i, j,
k) == value
assert cmn.scl_null_arg_impl[IMPL_META_AND_DATA][META_AND_DATA_META][
META_HALO_VALID]
assert cmn.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_1d_scalar(data, n_threads=1, halo=1):
# arrange
boundary_conditions = (Extrapolated(), )
field = ScalarField(data, halo, boundary_conditions)
traversals = Traversals(grid=field.grid,
halo=halo,
jit_flags=JIT_FLAGS,
n_threads=n_threads)
field.assemble(traversals)
meta_and_data, fill_halos = field.impl
sut = traversals._code['fill_halos_scalar'] # pylint:disable=protected-access
# act
thread_id = 0
sut(thread_id, *meta_and_data, *fill_halos)
# assert
print(field.data)
def test_formulae_upwind():
# 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=options.jit_flags,
n_threads=1)
upwind = make_upwind(options=options,
non_unit_g_factor=False,
traversals=traversals)
boundary_conditions = (Periodic(), )
psi = ScalarField(psi_data, halo, boundary_conditions)
psi.assemble(traversals)
psi_impl = psi.impl
flux = VectorField((flux_data, ), halo, boundary_conditions)
flux.assemble(traversals)
flux_impl = flux.impl
# Act
with warnings.catch_warnings():
warnings.simplefilter('ignore',
category=NumbaExperimentalFeatureWarning)
upwind(
traversals.null_impl,
_Impl(field=psi_impl[IMPL_META_AND_DATA], bc=psi_impl[IMPL_BC]),
_Impl(field=flux_impl[IMPL_META_AND_DATA], bc=flux_impl[IMPL_BC]),
_Impl(field=traversals.null_impl.scalar[IMPL_META_AND_DATA],
bc=traversals.null_impl.scalar[IMPL_BC]))
# Assert
np.testing.assert_array_equal(psi.get(), np.roll(psi_data, 1))
def test_scalar(data, halo, side, n_threads, dim):
n_dims = len(data.shape)
if n_dims == 1 and dim != INNER:
return
if n_dims == 2 and dim == MID3D:
return
if n_dims == 1 and n_threads > 1:
return
# arrange
field = ScalarField(data, halo,
tuple(Periodic() for _ in range(n_dims)))
traversals = make_traversals(grid=field.grid,
halo=halo,
n_threads=n_threads)
field.assemble(traversals)
meta_and_data, fill_halos = field.impl
sut = traversals._code['fill_halos_scalar'] # pylint:disable=protected-access
# act
for thread_id in range(
n_threads): # TODO #96: xfail if not all threads executed?
sut(thread_id, *meta_and_data, *fill_halos)
# assert
interior = (halo, -halo)
if side == LEFT:
np.testing.assert_array_equal(
field.data[shift(
indices((None, halo), interior, interior)[:n_dims], dim)],
data[shift(indices((-halo, None), ALL, ALL)[:n_dims], dim)])
else:
np.testing.assert_array_equal(
field.data[shift(
indices((-halo, None), interior, interior)[:n_dims], dim)],
data[shift(indices((None, halo), ALL, ALL)[:n_dims], dim)])