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_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_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 make_laplacian(non_unit_g_factor: bool, options: Options,
traversals: Traversals):
if not options.non_zero_mu_coeff:
@numba.njit(**options.jit_flags)
def apply(_1, _2, _3, _4, _5):
return
else:
idx = indexers[traversals.n_dims]
apply_vector = traversals.apply_vector()
formulae_laplacian = tuple([
__make_laplacian(options.jit_flags, idx.at[i], options.epsilon,
non_unit_g_factor)
if i < traversals.n_dims else None for i in range(MAX_DIM_NUM)
])
@numba.njit(**options.jit_flags)
def apply(advector, advectee, advectee_bc, null_vecfield_bc):
null_vecfield = advector
null_scalarfield = advectee
return apply_vector(*formulae_laplacian, *advector, *advectee,
*advectee_bc, *null_vecfield,
*null_vecfield_bc, *null_scalarfield,
*null_vecfield_bc)
return apply
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],
]), np.array([
[1, 5, 9],
[2, 6, 10],
[3, 7, 11],
[4, 8, 12],
]))
bc = (PeriodicBoundaryCondition(), PolarBoundaryCondition(grid, 0, 1))
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)
# assert
print(field.data[0].shape)
print(field.data[0])
print(field.data[1].shape)
print(field.data[1])
def make_laplacian(non_unit_g_factor: bool, options: Options,
traversals: Traversals):
if not options.non_zero_mu_coeff:
@numba.njit(**options.jit_flags)
def apply(_flux, _psi, _psi_bc, _GC_corr, _vec_bc):
return
else:
idx = indexers[traversals.n_dims]
apply_vector = traversals.apply_vector()
formulae_laplacian = tuple([
__make_laplacian(options.jit_flags, idx.at[i], options.epsilon,
non_unit_g_factor, traversals.n_dims)
if i < traversals.n_dims else null_vector_formula
for i in range(MAX_DIM_NUM)
])
@numba.njit(**options.jit_flags)
def apply(GC_phys, psi, psi_bc, null_vecfield_bc):
null_vecfield = GC_phys
null_scalarfield = psi
return apply_vector(*formulae_laplacian, *GC_phys, *psi, *psi_bc,
*null_vecfield, *null_vecfield_bc,
*null_scalarfield, *null_vecfield_bc)
return apply
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={})
sut, _ = traversals.make_boundary_conditions()
# act
sut(*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)))
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]
def test_scalar_2d(self, halo, n_threads):
# arrange
data = np.array([[1, 6], [2, 7], [3, 8], [4, 9]])
bc = (PeriodicBoundaryCondition(),
PolarBoundaryCondition(data.shape, 0, 1))
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[0] // 2, axis=0))
np.testing.assert_array_equal(
field.data[halo:-halo, -halo:],
np.roll(field.get()[:, -halo:], data.shape[0] // 2, axis=0))
def make_step_impl(options, non_unit_g_factor, grid, n_threads):
n_iters = options.n_iters
n_dims = len(grid)
halo = options.n_halo
non_zero_mu_coeff = options.non_zero_mu_coeff
flux_corrected_transport = options.flux_corrected_transport
traversals = Traversals(grid, halo, options.jit_flags, n_threads=n_threads)
upwind = make_upwind(options, non_unit_g_factor, traversals)
flux_first_pass = make_flux_first_pass(options, traversals)
flux_subsequent = make_flux_subsequent(options, traversals)
antidiff = make_antidiff(non_unit_g_factor, options, traversals)
laplacian = make_laplacian(non_unit_g_factor, options, traversals)
fct_psi_min = make_psi_extremum(min, options, traversals)
fct_psi_max = make_psi_extremum(max, options, traversals)
fct_beta_down = make_beta(min, non_unit_g_factor, options, traversals)
fct_beta_up = make_beta(max, non_unit_g_factor, options, traversals)
fct_correction = make_correction(options, traversals)
@numba.njit(**options.jit_flags)
def axpy(out_meta, out0, out1, a, x_meta, x0, x1, y_meta, y0, y1):
out0[:] = a[0] * x0[:] + y0[:]
if n_dims > 1:
out1[:] = a[1] * x1[:] + y1[:]
out_meta[meta_halo_valid] = False
@numba.njit(**options.jit_flags)
def step(nt, mu_coeff, advectee, advectee_bc, advector, advector_bc,
g_factor, g_factor_bc, vectmp_a, vectmp_a_bc, vectmp_b,
vectmp_b_bc, vectmp_c, vectmp_c_bc, psi_min, psi_min_bc, psi_max,
psi_max_bc, beta_up, beta_up_bc, beta_down, beta_down_bc):
vec_bc = advector_bc
null_vecfield = advector
null_vecfield_bc = vec_bc
time = clock()
for _ in range(nt):
if non_zero_mu_coeff:
advector_orig = advector
advector = vectmp_c
for it in range(n_iters):
if it == 0:
if flux_corrected_transport:
fct_psi_min(psi_min, advectee, advectee_bc,
null_vecfield, null_vecfield_bc)
fct_psi_max(psi_max, advectee, advectee_bc,
null_vecfield, null_vecfield_bc)
if non_zero_mu_coeff:
laplacian(advector, advectee, advectee_bc,
null_vecfield_bc)
axpy(*advector, mu_coeff, *advector, *advector_orig)
flux_first_pass(vectmp_a, advector, advectee, advectee_bc,
vec_bc)
flux = vectmp_a
else:
if it == 1:
advector_oscilatory = advector
advector_nonoscilatory = vectmp_a
flux = vectmp_b
elif it % 2 == 0:
advector_oscilatory = vectmp_a
advector_nonoscilatory = vectmp_b
flux = vectmp_a
else:
advector_oscilatory = vectmp_b
advector_nonoscilatory = vectmp_a
flux = vectmp_b
antidiff(advector_nonoscilatory, advectee, advectee_bc,
advector_oscilatory, vec_bc, g_factor,
g_factor_bc)
flux_subsequent(flux, advectee, advectee_bc,
advector_nonoscilatory, vec_bc)
if flux_corrected_transport:
fct_beta_down(beta_down, flux, vec_bc, advectee,
advectee_bc, psi_min, psi_min_bc,
g_factor, g_factor_bc)
fct_beta_up(beta_up, flux, vec_bc, advectee,
advectee_bc, psi_max, psi_max_bc, g_factor,
g_factor_bc)
fct_correction(advector_nonoscilatory, vec_bc,
beta_down, beta_down_bc, beta_up,
beta_up_bc)
flux_subsequent(flux, advectee, advectee_bc,
advector_nonoscilatory, vec_bc)
upwind(advectee, flux, vec_bc, g_factor, g_factor_bc)
if non_zero_mu_coeff:
advector = advector_orig
return (clock() - time) / nt
return step
def make_step_impl(options, non_unit_g_factor, grid):
n_iters = options.n_iters
n_dims = len(grid)
halo = options.n_halo
non_zero_mu_coeff = options.non_zero_mu_coeff
flux_corrected_transport = options.flux_corrected_transport
traversals = Traversals(grid, halo, options.jit_flags)
upwind = make_upwind(options, non_unit_g_factor, traversals)
flux_first_pass = make_flux_first_pass(options, traversals)
flux_subsequent = make_flux_subsequent(options, traversals)
antidiff = make_antidiff(non_unit_g_factor, options, traversals)
laplacian = make_laplacian(non_unit_g_factor, options, traversals)
fct_psi_min = make_psi_extremum(min, options, traversals)
fct_psi_max = make_psi_extremum(max, options, traversals)
fct_beta_down = make_beta(min, non_unit_g_factor, options, traversals)
fct_beta_up = make_beta(max, non_unit_g_factor, options, traversals)
fct_correction = make_correction(options, traversals)
@numba.njit(**options.jit_flags)
def axpy(out_meta, out0, out1, a, x_meta, x0, x1, y_meta, y0, y1):
out0[:] = a * x0[:] + y0[:]
if n_dims > 1:
out1[:] = a * x1[:] + y1[:]
out_meta[meta_halo_valid] = False
@numba.njit(**options.jit_flags)
def step(nt, mu_coeff,
psi, psi_bc,
GC_phys, GC_phys_bc,
g_factor, g_factor_bc,
vectmp_a, vectmp_a_bc,
vectmp_b, vectmp_b_bc,
vectmp_c, vectmp_c_bc,
psi_min, psi_min_bc,
psi_max, psi_max_bc,
beta_up, beta_up_bc,
beta_down, beta_down_bc
):
vec_bc = GC_phys_bc
null_vecfield = GC_phys
null_vecfield_bc = vec_bc
for _ in range(nt):
if non_zero_mu_coeff:
GC_orig = GC_phys
GC_phys = vectmp_c
for it in range(n_iters):
if it == 0:
if flux_corrected_transport:
fct_psi_min(psi_min, psi, psi_bc, null_vecfield, null_vecfield_bc)
fct_psi_max(psi_max, psi, psi_bc, null_vecfield, null_vecfield_bc)
if non_zero_mu_coeff:
laplacian(GC_phys, psi, psi_bc, null_vecfield_bc)
axpy(*GC_phys, mu_coeff, *GC_phys, *GC_orig)
flux_first_pass(vectmp_a, GC_phys, psi, psi_bc, vec_bc)
flux = vectmp_a
else:
if it == 1:
GC_unco = GC_phys
GC_corr = vectmp_a
flux = vectmp_b
elif it % 2 == 0:
GC_unco = vectmp_a
GC_corr = vectmp_b
flux = vectmp_a
else:
GC_unco = vectmp_b
GC_corr = vectmp_a
flux = vectmp_b
antidiff(GC_corr, psi, psi_bc, GC_unco, vec_bc, g_factor, g_factor_bc)
flux_subsequent(flux, psi, psi_bc, GC_corr, vec_bc)
if flux_corrected_transport:
fct_beta_down(beta_down, flux, vec_bc, psi, psi_bc, psi_min, psi_min_bc, g_factor, g_factor_bc)
fct_beta_up(beta_up, flux, vec_bc, psi, psi_bc, psi_max, psi_max_bc, g_factor, g_factor_bc)
fct_correction(GC_corr, vec_bc, beta_down, beta_down_bc, beta_up, beta_up_bc)
flux_subsequent(flux, psi, psi_bc, GC_corr, vec_bc)
upwind(psi, flux, vec_bc, g_factor, g_factor_bc)
if non_zero_mu_coeff:
GC_phys = GC_orig
return step