def test_streaming_w_shifting(self):
nt = 50
grid = Grid(shape=(5, 5))
time = grid.time_dim
factor = Constant(name='factor', value=5, dtype=np.int32)
t_sub = ConditionalDimension('t_sub', parent=time, factor=factor)
save_shift = Constant(name='save_shift', dtype=np.int32)
u = TimeFunction(name='u', grid=grid, time_order=0)
usave = TimeFunction(name='usave',
grid=grid,
time_order=0,
save=(int(nt // factor.data)),
time_dim=t_sub)
for i in range(usave.save):
usave.data[i, :] = i
eqns = Eq(u.forward, u + usave.subs(t_sub, t_sub - save_shift))
op = Operator(eqns, opt=('streaming', 'orchestrate'))
# From time_m=15 to time_M=35 with a factor=5 -- it means that, thanks
# to t_sub, we enter the Eq exactly (35-15)/5 + 1 = 5 times. We set
# save_shift=1 so instead of accessing the range usave[15/5:35/5+1],
# we rather access the range usave[15/5-1:35:5], which means accessing
# the usave values 2, 3, 4, 5, 6.
op.apply(time_m=15, time_M=35, save_shift=1)
assert np.allclose(u.data, 20)
# Again, but with a different shift
op.apply(time_m=15, time_M=35, save_shift=-2)
assert np.allclose(u.data, 20 + 35)
def test_const_change(self):
"""
Test that Constand.data can be set as required.
"""
n = 5
t = Constant(name='t', dtype=np.int32)
grid = Grid(shape=(2, 2))
x, y = grid.dimensions
f = TimeFunction(name='f', grid=grid, save=n+1)
f.data[:] = 0
eq = Eq(f.dt-1)
stencil = Eq(f.forward, solve(eq, f.forward))
op = Operator([stencil])
op.apply(time_m=0, time_M=n-1, dt=1)
check = Function(name='check', grid=grid)
eq_test = Eq(check, f[t, x, y])
op_test = Operator([eq_test])
for j in range(0, n+1):
t.data = j # Ensure constant is being updated correctly
op_test.apply(t=t)
assert(np.amax(check.data[:], axis=None) == j)
assert(np.amin(check.data[:], axis=None) == j)
def test_lm_ds(self, dat, dtype):
"""
Test logistic map with 2nd derivative term that should cancel.
"""
iterations = 10000
r = Constant(name='r', dtype=dtype)
r.data = dtype(0.5*dat)
s = dtype(0.1)
grid = Grid(shape=(2, 2), extent=(1, 1), dtype=dtype)
f0 = TimeFunction(name='f0', grid=grid, time_order=2, dtype=dtype)
f1 = TimeFunction(name='f1', grid=grid, time_order=2, save=iterations+2,
dtype=dtype)
initial_condition = dtype(0.7235)
lmap0 = Eq(f0.forward, -r*f0.dt2*s**2*(1.0-f0) +
r*(1.0-f0)*(f0.backward+f0.forward))
lmap1 = Eq(f1.forward, -r*f1.dt2*s**2*(1.0-f1) +
r*(1.0-f1)*(f1.backward+f1.forward))
f0.data[1, :, :] = initial_condition
f1.data[1, :, :] = initial_condition
op0 = Operator([Eq(f0.forward, dtype(0.0)), lmap0])
op1 = Operator(lmap1)
op0(time_m=1, time_M=iterations, dt=s)
op1(time_m=1, time_M=iterations, dt=s)
assert np.allclose(f0.data[np.mod(iterations+1, 3)], f1.data[iterations+1],
atol=0, rtol=0)
def test_lm_fb(self, dat, dtype):
"""
Test logistic map with forward and backward terms that should cancel.
"""
iterations = 10000
r = Constant(name='r', dtype=dtype)
r.data = dtype(dat)
s = dtype(0.1)
grid = Grid(shape=(2, 2), extent=(1, 1), dtype=dtype)
dt = grid.stepping_dim.spacing
print("dt = ", dt)
f0 = TimeFunction(name='f0', grid=grid, time_order=2, dtype=dtype)
f1 = TimeFunction(name='f1', grid=grid, time_order=2, save=iterations+2,
dtype=dtype)
initial_condition = dtype(0.7235)
lmap0 = Eq(f0.forward, r*f0*(1.0-f0+(1.0/s)*dt*f0.backward
- f0.backward+(1.0/s)*dt*f0.forward-f0.forward))
lmap1 = Eq(f1.forward, r*f1*(1.0-f1+(1.0/s)*dt*f1.backward
- f1.backward+(1.0/s)*dt*f1.forward-f1.forward))
f0.data[1, :, :] = initial_condition
f1.data[1, :, :] = initial_condition
op0 = Operator([Eq(f0.forward, dtype(0.0)), lmap0])
op1 = Operator(lmap1)
op0(time_m=1, time_M=iterations, dt=s)
op1(time_m=1, time_M=iterations, dt=s)
assert np.allclose(f0.data[np.mod(iterations+1, 3)], f1.data[iterations+1],
atol=0, rtol=0)
def test_const_change(self):
"""
Test that Constand.data can be set as required.
"""
n = 5
t = Constant(name='t', dtype=np.int32)
grid = Grid(shape=(2, 2))
x, y = grid.dimensions
f = TimeFunction(name='f', grid=grid, save=n+1)
f.data[:] = 0
eq = Eq(f.dt-1)
stencil = Eq(f.forward, solve(eq, f.forward))
op = Operator([stencil])
op.apply(time_m=0, time_M=n-1, dt=1)
check = Function(name='check', grid=grid)
eq_test = Eq(check, f[t, x, y])
op_test = Operator([eq_test])
for j in range(0, n+1):
t.data = j # Ensure constant is being updated correctly
op_test.apply(t=t)
assert(np.amax(check.data[:], axis=None) == j)
assert(np.amin(check.data[:], axis=None) == j)
def __init__(self, origin, spacing, shape, space_order, vp, vs, rho, nbpml=20,
dtype=np.float32):
super(ModelElastic, self).__init__(origin, spacing, shape, space_order,
nbpml=nbpml, dtype=dtype)
# Create dampening field as symbol `damp`
self.damp = Function(name="damp", grid=self.grid)
initialize_damp(self.damp, self.nbpml, self.spacing, mask=True)
# Create square slowness of the wave as symbol `m`
if isinstance(vp, np.ndarray):
self.vp = Function(name="vp", grid=self.grid, space_order=space_order)
initialize_function(self.vp, vp, self.nbpml)
else:
self.vp = Constant(name="vp", value=vp)
self._physical_parameters = ('vp',)
# Create square slowness of the wave as symbol `m`
if isinstance(vs, np.ndarray):
self.vs = Function(name="vs", grid=self.grid, space_order=space_order)
initialize_function(self.vs, vs, self.nbpml)
else:
self.vs = Constant(name="vs", value=vs)
self._physical_parameters += ('vs',)
# Create square slowness of the wave as symbol `m`
if isinstance(rho, np.ndarray):
self.rho = Function(name="rho", grid=self.grid, space_order=space_order)
initialize_function(self.rho, rho, self.nbpml)
else:
self.rho = Constant(name="rho", value=rho)
self._physical_parameters += ('rho',)
def test_argument_from_index_constant(self):
nx, ny = 30, 30
grid = Grid(shape=(nx, ny))
x, y = grid.dimensions
arbdim = Dimension('arb')
u = TimeFunction(name='u',
grid=grid,
save=None,
time_order=2,
space_order=0)
snap = Function(name='snap',
dimensions=(arbdim, x, y),
shape=(5, nx, ny),
space_order=0)
save_t = Constant(name='save_t', dtype=np.int32)
save_slot = Constant(name='save_slot', dtype=np.int32)
expr = Eq(snap.subs(arbdim, save_slot),
u.subs(grid.stepping_dim, save_t))
op = Operator(expr)
u.data[:] = 0.0
snap.data[:] = 0.0
u.data[0, 10, 10] = 1.0
op.apply(save_t=0, save_slot=1)
assert snap.data[1, 10, 10] == 1.0
def test_cache_constant_new():
"""Test that new u[x, y] instances don't cache"""
u0 = Constant(name='u')
u0.data = 6.
u1 = Constant(name='u')
u1.data = 2.
assert u0.data == 6.
assert u1.data == 2.
def test_constant():
c = Constant(name='c')
assert c.data == 0.
c.data = 1.
pkl_c = pickle.dumps(c)
new_c = pickle.loads(pkl_c)
# .data is initialized, so it should have been pickled too
assert np.all(c.data == 1.)
assert np.all(new_c.data == 1.)
def test_constant():
c = Constant(name='c')
assert c.data == 0.
c.data = 1.
pkl_c = pickle.dumps(c)
new_c = pickle.loads(pkl_c)
# .data is initialized, so it should have been pickled too
assert np.all(c.data == 1.)
assert np.all(new_c.data == 1.)
def test_everything():
nt = 50
grid = Grid(shape=(6, 6))
x, y = grid.dimensions
time = grid.time_dim
xi = SubDimension.middle(name='xi',
parent=x,
thickness_left=2,
thickness_right=2)
yi = SubDimension.middle(name='yi',
parent=y,
thickness_left=2,
thickness_right=2)
factor = Constant(name='factor', value=5, dtype=np.int32)
t_sub = ConditionalDimension('t_sub', parent=time, factor=factor)
save_shift = Constant(name='save_shift', dtype=np.int32)
u = TimeFunction(name='u', grid=grid, time_order=0)
u1 = TimeFunction(name='u', grid=grid, time_order=0)
va = TimeFunction(name='va',
grid=grid,
time_order=0,
save=(int(nt // factor.data)),
time_dim=t_sub)
vb = TimeFunction(name='vb',
grid=grid,
time_order=0,
save=(int(nt // factor.data)),
time_dim=t_sub)
for i in range(va.save):
va.data[i, :] = i
vb.data[i, :] = i * 2 - 1
vas = va.subs(t_sub, t_sub - save_shift)
vasb = va.subs(t_sub, t_sub - 1 - save_shift)
vasf = va.subs(t_sub, t_sub + 1 - save_shift)
eqns = [Eq(u.forward, u + (vasb + vas + vasf) * 2. + vb)]
eqns = [e.xreplace({x: xi, y: yi}) for e in eqns]
op0 = Operator(eqns, opt='noop')
op1 = Operator(eqns, opt='buffering')
# Check generated code
assert len([i for i in FindSymbols().visit(op1) if i.is_Array]) == 2
op0.apply(time_m=15, time_M=35, save_shift=0)
op1.apply(time_m=15, time_M=35, save_shift=0, u=u1)
assert np.all(u.data == u1.data)
def test_derive_constant_value(self):
"""Ensure that values for :class:`Constant` symbols are derived correctly."""
grid = Grid(shape=(5, 6))
f = Function(name='f', grid=grid)
a = Constant(name='a', value=3.)
Operator(Eq(f, a))()
assert np.allclose(f.data, 3.)
g = Function(name='g', grid=grid)
b = Constant(name='b')
op = Operator(Eq(g, b))
b.data = 4.
op()
assert np.allclose(g.data, 4.)
def test_constant_time_dense(self):
"""Test arithmetic between different data objects, namely Constant
and Function."""
i, j = dimify('i j')
const = Constant(name='truc', value=2.)
a = Function(name='a', shape=(20, 20), dimensions=(i, j))
a.data[:] = 2.
eqn = Eq(a, a + 2.*const)
op = Operator(eqn)
op.apply(a=a, truc=const)
assert(np.allclose(a.data, 6.))
# Applying a different constant still works
op.apply(a=a, truc=Constant(name='truc2', value=3.))
assert(np.allclose(a.data, 12.))
def test_save_w_shifting(self):
factor = 4
nt = 19
grid = Grid(shape=(11, 11))
time = grid.time_dim
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave',
grid=grid,
save=2,
time_dim=time_subsampled)
save_shift = Constant(name='save_shift', dtype=np.int32)
eqns = [
Eq(u.forward, u + 1.),
Eq(usave.subs(time_subsampled, time_subsampled - save_shift), u)
]
op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
# Starting at time_m=10, so time_subsampled - save_shift is in range
op.apply(time_m=10, time_M=nt - 2, save_shift=3)
assert np.all(np.allclose(u.data[0], 8))
assert np.all(
[np.allclose(usave.data[i], 2 + i * factor) for i in range(2)])
def test_save(self, opt, gpu_fit, async_degree):
nt = 10
grid = Grid(shape=(300, 300, 300))
time_dim = grid.time_dim
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub",
parent=time_dim,
factor=factor)
u = TimeFunction(name='u', 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
op = Operator(
[Eq(u.forward, u + 1), Eq(usave, u.forward)],
opt=(opt, {
'gpu-fit': usave if gpu_fit else None,
'buf-async-degree': async_degree
}))
op.apply(time_M=nt - 1)
assert all(
np.all(usave.data[i] == 2 * i + 1) for i in range(usave.save))
def run(shape=(50, 50, 50),
spacing=(20.0, 20.0, 20.0),
tn=1000.0,
space_order=4,
kernel='OT2',
nbl=40,
full_run=False,
fs=False,
autotune=False,
preset='layers-isotropic',
checkpointing=False,
**kwargs):
solver = acoustic_setup(shape=shape,
spacing=spacing,
nbl=nbl,
tn=tn,
space_order=space_order,
kernel=kernel,
fs=fs,
preset=preset,
**kwargs)
info("Applying Forward")
# Whether or not we save the whole time history. We only need the full wavefield
# with 'save=True' if we compute the gradient without checkpointing, if we use
# checkpointing, PyRevolve will take care of the time history
save = full_run and not checkpointing
# Define receiver geometry (spread across x, just below surface)
rec, u, summary = solver.forward(save=save, autotune=autotune)
# print(norm(rec))
print(norm(u))
if preset == 'constant':
# With a new m as Constant
v0 = Constant(name="v", value=2.0, dtype=np.float32)
solver.forward(save=save, vp=v0)
# With a new vp as a scalar value
solver.forward(save=save, vp=2.0)
if not full_run:
return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
# Smooth velocity
initial_vp = Function(name='v0',
grid=solver.model.grid,
space_order=space_order)
smooth(initial_vp, solver.model.vp)
dm = np.float32(initial_vp.data**(-2) - solver.model.vp.data**(-2))
info("Applying Adjoint")
solver.adjoint(rec, autotune=autotune)
info("Applying Born")
solver.jacobian(dm, autotune=autotune)
info("Applying Gradient")
solver.jacobian_adjoint(rec,
u,
autotune=autotune,
checkpointing=checkpointing)
return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
def test_intervals_subtract(self, x, y):
nullx = NullInterval(x)
# All nulls
assert nullx.subtract(nullx) == nullx
ix = Interval(x, 2, -2)
# Mixed nulls and defined on the same dimension
assert nullx.subtract(ix) == nullx
assert ix.subtract(ix) == Interval(x, 0, 0)
assert ix.subtract(nullx) == ix
ix2 = Interval(x, 4, -4)
ix3 = Interval(x, 6, -6)
# All defined same dimension
assert ix2.subtract(ix) == ix
assert ix.subtract(ix2) == Interval(x, -2, 2)
assert ix3.subtract(ix) == ix2
c = Constant(name='c')
ix4 = Interval(x, c + 2, c + 4)
ix5 = Interval(x, c + 1, c + 5)
# All defined symbolic
assert ix4.subtract(ix5) == Interval(x, 1, -1)
assert ix5.subtract(ix4) == Interval(x, -1, 1)
assert ix5.subtract(ix) == Interval(x, c - 1, c + 7)
def td_born_adjoint_op(model, geometry, time_order, space_order):
nt = geometry.nt
# Define the wavefields with the size of the model and the time dimension
u = TimeFunction(
name='u',
grid=model.grid,
time_order=time_order,
space_order=space_order,
save=nt
)
# Define the wave equation
pde = model.m * u.dt2 - u.laplace + model.damp * u.dt.T
# Use `solve` to rearrange the equation into a stencil expression
stencil = Eq(u.backward, solve(pde, u.backward), subdomain=model.grid.subdomains['physdomain'])
# Inject at receivers
born_data_rec = PointSource(
name='born_data_rec',
grid=model.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions
)
dt = Constant(name='dt')
rec_term = born_data_rec.inject(field=u.backward, expr=born_data_rec * (dt ** 2) / model.m)
return Operator([stencil] + rec_term, subs=model.spacing_map)
def forward_run(self,
wav,
src_coords,
rcv_coords,
save=False,
q=0,
v=None,
w=0,
grad=False):
# Computing residual
u = wavefield(self.model, self.space_order, save=save, nt=wav.shape[0])
kwargs = wf_kwargs(u)
rcv = Receiver(name="rcv",
grid=self.model.grid,
ntime=wav.shape[0],
coordinates=rcv_coords)
src = PointSource(name="src",
grid=self.model.grid,
ntime=wav.shape[0],
coordinates=src_coords)
src.data[:] = wav[:]
fwd = self.op_fwd(save=save, q=q, grad=grad)
if grad:
w = Constant(name="w", value=w)
gradm = Function(name="gradm", grid=self.model.grid)
kwargs.update({as_tuple(v)[0].name: as_tuple(v)[0]})
kwargs.update({w.name: w, gradm.name: gradm})
fwd(rcv=rcv, src=src, **kwargs)
if grad:
return rcv.data, u, gradm
return rcv.data, u
def test_shifted(self):
nt = 19
grid = Grid(shape=(11, 11))
time = grid.time_dim
u = TimeFunction(name='u', grid=grid)
assert(grid.stepping_dim in u.indices)
u2 = TimeFunction(name='u2', grid=grid, save=nt)
assert(time in u2.indices)
factor = 4
time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
usave = TimeFunction(name='usave', grid=grid, save=2, time_dim=time_subsampled)
assert(time_subsampled in usave.indices)
t_sub_shift = Constant(name='t_sub_shift', dtype=np.int32)
eqns = [Eq(u.forward, u + 1.), Eq(u2.forward, u2 + 1.),
Eq(usave.subs(time_subsampled, time_subsampled - t_sub_shift), u)]
op = Operator(eqns)
# Starting at time_m=10, so time_subsampled - t_sub_shift is in range
op.apply(time_m=10, time_M=nt-2, t_sub_shift=3)
assert np.all(np.allclose(u.data[0], 8))
assert np.all([np.allclose(u2.data[i], i - 10) for i in range(10, nt)])
assert np.all([np.allclose(usave.data[i], 2+i*factor) for i in range(2)])
def solver(I, w, dt, T):
"""
Solve u'=v, v' = - w**2*u for t in (0,T], u(0)=I and v(0)=0,
by an Euler-Cromer method.
"""
dt = float(dt)
Nt = int(round(T/dt))
t = Dimension('t', spacing=Constant('h_t'))
v = TimeFunction(name='v', dimensions=(t,), shape=(Nt+1,), space_order=2)
u = TimeFunction(name='u', dimensions=(t,), shape=(Nt+1,), space_order=2)
v.data[:] = 0
u.data[:] = I
eq_v = Eq(v.dt, -(w**2)*u)
eq_u = Eq(u.dt, v.forward)
stencil_v = solve(eq_v, v.forward)
stencil_u = solve(eq_u, u.forward)
update_v = Eq(v.forward, stencil_v)
update_u = Eq(u.forward, stencil_u)
op = Operator([update_v, update_u])
op.apply(h_t=dt, t_M=Nt-1)
return u.data, v.data, np.linspace(0, Nt*dt, Nt+1)
def test_save_w_nonaffine_time(self):
factor = 4
grid = Grid(shape=(11, 11))
x, y = grid.dimensions
t = grid.stepping_dim
time = grid.time_dim
time_subsampled = ConditionalDimension('t_sub',
parent=time,
factor=factor)
f = Function(name='f', grid=grid, dtype=np.int32)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave',
grid=grid,
save=2,
time_dim=time_subsampled)
save_shift = Constant(name='save_shift', dtype=np.int32)
eqns = [
Eq(u.forward, u[t, f[x, x], f[y, y]] + 1.),
Eq(usave.subs(time_subsampled, time_subsampled - save_shift), u)
]
op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
# We just check the generated code here
assert len([i for i in FindSymbols().visit(op)
if isinstance(i, Lock)]) == 1
assert len(op._func_table) == 2
def test_conddim_backwards():
nt = 10
grid = Grid(shape=(4, 4))
time_dim = grid.time_dim
x, y = grid.dimensions
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
u = TimeFunction(name='u', grid=grid, time_order=0, save=nt, time_dim=time_sub)
v = TimeFunction(name='v', grid=grid)
v1 = TimeFunction(name='v', grid=grid)
for i in range(u.save):
u.data[i, :] = i
eqns = [Eq(v.backward, v.backward + v + u + 1.)]
op0 = Operator(eqns, opt='noop')
op1 = Operator(eqns, opt='buffering')
# Check generated code
assert len(retrieve_iteration_tree(op1)) == 3
buffers = [i for i in FindSymbols().visit(op1) if i.is_Array]
assert len(buffers) == 1
op0.apply(time_m=1, time_M=9)
op1.apply(time_m=1, time_M=9, v=v1)
assert np.all(v.data == v1.data)
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 solver_v1(I, w, dt, T):
"""
Solve u'=v, v' = - w**2*u for t in (0,T], u(0)=I and v(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)
v = TimeFunction(name='v', dimensions=(t,), shape=(Nt+1,), space_order=2)
u.data[:] = I
v.data[:] = 0 - 0.5*dt*w**2*u.data[:]
eq_u = Eq(u.dt, v)
eq_v = Eq(v.dt, -(w**2)*u.forward)
stencil_u = solve(eq_u, u.forward)
stencil_v = solve(eq_v, v.forward)
update_u = Eq(u.forward, stencil_u)
update_v = Eq(v.forward, stencil_v)
op = Operator([update_u, update_v])
op.apply(h_t=dt, t_M=Nt-1)
t_mesh = np.linspace(0, Nt*dt, Nt+1) # mesh for u
t_v_mesh = (t_mesh + dt/2)[:-1] # mesh for v
return u.data, t_mesh, v.data, t_v_mesh
def test_symbolic_factor(self):
"""
Test ConditionalDimension with symbolic factor (provided as a Constant).
"""
g = Grid(shape=(4, 4, 4))
u = TimeFunction(name='u', grid=g, time_order=0)
fact = Constant(name='fact', dtype=np.int32, value=4)
tsub = ConditionalDimension(name='tsub',
parent=g.time_dim,
factor=fact)
usave = TimeFunction(name='usave', grid=g, time_dim=tsub, save=4)
op = Operator([Eq(u, u + 1), Eq(usave, u)])
op.apply(time=7) # Use `fact`'s default value, 4
assert np.all(usave.data[0] == 1)
assert np.all(usave.data[1] == 5)
u.data[:] = 0.
op.apply(time=7, fact=2)
assert np.all(usave.data[0] == 1)
assert np.all(usave.data[1] == 3)
assert np.all(usave.data[2] == 5)
assert np.all(usave.data[3] == 7)
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 test_save_w_subdims(self):
nt = 10
grid = Grid(shape=(10, 10))
x, y = grid.dimensions
time_dim = grid.time_dim
xi = SubDimension.middle(name='xi', parent=x, thickness_left=3, thickness_right=3)
yi = SubDimension.middle(name='yi', parent=y, thickness_left=3, thickness_right=3)
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
eqns = [Eq(u.forward, u + 1),
Eq(usave, u.forward)]
eqns = [e.xreplace({x: xi, y: yi}) for e in eqns]
op = Operator(eqns, opt=('buffering', 'tasking', 'orchestrate'))
op.apply(time_M=nt-1)
for i in range(usave.save):
assert np.all(usave.data[i, 3:-3, 3:-3] == 2*i + 1)
assert np.all(usave.data[i, :3, :] == 0)
assert np.all(usave.data[i, -3:, :] == 0)
assert np.all(usave.data[i, :, :3] == 0)
assert np.all(usave.data[i, :, -3:] == 0)
def test_save_multi_output(self):
nt = 10
grid = Grid(shape=(150, 150, 150))
time_dim = grid.time_dim
factor = Constant(name='factor', value=2, dtype=np.int32)
time_sub = ConditionalDimension(name="time_sub", parent=time_dim, factor=factor)
u = TimeFunction(name='u', grid=grid)
usave = TimeFunction(name='usave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
vsave = TimeFunction(name='vsave', grid=grid, time_order=0,
save=int(nt//factor.data), time_dim=time_sub)
eqns = [Eq(u.forward, u + 1),
Eq(usave, u.forward),
Eq(vsave, u.forward)]
op = Operator(eqns, opt=('buffering', 'tasking', 'topofuse', 'orchestrate'))
# Check generated code
assert len(op._func_table) == 4 # usave and vsave eqns are in separate tasks
op.apply(time_M=nt-1)
assert all(np.all(usave.data[i] == 2*i + 1) for i in range(usave.save))
assert all(np.all(vsave.data[i] == 2*i + 1) for i in range(vsave.save))
def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=1000.0,
time_order=2, space_order=4, nbpml=40, full_run=False,
autotune=False, constant=False, **kwargs):
solver = acoustic_setup(shape=shape, spacing=spacing, nbpml=nbpml, tn=tn,
space_order=space_order, time_order=time_order,
constant=constant, **kwargs)
initial_vp = smooth10(solver.model.m.data, solver.model.shape_domain)
dm = np.float32(initial_vp**2 - solver.model.m.data)
info("Applying Forward")
rec, u, summary = solver.forward(save=full_run, autotune=autotune)
if constant:
# With a new m as Constant
m0 = Constant(name="m", value=.25, dtype=np.float32)
solver.forward(save=full_run, m=m0)
# With a new m as a scalar value
solver.forward(save=full_run, m=.25)
if not full_run:
return summary.gflopss, summary.oi, summary.timings, [rec, u.data]
info("Applying Adjoint")
solver.adjoint(rec, autotune=autotune)
info("Applying Born")
solver.born(dm, autotune=autotune)
info("Applying Gradient")
solver.gradient(rec, u, autotune=autotune)
def test_custom_dimension():
symbolic_size = Constant(name='d_custom_size')
d = CustomDimension(name='d', symbolic_size=symbolic_size)
pkl_d = pickle.dumps(d)
new_d = pickle.loads(pkl_d)
assert d.name == new_d.name
assert d.symbolic_size.name == new_d.symbolic_size.name
def new_ops_gbl(self, c):
if c in self.ops_args:
return self.ops_args[c]
new_c = Constant(name='*%s' % c.name, dtype=c.dtype)
self.ops_args[c] = new_c
self.ops_params.append(new_c)
return new_c
def _gen_phys_param(self, field, name, space_order):
if isinstance(field, np.ndarray):
function = Function(name=name,
grid=self.grid,
space_order=space_order)
initialize_function(function, field, self.nbpml)
else:
function = Constant(name=name, value=field)
return function