def test_timing_2d(benchmark,
options,
dtype,
grid_static_str,
concurrency_str,
plot=False):
if grid_static_str == "static":
grid_static = True
elif grid_static_str == "dynamic":
grid_static = False
else:
raise ValueError()
if concurrency_str == "serial":
numba.set_num_threads(1)
else:
numba.set_num_threads(numba.config.NUMBA_NUM_THREADS)
setup = Setup(n_rotations=6)
_, __, z = from_pdf_2d(setup.pdf,
xrange=setup.xrange,
yrange=setup.yrange,
gridsize=setup.grid)
C = (-.5, .25)
grid = z.shape
advector_data = [
np.full((grid[0] + 1, grid[1]), C[0], dtype=options.dtype),
np.full((grid[0], grid[1] + 1), C[1], dtype=options.dtype)
]
advector = VectorField(advector_data,
halo=options.n_halo,
boundary_conditions=(PeriodicBoundaryCondition(),
PeriodicBoundaryCondition()))
advectee = ScalarField(data=z.astype(dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=(PeriodicBoundaryCondition(),
PeriodicBoundaryCondition()))
if grid_static:
stepper = Stepper(options=options, grid=grid)
else:
stepper = Stepper(options=options, n_dims=2)
solver = Solver(stepper=stepper, advectee=advectee, advector=advector)
def set_z():
solver.advectee.get()[:] = z
benchmark.pedantic(solver.advance, (setup.nt, ),
setup=set_z,
warmup_rounds=1,
rounds=3)
if options.n_iters == 1 or options.flux_corrected_transport:
np.testing.assert_almost_equal(np.amin(solver.advectee.get()), h0)
assert np.amax(solver.advectee.get()) < 10 * h
if plot:
pyplot.imshow(solver.advectee.get())
pyplot.colorbar()
pyplot.show()
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 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 test_diffusion_only_2d(
data0=np.array([[0, 0, 0], [0, 1., 0], [0, 0, 0]]), mu=(.1, .1), nt=1):
# Arrange
options = Options(non_zero_mu_coeff=True)
bc = [PeriodicBoundaryCondition()] * 2
advectee = ScalarField(data0, options.n_halo, bc)
advector = VectorField(data=(np.zeros(
(data0.shape[0] + 1, data0.shape[1])),
np.zeros(
(data0.shape[0], data0.shape[1] + 1))),
halo=options.n_halo,
boundary_conditions=bc)
solver = Solver(stepper=Stepper(options=options, grid=data0.shape),
advector=advector,
advectee=advectee)
# Act
solver.advance(nt=nt, mu_coeff=mu)
# Assert
data1 = solver.advectee.get()
np.testing.assert_almost_equal(actual=np.sum(data1), desired=np.sum(data0))
assert np.amax(data0) > np.amax(data1)
assert np.amin(data1) >= 0
assert np.count_nonzero(data1) == 5
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_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 test_1d_contiguous():
grid = (44, )
data = np.empty(grid)
bc = (PeriodicBoundaryCondition(), )
sut = ScalarField(data, halo=1, boundary_conditions=bc)
assert sut.get().data.contiguous
def test_2d_second_dim_contiguous():
grid = (44, 44)
data = np.empty(grid)
bc = (PeriodicBoundaryCondition(), PeriodicBoundaryCondition())
sut = ScalarField(data, halo=1, boundary_conditions=bc)
assert sut.get()[0, :].data.contiguous