def test_symbol_cache_aliasing_reverse():
"""Test to assert that removing he original u[x, y] instance does
not impede our alisaing cache or leaks memory.
"""
# Ensure a clean cache to start with
clear_cache()
# FIXME: Currently not working, presumably due to our
# failure to cache new instances?
# assert(len(_SymbolCache) == 0)
# Create first instance of u and fill its data
grid = Grid(shape=(3, 4))
u = Function(name='u', grid=grid)
u.data[:] = 6.
u_ref = weakref.ref(u.data)
# Create derivative and delete orignal u[x, y]
dx = u.dx
del u
clear_cache()
# We still have a references to u
# FIXME: Unreliable cache sizes
# assert len(_SymbolCache) == 1
# Ensure u[x + h, y] still holds valid data
assert np.allclose(dx.args[0].args[1].data, 6.)
del dx
clear_cache()
# FIXME: Unreliable cache sizes
# assert len(_SymbolCache) == 0 # We still have a reference to u_h
assert u_ref() is None
def unit_box(name='a', shape=(11, 11), grid=None):
"""Create a field with value 0. to 1. in each dimension"""
grid = grid or Grid(shape=shape)
a = Function(name=name, grid=grid)
dims = tuple([np.linspace(0., 1., d) for d in shape])
a.data[:] = np.meshgrid(*dims)[1]
return a
def test_arithmetic(self):
"""Test arithmetic operations involving Data objects."""
grid = Grid(shape=(16, 16, 16))
u = Function(name='yu3D', grid=grid, space_order=0)
u.data[:] = 1
# Simple arithmetic
assert np.all(u.data == 1)
assert np.all(u.data + 2. == 3.)
assert np.all(u.data - 2. == -1.)
assert np.all(u.data * 2. == 2.)
assert np.all(u.data / 2. == 0.5)
assert np.all(u.data % 2 == 1.)
# Increments and partial increments
u.data[:] += 2.
assert np.all(u.data == 3.)
u.data[9, :, :] += 1.
assert all(np.all(u.data[i, :, :] == 3.) for i in range(9))
assert np.all(u.data[9, :, :] == 4.)
# Right operations __rOP__
u.data[:] = 1.
arr = np.ndarray(shape=(16, 16, 16), dtype=np.float32)
arr.fill(2.)
assert np.all(arr - u.data == 1.)
def test_slicing(self):
grid = Grid(shape=(4, 4))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
myrank = grid.distributor.myrank
u = Function(name='u', grid=grid, space_order=0)
u.data[:] = myrank
# `u.data` is a view of the global data array restricted, on each rank,
# to the local rank domain, so it must be == myrank
assert np.all(u.data == myrank)
assert np.all(u.data._local == myrank)
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data[:2, :2] == myrank)
assert u.data[:2, 2:].size == u.data[2:, :2].size == u.data[2:, 2:].size == 0
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert np.all(u.data[:2, 2:] == myrank)
assert u.data[:2, :2].size == u.data[2:, :2].size == u.data[2:, 2:].size == 0
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data[2:, :2] == myrank)
assert u.data[:2, 2:].size == u.data[:2, :2].size == u.data[2:, 2:].size == 0
else:
assert np.all(u.data[2:, 2:] == myrank)
assert u.data[:2, 2:].size == u.data[2:, :2].size == u.data[:2, :2].size == 0
def test_indexing(self):
grid = Grid(shape=(4, 4))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
myrank = grid.distributor.myrank
u = Function(name='u', grid=grid, space_order=0)
u.data[:] = myrank
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert u.data[0, 0] == myrank
assert u.data[2, 2] is None
assert u.data[2].size == 0
assert u.data[:, 2].size == 0
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert u.data[0, 0] is None
assert u.data[2, 2] is None
assert u.data[2].size == 0
assert np.all(u.data[:, 2] == [myrank, myrank])
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert u.data[0, 0] is None
assert u.data[2, 2] is None
assert np.all(u.data[2] == [myrank, myrank])
assert u.data[:, 2].size == 0
else:
assert u.data[0, 0] is None
assert u.data[2, 2] == myrank
assert np.all(u.data[2] == [myrank, myrank])
assert np.all(u.data[:, 2] == [myrank, myrank])
def test_partial_offloading(self):
"""
Check that :class:`Function` objects not using any :class:`SpaceDimension`
are computed in Devito-land, rather than via YASK.
"""
shape = (4, 4, 4)
grid = Grid(shape=shape)
dx = Dimension(name='dx')
dy = Dimension(name='dy')
dz = Dimension(name='dz')
x, y, z = grid.dimensions
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
f = Function(name='f', dimensions=(dx, dy, dz), shape=shape)
u.data_with_halo[:] = 0.
f.data[:] = 0.
eqns = [Eq(u.forward, u + 1.),
Eq(f, u[1, dx, dy, dz] + 1.)]
op = Operator(eqns)
op(time=0)
assert np.all(u.data[0] == 0.)
assert np.all(u.data[1] == 1.)
assert np.all(f.data == 2.)
def a(shape=(11, 11)):
grid = Grid(shape=shape)
a = Function(name='a', grid=grid)
xarr = np.linspace(0., 1., shape[0])
yarr = np.linspace(0., 1., shape[1])
a.data[:] = np.meshgrid(xarr, yarr)[1]
return a
def test_localviews(self):
grid = Grid(shape=(4, 4))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
myrank = grid.distributor.myrank
u = Function(name='u', grid=grid)
u.data[:] = grid.distributor.myrank
assert u.data_ro_domain._local[0, 0] == grid.distributor.myrank
assert u.data_ro_domain._local[1, 1] == grid.distributor.myrank
assert u.data_ro_domain._local[-1, -1] == grid.distributor.myrank
assert u.data_ro_with_halo._local[1, 1] == grid.distributor.myrank
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data_ro_with_halo._local[1:, 1:] == myrank)
assert np.all(u.data_ro_with_halo._local[0] == 0.)
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert np.all(u.data_ro_with_halo._local[1:3, :2] == myrank)
assert np.all(u.data_ro_with_halo._local[0] == 0.)
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data_ro_with_halo._local[:2, 1:3] == myrank)
assert np.all(u.data_ro_with_halo._local[2] == 0.)
else:
assert np.all(u.data_ro_with_halo._local[:2, :2] == myrank)
assert np.all(u.data_ro_with_halo._local[2] == 0.)
def test_broadcasting(self):
"""
Test Data broadcasting, expected to behave as NumPy broadcasting.
Notes
-----
Refer to https://docs.scipy.org/doc/numpy-1.15.0/user/basics.broadcasting.html
for more info about NumPy broadcasting rules.
"""
grid = Grid(shape=(4, 4, 4))
u = Function(name='yu3D', grid=grid)
u.data[:] = 2.
# Assign from array with lower-dimensional shape
v = np.ones(shape=(4, 4), dtype=u.dtype)
u.data[:] = v
assert np.all(u.data == 1.)
# Assign from array with higher-dimensional shape causes a ValueError exception
v = np.zeros(shape=(4, 4, 4, 4), dtype=u.dtype)
try:
u.data[:] = v
except ValueError:
assert True
except:
assert False
# Assign from array having shape with some 1-valued entries
v = np.zeros(shape=(4, 1, 4), dtype=u.dtype)
u.data[:] = v
assert np.all(u.data == 0.)
def test_fd_space_staggered(self, space_order, stagger):
"""
This test compares the discrete finite-difference scheme against polynomials
For a given order p, the finite difference scheme should
be exact for polynomials of order p
:param derivative: name of the derivative to be tested
:param space_order: space order of the finite difference stencil
"""
clear_cache()
# dummy axis dimension
nx = 100
xx = np.linspace(-1, 1, nx)
dx = xx[1] - xx[0]
# Symbolic data
grid = Grid(shape=(nx,), dtype=np.float32)
x = grid.dimensions[0]
# Location of the staggered function
if stagger == left:
off = -.5
side = -x
xx2 = xx - off * dx
elif stagger == right:
off = .5
side = x
xx2 = xx[:-1] - off * dx
else:
off = 0
side = NODE
xx2 = xx
u = Function(name="u", grid=grid, space_order=space_order, staggered=(side,))
du = Function(name="du", grid=grid, space_order=space_order)
# Define polynomial with exact fd
coeffs = np.ones((space_order-1,), dtype=np.float32)
polynome = sum([coeffs[i]*x**i for i in range(0, space_order-1)])
polyvalues = np.array([polynome.subs(x, xi) for xi in xx2], np.float32)
# Fill original data with the polynomial values
u.data[:] = polyvalues
# True derivative of the polynome
Dpolynome = diff(polynome)
Dpolyvalues = np.array([Dpolynome.subs(x, xi) for xi in xx], np.float32)
# FD derivative, symbolic
u_deriv = generic_derivative(u, deriv_order=1, fd_order=space_order,
dim=x, stagger=stagger)
# Compute numerical FD
stencil = Eq(du, u_deriv)
op = Operator(stencil, subs={x.spacing: dx})
op.apply()
# Check exactness of the numerical derivative except inside space_brd
space_border = space_order
error = abs(du.data[space_border:-space_border] -
Dpolyvalues[space_border:-space_border])
assert np.isclose(np.mean(error), 0., atol=1e-3)
def test_trivial_insertion(self):
grid = Grid(shape=(4, 4))
u = Function(name='u', grid=grid, space_order=0)
v = Function(name='v', grid=grid, space_order=1)
u.data[:] = 1.
assert np.all(u.data == 1.)
assert np.all(u.data._local == 1.)
v.data_with_halo[:] = 1.
assert v.data_with_halo[:].shape == (3, 3)
assert np.all(v.data_with_halo == 1.)
assert np.all(v.data_with_halo[:] == 1.)
assert np.all(v.data_with_halo._local == 1.)
def test_function():
grid = Grid(shape=(3, 3, 3))
f = Function(name='f', grid=grid)
f.data[0] = 1.
pkl_f = pickle.dumps(f)
new_f = pickle.loads(pkl_f)
# .data is initialized, so it should have been pickled too
assert np.all(f.data[0] == 1.)
assert np.all(new_f.data[0] == 1.)
assert f.space_order == new_f.space_order
assert f.dtype == new_f.dtype
assert f.shape == new_f.shape
def test_inject_from_field(shape, coords, result, npoints=19):
"""Test point injection from a second field along a line
through the middle of the grid.
"""
a = unit_box(shape=shape)
a.data[:] = 0.
b = Function(name='b', grid=a.grid)
b.data[:] = 1.
p = points(a.grid, ranges=coords, npoints=npoints)
expr = p.inject(field=a, expr=b)
Operator(expr)(a=a, b=b)
indices = [slice(4, 6, 1) for _ in coords]
indices[0] = slice(1, -1, 1)
assert np.allclose(a.data[indices], result, rtol=1.e-5)
def test_halo_indexing(self):
"""Test data packing/unpacking in presence of a halo region."""
domain_shape = (16, 16, 16)
grid = Grid(shape=domain_shape)
u = Function(name='yu3D', grid=grid, space_order=2)
assert u.shape == u.data.shape == domain_shape
assert u._shape_with_inhalo == u.data_with_halo.shape == (20, 20, 20)
assert u.shape_with_halo == u._shape_with_inhalo # W/o MPI, these two coincide
# Test simple insertion and extraction
u.data_with_halo[0, 0, 0] = 1.
u.data[0, 0, 0] = 2.
assert u.data_with_halo[0, 0, 0] == 1.
assert u.data[0, 0, 0] == 2.
assert u.data_with_halo[2, 2, 2] == 2.
# Test negative indices
u.data_with_halo[-1, -1, -1] = 3.
assert u.data[-1, -1, -1] == 0.
assert u.data_with_halo[-1, -1, -1] == 3.
def test_symbol_cache_aliasing():
"""Test to assert that our aliasing cache isn't defeated by sympys
non-aliasing symbol cache.
For further explanation consider the symbol u[x, y] and it's first
derivative in x, which includes the symbols u[x, y] and u[x + h, y].
The two functions are aliased in devito's caching mechanism to allow
multiple stencil indices pointing at the same data object u, but
SymPy treats these two instances as separate functions and thus is
allowed to delete one or the other when the cache is cleared.
The test below asserts that u[x + h, y] is deleted, the data on u
is still intact through our own caching mechanism."""
# Ensure a clean cache to start with
clear_cache()
# FIXME: Currently not working, presumably due to our
# failure to cache new instances?
# assert(len(_SymbolCache) == 0)
# Create first instance of u and fill its data
grid = Grid(shape=(3, 4))
u = Function(name='u', grid=grid)
u.data[:] = 6.
u_ref = weakref.ref(u.data)
# Create u[x + h, y] and delete it again
dx = u.dx # Contains two u symbols: u[x, y] and u[x + h, y]
del dx
clear_cache()
# FIXME: Unreliable cache sizes
# assert len(_SymbolCache) == 1 # We still have a reference to u
assert np.allclose(u.data, 6.) # u.data is alive and well
# Remove the final instance and ensure u.data got deallocated
del u
clear_cache()
assert u_ref() is None
def test_from_replicated_to_distributed(self):
shape = (4, 4)
grid = Grid(shape=shape)
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
u = Function(name='u', grid=grid, space_order=0) # distributed
v = Function(name='v', grid=grid, space_order=0) # distributed
a = np.arange(16).reshape(shape) # replicated
# Full array
u.data[:] = a
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data == [[0, 1], [4, 5]])
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert np.all(u.data == [[2, 3], [6, 7]])
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data == [[8, 9], [12, 13]])
else:
assert np.all(u.data == [[10, 11], [14, 15]])
# Subsection (all ranks touched)
u.data[:] = 0
u.data[1:3, 1:3] = a[1:3, 1:3]
# Same as above but with negative indices
v.data[:] = 0
v.data[1:-1, 1:-1] = a[1:-1, 1:-1]
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data == [[0, 0], [0, 5]])
assert np.all(v.data == [[0, 0], [0, 5]])
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert np.all(u.data == [[0, 0], [6, 0]])
assert np.all(v.data == [[0, 0], [6, 0]])
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(u.data == [[0, 9], [0, 0]])
assert np.all(v.data == [[0, 9], [0, 0]])
else:
assert np.all(u.data == [[10, 0], [0, 0]])
assert np.all(v.data == [[10, 0], [0, 0]])
# The assigned data must have same shape as the one of the distributed array,
# otherwise an exception is expected
try:
u.data[1:3, 1:3] = a[1:2, 1:2]
except ValueError:
assert True
except:
assert False
def _new_operator1(shape, blockshape=None, dle=None):
blockshape = as_tuple(blockshape)
grid = Grid(shape=shape, dtype=np.int32)
infield = Function(name='infield', grid=grid)
infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
outfield = Function(name='outfield', grid=grid)
stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0)
op = Operator(stencil, dle=dle)
blocksizes = get_blocksizes(op, dle, grid, blockshape)
op(infield=infield, outfield=outfield, **blocksizes)
return outfield, op
def _gen_phys_param(self,
field,
name,
space_order,
is_param=False,
default_value=0,
init_empty=False):
if field is None and not init_empty:
return default_value
if isinstance(field, np.ndarray) or init_empty:
function = Function(name=name,
grid=self.grid,
space_order=space_order,
parameter=is_param)
if not init_empty:
if name is 'rho':
initialize_function(function, 1 / field, self.nbpml)
else:
initialize_function(function, field, self.nbpml)
else:
function = Constant(name=name, value=field)
self._physical_parameters.append(name)
return function
def test_function_wo(self):
grid = Grid(shape=(3, 3, 3))
i = Dimension(name='i')
f = Function(name='f', shape=(1,), dimensions=(i,), grid=grid)
u = TimeFunction(name='u', grid=grid)
eqns = [Eq(u.forward, u + 1),
Eq(f[0], u[0, 0, 0, 0])]
op = Operator(eqns, opt='noop', language='openmp')
assert len(op.body.maps) == 1
assert op.body.maps[0].pragmas[0].value ==\
('omp target enter data map(to: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert len(op.body.unmaps) == 2
assert op.body.unmaps[0].pragmas[0].value ==\
('omp target update from(u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert op.body.unmaps[1].pragmas[0].value ==\
('omp target exit data map(release: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]]) if(devicerm)')
def test_array_rw(self):
grid = Grid(shape=(3, 3, 3))
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid, space_order=2)
eqn = Eq(u.forward, u*cos(f*2))
op = Operator(eqn, language='openmp')
assert len(op.body.allocs) == 1
assert str(op.body.allocs[0]) ==\
('float *r0_vec = (float*) '
'omp_target_alloc(sizeof(float[x_size*y_size*z_size]), '
'omp_get_default_device());')
assert len(op.body.maps) == 2
assert all('r0' not in str(i) for i in op.body.maps)
assert len(op.body.frees) == 1
assert str(op.body.frees[0]) ==\
'omp_target_free(r0_vec, omp_get_default_device());'
assert len(op.body.unmaps) == 3
assert all('r0' not in str(i) for i in op.body.unmaps)
def test_if_parallel(self):
a = np.arange(36).reshape((6, 6))
grid = Grid(shape=(18, 18))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
f = Function(name='f',
grid=grid,
halo=((3, 3), (3, 3)),
dtype=np.int32)
initialize_function(f, a, 6, mode='reflect')
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(a[::-1, 0:3] - np.array(f.data[0:6, 6:9]) == 0)
assert np.all(a[0:3, ::-1] - np.array(f.data[6:9, 0:6]) == 0)
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert np.all(a[::-1, 3:6] - np.array(f.data[0:6, 9:12]) == 0)
assert np.all(a[0:3, ::-1] - np.array(f.data[6:9, 12:18]) == 0)
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert np.all(a[::-1, 0:3] - np.array(f.data[12:18, 6:9]) == 0)
assert np.all(a[3:6, ::-1] - np.array(f.data[9:12, 0:6]) == 0)
else:
assert np.all(a[::-1, 3:6] - np.array(f.data[12:18, 9:12]) == 0)
assert np.all(a[3:6, ::-1] - np.array(f.data[9:12, 12:18]) == 0)
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_misc_dims(self):
"""
Tests grid-independent :class:`Function`s, which require YASK's "misc"
dimensions.
"""
dx = Dimension(name='dx')
grid = Grid(shape=(10, 10))
x, y = grid.dimensions
time = grid.time_dim
u = TimeFunction(name='u', grid=grid, time_order=1, space_order=4, save=4)
c = Function(name='c', dimensions=(x, dx), shape=(10, 5))
step = Eq(u.forward, (
u[time, x-2, y] * c[x, 0]
+ u[time, x-1, y] * c[x, 1]
+ u[time, x, y] * c[x, 2]
+ u[time, x+1, y] * c[x, 3]
+ u[time, x+2, y] * c[x, 4]))
for i in range(10):
c.data[i, 0] = 1.0+i
c.data[i, 1] = 1.0+i
c.data[i, 2] = 3.0+i
c.data[i, 3] = 6.0+i
c.data[i, 4] = 5.0+i
u.data[:] = 0.0
u.data[0, 2, :] = 2.0
op = Operator(step)
assert 'run_solution' in str(op)
op(time_m=0, time_M=0)
assert(np.all(u.data[1, 0, :] == 10.0))
assert(np.all(u.data[1, 1, :] == 14.0))
assert(np.all(u.data[1, 2, :] == 10.0))
assert(np.all(u.data[1, 3, :] == 8.0))
assert(np.all(u.data[1, 4, :] == 10.0))
assert(np.all(u.data[1, 5:10, :] == 0.0))
def test_arrays_defined_over_subdims(self):
"""
Check code generation when an Array uses a SubDimension.
"""
grid = Grid(shape=(3, ))
x, = grid.dimensions
xi, = grid.interior.dimensions
f = Function(name='f', grid=grid)
a = Array(name='a', dimensions=(xi, ), dtype=grid.dtype)
op = Operator(
[Eq(a[xi], 1),
Eq(f, f + a[xi + 1], subdomain=grid.interior)],
dle=('advanced', {
'openmp': False
}))
assert len(op.parameters) == 6
# neither `x_size` nor `xi_size` are expected here
assert not any(i.name in ('x_size', 'xi_size') for i in op.parameters)
# Try running it -- regardless of what it will produce, this should run
# ie, this checks this error isn't raised:
# "ValueError: No value found for parameter xi_size"
op()
def test_default_functions(self):
"""
Test the default argument derivation for functions.
"""
grid = Grid(shape=(5, 6, 7))
f = TimeFunction(name='f', grid=grid)
g = Function(name='g', grid=grid)
op = Operator(Eq(g, g + f))
expected = {
'x_size': 5,
'x_s': 0,
'x_e': 5,
'y_size': 6,
'y_s': 0,
'y_e': 6,
'z_size': 7,
'z_s': 0,
'z_e': 7,
'f': f.data,
'g': g.data,
}
self.verify_arguments(op.arguments(), expected)
def test_operator_leakage_sparse(self):
"""
Test to ensure that Operator creation does not cause memory leaks for
SparseTimeFunctions.
"""
grid = Grid(shape=(5, 6))
a = Function(name='a', grid=grid)
s = SparseTimeFunction(name='s', grid=grid, npoint=1, nt=1)
w_a = weakref.ref(a)
w_s = weakref.ref(s)
# Create operator and delete everything again
op = Operator(s.interpolate(a))
w_op = weakref.ref(op)
del op
del s
del a
clear_cache()
# Test whether things are still hanging around
assert w_a() is None
assert w_s() is None
assert w_op() is None
def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=1000.0,
space_order=4, kernel='OT2', nbpml=40, full_run=False,
autotune=False, constant=False, checkpointing=False, **kwargs):
solver = acoustic_setup(shape=shape, spacing=spacing, nbpml=nbpml, tn=tn,
space_order=space_order, kernel=kernel,
constant=constant, **kwargs)
# Smooth velocity
initial_vp = Function(name='v0', grid=solver.model.grid, space_order=space_order)
smooth(initial_vp, solver.model.m)
dm = np.float32(initial_vp.data**2 - solver.model.m.data)
info("Applying Forward")
# Whether or not we save the whole time history. We only need the full wavefield
# with 'save=True' if we compute the gradient without checkpointing, if we use
# checkpointing, PyRevolve will take care of the time history
save = full_run and not checkpointing
# Define receiver geometry (spread across x, just below surface)
rec, u, summary = solver.forward(save=save, autotune=autotune)
if constant:
# With a new m as Constant
m0 = Constant(name="m", value=.25, dtype=np.float32)
solver.forward(save=save, m=m0)
# With a new m as a scalar value
solver.forward(save=save, m=.25)
if not full_run:
return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
info("Applying Adjoint")
solver.adjoint(rec, autotune=autotune)
info("Applying Born")
solver.born(dm, autotune=autotune)
info("Applying Gradient")
solver.gradient(rec, u, autotune=autotune, checkpointing=checkpointing)
def fwi_gradient(m_in):
# Create symbols to hold the gradient and residual
grad = Function(name="grad", grid=model.grid)
residual = Receiver(name='rec',
grid=model.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)
objective = 0.
# Creat forward wavefield to reuse to avoid memory overload
u0 = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=4,
save=geometry.nt)
for i in range(nshots):
# Important: We force previous wavefields to be destroyed,
# so that we may reuse the memory.
clear_cache()
# Update source location
geometry.src_positions[0, :] = source_locations[i, :]
# Generate synthetic data from true model
true_d, _, _ = solver.forward(m=model.m)
# Compute smooth data and full forward wavefield u0
u0.data.fill(0.)
smooth_d, _, _ = solver.forward(m=m_in, save=True, u=u0)
# Compute gradient from data residual and update objective function
residual.data[:] = smooth_d.data[:] - true_d.data[:]
objective += .5 * np.linalg.norm(residual.data.flatten())**2
solver.gradient(rec=residual, u=u0, m=m_in, grad=grad)
return objective, grad.data
def gradient(self, rec, u, v=None, grad=None, m=None, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
:param recin: Receiver data as a numpy array
:param u: Symbol for full wavefield `u` (created with save=True)
:param v: (Optional) Symbol to store the computed wavefield
:param grad: (Optional) Symbol to store the gradient field
:returns: Gradient field and performance summary
"""
# Gradient symbol
if grad is None:
grad = Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
if v is None:
v = TimeFunction(name='v',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
# Pick m from model unless explicitly provided
if m is None:
m = m or self.model.m
summary = self.op_grad().apply(rec=rec,
grad=grad,
v=v,
u=u,
m=m,
dt=kwargs.pop('dt', self.dt),
**kwargs)
return grad, summary
def test_mpi_fullmode_objects():
grid = Grid(shape=(4, 4, 4))
x, y, _ = grid.dimensions
# Message
f = Function(name='f', grid=grid)
obj = MPIMsgEnriched('msg', f, [Halo(x, LEFT)])
pkl_obj = pickle.dumps(obj)
new_obj = pickle.loads(pkl_obj)
assert obj.name == new_obj.name
assert obj.function.name == new_obj.function.name
assert all(obj.function.dimensions[i].name == new_obj.function.dimensions[i].name
for i in range(grid.dim))
assert new_obj.function.dimensions[0] is new_obj.halos[0].dim
# Region
x_m, x_M = x.symbolic_min, x.symbolic_max
y_m, y_M = y.symbolic_min, y.symbolic_max
obj = MPIRegion('reg', 1, [y, x],
[(((x, OWNED, LEFT),), {x: (x_m, Min(x_M, x_m))}),
(((y, OWNED, LEFT),), {y: (y_m, Min(y_M, y_m))})])
pkl_obj = pickle.dumps(obj)
new_obj = pickle.loads(pkl_obj)
assert obj.prefix == new_obj.prefix
assert obj.key == new_obj.key
assert obj.name == new_obj.name
assert len(new_obj.arguments) == 2
assert all(d0.name == d1.name for d0, d1 in zip(obj.arguments, new_obj.arguments))
assert all(new_obj.arguments[i] is new_obj.owned[i][0][0][0] # `x` and `y`
for i in range(2))
assert new_obj.owned[0][0][0][1] is new_obj.owned[1][0][0][1] # `OWNED`
assert new_obj.owned[0][0][0][2] is new_obj.owned[1][0][0][2] # `LEFT`
for n, i in enumerate(new_obj.owned):
d, v = list(i[1].items())[0]
assert d is new_obj.arguments[n]
assert v[0] is d.symbolic_min
assert v[1] == Min(d.symbolic_max, d.symbolic_min)
def test_multiple_subnests_v0(self):
grid = Grid(shape=(3, 3, 3))
x, y, z = grid.dimensions
t = grid.stepping_dim
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid, space_order=3)
eqn = Eq(
u.forward,
_R(
_R(u[t, x, y, z] + u[t, x + 1, y + 1, z + 1]) * 3. * f +
_R(u[t, x + 2, y + 2, z + 2] + u[t, x + 3, y + 3, z + 3]) *
3. * f) + 1.)
op = Operator(eqn,
opt=('advanced', {
'openmp': True,
'cire-mingain': 0,
'par-nested': 0,
'par-collapse-ncores': 1,
'par-dynamic-work': 0
}))
bns, _ = assert_blocking(op, {'x0_blk0'})
trees = retrieve_iteration_tree(bns['x0_blk0'])
assert len(trees) == 2
assert trees[0][0] is trees[1][0]
assert trees[0][0].pragmas[0].value ==\
'omp for collapse(2) schedule(dynamic,1)'
assert trees[0][2].pragmas[0].value == ('omp parallel for collapse(2) '
'schedule(dynamic,1) '
'num_threads(nthreads_nested)')
assert trees[1][2].pragmas[0].value == ('omp parallel for collapse(2) '
'schedule(dynamic,1) '
'num_threads(nthreads_nested)')
def run(shape=(50, 50, 50),
spacing=(10.0, 10.0, 10.0),
tn=1000.0,
space_order=4,
nbl=40,
full_run=False,
autotune=False,
**kwargs):
solver = acoustic_ssa_setup(shape=shape,
spacing=spacing,
nbl=nbl,
tn=tn,
space_order=space_order,
**kwargs)
info("Applying Forward")
# Define receiver geometry (spread across x, just below surface)
rec, u, summary = solver.forward(save=full_run, autotune=autotune)
if not full_run:
return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
# Smooth velocity
initial_vp = Function(name='v0',
grid=solver.model.grid,
space_order=space_order)
smooth(initial_vp, solver.model.vp)
dm = solver.model.vp - initial_vp
info("Applying Adjoint")
solver.adjoint(rec, autotune=autotune)
info("Applying Born")
solver.jacobian(dm, autotune=autotune)
info("Applying Gradient")
solver.jacobian_adjoint(rec, u, autotune=autotune)
return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
def test_argument_derivation_order(self, nt=100):
""" Ensure the precedence order of arguments is respected
Defaults < (overriden by) Tensor Arguments < Dimensions < Scalar Arguments
"""
i, j, k = dimify('i j k')
shape = (10, 10, 10)
grid = Grid(shape=shape, dimensions=(i, j, k))
a = Function(name='a', grid=grid).indexed
b_function = TimeFunction(name='b', grid=grid, save=True, time_dim=nt)
b = b_function.indexed
time = b_function.indices[0]
b1 = TimeFunction(name='b1', grid=grid, save=True,
time_dim=nt + 1).indexed
eqn = Eq(b[time, i, j, k], a[i, j, k])
op = Operator(eqn)
# Simple case, same as that tested above.
# Repeated here for clarity of further tests.
op_arguments, _ = op.arguments()
assert (op_arguments[time.start_name] == 0)
assert (op_arguments[time.end_name] == nt)
# Providing a tensor argument should infer the dimension size from its shape
op_arguments, _ = op.arguments(b=b1)
assert (op_arguments[time.start_name] == 0)
assert (op_arguments[time.end_name] == nt + 1)
# Providing a dimension size explicitly should override the automatically inferred
op_arguments, _ = op.arguments(b=b1, time=nt - 1)
assert (op_arguments[time.start_name] == 0)
assert (op_arguments[time.end_name] == nt - 1)
# Providing a scalar argument explicitly should override the automatically\
# inferred
op_arguments, _ = op.arguments(b=b1, time=nt - 1, time_e=nt - 2)
assert (op_arguments[time.start_name] == 0)
assert (op_arguments[time.end_name] == nt - 2)
def test_expr_like_lowering(self):
"""
Test the lowering of an expr-like ConditionalDimension's condition.
This test makes an Operator that should indexify and lower the condition
passed in the Conditional Dimension
"""
grid = Grid(shape=(3, 3))
g1 = Function(name='g1', grid=grid)
g2 = Function(name='g2', grid=grid)
g1.data[:] = 0.49
g2.data[:] = 0.49
x, y = grid.dimensions
ci = ConditionalDimension(name='ci',
parent=y,
condition=Le((g1 + g2), 1.01 * (g1 + g2)))
f = Function(name='f', shape=grid.shape, dimensions=(x, ci))
Operator(Eq(f, g1 + g2)).apply()
assert np.all(f.data[:] == g1.data[:] + g2.data[:])
def test_default_rules_deriv_offset(self, so, offset):
"""
Test that default_rules generates correct derivatives when derivatives
are evaluated at offset x0.
"""
grid = Grid(shape=(11, ), extent=(10., ))
x = grid.dimensions[0]
h_x = x.spacing
f_std = Function(name='f', grid=grid, space_order=so)
f_sym = Function(name='f',
grid=grid,
space_order=so,
coefficients='symbolic')
g = Function(name='g', grid=grid, space_order=so)
eval_std = str(
Eq(g,
f_std.dx(x0=x + offset * h_x / 2)).evaluate.evalf(_PRECISION))
eval_sym = str(
Eq(g,
f_sym.dx(x0=x + offset * h_x / 2)).evaluate.evalf(_PRECISION))
assert eval_std == eval_sym
def __init__(self,
origin,
spacing,
shape,
space_order,
nbl=20,
dtype=np.float32,
subdomains=(),
damp_mask=False):
self.shape = shape
self.nbl = int(nbl)
self.origin = tuple([dtype(o) for o in origin])
# Origin of the computational domain with boundary to inject/interpolate
# at the correct index
origin_pml = tuple(
[dtype(o - s * nbl) for o, s in zip(origin, spacing)])
phydomain = PhysicalDomain(self.nbl)
subdomains = subdomains + (phydomain, )
shape_pml = np.array(shape) + 2 * self.nbl
# Physical extent is calculated per cell, so shape - 1
extent = tuple(np.array(spacing) * (shape_pml - 1))
self.grid = Grid(extent=extent,
shape=shape_pml,
origin=origin_pml,
dtype=dtype,
subdomains=subdomains)
if self.nbl != 0:
# Create dampening field as symbol `damp`
self.damp = Function(name="damp", grid=self.grid)
initialize_damp(self.damp, self.nbl, self.spacing, mask=damp_mask)
self._physical_parameters = ['damp']
else:
self.damp = 1 if damp_mask else 0
self._physical_parameters = []
def test_operator_leakage_function():
"""
Test to ensure that Operator creation does not cause memory leaks for (Time)Functions.
"""
grid = Grid(shape=(5, 6))
f = Function(name='f', grid=grid)
g = TimeFunction(name='g', grid=grid)
# Take weakrefs to test whether symbols are dead or alive
w_f = weakref.ref(f)
w_g = weakref.ref(g)
# Create operator and delete everything again
op = Operator(Eq(f, 2 * g))
w_op = weakref.ref(op)
del op
del f
del g
clear_cache()
# Test whether things are still hanging around
assert w_f() is None
assert w_g() is None
assert w_op() is None
def GradientOperator(model, source, receiver, space_order=4, save=True,
kernel='OT2', **kwargs):
"""
Constructor method for the gradient operator in an acoustic media
:param model: :class:`Model` object containing the physical parameters
:param source: :class:`PointData` object containing the source geometry
:param receiver: :class:`PointData` object containing the acquisition geometry
:param time_order: Time discretization order
:param space_order: Space discretization order
"""
m, damp = model.m, model.damp
# Gradient symbol and wavefield symbols
grad = Function(name='grad', grid=model.grid)
u = TimeFunction(name='u', grid=model.grid, save=source.nt if save
else None, time_order=2, space_order=space_order)
v = TimeFunction(name='v', grid=model.grid, save=None,
time_order=2, space_order=space_order)
rec = Receiver(name='rec', grid=model.grid,
time_range=receiver.time_range, npoint=receiver.npoint)
s = model.grid.stepping_dim.spacing
eqn = iso_stencil(v, m, s, damp, kernel, forward=False)
if kernel == 'OT2':
gradient_update = Inc(grad, - u.dt2 * v)
elif kernel == 'OT4':
gradient_update = Inc(grad, - (u.dt2 + s**2 / 12.0 * u.laplace2(m**(-2))) * v)
# Add expression for receiver injection
receivers = rec.inject(field=v.backward, expr=rec * s**2 / m,
offset=model.nbpml)
# Substitute spacing terms to reduce flops
return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map,
name='Gradient', **kwargs)
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_function_wo(self):
grid = Grid(shape=(3, 3, 3))
i = Dimension(name='i')
f = Function(name='f', shape=(1, ), dimensions=(i, ), grid=grid)
u = TimeFunction(name='u', grid=grid)
eqns = [Eq(u.forward, u + 1), Eq(f[0], u[0, 0, 0, 0])]
op = Operator(eqns, opt='noop')
assert len(op.body[1].header) == 2
assert len(op.body[1].footer) == 2
assert op.body[1].header[0].value ==\
('omp target enter data map(to: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert str(op.body[1].header[1]) == ''
assert str(op.body[1].footer[0]) == ''
assert op.body[1].footer[1].contents[0].value ==\
('omp target update from(u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert op.body[1].footer[1].contents[1].value ==\
('omp target exit data map(release: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
def test_sdf_distance(self, space_order):
"""Check the mean error in SDF values"""
# Error threshold
thres = 0.11 # 0.11 grid increments
# Sphere of radius 40, center (50, 50, 50)
sphere = 'tests/trial_surfaces/sphere.ply'
# Grid configuration
extent = (100., 100., 100.)
shape = (101, 101, 101)
origin = (0., 0., 0.)
grid = Grid(shape=shape, extent=extent, origin=origin)
x, y, z = grid.dimensions
h_x, h_y, h_z = grid.spacing
f = Function(name='f', grid=grid, space_order=space_order)
# Create signed distance function
sig = SignedDistanceFunction(f, sphere)
pos_x, pos_y, pos_z = np.where(sig.sdf.data != -space_order//2-1)
def true_dist(pos_1, pos_2, pos_3):
return 40 - np.sqrt((h_x*pos_1 - 50)**2
+ (h_y*pos_2 - 50)**2
+ (h_z*pos_3 - 50)**2)
err = np.absolute(sig.sdf.data[pos_x, pos_y, pos_z]
- true_dist(pos_x, pos_y, pos_z))
avg = np.mean(err)
if avg > thres:
err_message = "Mean error in measured distances is {:.6}"
raise ValueError(err_message.format(avg))
def test_precomputed_interpolation():
""" Test interpolation with PrecomputedSparseFunction which accepts
precomputed values for interpolation coefficients
"""
shape = (101, 101)
points = [(.05, .9), (.01, .8), (0.07, 0.84)]
origin = (0, 0)
grid = Grid(shape=shape, origin=origin)
r = 2 # Constant for linear interpolation
# because we interpolate across 2 neighbouring points in each dimension
def init(data):
for i in range(data.shape[0]):
for j in range(data.shape[1]):
data[i,
j] = sin(grid.spacing[0] * i) + sin(grid.spacing[1] * j)
return data
m = Function(name='m', grid=grid, initializer=init, space_order=0)
gridpoints, interpolation_coeffs = precompute_linear_interpolation(
points, grid, origin)
sf = PrecomputedSparseFunction(name='s',
grid=grid,
r=r,
npoint=len(points),
gridpoints=gridpoints,
interpolation_coeffs=interpolation_coeffs)
eqn = sf.interpolate(m)
op = Operator(eqn)
op()
expected_values = [sin(point[0]) + sin(point[1]) for point in points]
assert (all(np.isclose(sf.data, expected_values, rtol=1e-6)))
def test_array_rw(self):
grid = Grid(shape=(3, 3, 3))
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid, space_order=2)
eqn = Eq(u.forward, u * cos(f * 2))
op = Operator(eqn)
assert len(op.body[1].header) == 7
assert str(op.body[1].header[0]) == 'float (*r1)[y_size][z_size];'
assert op.body[1].header[1].text ==\
'posix_memalign((void**)&r1, 64, sizeof(float[x_size][y_size][z_size]))'
assert op.body[1].header[2].value ==\
('omp target enter data map(alloc: r1[0:x_size][0:y_size][0:z_size])'
'')
assert len(op.body[1].footer) == 6
assert str(op.body[1].footer[0]) == ''
assert op.body[1].footer[1].value ==\
('omp target exit data map(delete: r1[0:x_size][0:y_size][0:z_size])'
' if((x_size != 0) && (y_size != 0) && (z_size != 0))')
assert op.body[1].footer[2].text == 'free(r1)'
def test_override_function_size(self):
"""
Test runtime size overrides for :class:`Function` dimensions.
Note: The current behaviour for size-only arguments seems
ambiguous (eg. op(x=3, y=4), as it sets `dim_size` as well as
`dim_end`. Since `dim_size` is used for the cast, we can get
garbage results if it does not agree with the shape of the
provided data. This should error out, or potentially we could
set the corresponding size, while aliasing `dim` to `dim_e`?
The same should be tested for :class:`TimeFunction` once fixed.
"""
grid = Grid(shape=(5, 6, 7))
g = Function(name='g', grid=grid)
op = Operator(Eq(g, 1.))
args = {'x': 3, 'y': 4, 'z': 5}
arguments = op.arguments(**args)
expected = {
'x_size': 5,
'x_s': 0,
'x_e': 3,
'y_size': 6,
'y_s': 0,
'y_e': 4,
'z_size': 7,
'z_s': 0,
'z_e': 5,
'g': g.data
}
self.verify_arguments(arguments, expected)
# Verify execution
op(**args)
assert (g.data[3:, 4:, 5:] == 0.).all()
assert (g.data[:3, :4, :5] == 1.).all()
def test_tasking_fused(self):
nt = 10
bundle0 = Bundle()
grid = Grid(shape=(10, 10, 10), subdomains=bundle0)
tmp = Function(name='tmp', grid=grid)
u = TimeFunction(name='u', grid=grid, save=nt)
v = TimeFunction(name='v', grid=grid, save=nt)
w = TimeFunction(name='w', grid=grid)
eqns = [Eq(w.forward, w + 1),
Eq(tmp, w.forward),
Eq(u.forward, tmp, subdomain=bundle0),
Eq(v.forward, tmp, subdomain=bundle0)]
op = Operator(eqns, opt=('tasking', 'fuse', 'orchestrate'))
# Check generated code
assert len(retrieve_iteration_tree(op)) == 4
locks = [i for i in FindSymbols().visit(op) if isinstance(i, Lock)]
assert len(locks) == 1 # Only 1 because it's only `tmp` that needs protection
assert len(op._func_table) == 2
exprs = FindNodes(Expression).visit(op._func_table['copy_device_to_host0'].root)
assert len(exprs) == 19
assert str(exprs[12]) == 'int id = sdata0->id;'
assert str(exprs[13]) == 'const int time = sdata0->time;'
assert str(exprs[14]) == 'lock0[0] = 1;'
assert exprs[15].write is u
assert exprs[16].write is v
assert str(exprs[17]) == 'lock0[0] = 2;'
assert str(exprs[18]) == 'sdata0->flag = 1;'
op.apply(time_M=nt-2)
assert np.all(u.data[nt-1] == 9)
assert np.all(v.data[nt-1] == 9)
def _initialize_bcs(self, bcs="damp"):
# Create dampening field as symbol `damp`
if self.nbl == 0:
self.damp = 1 if bcs == "mask" else 0
return
# First initialization
init = self.damp is None
# Get current Function if alread yinitialized
self.damp = self.damp or Function(name="damp", grid=self.grid)
if callable(bcs):
bcs(self.damp, self.nbl)
else:
re_init = ((bcs == "mask" and mmin(self.damp) == 0) or
(bcs == "damp" and mmax(self.damp) == 1))
if init or re_init:
if re_init and not init:
bcs_o = "damp" if bcs == "mask" else "mask"
warning("Re-initializing damp profile from %s to %s" % (bcs_o, bcs))
warning("Model has to be created with `bcs=\"%s\"`"
"for this WaveSolver" % bcs)
initialize_damp(self.damp, self.padsizes, self.spacing,
abc_type=bcs, fs=self.fs)
self._physical_parameters.update(['damp'])
def test_multiple_subnests(self):
grid = Grid(shape=(3, 3, 3))
x, y, z = grid.dimensions
t = grid.stepping_dim
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid)
eqn = Eq(u.forward, ((u[t, x, y, z] + u[t, x+1, y+1, z+1])*3*f +
(u[t, x+2, y+2, z+2] + u[t, x+3, y+3, z+3])*3*f + 1))
op = Operator(eqn, dse='aggressive', dle=('advanced', {'openmp': True}))
trees = retrieve_iteration_tree(op._func_table['bf0'].root)
assert len(trees) == 2
assert trees[0][0] is trees[1][0]
assert trees[0][0].pragmas[0].value ==\
'omp for collapse(1) schedule(dynamic,1)'
assert trees[0][2].pragmas[0].value == ('omp parallel for collapse(1) '
'schedule(dynamic,1) '
'num_threads(nthreads_nested)')
assert trees[1][2].pragmas[0].value == ('omp parallel for collapse(1) '
'schedule(dynamic,1) '
'num_threads(nthreads_nested)')
def test_misc_dims(self):
"""
Tests grid-independent :class:`Function`s, which require YASK's "misc"
dimensions.
"""
dx = Dimension(name='dx')
grid = Grid(shape=(10, 10))
x, y = grid.dimensions
time = grid.time_dim
u = TimeFunction(name='u', grid=grid, time_order=1, space_order=4, save=4)
c = Function(name='c', dimensions=(x, dx), shape=(10, 5))
step = Eq(u.forward, (
u[time, x-2, y] * c[x, 0]
+ u[time, x-1, y] * c[x, 1]
+ u[time, x, y] * c[x, 2]
+ u[time, x+1, y] * c[x, 3]
+ u[time, x+2, y] * c[x, 4]))
for i in range(10):
c.data[i, 0] = 1.0+i
c.data[i, 1] = 1.0+i
c.data[i, 2] = 3.0+i
c.data[i, 3] = 6.0+i
c.data[i, 4] = 5.0+i
u.data[:] = 0.0
u.data[0, 2, :] = 2.0
op = Operator(step)
assert 'run_solution' in str(op)
op(time_m=0, time_M=0)
assert(np.all(u.data[1, 0, :] == 10.0))
assert(np.all(u.data[1, 1, :] == 14.0))
assert(np.all(u.data[1, 2, :] == 10.0))
assert(np.all(u.data[1, 3, :] == 8.0))
assert(np.all(u.data[1, 4, :] == 10.0))
assert(np.all(u.data[1, 5:10, :] == 0.0))
def test_misc_data(self):
"""
Test data insertion/indexing for Functions with mixed
distributed/replicated Dimensions.
"""
dx = Dimension(name='dx')
grid = Grid(shape=(4, 4))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
# Note: `grid` must be passed to `c` since `x` is a distributed dimension,
# and `grid` carries the `x` decomposition
c = Function(name='c', grid=grid, dimensions=(x, dx), shape=(4, 5))
# Data insertion
for i in range(4):
c.data[i, 0] = 1.0+i
c.data[i, 1] = 1.0+i
c.data[i, 2] = 3.0+i
c.data[i, 3] = 6.0+i
c.data[i, 4] = 5.0+i
# Data indexing
if LEFT in glb_pos_map[x]:
assert(np.all(c.data[0] == [1., 1., 3., 6., 5.]))
assert(np.all(c.data[1] == [2., 2., 4., 7., 6.]))
else:
assert(np.all(c.data[2] == [3., 3., 5., 8., 7.]))
assert(np.all(c.data[3] == [4., 4., 6., 9., 8.]))
# Same as before, but with negative indices and non-trivial slices
if LEFT in glb_pos_map[x]:
assert(np.all(c.data[0:-3] == [1., 1., 3., 6., 5.]))
assert(np.all(c.data[-3:-2] == [2., 2., 4., 7., 6.]))
else:
assert(np.all(c.data[-2:-1] == [3., 3., 5., 8., 7.]))
assert(np.all(c.data[-1] == [4., 4., 6., 9., 8.]))
def test_at_is_actually_working(shape, expected):
"""
Check that autotuning is actually running when switched on,
in both 2D and 3D operators.
"""
grid = Grid(shape=shape)
infield = Function(name='infield', grid=grid)
infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
outfield = Function(name='outfield', grid=grid)
stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0)
op = Operator(stencil, dle=('blocking', {'openmp': False,
'blockinner': True,
'blockalways': True}))
# Run with whatever `configuration` says (by default, basic+preemptive)
op(infield=infield, outfield=outfield, autotune=True)
assert op._state['autotuning'][-1]['runs'] == 4
assert op._state['autotuning'][-1]['tpr'] == 1
# Now try `aggressive` autotuning
configuration['autotuning'] = 'aggressive'
op(infield=infield, outfield=outfield, autotune=True)
assert op._state['autotuning'][-1]['runs'] == expected
assert op._state['autotuning'][-1]['tpr'] == 1
configuration['autotuning'] = configuration._defaults['autotuning']
# Try again, but using the Operator API directly
op(infield=infield, outfield=outfield, autotune='aggressive')
assert op._state['autotuning'][-1]['runs'] == expected
assert op._state['autotuning'][-1]['tpr'] == 1
# Similar to above
op(infield=infield, outfield=outfield, autotune=('aggressive', 'preemptive'))
assert op._state['autotuning'][-1]['runs'] == expected
assert op._state['autotuning'][-1]['tpr'] == 1
def test_simple_indexing(self):
"""Test data packing/unpacking via basic indexing."""
grid = Grid(shape=(16, 16, 16))
u = Function(name='yu3D', grid=grid, space_order=0)
# Test simple insertion and extraction
u.data[0, 1, 1] = 1.
assert u.data[0, 0, 0] == 0.
assert u.data[0, 1, 1] == 1.
assert np.all(u.data == u.data[:, :, :])
assert 1. in u.data[0]
assert 1. in u.data[0, 1]
# Test negative indices
assert u.data[0, -15, -15] == 1.
u.data[6, 0, 0] = 1.
assert u.data[-10, :, :].sum() == 1.
# Test setting whole array to given value
u.data[:] = 3.
assert np.all(u.data == 3.)
# Test insertion of single value into block
u.data[5, :, 5] = 5.
assert np.all(u.data[5, :, 5] == 5.)
# Test extraction of block with negative indices
sliced = u.data[-11, :, -11]
assert sliced.shape == (16,)
assert np.all(sliced == 5.)
# Test insertion of block into block
block = np.ndarray(shape=(1, 16, 1), dtype=np.float32)
block.fill(4.)
u.data[4:5, :, 4:5] = block
assert np.all(u.data[4, :, 4] == block)
def test_timesteps_per_at_run():
"""
Check that each autotuning run (ie with a given block shape) takes
``autotuning.core.options['squeezer']`` timesteps, for an operator
performing the increment ``a[t + timeorder, ...] = f(a[t, ...], ...)``.
"""
shape = (30, 30, 30)
grid = Grid(shape=shape)
x, y, z = grid.dimensions
t = grid.stepping_dim
# Function
infield = Function(name='infield', grid=grid)
infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
outfield = Function(name='outfield', grid=grid)
stencil = Eq(outfield.indexify(), outfield.indexify() + infield.indexify()*3.0)
op = Operator(stencil, dle=('blocking', {'openmp': False, 'blockalways': True}))
op(infield=infield, outfield=outfield, autotune=True)
assert op._state['autotuning'][-1]['runs'] == 4
assert op._state['autotuning'][-1]['tpr'] == 1
# TimeFunction with increasing time order; increasing the time order
# shouldn't affect how many iterations the autotuner is gonna run
for to in [1, 2, 4]:
infield = TimeFunction(name='infield', grid=grid, time_order=to)
infield.data[:] = np.arange(reduce(mul, infield.shape),
dtype=np.int32).reshape(infield.shape)
outfield = TimeFunction(name='outfield', grid=grid, time_order=to)
stencil = Eq(outfield[t + to, x, y, z],
outfield.indexify() + infield.indexify()*3.0)
op = Operator(stencil, dle=('blocking', {'openmp': False, 'blockalways': True}))
op(infield=infield, outfield=outfield, time=20, autotune=True)
assert op._state['autotuning'][-1]['runs'] == 4
assert op._state['autotuning'][-1]['tpr'] == options['squeezer'] + 1
def test_index_alignment(const):
""" A much simpler test meant to ensure that the forward and reverse indices are
correctly aligned (i.e. u * v , where u is the forward field and v the reverse field
corresponds to the correct timesteps in u and v). The forward operator does u = u + 1
which means that the field a will be equal to nt (0 -> 1 -> 2 -> 3), the number of
timesteps this operator is run for. The field at the last time step of the forward is
used to initialise the field v for the reverse pass. The reverse operator does
v = v - 1, which means that if the reverse operator is run for the same number of
timesteps as the forward operator, v should be 0 at the last time step
(3 -> 2 -> 1 -> 0). There is also a grad = grad + u * v accumulator in the reverse
operator. If the alignment is correct, u and v should have the same value at every
time step:
0 -> 1 -> 2 -> 3 u
0 <- 1 <- 2 <- 3 v
and hence grad = 0*0 + 1*1 + 2*2 + 3*3 = sum(n^2) where n -> [0, nt]
If the test fails, the resulting number can tell you how the fields are misaligned
"""
n = 4
grid = Grid(shape=(2, 2))
order_of_eqn = 1
modulo_factor = order_of_eqn + 1
nt = n - order_of_eqn
u = TimeFunction(name='u', grid=grid, save=n)
# Increment one in the forward pass 0 -> 1 -> 2 -> 3
fwd_op = Operator(Eq(u.forward, u + 1.*const))
# Invocation 1
fwd_op(time=nt-1, constant=1)
last_time_step_v = nt % modulo_factor
# Last time step should be equal to the number of timesteps we ran
assert(np.allclose(u.data[nt, :, :], nt))
v = TimeFunction(name='v', grid=grid, save=None)
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Decrement one in the reverse pass 3 -> 2 -> 1 -> 0
adj_eqn = Eq(v, v.forward - 1.*const)
adj_op = Operator(adj_eqn)
# Invocation 2
adj_op(time=nt-1, constant=1)
# Last time step should be back to 0
assert(np.allclose(v.data[0, :, :], 0))
# Reset v to run the backward again
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
prod = Function(name="prod", grid=grid)
# Multiply u and v and add them
# = 3*3 + 2*2 + 1*1 + 0*0
prod_eqn = Eq(prod, prod + u * v)
comb_op = Operator([adj_eqn, prod_eqn])
# Invocation 3
comb_op(time=nt-1, constant=1)
final_value = sum([n**2 for n in range(nt)])
# Final value should be sum of squares of first nt natural numbers
assert(np.allclose(prod.data, final_value))
# Now reset to repeat all the above tests with checkpointing
prod.data[:] = 0
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Checkpointed version doesn't require to save u
u_nosave = TimeFunction(name='u_n', grid=grid)
# change equations to use new symbols
fwd_eqn_2 = Eq(u_nosave.forward, u_nosave + 1.*const)
fwd_op_2 = Operator(fwd_eqn_2)
cp = DevitoCheckpoint([u_nosave])
wrap_fw = CheckpointOperator(fwd_op_2, constant=1)
prod_eqn_2 = Eq(prod, prod + u_nosave * v)
comb_op_2 = Operator([adj_eqn, prod_eqn_2])
wrap_rev = CheckpointOperator(comb_op_2, constant=1)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, nt)
# Invocation 4
wrp.apply_forward()
assert(np.allclose(u_nosave.data[last_time_step_v, :, :], nt))
# Invocation 5
wrp.apply_reverse()
assert(np.allclose(v.data[0, :, :], 0))
assert(np.allclose(prod.data, final_value))
def test_indexing_in_views(self):
grid = Grid(shape=(4, 4))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
myrank = grid.distributor.myrank
u = Function(name='u', grid=grid, space_order=0)
u.data[:] = myrank
# Note that the `1`s are global indices
view = u.data[1:, 1:]
assert np.all(view[:] == myrank)
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert view.shape == (1, 1)
assert np.all(view == 0.)
assert view[0, 0] == 0.
assert view[1, 1] is None
assert view[1].shape == (0, 1)
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert view.shape == (1, 2)
assert np.all(view == 1.)
assert view[0, 0] is None
assert view[1, 1] is None
assert view[1].shape == (0, 2)
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert view.shape == (2, 1)
assert np.all(view == 2.)
assert view[0, 0] is None
assert view[1, 1] is None
assert view[1].shape == (1,)
assert np.all(view[1] == 2.)
else:
assert view.shape == (2, 2)
assert np.all(view == 3.)
assert view[0, 0] is None
assert view[1, 1] == 3.
assert view[1].shape == (2,)
assert np.all(view[1] == 3.)
# Now we further slice into `view`
view2 = view[1:, 1:]
assert np.all(view2[:] == myrank)
if LEFT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert view2.shape == (0, 0)
elif LEFT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert view2.shape == (0, 2)
elif RIGHT in glb_pos_map[x] and LEFT in glb_pos_map[y]:
assert view2.shape == (2, 0)
else:
assert view2.shape == (2, 2)
# Now a change in `view2` by the only rank that "sees" it should affect
# both `view` and `u.data`
view2[:] += 1
if RIGHT in glb_pos_map[x] and RIGHT in glb_pos_map[y]:
assert np.all(u.data[:] == myrank + 1)
assert np.all(view[:] == myrank + 1)
assert np.all(view2[:] == myrank + 1)
else:
assert np.all(view[:] == myrank)
assert np.all(view2[:] == myrank)
assert view2.size == 0
