def test_haloupdate_not_requried(self):
grid = Grid(shape=(4, 4))
u = TimeFunction(name='u',
grid=grid,
space_order=4,
time_order=2,
save=None)
v = TimeFunction(name='v',
grid=grid,
space_order=0,
time_order=0,
save=5)
g = Function(name='g', grid=grid, space_order=0)
i = Function(name='i', grid=grid, space_order=0)
shift = Constant(name='shift', dtype=np.int32)
step = Eq(u.forward, u - u.backward + 1)
g_inc = Inc(g, u * v.subs(grid.time_dim, grid.time_dim - shift))
i_inc = Inc(i, (v * v).subs(grid.time_dim, grid.time_dim - shift))
op = Operator([step, g_inc, i_inc])
# No stencil in the expressions, so no halo update required!
calls = FindNodes(Call).visit(op)
assert len(calls) == 0
def GradientOperator(model, source, receiver, space_order=4, save=True,
kernel='OT2', **kwargs):
"""
Constructor method for the gradient operator in an acoustic media
:param model: :class:`Model` object containing the physical parameters
:param source: :class:`PointData` object containing the source geometry
:param receiver: :class:`PointData` object containing the acquisition geometry
:param time_order: Time discretization order
:param space_order: Space discretization order
"""
m, damp = model.m, model.damp
# Gradient symbol and wavefield symbols
grad = Function(name='grad', grid=model.grid)
u = TimeFunction(name='u', grid=model.grid, save=source.nt if save
else None, time_order=2, space_order=space_order)
v = TimeFunction(name='v', grid=model.grid, save=None,
time_order=2, space_order=space_order)
rec = Receiver(name='rec', grid=model.grid,
time_range=receiver.time_range, npoint=receiver.npoint)
s = model.grid.stepping_dim.spacing
eqn = iso_stencil(v, m, s, damp, kernel, forward=False)
if kernel == 'OT2':
gradient_update = Inc(grad, - u.dt2 * v)
elif kernel == 'OT4':
gradient_update = Inc(grad, - (u.dt2 + s**2 / 12.0 * u.laplace2(m**(-2))) * v)
# Add expression for receiver injection
receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)
# Substitute spacing terms to reduce flops
return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map,
name='Gradient', **kwargs)
def test_incs_no_atomic(self):
"""
Test that `Inc`'s don't get a `#pragma omp atomic` if performing
an increment along a fully parallel loop.
"""
grid = Grid(shape=(8, 8, 8))
x, y, z = grid.dimensions
t = grid.stepping_dim
f = Function(name='f', grid=grid)
u = TimeFunction(name='u', grid=grid)
v = TimeFunction(name='v', grid=grid)
# Format: u(t, x, nastyness) += 1
uf = u[t, x, f, z]
# All loops get collapsed, but the `y` and `z` loops are PARALLEL_IF_ATOMIC,
# hence an atomic pragma is expected
op0 = Operator(Inc(uf, 1),
opt=('advanced', {
'openmp': True,
'par-collapse-ncores': 1
}))
assert 'collapse(3)' in str(op0)
assert 'atomic' in str(op0)
# Now only `x` is parallelized
op1 = Operator([Eq(v[t, x, 0, 0], v[t, x, 0, 0] + 1),
Inc(uf, 1)],
opt=('advanced', {
'openmp': True,
'par-collapse-ncores': 1
}))
assert 'collapse(1)' in str(op1)
assert 'atomic' not in str(op1)
def GradientOperator(model,
geometry,
space_order=4,
save=True,
kernel='OT2',
**kwargs):
"""
Construct a gradient 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.
space_order : int, optional
Space discretization order.
save : int or Buffer, optional
Option to store the entire (unrolled) wavefield.
kernel : str, optional
Type of discretization, centered or shifted.
"""
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.biharmonic(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 forward_freq_modeling(model, src_coords, wavelet, rec_coords, freq, space_order=8, nb=40, dt=None, factor=None):
# Forward modeling with on-the-fly DFT of forward wavefields
clear_cache()
# Parameters
nt = wavelet.shape[0]
if dt is None:
dt = model.critical_dt
m, rho, damp = model.m, model.rho, model.damp
freq_dim = Dimension(name='freq_dim')
time = model.grid.time_dim
if factor is None:
factor = int(1 / (dt*4*np.max(freq)))
tsave = ConditionalDimension(name='tsave', parent=model.grid.time_dim, factor=factor)
if factor==1:
tsave = time
else:
tsave = ConditionalDimension(name='tsave', parent=model.grid.time_dim, factor=factor)
print("DFT subsampling factor: ", factor)
# Create wavefields
nfreq = freq.shape[0]
u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order)
f = Function(name='f', dimensions=(freq_dim,), shape=(nfreq,))
f.data[:] = freq[:]
ufr = Function(name='ufr', dimensions=(freq_dim,) + u.indices[1:], shape=(nfreq,) + model.shape_domain)
ufi = Function(name='ufi', dimensions=(freq_dim,) + u.indices[1:], shape=(nfreq,) + model.shape_domain)
ulaplace, rho = acoustic_laplacian(u, rho)
# Set up PDE and rearrange
stencil = damp * (2.0 * u - damp * u.backward + dt**2 * rho / m * ulaplace)
expression = [Eq(u.forward, stencil)]
expression += [Inc(ufr, factor*u*cos(2*np.pi*f*tsave*factor*dt))]
expression += [Inc(ufi, -factor*u*sin(2*np.pi*f*tsave*factor*dt))]
# Source symbol with input wavelet
src = PointSource(name='src', grid=model.grid, ntime=nt, coordinates=src_coords)
src.data[:] = wavelet[:]
src_term = src.inject(field=u.forward, expr=src * dt**2 / m)
# Data is sampled at receiver locations
rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords)
rec_term = rec.interpolate(expr=u)
# Create operator and run
set_log_level('ERROR')
expression += src_term + rec_term
subs = model.spacing_map
subs[u.grid.time_dim.spacing] = dt
op = Operator(expression, subs=subs, dse='advanced', dle='advanced',
name="Forward%s" % randint(1e5))
op()
return rec.data, ufr, ufi
def otf_dft(u, freq, dt, factor=None):
"""
On the fly DFT wavefield (frequency slices) and expression
Parameters
----------
u: TimeFunction or Tuple
Forward wavefield
freq: Array
Array of frequencies for on-the-fly DFT
factor: int
Subsampling factor for DFT
"""
if freq is None:
return [], None
# init
dft_modes = []
# Subsampled dft time axis
time = as_tuple(u)[0].grid.time_dim
tsave, factor = sub_time(time, factor, dt=dt, freq=freq)
# Frequencies
nfreq = np.shape(freq)[0]
freq_dim = DefaultDimension(name='freq_dim', default_value=nfreq)
f = Function(name='f', dimensions=(freq_dim, ), shape=(nfreq, ))
f.data[:] = np.array(freq[:])
# Get fourier atoms to avoid shitty perf
cf = TimeFunction(name="cf",
dimensions=(as_tuple(u)[0].grid.stepping_dim, freq_dim),
shape=(3, nfreq),
time_order=2)
sf = TimeFunction(name="sf",
dimensions=(as_tuple(u)[0].grid.stepping_dim, freq_dim),
shape=(3, nfreq),
time_order=2)
# Pulsation
omega_t = 2 * np.pi * f * tsave * factor * dt
dft = [Eq(cf, cos(omega_t)), Eq(sf, sin(omega_t))]
for wf in as_tuple(u):
ufr = Function(name='ufr%s' % wf.name,
dimensions=(freq_dim, ) + wf.indices[1:],
grid=wf.grid,
shape=(nfreq, ) + wf.shape[1:])
ufi = Function(name='ufi%s' % wf.name,
dimensions=(freq_dim, ) + wf.indices[1:],
grid=wf.grid,
shape=(nfreq, ) + wf.shape[1:])
dft += [Inc(ufr, factor * cf * wf)]
dft += [Inc(ufi, -factor * sf * wf)]
dft_modes += [(ufr, ufi)]
return dft, dft_modes
def grad(nx, ny, nchi, ncho, n, m, xdat, dydat):
# Image size
dt = np.float32
x, y, ci, co = (SpaceDimension("x"), SpaceDimension("y"), Dimension("ci"),
Dimension("co"))
grid = Grid((nx, ny), dtype=dt, dimensions=(x, y))
# Image
X = Function(name="xin", dimensions=(ci, x, y),
shape=(nchi, nx, ny), grid=grid, space_order=n//2)
# Output
dy = Function(name="dy", dimensions=(co, x, y),
shape=(ncho, nx, ny), grid=grid, space_order=n//2)
# Weights
i, j = Dimension("i"), Dimension("j")
dW = Function(name="dW", dimensions=(co, ci, i, j),
shape=(ncho, nchi, n, m), grid=grid)
# Gradient
grad_eq = Inc(dW[co, ci, i, j], dy[co, x, y]*X[ci, x+i-n//2, y+j-m//2])
op = Operator(grad_eq)
op.cfunction
X.data[:] = xdat[:]
dy.data[:] = dydat[:]
return op
def test_array_reduction(self, so, dim):
"""
Test generation of OpenMP reduction clauses involving Function's.
"""
grid = Grid(shape=(3, 3, 3))
d = grid.dimensions[dim]
f = Function(name='f',
shape=(3, ),
dimensions=(d, ),
grid=grid,
space_order=so)
u = TimeFunction(name='u', grid=grid)
op = Operator(Inc(f, u + 1),
opt=('openmp', {
'par-collapse-ncores': 1
}))
iterations = FindNodes(Iteration).visit(op)
assert "reduction(+:f[0:f_vec->size[0]])" in iterations[1].pragmas[
0].value
try:
op(time_M=1)
except:
# Older gcc <6.1 don't support reductions on array
info("Un-supported older gcc version for array reduction")
assert True
return
assert np.allclose(f.data, 18)
def GradientOperator(model,
geometry,
space_order=4,
kernel='sls',
time_order=2,
save=True,
**kwargs):
"""
Construct a gradient 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.
space_order : int, optional
Space discretization order.
save : int or Buffer, optional
Option to store the entire (unrolled) wavefield.
kernel : selects a visco-acoustic equation from the options below:
sls (Standard Linear Solid) :
1st order - Blanch and Symes (1995) / Dutta and Schuster (2014)
viscoacoustic equation
2nd order - Bai et al. (2014) viscoacoustic equation
ren - Ren et al. (2014) viscoacoustic equation
deng_mcmechan - Deng and McMechan (2007) viscoacoustic equation
Defaults to sls 2nd order.
"""
# Gradient symbol and wavefield symbols
save_t = geometry.nt if save else None
grad = Function(name='grad', grid=model.grid)
p = TimeFunction(name='p',
grid=model.grid,
time_order=2,
space_order=space_order,
save=save_t,
staggered=NODE)
pa = TimeFunction(name='pa',
grid=model.grid,
time_order=time_order,
space_order=space_order,
staggered=NODE)
# Equations kernels
eq_kernel = kernels[kernel]
eqn = eq_kernel(model, geometry, pa, forward=False, save=False, **kwargs)
gradient_update = Inc(grad, -p.dt2 * pa)
# Add expression for receiver injection
_, recterm = src_rec(pa, model, geometry, forward=False)
# Substitute spacing terms to reduce flops
return Operator(eqn + recterm + [gradient_update],
subs=model.spacing_map,
name='Gradient',
**kwargs)
def test_xcor_from_saved(self, opt, gpu_fit):
nt = 10
grid = Grid(shape=(300, 300, 300))
time_dim = grid.time_dim
period = 2
factor = Constant(name='factor', value=period, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
g = Function(name='g', grid=grid)
v = TimeFunction(name='v', grid=grid)
usave = TimeFunction(name='usave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
# For the given `nt` and grid shape, `usave` is roughly 4*5*300**3=~ .5GB of data
for i in range(int(nt//period)):
usave.data[i, :] = i
v.data[:] = i*2 + 1
# Assuming nt//period=5, we are computing, over 5 iterations:
# g = 4*4 [time=8] + 3*3 [time=6] + 2*2 [time=4] + 1*1 [time=2]
op = Operator([Eq(v.backward, v - 1), Inc(g, usave*(v/2))],
opt=(opt, {'gpu-fit': usave if gpu_fit else None}))
op.apply(time_M=nt-1)
assert np.all(g.data == 30)
def otf_dft(u, freq, dt, factor=None):
"""
On the fly DFT wavefield (frequency slices) and expression
Parameters
----------
u: TimeFunction or Tuple
Forward wavefield
freq: Array
Array of frequencies for on-the-fly DFT
factor: int
Subsampling factor for DFT
"""
if freq is None:
return [], None
# init
dft = []
dft_modes = []
# Subsampled dft time axis
time = as_tuple(u)[0].grid.time_dim
tsave, factor = sub_time(time, factor, dt=dt, freq=freq)
# Frequencies
nfreq = freq.shape[0]
freq_dim = DefaultDimension(name='freq_dim', default_value=nfreq)
f = Function(name='f', dimensions=(freq_dim, ), shape=(nfreq, ))
f.data[:] = freq[:]
# Pulsation
omega_t = 2 * np.pi * f * tsave * factor * dt
for wf in as_tuple(u):
ufr = Function(name='ufr%s' % wf.name,
dimensions=(freq_dim, ) + wf.indices[1:],
grid=wf.grid,
shape=(nfreq, ) + wf.shape[1:])
ufi = Function(name='ufi%s' % wf.name,
dimensions=(freq_dim, ) + wf.indices[1:],
grid=wf.grid,
shape=(nfreq, ) + wf.shape[1:])
dft += [Inc(ufr, factor * cos(omega_t) * wf)]
dft += [Inc(ufi, -factor * sin(omega_t) * wf)]
dft_modes += [(ufr, ufi)]
return dft, dft_modes
def test_from_different_nests(self):
"""
Check that aliases arising from two sets of equations A and B,
characterized by a flow dependence, are scheduled within A's and B's
loop nests respectively.
"""
grid = Grid(shape=(3, 3, 3))
x, y, z = grid.dimensions # noqa
t = grid.stepping_dim
i = Dimension(name='i')
f = Function(name='f', grid=grid)
f.data_with_halo[:] = 1.
g = Function(name='g', shape=(3, ), dimensions=(i, ))
g.data[:] = 2.
u = TimeFunction(name='u', grid=grid, space_order=3)
v = TimeFunction(name='v', grid=grid, space_order=3)
# Leads to 3D aliases
eqns = [
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)),
Inc(u[t + 1, i, i, i], g + 1),
Eq(v.forward,
((v[t, x, y, z] + v[t, x + 1, y + 1, z + 1]) * 3 * u.forward +
(v[t, x + 2, y + 2, z + 2] + v[t, x + 3, y + 3, z + 3]) * 3 *
u.forward + 1))
]
op0 = Operator(eqns, dse='noop', dle=('noop', {'openmp': True}))
op1 = Operator(eqns,
dse='aggressive',
dle=('advanced', {
'openmp': True
}))
# Check code generation
assert 'bf0' in op1._func_table
assert 'bf1' in op1._func_table
trees = retrieve_iteration_tree(op1._func_table['bf0'].root)
assert len(trees) == 2
assert trees[0][-1].nodes[0].body[0].write.is_Array
assert trees[1][-1].nodes[0].body[0].write is u
trees = retrieve_iteration_tree(op1._func_table['bf1'].root)
assert len(trees) == 2
assert trees[0][-1].nodes[0].body[0].write.is_Array
assert trees[1][-1].nodes[0].body[0].write is v
# 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 initialize_damp(damp, padsizes, spacing, abc_type="damp", fs=False):
"""
Initialize damping field with an absorbing boundary layer.
Parameters
----------
damp : Function
The damping field for absorbing boundary condition.
nbl : int
Number of points in the damping layer.
spacing :
Grid spacing coefficient.
mask : bool, optional
whether the dampening is a mask or layer.
mask => 1 inside the domain and decreases in the layer
not mask => 0 inside the domain and increase in the layer
"""
eqs = [Eq(damp, 1.0 if abc_type == "mask" else 0.0)]
for (nbl, nbr), d in zip(padsizes, damp.dimensions):
if not fs or d is not damp.dimensions[-1]:
dampcoeff = 1.5 * np.log(1.0 / 0.001) / (nbl)
# left
dim_l = SubDimension.left(name='abc_%s_l' % d.name,
parent=d,
thickness=nbl)
pos = Abs((nbl - (dim_l - d.symbolic_min) + 1) / float(nbl))
val = dampcoeff * (pos - sin(2 * np.pi * pos) / (2 * np.pi))
val = -val if abc_type == "mask" else val
eqs += [Inc(damp.subs({d: dim_l}), val / d.spacing)]
# right
dampcoeff = 1.5 * np.log(1.0 / 0.001) / (nbr)
dim_r = SubDimension.right(name='abc_%s_r' % d.name,
parent=d,
thickness=nbr)
pos = Abs((nbr - (d.symbolic_max - dim_r) + 1) / float(nbr))
val = dampcoeff * (pos - sin(2 * np.pi * pos) / (2 * np.pi))
val = -val if abc_type == "mask" else val
eqs += [Inc(damp.subs({d: dim_r}), val / d.spacing)]
Operator(eqs, name='initdamp')()
def initialize_damp(damp, nbl, spacing, mask=False):
"""
Initialise damping field with an absorbing boundary layer.
Parameters
----------
damp : Function
The damping field for absorbing boundary condition.
nbl : int
Number of points in the damping layer.
spacing :
Grid spacing coefficient.
mask : bool, optional
whether the dampening is a mask or layer.
mask => 1 inside the domain and decreases in the layer
not mask => 0 inside the domain and increase in the layer
"""
dampcoeff = 1.5 * np.log(1.0 / 0.001) / (40)
eqs = [Eq(damp, 1.0)] if mask else []
for d in damp.dimensions:
# left
dim_l = SubDimension.left(name='abc_%s_l' % d.name,
parent=d,
thickness=nbl)
pos = Abs((nbl - (dim_l - d.symbolic_min) + 1) / float(nbl))
val = dampcoeff * (pos - sin(2 * np.pi * pos) / (2 * np.pi))
val = -val if mask else val
eqs += [Inc(damp.subs({d: dim_l}), val / d.spacing)]
# right
dim_r = SubDimension.right(name='abc_%s_r' % d.name,
parent=d,
thickness=nbl)
pos = Abs((nbl - (d.symbolic_max - dim_r) + 1) / float(nbl))
val = dampcoeff * (pos - sin(2 * np.pi * pos) / (2 * np.pi))
val = -val if mask else val
eqs += [Inc(damp.subs({d: dim_r}), val / d.spacing)]
# TODO: Figure out why yask doesn't like it with dse/dle
Operator(eqs, name='initdamp', dse='noop', dle='noop')()
def JacobianAdjOperator(model, geometry, space_order=4,
save=True, **kwargs):
"""
Construct a linearized JacobianAdjoint modeling Operator in a TTI media.
Parameters
----------
model : Model
Object containing the physical parameters.
geometry : AcquisitionGeometry
Geometry object that contains the source (SparseTimeFunction) and
receivers (SparseTimeFunction) and their position.
space_order : int, optional
Space discretization order.
save : int or Buffer, optional
Option to store the entire (unrolled) wavefield.
"""
dt = model.grid.stepping_dim.spacing
m = model.m
time_order = 2
# Gradient symbol and wavefield symbols
u0 = TimeFunction(name='u0', grid=model.grid, save=geometry.nt if save
else None, time_order=time_order, space_order=space_order)
v0 = TimeFunction(name='v0', grid=model.grid, save=geometry.nt if save
else None, time_order=time_order, space_order=space_order)
du = TimeFunction(name="du", grid=model.grid, save=None,
time_order=time_order, space_order=space_order)
dv = TimeFunction(name="dv", grid=model.grid, save=None,
time_order=time_order, space_order=space_order)
dm = Function(name="dm", grid=model.grid)
rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis,
npoint=geometry.nrec)
# FD kernels of the PDE
FD_kernel = kernels[('centered', len(model.shape))]
eqn = FD_kernel(model, du, dv, space_order, forward=False)
dm_update = Inc(dm, - (u0 * du.dt2 + v0 * dv.dt2))
# Add expression for receiver injection
rec_term = rec.inject(field=du.backward, expr=rec * dt**2 / m)
rec_term += rec.inject(field=dv.backward, expr=rec * dt**2 / m)
# Substitute spacing terms to reduce flops
return Operator(eqn + rec_term + [dm_update], subs=model.spacing_map,
name='GradientTTI', **kwargs)
def test_nofission_as_illegal():
"""
Test there's no fission if dependencies would break.
"""
grid = Grid(shape=(20, 20))
x, y = grid.dimensions
f = Function(name='f', grid=grid, dimensions=(y, ), shape=(20, ))
u = TimeFunction(name='u', grid=grid)
v = TimeFunction(name='v', grid=grid)
eqns = [Inc(f, v + 1.), Eq(u.forward, f[y + 1] + 1.)]
op = Operator(eqns, opt='fission')
assert_structure(op, ['t,x,y', 't,x,y'], 't,x,y,y')
def test_avoid_redundant_haloupdate(self):
grid = Grid(shape=(12,))
x = grid.dimensions[0]
t = grid.stepping_dim
i = Dimension(name='i')
j = Dimension(name='j')
f = TimeFunction(name='f', grid=grid)
g = Function(name='g', grid=grid)
op = Operator([Eq(f.forward, f[t, x-1] + f[t, x+1] + 1.),
Inc(f[t+1, i], 1.), # no halo update as it's an Inc
Eq(g, f[t, j] + 1)]) # access `f` at `t`, not `t+1`!
calls = FindNodes(Call).visit(op)
assert len(calls) == 1
def test_array_reduction(self, so, dim):
"""
Test generation of OpenMP reduction clauses involving Function's.
"""
grid = Grid(shape=(3, 3, 3))
d = grid.dimensions[dim]
f = Function(name='f',
shape=(3, ),
dimensions=(d, ),
grid=grid,
space_order=so)
u = TimeFunction(name='u', grid=grid)
op = Operator(Inc(f, u + 1),
opt=('openmp', {
'par-collapse-ncores': 1
}))
iterations = FindNodes(Iteration).visit(op)
parallelized = iterations[dim + 1]
assert parallelized.pragmas
if parallelized is iterations[-1]:
# With the `f[z] += u[t0][x + 1][y + 1][z + 1] + 1` expr, the innermost
# `z` Iteration gets parallelized, nothing is collapsed, hence no
# reduction is required
assert "reduction" not in parallelized.pragmas[0].value
elif Ompizer._support_array_reduction(configuration['compiler']):
assert "reduction(+:f[0:f_vec->size[0]])" in parallelized.pragmas[
0].value
else:
# E.g. old GCC's
assert "atomic update" in str(iterations[-1])
try:
op(time_M=1)
except:
# Older gcc <6.1 don't support reductions on array
info("Un-supported older gcc version for array reduction")
assert True
return
assert np.allclose(f.data, 18)
def test_timeparallel_reduction(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)
op = Operator(Inc(f[0], u + 1), opt='noop')
trees = retrieve_iteration_tree(op)
assert len(trees) == 1
tree = trees[0]
assert all(i.is_ParallelRelaxed and not i.is_Parallel for i in tree)
# The time loop is not in OpenMP canonical form, so it won't be parallelized
assert not tree.root.pragmas
assert len(tree[1].pragmas) == 1
assert tree[1].pragmas[0].value ==\
'omp target teams distribute parallel for collapse(3) reduction(+:f[0])'
def test_simd_space_invariant(self):
"""
Similar to test_space_invariant_v3, testing simd vectorization happens
in the correct place.
"""
grid = Grid(shape=(10, 10, 10))
x, y, z = grid.dimensions
f = Function(name='f', grid=grid)
eq = Inc(f, cos(x * y) + cos(x * z))
op = Operator(eq, opt=('advanced', {'openmp': True}))
iterations = FindNodes(Iteration).visit(op)
assert 'omp for collapse(1) schedule(static,1)' in iterations[
0].pragmas[0].value
assert 'omp simd' in iterations[1].pragmas[0].value
assert 'omp simd' in iterations[3].pragmas[0].value
op.apply()
assert np.isclose(np.linalg.norm(f.data), 37.1458, rtol=1e-5)
def weighted_norm(u, weight=None):
"""
Space-time norm of a wavefield, split into norm in time first then in space to avoid
breaking loops
Parameters
----------
u: TimeFunction or Tuple of TimeFunction
Wavefield to take the norm of
weight: String
Spacial weight to apply
"""
grid = as_tuple(u)[0].grid
expr = grid.time_dim.spacing * sum(uu**2 for uu in as_tuple(u))
# Norm in time
norm_vy2_t = Function(name="nvy2t", grid=grid, space_order=0)
n_t = [Eq(norm_vy2_t, norm_vy2_t + expr)]
# Then norm in space
i = Dimension(name="i", )
norm_vy2 = Function(name="nvy2", shape=(1, ), dimensions=(i, ), grid=grid)
w = weight or 1
n_s = [Inc(norm_vy2[0], norm_vy2_t / w**2)]
# Return norm object and expr
return norm_vy2, (n_t, n_s)
def subsampled_gradient(factor=1, tn=2000.):
t0 = 0. # Simulation starts a t=0
shape = (100, 100)
origin = (0., 0.)
spacing = (15., 15.)
space_order = 4
vp = np.empty(shape, dtype=np.float64)
vp[:, :51] = 1.5
vp[:, 51:] = 2.5
model = Model(vp=vp, origin=origin, shape=shape, spacing=spacing,
space_order=space_order, nbl=10)
dt = model.critical_dt # Time step from model grid spacing
time_range = TimeAxis(start=t0, stop=tn, step=dt)
nt = time_range.num # number of time steps
f0 = 0.010 # Source peak frequency is 10Hz (0.010 kHz)
src = RickerSource(
name='src',
grid=model.grid,
f0=f0,
time_range=time_range)
src.coordinates.data[0, :] = np.array(model.domain_size) * .5
src.coordinates.data[0, -1] = 20. # Depth is 20m
rec = Receiver(
name='rec',
grid=model.grid,
npoint=101,
time_range=time_range) # new
rec.coordinates.data[:, 0] = np.linspace(0, model.domain_size[0], num=101)
rec.coordinates.data[:, 1] = 20. # Depth is 20m
save_elements = (nt + factor - 1) // factor
print(save_elements)
time_subsampled = ConditionalDimension(
't_sub', parent=model.grid.time_dim, factor=factor)
usave = TimeFunction(name='usave', grid=model.grid, time_order=2,
space_order=space_order, save=save_elements,
time_dim=time_subsampled)
u = TimeFunction(name="u", grid=model.grid, time_order=2,
space_order=space_order)
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,
offset=model.nbl)
rec_term = rec.interpolate(expr=u, offset=model.nbl)
fwd_op = Operator([stencil] + src_term + [Eq(usave, u)] + rec_term,
subs=model.spacing_map) # operator with snapshots
v = TimeFunction(name='v', grid=model.grid, save=None,
time_order=2, space_order=space_order)
grad = Function(name='grad', grid=model.grid)
rev_pde = model.m * v.dt2 - v.laplace + model.damp * v.dt.T
rev_stencil = Eq(v.backward, solve(rev_pde, v.backward))
gradient_update = Inc(grad, - usave.dt2 * v)
s = model.grid.stepping_dim.spacing
receivers = rec.inject(field=v.backward, expr=rec*s**2/model.m)
rev_op = Operator([rev_stencil] + receivers + [gradient_update],
subs=model.spacing_map)
fwd_op(time=nt - 2, dt=model.critical_dt)
rev_op(dt=model.critical_dt, time=nt-16)
return grad.data
zind = Scalar(name='zind', dtype=np.int32)
xb_size = Scalar(name='xb_size', dtype=np.int32)
yb_size = Scalar(name='yb_size', dtype=np.int32)
x0_blk0_size = Scalar(name='x0_blk0_size', dtype=np.int32)
y0_blk0_size = Scalar(name='y0_blk0_size', dtype=np.int32)
eq0 = Eq(sp_zi.symbolic_max,
nnz_sp_source_mask[x, y] - 1,
implicit_dims=(time, x, y))
eq1 = Eq(zind, sp_source_mask[x, y, sp_zi], implicit_dims=(time, x, y, sp_zi))
myexpr = source_mask[x, y, zind] * save_src[time, source_id[x, y, zind]]
eq2 = Inc(usol.forward[t + 1, x, y, zind],
myexpr,
implicit_dims=(time, x, y, sp_zi))
pde_2 = model.m * usol.dt2 - usol.laplace + model.damp * usol.dt
stencil_2 = Eq(usol.forward, solve(pde_2, usol.forward))
block_sizes = Function(name='block_sizes',
shape=(4, ),
dimensions=(b_dim, ),
space_order=0,
dtype=np.int32)
# import pdb; pdb.set_trace()
block_sizes.data[:] = args.bsizes
# import pdb; pdb.set_trace()
def forward(self,
src=None,
rec=None,
u=None,
v=None,
vp=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.
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.
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 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)
print("Initial Norm u", norm(u))
print("Initial Norm v", norm(v))
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 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
op = self.op_fwd(kernel, save)
print(kwargs)
summary = op.apply(src=src,
u=u,
v=v,
dt=kwargs.pop('dt', self.dt),
**kwargs)
regnormu = norm(u)
regnormv = norm(v)
print("Norm u:", regnormu)
print("Norm v:", regnormv)
if 0:
cmap = plt.cm.get_cmap("viridis")
values = u.data[0, :, :, :]
vistagrid = pv.UniformGrid()
vistagrid.dimensions = np.array(values.shape) + 1
vistagrid.spacing = (1, 1, 1)
vistagrid.origin = (0, 0, 0
) # The bottom left corner of the data set
vistagrid.cell_arrays["values"] = values.flatten(order="F")
vistaslices = vistagrid.slice_orthogonal()
vistagrid.plot(show_edges=True)
vistaslices.plot(cmap=cmap)
print("=========================================")
s_u = TimeFunction(name='s_u',
grid=self.model.grid,
space_order=self.space_order,
time_order=1)
s_v = TimeFunction(name='s_v',
grid=self.model.grid,
space_order=self.space_order,
time_order=1)
src_u = src.inject(field=s_u.forward,
expr=src * self.model.grid.time_dim.spacing**2 /
self.model.m)
src_v = src.inject(field=s_v.forward,
expr=src * self.model.grid.time_dim.spacing**2 /
self.model.m)
op_f = Operator([src_u, src_v])
op_f.apply(src=src, dt=kwargs.pop('dt', self.dt))
print("Norm s_u", norm(s_u))
print("Norm s_v", norm(s_v))
# Get the nonzero indices
nzinds = np.nonzero(s_u.data[0]) # nzinds is a tuple
assert len(nzinds) == len(self.model.grid.shape)
shape = self.model.grid.shape
x, y, z = self.model.grid.dimensions
time = self.model.grid.time_dim
t = self.model.grid.stepping_dim
source_mask = Function(name='source_mask',
shape=self.model.grid.shape,
dimensions=(x, y, z),
space_order=0,
dtype=np.int32)
source_id = Function(name='source_id',
shape=shape,
dimensions=(x, y, z),
space_order=0,
dtype=np.int32)
print("source_id data indexes start from 0 now !!!")
# source_id.data[nzinds[0], nzinds[1], nzinds[2]] = tuple(np.arange(1, len(nzinds[0])+1))
source_id.data[nzinds[0], nzinds[1],
nzinds[2]] = tuple(np.arange(len(nzinds[0])))
source_mask.data[nzinds[0], nzinds[1], nzinds[2]] = 1
# plot3d(source_mask.data, model)
# import pdb; pdb.set_trace()
print("Number of unique affected points is: %d", len(nzinds[0]))
# Assert that first and last index are as expected
assert (source_id.data[nzinds[0][0], nzinds[1][0], nzinds[2][0]] == 0)
assert (source_id.data[nzinds[0][-1], nzinds[1][-1],
nzinds[2][-1]] == len(nzinds[0]) - 1)
assert (source_id.data[nzinds[0][len(nzinds[0]) - 1],
nzinds[1][len(nzinds[0]) - 1],
nzinds[2][len(nzinds[0]) -
1]] == len(nzinds[0]) - 1)
assert (np.all(np.nonzero(source_id.data)) == np.all(
np.nonzero(source_mask.data)))
assert (np.all(np.nonzero(source_id.data)) == np.all(
np.nonzero(s_u.data[0])))
print(
"-At this point source_mask and source_id have been popoulated correctly-"
)
nnz_shape = (self.model.grid.shape[0], self.model.grid.shape[1])
nnz_sp_source_mask = Function(name='nnz_sp_source_mask',
shape=(list(nnz_shape)),
dimensions=(x, y),
space_order=0,
dtype=np.int32)
nnz_sp_source_mask.data[:, :] = source_mask.data[:, :, :].sum(2)
inds = np.where(source_mask.data == 1.)
print("Grid - source positions:", inds)
maxz = len(np.unique(inds[-1]))
# Change only 3rd dim
sparse_shape = (self.model.grid.shape[0], self.model.grid.shape[1],
maxz)
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')
save_src_u = TimeFunction(name='save_src_u',
shape=(src.shape[0], nzinds[1].shape[0]),
dimensions=(src.dimensions[0], id_dim))
save_src_v = TimeFunction(name='save_src_v',
shape=(src.shape[0], nzinds[1].shape[0]),
dimensions=(src.dimensions[0], id_dim))
save_src_u_term = src.inject(
field=save_src_u[src.dimensions[0], source_id],
expr=src * self.model.grid.time_dim.spacing**2 / self.model.m)
save_src_v_term = src.inject(
field=save_src_v[src.dimensions[0], source_id],
expr=src * self.model.grid.time_dim.spacing**2 / self.model.m)
print("Injecting to empty grids")
op1 = Operator([save_src_u_term, save_src_v_term])
op1.apply(src=src, dt=kwargs.pop('dt', self.dt))
print("Injecting to empty grids finished")
sp_zi = Dimension(name='sp_zi')
sp_source_mask = Function(name='sp_source_mask',
shape=(list(sparse_shape)),
dimensions=(x, y, sp_zi),
space_order=0,
dtype=np.int32)
# Now holds IDs
sp_source_mask.data[inds[0], inds[1], :] = tuple(
inds[-1][:len(np.unique(inds[-1]))])
assert (np.count_nonzero(sp_source_mask.data) == len(nzinds[0]))
assert (len(sp_source_mask.dimensions) == 3)
# import pdb; pdb.set_trace() .
zind = Scalar(name='zind', dtype=np.int32)
xb_size = Scalar(name='xb_size', dtype=np.int32)
yb_size = Scalar(name='yb_size', dtype=np.int32)
x0_blk0_size = Scalar(name='x0_blk0_size', dtype=np.int32)
y0_blk0_size = Scalar(name='y0_blk0_size', dtype=np.int32)
block_sizes = Function(name='block_sizes',
shape=(4, ),
dimensions=(b_dim, ),
space_order=0,
dtype=np.int32)
bsizes = (8, 8, 32, 32)
block_sizes.data[:] = bsizes
# eqxb = Eq(xb_size, block_sizes[0])
# eqyb = Eq(yb_size, block_sizes[1])
# eqxb2 = Eq(x0_blk0_size, block_sizes[2])
# eqyb2 = Eq(y0_blk0_size, block_sizes[3])
eq0 = Eq(sp_zi.symbolic_max,
nnz_sp_source_mask[x, y] - 1,
implicit_dims=(time, x, y))
# eq1 = Eq(zind, sp_source_mask[x, sp_zi], implicit_dims=(time, x, sp_zi))
eq1 = Eq(zind,
sp_source_mask[x, y, sp_zi],
implicit_dims=(time, x, y, sp_zi))
inj_u = source_mask[x, y, zind] * save_src_u[time, source_id[x, y,
zind]]
inj_v = source_mask[x, y, zind] * save_src_v[time, source_id[x, y,
zind]]
eq_u = Inc(u.forward[t + 1, x, y, zind],
inj_u,
implicit_dims=(time, x, y, sp_zi))
eq_v = Inc(v.forward[t + 1, x, y, zind],
inj_v,
implicit_dims=(time, x, y, sp_zi))
# The additional time-tiling equations
# tteqs = (eqxb, eqyb, eqxb2, eqyb2, eq0, eq1, eq_u, eq_v)
performance_map = np.array([[0, 0, 0, 0, 0]])
bxstart = 4
bxend = 17
bystart = 4
byend = 17
bstep = 16
txstart = 8
txend = 9
tystart = 8
tyend = 9
tstep = 16
# Temporal autotuning
for tx in range(txstart, txend, tstep):
# import pdb; pdb.set_trace()
for ty in range(tystart, tyend, tstep):
for bx in range(bxstart, bxend, bstep):
for by in range(bystart, byend, bstep):
block_sizes.data[:] = [tx, ty, bx, by]
eqxb = Eq(xb_size, block_sizes[0])
eqyb = Eq(yb_size, block_sizes[1])
eqxb2 = Eq(x0_blk0_size, block_sizes[2])
eqyb2 = Eq(y0_blk0_size, block_sizes[3])
u.data[:] = 0
v.data[:] = 0
print("-----")
tteqs = (eqxb, eqyb, eqxb2, eqyb2, eq0, eq1, eq_u,
eq_v)
op_tt = self.op_fwd(kernel, save, tteqs)
summary_tt = op_tt.apply(u=u,
v=v,
dt=kwargs.pop('dt', self.dt),
**kwargs)
norm_tt_u = norm(u)
norm_tt_v = norm(v)
print("Norm u:", regnormu)
print("Norm v:", regnormv)
print("Norm(tt_u):", norm_tt_u)
print("Norm(tt_v):", norm_tt_v)
print(
"===Temporal blocking======================================"
)
performance_map = np.append(performance_map, [[
tx, ty, bx, by,
summary_tt.globals['fdlike'].gflopss
]], 0)
print(performance_map)
# tids = np.unique(performance_map[:, 0])
#for tid in tids:
bids = np.where((performance_map[:, 0] == tx)
& (performance_map[:, 1] == ty))
bx_data = np.unique(performance_map[bids, 2])
by_data = np.unique(performance_map[bids, 3])
gptss_data = performance_map[bids, 4]
gptss_data = gptss_data.reshape(len(bx_data), len(by_data))
fig, ax = plt.subplots()
im = ax.imshow(gptss_data)
pause(2)
# We want to show all ticks...
ax.set_xticks(np.arange(len(bx_data)))
ax.set_yticks(np.arange(len(by_data)))
# ... and label them with the respective list entries
ax.set_xticklabels(bx_data)
ax.set_yticklabels(by_data)
ax.set_title(
"Gpts/s for fixed tile size. (Sweeping block sizes)")
fig.tight_layout()
fig.colorbar(im, ax=ax)
# ax = sns.heatmap(gptss_data, linewidth=0.5)
plt.savefig(
str(shape[0]) + str(np.int32(tx)) + str(np.int32(ty)) +
".pdf")
if 0:
cmap = plt.cm.get_cmap("viridis")
values = u.data[0, :, :, :]
vistagrid = pv.UniformGrid()
vistagrid.dimensions = np.array(values.shape) + 1
vistagrid.spacing = (1, 1, 1)
vistagrid.origin = (0, 0, 0
) # The bottom left corner of the data set
vistagrid.cell_arrays["values"] = values.flatten(order="F")
vistaslices = vistagrid.slice_orthogonal()
vistagrid.plot(show_edges=True)
vistaslices.plot(cmap=cmap)
return rec, u, v, summary
def chain_contractions(A, B, C, D, E, F, optimize):
"""``AB + AC = D, DE = F``."""
op = Operator([Inc(D, A * B + A * C), Inc(F, D * E)], dle=optimize)
op.apply()
info('Executed `AB + AC = D, DE = F`')
def mat_mat_sum(A, B, C, D, optimize):
"""``AB + AC = D``."""
op = Operator(Inc(D, A * B + A * C), dle=optimize)
op.apply()
info('Executed `AB + AC = D`')
def mat_mat(A, B, C, optimize):
"""``AB = C``."""
op = Operator(Inc(C, A * B), dle=optimize)
op.apply()
info('Executed `AB = C`')
def transpose_mat_vec(A, x, b, optimize):
"""``A -> A^T, A^Tx = b``."""
i, j = A.indices
op = Operator([Inc(b, A[j, i] * x)], dle=optimize)
op.apply()
info('Executed `A^Tx = b`')
def mat_vec(A, x, b, optimize):
"""``Ax = b``."""
op = Operator(Inc(b, A * x), dle=optimize)
op.apply()
info('Executed `Ax = b`')
def fwi_gradient(vp_in):
# AUTO
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.
u0 = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=4,
save=geometry.nt)
# MANUAL
grad_manual = Function(name="grad", grid=model.grid)
residual_man = Receiver(name='rec',
grid=model.grid,
time_range=time_axis,
coordinates=rec_coordinates)
objective_manual = 0.
u0_man = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=4,
save=nt)
for i in range(9):
# AUTO
clear_cache()
geometry.src_positions[0, :] = source_locations[i, :]
true_d, _, _ = solver.forward(vp=model.vp)
u0.data.fill(0.)
smooth_d, _, _ = solver.forward(vp=vp_in, save=True, u=u0)
residual.data[:] = smooth_d.data[:] - true_d.data[:]
objective += .5 * np.linalg.norm(residual.data.flatten())**2
solver.gradient(rec=residual, u=u0, vp=vp_in, grad=grad)
# MANUAL
# source
src_true = RickerSource(name='src',
grid=model.grid,
time_range=time_axis,
coordinates=source_locations[i, :],
npoint=1,
f0=f0)
src_term = src_true.inject(
field=u.forward,
expr=src_true * model.grid.stepping_dim.spacing**2 / model.m)
# receiver
rec_true = Receiver(name='rec',
grid=model.grid,
time_range=time_axis,
coordinates=rec_coordinates,
npoint=nreceivers)
rec_term = rec_true.interpolate(expr=u)
# operator
op_fwd = Operator(eqn + src_term + rec_term,
subs=model.spacing_map,
name='Forward')
op_fwd.apply(src=src_true,
rec=rec_true,
u=u0_man,
vp=model.vp,
dt=model.critical_dt)
u0_man.data.fill(0.)
rec_smooth = Receiver(name='rec',
grid=model.grid,
time_range=time_axis,
coordinates=rec_coordinates,
npoint=nreceivers)
op_fwd.apply(src=src_true,
rec=rec_smooth,
u=u0_man,
vp=vp_in,
dt=model.critical_dt)
# back-receiver
rec_back = Receiver(name='rec',
grid=model.grid,
time_range=time_axis,
coordinates=rec_coordinates,
npoint=nreceivers)
rec_back_term = rec_back.inject(
field=v.backward,
expr=rec_back * model.grid.stepping_dim.spacing**2 / model.m)
# gradient
gradient_update = Inc(grad_manual, -u.dt2 * v)
op_grad = Operator(eqn_back + rec_back_term + [gradient_update],
subs=model.spacing_map,
name='Gradient')
residual_man.data[:] = rec_smooth.data[:] - rec_true.data[:]
objective_manual += .5 * np.linalg.norm(residual_man.data.flatten())**2
op_grad.apply(rec=residual_man,
u=u0_man,
vp=vp_in,
dt=model.critical_dt,
grad=grad_manual)
# sanity-check -> expect for 0!
# plot_shotrecord(true_d.data[:] - rec_true.data[:], model, t0, tn)
# plot_shotrecord(smooth_d.data[:] - rec_smooth.data[:], model, t0, tn)
return objective, -grad.data, objective_manual, -grad_manual.data
