def test_vector_2d(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))
boundary_conditions = (Periodic(),
Polar(grid=grid,
longitude_idx=OUTER,
latitude_idx=INNER))
field = VectorField(data, halo, boundary_conditions)
traversals = Traversals(grid=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_vector'] # 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)
def test_1d_vector(data, n_threads=1, halo=2):
# arrange
boundary_condition = (Extrapolated(), )
field = VectorField((data, ), halo, boundary_condition)
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_vector'] # 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_vector(data, halo, side, n_threads, comp, dim_offset):
n_dims = len(data)
if n_dims == 1 and n_threads > 1:
return
if n_dims == 1 and (comp != INNER or dim_offset != 0):
return
if n_dims == 2 and (comp == MID3D or dim_offset == 2):
return
# arrange
field = VectorField(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_vector'] # 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 n_dims == 1 and halo == 1:
np.testing.assert_array_equal(field.data[comp], data[comp])
if side == LEFT:
if dim_offset == 1:
np.testing.assert_array_equal(
field.data[comp][shift(
indices((None, halo), (halo - 1, -(halo - 1)),
interior)[:n_dims], -comp + dim_offset)],
data[comp][shift(
indices((-halo, None), ALL, ALL)[:n_dims],
-comp + dim_offset)])
elif dim_offset == 2:
np.testing.assert_array_equal(
field.data[comp][shift(
indices((None, halo), interior,
(halo - 1, -(halo - 1)))[:n_dims],
-comp + dim_offset)], data[comp][shift(
indices((-halo, None), ALL, ALL)[:n_dims],
-comp + dim_offset)])
elif dim_offset == 0:
np.testing.assert_array_equal(
field.data[comp][shift(
indices((None, halo - 1), interior, interior)[:n_dims],
-comp + dim_offset)], data[comp][shift(
indices((-(halo - 1), None), ALL, ALL)[:n_dims],
-comp + dim_offset)])
else:
if dim_offset == 1:
np.testing.assert_array_equal(
field.data[comp][shift(
indices((-halo, None), (halo - 1, -(halo - 1)),
interior)[:n_dims], -comp + dim_offset)],
data[comp][shift(
indices((None, halo), ALL, ALL)[:n_dims],
-comp + dim_offset)])
elif dim_offset == 2:
np.testing.assert_array_equal(
field.data[comp][shift(
indices((-halo, None), interior,
(halo - 1, -(halo - 1)))[:n_dims],
-comp + dim_offset)], data[comp][shift(
indices((None, halo), ALL, ALL)[:n_dims],
-comp + dim_offset)])
elif dim_offset == 0:
np.testing.assert_array_equal(
field.data[comp][shift(
indices((-(halo - 1), None), interior,
interior)[:n_dims], -comp + dim_offset)],
data[comp][shift(
indices((None, halo - 1), ALL, ALL)[:n_dims],
-comp + dim_offset)])
def test_apply_vector(n_threads: int, halo: int, grid: tuple):
if len(grid) == 1 and n_threads > 1:
return
cmn = make_commons(grid, halo, n_threads)
# arrange
sut = cmn.traversals.apply_vector()
data = {
1:
lambda: (np.zeros(grid[0] + 1), ),
2:
lambda: (np.zeros(
(grid[0] + 1, grid[1])), np.zeros((grid[0], grid[1] + 1))),
3:
lambda: (
np.zeros((grid[0] + 1, grid[1], grid[2])),
np.zeros((grid[0], grid[1] + 1, grid[2])),
np.zeros((grid[0], grid[1], grid[2] + 1)),
)
}[cmn.n_dims]()
out = VectorField(data, halo, tuple([Constant(np.nan)] * cmn.n_dims))
out.assemble(cmn.traversals)
# act
sut(*[_cell_id_vector] * MAX_DIM_NUM, *out.impl[IMPL_META_AND_DATA],
*cmn.scl_null_arg_impl[IMPL_META_AND_DATA],
*cmn.scl_null_arg_impl[IMPL_BC],
*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])
# assert
dims = {
1: (INNER, ),
2: (OUTER, INNER),
3: (OUTER, MID3D, INNER)
}[cmn.n_dims]
for dim in dims:
data = out.get_component(dim)
focus = tuple(-cmn.halos[dim][i] for i in range(MAX_DIM_NUM))
for i in range(cmn.halos[dim][OUTER], cmn.halos[dim][OUTER] +
data.shape[OUTER]) if cmn.n_dims > 1 else (
INVALID_INDEX, ):
for j in range(cmn.halos[dim][MID3D], cmn.halos[dim][MID3D] +
data.shape[MID3D]) if cmn.n_dims > 2 else (
INVALID_INDEX, ):
for k in range(cmn.halos[dim][INNER],
cmn.halos[dim][INNER] + data.shape[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 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]