def test_single_timestep(options):
# Arrange
stepper = Stepper(options=options, grid=GRID, non_unit_g_factor=True)
advector = nondivergent_vector_field_2d(GRID, SIZE, TIMESTEP,
stream_function, options.n_halo)
g_factor = ScalarField(RHOD.astype(dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()))
mpdatas = {}
for key, value in VALUES.items():
advectee = ScalarField(np.full(GRID, value, dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()))
mpdatas[key] = Solver(stepper=stepper,
advectee=advectee,
advector=advector,
g_factor=g_factor)
# Act
for mpdata in mpdatas.values():
mpdata.advance(n_steps=1)
# Assert
for value in mpdatas.values():
assert np.isfinite(value.advectee.get()).all()
def test_upwind(shape, ij0, out, courant_number):
value = 44
scalar_field_init = np.zeros(shape)
scalar_field_init[ij0] = value
vector_field_init = (
np.full((shape[0] + 1, shape[1]), courant_number[0]),
np.full((shape[0], shape[1] + 1), courant_number[1])
)
options = Options(n_iters=1)
bcs = (Periodic(), Periodic())
advectee = ScalarField(scalar_field_init, halo=options.n_halo, boundary_conditions=bcs)
advector = VectorField(vector_field_init, halo=options.n_halo, boundary_conditions=bcs)
mpdata = Solver(
stepper=Stepper(options=options, grid=shape, n_threads=1),
advector=advector,
advectee=advectee
)
mpdata.advance(n_steps=1)
np.testing.assert_array_equal(
mpdata.advectee.get(),
out
)
def test_2d_second_dim_contiguous():
grid = (44, 44)
data = np.empty(grid)
boundary_conditions = (Periodic(), Periodic())
sut = ScalarField(data,
halo=1,
boundary_conditions=boundary_conditions)
assert sut.get()[0, :].data.contiguous
def test_timing_2d(benchmark,
options,
grid_static_str,
num_threads,
plot=False):
if grid_static_str == "static":
grid_static = True
elif grid_static_str == "dynamic":
grid_static = False
else:
raise ValueError()
numba.set_num_threads(num_threads)
settings = Settings(n_rotations=6)
_, __, psi = from_pdf_2d(settings.pdf,
xrange=settings.xrange,
yrange=settings.yrange,
gridsize=settings.grid)
advectee = ScalarField(data=psi.astype(dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()))
advector = VectorField(data=(np.full(
(advectee.grid[0] + 1, advectee.grid[1]),
COURANT[0],
dtype=options.dtype),
np.full(
(advectee.grid[0], advectee.grid[1] + 1),
COURANT[1],
dtype=options.dtype)),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()))
if grid_static:
stepper = Stepper(options=options, grid=psi.shape)
else:
stepper = Stepper(options=options, n_dims=2)
solver = Solver(stepper=stepper, advectee=advectee, advector=advector)
def set_z():
solver.advectee.get()[:] = psi
benchmark.pedantic(solver.advance, (settings.n_steps, ),
setup=set_z,
warmup_rounds=1,
rounds=3)
if options.n_iters == 1 or options.nonoscillatory:
np.testing.assert_almost_equal(np.amin(solver.advectee.get()), H_0)
assert np.amax(solver.advectee.get()) < 10 * H
if plot:
pyplot.imshow(solver.advectee.get())
pyplot.colorbar()
pyplot.show()
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_double_pass_donor_cell(n_iters):
courant = .5
options = Options(n_iters=n_iters, DPDC=True, nonoscillatory=True)
state = np.array([0, 1, 0], dtype=options.dtype)
boundary_conditions = (Periodic(),)
mpdata = Solver(
stepper=Stepper(options=options, n_dims=state.ndim, non_unit_g_factor=False),
advectee=ScalarField(
state,
halo=options.n_halo,
boundary_conditions=boundary_conditions
),
advector=VectorField(
(np.full(state.shape[0] + 1, courant, dtype=options.dtype),),
halo=options.n_halo,
boundary_conditions=boundary_conditions
)
)
steps = 1
conserved = np.sum(mpdata.advectee.get())
mpdata.advance(steps)
assert np.sum(mpdata.advectee.get()) == conserved
def test_diffusion_only_2d(data0=np.array([[0, 0, 0], [0, 1., 0], [0, 0, 0]]),
mu_coeff=(.1, .1),
n_steps=1):
# Arrange
options = Options(non_zero_mu_coeff=True)
boundary_conditions = tuple([Periodic()] * 2)
advectee = ScalarField(data0, options.n_halo, boundary_conditions)
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=boundary_conditions)
solver = Solver(stepper=Stepper(options=options, grid=data0.shape),
advector=advector,
advectee=advectee)
# Act
solver.advance(n_steps=n_steps, mu_coeff=mu_coeff)
# 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_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_1d_contiguous():
grid = (44, )
data = np.empty(grid)
boundary_conditions = (Periodic(), )
sut = ScalarField(data,
halo=1,
boundary_conditions=boundary_conditions)
assert sut.get().data.contiguous
def test_upwind_1d():
state = np.array([0, 1, 0])
courant = 1
options = Options(n_iters=1)
mpdata = Solver(
stepper=Stepper(options=options, n_dims=len(state.shape), non_unit_g_factor=False),
advectee=ScalarField(
state.astype(options.dtype),
halo=options.n_halo,
boundary_conditions=(Periodic(),)
),
advector=VectorField(
(np.full(state.shape[0] + 1, courant, dtype=options.dtype),),
halo=options.n_halo,
boundary_conditions=(Periodic(),)
)
)
n_steps = 5
conserved = np.sum(mpdata.advectee.get())
mpdata.advance(n_steps)
assert np.sum(mpdata.advectee.get()) == conserved
def test_shared_advector():
n_x = 100
arr = np.zeros(n_x)
opt1 = Options(n_iters=2, DPDC=True)
opt2 = Options(n_iters=2)
b_c = (Periodic(),)
halo = opt1.n_halo
assert opt2.n_halo == halo
advector = VectorField(data=(np.zeros(n_x + 1),), halo=halo, boundary_conditions=b_c)
_ = Solver(
stepper=Stepper(options=opt1, grid=(n_x,)),
advectee=ScalarField(data=arr, halo=halo, boundary_conditions=b_c),
advector=advector
)
solver = Solver(
stepper=Stepper(options=opt2, grid=(n_x,)),
advectee=ScalarField(data=arr, halo=halo, boundary_conditions=b_c),
advector=advector
)
solver.advance(1)
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)])
# pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
import numpy as np
import pytest
from PyMPDATA import Solver, Stepper, ScalarField, VectorField, Options
from PyMPDATA.boundary_conditions import Periodic
BCS = (Periodic(),)
@pytest.mark.parametrize("case", (
{'g_factor': None, 'non_zero_mu_coeff': True, 'mu': None},
{'g_factor': None, 'non_zero_mu_coeff': True, 'mu': (0,)},
pytest.param({
'g_factor': None,
'non_zero_mu_coeff': False,
'mu': (0,)
}, marks=pytest.mark.xfail(strict=True)),
pytest.param({
'g_factor': ScalarField(np.asarray([1., 1]), Options().n_halo, BCS),
'non_zero_mu_coeff': True,
'mu': None
}, marks=pytest.mark.xfail(strict=True))
))
def test_mu_arg_handling(case):
opt = Options(non_zero_mu_coeff=case['non_zero_mu_coeff'])
advector = VectorField((np.asarray([1., 2, 3]),), opt.n_halo, BCS)
advectee = ScalarField(np.asarray([4., 5]), opt.n_halo, BCS)
stepper = Stepper(options=opt, n_dims=1)
sut = Solver(stepper, advectee, advector, case['g_factor'])
sut.advance(1, mu_coeff=case['mu'])
"nt": case_data[4],
"ni": case_data[5],
"dimsplit": case_data[6],
"input": case_data[7],
"output": case_data[8]
}
# Arrange
data = case["input"].reshape((case["nx"], case["ny"]))
courant = [case["Cx"], case["Cy"]]
options = Options(n_iters=case["ni"], dimensionally_split=case["dimsplit"])
grid = data.shape
advector_data = [
np.full((grid[0] + 1, grid[1]), courant[0], dtype=options.dtype),
np.full((grid[0], grid[1] + 1), courant[1], dtype=options.dtype)
]
bcs = (Periodic(), Periodic())
advector = VectorField(advector_data,
halo=options.n_halo,
boundary_conditions=bcs)
advectee = ScalarField(data=data.astype(dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=bcs)
stepper = Stepper(options=options, grid=grid, non_unit_g_factor=False)
mpdata = Solver(stepper=stepper, advectee=advectee, advector=advector)
sut = mpdata
# Act
sut.advance(n_steps=case["nt"])
# Assert
np.testing.assert_almost_equal(sut.advectee.get(),
def __init__(self,
*,
advectees,
stream_function,
rhod_of_zZ,
dt,
grid,
size,
displacement,
n_iters=2,
infinite_gauge=True,
nonoscillatory=True,
third_order_terms=False):
self.grid = grid
self.size = size
self.dt = dt
self.stream_function = stream_function
self.stream_function_time_dependent = (
"t" in inspect.signature(stream_function).parameters)
self.asynchronous = False
self.thread: (Thread, None) = None
self.displacement = displacement
self.t = 0
options = Options(
n_iters=n_iters,
infinite_gauge=infinite_gauge,
nonoscillatory=nonoscillatory,
third_order_terms=third_order_terms,
)
disable_threads_if_needed = {}
if not conf.JIT_FLAGS["parallel"]:
disable_threads_if_needed["n_threads"] = 1
stepper = Stepper(options=options,
grid=self.grid,
non_unit_g_factor=True,
**disable_threads_if_needed)
advector_impl = VectorField(
(
np.full((grid[0] + 1, grid[1]), np.nan),
np.full((grid[0], grid[1] + 1), np.nan),
),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()),
)
g_factor = make_rhod(self.grid, rhod_of_zZ)
g_factor_impl = ScalarField(
g_factor.astype(dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()),
)
self.g_factor_vec = (
rhod_of_zZ(zZ=x_vec_coord(self.grid)[-1]),
rhod_of_zZ(zZ=z_vec_coord(self.grid)[-1]),
)
self.mpdatas = {}
for k, v in advectees.items():
advectee_impl = ScalarField(
np.asarray(v, dtype=options.dtype),
halo=options.n_halo,
boundary_conditions=(Periodic(), Periodic()),
)
self.mpdatas[k] = Solver(
stepper=stepper,
advectee=advectee_impl,
advector=advector_impl,
g_factor=g_factor_impl,
)
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)])