def freesurface(model, eq):
"""
Generate the stencil that mirrors the field as a free surface modeling for
the acoustic wave equation.
Parameters
----------
model : Model
Physical model.
eq : Eq
Time-stepping stencil (time update) to mirror at the freesurface.
"""
lhs, rhs = eq.evaluate.args
# Get vertical dimension and corresponding subdimension
zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
z = zfs.parent
# Functions present in the stencil
funcs = retrieve_functions(rhs)
mapper = {}
# Antisymmetric mirror at negative indices
# TODO: Make a proper "mirror_indices" tool function
for f in funcs:
zind = f.indices[-1]
if (zind - z).as_coeff_Mul()[0] < 0:
s = sign(zind.subs({z: zfs, z.spacing: 1}))
mapper.update({f: s * f.subs({zind: INT(abs(zind))})})
return Eq(lhs, rhs.subs(mapper), subdomain=model.grid.subdomains['fsdomain'])
def _add_implicit(cls, expressions):
"""
Create and add any associated implicit expressions.
Implicit expressions are those not explicitly defined by the user
but instead are requisites of some specified functionality.
"""
processed = []
for e in expressions:
if e.subdomain:
try:
dims = [
d.root for d in e.free_symbols
if isinstance(d, Dimension)
]
sub_dims = [d.root for d in e.subdomain.dimensions]
sub_dims.append(e.subdomain.implicit_dimension)
dims = [d for d in dims if d not in frozenset(sub_dims)]
dims.append(e.subdomain.implicit_dimension)
grid = list(retrieve_functions(e, mode='unique'))[0].grid
processed.extend([
i.func(*i.args, implicit_dims=dims)
for i in e.subdomain._create_implicit_exprs(grid)
])
dims.extend(e.subdomain.dimensions)
new_e = Eq(e.lhs,
e.rhs,
subdomain=e.subdomain,
implicit_dims=dims)
processed.append(new_e)
except AttributeError:
processed.append(e)
else:
processed.append(e)
return processed
def _lower_exprs(cls, expressions, **kwargs):
"""
Expression lowering:
* Form and gather any required implicit expressions;
* Evaluate derivatives;
* Flatten vectorial equations;
* Indexify Functions;
* Apply substitution rules;
* Specialize (e.g., index shifting)
"""
# Add in implicit expressions, e.g., induced by SubDomains
expressions = cls._add_implicit(expressions)
# Unfold lazyiness
expressions = flatten([i.evaluate for i in expressions])
# Scalarize tensor expressions
expressions = [j for i in expressions for j in i._flatten]
# Indexification
# E.g., f(x - 2*h_x, y) -> f[xi + 2, yi + 4] (assuming halo_size=4)
processed = []
for expr in expressions:
if expr.subdomain:
dimension_map = expr.subdomain.dimension_map
else:
dimension_map = {}
# Handle Functions (typical case)
mapper = {
f: f.indexify(lshift=True, subs=dimension_map)
for f in retrieve_functions(expr)
}
# Handle Indexeds (from index notation)
for i in retrieve_indexed(expr):
f = i.function
# Introduce shifting to align with the computational domain
indices = [(a + o)
for a, o in zip(i.indices, f._size_nodomain.left)]
# Apply substitutions, if necessary
if dimension_map:
indices = [j.xreplace(dimension_map) for j in indices]
mapper[i] = f.indexed[indices]
subs = kwargs.get('subs')
if subs:
# Include the user-supplied substitutions, and use
# `xreplace` for constant folding
processed.append(expr.xreplace({**mapper, **subs}))
else:
processed.append(uxreplace(expr, mapper))
processed = cls._specialize_exprs(processed)
return processed
def Gzz(field, costheta, sintheta, cosphi, sinphi, rho):
"""
3D rotated second order derivative in the direction z
:param field: symbolic data whose derivative we are computing
:param costheta: cosine of the tilt angle
:param sintheta: sine of the tilt angle
:param cosphi: cosine of the azymuth angle
:param sinphi: sine of the azymuth angle
:param space_order: discretization order
:return: rotated second order derivative wrt z
"""
order1 = field.space_order // 2
func = list(retrieve_functions(field))[0]
if func.grid.dim == 2:
return Gzz2d(field, costheta, sintheta, space_order)
x, y, z = func.space_dimensions
Gz = -(sintheta * cosphi * first_derivative(field, dim=x, side=centered,
fd_order=order1) +
sintheta * sinphi * first_derivative(field, dim=y, side=centered,
fd_order=order1) +
costheta * first_derivative(field, dim=z, side=centered,
fd_order=order1))
Gzz = (first_derivative(Gz * sintheta * cosphi * rho,
dim=x, side=centered, fd_order=order1,
matvec=transpose) +
first_derivative(Gz * sintheta * sinphi * rho,
dim=y, side=centered, fd_order=order1,
matvec=transpose) +
first_derivative(Gz * costheta * rho,
dim=z, side=centered, fd_order=order1,
matvec=transpose))
return Gzz
def test_is_on_grid():
grid = Grid((10,))
x = grid.dimensions[0]
x0 = x + .5 * x.spacing
u = Function(name="u", grid=grid, space_order=2)
assert u._is_on_grid
assert not u.subs({x: x0})._is_on_grid
assert all(uu._is_on_grid for uu in retrieve_functions(u.subs({x: x0}).evaluate))
def lower_exprs(expressions, **kwargs):
"""
Lowering an expression consists of the following passes:
* Indexify functions;
* Align Indexeds with the computational domain;
* Apply user-provided substitution;
Examples
--------
f(x - 2*h_x, y) -> f[xi + 2, yi + 4] (assuming halo_size=4)
"""
processed = []
for expr in as_tuple(expressions):
try:
dimension_map = expr.subdomain.dimension_map
except AttributeError:
# Some Relationals may be pure SymPy objects, thus lacking the subdomain
dimension_map = {}
# Handle Functions (typical case)
mapper = {
f: f.indexify(lshift=True, subs=dimension_map)
for f in retrieve_functions(expr)
}
# Handle Indexeds (from index notation)
for i in retrieve_indexed(expr):
f = i.function
# Introduce shifting to align with the computational domain
indices = [(lower_exprs(a) + o)
for a, o in zip(i.indices, f._size_nodomain.left)]
# Apply substitutions, if necessary
if dimension_map:
indices = [j.xreplace(dimension_map) for j in indices]
mapper[i] = f.indexed[indices]
subs = kwargs.get('subs')
# Add dimensions map to the mapper in case dimensions are used
# as an expression, i.e. Eq(u, x, subdomain=xleft)
mapper.update(dimension_map)
if subs:
# Include the user-supplied substitutions, and use
# `xreplace` for constant folding
processed.append(expr.xreplace({**mapper, **subs}))
else:
processed.append(uxreplace(expr, mapper))
if isinstance(expressions, Iterable):
return processed
else:
assert len(processed) == 1
return processed.pop()
def test_time_subsampling_fd(self):
nt = 19
grid = Grid(shape=(11, 11))
x, y = grid.dimensions
time = grid.time_dim
factor = 4
time_subsampled = ConditionalDimension('t_sub', parent=time, factor=factor)
usave = TimeFunction(name='usave', grid=grid, save=(nt+factor-1)//factor,
time_dim=time_subsampled, time_order=2)
dx2 = [indexify(i) for i in retrieve_functions(usave.dt2.evaluate)]
assert dx2 == [usave[time_subsampled - 1, x, y],
usave[time_subsampled + 1, x, y],
usave[time_subsampled, x, y]]
def generate_implicit_exprs(expressions):
"""
Create and add implicit expressions.
Implicit expressions are those not explicitly defined by the user
but instead are requisites of some specified functionality.
Currently, implicit expressions stem from the following:
* MultiSubDomains attached to input equations.
"""
found = {}
processed = []
for e in expressions:
if e.subdomain:
try:
dims = [
d.root for d in e.free_symbols if isinstance(d, Dimension)
]
sub_dims = [d.root for d in e.subdomain.dimensions]
sub_dims.extend(e.subdomain.implicit_dimensions)
dims = [d for d in dims if d not in frozenset(sub_dims)]
dims.extend(e.subdomain.implicit_dimensions)
if e.subdomain not in found:
grid = list(retrieve_functions(e, mode='unique'))[0].grid
found[e.subdomain] = [
i.func(*i.args, implicit_dims=dims)
for i in e.subdomain._create_implicit_exprs(grid)
]
processed.extend(found[e.subdomain])
dims.extend(e.subdomain.dimensions)
new_e = Eq(e.lhs,
e.rhs,
subdomain=e.subdomain,
implicit_dims=dims)
processed.append(new_e)
except AttributeError:
# Not a MultiSubDomain
processed.append(e)
else:
processed.append(e)
return processed
def Gxxyy(field, costheta, sintheta, cosphi, sinphi, rho):
"""
Sum of the 3D rotated second order derivative in the direction x and y.
As the Laplacian is rotation invariant, it is computed as the conventional
Laplacian minus the second order rotated second order derivative in the direction z
Gxx + Gyy = field.laplace - Gzz
:param field: symbolic data whose derivative we are computing
:param costheta: cosine of the tilt angle
:param sintheta: sine of the tilt angle
:param cosphi: cosine of the azymuth angle
:param sinphi: sine of the azymuth angle
:param space_order: discretization order
:return: Sum of the 3D rotated second order derivative in the direction x and y
"""
lap = laplacian(field, rho)
func = list(retrieve_functions(field))[0]
if func.grid.dim == 2:
Gzzr = Gzz2d(field, costheta, sintheta, rho)
else:
Gzzr = Gzz(field, costheta, sintheta, cosphi, sinphi, rho)
return lap - Gzzr
def freesurface(model, eq):
"""
Generate the stencil that mirrors the field as a free surface modeling for
the acoustic wave equation
Parameters
----------
model: Model
Physical model
eq: Eq or List of Eq
Equation to apply mirror to
"""
fs_eq = []
for eq_i in eq:
for p in eq_i._flatten:
lhs, rhs = p.evaluate.args
# Add modulo replacements to to rhs
zfs = model.grid.subdomains['fsdomain'].dimensions[-1]
z = zfs.parent
funcs = retrieve_functions(rhs.evaluate)
mapper = {}
for f in funcs:
zind = f.indices[-1]
if (zind - z).as_coeff_Mul()[0] < 0:
s = sign((zind - z.symbolic_min).subs({
z: zfs,
z.spacing: 1
}))
mapper.update({f: s * f.subs({zind: INT(abs(zind))})})
fs_eq.append(
Eq(lhs,
sign(lhs.indices[-1] - z.symbolic_min) * rhs.subs(mapper),
subdomain=model.grid.subdomains['fsdomain']))
return fs_eq
def _lower_exprs(cls, expressions, **kwargs):
"""
Expression lowering:
* Form and gather any required implicit expressions;
* Evaluate derivatives;
* Flatten vectorial equations;
* Indexify Functions;
* Apply substitution rules;
* Specialize (e.g., index shifting)
"""
# Add in implicit expressions, e.g., induced by SubDomains
expressions = cls._add_implicit(expressions)
# Unfold lazyiness
expressions = flatten([i.evaluate for i in expressions])
# Scalarize tensor expressions
expressions = [j for i in expressions for j in i._flatten]
# Indexification
# E.g., f(x - 2*h_x, y) -> f[xi + 2, yi + 4] (assuming halo_size=4)
processed = []
for expr in expressions:
if expr.subdomain:
dimension_map = expr.subdomain.dimension_map
else:
dimension_map = {}
mapper = {}
# Handle Functions (typical case)
for f in retrieve_functions(expr):
# Get spacing symbols for replacement
spacings = [i.spacing for i in f.dimensions]
# Only keep the ones used as indices
spacings = [
s for i, s in enumerate(spacings)
if s.free_symbols.intersection(f.args[i].free_symbols)
]
# Substitution for each index
subs = {**{s: 1 for s in spacings}, **dimension_map}
# Introduce shifting to align with the computational domain,
# and apply substitutions
indices = [
(a - i + o).xreplace(subs)
for a, i, o in zip(f.args, f.origin, f._size_nodomain.left)
]
mapper[f] = f.indexed[indices]
# Handle Indexeds (from index notation)
for i in retrieve_indexed(expr):
f = i.function
# Introduce shifting to align with the computational domain
indices = [(a + o)
for a, o in zip(i.indices, f._size_nodomain.left)]
# Apply substitutions, if necessary
if dimension_map:
indices = [j.xreplace(dimension_map) for j in indices]
mapper[i] = f.indexed[indices]
subs = kwargs.get('subs')
if subs:
# Include the user-supplied substitutions, and use
# `xreplace` for constant folding
processed.append(expr.xreplace({**mapper, **subs}))
else:
processed.append(uxreplace(expr, mapper))
processed = cls._specialize_exprs(processed)
return processed
def lower_exprs(expressions, **kwargs):
"""
Lowering an expression consists of the following passes:
* Indexify functions;
* Align Indexeds with the computational domain;
* Apply user-provided substitution;
Examples
--------
f(x - 2*h_x, y) -> f[xi + 2, yi + 4] (assuming halo_size=4)
"""
# Normalize subs
subs = {k: sympify(v) for k, v in kwargs.get('subs', {}).items()}
processed = []
for expr in as_tuple(expressions):
try:
dimension_map = expr.subdomain.dimension_map
except AttributeError:
# Some Relationals may be pure SymPy objects, thus lacking the subdomain
dimension_map = {}
# Handle Functions (typical case)
mapper = {
f: lower_exprs(f.indexify(subs=dimension_map), **kwargs)
for f in retrieve_functions(expr)
}
# Handle Indexeds (from index notation)
for i in retrieve_indexed(expr):
f = i.function
# Introduce shifting to align with the computational domain
indices = [(lower_exprs(a) + o)
for a, o in zip(i.indices, f._size_nodomain.left)]
# Substitute spacing (spacing only used in own dimension)
indices = [
i.xreplace({
d.spacing: 1,
-d.spacing: -1
}) for i, d in zip(indices, f.dimensions)
]
# Apply substitutions, if necessary
if dimension_map:
indices = [j.xreplace(dimension_map) for j in indices]
mapper[i] = f.indexed[indices]
# Add dimensions map to the mapper in case dimensions are used
# as an expression, i.e. Eq(u, x, subdomain=xleft)
mapper.update(dimension_map)
# Add the user-supplied substitutions
mapper.update(subs)
# Apply mapper to expression
processed.append(uxreplace(expr, mapper))
if isinstance(expressions, Iterable):
return processed
else:
assert len(processed) == 1
return processed.pop()
