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 test_special_symbols(self):
"""
This test checks the singletonization, through the caching infrastructure,
of the special symbols that an Operator may generate (e.g., `nthreads`).
"""
grid = Grid(shape=(4, 4, 4))
f = TimeFunction(name='f', grid=grid)
sf = SparseTimeFunction(name='sf', grid=grid, npoint=1, nt=10)
eqns = [Eq(f.forward, f + 1.)] + sf.inject(field=f.forward, expr=sf)
opt = ('advanced', {'par-nested': 0, 'openmp': True})
op0 = Operator(eqns, opt=opt)
op1 = Operator(eqns, opt=opt)
nthreads0, nthreads_nested0, nthreads_nonaffine0 =\
[i for i in op0.input if isinstance(i, NThreadsBase)]
nthreads1, nthreads_nested1, nthreads_nonaffine1 =\
[i for i in op1.input if isinstance(i, NThreadsBase)]
assert nthreads0 is nthreads1
assert nthreads_nested0 is nthreads_nested1
assert nthreads_nonaffine0 is nthreads_nonaffine1
tid0 = ThreadID(op0.nthreads)
tid1 = ThreadID(op0.nthreads)
assert tid0 is tid1
did0 = DeviceID()
did1 = DeviceID()
assert did0 is did1
npt0 = NPThreads(name='npt', size=3)
npt1 = NPThreads(name='npt', size=3)
npt2 = NPThreads(name='npt', size=4)
assert npt0 is npt1
assert npt0 is not npt2
def test_bcs_basic(self):
"""
Test MPI in presence of boundary condition loops. Here, no halo exchange
is expected (as there is no stencil in the computed expression) but we
check that:
* the left BC loop is computed by the leftmost rank only
* the right BC loop is computed by the rightmost rank only
"""
grid = Grid(shape=(20,))
x = grid.dimensions[0]
t = grid.stepping_dim
thickness = 4
u = TimeFunction(name='u', grid=grid, time_order=1)
xleft = SubDimension.left(name='xleft', parent=x, thickness=thickness)
xi = SubDimension.middle(name='xi', parent=x,
thickness_left=thickness, thickness_right=thickness)
xright = SubDimension.right(name='xright', parent=x, thickness=thickness)
t_in_centre = Eq(u[t+1, xi], 1)
leftbc = Eq(u[t+1, xleft], u[t+1, xleft+1] + 1)
rightbc = Eq(u[t+1, xright], u[t+1, xright-1] + 1)
op = Operator([t_in_centre, leftbc, rightbc])
op.apply(time_m=1, time_M=1)
glb_pos_map = u.grid.distributor.glb_pos_map
if LEFT in glb_pos_map[x]:
assert np.all(u.data_ro_domain[0, thickness:] == 1.)
assert np.all(u.data_ro_domain[0, :thickness] == range(thickness+1, 1, -1))
else:
assert np.all(u.data_ro_domain[0, :-thickness] == 1.)
assert np.all(u.data_ro_domain[0, -thickness:] == range(2, thickness+2))
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 adjoint(self, rec, srca=None, v=None, m=None, **kwargs):
"""
Adjoint modelling function that creates the necessary
data objects for running an adjoint modelling operator.
:param rec: Symbol with stored receiver data. Please note that
these act as the source term in the adjoint run.
:param srca: Symbol to store the resulting data for the
interpolated at the original source location.
:param v: (Optional) Symbol to store the computed wavefield
:param m: (Optional) Symbol for the time-constant square slowness
:returns: Adjoint source, wavefield and performance summary
"""
# Create a new adjoint source and receiver symbol
srca = srca or PointSource(name='srca',
grid=self.model.grid,
time_range=self.geometry.time_axis,
coordinates=self.geometry.src_positions)
# Create the adjoint wavefield if not provided
v = v or TimeFunction(name='v',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
# Pick m from model unless explicitly provided
m = m or self.model.m
# Execute operator and return wavefield and receiver data
summary = self.op_adj().apply(srca=srca,
rec=rec,
v=v,
m=m,
dt=kwargs.pop('dt', self.dt),
**kwargs)
return srca, v, summary
def random_walks1D_vec(x0, N, p, num_walks=1, num_times=1, random=random):
"""Vectorized version of random_walks1D."""
position = np.zeros(N + 1) # Accumulated positions
position2 = np.zeros(N + 1) # Accumulated positions**2
# Histogram at num_times selected time points
pos_hist = np.zeros((num_walks, num_times))
pos_hist_times = [(N // num_times) * i for i in range(num_times)]
# Create and initialise our TimeFunction r
x_d = Dimension(name='x_d')
t_d = Dimension(name='t_d')
r = TimeFunction(name='r', dimensions=(x_d, t_d), shape=(N + 1, num_walks))
# Setting each walk's first element to x0
r.data[0, :] = x0
steps = Function(name='steps',
dimensions=(x_d, t_d),
shape=(N + 1, num_walks))
# Populating steps with -1 if value in rs <= p at that point and 1 otherwise
rs = random.uniform(0, 1, size=N * num_walks).reshape(num_walks, N)
for n in range(num_walks):
steps.data[:N, n] = np.where(rs[n] <= p, -1, 1)
# Creating and applying operator that contains equation for adding steps
eq = Eq(r.forward, r + steps)
op = Operator(eq)
op.apply()
# Summing over data to give positions
position = np.sum(r.data, axis=1) # Accumulated positions
position2 = np.sum(r.data**2, axis=1) # Accumulated positions**2
pos_hist[:, :] = np.transpose(r.data[pos_hist_times, :])
return position, position2, pos_hist, np.array(pos_hist_times)
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_precomputed_interpolation_time():
""" Test interpolation with PrecomputedSparseFunction which accepts
precomputed values for interpolation coefficients, but this time
with a TimeFunction
"""
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
u = TimeFunction(name='u', grid=grid, space_order=0, save=5)
for it in range(5):
u.data[it, :] = it
gridpoints, interpolation_coeffs = precompute_linear_interpolation(
points, grid, origin)
sf = PrecomputedSparseTimeFunction(
name='s',
grid=grid,
r=r,
npoint=len(points),
nt=5,
gridpoints=gridpoints,
interpolation_coeffs=interpolation_coeffs)
assert sf.data.shape == (5, 3)
eqn = sf.interpolate(u)
op = Operator(eqn)
op(time_m=0, time_M=4)
for it in range(5):
assert np.allclose(sf.data[it, :], it)
def test_contracted_shape_after_blocking(self):
"""
Like `test_full_alias_shape_after_blocking`, but a different
Operator is used, leading to contracted Arrays (2D instead of 3D).
"""
grid = Grid(shape=(3, 3, 3))
x, y, z = grid.dimensions # noqa
t = grid.stepping_dim
f = Function(name='f', grid=grid)
f.data_with_halo[:] = 1.
u = TimeFunction(name='u', grid=grid, space_order=3)
u.data_with_halo[:] = 0.
# Leads to 2D aliases
eqn = Eq(u.forward, ((u[t, x, y, z] + u[t, x, y+1, z+1])*3*f +
(u[t, x, y+2, z+2] + u[t, x, y+3, z+3])*3*f + 1))
op0 = Operator(eqn, dse='noop', dle=('advanced', {'openmp': True}))
op1 = Operator(eqn, dse='aggressive', dle=('advanced', {'openmp': True}))
y0_blk_size = op1.parameters[-2]
z_size = op1.parameters[3]
arrays = [i for i in FindSymbols().visit(op1._func_table['bf0'].root)
if i.is_Array]
assert len(arrays) == 1
a = arrays[0]
assert len(a.dimensions) == 2
assert a.halo == ((1, 1), (1, 1))
assert Add(*a.symbolic_shape[0].args) == y0_blk_size + 2
assert Add(*a.symbolic_shape[1].args) == z_size + 2
# Check numerical output
op0(time_M=1)
exp = np.copy(u.data[:])
u.data_with_halo[:] = 0.
op1(time_M=1)
assert np.all(u.data == exp)
def wavefield_subsampled(model, u, nt, t_sub, space_order=8):
"""
Create a subsampled wavefield
Parameters
----------
model : Model
Physical model
u : TimeFunction
Forward wavefield for modeling
nt : int
Number of time steps on original time axis
t_sub : int
Factor for time-subsampling
space_order: int
Spatial discretization order
"""
wf_s = []
eq_save = []
if t_sub > 1:
time_subsampled = ConditionalDimension(name='t_sub',
parent=model.grid.time_dim,
factor=t_sub)
nsave = (nt - 1) // t_sub + 2
else:
return None, []
for wf in as_tuple(u):
usave = TimeFunction(name='us_%s' % wf.name,
grid=model.grid,
time_order=2,
space_order=space_order,
time_dim=time_subsampled,
save=nsave)
wf_s.append(usave)
eq_save.append(Eq(usave, wf))
return wf_s, eq_save
def test_cache_blocking_structure_subdims():
"""
Test that:
* With local SubDimensions no-blocking is expected.
* With non-local SubDimensions, blocking is expected.
"""
grid = Grid(shape=(4, 4, 4))
x, y, z = grid.dimensions
xi, yi, zi = grid.interior.dimensions
t = grid.stepping_dim
xl = SubDimension.left(name='xl', parent=x, thickness=4)
f = TimeFunction(name='f', grid=grid)
assert xl.local
# Local SubDimension -> no blocking expected
op = Operator(Eq(f[t + 1, xl, y, z], f[t, xl, y, z] + 1))
assert len(op._func_table) == 0
# Non-local SubDimension -> blocking expected
op = Operator(Eq(f.forward, f + 1, subdomain=grid.interior))
trees = retrieve_iteration_tree(op._func_table['bf0'].root)
assert len(trees) == 1
tree = trees[0]
assert len(tree) == 5
assert tree[
0].dim.is_Incr and tree[0].dim.parent is xi and tree[0].dim.root is x
assert tree[
1].dim.is_Incr and tree[1].dim.parent is yi and tree[1].dim.root is y
assert tree[2].dim.is_Incr and tree[2].dim.parent is tree[0].dim and\
tree[2].dim.root is x
assert tree[3].dim.is_Incr and tree[3].dim.parent is tree[1].dim and\
tree[3].dim.root is y
assert not tree[
4].dim.is_Incr and tree[4].dim is zi and tree[4].dim.parent is z
def test_dynamic_nthreads(self):
grid = Grid(shape=(16, 16, 16))
f = TimeFunction(name='f', grid=grid)
sf = SparseTimeFunction(name='sf', grid=grid, npoint=1, nt=5)
eqns = [Eq(f.forward, f + 1)]
eqns += sf.interpolate(f)
op = Operator(eqns, dle='openmp')
parregions = FindNodes(ParallelRegion).visit(op)
assert len(parregions) == 2
# Check suitable `num_threads` appear in the generated code
# Not very elegant, but it does the trick
assert 'num_threads(nthreads)' in str(parregions[0].header[0])
assert 'num_threads(nthreads_nonaffine)' in str(
parregions[1].header[0])
# Check `op` accepts the `nthreads*` kwargs
op.apply(time=0)
op.apply(time_m=1, time_M=1, nthreads=4)
op.apply(time_m=1, time_M=1, nthreads=4, nthreads_nonaffine=2)
op.apply(time_m=1, time_M=1, nthreads_nonaffine=2)
assert np.all(f.data[0] == 2.)
# Check the actual value assumed by `nthreads` and `nthreads_nonaffine`
assert op.arguments(time=0)['nthreads'] == NThreads.default_value()
assert op.arguments(time=0)['nthreads_nonaffine'] == \
NThreadsNonaffine.default_value()
# Again, but with user-supplied values
assert op.arguments(time=0, nthreads=123)['nthreads'] == 123
assert op.arguments(
time=0, nthreads_nonaffine=100)['nthreads_nonaffine'] == 100
# Again, but with the aliases
assert op.arguments(time=0, nthreads0=123)['nthreads'] == 123
assert op.arguments(time=0, nthreads2=123)['nthreads_nonaffine'] == 123
def adjoint_wf_src(model, u, src_coords, space_order=8):
"""
Adjoint/backward modeling of a full wavefield (full wavefield as adjoint source)
Ps*F^T*u.
Parameters
----------
model: Model
Physical model
u: Array or TimeFunction
Time-space dependent source
src_coords: Array
Source coordinates
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Shot record (sampled at source position(s))
"""
wf_src = TimeFunction(name='wf_src',
grid=model.grid,
time_order=2,
space_order=space_order,
save=u.shape[0])
if isinstance(u, TimeFunction):
wf_src._data = u._data
else:
wf_src.data[:] = u[:]
rec, _, _ = adjoint(model,
None,
src_coords,
None,
space_order=space_order,
q=wf_src)
return rec.data
def adjoint_wf_src_norec(model, u, space_order=8):
"""
Adjoint/backward modeling of a full wavefield (full wavefield as adjoint source)
F^T*u.
Parameters
----------
model: Model
Physical model
u: Array or TimeFunction
Time-space dependent source
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Adjoint wavefield
"""
wf_src = TimeFunction(name='wf_src',
grid=model.grid,
time_order=2,
space_order=space_order,
save=u.shape[0])
if isinstance(u, TimeFunction):
wf_src._data = u._data
else:
wf_src.data[:] = u[:]
_, v, _ = adjoint(model,
None,
None,
None,
space_order=space_order,
save=True,
q=wf_src)
return v.data
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_bcs(self):
"""
Tests application of an Operator consisting of multiple equations
defined over different sub-regions, explicitly created through the
use of :class:`SubDimension`s.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
t = grid.stepping_dim
thickness = 4
u = TimeFunction(name='u', save=None, grid=grid, space_order=0, time_order=1)
xleft = SubDimension.left(name='xleft', parent=x, thickness=thickness)
xi = SubDimension.middle(name='xi', parent=x,
thickness_left=thickness, thickness_right=thickness)
xright = SubDimension.right(name='xright', parent=x, thickness=thickness)
yi = SubDimension.middle(name='yi', parent=y,
thickness_left=thickness, thickness_right=thickness)
t_in_centre = Eq(u[t+1, xi, yi], 1)
leftbc = Eq(u[t+1, xleft, yi], u[t+1, xleft+1, yi] + 1)
rightbc = Eq(u[t+1, xright, yi], u[t+1, xright-1, yi] + 1)
op = Operator([t_in_centre, leftbc, rightbc])
op.apply(time_m=1, time_M=1)
assert np.all(u.data[0, :, 0:thickness] == 0.)
assert np.all(u.data[0, :, -thickness:] == 0.)
assert all(np.all(u.data[0, i, thickness:-thickness] == (thickness+1-i))
for i in range(thickness))
assert all(np.all(u.data[0, -i, thickness:-thickness] == (thickness+2-i))
for i in range(1, thickness + 1))
assert np.all(u.data[0, thickness:-thickness, thickness:-thickness] == 1.)
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_subdimleft_notparallel(self):
"""
Tests application of an Operator consisting of a subdimension
defined over different sub-regions, explicitly created through the
use of :class:`SubDimension`s.
This tests that flow direction is not being automatically inferred
from whether the subdimension is on the left or right boundary.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
t = grid.stepping_dim
thickness = 4
u = TimeFunction(name='u', save=None, grid=grid, space_order=1, time_order=0)
xl = SubDimension.left(name='xl', parent=x, thickness=thickness)
yi = SubDimension.middle(name='yi', parent=y,
thickness_left=thickness, thickness_right=thickness)
# Flows inward (i.e. forward) rather than outward
eq = Eq(u[t+1, xl, yi], u[t+1, xl-1, yi] + 1)
op = Operator([eq])
iterations = FindNodes(Iteration).visit(op)
assert all(i.is_Affine and i.is_Sequential for i in iterations if i.dim == xl)
assert all(i.is_Affine and i.is_Parallel for i in iterations if i.dim == yi)
op.apply(time_m=0, time_M=0)
assert all(np.all(u.data[0, :thickness, thickness+i] == [1, 2, 3, 4])
for i in range(12))
assert np.all(u.data[0, thickness:] == 0)
assert np.all(u.data[0, :, thickness+12:] == 0)
def ForwardOperator(model, source, receiver, space_order=4,
save=False, kernel='OT2', **kwargs):
"""
Constructor method for the forward modelling 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
:param save: Saving flag, True saves all time steps, False only the three
"""
m, damp = model.m, model.damp
# Create symbols for forward wavefield, source and receivers
u = TimeFunction(name='u', grid=model.grid,
save=source.nt if save else None,
time_order=2, space_order=space_order)
src = PointSource(name='src', grid=model.grid, ntime=source.nt,
npoint=source.npoint)
rec = Receiver(name='rec', grid=model.grid, ntime=receiver.nt,
npoint=receiver.npoint)
s = model.grid.stepping_dim.spacing
eqn = iso_stencil(u, m, s, damp, kernel)
# Construct expression to inject source values
src_term = src.inject(field=u.forward, expr=src * s**2 / m,
offset=model.nbpml)
# Create interpolation expression for receivers
rec_term = rec.interpolate(expr=u, offset=model.nbpml)
# Substitute spacing terms to reduce flops
return Operator(eqn + src_term + rec_term, subs=model.spacing_map,
name='Forward', **kwargs)
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_iteration_parallelism_3d(self, exprs, atomic, parallel):
"""Tests detection of PARALLEL_* properties."""
grid = Grid(shape=(10, 10, 10))
time = grid.time_dim # noqa
t = grid.stepping_dim # noqa
x, y, z = grid.dimensions # noqa
p = Dimension(name='p')
d = Dimension(name='d')
rx = Dimension(name='rx')
ry = Dimension(name='ry')
rz = Dimension(name='rz')
u = Function(name='u', grid=grid) # noqa
v = TimeFunction(name='v', grid=grid, save=None) # noqa
cx = Function(name='coeff_x', dimensions=(p, rx), shape=(1, 2)) # noqa
cy = Function(name='coeff_y', dimensions=(p, ry), shape=(1, 2)) # noqa
cz = Function(name='coeff_z', dimensions=(p, rz), shape=(1, 2)) # noqa
gp = Function(name='gridpoints', dimensions=(p, d), shape=(1, 3)) # noqa
src = Function(name='src', dimensions=(p,), shape=(1,)) # noqa
rcv = Function(name='rcv', dimensions=(time, p), shape=(100, 1), space_order=0) # noqa
# List comprehension would need explicit locals/globals mappings to eval
for i, e in enumerate(list(exprs)):
exprs[i] = eval(e)
# Use 'openmp' here instead of 'advanced' to disable loop blocking
op = Operator(exprs, dle='openmp')
iters = FindNodes(Iteration).visit(op)
assert all([i.is_ParallelAtomic for i in iters if i.dim.name in atomic])
assert all([not i.is_ParallelAtomic for i in iters if i.dim.name not in atomic])
assert all([i.is_Parallel for i in iters if i.dim.name in parallel])
assert all([not i.is_Parallel for i in iters if i.dim.name not in parallel])
def forward_wf_src(model, u, rec_coords, space_order=8):
"""
Forward modeling of a full wavefield source Pr*F*u.
Parameters
----------
model: Model
Physical model
u: TimeFunction or Array
Time-space dependent wavefield
rec_coords: Array
Coordiantes of the receiver(s)
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Shot record
"""
wf_src = TimeFunction(name='wf_src',
grid=model.grid,
time_order=2,
space_order=space_order,
save=u.shape[0])
if isinstance(u, TimeFunction):
wf_src._data = u._data
else:
wf_src.data[:] = u[:]
rec, _, _ = forward(model,
None,
rec_coords,
None,
space_order=space_order,
q=wf_src)
return rec.data
def forward(self, src=None, rec=None, u=None, m=None, save=None, **kwargs):
"""
Forward modelling function that creates the necessary
data objects for running a forward modelling operator.
:param src: Symbol with time series data for the injected source term
:param rec: Symbol to store interpolated receiver data
:param u: (Optional) Symbol to store the computed wavefield
:param m: (Optional) Symbol for the time-constant square slowness
:param save: Option to store the entire (unrolled) wavefield
:returns: Receiver, wavefield and performance summary
"""
# Source term is read-only, so re-use the default
if src is None:
src = self.source
# Create a new receiver object to store the result
if rec is None:
rec = Receiver(name='rec', grid=self.model.grid,
ntime=self.receiver.nt,
coordinates=self.receiver.coordinates.data)
# Create the forward wavefield if not provided
if u is None:
u = TimeFunction(name='u', grid=self.model.grid,
save=self.source.nt if save else None,
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
# Execute operator and return wavefield and receiver data
summary = self.op_fwd(save).apply(src=src, rec=rec, u=u, m=m,
dt=kwargs.pop('dt', self.dt), **kwargs)
return rec, u, summary
def test_subdim_middle(self):
"""
Tests that instantiating SubDimensions using the classmethod
constructors works correctly.
"""
grid = Grid(shape=(4, 4, 4))
x, y, z = grid.dimensions
t = grid.stepping_dim # noqa
u = TimeFunction(name='u', grid=grid) # noqa
xi = SubDimension.middle(name='xi',
parent=x,
thickness_left=1,
thickness_right=1)
eqs = [Eq(u.forward, u + 1)]
eqs = [e.subs(x, xi) for e in eqs]
op = Operator(eqs)
u.data[:] = 1.0
op.apply(time_M=1)
assert np.all(u.data[1, 0, :, :] == 1)
assert np.all(u.data[1, -1, :, :] == 1)
assert np.all(u.data[1, 1:3, :, :] == 2)
def test_empty_arrays(self):
"""
MFE for issue #1641.
"""
grid = Grid(shape=(4, 4), extent=(3.0, 3.0))
f = TimeFunction(name='f', grid=grid, space_order=0)
f.data[:] = 1.
sf1 = SparseTimeFunction(name='sf1', grid=grid, npoint=0, nt=10)
sf2 = SparseTimeFunction(name='sf2', grid=grid, npoint=0, nt=10)
assert sf1.size == 0
assert sf2.size == 0
eqns = sf1.inject(field=f, expr=sf1 + sf2 + 1.)
op = Operator(eqns)
op.apply()
assert np.all(f.data == 1.)
# Again, but with a MatrixSparseTimeFunction
mat = scipy.sparse.coo_matrix((0, 0), dtype=np.float32)
sf = MatrixSparseTimeFunction(name="s",
grid=grid,
r=2,
matrix=mat,
nt=10)
assert sf.size == 0
eqns = sf.interpolate(f)
op = Operator(eqns)
sf.manual_scatter()
op(time_m=0, time_M=9)
sf.manual_gather()
assert np.all(f.data == 1.)
def test_injection_wodup_wtime(self):
"""
Just like ``test_injection_wodup``, but using a SparseTimeFunction
instead of a SparseFunction. Hence, the data scattering/gathering now
has to correctly pack/unpack multidimensional arrays.
"""
grid = Grid(shape=(4, 4), extent=(3.0, 3.0))
save = 3
f = TimeFunction(name='f', grid=grid, save=save, space_order=0)
f.data[:] = 0.
coords = np.array([(0.5, 0.5), (0.5, 2.5), (2.5, 0.5), (2.5, 2.5)])
sf = SparseTimeFunction(name='sf', grid=grid, nt=save,
npoint=len(coords), coordinates=coords)
sf.data[0, :] = 4.
sf.data[1, :] = 8.
sf.data[2, :] = 12.
op = Operator(sf.inject(field=f, expr=sf + 1))
op.apply()
assert np.all(f.data[0] == 1.25)
assert np.all(f.data[1] == 2.25)
assert np.all(f.data[2] == 3.25)
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_codegen_quality1():
grid = Grid(shape=(4, 4, 4))
u = TimeFunction(name='u', grid=grid, space_order=2)
eqn = Eq(u.forward, u.dy.dy + 1.)
op = Operator(eqn,
opt=('advanced', {
'linearize': True,
'cire-mingain': 0
}))
assert 'uL0' in str(op)
# 11 expressions in total are expected, 8 of which are for the linearized accesses
exprs = FindNodes(Expression).visit(op)
assert len(exprs) == 11
assert all('const long' in str(i) for i in exprs[:-3])
assert all('const long' not in str(i) for i in exprs[-3:])
# Only two access macros necessary, namely `uL0` and `r1L0` (the other five
# obviously are _POSIX_C_SOURCE, MIN, MAX, START_TIMER, STOP_TIMER)
assert len(op._headers) == 7
def test_equations_mixed_densedata_timedata(self, shape, dimensions):
"""
Test that equations using a mixture of Function and TimeFunction objects
are embedded within the same time loop.
"""
grid = Grid(shape=shape, dimensions=dimensions, time_dimension=time)
a = TimeFunction(name='a', grid=grid, time_order=2, space_order=2)
p_aux = Dimension(name='p_aux')
b = Function(name='b', shape=shape + (10,), dimensions=dimensions + (p_aux,),
space_order=2)
b.data_allocated[:] = 1.0
b2 = Function(name='b2', shape=(10,) + shape, dimensions=(p_aux,) + dimensions,
space_order=2)
b2.data_allocated[:] = 1.0
eqns = [Eq(a.forward, a.laplace + 1.),
Eq(b, time*b*a + b)]
eqns2 = [Eq(a.forward, a.laplace + 1.),
Eq(b2, time*b2*a + b2)]
subs = {x.spacing: 2.5, y.spacing: 1.5, z.spacing: 2.0}
op = Operator(eqns, subs=subs, dle='noop')
trees = retrieve_iteration_tree(op)
assert len(trees) == 2
assert all(trees[0][i] is trees[1][i] for i in range(3))
op2 = Operator(eqns2, subs=subs, dle='noop')
trees = retrieve_iteration_tree(op2)
assert len(trees) == 2
# Verify both operators produce the same result
op(time=10)
a.data_allocated[:] = 0.
op2(time=10)
for i in range(10):
assert(np.allclose(b2.data[i, ...].reshape(-1) -
b.data[..., i].reshape(-1), 0.))
def test_basic(self):
grid = Grid(shape=(3, 3, 3))
u = TimeFunction(name='u', grid=grid)
op = Operator(Eq(u.forward, u + 1),
platform='nvidiaX',
language='openacc')
trees = retrieve_iteration_tree(op)
assert len(trees) == 1
assert trees[0][1].pragmas[0].value ==\
'acc parallel loop collapse(3) present(u)'
assert op.body[2].header[0].value ==\
('acc enter data copyin(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[2].footer[0]) == ''
assert op.body[2].footer[1].contents[0].value ==\
('acc exit data copyout(u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert op.body[2].footer[1].contents[1].value ==\
('acc exit data delete(u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
