def test_compilerfunction():
grid = Grid(shape=(3, 3))
d = Dimension(name='d')
cf = TempFunction(name='f', dtype=np.float64, dimensions=grid.dimensions,
halo=((1, 1), (1, 1), (1, 1)))
pkl_cf = pickle.dumps(cf)
new_cf = pickle.loads(pkl_cf)
assert new_cf.name == cf.name
assert new_cf.dtype is np.float64
assert new_cf.halo == ((1, 1), (1, 1), (1, 1))
assert new_cf.ndim == cf.ndim
assert new_cf.dim is None
pcf = cf._make_pointer(d)
pkl_pcf = pickle.dumps(pcf)
new_pcf = pickle.loads(pkl_pcf)
assert new_pcf.name == pcf.name
assert new_pcf.dim.name == 'd'
assert new_pcf.ndim == cf.ndim + 1
def test_conditional_dimension(self):
"""
Test that ConditionalDimensions with same name but different attributes do not
alias to the same ConditionalDimension. Conversely, if the name and the attributes
are the same, they must alias to the same ConditionalDimension.
"""
i = Dimension(name='i')
ci0 = ConditionalDimension(name='ci', parent=i, factor=4)
ci1 = ConditionalDimension(name='ci', parent=i, factor=4)
assert ci0 is ci1
ci2 = ConditionalDimension(name='ci', parent=i, factor=8)
assert ci2 is not ci1
ci3 = ConditionalDimension(name='ci', parent=i, factor=4, indirect=True)
assert ci3 is not ci1
s = Scalar(name='s')
ci4 = ConditionalDimension(name='ci', parent=i, factor=4, condition=s > 3)
assert ci4 is not ci1
ci5 = ConditionalDimension(name='ci', parent=i, factor=4, condition=s > 3)
assert ci5 is ci4
def test_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 tree.root.is_Sequential
assert all(i.is_ParallelRelaxed and not i.is_Parallel
for i in tree[1:])
# 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 weighted_norm(u, weight=None):
"""
Space-time nor of a wavefield, split into norm in time first then in space to avoid
breaking loops
"""
if type(u) is tuple:
expr = u[0].grid.time_dim.spacing * (u[0]**2 + u[1]**2)
grid = u[0].grid
else:
expr = u.grid.time_dim.spacing * u**2
grid = u.grid
# Norm in time
norm_vy2_t = Function(name="nvy2t", grid=grid)
n_v = [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)
if weight is None:
n_v += [Eq(norm_vy2[0], norm_vy2[0] + norm_vy2_t)]
else:
n_v += [Eq(norm_vy2[0], norm_vy2[0] + norm_vy2_t / weight**2)]
return norm_vy2, n_v
def test_function_wo(self):
grid = Grid(shape=(3, 3, 3))
i = Dimension(name='i')
f = Function(name='f', shape=(1, ), dimensions=(i, ), grid=grid)
u = TimeFunction(name='u', grid=grid)
eqns = [Eq(u.forward, u + 1), Eq(f[0], u[0, 0, 0, 0])]
op = Operator(eqns, opt='noop')
assert len(op.body[2].header) == 1
assert len(op.body[2].footer) == 1
assert op.body[2].header[0].value ==\
('omp target enter data map(to: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert op.body[2].footer[0].contents[0].value ==\
('omp target update from(u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert op.body[2].footer[0].contents[1].value ==\
('omp target exit data map(release: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
def solver_adjust_w(I, w, dt, T, adjust_w=True):
"""
Solve u'' + w**2*u = 0 for t in (0,T], u(0)=I and u'(0)=0,
by a central finite difference method with time step dt.
"""
dt = float(dt)
Nt = int(round(T / dt))
t = Dimension('t', spacing=Constant('h_t'))
u = TimeFunction(name='u',
dimensions=(t, ),
shape=(Nt + 1, ),
space_order=2)
w_adj = w * (1 - w**2 * dt**2 / 24.) if adjust_w else w
u.data[:] = I
eqn = u.dt2 + (w**2) * u
stencil = Eq(u.forward, solve(eqn, u.forward))
op = Operator(stencil)
op.apply(h_t=dt, t_M=Nt - 1)
return u.data, np.linspace(0, Nt * dt, Nt + 1)
def test_misc_dims(self):
"""
Tests grid-independent Functions, 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_equations_mixed_densedata_timedata(self, shape, dimensions):
"""
Test that equations using a mixture of Function and TimeFunction objects
are embedded within the same time loop.
"""
grid = Grid(shape=shape, dimensions=dimensions, time_dimension=time)
a = TimeFunction(name='a', grid=grid, time_order=2, space_order=2)
p_aux = Dimension(name='p_aux', size=10)
b = Function(name='b',
shape=shape + (10, ),
dimensions=dimensions + (p_aux, ),
space_order=2)
b.data[:] = 1.0
b2 = Function(name='b2',
shape=(10, ) + shape,
dimensions=(p_aux, ) + dimensions,
space_order=2)
b2.data[:] = 1.0
eqns = [Eq(a.forward, a.laplace + 1.), Eq(b, time * b * a + b)]
eqns2 = [Eq(a.forward, a.laplace + 1.), Eq(b2, time * b2 * a + b2)]
subs = {x.spacing: 2.5, y.spacing: 1.5, z.spacing: 2.0}
op = Operator(eqns, subs=subs, dle='noop')
trees = retrieve_iteration_tree(op)
assert len(trees) == 2
assert all(trees[0][i] is trees[1][i] for i in range(3))
op2 = Operator(eqns2, subs=subs, dle='noop')
trees = retrieve_iteration_tree(op2)
assert len(trees) == 2
# Verify both operators produce the same result
op(time=10)
a.data[:] = 0.
op2(time=10)
for i in range(10):
assert (np.allclose(
b2.data[i, ...].reshape(-1) - b.data[..., i].reshape(-1), 0.))
def test_function_wo(self):
grid = Grid(shape=(3, 3, 3))
i = Dimension(name='i')
f = Function(name='f', shape=(1,), dimensions=(i,), grid=grid)
u = TimeFunction(name='u', grid=grid)
eqns = [Eq(u.forward, u + 1),
Eq(f[0], u[0, 0, 0, 0])]
op = Operator(eqns, opt='noop', language='openmp')
assert len(op.body.maps) == 1
assert op.body.maps[0].pragmas[0].value ==\
('omp target enter data map(to: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert len(op.body.unmaps) == 2
assert op.body.unmaps[0].pragmas[0].value ==\
('omp target update from(u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]])')
assert op.body.unmaps[1].pragmas[0].value ==\
('omp target exit data map(release: u[0:u_vec->size[0]]'
'[0:u_vec->size[1]][0:u_vec->size[2]][0:u_vec->size[3]]) if(devicerm)')
def test_misc_data(self):
"""
Test data insertion/indexing for Functions with mixed
distributed/replicated Dimensions.
"""
dx = Dimension(name='dx')
grid = Grid(shape=(4, 4))
x, y = grid.dimensions
glb_pos_map = grid.distributor.glb_pos_map
# Note: `grid` must be passed to `c` since `x` is a distributed dimension,
# and `grid` carries the `x` decomposition
c = Function(name='c', grid=grid, dimensions=(x, dx), shape=(4, 5))
# Data insertion
for i in range(4):
c.data[i, 0] = 1.0+i
c.data[i, 1] = 1.0+i
c.data[i, 2] = 3.0+i
c.data[i, 3] = 6.0+i
c.data[i, 4] = 5.0+i
# Data indexing
if LEFT in glb_pos_map[x]:
assert(np.all(c.data[0] == [1., 1., 3., 6., 5.]))
assert(np.all(c.data[1] == [2., 2., 4., 7., 6.]))
else:
assert(np.all(c.data[2] == [3., 3., 5., 8., 7.]))
assert(np.all(c.data[3] == [4., 4., 6., 9., 8.]))
# Same as before, but with negative indices and non-trivial slices
if LEFT in glb_pos_map[x]:
assert(np.all(c.data[0:-3] == [1., 1., 3., 6., 5.]))
assert(np.all(c.data[-3:-2] == [2., 2., 4., 7., 6.]))
else:
assert(np.all(c.data[-2:-1] == [3., 3., 5., 8., 7.]))
assert(np.all(c.data[-1] == [4., 4., 6., 9., 8.]))
def 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 __new__(cls, *args, **kwargs):
options = kwargs.get('options', {})
if cls in _SymbolCache:
obj = sympy.Function.__new__(cls, *args, **options)
obj._cached_init()
else:
p_dim = kwargs.get('dimension',
Dimension('p_%s' % kwargs.get("name")))
npoint = kwargs.get("npoint")
coords = kwargs.get("coordinates")
if npoint is None:
if coords is None:
raise TypeError("Need either `npoint` or `coordinates`")
else:
npoint = coords.shape[0]
name = kwargs.get("name")
grid = kwargs.get("grid")
ntime = kwargs.get("ntime")
if kwargs.get("data") is None:
if ntime is None:
error('Either data or ntime are required to'
'initialise source/receiver objects')
else:
ntime = kwargs.get("ntime") or kwargs.get("data").shape[0]
# Create the underlying SparseTimeFunction object
kwargs["nt"] = ntime
kwargs['npoint'] = npoint
obj = SparseTimeFunction.__new__(cls,
dimensions=[grid.time_dim, p_dim],
**kwargs)
# If provided, copy initial data into the allocated buffer
if kwargs.get("data") is not None:
obj.data[:] = kwargs.get("data")
return obj
def test_dimension_cache():
"""
Test that :class:`Dimension`s with same name but different attributes do not
alias to the same Dimension.
"""
d0 = Dimension(name='d')
d1 = Dimension(name='d')
assert d0 is d1
s0 = Scalar(name='s0')
s1 = Scalar(name='s1')
d2 = Dimension(name='d', spacing=s0)
d3 = Dimension(name='d', spacing=s1)
assert d2 is not d3
d4 = Dimension(name='d', spacing=s1)
assert d3 is d4
d5 = Dimension(name='d', spacing=Constant(name='s1'))
assert d2 is not d5
def test_dimension(self):
"""
Test that Dimensions with same name but different attributes do not alias to
the same Dimension. Conversely, if the name and the attributes are the same,
they must alias to the same Dimension.
"""
d0 = Dimension(name='d')
d1 = Dimension(name='d')
assert d0 is d1
s0 = Scalar(name='s0')
s1 = Scalar(name='s1')
d2 = Dimension(name='d', spacing=s0)
d3 = Dimension(name='d', spacing=s1)
assert d2 is not d3
d4 = Dimension(name='d', spacing=s1)
assert d3 is d4
d5 = Dimension(name='d', spacing=Constant(name='s1'))
assert d2 is not d5
def test_zero_handling_staggered(self, side, order, spec):
"""
Check that stencils with distances of zero evaluate correctly for staggered
systems.
"""
# Unpack the spec
bc_type = spec['bcs']
deriv = spec['deriv']
goffset = spec['goffset']
eoffset = spec['eoffset']
if bc_type == 'even':
bcs = BoundaryConditions({2 * i: 0
for i in range(1 + order // 2)}, order)
else:
bcs = BoundaryConditions(
{2 * i + 1: 0
for i in range(1 + order // 2)}, order)
cache = os.path.dirname(
__file__) + '/../devitoboundary/extrapolation_cache.dat'
# stencils_lambda = get_stencils_lambda(deriv, eoffset, bcs, cache=cache)
stencils = StencilSet(deriv, eoffset, bcs, cache=cache)
lambdas = stencils.lambdaify
max_ext_points = stencils.max_ext_points
distances = np.full((10, 1, 1), -2 * order, dtype=float)
if goffset == 0.5:
distances[4, :, :] = 0.5
else:
distances[4, :, :] = 0
data = get_data_inc_reciprocals(distances, 1, 'x', goffset, eoffset)
add_distance_column(data)
right_dist = pd.notna(data.eta_l)
left_dist = pd.notna(data.eta_r)
data.loc[right_dist, 'dist'] = order
data.loc[left_dist, 'dist'] = -order
first = data.loc[left_dist]
last = data.loc[right_dist]
first = shift_grid_endpoint(first, 'x', goffset, eoffset)
last = shift_grid_endpoint(last, 'x', goffset, eoffset)
first = get_n_pts(first, 'first', order, eoffset)
last = get_n_pts(last, 'last', order, eoffset)
grid = Grid(shape=(10, 1, 1), extent=(9, 0, 0))
s_dim = Dimension(name='s')
ncoeffs = 2 * max_ext_points + 1
w_shape = grid.shape + (ncoeffs, )
w_dims = grid.dimensions + (s_dim, )
w = Function(name='w', dimensions=w_dims, shape=w_shape)
# Fill the weights using the standard stencil (for interior)
if max_ext_points > order // 2:
# Need to zero pad the standard stencil
zero_pad = max_ext_points - order // 2
w.data[:, :, :,
zero_pad:-zero_pad] = standard_stencil(deriv,
order,
offset=eoffset)
else:
w.data[:] = standard_stencil(deriv, order, offset=eoffset)
if side == 'first':
w.data[4:] = 0
fill_stencils(first, 'first', max_ext_points, lambdas, w, 10, 'x')
if side == 'last':
w.data[:5] = 0
fill_stencils(last, 'last', max_ext_points, lambdas, w, 10, 'x')
# Check against a saved correct version
# Generate filename
filename = side + '_' + str(order) + '_' + bc_type + '_' + str(
deriv) + '_' + str(goffset) + '_' + str(eoffset)
# Load reference data
reference = np.load(
os.path.dirname(__file__) + '/zero_handling_stencils/' + filename +
'.npy')
assert np.all(np.absolute(w.data - reference) < np.finfo(float).eps)
def test_fill_stencils_offset(self, offset, point_type, order, spacing):
"""
Check that offsetting the grid and boundary by the same amount results
in identical stencils for both cases. This is checked on both sides of
the boundary.
"""
spec = {2 * i: 0 for i in range(1 + order // 2)}
bcs = BoundaryConditions(spec, order)
cache = os.path.dirname(
__file__) + '/../devitoboundary/extrapolation_cache.dat'
stencils = StencilSet(2, 0, bcs, cache=cache)
lambdas = stencils.lambdaify
max_ext_points = stencils.max_ext_points
distances = np.full((10, 1, 10), -2 * order * spacing, dtype=float)
distances[4, :, :] = np.linspace(0, 0.9 * spacing, 10)
offset_distances = np.full((10, 1, 10),
-2 * order * spacing,
dtype=float)
if offset == 0.5:
# +ve stagger
offset_distances[4, :, :5] = np.linspace(0.5 * spacing,
0.9 * spacing, 5)
offset_distances[5, :, 5:] = np.linspace(0, 0.4 * spacing, 5)
else:
# -ve stagger
offset_distances[4, :, :] = np.linspace(-0.5 * spacing,
0.4 * spacing, 10)
data = get_data_inc_reciprocals(distances, spacing, 'x', 0, 0)
offset_data = get_data_inc_reciprocals(offset_distances, spacing, 'x',
offset, 0)
dmask = np.full(21, True, dtype=bool)
dmask[1] = False
data = data[dmask]
offset_data = offset_data[dmask]
add_distance_column(data)
add_distance_column(offset_data)
if point_type == 'first':
data = data[::2]
data.dist = -order // 2
offset_data = offset_data[::2]
offset_data.dist = -order // 2
# No need to drop points or shift grid endpoint, as that is done here
else:
data = data[1::2]
data.dist = order // 2
offset_data = offset_data[1::2]
offset_data.dist = order // 2
# No need to drop points or shift grid endpoint, as that is done here
# Set n_pts
data['n_pts'] = order // 2
offset_data['n_pts'] = order // 2
grid = Grid(shape=(10, 1, 10), extent=(9 * spacing, 0, 9 * spacing))
s_dim = Dimension(name='s')
ncoeffs = order + 1
w_shape = grid.shape + (ncoeffs, )
w_dims = grid.dimensions + (s_dim, )
w_normal = Function(name='w_n', dimensions=w_dims, shape=w_shape)
w_offset = Function(name='w_o', dimensions=w_dims, shape=w_shape)
fill_stencils(data, point_type, max_ext_points, lambdas, w_normal, 10,
'x')
fill_stencils(offset_data, point_type, max_ext_points, lambdas,
w_offset, 10, 'x')
if point_type == 'first':
assert np.all(np.isclose(w_normal.data[2:5], w_offset.data[2:5]))
else:
assert np.all(np.isclose(w_normal.data[5:7], w_offset.data[5:7]))
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 = (model.grid.shape[0], 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 = TimeFunction(name='save_src',
shape=(src.shape[0], nzinds[1].shape[0]),
dimensions=(src.dimensions[0], id_dim))
save_src_term = src.inject(field=save_src[src.dimensions[0], source_id],
expr=src * dt**2 / model.m)
op1 = Operator([save_src_term])
op1.apply(time=time_range.num - 1)
usol = TimeFunction(name="usol", grid=model.grid, space_order=so, time_order=2)
sp_zi = Dimension(name='sp_zi')
def x(xdim=4):
return Dimension(name='x', size=xdim)
def adjoint_y(model,
y,
src_coords,
rcv_coords,
weight_fun_pars=None,
dt=None,
space_order=8,
save=False):
"Compute adjoint wavefield v = adjoint(F(m))*y and related quantities (||v||_w, v(xsrc))"
clear_cache()
# Setting time sampling
if dt is None:
dt = model.critical_dt
# Physical parameters
m, rho, damp = model.m, model.rho, model.damp
# Setting adjoint wavefield
nt = y.shape[0]
v = TimeFunction(name="v",
grid=model.grid,
time_order=2,
space_order=space_order,
save=None if not save else nt)
# Set up PDE expression and rearrange
vlaplace, rho = laplacian(v, rho)
stencil = damp * (2.0 * v - damp * v.forward + dt**2 * rho / m * vlaplace)
expression = [Eq(v.backward, stencil)]
# Setup adjoint source injected at receiver locations
rcv = Receiver(name="rcv",
grid=model.grid,
ntime=nt,
coordinates=rcv_coords)
rcv.data[:] = y[:]
adj_src = rcv.inject(field=v.backward, expr=rcv * rho * dt**2 / m)
expression += adj_src
# Setup adjoint wavefield sampling at source locations
src = PointSource(name="src",
grid=model.grid,
ntime=nt,
coordinates=src_coords)
adj_rcv = src.interpolate(expr=v)
expression += adj_rcv
# Setup ||v||_w computation
norm_vy2_t = Function(name="nvy2t", grid=model.grid)
expression += [Inc(norm_vy2_t, Pow(v, 2))]
i = Dimension(name="i", )
norm_vy2 = Function(name="nvy2",
shape=(1, ),
dimensions=(i, ),
grid=model.grid)
if weight_fun_pars is None:
expression += [Inc(norm_vy2[0], norm_vy2_t)]
else:
weight = weight_fun(weight_fun_pars, model, src_coords)
expression += [Inc(norm_vy2[0], norm_vy2_t / weight**2)]
# Create operator and run
subs = model.spacing_map
subs[v.grid.time_dim.spacing] = dt
op = Operator(expression,
subs=subs,
dse="advanced",
dle="advanced",
name="adjoint_y")
op()
# Output
if save:
return norm_vy2.data[0], src.data, v
else:
return norm_vy2.data[0], src.data, None
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, damp = model.m, 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)
# Set up PDE and rearrange
eqn = m * u.dt2 - u.laplace + damp * u.dt
stencil = solve(eqn, u.forward, simplify=False, rational=False)[0]
expression = [Eq(u.forward, stencil)]
expression += [
Eq(ufr, ufr + factor * u * cos(2 * np.pi * f * tsave * factor * dt))
]
expression += [
Eq(ufi, 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,
offset=model.nbpml,
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, offset=model.nbpml)
# 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 forward_freq_modeling(model,
src_coords,
wavelet,
rec_coords,
freq,
space_order=8,
dt=None,
factor=None,
free_surface=False):
# 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 += [
Eq(ufr, ufr + factor * u * cos(2 * np.pi * f * tsave * factor * dt))
]
expression += [
Eq(ufi, 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
expression += src_term + rec_term
# Free surface
if free_surface is True:
expression += freesurface(u, space_order // 2, model.nbpml)
subs = model.spacing_map
subs[u.grid.time_dim.spacing] = dt
op = Operator(expression, subs=subs, dse='advanced', dle='advanced')
cf = op.cfunction
op()
return rec.data, ufr, ufi
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 test_make_cpp_parfor():
"""
Test construction of a CPP parallel for. This excites the IET construction
machinery in several ways, in particular by using Lambda nodes (to generate
C++ lambda functions) and nested Calls.
"""
class STDVectorThreads(LocalObject):
dtype = type('std::vector<std::thread>', (c_void_p, ), {})
def __init__(self):
self.name = 'threads'
class STDThread(LocalObject):
dtype = type('std::thread&', (c_void_p, ), {})
def __init__(self, name):
self.name = name
class FunctionType(LocalObject):
dtype = type('FuncType&&', (c_void_p, ), {})
def __init__(self, name):
self.name = name
# Basic symbols
nthreads = Symbol(name='nthreads', is_const=True)
threshold = Symbol(name='threshold', is_const=True)
last = Symbol(name='last', is_const=True)
first = Symbol(name='first', is_const=True)
portion = Symbol(name='portion', is_const=True)
# Composite symbols
threads = STDVectorThreads()
# Iteration helper symbols
begin = Symbol(name='begin')
l = Symbol(name='l')
end = Symbol(name='end')
# Functions
stdmax = sympy.Function('std::max')
# Construct the parallel-for body
func = FunctionType('func')
i = Dimension(name='i')
threadobj = Call(
'std::thread',
Lambda(
Iteration(Call(func.name, i), i, (begin, end - 1, 1)),
['=', Byref(func.name)],
))
threadpush = Call(FieldFromComposite('push_back', threads), threadobj)
it = Dimension(name='it')
iteration = Iteration([
LocalExpression(DummyEq(begin, it)),
LocalExpression(DummyEq(l, it + portion)),
LocalExpression(DummyEq(end, InlineIf(l > last, last, l))), threadpush
], it, (first, last, portion))
thread = STDThread('x')
waitcall = Call('std::for_each', [
Call(FieldFromComposite('begin', threads)),
Call(FieldFromComposite('end', threads)),
Lambda(Call(FieldFromComposite('join', thread.name)), [], [thread])
])
body = [
LocalExpression(DummyEq(threshold, 1)),
LocalExpression(
DummyEq(portion, stdmax(threshold, (last - first) / nthreads))),
Call(FieldFromComposite('reserve', threads), nthreads), iteration,
waitcall
]
parfor = ElementalFunction('parallel_for', body, 'void',
[first, last, func, nthreads])
assert str(parfor) == """\
def get_component_weights(data, axis, function, deriv, lambdas, interior,
max_span, eval_offset):
"""
Take a component of the distance field and return the associated weight
function.
Parameters
----------
data : ndarray
The field of the axial distance function for the specified axis
axis : int
The axis along which the stencils are orientated. Can be 0, 1, or 2
function : devito Function
The function for which stencils should be calculated
deriv : int
The order of the derivative to which the stencils pertain
lambdas : dict
The functions for stencils to be evaluated
interior : ndarray
The interior-exterior segmentation of the domain
max_span : int
The maximum span of the stencil from the center point
eval_offset : float
The relative offset at which the derivative should be evaluated.
Used for setting the default fill stencil.
Returns
-------
w : devito Function
Function containing the stencil coefficients
"""
grid_offset = get_grid_offset(function, axis)
f_grid = function.grid
dim_limit = f_grid.shape[axis]
axis_dim = 'x' if axis == 0 else 'y' if axis == 1 else 'z'
# Additional dimension for storing weights
# This will be dependent on the number of extrapolation points required
# and the space order.
dim_size = max(function.space_order // 2, max_span)
s_dim = Dimension(name='s' + str(2 * dim_size))
ncoeffs = 2 * dim_size + 1
w_shape = f_grid.shape + (ncoeffs, )
w_dims = f_grid.dimensions + (s_dim, )
w = Function(name='w_' + function.name + '_' + axis_dim + '_' + str(deriv),
dimensions=w_dims,
shape=w_shape)
# Do the initial stencil fill, padding where needs be
if max_span > function.space_order // 2:
# Need to zero pad the standard stencil
zero_pad = max_span - function.space_order // 2
w.data[:, :, :,
zero_pad:-zero_pad] = standard_stencil(deriv,
function.space_order,
offset=eval_offset)
# Needs to return a warning if padding is used for the time being
# Will need a dummy function to create the substitutions, with a higher
# order function to substitute into
warning(
"Generated stencils have been padded due to required number of"
" extrapolation points. A dummy function will be needed to"
" create the substitutions. The required order for substitution"
" is {}".format(2 * max_span))
else:
w.data[:] = standard_stencil(deriv,
function.space_order,
offset=eval_offset)
w.data[~interior] = 0
fill_val = np.amin(data)
# Still want to zero above boundary if no points in need of modification
# So skip this if no points are available
if np.any(data != fill_val):
full_data = get_data_inc_reciprocals(data, f_grid.spacing[axis],
axis_dim, grid_offset,
eval_offset)
add_distance_column(full_data)
first, last, double, paired_left, paired_right \
= split_types(full_data, axis_dim, f_grid.shape[axis])
# Need to drop exterior points and shift grid endpoint
first = drop_outside_points(first, interior)
last = drop_outside_points(last, interior)
double = drop_outside_points(double, interior)
paired_left = drop_outside_points(paired_left, interior)
paired_right = drop_outside_points(paired_right, interior)
first = shift_grid_endpoint(first, axis_dim, grid_offset, eval_offset)
last = shift_grid_endpoint(last, axis_dim, grid_offset, eval_offset)
double = shift_grid_endpoint(double, axis_dim, grid_offset,
eval_offset)
paired_left = shift_grid_endpoint(paired_left, axis_dim, grid_offset,
eval_offset)
paired_right = shift_grid_endpoint(paired_right, axis_dim, grid_offset,
eval_offset)
double = apply_dist(double, 'double')
paired_left = apply_dist(paired_left, 'paired_left')
paired_right = apply_dist(paired_right, 'paired_right')
# Fill the stencils
if len(first.index) != 0:
first = get_n_pts(first, 'first', function.space_order,
eval_offset)
fill_stencils(first, 'first', max_span, lambdas, w, dim_limit,
axis_dim)
if len(last.index) != 0:
last = get_n_pts(last, 'last', function.space_order, eval_offset)
fill_stencils(last, 'last', max_span, lambdas, w, dim_limit,
axis_dim)
if len(double.index) != 0:
double = get_n_pts(double, 'double', function.space_order,
eval_offset)
fill_stencils(double, 'double', max_span, lambdas, w, dim_limit,
axis_dim)
if len(paired_left.index) != 0:
paired_left = get_n_pts(paired_left, 'paired_left',
function.space_order, eval_offset)
fill_stencils(paired_left, 'paired_left', max_span, lambdas, w,
dim_limit, axis_dim)
if len(paired_right.index) != 0:
paired_right = get_n_pts(paired_right, 'paired_right',
function.space_order, eval_offset)
fill_stencils(paired_right, 'paired_right', max_span, lambdas, w,
dim_limit, axis_dim)
w.data[:] /= f_grid.spacing[
axis]**deriv # Divide everything through by spacing
return w
def GradientOperator(model,
v,
grad,
rec,
u,
data,
time_order=2,
spc_order=6,
tsave=4.0,
free_surface=False,
**kwargs):
"""
Class to setup the gradient 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: recin : receiver data for the adjoint source
:param: time_order: Time discretization order
:param: spc_order: Space discretization order
"""
nt = data.shape[0]
s = t.spacing
dt = model.critical_dt
m, damp, rho = model.m, model.damp, model.rho
Lap, rho = acoustic_laplacian(v, rho)
# Derive stencil from symbolic equation
eqn = m / rho * v.dt2 - Lap - damp * v.dt
stencil = solve(eqn, v.backward, rational=False)[0]
nsave = int(nt / (tsave / dt) + 1)
rate = int(nt / nsave) + 1
gradient_update = Eq(
grad, grad - ((time % (Function('INT')(rate))) < 1) *
u.subs(u.indices[0],
Function('INT')(time / rate)) * v.dt2 / rho)
# Add substitutions for spacing (temporal and spatial)
subs = dict([(s, dt)] +
[(i.spacing, model.get_spacing()[j])
for i, j in zip(v.indices[1:], range(len(model.shape)))])
dse = kwargs.get('dse', 'advanced')
dle = kwargs.get('dle', 'advanced')
# Create stencil expressions for operator, source and receivers
eqn = Eq(v.backward, stencil)
# Add expression for receiver injection
ti = v.indices[0]
receivers = rec.inject(field=v.backward,
offset=model.nbpml,
expr=rho * rec * dt * dt / m)
stencils = [eqn] + receivers + [gradient_update]
if free_surface:
fs = Dimension(name="fs", size=model.nbpml)
stencils += [
Eq(v.backward.subs({v.indices[-1]: fs}),
-v.backward.subs({v.indices[-1]: 2 * model.nbpml - fs}))
]
op = Operator(stencils,
subs=subs,
dse=dse,
dle=dle,
time_axis=Backward,
name="Gradient%s" % randint(1e5),
profiler=False,
external=True)
return op
def y(ydim=6):
return Dimension(name='y', size=ydim)
def BornOperator(model,
u,
du,
src,
Linrec,
dm,
data,
time_order=2,
spc_order=6,
save=False,
free_surface=False,
**kwargs):
"""
Class to setup the linearized 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: dmin : square slowness perturbation
:param: recin : receiver data for the adjoint source
:param: time_order: Time discretization order
:param: spc_order: Space discretization order
"""
nt = data.shape[0]
s = t.spacing
dt = model.critical_dt
m, damp, rho = model.m, model.damp, model.rho
Lap, rho = acoustic_laplacian(u, rho)
LapU, _ = acoustic_laplacian(du, rho)
# Derive stencils from symbolic equation
first_eqn = m / rho * u.dt2 - Lap + damp * u.dt
first_stencil = solve(first_eqn, u.forward, rational=False)[0]
second_eqn = m / rho * du.dt2 - LapU + damp * du.dt + dm / rho * u.dt2
second_stencil = solve(second_eqn, du.forward, rational=False)[0]
# Add substitutions for spacing (temporal and spatial)
subs = dict([(s, dt)] +
[(i.spacing, model.get_spacing()[j])
for i, j in zip(u.indices[1:], range(len(model.shape)))])
# Add Born-specific updates and resets
dse = kwargs.get('dse', 'advanced')
dle = kwargs.get('dle', 'advanced')
# Create stencil expressions for operator, source and receivers
eqn1 = [Eq(u.forward, first_stencil)]
eqn2 = [Eq(du.forward, second_stencil)]
# Add source term expression for u
ti = u.indices[0]
source = src.inject(field=u.forward,
offset=model.nbpml,
expr=rho * src * dt * dt / m)
# Create receiver interpolation expression from U
receivers = Linrec.interpolate(expr=du, offset=model.nbpml)
stencils = eqn1 + source + eqn2 + receivers
if free_surface:
fs = Dimension(name="fs", size=model.nbpml)
stencils += [
Eq(u.forward.subs({u.indices[-1]: fs}),
-u.forward.subs({u.indices[-1]: 2 * model.nbpml - fs}))
]
stencils += [
Eq(du.forward.subs({du.indices[-1]: fs}),
-du.forward.subs({du.indices[-1]: 2 * model.nbpml - fs}))
]
op = Operator(stencils,
subs=subs,
dse=dse,
dle=dle,
time_axis=Forward,
name="Born%s" % randint(1e5),
profiler=False,
external=True)
return op
def ForwardOperator(model,
u,
src,
rec,
data,
q,
time_order=2,
spc_order=6,
save=False,
tsave=4.0,
free_surface=False,
**kwargs):
nt = data.shape[0]
dt = model.critical_dt
s = t.spacing
m, damp, rho = model.m, model.damp, model.rho
Lap, rho = acoustic_laplacian(u, rho)
# Derive stencil from symbolic equation
eqn = m / rho * u.dt2 - Lap + damp * u.dt + q
# stencil = solve(eqn, u.forward)[0]
stencil = solve(eqn, u.forward, rational=False)[0]
# Add substitutions for spacing (temporal and spatial)
subs = dict([(s, dt)] +
[(i.spacing, model.get_spacing()[j])
for i, j in zip(u.indices[1:], range(len(model.shape)))])
stencils = [Eq(u.forward, stencil)]
# Create stencil expressions for operator, source and receivers
ti = u.indices[0]
src_term = src.inject(field=u.forward,
offset=model.nbpml,
expr=rho * src * dt**2 / m)
# Create interpolation expression for receivers
rec_term = rec.interpolate(expr=u, offset=model.nbpml)
stencils = stencils + src_term + rec_term
if save:
nsave = int(nt / (tsave / dt) + 1)
rate = int(nt / nsave) + 1
usave = TimeData(name="usave",
shape=model.shape_domain,
time_dim=nt,
time_order=2,
space_order=spc_order,
save=True,
dtype=model.dtype)
stencils += [
Eq(usave.subs(usave.indices[0],
Function('INT')(time / rate)), u)
]
if free_surface:
fs = Dimension(name="fs", size=model.nbpml)
stencils += [
Eq(u.forward.subs({u.indices[-1]: fs}),
-u.forward.subs({u.indices[-1]: 2 * model.nbpml - fs}))
]
dse = kwargs.get('dse', 'advanced')
dle = kwargs.get('dle', 'advanced')
op = Operator(stencils,
subs=subs,
dse=dse,
dle=dle,
time_axis=Forward,
name="Forward%s" % randint(1e5),
profiler=False,
external=True)
return op
def dimify(dimensions):
assert isinstance(dimensions, str)
return tuple(Dimension(name=i) for i in dimensions.split())
def forward(self, src=None, rec=None, u=None, vp=None, save=None, **kwargs):
"""
Forward modelling function that creates the necessary
data objects for running a forward 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
Stores the computed wavefield.
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 the forward wavefield if not provided
u = u or TimeFunction(name='u', grid=self.model.grid,
save=self.geometry.nt if save else None,
time_order=2, space_order=self.space_order)
# Pick vp from model unless explicitly provided
vp = vp or self.model.vp
print("====Forward norm(u)", norm(u))
# Execute operator and return wavefield and receiver data
# summary = self.op_fwd(save).apply(src=src, rec=rec, u=u, vp=vp,
summary = self.op_fwd(save).apply(src=src, u=u, vp=vp,
dt=kwargs.pop('dt', self.dt), **kwargs)
print("====Forward norm(u)", norm(u))
regnormu = norm(u)
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("Norm u:", regnormu)
s_u = TimeFunction(name='s_u', grid=self.model.grid, space_order=self.space_order, time_order=2)
src_u = src.inject(field=s_u.forward, expr=src * self.model.grid.time_dim.spacing**2 / self.model.m)
op_f = Operator([src_u])
op_f.apply(src=src, dt=kwargs.pop('dt', self.dt))
# import pdb;pdb.set_trace()
print("Norm s_u", norm(s_u))
# 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
print("Number of unique affected points is:", 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 populated 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_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)
print("Injecting to empty grids")
op1 = Operator([save_src_u_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_id = Function(name='sp_source_id', shape=(list(sparse_shape)),
dimensions=(x, y, sp_zi), space_order=0, dtype=np.int32)
# Now holds IDs
sp_source_id.data[inds[0], inds[1], :] = tuple(inds[-1][:len(np.unique(inds[-1]))])
assert(np.count_nonzero(sp_source_id.data) == len(nzinds[0]))
assert(len(sp_source_id.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_id[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]]
# Is source_mask needed /
inj_u = save_src_u[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))
# 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 = 9
bystart = 4
byend = 9
bstep = 4
txstart = 32
txend = 65
tystart = 32
tyend = 65
tstep = 32
# 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
print("-----")
tteqs = (eqxb, eqyb, eqxb2, eqyb2, eq0, eq1, eq_u)
# import pdb; pdb.set_trace()
# Execute operator and return wavefield and receiver data
print("TT====Forward norm(u)", norm(u))
summary_tt = self.op_fwd(save, tteqs).apply(u=u, vp=vp,
dt=kwargs.pop('dt', self.dt), **kwargs)
print("TT====Forward norm(u)", norm(u))
# op_tt = self.op_fwd(save, tteqs)
# Execute operator and return wavefield and receiver data
#summary_tt = self.op_fwd(save).apply(src=src, rec=rec, u=u, vp=vp,
# dt=kwargs.pop('dt', self.dt), **kwargs)
# op_tt = self.op_fwd(kernel, save, tteqs)
# summary_tt = op_tt.apply(u=u, dt=kwargs.pop('dt', self.dt), **kwargs)
configuration['jit-backdoor'] = False
norm_tt_u = norm(u)
print("Norm u:", regnormu)
print("Norm(tt_u):", norm_tt_u)
configuration['jit-backdoor'] = True
print("===Temporal blocking======================================")
performance_map = np.append(performance_map, [[tx, ty, bx, by, summary_tt.globals['fdlike'].gpointss]], 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 1:
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)
# import pdb;pdb.set_trace()
return rec, u, summary
