def test_subsampling(self):
grid = Grid(shape=(40,))
x = grid.dimensions[0]
t = grid.stepping_dim
time = grid.time_dim
nt = 9
f = TimeFunction(name='f', grid=grid)
f.data_with_halo[:] = 1.
# Setup subsampled function
factor = 4
nsamples = (nt+factor-1)//factor
times = ConditionalDimension('t_sub', parent=time, factor=factor)
fsave = TimeFunction(name='fsave', grid=grid, save=nsamples, time_dim=times)
eqns = [Eq(f.forward, f[t, x-1] + f[t, x+1]), Eq(fsave, f)]
op = Operator(eqns)
op.apply(time=nt-1)
assert np.all(f.data_ro_domain[0] == fsave.data_ro_domain[nsamples-1])
glb_pos_map = f.grid.distributor.glb_pos_map
if LEFT in glb_pos_map[x]:
assert np.all(fsave.data_ro_domain[nsamples-1, nt-1:] == 256.)
else:
assert np.all(fsave.data_ro_domain[nsamples-1, :-(nt-1)] == 256.)
# Also check there are no redundant halo exchanges
calls = FindNodes(Call).visit(op)
assert len(calls) == 1
# In particular, there is no need for a halo exchange within the conditional
conditional = FindNodes(Conditional).visit(op)
assert len(conditional) == 1
assert len(FindNodes(Call).visit(conditional[0])) == 0
def test_save_w_subdims(self):
nt = 10
grid = Grid(shape=(10, 10))
x, y = grid.dimensions
time_dim = grid.time_dim
xi = SubDimension.middle(name='xi', parent=x, thickness_left=3, thickness_right=3)
yi = SubDimension.middle(name='yi', parent=y, thickness_left=3, thickness_right=3)
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
eqns = [Eq(u.forward, u + 1),
Eq(usave, u.forward)]
eqns = [e.xreplace({x: xi, y: yi}) for e in eqns]
op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
op.apply(time_M=nt-1)
for i in range(usave.save):
assert np.all(usave.data[i, 3:-3, 3:-3] == 2*i + 1)
assert np.all(usave.data[i, :3, :] == 0)
assert np.all(usave.data[i, -3:, :] == 0)
assert np.all(usave.data[i, :, :3] == 0)
assert np.all(usave.data[i, :, -3:] == 0)
def test_spacial_subsampling(self):
"""
Test conditional dimension for the spatial ones.
This test saves u every two grid points :
u2[x, y] = u[2*x, 2*y]
"""
nt = 19
grid = Grid(shape=(12, 12))
time = grid.time_dim
u = TimeFunction(name='u', grid=grid, save=nt)
assert (grid.time_dim in u.indices)
# Creates subsampled spatial dimensions and accordine grid
dims = tuple([
ConditionalDimension(d.name + 'sub', parent=d, factor=2)
for d in u.grid.dimensions
])
grid2 = Grid((6, 6), dimensions=dims, time_dimension=time)
u2 = TimeFunction(name='u2', grid=grid2, save=nt)
assert (time in u2.indices)
eqns = [Eq(u.forward, u + 1.), Eq(u2, u)]
op = Operator(eqns)
op.apply(time_M=nt - 2)
# Verify that u2[x,y]= u[2*x, 2*y]
assert np.allclose(u.data[:-1, 0:-1:2, 0:-1:2], u2.data[:-1, :, :])
def test_shifted(self):
nt = 19
grid = Grid(shape=(11, 11))
time = grid.time_dim
u = TimeFunction(name='u', grid=grid)
assert(grid.stepping_dim in u.indices)
u2 = TimeFunction(name='u2', grid=grid, save=nt)
assert(time in u2.indices)
factor = 4
time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
usave = TimeFunction(name='usave', grid=grid, save=2, time_dim=time_subsampled)
assert(time_subsampled in usave.indices)
t_sub_shift = Constant(name='t_sub_shift', dtype=np.int32)
eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.),
Eq(usave.subs(time_subsampled, time_subsampled - t_sub_shift), u)]
op = Operator(eqns)
# Starting at time_m=10, so time_subsampled - t_sub_shift is in range
op.apply(time_m=10, time_M=nt-2, t_sub_shift=3)
assert np.all(np.allclose(u.data[0], 8))
assert np.all([np.allclose(u2.data[i], i - 10) for i in range(10, nt)])
assert np.all([np.allclose(usave.data[i], 2+i*factor) for i in range(2)])
def test_no_fusion_convoluted(self):
"""
Conceptually like `test_no_fusion_simple`, but with more expressions
and non-trivial data flow.
"""
grid = Grid(shape=(4, 4, 4))
time = grid.time_dim
f = TimeFunction(name='f', grid=grid)
g = Function(name='g', grid=grid)
h = Function(name='h', grid=grid)
ctime = ConditionalDimension(name='ctime', parent=time, condition=time > 4)
eqns = [Eq(f.forward, f + 1),
Eq(h, f + 1),
Eq(g, f + 1, implicit_dims=[ctime]),
Eq(f.forward, f + 1, implicit_dims=[ctime]),
Eq(f.forward, f + 1),
Eq(g, f + 1)]
op = Operator(eqns)
exprs = FindNodes(Expression).visit(op._func_table['bf0'].root)
assert len(exprs) == 3
assert exprs[1].expr.rhs is exprs[0].output
assert exprs[2].expr.rhs is exprs[0].output
exprs = FindNodes(Expression).visit(op._func_table['bf1'].root)
assert len(exprs) == 2
exprs = FindNodes(Expression).visit(op._func_table['bf2'].root)
assert len(exprs) == 3
assert exprs[1].expr.rhs is exprs[0].output
assert exprs[2].expr.rhs is exprs[0].output
def test_xcor_from_saved(self, opt, gpu_fit):
nt = 10
grid = Grid(shape=(300, 300, 300))
time_dim = grid.time_dim
period = 2
factor = Constant(name='factor', value=period, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
g = Function(name='g', grid=grid)
v = TimeFunction(name='v', grid=grid)
usave = TimeFunction(name='usave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
# For the given `nt` and grid shape, `usave` is roughly 4*5*300**3=~ .5GB of data
for i in range(int(nt//period)):
usave.data[i, :] = i
v.data[:] = i*2 + 1
# Assuming nt//period=5, we are computing, over 5 iterations:
# g = 4*4 [time=8] + 3*3 [time=6] + 2*2 [time=4] + 1*1 [time=2]
op = Operator([Eq(v.backward, v - 1), Inc(g, usave*(v/2))],
opt=(opt, {'gpu-fit': usave if gpu_fit else None}))
op.apply(time_M=nt-1)
assert np.all(g.data == 30)
def test_conddim_backwards():
nt = 10
grid = Grid(shape=(4, 4))
time_dim = grid.time_dim
x, y = grid.dimensions
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
u = TimeFunction(name='u', grid=grid, time_order=0, save=nt, time_dim=time_sub)
v = TimeFunction(name='v', grid=grid)
v1 = TimeFunction(name='v', grid=grid)
for i in range(u.save):
u.data[i, :] = i
eqns = [Eq(v.backward, v.backward + v + u + 1.)]
op0 = Operator(eqns, opt='noop')
op1 = Operator(eqns, opt='buffering')
# Check generated code
assert len(retrieve_iteration_tree(op1)) == 3
buffers = [i for i in FindSymbols().visit(op1) if i.is_Array]
assert len(buffers) == 1
op0.apply(time_m=1, time_M=9)
op1.apply(time_m=1, time_M=9, v=v1)
assert np.all(v.data == v1.data)
def test_streaming_w_shifting(self):
nt = 50
grid = Grid(shape=(5, 5))
time = grid.time_dim
factor = Constant(name='factor', value=5, dtype=np.int32)
t_sub = ConditionalDimension('t_sub', parent=time, factor=factor)
save_shift = Constant(name='save_shift', dtype=np.int32)
u = TimeFunction(name='u', grid=grid, time_order=0)
usave = TimeFunction(name='usave',
grid=grid,
time_order=0,
save=(int(nt // factor.data)),
time_dim=t_sub)
for i in range(usave.save):
usave.data[i, :] = i
eqns = Eq(u.forward, u + usave.subs(t_sub, t_sub - save_shift))
op = Operator(eqns, opt=('streaming', 'orchestrate'))
# From time_m=15 to time_M=35 with a factor=5 -- it means that, thanks
# to t_sub, we enter the Eq exactly (35-15)/5 + 1 = 5 times. We set
# save_shift=1 so instead of accessing the range usave[15/5:35/5+1],
# we rather access the range usave[15/5-1:35:5], which means accessing
# the usave values 2, 3, 4, 5, 6.
op.apply(time_m=15, time_M=35, save_shift=1)
assert np.allclose(u.data, 20)
# Again, but with a different shift
op.apply(time_m=15, time_M=35, save_shift=-2)
assert np.allclose(u.data, 20 + 35)
def test_basic(self):
nt = 19
grid = Grid(shape=(11, 11))
time = grid.time_dim
u = TimeFunction(name='u', grid=grid)
assert (grid.stepping_dim in u.indices)
u2 = TimeFunction(name='u2', grid=grid, save=nt)
assert (time in u2.indices)
factor = 4
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
usave = TimeFunction(name='usave',
grid=grid,
save=(nt + factor - 1) // factor,
time_dim=time_subsampled)
assert (time_subsampled in usave.indices)
eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.), Eq(usave, u)]
op = Operator(eqns)
op.apply(t_M=nt - 2)
assert np.all(np.allclose(u.data[(nt - 1) % 3], nt - 1))
assert np.all([np.allclose(u2.data[i], i) for i in range(nt)])
assert np.all([
np.allclose(usave.data[i], i * factor)
for i in range((nt + factor - 1) // factor)
])
def test_save_w_nonaffine_time(self):
factor = 4
grid = Grid(shape=(11, 11))
x, y = grid.dimensions
t = grid.stepping_dim
time = grid.time_dim
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
f = Function(name='f', grid=grid, dtype=np.int32)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave',
grid=grid,
save=2,
time_dim=time_subsampled)
save_shift = Constant(name='save_shift', dtype=np.int32)
eqns = [
Eq(u.forward, u[t, f[x, x], f[y, y]] + 1.),
Eq(usave.subs(time_subsampled, time_subsampled - save_shift), u)
]
op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
# We just check the generated code here
assert len([i for i in FindSymbols().visit(op)
if isinstance(i, Lock)]) == 1
assert len(op._func_table) == 2
def test_save_w_shifting(self):
factor = 4
nt = 19
grid = Grid(shape=(11, 11))
time = grid.time_dim
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave',
grid=grid,
save=2,
time_dim=time_subsampled)
save_shift = Constant(name='save_shift', dtype=np.int32)
eqns = [
Eq(u.forward, u + 1.),
Eq(usave.subs(time_subsampled, time_subsampled - save_shift), u)
]
op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
# Starting at time_m=10, so time_subsampled - save_shift is in range
op.apply(time_m=10, time_M=nt - 2, save_shift=3)
assert np.all(np.allclose(u.data[0], 8))
assert np.all(
[np.allclose(usave.data[i], 2 + i * factor) for i in range(2)])
def test_save(self, opt, gpu_fit, async_degree):
nt = 10
grid = Grid(shape=(300, 300, 300))
time_dim = grid.time_dim
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub",
parent=time_dim,
factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave',
grid=grid,
time_order=0,
save=int(nt // factor.data),
time_dim=time_sub)
# For the given `nt` and grid shape, `usave` is roughly 4*5*300**3=~ .5GB of data
op = Operator(
[Eq(u.forward, u + 1), Eq(usave, u.forward)],
opt=(opt, {
'gpu-fit': usave if gpu_fit else None,
'buf-async-degree': async_degree
}))
op.apply(time_M=nt - 1)
assert all(
np.all(usave.data[i] == 2 * i + 1) for i in range(usave.save))
def test_save_multi_output(self):
nt = 10
grid = Grid(shape=(150, 150, 150))
time_dim = grid.time_dim
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
vsave = TimeFunction(name='vsave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
eqns = [Eq(u.forward, u + 1),
Eq(usave, u.forward),
Eq(vsave, u.forward)]
op = Operator(eqns, opt=('buffering', 'tasking', 'topofuse', 'orchestrate'))
# Check generated code
assert len(op._func_table) == 4 # usave and vsave eqns are in separate tasks
op.apply(time_M=nt-1)
assert all(np.all(usave.data[i] == 2*i + 1) for i in range(usave.save))
assert all(np.all(vsave.data[i] == 2*i + 1) for i in range(vsave.save))
def test_no_index(self):
"""Test behaviour when the ConditionalDimension is used as a symbol in
an expression."""
nt = 19
grid = Grid(shape=(11, 11))
time = grid.time_dim
u = TimeFunction(name='u', grid=grid)
assert (grid.stepping_dim in u.indices)
v = Function(name='v', grid=grid)
factor = 4
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
eqns = [Eq(u.forward, u + 1), Eq(v, v + u * u * time_subsampled)]
op = Operator(eqns)
op.apply(t_M=nt - 2)
assert np.all(np.allclose(u.data[(nt - 1) % 3], nt - 1))
# expected result is 1024
# v = u[0]**2 * 0 + u[4]**2 * 1 + u[8]**2 * 2 + u[12]**2 * 3 + u[16]**2 * 4
# with u[t] = t
# v = 16 * 1 + 64 * 2 + 144 * 3 + 256 * 4 = 1600
assert np.all(np.allclose(v.data, 1600))
def test_subsampled_fd(self):
"""
Test that the FD shortcuts are handled correctly with ConditionalDimensions
"""
grid = Grid(shape=(21, 21))
time = grid.time_dim
# Creates subsampled spatial dimensions and accordine grid
dims = tuple([
ConditionalDimension(d.name + 'sub', parent=d, factor=2)
for d in grid.dimensions
])
grid2 = Grid((6, 6), dimensions=dims, time_dimension=time)
u2 = TimeFunction(name='u2', grid=grid2, space_order=2, time_order=1)
u2.data.fill(2.)
eqns = [Eq(u2.forward, u2.dx + u2.dy)]
op = Operator(eqns)
op.apply(time_M=0, x_M=11, y_M=11)
# Verify that u2 contains subsampled fd values
assert np.all(u2.data[0, :, :] == 2.)
assert np.all(u2.data[1, 0, 0] == 20.)
assert np.all(u2.data[1, -1, -1] == -20.)
assert np.all(u2.data[1, 0, -1] == 0.)
assert np.all(u2.data[1, -1, 0] == -0.)
assert np.all(u2.data[1, 1:-1, 0] == 10.)
assert np.all(u2.data[1, 0, 1:-1] == 10.)
assert np.all(u2.data[1, 1:-1, -1] == -10.)
assert np.all(u2.data[1, -1, 1:-1] == -10.)
assert np.all(u2.data[1, 1:4, 1:4] == 0.)
def test_symbolic_factor(self):
"""
Test ConditionalDimension with symbolic factor (provided as a Constant).
"""
g = Grid(shape=(4, 4, 4))
u = TimeFunction(name='u', grid=g, time_order=0)
fact = Constant(name='fact', dtype=np.int32, value=4)
tsub = ConditionalDimension(name='tsub',
parent=g.time_dim,
factor=fact)
usave = TimeFunction(name='usave', grid=g, time_dim=tsub, save=4)
op = Operator([Eq(u, u + 1), Eq(usave, u)])
op.apply(time=7) # Use `fact`'s default value, 4
assert np.all(usave.data[0] == 1)
assert np.all(usave.data[1] == 5)
u.data[:] = 0.
op.apply(time=7, fact=2)
assert np.all(usave.data[0] == 1)
assert np.all(usave.data[1] == 3)
assert np.all(usave.data[2] == 5)
assert np.all(usave.data[3] == 7)
def test_topofusion_w_subdims_conddims_v2(self):
"""
Like `test_topofusion_w_subdims_conddims` but with more SubDomains,
so we expect fewer loop nests.
"""
grid = Grid(shape=(4, 4, 4))
time = grid.time_dim
f = TimeFunction(name='f', grid=grid, time_order=2)
g = TimeFunction(name='g', grid=grid, time_order=2)
h = TimeFunction(name='h', grid=grid, time_order=2)
fsave = TimeFunction(name='fsave', grid=grid, time_order=2, save=5)
gsave = TimeFunction(name='gsave', grid=grid, time_order=2, save=5)
ctime = ConditionalDimension(name='ctime',
parent=time,
condition=time > 4)
eqns = [
Eq(f.forward, f + 1, subdomain=grid.interior),
Eq(g.forward, g + 1, subdomain=grid.interior),
Eq(fsave, f.dt2, implicit_dims=[ctime]),
Eq(h, f + g, subdomain=grid.interior),
Eq(gsave, g.dt2, implicit_dims=[ctime])
]
op = Operator(eqns)
# Check generated code -- expect the gsave equation to be scheduled together
# in the same loop nest with the fsave equation
assert len(op._func_table) == 2
assert len(FindNodes(Expression).visit(
op._func_table['bf0'].root)) == 3
assert len(FindNodes(Expression).visit(
op._func_table['bf1'].root)) == 2 + 1 # r0
def test_subsampling(self):
"""
Tests (time) subsampling support. This stresses the compiler as two
different YASK kernels need to be generated.
"""
grid = Grid(shape=(8, 8))
time = grid.time_dim
nt = 9
u = TimeFunction(name='u', grid=grid)
u.data_with_halo[:] = 0.
# Setup subsampled function
factor = 4
nsamples = (nt+factor-1)//factor
times = ConditionalDimension('t_sub', parent=time, factor=factor)
usave = TimeFunction(name='usave', grid=grid, save=nsamples, time_dim=times)
eqns = [Eq(u.forward, u + 1.), Eq(usave, u)]
op = Operator(eqns)
op.apply(time=nt-1)
# Check numerical correctness
assert np.all(usave.data[0] == 0.)
assert np.all(usave.data[1] == 4.)
assert np.all(usave.data[2] == 8.)
# Check code generation
solns = FindNodes(ForeignExpression).visit(op)
assert len(solns) == 2
assert all('run_solution' in str(i) for i in solns)
def test_stepping_dim_in_condition_lowering(self):
"""
Check that the compiler performs lowering on conditions
with TimeDimensions and generates the expected code::
if (g[t][x + 1][y + 1] <= 10){ if (g[t0][x + 1][y + 1] <= 10){
... --> ...
} }
This test increments a function by one at every timestep until it is
less-or-equal to 10 (g<=10) while although operator runs for 13 timesteps.
"""
grid = Grid(shape=(4, 4))
_, y = grid.dimensions
ths = 10
g = TimeFunction(name='g', grid=grid)
ci = ConditionalDimension(name='ci', parent=y, condition=Le(g, ths))
op = Operator(Eq(g.forward, g + 1, implicit_dims=ci))
op.apply(time_M=ths + 3)
assert np.all(g.data[0, :, :] == ths)
assert np.all(g.data[1, :, :] == ths + 1)
assert 'if (g[t0][x + 1][y + 1] <= 10)\n'
'{\n g[t1][x + 1][y + 1] = g[t0][x + 1][y + 1] + 1' in str(op.ccode)
def test_relational_classes(self, setup_rel, rhs, c1, c2, c3, c4):
"""
Test ConditionalDimension using conditions based on Relations over SubDomains.
"""
class InnerDomain(SubDomain):
name = 'inner'
def define(self, dimensions):
return {d: ('middle', 2, 2) for d in dimensions}
inner_domain = InnerDomain()
grid = Grid(shape=(8, 8), subdomains=(inner_domain, ))
g = Function(name='g', grid=grid)
g2 = Function(name='g2', grid=grid)
for i in [g, g2]:
i.data[:4, :4] = 1
i.data[4:, :4] = 2
i.data[4:, 4:] = 3
i.data[:4, 4:] = 4
xi, yi = grid.subdomains['inner'].dimensions
cond = setup_rel(0.25 * g + 0.75 * g2,
rhs,
subdomain=grid.subdomains['inner'])
ci = ConditionalDimension(name='ci', parent=yi, condition=cond)
f = Function(name='f', shape=grid.shape, dimensions=(xi, ci))
eq1 = Eq(f, 0.4 * g + 0.6 * g2)
eq2 = Eq(f, 5)
Operator([eq1, eq2]).apply()
assert np.all(f.data[2:6, c1:c2] == 5.)
assert np.all(f.data[:, c3:c4] < 5.)
def test_implicit_dims(self):
"""
Test ConditionalDimension as an implicit dimension for an equation.
"""
# This test makes an Operator that should create a vector of increasing
# integers, but stop incrementing when a certain stop value is reached
shape = (50, )
stop_value = 20
time = Dimension(name='time')
f = TimeFunction(name='f', shape=shape, dimensions=[time])
# The condition to stop incrementing
cond = ConditionalDimension(name='cond',
parent=time,
condition=f[time] < stop_value)
eqs = [
Eq(f.forward, f),
Eq(f.forward, f.forward + 1, implicit_dims=[cond])
]
op = Operator(eqs)
op.apply(time_M=shape[0] - 2)
# Make the same calculation in python to assert the result
F = np.zeros(shape[0])
for i in range(shape[0]):
F[i] = i if i < stop_value else stop_value
assert np.all(f.data == F)
def test_laplace(self):
grid = Grid(shape=(20, 20, 20))
x, y, z = grid.dimensions
time = grid.time_dim
t = grid.stepping_dim
tsave = ConditionalDimension(name='tsave', parent=time, factor=2)
u = TimeFunction(name='u', grid=grid, save=None, time_order=2)
usave = TimeFunction(name='usave',
grid=grid,
time_dim=tsave,
time_order=0,
space_order=0)
steps = []
# save of snapshot
steps.append(Eq(usave, u))
# standard laplace-like thing
steps.append(
Eq(
u[t + 1, x, y, z], u[t, x, y, z] - u[t - 1, x, y, z] +
u[t, x - 1, y, z] + u[t, x + 1, y, z] + u[t, x, y - 1, z] +
u[t, x, y + 1, z] + u[t, x, y, z - 1] + u[t, x, y, z + 1]))
op = Operator(steps)
u.data[:] = 0.0
u.data[0, 10, 10, 10] = 1.0
op.apply(time_m=0, time_M=0)
assert np.sum(u.data[0, :, :, :]) == 1.0
assert np.sum(u.data[1, :, :, :]) == 7.0
assert np.all(usave.data[0, :, :, :] == u.data[0, :, :, :])
def test_subsampled_fd(self):
"""
Test that the symbolic interface is working for space subsampled
functions.
"""
nt = 19
grid = Grid(shape=(12, 12), extent=(11, 11))
u = TimeFunction(name='u', grid=grid, save=nt, space_order=2)
assert(grid.time_dim in u.indices)
# Creates subsampled spatial dimensions and according grid
dims = tuple([ConditionalDimension(d.name+'sub', parent=d, factor=2)
for d in u.grid.dimensions])
grid2 = Grid((6, 6), dimensions=dims)
u2 = TimeFunction(name='u2', grid=grid2, save=nt, space_order=1)
for i in range(nt):
for j in range(u2.data_with_halo.shape[2]):
u2.data_with_halo[i, :, j] = np.arange(u2.data_with_halo.shape[2])
eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2.dx)]
op = Operator(eqns, dse="advanced")
op.apply(time_M=nt-2)
# Verify that u2[1, x,y]= du2/dx[0, x, y]
assert np.allclose(u.data[-1], nt-1)
assert np.allclose(u2.data[1], 0.5)
def test_nothing_in_negative(self):
"""Test the case where when the condition is false, there is nothing to do."""
nt = 4
grid = Grid(shape=(11, 11))
time = grid.time_dim
u = TimeFunction(name='u', save=nt, grid=grid)
assert (grid.time_dim in u.indices)
factor = 4
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
usave = TimeFunction(name='usave',
grid=grid,
save=(nt + factor - 1) // factor,
time_dim=time_subsampled)
assert (time_subsampled in usave.indices)
eqns = [Eq(usave, u)]
op = Operator(eqns)
u.data[:] = 1.0
usave.data[:] = 0.0
op.apply(time_m=1, time_M=1)
assert np.allclose(usave.data, 0.0)
op.apply(time_m=0, time_M=0)
assert np.allclose(usave.data, 1.0)
def test_fission(self):
nt = 20
grid = Grid(shape=(10, 10))
time = grid.time_dim
usave = TimeFunction(name='usave', grid=grid, save=nt, time_order=2)
vsave = TimeFunction(name='vsave', grid=grid, save=nt, time_order=2)
ctime0 = ConditionalDimension(name='ctime',
parent=time,
condition=time > 4)
ctime1 = ConditionalDimension(name='ctime',
parent=time,
condition=time <= 4)
eqns = [
Eq(usave, time + 0.2, implicit_dims=[ctime0]),
Eq(vsave, time + 0.2, implicit_dims=[ctime1])
]
op = Operator(eqns)
# Check generated code
trees = retrieve_iteration_tree(op, mode='superset')
assert len(trees) == 2
assert trees[0].root is trees[1].root
assert trees[0][1] is not trees[1][1]
assert trees[0].root.dim is time
assert not trees[0].root.pragmas
assert trees[0][1].pragmas
assert trees[1][1].pragmas
op.apply()
expected = np.full(shape=(20, 10, 10),
fill_value=0.2,
dtype=np.float32)
for i in range(nt):
expected[i] += i
assert np.all(usave.data[5:] == expected[5:])
assert np.all(vsave.data[:5] == expected[:5])
def test_issue_1592(self):
grid = Grid(shape=(11, 11))
time = grid.time_dim
time_sub = ConditionalDimension('t_sub', parent=time, factor=2)
v = TimeFunction(name="v", grid=grid, space_order=4, time_dim=time_sub, save=5)
w = Function(name="w", grid=grid, space_order=4)
Operator(Eq(w, v.dx))(time=6)
op = Operator(Eq(v.forward, v.dx))
op.apply(time=6)
exprs = FindNodes(Expression).visit(op)
assert exprs[-1].expr.lhs.indices[0] == IntDiv(time, 2) + 1
def test_everything():
nt = 50
grid = Grid(shape=(6, 6))
x, y = grid.dimensions
time = grid.time_dim
xi = SubDimension.middle(name='xi',
parent=x,
thickness_left=2,
thickness_right=2)
yi = SubDimension.middle(name='yi',
parent=y,
thickness_left=2,
thickness_right=2)
factor = Constant(name='factor', value=5, dtype=np.int32)
t_sub = ConditionalDimension('t_sub', parent=time, factor=factor)
save_shift = Constant(name='save_shift', dtype=np.int32)
u = TimeFunction(name='u', grid=grid, time_order=0)
u1 = TimeFunction(name='u', grid=grid, time_order=0)
va = TimeFunction(name='va',
grid=grid,
time_order=0,
save=(int(nt // factor.data)),
time_dim=t_sub)
vb = TimeFunction(name='vb',
grid=grid,
time_order=0,
save=(int(nt // factor.data)),
time_dim=t_sub)
for i in range(va.save):
va.data[i, :] = i
vb.data[i, :] = i * 2 - 1
vas = va.subs(t_sub, t_sub - save_shift)
vasb = va.subs(t_sub, t_sub - 1 - save_shift)
vasf = va.subs(t_sub, t_sub + 1 - save_shift)
eqns = [Eq(u.forward, u + (vasb + vas + vasf) * 2. + vb)]
eqns = [e.xreplace({x: xi, y: yi}) for e in eqns]
op0 = Operator(eqns, opt='noop')
op1 = Operator(eqns, opt='buffering')
# Check generated code
assert len([i for i in FindSymbols().visit(op1) if i.is_Array]) == 2
op0.apply(time_m=15, time_M=35, save_shift=0)
op1.apply(time_m=15, time_M=35, save_shift=0, u=u1)
assert np.all(u.data == u1.data)
def test_conditional_dimension():
d = Dimension(name='d')
s = Scalar(name='s')
cd = ConditionalDimension(name='ci', parent=d, factor=4, condition=s > 3)
pkl_cd = pickle.dumps(cd)
new_cd = pickle.loads(pkl_cd)
assert cd.name == new_cd.name
assert cd.parent == new_cd.parent
assert cd.factor == new_cd.factor
assert cd.condition == new_cd.condition
def test_conditional_dimension(self):
"""
Test that ConditionalDimensions with same name but different attributes do not
alias to the same ConditionalDimension. Conversely, if the name and the attributes
are the same, they must alias to the same ConditionalDimension.
"""
i = Dimension(name='i')
ci0 = ConditionalDimension(name='ci', parent=i, factor=4)
ci1 = ConditionalDimension(name='ci', parent=i, factor=4)
assert ci0 is ci1
ci2 = ConditionalDimension(name='ci', parent=i, factor=8)
assert ci2 is not ci1
ci3 = ConditionalDimension(name='ci',
parent=i,
factor=4,
indirect=True)
assert ci3 is not ci1
s = Scalar(name='s')
ci4 = ConditionalDimension(name='ci',
parent=i,
factor=4,
condition=s > 3)
assert ci4 is not ci1
ci5 = ConditionalDimension(name='ci',
parent=i,
factor=4,
condition=s > 3)
assert ci5 is ci4
def test_conditional_dimension(self):
"""Test that different ConditionalDimensions have different hash value."""
d0 = Dimension(name='d')
s0 = Scalar(name='s')
d1 = Dimension(name='d', spacing=s0)
cd0 = ConditionalDimension(name='cd', parent=d0, factor=4)
cd1 = ConditionalDimension(name='cd', parent=d0, factor=5)
assert cd0 is not cd1
assert hash(cd0) != hash(cd1)
cd2 = ConditionalDimension(name='cd',
parent=d0,
factor=4,
indirect=True)
assert hash(cd0) != hash(cd2)
cd3 = ConditionalDimension(name='cd', parent=d1, factor=4)
assert hash(cd0) != hash(cd3)
s1 = Scalar(name='s', dtype=np.int32)
cd4 = ConditionalDimension(name='cd',
parent=d0,
factor=4,
condition=s0 > 3)
assert hash(cd0) != hash(cd4)
cd5 = ConditionalDimension(name='cd',
parent=d0,
factor=4,
condition=s1 > 3)
assert hash(cd0) != hash(cd5)
assert hash(cd4) != hash(cd5)