def test_staggered_function(self, order):
"""Test custom function coefficients on a staggered grid"""
grid = Grid(shape=(11, ), extent=(10., ))
x = grid.dimensions[0]
f = Function(name='f',
grid=grid,
space_order=order,
coefficients='symbolic')
g = Function(name='g',
grid=grid,
space_order=order,
coefficients='symbolic',
staggered=x)
f.data[::2] = 1
g.data[::2] = 1
s = Dimension(name='s')
ncoeffs = order + 1
wshape = grid.shape + (ncoeffs, )
wdims = grid.dimensions + (s, )
w = Function(name='w', dimensions=wdims, shape=wshape)
w.data[:] = 1.0 / grid.spacing[0]**2
coeffs_f = Coefficient(2, f, x, w)
coeffs_g = Coefficient(2, g, x, w)
eq_f = Eq(f, f.dx2, coefficients=Substitutions(coeffs_f))
eq_g = Eq(g, g.dx2, coefficients=Substitutions(coeffs_g))
Operator([eq_f, eq_g])()
assert np.allclose(f.data, g.data, atol=1e-7)
def test_staggered_array(self, order):
"""Test custom coefficients provided as an array on a staggered grid"""
grid = Grid(shape=(11, ), extent=(10., ))
x = grid.dimensions[0]
f = Function(name='f',
grid=grid,
space_order=order,
coefficients='symbolic')
g = Function(name='g',
grid=grid,
space_order=order,
coefficients='symbolic',
staggered=x)
f.data[::2] = 1
g.data[::2] = 1
weights = np.ones(order + 1) / grid.spacing[0]**2
coeffs_f = Coefficient(2, f, x, weights)
coeffs_g = Coefficient(2, g, x, weights)
eq_f = Eq(f, f.dx2, coefficients=Substitutions(coeffs_f))
eq_g = Eq(g, g.dx2, coefficients=Substitutions(coeffs_g))
Operator([eq_f, eq_g])()
assert np.allclose(f.data, g.data, atol=1e-7)
def test_with_timefunction(self, stagger):
"""Check compatibility of custom coefficients and TimeFunctions"""
grid = Grid(shape=(11, ), extent=(10., ))
x = grid.dimensions[0]
if stagger:
staggered = x
else:
staggered = None
f = TimeFunction(name='f',
grid=grid,
space_order=2,
staggered=staggered)
g = TimeFunction(name='g',
grid=grid,
space_order=2,
staggered=staggered,
coefficients='symbolic')
f.data[:, ::2] = 1
g.data[:, ::2] = 1
weights = np.array([-1, 2, -1]) / grid.spacing[0]**2
coeffs = Coefficient(2, g, x, weights)
eq_f = Eq(f.forward, f.dx2)
eq_g = Eq(g.forward, g.dx2, coefficients=Substitutions(coeffs))
Operator([eq_f, eq_g])(t_m=0, t_M=1)
assert np.allclose(f.data[-1], -g.data[-1], atol=1e-7)
def test_function_coefficients(self):
"""Test that custom function coefficients return the expected result"""
so = 2
grid = Grid(shape=(4, 4))
f0 = TimeFunction(name='f0',
grid=grid,
space_order=so,
coefficients='symbolic')
f1 = TimeFunction(name='f1', grid=grid, space_order=so)
x, y = grid.dimensions
s = Dimension(name='s')
ncoeffs = so + 1
wshape = list(grid.shape)
wshape.append(ncoeffs)
wshape = as_tuple(wshape)
wdims = list(grid.dimensions)
wdims.append(s)
wdims = as_tuple(wdims)
w = Function(name='w', dimensions=wdims, shape=wshape)
w.data[:, :, 0] = 0.0
w.data[:, :, 1] = -1.0 / grid.spacing[0]
w.data[:, :, 2] = 1.0 / grid.spacing[0]
f_x_coeffs = Coefficient(1, f0, x, w)
subs = Substitutions(f_x_coeffs)
eq0 = Eq(f0.dt + f0.dx, 1, coefficients=subs)
eq1 = Eq(f1.dt + f1.dx, 1)
stencil0 = solve(eq0.evaluate, f0.forward)
stencil1 = solve(eq1.evaluate, f1.forward)
op0 = Operator(Eq(f0.forward, stencil0))
op1 = Operator(Eq(f1.forward, stencil1))
op0(time_m=0, time_M=5, dt=1.0)
op1(time_m=0, time_M=5, dt=1.0)
assert np.all(
np.isclose(f0.data[:] - f1.data[:], 0.0, atol=1e-5, rtol=0))
def test_coefficients(self, expr, sorder, dorder, dim, weights, expected):
"""Test that custom coefficients return the expected result"""
grid = Grid(shape=(10, 10))
u = Function(name='u',
grid=grid,
space_order=sorder,
coefficients='symbolic')
x = grid.dimensions
order = dorder
dim = x[dim]
weights = np.array(weights)
coeffs = Coefficient(order, u, dim, weights)
eq = Eq(eval(expr), coefficients=Substitutions(coeffs))
assert isinstance(eq.lhs, Differentiable)
assert expected == str(eq.evaluate.lhs)
def test_staggered_equation(self):
"""
Check that expressions with substitutions are consistent with
those without
"""
grid = Grid(shape=(11, ), extent=(10., ))
x = grid.dimensions[0]
f = Function(name='f',
grid=grid,
space_order=2,
coefficients='symbolic',
staggered=x)
weights = np.array([1, -2, 1]) / grid.spacing[0]**2
coeffs_f = Coefficient(2, f, x, weights)
eq_f = Eq(f, f.dx2, coefficients=Substitutions(coeffs_f))
expected = 'Eq(f(x + h_x/2), f(x - h_x/2) - 2.0*f(x + h_x/2) + f(x + 3*h_x/2))'
assert (str(eq_f.evaluate) == expected)
def test_coefficients_w_xreplace(self):
"""Test custom coefficients with an xreplace before they are applied"""
grid = Grid(shape=(4, 4))
u = Function(name='u',
grid=grid,
space_order=2,
coefficients='symbolic')
x = grid.dimensions[0]
dorder = 1
weights = np.array([-0.6, 0.1, 0.6])
coeffs = Coefficient(dorder, u, x, weights)
c = sp.Symbol('c')
eq = Eq(u.dx + c, coefficients=Substitutions(coeffs))
eq = eq.xreplace({c: 2})
expected = '0.1*u(x, y) - 0.6*u(x - h_x, y) + 0.6*u(x + h_x, y) + 2'
assert expected == str(eq.evaluate.lhs)
def get_weights(data,
function,
deriv,
bcs,
interior,
fill_function=None,
eval_offsets=(0., 0., 0.)):
"""
Get the modified stencil weights for a function and derivative given the
axial distances.
Parameters
----------
data : devito VectorFunction
The axial distance function
function : devito Function
The function for which stencils should be calculated
deriv : int
The order of the derivative to which the stencils pertain
bcs : list of devito Eq
The boundary conditions which should hold at the surface
interior : ndarray
The interior-exterior segmentation of the domain
fill_function : devito Function
A secondary function to use when creating the Coefficient objects. If none
then the Function supplied with the function argument will be used instead.
eval_offsets : tuple of float
The relative offsets at which derivatives should be evaluated for each
axis.
Returns
-------
substitutions : tuple of Coefficient
The substitutions to be included in the devito equation
"""
cache = os.path.dirname(__file__) + '/extrapolation_cache.dat'
# This wants to start as an empty list
weights = []
for axis in range(3):
print(deriv, function, function.grid.dimensions[axis])
stencils = StencilSet(deriv, eval_offsets[axis], bcs, cache=cache)
lambdas = stencils.lambdaify
max_span = stencils.max_span
axis_weights = get_component_weights(data[axis].data, axis, function,
deriv, lambdas, interior,
max_span, eval_offsets[axis])
if fill_function is None:
weights.append(
Coefficient(deriv, function, function.grid.dimensions[axis],
axis_weights))
else:
weights.append(
Coefficient(deriv, fill_function,
fill_function.grid.dimensions[axis], axis_weights))
# Raise error if list is empty
if len(weights) == 0:
raise ValueError("No boundary-adjacent points in provided fields")
return tuple(weights)