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]