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 GradientOperator(model, geometry, 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=geometry.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=geometry.time_axis,
npoint=geometry.nrec)
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)
# Substitute spacing terms to reduce flops
return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map,
name='Gradient', **kwargs)
def test_multiple_loop_nests(self):
"""
Compute a simple stencil S, preceded by an "initialization loop" I and
followed by a "random loop" R.
* S is the trivial equation ``u[t+1,x,y,z] = u[t,x,y,z] + 1``;
* I initializes ``u`` to 0;
* R adds 2 to another field ``v`` along the ``z`` dimension but only
over the planes ``[x=0, y=2]`` and ``[x=0, y=5]``.
Out of these three loop nests, only S should be "offloaded" to YASK; indeed,
I is outside the time loop, while R does not loop over space dimensions.
This test checks that S is the only loop nest "offloaded" to YASK, and
that the numerical output is correct.
"""
grid = Grid(shape=(12, 12, 12))
x, y, z = grid.dimensions
t = grid.stepping_dim
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
v = TimeFunction(name='yv4D', grid=grid, space_order=0)
v.data[:] = 0.
eqs = [Eq(u[0, x, y, z], 0),
Eq(u[1, x, y, z], 0),
Eq(u.forward, u + 1.),
Eq(v[t + 1, 0, 2, z], v[t + 1, 0, 2, z] + 2.),
Eq(v[t + 1, 0, 5, z], v[t + 1, 0, 5, z] + 2.)]
op = Operator(eqs)
op(yu4D=u, yv4D=v, time=0)
assert 'run_solution' in str(op)
assert len(retrieve_iteration_tree(op)) == 3
assert np.all(u.data[0] == 0.)
assert np.all(u.data[1] == 1.)
assert np.all(v.data[0] == 0.)
assert np.all(v.data[1, 0, 2] == 2.)
assert np.all(v.data[1, 0, 5] == 2.)
def test_const_change(self):
"""
Test that Constand.data can be set as required.
"""
n = 5
t = Constant(name='t', dtype=np.int32)
grid = Grid(shape=(2, 2))
x, y = grid.dimensions
f = TimeFunction(name='f', grid=grid, save=n+1)
f.data[:] = 0
eq = Eq(f.dt-1)
stencil = Eq(f.forward, solve(eq, f.forward))
op = Operator([stencil])
op.apply(time_m=0, time_M=n-1, dt=1)
check = Function(name='check', grid=grid)
eq_test = Eq(check, f[t, x, y])
op_test = Operator([eq_test])
for j in range(0, n+1):
t.data = j # Ensure constant is being updated correctly
op_test.apply(t=t)
assert(np.amax(check.data[:], axis=None) == j)
assert(np.amin(check.data[:], axis=None) == j)
def test_no_omp_if_offloaded(self):
grid = Grid(shape=(4, 4, 4))
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
u.data[:] = 0.
op = Operator(Eq(u, u + 1.))
assert 'run_solution' in str(op)
assert 'pragma omp' not in str(op)
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 at(shape=(11, 11)):
grid = Grid(shape=shape)
a = TimeFunction(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_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 all(np.isclose(sf.data[it, :], it))
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_segmented_averaging():
"""
Test for segmented operator execution of a two-sided, second order
function (averaging) in space. The corresponding set of stencil
offsets in the x domain are (1, 1).
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
t = grid.stepping_dim
f = TimeFunction(name='f', grid=grid)
fi = f.indexed
op = Operator(Eq(f, f.backward + (fi[t-1, x+1, y] + fi[t-1, x-1, y]) / 2.))
# We add the average to the point itself, so the grid "interior"
# (domain) is updated only.
f_ref = TimeFunction(name='f', grid=grid)
f_ref.data_with_halo[:] = 1.
op(f=f_ref, time=1)
assert (f_ref.data[1, :] == 2.).all()
assert (f_ref.data_with_halo[1, 0] == 1.).all()
assert (f_ref.data_with_halo[1, -1] == 1.).all()
# Now we sweep the x direction in 4 segmented steps of 5 iterations each
nsteps = 5
f.data_with_halo[:] = 1.
for i in range(4):
op(f=f, time=1, x_m=i*nsteps, x_M=(i+1)*nsteps-1)
assert (f_ref.data[1, :] == 2.).all()
assert (f_ref.data_with_halo[1, 0] == 1.).all()
assert (f_ref.data_with_halo[1, -1] == 1.).all()
def test_segmented_fibonacci():
"""
Test for segmented operator execution of a one-sided second order
function (fibonacci). The corresponding set of stencil offsets in
the time domain is (2, 0)
"""
# Reference Fibonacci solution from:
# https://stackoverflow.com/questions/4935957/fibonacci-numbers-with-an-one-liner-in-python-3
fib = lambda n: reduce(lambda x, n: [x[1], x[0] + x[1]], range(n), [0, 1])[0]
grid = Grid(shape=(5, 5))
x, y = grid.dimensions
t = grid.stepping_dim
f = TimeFunction(name='f', grid=grid, time_order=2)
fi = f.indexed
op = Operator(Eq(fi[t, x, y], fi[t-1, x, y] + fi[t-2, x, y]))
# Reference solution with a single invocation, up to timestep=12 (included)
# =========================================================================
# Developer note: the i-th Fibonacci number resides at logical index i-1
f_ref = TimeFunction(name='f', grid=grid, time_order=2)
f_ref.data[:] = 1.
op(f=f_ref, time=12)
assert (f_ref.data[11] == fib(12)).all()
assert (f_ref.data[12] == fib(13)).all()
# Now run with 2 invocations of 5 timesteps each
nsteps = 5
f.data[:] = 1.
for i in range(2):
op(f=f, time_m=2+i*nsteps, time_M=2+(i+1)*nsteps)
assert (f.data[11] == fib(12)).all()
assert (f.data[12] == fib(13)).all()
def test_fixed_halo_w_ofs(self, space_order):
"""
Compute an N-point stencil sum, where N is the number of points sorrounding
an inner (i.e., non-border) grid point.
For example (in 2D view):
1 1 1 ... 1 1
1 4 4 ... 4 1
1 4 4 ... 4 1
1 4 4 ... 4 1
1 1 1 ... 1 1
"""
grid = Grid(shape=(16, 16, 16))
v = TimeFunction(name='yv4D', grid=grid, space_order=space_order)
v.data_with_halo[:] = 1.
op = Operator(Eq(v.forward, v.laplace + 6*v), subs=grid.spacing_map)
op(yv4D=v, time=0)
assert 'run_solution' in str(op)
# Chech that the domain size has actually been written to
assert np.all(v.data[1] == 6.)
# Check that the halo planes are untouched
assert all(np.all(v.data_with_halo[1, i, :, :] == 1)
for i in range(v._size_halo.left[1]))
assert all(np.all(v.data_with_halo[1, :, i, :] == 1)
for i in range(v._size_halo.left[2]))
assert all(np.all(v.data_with_halo[1, :, :, i] == 1)
for i in range(v._size_halo.left[3]))
def unit_box_time(name='a', shape=(11, 11)):
"""Create a field with value 0. to 1. in each dimension"""
grid = Grid(shape=shape)
a = TimeFunction(name=name, grid=grid, time_order=1)
dims = tuple([np.linspace(0., 1., d) for d in shape])
a.data[0, :] = np.meshgrid(*dims)[1]
a.data[1, :] = np.meshgrid(*dims)[1]
return a
def test_second_derivatives_space(self, derivative, dim, order):
"""Test second derivative expressions against native sympy"""
dim = dim(self.grid)
u = TimeFunction(name='u', grid=self.grid, time_order=2, space_order=order)
expr = getattr(u, derivative)
# Establish native sympy derivative expression
width = int(order / 2)
indices = [(dim + i * dim.spacing) for i in range(-width, width + 1)]
s_expr = u.diff(dim, dim).as_finite_difference(indices).evalf(_PRECISION)
assert(simplify(expr - s_expr) == 0) # Symbolic equality
assert(expr == s_expr) # Exact equailty
def test_advanced_indexing(self):
"""Test data packing/unpacking via advanced indexing."""
grid = Grid(shape=(4, 4, 4))
u = TimeFunction(name='yu4D', grid=grid, space_order=0, time_order=1)
u.data[:] = 0.
# Test slicing w/ negative indices, combined to explicit indexing
u.data[1, 1:-1, 1:-1, 1:-1] = 6.
assert np.all(u.data[0] == 0.)
assert np.all(u.data[1, 1:-1, 1:-1, 1:-1] == 6.)
assert np.all(u.data[1, :, 0] == 0.)
assert np.all(u.data[1, :, -1] == 0.)
assert np.all(u.data[1, :, :, 0] == 0.)
assert np.all(u.data[1, :, :, -1] == 0.)
def test_logic_indexing(self):
"""Test logic indexing along stepping dimensions."""
grid = Grid(shape=(4, 4, 4))
v_mod = TimeFunction(name='v_mod', grid=grid)
v_mod.data[0] = 1.
v_mod.data[1] = 2.
assert np.all(v_mod.data[0] == 1.)
assert np.all(v_mod.data[1] == 2.)
assert np.all(v_mod.data[2] == v_mod.data[0])
assert np.all(v_mod.data[4] == v_mod.data[0])
assert np.all(v_mod.data[3] == v_mod.data[1])
assert np.all(v_mod.data[-1] == v_mod.data[1])
assert np.all(v_mod.data[-2] == v_mod.data[0])
def test_increasing_multi_steps(self):
"""
Apply the trivial equation ``u[t+1,x,y,z] = u[t,x,y,z] + 1`` for 11
timesteps and check that all grid domain values are equal to 11 within
``u[1]`` and equal to 10 within ``u[0]``.
"""
grid = Grid(shape=(8, 8, 8))
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
u.data_with_halo[:] = 0.
op = Operator(Eq(u.forward, u + 1.))
op(yu4D=u, time=10)
assert 'run_solution' in str(op)
assert np.all(u.data[0] == 10.)
assert np.all(u.data[1] == 11.)
def test_reverse_time_loop(self):
"""
Check that YASK evaluates stencil equations correctly when iterating in the
reverse time direction.
"""
grid = Grid(shape=(4, 4, 4))
u = TimeFunction(name='yu4D', grid=grid, space_order=0, time_order=2)
u.data[:] = 2.
eq = Eq(u.backward, u - 1.)
op = Operator(eq)
op(yu4D=u, time=2)
assert 'run_solution' in str(op)
assert np.all(u.data[2] == 2.)
assert np.all(u.data[1] == 1.)
assert np.all(u.data[0] == 0.)
def _new_operator2(shape, time_order, blockshape=None, dle=None):
blockshape = as_tuple(blockshape)
grid = Grid(shape=shape, dtype=np.int32)
infield = TimeFunction(name='infield', grid=grid, time_order=time_order)
infield.data[:] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
outfield = TimeFunction(name='outfield', grid=grid, time_order=time_order)
stencil = Eq(outfield.forward.indexify(),
outfield.indexify() + infield.indexify()*3.0)
op = Operator(stencil, dle=dle)
blocksizes = get_blocksizes(op, dle, grid, blockshape)
op(infield=infield, outfield=outfield, t=10, **blocksizes)
return outfield, op
def test_increasing_halo_wo_ofs(self, space_order, nosimd):
"""
Apply the trivial equation ``u[t+1,x,y,z] = u[t,x,y,z] + 1`` and check
that increasing space orders lead to proportionately larger halo regions,
which are *not* written by the Operator.
For example, with ``space_order = 0``, produce (in 2D view):
1 1 1 ... 1 1
1 1 1 ... 1 1
1 1 1 ... 1 1
1 1 1 ... 1 1
1 1 1 ... 1 1
With ``space_order = 1``, produce:
0 0 0 0 0 0 0 0 0
0 1 1 1 ... 1 1 0
0 1 1 1 ... 1 1 0
0 1 1 1 ... 1 1 0
0 1 1 1 ... 1 1 0
0 1 1 1 ... 1 1 0
0 0 0 0 0 0 0 0 0
And so on and so forth.
"""
# SIMD on/off
configuration['develop-mode'] = nosimd
grid = Grid(shape=(16, 16, 16))
u = TimeFunction(name='yu4D', grid=grid, space_order=space_order)
u.data_with_halo[:] = 0.
op = Operator(Eq(u.forward, u + 1.))
op(yu4D=u, time=0)
assert 'run_solution' in str(op)
# Chech that the domain size has actually been written to
assert np.all(u.data[1] == 1.)
# Check that the halo planes are still 0
assert all(np.all(u.data_with_halo[1, i, :, :] == 0)
for i in range(u._size_halo.left[1]))
assert all(np.all(u.data_with_halo[1, :, i, :] == 0)
for i in range(u._size_halo.left[2]))
assert all(np.all(u.data_with_halo[1, :, :, i] == 0)
for i in range(u._size_halo.left[3]))
def test_mixed_space_order(self):
"""
Make sure that no matter whether data objects have different space order,
as long as they have same domain, the Operator will be executed correctly.
"""
grid = Grid(shape=(8, 8, 8))
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
v = TimeFunction(name='yv4D', grid=grid, space_order=1)
u.data_with_halo[:] = 1.
v.data_with_halo[:] = 2.
op = Operator(Eq(v.forward, u + v))
op(yu4D=u, yv4D=v, time=0)
assert 'run_solution' in str(op)
# Chech that the domain size has actually been written to
assert np.all(v.data[1] == 3.)
# Check that the halo planes are untouched
assert np.all(v.data_with_halo[1, 0, :, :] == 2)
assert np.all(v.data_with_halo[1, :, 0, :] == 2)
assert np.all(v.data_with_halo[1, :, :, 0] == 2)
def test_at_w_mpi():
"""Make sure autotuning works in presence of MPI. MPI ranks work
in isolation to determine the best block size, locally."""
grid = Grid(shape=(8, 8))
t = grid.stepping_dim
x, y = grid.dimensions
f = TimeFunction(name='f', grid=grid, time_order=1)
f.data_with_halo[:] = 1.
eq = Eq(f.forward, f[t, x, y-1] + f[t, x, y+1])
op = Operator(eq, dle=('advanced', {'openmp': False, 'blockinner': True}))
op.apply(time=-1, autotune=('basic', 'runtime'))
# Nothing really happened, as not enough timesteps
assert np.all(f.data_ro_domain[0] == 1.)
assert np.all(f.data_ro_domain[1] == 1.)
# The 'destructive' mode writes directly to `f` for whatever timesteps required
# to perform the autotuning. Eventually, the result is complete garbage; note
# also that this autotuning mode disables the halo exchanges
op.apply(time=-1, autotune=('basic', 'destructive'))
assert np.all(f._data_ro_with_inhalo.sum() == 904)
# Check the halo hasn't been touched during AT
glb_pos_map = grid.distributor.glb_pos_map
if LEFT in glb_pos_map[y]:
assert np.all(f._data_ro_with_inhalo[:, :, -1] == 1)
else:
assert np.all(f._data_ro_with_inhalo[:, :, 0] == 1)
# Finally, try running w/o AT, just to be sure nothing was broken
f.data_with_halo[:] = 1.
op.apply(time=2)
if LEFT in glb_pos_map[y]:
assert np.all(f.data_ro_domain[1, :, 0] == 5.)
assert np.all(f.data_ro_domain[1, :, 1] == 7.)
assert np.all(f.data_ro_domain[1, :, 2:4] == 8.)
else:
assert np.all(f.data_ro_domain[1, :, 4:6] == 8)
assert np.all(f.data_ro_domain[1, :, 6] == 7)
assert np.all(f.data_ro_domain[1, :, 7] == 5)
def test_mpi_operator():
grid = Grid(shape=(4,))
f = TimeFunction(name='f', grid=grid)
g = TimeFunction(name='g', grid=grid)
# Using `sum` creates a stencil in `x`, which in turn will
# trigger the generation of code for MPI halo exchange
op = Operator(Eq(f.forward, f.sum() + 1))
op.apply(time=2)
pkl_op = pickle.dumps(op)
new_op = pickle.loads(pkl_op)
assert str(op) == str(new_op)
new_op.apply(time=2, f=g)
assert np.all(f.data[0] == [2., 3., 3., 3.])
assert np.all(f.data[1] == [3., 6., 7., 7.])
assert np.all(g.data[0] == f.data[0])
assert np.all(g.data[1] == f.data[1])
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_irregular_write(self):
"""
Compute a simple stencil S w/o offloading it to YASK because of the presence
of indirect write accesses (e.g. A[B[i]] = ...); YASK grid functions are however
used in the generated code to access the data at the right location. This
test checks that the numerical output is correct after this transformation.
Initially, the input array (a YASK grid, under the hood), at t=0 is (2D view):
0 1 2 3
0 1 2 3
0 1 2 3
0 1 2 3
Then, the Operator "flips" its content, and at timestep t=1 we get (2D view):
3 2 1 0
3 2 1 0
3 2 1 0
3 2 1 0
"""
grid = Grid(shape=(4, 4, 4))
x, y, z = grid.dimensions
t = grid.stepping_dim
p = SparseTimeFunction(name='points', grid=grid, nt=1, npoint=4)
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
for i in range(4):
for j in range(4):
for k in range(4):
u.data[0, i, j, k] = k
ind = lambda i: p[0, i]
eqs = [Eq(p[0, 0], 3.), Eq(p[0, 1], 2.),
Eq(p[0, 2], 1.), Eq(p[0, 3], 0.),
Eq(u[t + 1, ind(x), ind(y), ind(z)], u[t, x, y, z])]
op = Operator(eqs, subs=grid.spacing_map)
op(yu4D=u, time=0)
assert 'run_solution' not in str(op)
assert all(np.all(u.data[1, :, :, i] == 3 - i) for i in range(4))
def _new_operator3(shape, blockshape=None, dle=None):
blockshape = as_tuple(blockshape)
grid = Grid(shape=shape)
spacing = 0.1
a = 0.5
c = 0.5
dx2, dy2 = spacing**2, spacing**2
dt = dx2 * dy2 / (2 * a * (dx2 + dy2))
# Allocate the grid and set initial condition
# Note: This should be made simpler through the use of defaults
u = TimeFunction(name='u', grid=grid, time_order=1, space_order=(2, 2, 2))
u.data[0, :] = np.arange(reduce(mul, shape), dtype=np.int32).reshape(shape)
# Derive the stencil according to devito conventions
eqn = Eq(u.dt, a * (u.dx2 + u.dy2) - c * (u.dxl + u.dyl))
stencil = solve(eqn, u.forward)
op = Operator(Eq(u.forward, stencil), dle=dle)
blocksizes = get_blocksizes(op, dle, grid, blockshape)
op.apply(u=u, t=10, dt=dt, **blocksizes)
return u.data[1, :], op
def test_constants(self):
"""
Check that :class:`Constant` objects are treated correctly.
"""
grid = Grid(shape=(4, 4, 4))
c = Constant(name='c', value=2., dtype=grid.dtype)
p = SparseTimeFunction(name='points', grid=grid, nt=1, npoint=1)
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
u.data[:] = 0.
op = Operator([Eq(u.forward, u + c), Eq(p[0, 0], 1. + c)])
assert 'run_solution' in str(op)
op.apply(yu4D=u, c=c, time=9)
# Check YASK did its job and could read constant grids w/o problems
assert np.all(u.data[0] == 20.)
# Check the Constant could be read correctly even in Devito-land, i.e.,
# outside of run_solution
assert p.data[0][0] == 3.
# Check re-executing with another constant gives the correct result
c2 = Constant(name='c', value=5.)
op.apply(yu4D=u, c=c2, time=2)
assert np.all(u.data[0] == 30.)
assert np.all(u.data[1] == 35.)
assert p.data[0][0] == 6.
def execute_devito(ui, spacing=0.01, a=0.5, timesteps=500):
"""Execute diffusion stencil using the devito Operator API."""
nx, ny = ui.shape
dx2, dy2 = spacing**2, spacing**2
dt = dx2 * dy2 / (2 * a * (dx2 + dy2))
# Allocate the grid and set initial condition
# Note: This should be made simpler through the use of defaults
grid = Grid(shape=(nx, ny))
u = TimeFunction(name='u', grid=grid, time_order=1, space_order=2)
u.data[0, :] = ui[:]
# Derive the stencil according to devito conventions
eqn = Eq(u.dt, a * (u.dx2 + u.dy2))
stencil = solve(eqn, u.forward)
op = Operator(Eq(u.forward, stencil))
# Execute the generated Devito stencil operator
tstart = time.time()
op.apply(u=u, t=timesteps, dt=dt)
runtime = time.time() - tstart
log("Devito: Diffusion with dx=%0.4f, dy=%0.4f, executed %d timesteps in %f seconds"
% (spacing, spacing, timesteps, runtime))
return u.data[1, :], runtime
def test_repeated_op_calls(self):
"""
Tests that calling the same Operator with different input data
produces the expected results.
"""
grid = Grid(shape=(4, 4, 4))
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
u.data[:] = 0.
op = Operator(Eq(u.forward, u + 1.))
# First run
op(time=0)
assert np.all(u.data[1] == 1.)
assert u.data[:].sum() == np.prod(grid.shape)
# Nothing should have changed at this point
op(time=0, yu4D=u)
assert np.all(u.data[1] == 1.)
assert u.data[:].sum() == np.prod(grid.shape)
# Now try with a different grid
grid = Grid(shape=(3, 3, 3))
u = TimeFunction(name='yu4D', grid=grid, space_order=0)
u.data[:] = 0.
op(time=0, yu4D=u)
assert np.all(u.data[1] == 1.)
assert u.data[:].sum() == np.prod(grid.shape)
def forward(self,
src=None,
rec=None,
u=None,
v=None,
m=None,
epsilon=None,
delta=None,
theta=None,
phi=None,
save=False,
kernel='centered',
**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
:param kernel: type of discretization, centered or shifted
:returns: Receiver, wavefield and performance summary
"""
# Space order needs to be halved in the shifted case to have an
# overall space_order discretization
self.space_order = self.space_order // 2 if kernel == 'shifted' \
else self.space_order
# 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=save,
time_dim=self.source.nt if save else None,
time_order=self.time_order,
space_order=self.space_order)
# Create the forward wavefield if not provided
if v is None:
v = TimeFunction(name='v',
grid=self.model.grid,
save=save,
time_dim=self.source.nt if save else None,
time_order=self.time_order,
space_order=self.space_order)
# Pick m from model unless explicitly provided
if m is None:
m = m or self.model.m
if epsilon is None:
epsilon = epsilon or self.model.epsilon
if delta is None:
delta = delta or self.model.delta
if theta is None:
theta = theta or self.model.theta
if phi is None:
phi = phi or self.model.phi
# Execute operator and return wavefield and receiver data
op = self.op_fwd(kernel, save)
if len(m.shape) == 2:
summary = op.apply(src=src,
rec=rec,
u=u,
v=v,
m=m,
epsilon=epsilon,
delta=delta,
theta=theta,
dt=self.dt,
**kwargs)
else:
summary = op.apply(src=src,
rec=rec,
u=u,
v=v,
m=m,
epsilon=epsilon,
delta=delta,
theta=theta,
phi=phi,
dt=self.dt,
**kwargs)
return rec, u, v, summary
def u(self, model, space_order, kernel):
return TimeFunction(name='u',
grid=model.grid,
space_order=space_order,
time_order=2)
def td_born_adjoint(born_data, model_pert, src_coords, vel, geometry, solver, params, dt=None):
"""
@Params
born_data: float32 numpy array of size (Ns, Nt, Nr).
model_pert: float32 numpy array of size (Nt, Nx, Nz). This will be the output. Assumed to be zeros.
src_coords: float32 numpy array of size (Ns, 2) with the source coordinates.
vel: background velocity model object (The object should be the same type as returned by create_model() function)
geometry: geometry object
solver: acoustic solver object
params: python dict of parameters
Note: The receiver coordinates are not needed since it is assumed that they remain fixed for all sources.
So the receiver information should already be available to the solver object.
"""
# Get params
nbl = params["nbl"]
nx = params["Nx"]
nz = params["Nz"]
ns = params["Ns"]
nt = params["Nt"]
space_order = params["so"]
time_order = params["to"]
offset = nbl + space_order
if dt is None:
dt = vel.critical_dt
# Allocate time function to store adjoint wavefield
u = TimeFunction(
name='u',
grid=vel.grid,
time_order=time_order,
space_order=space_order,
save=nt
)
# Create time dependent Born modeling operator
op = td_born_adjoint_op(model=vel, geometry=geometry, time_order=time_order, space_order=space_order)
# Second derivative filter stencil
laplacian_filter = np.asarray([1, -2, 1], dtype=np.float32) / (dt ** 2.0)
# Perform adjoint Born modeling
for i in range(ns):
# Update source location
geometry.src_positions[0, :] = src_coords[i, :]
# Get Born modeled data for the current shot and update born_data array appropriately
_, u0, _ = solver.forward(vp=vel.vp, save=True)
spim.convolve1d(
input=u0.data_with_halo,
weights=laplacian_filter,
axis=0,
output=u0.data_with_halo[:, :, :],
mode='nearest'
)
op.apply(u=u, born_data_rec=born_data[i, :, :], dt=dt)
# Add to model_pert
model_pert += \
u.data_with_halo[:, offset:nx + offset, offset:nz + offset] * \
u0.data_with_halo[:, offset:nx + offset, offset:nz + offset]
def test_tti(shape, space_order, kernel):
"""
This first test compare the solution of the acoustic wave-equation and the
TTI wave-eqatuon with all anisotropy parametrs to 0. The two solutions should
be the same.
"""
if kernel == 'shifted':
space_order *= 2
to = 2
so = space_order
nbpml = 10
origin = [0. for _ in shape]
spacing = [10. for _ in shape]
vp = 1.5 * np.ones(shape)
nrec = shape[0]
# Constant model for true velocity
model = Model(origin=origin, shape=shape, vp=vp,
spacing=spacing, nbpml=nbpml, space_order=space_order,
epsilon=np.zeros(shape), delta=np.zeros(shape),
theta=np.zeros(shape), phi=np.zeros(shape))
# Source and receiver geometries
src_coordinates = np.empty((1, len(spacing)))
src_coordinates[0, :] = np.array(model.domain_size) * .5
src_coordinates[0, -1] = model.origin[-1] + 2 * spacing[-1]
rec_coordinates = np.empty((nrec, len(spacing)))
rec_coordinates[:, 0] = np.linspace(0., model.domain_size[0], num=nrec)
rec_coordinates[:, 1] = np.array(model.domain_size)[1] * .5
rec_coordinates[:, -1] = model.origin[-1] + 2 * spacing[-1]
geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates,
t0=0.0, tn=350., src_type='Ricker', f0=0.010)
acoustic = AcousticWaveSolver(model, geometry, time_order=2, space_order=so)
rec, u1, _ = acoustic.forward(save=False)
# Solvers
solver_tti = AnisotropicWaveSolver(model, geometry, time_order=2,
space_order=space_order)
# zero src
src = geometry.src
src.data.fill(0.)
# last time index
nt = geometry.nt
last = (nt - 2) % 3
indlast = [(last + 1) % 3, last % 3, (last-1) % 3]
# Create new wavefield object restart forward computation
u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=so)
u.data[0:3, :] = u1.data[indlast, :]
acoustic.forward(save=False, u=u, time_M=10, src=src)
utti = TimeFunction(name='u', grid=model.grid, time_order=to, space_order=so)
vtti = TimeFunction(name='v', grid=model.grid, time_order=to, space_order=so)
utti.data[0:to+1, :] = u1.data[indlast[:to+1], :]
vtti.data[0:to+1, :] = u1.data[indlast[:to+1], :]
solver_tti.forward(u=utti, v=vtti, kernel=kernel, time_M=10, src=src)
normal_u = u.data[:]
normal_utti = .5 * utti.data[:]
normal_vtti = .5 * vtti.data[:]
res = linalg.norm((normal_u - normal_utti - normal_vtti).reshape(-1))**2
res /= np.linalg.norm(normal_u.reshape(-1))**2
log("Difference between acoustic and TTI with all coefficients to 0 %2.4e" % res)
assert np.isclose(res, 0.0, atol=1e-4)
def heat3d_devito(in_field, alpha, Tcool, dt, nt, platform='default', result='time'):
'''
Solves the 3 dimensional heat equation using devito.
in_field -- input field (nx x ny x nz)
alpha -- thermal diffusivity
Tcool -- initial temperature
dt -- time step
nt -- number of iterations
platform -- the platform on which the equation is to be solved
result -- either 'time', 'field' or 'both', returning either the total time the computation took, the resulting field or both
'''
nx = in_field.shape[0]
ny = in_field.shape[1]
nz = in_field.shape[2]
grid = Grid(shape = (nx, ny, nz), subdomains = (mid, ), extent = (2., 2., 2.))
u = TimeFunction(name='u', grid=grid, time_order=1, space_order=2, dtype=np.float64)
u.data[:, :, :, :] = in_field
x, y, z = grid.dimensions
t = grid.stepping_dim
#boundary conditions
bc = [Eq(u[t+1, 0 , y, z], Tcool)] # left
bc += [Eq(u[t+1, nx-1, y, z], Tcool)] # right
bc += [Eq(u[t+1, x, ny-1, z], Tcool)] # top
bc += [Eq(u[t+1, x, 0, z], Tcool)] # bottom
bc += [Eq(u[t+1, x, y, 0], Tcool)] # top
bc += [Eq(u[t+1, x, y, nz-1], Tcool)] # bottom
#define heat equation in 3d
eq = Eq(u.dt, alpha * ((u.dx2)+(u.dy2)+(u.dz2)))
#solve equation
stencil = solve(eq, u.forward)
#create stencil
eq_stencil = Eq(u.forward, stencil,subdomain = grid.subdomains['middle'])
eq_stencil
#solve
if platform == 'default':
op = Operator([eq_stencil]+bc)
else:
op = Operator([eq_stencil]+bc, platform=platform)
#print(op.ccode)
tic = get_time()
op(time=nt, dt=dt)
#toc = time.time()
elapsed_time = get_time() - tic
out_field = u.data[0,:,:,:]
if result == 'time':
return elapsed_time
if result == 'field':
return out_field
if result == 'both':
return elapsed_time, out_field
from devito import Grid, TimeFunction, Eq, solve, Operator
so_dev = 6
to_dev = 2
extent = (150., 150.)
shape = (151, 151)
dt = 0.2
grid = Grid(extent=extent, shape=shape)
t_end = 75
ns = int(t_end / dt) # Number of timesteps = total time/timestep size.
# Set up function and stencil
u_dev = TimeFunction(name="u_dev",
grid=grid,
space_order=so_dev,
time_order=to_dev)
u_dev.data[:] = 0.
time = grid.time_dim
x = grid.dimensions[0]
y = grid.dimensions[1]
h_x = grid.spacing[0]
h_y = grid.spacing[1]
# Optimized stencil coefficients
a_0 = -2.81299833
a_1 = 1.56808208
a_2 = -0.17723283
a_3 = 0.01564992
def forward(self,
src=None,
rec=None,
v=None,
r=None,
p=None,
qp=None,
b=None,
vp=None,
save=None,
**kwargs):
"""
Forward modelling function that creates the necessary
data objects for running a forward modelling operator.
Parameters
----------
geometry : AcquisitionGeometry
Geometry object that contains the source (SparseTimeFunction) and
receivers (SparseTimeFunction) and their position.
v : VectorTimeFunction, optional
The computed particle velocity.
r : TimeFunction, optional
The computed memory variable.
p : TimeFunction, optional
Stores the computed wavefield.
qp : Function, optional
The P-wave quality factor.
b : Function, optional
The time-constant inverse density.
vp : Function or float, optional
The time-constant velocity.
save : bool, optional
Whether or not to save the entire (unrolled) wavefield.
Returns
-------
Receiver, wavefield and performance summary
"""
# Source term is read-only, so re-use the default
src = src or self.geometry.src
# Create a new receiver object to store the result
rec = rec or self.geometry.rec
# Create all the fields v, p, r
save_t = src.nt if save else None
v = v or VectorTimeFunction(name="v",
grid=self.model.grid,
save=save_t,
time_order=1,
space_order=self.space_order)
# Create the forward wavefield if not provided
p = p or TimeFunction(name="p",
grid=self.model.grid,
save=save_t,
time_order=1,
space_order=self.space_order,
staggered=NODE)
# Memory variable:
r = r or TimeFunction(name="r",
grid=self.model.grid,
save=save_t,
time_order=1,
space_order=self.space_order,
staggered=NODE)
kwargs.update({k.name: k for k in v})
# Pick physical parameters from model unless explicitly provided
b = b or self.model.b
qp = qp or self.model.qp
# Pick vp from model unless explicitly provided
vp = vp or self.model.vp
if self.kernel == 'blanch_symes':
# Execute operator and return wavefield and receiver data
# With Memory variable
summary = self.op_fwd(save).apply(src=src,
rec=rec,
qp=qp,
r=r,
p=p,
b=b,
vp=vp,
dt=kwargs.pop('dt', self.dt),
**kwargs)
else:
# Execute operator and return wavefield and receiver data
# Without Memory variable
summary = self.op_fwd(save).apply(src=src,
rec=rec,
qp=qp,
p=p,
b=b,
vp=vp,
dt=kwargs.pop('dt', self.dt),
**kwargs)
return rec, p, summary
def adjoint(self,
rec,
srca=None,
p=None,
r=None,
model=None,
save=None,
**kwargs):
"""
Adjoint modelling function that creates the necessary
data objects for running an adjoint modelling operator.
Parameters
----------
geometry : AcquisitionGeometry
Geometry object that contains the source (SparseTimeFunction) and
receivers (SparseTimeFunction) and their position.
p : TimeFunction, optional
The computed wavefield first component.
r : TimeFunction, optional
The computed wavefield second component.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
Returns
-------
Adjoint source, wavefield and performance summary.
"""
if self.kernel == 'staggered':
time_order = 1
dims = self.model.space_dimensions
stagg_p = (-dims[-1])
stagg_r = (-dims[0],
-dims[1]) if self.model.grid.dim == 3 else (-dims[0])
else:
time_order = 2
stagg_p = stagg_r = None
# Source term is read-only, so re-use the default
srca = srca or self.geometry.new_src(name='srca', src_type=None)
# Create the wavefield if not provided
if p is None:
p = TimeFunction(name='p',
grid=self.model.grid,
staggered=stagg_p,
time_order=time_order,
space_order=self.space_order)
# Create the wavefield if not provided
if r is None:
r = TimeFunction(name='r',
grid=self.model.grid,
staggered=stagg_r,
time_order=time_order,
space_order=self.space_order)
if self.kernel == 'staggered':
vx, vz, vy = particle_velocity_fields(self.model, self.space_order)
kwargs["vx"] = vx
kwargs["vz"] = vz
if vy is not None:
kwargs["vy"] = vy
model = model or self.model
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if self.model.dim < 3:
kwargs.pop('phi', None)
# Execute operator and return wavefield and receiver data
summary = self.op_adj().apply(srca=srca,
rec=rec,
p=p,
r=r,
dt=kwargs.pop('dt', self.dt),
time_m=0 if time_order == 1 else None,
**kwargs)
return srca, p, r, summary
def jacobian(self,
dm,
src=None,
rec=None,
u0=None,
v0=None,
du=None,
dv=None,
model=None,
save=None,
kernel='centered',
**kwargs):
"""
Linearized Born modelling function that creates the necessary
data objects for running an adjoint modelling operator.
Parameters
----------
src : SparseTimeFunction or array_like, optional
Time series data for the injected source term.
rec : SparseTimeFunction or array_like, optional
The interpolated receiver data.
u : TimeFunction, optional
The computed background wavefield first component.
v : TimeFunction, optional
The computed background wavefield second component.
du : TimeFunction, optional
The computed perturbed wavefield first component.
dv : TimeFunction, optional
The computed perturbed wavefield second component.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
"""
if kernel != 'centered':
raise ValueError(
'Only centered kernel is supported for the jacobian')
dt = kwargs.pop('dt', self.dt)
# Source term is read-only, so re-use the default
src = src or self.geometry.src
# Create a new receiver object to store the result
rec = rec or self.geometry.rec
# Create the forward wavefields u, v du and dv if not provided
u0 = u0 or TimeFunction(name='u0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
v0 = v0 or TimeFunction(name='v0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
du = du or TimeFunction(name='du',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
dv = dv or TimeFunction(name='dv',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
model = model or self.model
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if self.model.dim < 3:
kwargs.pop('phi', None)
# Execute operator and return wavefield and receiver data
summary = self.op_jac().apply(dm=dm,
u0=u0,
v0=v0,
du=du,
dv=dv,
src=src,
rec=rec,
dt=dt,
**kwargs)
return rec, u0, v0, du, dv, summary
def jacobian_adjoint(self,
rec,
u0,
v0,
du=None,
dv=None,
dm=None,
model=None,
checkpointing=False,
kernel='centered',
**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.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u0 : TimeFunction
The computed background wavefield.
v0 : TimeFunction, optional
The computed background wavefield.
du : Function or float
The computed perturbed wavefield.
dv : Function or float
The computed perturbed wavefield.
dm : Function, optional
Stores the gradient field.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
Returns
-------
Gradient field and performance summary.
"""
if kernel != 'centered':
raise ValueError(
'Only centered kernel is supported for the jacobian_adj')
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
dm = dm or Function(name='dm', grid=self.model.grid)
# Create the perturbation wavefields if not provided
du = du or TimeFunction(name='du',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
dv = dv or TimeFunction(name='dv',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
model = model or self.model
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if self.model.dim < 3:
kwargs.pop('phi', None)
if checkpointing:
u0 = TimeFunction(name='u0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
v0 = TimeFunction(name='v0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
cp = DevitoCheckpoint([u0, v0])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src,
u=u0,
v=v0,
dt=dt,
**kwargs)
wrap_rev = CheckpointOperator(self.op_jacadj(save=False),
u0=u0,
v0=v0,
du=du,
dv=dv,
rec=rec,
dm=dm,
dt=dt,
**kwargs)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
rec.data.shape[0] - 2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_jacadj().apply(rec=rec,
dm=dm,
u0=u0,
v0=v0,
du=du,
dv=dv,
dt=dt,
**kwargs)
return dm, summary
def forward(self,
src=None,
rec=None,
u=None,
v=None,
model=None,
save=False,
**kwargs):
"""
Forward modelling function that creates the necessary
data objects for running a forward modelling operator.
Parameters
----------
geometry : AcquisitionGeometry
Geometry object that contains the source (SparseTimeFunction) and
receivers (SparseTimeFunction) and their position.
u : TimeFunction, optional
The computed wavefield first component.
v : TimeFunction, optional
The computed wavefield second component.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
save : bool, optional
Whether or not to save the entire (unrolled) wavefield.
kernel : str, optional
Type of discretization, centered or shifted.
Returns
-------
Receiver, wavefield and performance summary.
"""
if self.kernel == 'staggered':
time_order = 1
dims = self.model.space_dimensions
stagg_u = (-dims[-1])
stagg_v = (-dims[0],
-dims[1]) if self.model.grid.dim == 3 else (-dims[0])
else:
time_order = 2
stagg_u = stagg_v = None
# Source term is read-only, so re-use the default
src = src or self.geometry.src
# Create a new receiver object to store the result
rec = rec or self.geometry.rec
# Create the forward wavefield if not provided
if u is None:
u = TimeFunction(name='u',
grid=self.model.grid,
staggered=stagg_u,
save=self.geometry.nt if save else None,
time_order=time_order,
space_order=self.space_order)
# Create the forward wavefield if not provided
if v is None:
v = TimeFunction(name='v',
grid=self.model.grid,
staggered=stagg_v,
save=self.geometry.nt if save else None,
time_order=time_order,
space_order=self.space_order)
if self.kernel == 'staggered':
vx, vz, vy = particle_velocity_fields(self.model, self.space_order)
kwargs["vx"] = vx
kwargs["vz"] = vz
if vy is not None:
kwargs["vy"] = vy
model = model or self.model
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if self.model.dim < 3:
kwargs.pop('phi', None)
# Execute operator and return wavefield and receiver data
summary = self.op_fwd(save).apply(src=src,
rec=rec,
u=u,
v=v,
dt=kwargs.pop('dt', self.dt),
**kwargs)
return rec, u, v, summary
def timefunction(name, space_order=1):
return TimeFunction(name=name, grid=grid, space_order=space_order)
def ForwardOperator(model,
source,
receiver,
space_order=4,
save=False,
kernel='centered',
**kwargs):
"""
Constructor method for the forward modelling operator in an acoustic media
:param model: :class:`Model` object containing the physical parameters
:param src: None ot IShot() (not currently supported properly)
:param data: IShot() object containing the acquisition geometry and field data
:param: time_order: Time discretization order
:param: spc_order: Space discretization order
"""
dt = model.grid.time_dim.spacing
m, damp, epsilon, delta, theta, phi = (model.m, model.damp, model.epsilon,
model.delta, model.theta, model.phi)
# 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)
v = TimeFunction(name='v',
grid=model.grid,
save=source.nt if save else None,
time_order=2,
space_order=space_order)
src = PointSource(name='src',
grid=model.grid,
time_range=source.time_range,
npoint=source.npoint)
rec = Receiver(name='rec',
grid=model.grid,
time_range=receiver.time_range,
npoint=receiver.npoint)
# Tilt and azymuth setup
ang0 = cos(theta)
ang1 = sin(theta)
ang2 = 0
ang3 = 0
if len(model.shape) == 3:
ang2 = cos(phi)
ang3 = sin(phi)
FD_kernel = kernels[(kernel, len(model.shape))]
H0, Hz = FD_kernel(u, v, ang0, ang1, ang2, ang3, space_order)
# Stencils
s = model.grid.stepping_dim.spacing
stencilp = 1.0 / (2.0 * m + s * damp) * \
(4.0 * m * u + (s * damp - 2.0 * m) *
u.backward + 2.0 * s ** 2 * (epsilon * H0 + delta * Hz))
stencilr = 1.0 / (2.0 * m + s * damp) * \
(4.0 * m * v + (s * damp - 2.0 * m) *
v.backward + 2.0 * s ** 2 * (delta * H0 + Hz))
first_stencil = Eq(u.forward, stencilp)
second_stencil = Eq(v.forward, stencilr)
stencils = [first_stencil, second_stencil]
# Source and receivers
stencils += src.inject(field=u.forward,
expr=src * dt * dt / m,
offset=model.nbpml)
stencils += src.inject(field=v.forward,
expr=src * dt * dt / m,
offset=model.nbpml)
stencils += rec.interpolate(expr=u + v, offset=model.nbpml)
# Substitute spacing terms to reduce flops
return Operator(stencils,
subs=model.spacing_map,
name='ForwardTTI',
**kwargs)
def forward(self,
src=None,
rec=None,
u=None,
v=None,
m=None,
epsilon=None,
delta=None,
theta=None,
phi=None,
save=False,
kernel='centered',
**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 (u+v)
:param u: (Optional) Symbol to store the computed wavefield first component
:param v: (Optional) Symbol to store the computed wavefield second component
:param m: (Optional) Symbol for the time-constant square slowness
:param epsilon: (Optional) Symbol for the time-constant first Thomsen parameter
:param delta: (Optional) Symbol for the time-constant second Thomsen parameter
:param theta: (Optional) Symbol for the time-constant Dip angle (radians)
:param phi: (Optional) Symbol for the time-constant Azimuth angle (radians)
:param save: Option to store the entire (unrolled) wavefield
:param kernel: type of discretization, centered or shifted
:returns: Receiver, wavefield and performance summary
"""
if kernel == 'staggered':
time_order = 1
dims = self.model.space_dimensions
stagg_u = (-dims[-1])
stagg_v = (-dims[0],
-dims[1]) if self.model.grid.dim == 3 else (-dims[0])
else:
time_order = 2
stagg_u = stagg_v = None
# Source term is read-only, so re-use the default
src = src or self.geometry.src
# Create a new receiver object to store the result
rec = rec or Receiver(name='rec',
grid=self.model.grid,
time_range=self.geometry.time_axis,
coordinates=self.geometry.rec_positions)
# Create the forward wavefield if not provided
if u is None:
u = TimeFunction(name='u',
grid=self.model.grid,
staggered=stagg_u,
save=self.geometry.nt if save else None,
time_order=time_order,
space_order=self.space_order)
# Create the forward wavefield if not provided
if v is None:
v = TimeFunction(name='v',
grid=self.model.grid,
staggered=stagg_v,
save=self.geometry.nt if save else None,
time_order=time_order,
space_order=self.space_order)
if kernel == 'staggered':
vx, vz, vy = particle_velocity_fields(self.model, self.space_order)
kwargs["vx"] = vx
kwargs["vz"] = vz
if vy is not None:
kwargs["vy"] = vy
# Pick m from model unless explicitly provided
kwargs.update(
self.model.physical_params(m=m,
epsilon=epsilon,
delta=delta,
theta=theta,
phi=phi))
# Execute operator and return wavefield and receiver data
op = self.op_fwd(kernel, save)
summary = op.apply(src=src,
rec=rec,
u=u,
v=v,
dt=kwargs.pop('dt', self.dt),
**kwargs)
return rec, u, v, summary
def test_gradient_equivalence(self, shape, kernel, space_order, preset,
nbl, dtype, tolerance, spacing, tn):
""" This test asserts that the gradient calculated through the following three
expressions should match within floating-point precision:
- grad = sum(-u.dt2 * v)
- grad = sum(-u * v.dt2)
- grad = sum(-u.dt * v.dt)
The computation has the following number of operations:
u.dt2 (5 ops) * v = 6ops * 500 (nt) ~ 3000 ops ~ 1e4 ops
Hence tolerances are eps * ops = 1e-4 (sp) and 1e-13 (dp)
"""
model = demo_model(preset,
space_order=space_order,
shape=shape,
nbl=nbl,
dtype=dtype,
spacing=spacing)
m = model.m
v_true = model.vp
geometry = setup_geometry(model, tn)
dt = model.critical_dt
src = geometry.src
rec = geometry.rec
rec_true = geometry.rec
rec0 = geometry.rec
s = model.grid.stepping_dim.spacing
u = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=space_order,
save=geometry.nt)
eqn_fwd = iso_stencil(u, model, kernel)
src_term = src.inject(field=u.forward, expr=src * s**2 / m)
rec_term = rec.interpolate(expr=u)
fwd_op = Operator(eqn_fwd + src_term + rec_term,
subs=model.spacing_map,
name='Forward')
v0 = Function(name='v0',
grid=model.grid,
space_order=space_order,
dtype=dtype)
smooth(v0, model.vp)
grad_u = Function(name='gradu', grid=model.grid)
grad_v = Function(name='gradv', grid=model.grid)
grad_uv = Function(name='graduv', grid=model.grid)
v = TimeFunction(name='v',
grid=model.grid,
save=None,
time_order=2,
space_order=space_order)
s = model.grid.stepping_dim.spacing
eqn_adj = iso_stencil(v, model, kernel, forward=False)
receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)
gradient_update_v = Eq(grad_v, grad_v - u * v.dt2)
grad_op_v = Operator(eqn_adj + receivers + [gradient_update_v],
subs=model.spacing_map,
name='GradientV')
gradient_update_u = Eq(grad_u, grad_u - u.dt2 * v)
grad_op_u = Operator(eqn_adj + receivers + [gradient_update_u],
subs=model.spacing_map,
name='GradientU')
gradient_update_uv = Eq(grad_uv, grad_uv + u.dt * v.dt)
grad_op_uv = Operator(eqn_adj + receivers + [gradient_update_uv],
subs=model.spacing_map,
name='GradientUV')
fwd_op.apply(dt=dt, vp=v_true, rec=rec_true)
fwd_op.apply(dt=dt, vp=v0, rec=rec0)
residual = Receiver(name='rec',
grid=model.grid,
data=(rec0.data - rec_true.data),
time_range=geometry.time_axis,
coordinates=geometry.rec_positions,
dtype=dtype)
grad_op_u.apply(dt=dt, vp=v0, rec=residual)
# Reset v before calling the second operator since the object is shared
v.data[:] = 0.
grad_op_v.apply(dt=dt, vp=v0, rec=residual)
v.data[:] = 0.
grad_op_uv.apply(dt=dt, vp=v0, rec=residual)
assert (np.allclose(grad_u.data,
grad_v.data,
rtol=tolerance,
atol=tolerance))
assert (np.allclose(grad_u.data,
grad_uv.data,
rtol=tolerance,
atol=tolerance))
#### SOURCE SETTING ####
t0 = 0. # Simulation starts a t=0
tn = np.ceil(TIME * 1000) # Simulation last 1 second (1000 ms)
dt = model.critical_dt # Time step from model grid spacing
time_range = TimeAxis(start=t0, stop=tn, step=dt)
f0 = 0.010 # Source peak frequency is 10Hz (0.010 kHz)
src = RickerSource(name='src', grid=model.grid, f0=f0,
npoint=1, time_range=time_range)
src.coordinates.data[0, :] = np.array([SRC_LOC_X, SRC_LOC_Y, SRC_LOC_Z])
#src.show()
#### RECEIVER SETTING ####
rec = Receiver(name='rec', grid=model.grid, npoint=REC_LOC.shape[0], time_range=time_range)
rec.coordinates.data[:,:] = REC_LOC
# ACOUSTIC MODEL
from devito import TimeFunction
u = TimeFunction(name="u", grid=model.grid, time_order=2, space_order=2)
pde = model.m * u.dt2 - u.laplace + model.damp * u.dt
stencil = Eq(u.forward, solve(pde, u.forward))
src_term = src.inject(field=u.forward, expr=src * dt**2 / model.m)
rec_term = rec.interpolate(expr=u.forward)
op = Operator([stencil] + src_term + rec_term, subs=model.spacing_map)
op(time=time_range.num-1, dt=model.critical_dt)
plot_shotrecord(rec.data, model, t0, tn)
def test_streaming_complete(self):
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)
u2 = 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', 'streaming', 'orchestrate'))
op2 = Operator(eqns,
opt=('buffering', 'streaming', 'fuse', 'orchestrate'))
# Check generated code
assert len(op1._func_table) == 6
assert len([i for i in FindSymbols().visit(op1) if i.is_Array]) == 2
assert len(op2._func_table) == 4
assert len([i for i in FindSymbols().visit(op2) 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)
op2.apply(time_m=15, time_M=35, save_shift=0, u=u2)
assert np.all(u.data == u1.data)
assert np.all(u.data == u2.data)
def ForwardOperator(model,
geometry,
space_order=4,
save=False,
kernel='centered',
**kwargs):
"""
Construct an forward modelling operator in an acoustic media.
Parameters
----------
model : Model
Object containing the physical parameters.
geometry : AcquisitionGeometry
Geometry object that contains the source (SparseTimeFunction) and
receivers (SparseTimeFunction) and their position.
data : ndarray
IShot() object containing the acquisition geometry and field data.
time_order : int
Time discretization order.
space_order : int
Space discretization order.
"""
dt = model.grid.time_dim.spacing
m = model.m
time_order = 1 if kernel == 'staggered' else 2
if kernel == 'staggered':
dims = model.space_dimensions
stagg_u = (-dims[-1])
stagg_v = (-dims[0], -dims[1]) if model.grid.dim == 3 else (-dims[0])
else:
stagg_u = stagg_v = None
# Create symbols for forward wavefield, source and receivers
u = TimeFunction(name='u',
grid=model.grid,
staggered=stagg_u,
save=geometry.nt if save else None,
time_order=time_order,
space_order=space_order)
v = TimeFunction(name='v',
grid=model.grid,
staggered=stagg_v,
save=geometry.nt if save else None,
time_order=time_order,
space_order=space_order)
src = PointSource(name='src',
grid=model.grid,
time_range=geometry.time_axis,
npoint=geometry.nsrc)
rec = Receiver(name='rec',
grid=model.grid,
time_range=geometry.time_axis,
npoint=geometry.nrec)
# FD kernels of the PDE
FD_kernel = kernels[(kernel, len(model.shape))]
stencils = FD_kernel(model, u, v, space_order)
# Source and receivers
stencils += src.inject(field=u.forward, expr=src * dt**2 / m)
stencils += src.inject(field=v.forward, expr=src * dt**2 / m)
stencils += rec.interpolate(expr=u + v)
# Substitute spacing terms to reduce flops
return Operator(stencils,
subs=model.spacing_map,
name='ForwardTTI',
**kwargs)
assert(len(nnz_sp_source_mask.dimensions) == (len(source_mask.dimensions)-1))
# Note:sparse_source_id is not needed as long as sparse info is kept in mask
# sp_source_id.data[inds[0],inds[1],:] = inds[2][:maxz]
id_dim = Dimension(name='id_dim')
b_dim = Dimension(name='b_dim')
# import pdb; pdb.set_trace()
v_sol = VectorTimeFunction(name='v_sol', grid=grid, space_order=so, time_order=1)
tau_sol = TensorTimeFunction(name='tau_sol', grid=grid, space_order=so, time_order=1)
save_src_fxx = TimeFunction(name='save_src_fxx', shape=(src.shape[0],
nzinds[1].shape[0]), dimensions=(src.dimensions[0], id_dim))
save_src_fyy = TimeFunction(name='save_src_fyy', shape=(src.shape[0],
nzinds[1].shape[0]), dimensions=(src.dimensions[0], id_dim))
save_src_fzz = TimeFunction(name='save_src_fzz', shape=(src.shape[0],
nzinds[1].shape[0]), dimensions=(src.dimensions[0], id_dim))
src_fxx = src.inject(field=tau_sol.forward[0, 0], expr=src)
src_fyy = src.inject(field=tau_sol.forward[1, 1], expr=src)
src_fzz = src.inject(field=tau_sol.forward[2, 2], expr=src)
save_src_fxx_term = src.inject(field=save_src_fxx[src.dimensions[0], source_id], expr=src)
save_src_fyy_term = src.inject(field=save_src_fyy[src.dimensions[0], source_id], expr=src)
save_src_fzz_term = src.inject(field=save_src_fzz[src.dimensions[0], source_id], expr=src)
op1 = Operator(save_src_fxx_term + save_src_fyy_term + save_src_fzz_term)
def adjoint(self,
rec,
srca=None,
p=None,
r=None,
vp=None,
epsilon=None,
delta=None,
theta=None,
phi=None,
save=None,
kernel='centered',
**kwargs):
"""
Adjoint modelling function that creates the necessary
data objects for running an adjoint modelling operator.
Parameters
----------
geometry : AcquisitionGeometry
Geometry object that contains the source (SparseTimeFunction) and
receivers (SparseTimeFunction) and their position.
p : TimeFunction, optional
The computed wavefield first component.
r : TimeFunction, optional
The computed wavefield second component.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
Returns
-------
Adjoint source, wavefield and performance summary.
"""
if kernel != 'centered':
raise ValueError(
'Only centered kernel is supported for the adjoint')
time_order = 2
stagg_p = stagg_r = None
# Source term is read-only, so re-use the default
srca = srca or self.geometry.new_src(name='srca', src_type=None)
# Create the wavefield if not provided
if p is None:
p = TimeFunction(name='p',
grid=self.model.grid,
staggered=stagg_p,
time_order=time_order,
space_order=self.space_order)
# Create the wavefield if not provided
if r is None:
r = TimeFunction(name='r',
grid=self.model.grid,
staggered=stagg_r,
time_order=time_order,
space_order=self.space_order)
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(
self.model.physical_params(vp=vp,
epsilon=epsilon,
delta=delta,
theta=theta,
phi=phi))
if self.model.dim < 3:
kwargs.pop('phi', None)
# Execute operator and return wavefield and receiver data
summary = self.op_adj().apply(srca=srca,
rec=rec,
p=p,
r=r,
dt=kwargs.pop('dt', self.dt),
**kwargs)
return srca, p, r, summary
def td_born_hessian(model_pert_in, model_pert_out, src_coords, vel, geometry, solver, params, dt=None):
"""
@Params
model_pert_in: float32 numpy array of size (Nx, Nz).
model_pert_out: float32 numpy array of size (Nx, Nz). This will be the output. Assumed to be zeros.
src_coords: float32 numpy array of size (Ns, 2) with the source coordinates.
vel: background velocity model object (The object should be the same type as returned by create_model() function)
geometry: geometry object
solver: solver object
params: python dict of parameters
Note: The receiver coordinates are not needed since it is assumed that they remain fixed for all sources.
So the receiver information should already be available to the solver object.
"""
# Get params
nbl = params["nbl"]
nx = params["Nx"]
nz = params["Nz"]
ns = params["Ns"]
nr = params["Nr"]
nt = params["Nt"]
space_order = params["so"]
time_order = params["to"]
offset = nbl + space_order
if dt is None:
dt = vel.critical_dt
# Create padded model perturbation with halo and copy model_pert_in into model_pert_padded
model_pert_padded = np.zeros((nt, nx + 2 * offset, nz + 2 * offset), dtype=np.float32)
model_pert_padded[:, offset:nx + offset, offset:nz + offset] = model_pert_in
# 1. Allocate space for Born modeled data for 1 source
# 2. Allocate time function to store intermediate wavefields
born_data = np.zeros((1, nt, nr), dtype=np.float32)
u1 = TimeFunction(
name='u',
grid=vel.grid,
time_order=time_order,
space_order=space_order,
save=nt
)
u2 = TimeFunction(
name='u',
grid=vel.grid,
time_order=time_order,
space_order=space_order,
save=nt
)
# Second derivative filter stencil
laplacian_filter = np.asarray([1, -2, 1], dtype=np.float32) / (dt ** 2.0)
# Initialize forward and adjoint operators
op_fwd = td_born_forward_op(
model=vel,
geometry=geometry,
time_order=time_order,
space_order=space_order
)
op_adjoint = td_born_adjoint_op(
model=vel,
geometry=geometry,
time_order=time_order,
space_order=space_order
)
for i in range(ns):
# Update source location
geometry.src_positions[0, :] = src_coords[i, :]
# Get Born modeled data for the current shot
_, u0, _ = solver.forward(vp=vel.vp, save=True)
spim.convolve1d(
input=u0.data_with_halo,
weights=laplacian_filter,
axis=0,
output=u0.data_with_halo[:, :, :],
mode='nearest'
)
op_fwd.apply(u0=u0, u=u1, dm=model_pert_padded, born_data_rec=born_data[0, :, :], dt=dt)
# Create adjoint image
op_adjoint.apply(u=u2, born_data_rec=born_data[0, :, :], dt=dt)
# Add to model_pert_out
model_pert_out += \
u2.data_with_halo[:, offset:nx + offset, offset:nz + offset] * \
u0.data_with_halo[:, offset:nx + offset, offset:nz + offset]
def forward(model, save=False, space_order=12, sub=None, norec=False, fs=False):
clear_cache()
# Parameters
s = model.grid.stepping_dim.spacing
nt = 10
time_range = TimeAxis(start=0, num=nt, step=1)
m, damp, epsilon, delta, theta, phi, rho = (model.m, model.damp, model.epsilon,
model.delta, model.theta, model.phi,
model.rho)
m = m * rho
# Tilt and azymuth setup
ang0 = cos(theta)
ang1 = sin(theta)
ang2 = cos(phi)
ang3 = sin(phi)
# Create the forward wavefield
if sub is not None and (sub[0] > 1 or sub[1] > 1):
usave, vsave = subsampled(model, nt, space_order, t_sub=sub[0], space_sub=sub[1])
u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order)
v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order)
eq_save = [Eq(usave, u), Eq(vsave, v)]
elif save:
u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order,
save=nt)
v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order,
save=nt)
eq_save = []
else:
u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order)
v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order)
eq_save = []
H0, H1 = kernel_zhang_fwd(u, v, ang0, ang1, ang2, ang3, epsilon, delta, rho)
# Stencils
s = model.grid.stepping_dim.spacing
stencilp = damp * (2 * u - damp * u.backward + s**2 / m * H0)
stencilr = damp * (2 * v - damp * v.backward + s**2 / m * H1)
first_stencil = Eq(u.forward, stencilp)
second_stencil = Eq(v.forward, stencilr)
expression = [first_stencil, second_stencil]
# Source symbol with input wavelet
src = Receiver(name='src', grid=model.grid, time_range=time_range, npoint=1)
src_term = src.inject(field=u.forward, expr=src.dt * s**2 / m)
src_term += src.inject(field=v.forward, expr=src.dt * s**2 / m)
expression += src_term
if fs:
expression += freesurface(u, model.nbpml)
expression += freesurface(v, model.nbpml)
if not norec:
rec = Receiver(name='rec', grid=model.grid, time_range=time_range, npoint=2)
expression += rec.interpolate(expr=u + v)
kwargs = {'dse': 'aggressive', 'dle': 'advanced'}
op = Operator(expression + eq_save, subs=model.spacing_map,
name="forward", **kwargs)
return op
def test_tti(shape, space_order):
nbpml = 10
ndim = len(shape)
origin = [0. for _ in shape]
spacing = [10. for _ in shape]
# Source location
location = np.zeros((1, ndim), dtype=np.float32)
location[0, :-1] = [
origin[i] + shape[i] * spacing[i] * .5 for i in range(ndim - 1)
]
location[0, -1] = origin[-1] + 2 * spacing[-1]
# Receivers locations
receiver_coords = np.zeros((shape[0], ndim), dtype=np.float32)
receiver_coords[:,
0] = np.linspace(0,
origin[0] + (shape[0] - 1) * spacing[0],
num=shape[0])
receiver_coords[:, 1:] = location[0, 1:]
# Two layer model for true velocity
model = demo_model('layers-isotropic',
ratio=3,
shape=shape,
spacing=spacing,
nbpml=nbpml,
space_order=space_order,
epsilon=np.zeros(shape),
delta=np.zeros(shape),
theta=np.zeros(shape),
phi=np.zeros(shape))
# Define seismic data and parameters
f0 = .010
dt = model.critical_dt
t0 = 0.0
tn = 350.0
time_range = TimeAxis(start=t0, stop=tn, step=dt)
nt = time_range.num
last = (nt - 2) % 3
indlast = [(last + 1) % 3, last % 3, (last - 1) % 3]
# Adjoint test
source = RickerSource(name='src',
grid=model.grid,
f0=f0,
time_range=time_range)
receiver = Receiver(name='rec',
grid=model.grid,
time_range=time_range,
coordinates=receiver_coords)
acoustic = AcousticWaveSolver(model,
source=source,
receiver=receiver,
time_order=2,
space_order=space_order)
rec, u1, _ = acoustic.forward(save=False)
tn = 100.0
time_range = TimeAxis(start=t0, stop=tn, step=dt)
nt = time_range.num
# Source geometry
time_series = np.zeros((nt, 1))
source = PointSource(name='src',
grid=model.grid,
time_range=time_range,
data=time_series,
coordinates=location)
receiver = Receiver(name='rec',
grid=model.grid,
time_range=time_range,
coordinates=receiver_coords)
acoustic = AcousticWaveSolver(model,
source=source,
receiver=receiver,
time_order=2,
space_order=space_order)
solver_tti = AnisotropicWaveSolver(model,
source=source,
receiver=receiver,
time_order=2,
space_order=space_order)
# Create new wavefield object restart forward computation
u = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=space_order,
dtype=model.dtype)
u.data[0:3, :] = u1.data[indlast, :]
rec, _, _ = acoustic.forward(save=False, u=u)
utti = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=space_order,
dtype=model.dtype)
vtti = TimeFunction(name='v',
grid=model.grid,
time_order=2,
space_order=space_order,
dtype=model.dtype)
utti.data[0:3, :] = u1.data[indlast, :]
vtti.data[0:3, :] = u1.data[indlast, :]
rec_tti, u_tti, v_tti, _ = solver_tti.forward(u=utti, v=vtti)
res = linalg.norm(
u.data.reshape(-1) - .5 * u_tti.data.reshape(-1) -
.5 * v_tti.data.reshape(-1))
res /= linalg.norm(u.data.reshape(-1))
log("Difference between acoustic and TTI with all coefficients to 0 %f" %
res)
assert np.isclose(res, 0.0, atol=1e-4)
def test_cache_blocking_imperfect_nest(blockinner):
"""
Test that a non-perfect Iteration nest is blocked correctly.
"""
grid = Grid(shape=(4, 4, 4), dtype=np.float64)
u = TimeFunction(name='u', grid=grid, space_order=2)
v = TimeFunction(name='v', grid=grid, space_order=2)
eqns = [Eq(u.forward, v.laplace),
Eq(v.forward, u.forward.dz)]
op0 = Operator(eqns, opt='noop')
op1 = Operator(eqns, opt=('advanced', {'blockinner': blockinner}))
# First, check the generated code
trees = retrieve_iteration_tree(op1._func_table['bf0'].root)
assert len(trees) == 2
assert len(trees[0]) == len(trees[1])
assert all(i is j for i, j in zip(trees[0][:4], trees[1][:4]))
assert trees[0][4] is not trees[1][4]
assert trees[0].root.dim.is_Incr
assert trees[1].root.dim.is_Incr
assert op1.parameters[7] is trees[0][0].step
assert op1.parameters[10] is trees[0][1].step
u.data[:] = 0.2
v.data[:] = 1.5
op0(time_M=0)
u1 = TimeFunction(name='u1', grid=grid, space_order=2)
v1 = TimeFunction(name='v1', grid=grid, space_order=2)
u1.data[:] = 0.2
v1.data[:] = 1.5
op1(u=u1, v=v1, time_M=0)
assert np.all(u.data == u1.data)
assert np.all(v.data == v1.data)
def jacobian_adjoint(self,
rec,
u,
v=None,
grad=None,
precon=None,
vp=None,
checkpointing=False,
**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.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u : TimeFunction
Full wavefield `u` (created with save=True).
v : TimeFunction, optional
Stores the computed wavefield.
grad : Function, optional
Stores the gradient field.
vp : Function or float, optional
The time-constant velocity.
Returns
-------
Gradient field and performance summary.
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
v = v or TimeFunction(name='v',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
# Pick vp from model unless explicitly provided
vp = vp or self.model.vp
if checkpointing:
u = TimeFunction(name='u',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src,
u=u,
vp=vp,
dt=dt)
wrap_rev = CheckpointOperator(self.op_grad(save=False),
u=u,
v=v,
vp=vp,
rec=rec,
dt=dt,
grad=grad)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
rec.data.shape[0] - 2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
if precon is not None:
precon = precon #don't know if this is needded but I'll just follow style
summary = self.op_grad_precon().apply(rec=rec,
grad=grad,
precon=precon,
v=v,
u=u,
vp=vp,
dt=dt,
**kwargs)
else:
summary = self.op_grad().apply(rec=rec,
grad=grad,
v=v,
u=u,
vp=vp,
dt=dt,
**kwargs)
return grad, summary
# dt is defined using Courant condition (c = 1)
# dt = 0.2*(L/(shape[0]-1)) # Timestep is half critical dt (0.0025)
# dt = 5e-4
dt = 0.2 * L / (l - 1)
t_end = L # Standing wave will cycle twice in time L
ns = int(t_end / dt) # Number of timesteps = total time/timestep size
grid = Grid(shape=(l), extent=(L))
time = grid.time_dim
x = grid.dimensions[0]
# Set up function and stencil
u_dev = TimeFunction(name="u_dev",
grid=grid,
space_order=so_dev,
time_order=to_dev,
save=ns + 1)
# Functions for initalizing standing square wave
def b_n_inner(
x,
n): # Inner part of b_n to allow for scipy.integrate.quad to be used
return square_init(x) * sin(2. * n * pi * x / L)
def b_n_calc(n):
return (4. / L) * quad(b_n_inner, 0, L / 2., args=(n))[0]