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()