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_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_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_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_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 __init__(self, nz, dt, advector_of_t, advectee_of_zZ_at_t0,
g_factor_of_zZ, mpdata_settings):
self.t = 0
self.dt = dt
self.advector_of_t = advector_of_t
grid = (nz, )
options = Options(n_iters=mpdata_settings['n_iters'],
infinite_gauge=mpdata_settings['iga'],
flux_corrected_transport=mpdata_settings['fct'],
third_order_terms=mpdata_settings['tot'])
stepper = Stepper(options=options, grid=grid, non_unit_g_factor=True)
bcs = (ExtrapolatedBoundaryCondition(), )
g_factor = ScalarField(data=g_factor_of_zZ(
arakawa_c.z_scalar_coord(grid)),
halo=options.n_halo,
boundary_conditions=bcs)
advector = VectorField(data=(np.full(nz + 1, advector_of_t(0)), ),
halo=options.n_halo,
boundary_conditions=bcs)
self.advectee = ScalarField(data=advectee_of_zZ_at_t0(
arakawa_c.z_scalar_coord(grid)),
halo=options.n_halo,
boundary_conditions=bcs)
self.solver = Solver(stepper=stepper,
advectee=self.advectee,
advector=advector,
g_factor=g_factor)
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_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 __init__(self, *, fields,
n_iters=2, infinite_gauge=True,
flux_corrected_transport=True, third_order_terms=False):
self.grid = fields.g_factor.shape
self.asynchronous = False
self.thread: (Thread, None) = None
options = Options(
n_iters=n_iters,
infinite_gauge=infinite_gauge,
flux_corrected_transport=flux_corrected_transport,
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)
# CFL condition
for d in range(len(fields.advector)):
np.testing.assert_array_less(np.abs(fields.advector[d]), 1)
self.advector = fields.advector
advector_impl = VectorField(fields.advector, halo=options.n_halo,
boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition()))
self.g_factor = fields.g_factor
g_factor_impl = ScalarField(fields.g_factor.astype(dtype=options.dtype), halo=options.n_halo,
boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition()))
self.mpdatas = {}
for k, v in fields.advectees.items():
advectee = ScalarField(np.full(self.grid, v, dtype=options.dtype), halo=options.n_halo,
boundary_conditions=(PeriodicBoundaryCondition(), PeriodicBoundaryCondition()))
self.mpdatas[k] = Solver(stepper=stepper, advectee=advectee, advector=advector_impl, g_factor=g_factor_impl)
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'])
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_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 make_condensational_growth_solver(nr, r_min, r_max, GC_max,
grid_layout, psi_coord, pdf_of_r,
drdt_of_r, opts: Options):
# psi = psi(p)
dp_dr = psi_coord.dx_dr
dx_dr = grid_layout.dx_dr
xh, dx = np.linspace(grid_layout.x(r_min),
grid_layout.x(r_max),
nr + 1,
retstep=True)
rh = grid_layout.r(xh)
x = np.linspace(xh[0] + dx / 2, xh[-1] - dx / 2, nr)
r = grid_layout.r(x)
def pdf_of_r_over_psi(r):
return pdf_of_r(r) / psi_coord.dx_dr(r)
psi = discretised_analytical_solution(rh,
pdf_of_r_over_psi,
midpoint_value=True,
r=r)
dp_dt = drdt_of_r(rh) * dp_dr(rh)
G = dp_dr(r) / dx_dr(r)
# C = dr_dt * dt / dr
# GC = dp_dr / dx_dr * dr_dt * dt / dr =
# \ \_____ / _..____/
# \_____.._____/ \_ dt/dx
# |
# dp_dt
dt = GC_max * dx / np.amax(dp_dt)
GCh = dp_dt * dt / dx
# CFL condition
np.testing.assert_array_less(np.abs(GCh), 1)
g_factor = ScalarField(
G.astype(dtype=opts.dtype),
halo=opts.n_halo,
boundary_conditions=(ExtrapolatedBoundaryCondition(), ))
state = ScalarField(
psi.astype(dtype=opts.dtype),
halo=opts.n_halo,
boundary_conditions=(ConstantBoundaryCondition(0), ))
GC_field = VectorField(
[GCh.astype(dtype=opts.dtype)],
halo=opts.n_halo,
boundary_conditions=(ConstantBoundaryCondition(0), ))
stepper = Stepper(options=opts, n_dims=1, non_unit_g_factor=True)
return (Solver(stepper=stepper,
g_factor=g_factor,
advectee=state,
advector=GC_field), r, rh, dx, dt, g_factor)
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_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 make_commons(grid, halo, num_threads):
traversals = make_traversals(grid, halo, num_threads)
n_dims = len(grid)
halos = ((halo - 1, halo, halo), (halo, halo - 1, halo), (halo, halo,
halo - 1))
_Commons = namedtuple('_Commons',
('n_dims', 'traversals', 'scl_null_arg_impl',
'vec_null_arg_impl', 'halos'))
return _Commons(
n_dims=n_dims,
traversals=traversals,
scl_null_arg_impl=ScalarField.make_null(n_dims, traversals).impl,
vec_null_arg_impl=VectorField.make_null(n_dims, traversals).impl,
halos=halos)
def __init__(
self,
nz,
dt,
advector_of_t,
advectee_of_zZ_at_t0,
g_factor_of_zZ,
mpdata_settings,
):
self.__t = 0
self.dt = dt
self.advector_of_t = advector_of_t
grid = (nz, )
options = Options(
n_iters=mpdata_settings["n_iters"],
infinite_gauge=mpdata_settings["iga"],
nonoscillatory=mpdata_settings["fct"],
third_order_terms=mpdata_settings["tot"],
)
stepper = Stepper(options=options, grid=grid, non_unit_g_factor=True)
bcs = (Extrapolated(), )
zZ_scalar = arakawa_c.z_scalar_coord(grid) / nz
g_factor = ScalarField(
data=g_factor_of_zZ(zZ_scalar),
halo=options.n_halo,
boundary_conditions=bcs,
)
advector = VectorField(
data=(np.full(nz + 1, advector_of_t(0)), ),
halo=options.n_halo,
boundary_conditions=bcs,
)
self.advectee = ScalarField(
data=advectee_of_zZ_at_t0(zZ_scalar),
halo=options.n_halo,
boundary_conditions=bcs,
)
self.solver = Solver(
stepper=stepper,
advectee=self.advectee,
advector=advector,
g_factor=g_factor,
)
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_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_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)])
"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(),
case["output"].reshape(
case["nx"], case["ny"]),
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]
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]
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,
)