Skip to content

Commit

Permalink
Merge pull request #2434 from devitocodes/coeffs-revamp
Browse files Browse the repository at this point in the history 
api: revamp custom coefficients API
  • Loading branch information
@mloubout
mloubout authored Sep 24, 2024
2 parents cec0542 + 77e0021 commit 4c8be71
Show file tree
Hide file tree
Showing 16 changed files with 291 additions and 862 deletions.
18 changes: 4 additions & 14 deletions devito/builtins/initializers.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,14 +144,7 @@ def define(self, dimensions):
def create_gaussian_weights(sigma, lw):
weights = [w/w.sum() for w in (np.exp(-0.5/s**2*(np.linspace(-l, l, 2*l+1))**2)
for s, l in zip(sigma, lw))]
processed = []
for w in weights:
temp = list(w)
while len(temp) < 2*max(lw)+1:
temp.insert(0, 0)
temp.append(0)
processed.append(np.array(temp))
return as_tuple(processed)
return as_tuple(np.array(w) for w in weights)

def fset(f, g):
indices = [slice(l, -l, 1) for _, l in zip(g.dimensions, lw)]
Expand Down Expand Up @@ -188,8 +181,7 @@ def fset(f, g):
shape_padded = tuple([np.array(s) + 2*l for s, l in zip(shape, lw)])
grid = dv.Grid(shape=shape_padded, subdomains=objective_domain)

f_c = dv.Function(name='f_c', grid=grid, space_order=2*max(lw),
coefficients='symbolic', dtype=dtype)
f_c = dv.Function(name='f_c', grid=grid, space_order=2*max(lw), dtype=dtype)
f_o = dv.Function(name='f_o', grid=grid, dtype=dtype)

weights = create_gaussian_weights(sigma, lw)
Expand All @@ -201,10 +193,8 @@ def fset(f, g):
options = []

lhs.append(f_o)
rhs.append(dv.generic_derivative(f_c, d, 2*l, 1))
coeffs = dv.Coefficient(1, f_c, d, w)
options.append({'coefficients': dv.Substitutions(coeffs),
'subdomain': grid.subdomains['objective_domain']})
rhs.append(dv.generic_derivative(f_c, d, 2*l, 1, weights=w))
options.append({'subdomain': grid.subdomains['objective_domain']})

lhs.append(f_c)
rhs.append(f_o)
Expand Down
15 changes: 15 additions & 0 deletions devito/deprecations.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
from functools import cached_property
from warnings import warn


class DevitoDeprecation():

@cached_property
def coeff_warn(self):
warn("The Coefficient API is deprecated and will be removed, coefficients should"
"be passed directly to the derivative object `u.dx(weights=...)",
DeprecationWarning, stacklevel=2)
return


deprecations = DevitoDeprecation()
289 changes: 5 additions & 284 deletions devito/finite_differences/coefficients.py
View file
Original file line number Diff line number Diff line change
@@ -1,73 +1,15 @@
from functools import cached_property
from devito.deprecations import deprecations

import numpy as np

from devito.finite_differences import Weights, generate_indices
from devito.finite_differences.tools import numeric_weights
from devito.tools import filter_ordered, as_tuple

__all__ = ['Coefficient', 'Substitutions', 'default_rules']
__all__ = ['Coefficient', 'Substitutions']


class Coefficient:
"""
Prepare custom coefficients to pass to a Substitutions object.
Parameters
----------
deriv_order : int
The order of the derivative being taken.
function : Function
The function for which the supplied coefficients
will be used.
dimension : Dimension
The dimension with respect to which the
derivative is being taken.
weights : np.ndarray
The set of finite difference weights
intended to be used in place of the standard
weights (obtained from a Taylor expansion).
Examples
--------
>>> import numpy as np
>>> from devito import Grid, Function, Coefficient
>>> grid = Grid(shape=(4, 4))
>>> u = Function(name='u', grid=grid, space_order=2, coefficients='symbolic')
>>> x, y = grid.dimensions
Now define some partial d/dx FD coefficients of the Function u:
>>> u_x_coeffs = Coefficient(1, u, x, np.array([-0.6, 0.1, 0.6]))
And some partial d^2/dy^2 FD coefficients:
>>> u_y2_coeffs = Coefficient(2, u, y, np.array([0.0, 0.0, 0.0]))
"""

def __init__(self, deriv_order, function, dimension, weights):

self._check_input(deriv_order, function, dimension, weights)

# Ensure the given set of weights is the correct length
try:
wl = weights.shape[-1]-1
except AttributeError:
wl = len(weights)-1
if dimension.is_Time:
if wl != function.time_order:
raise ValueError("Number of FD weights provided does not "
"match the functions space_order")
elif dimension.is_Space:
if wl != function.space_order:
raise ValueError("Number of FD weights provided does not "
"match the functions space_order")

deprecations.coeff_warn
self._weights = weights
self._deriv_order = deriv_order
self._function = function
self._dimension = dimension
self._weights = weights

@property
def deriv_order(self):
Expand All @@ -84,242 +26,21 @@ def dimension(self):
"""The dimension to which the coefficients will be applied."""
return self._dimension

@property
def index(self):
"""
The dimension to which the coefficients will be applied plus the offset
in that dimension if the function is staggered.
"""
return self._dimension

@property
def weights(self):
"""The set of weights."""
return self._weights

def _check_input(self, deriv_order, function, dimension, weights):
if not isinstance(deriv_order, int):
raise TypeError("Derivative order must be an integer")
try:
if not function.is_Function:
raise TypeError("Object is not of type Function")
except AttributeError:
raise TypeError("Object is not of type Function")
try:
if not dimension.is_Dimension:
raise TypeError("Coefficients must be attached to a valid dimension")
except AttributeError:
raise TypeError("Coefficients must be attached to a valid dimension")
try:
weights.is_Function is True
except AttributeError:
if not isinstance(weights, np.ndarray):
raise TypeError("Weights must be of type np.ndarray or a Devito Function")
return


class Substitutions:
"""
Devito class to convert Coefficient objects into replacent rules
to be applied when constructing a Devito Eq.
Examples
--------
>>> from devito import Grid, TimeFunction, Coefficient
>>> grid = Grid(shape=(4, 4))
>>> u = TimeFunction(name='u', grid=grid, space_order=2, coefficients='symbolic')
>>> x, y = grid.dimensions
Now define some partial d/dx FD coefficients of the Function u:
>>> u_x_coeffs = Coefficient(2, u, x, np.array([-0.6, 0.1, 0.6]))
And now create our Substitutions object to pass to equation:
>>> from devito import Substitutions
>>> subs = Substitutions(u_x_coeffs)
Now create a Devito equation and pass to it 'subs'
>>> from devito import Eq
>>> eq = Eq(u.dt+u.dx2, coefficients=subs)
When evaluated, the derivatives will use the custom coefficients. We can
check that by
>>> eq.evaluate
Eq(-u(t, x, y)/dt + u(t + dt, x, y)/dt + 0.1*u(t, x, y) - \
0.6*u(t, x - h_x, y) + 0.6*u(t, x + h_x, y), 0)
Notes
-----
If a Function is declared with 'symbolic' coefficients and no
replacement rules for any derivative appearing in a Devito equation,
the coefficients will revert to those of the 'default' Taylor
expansion.
"""

def __init__(self, *args):

deprecations.coeff_warn
if any(not isinstance(arg, Coefficient) for arg in args):
raise TypeError("Non Coefficient object within input")

self._coefficients = args
self._function_list = self.function_list
self._rules = self.rules

@property
def coefficients(self):
"""The Coefficient objects passed."""
return self._coefficients

@cached_property
def function_list(self):
return filter_ordered((i.function for i in self.coefficients), lambda i: i.name)

@cached_property
def rules(self):

def generate_subs(i):

deriv_order = i.deriv_order
function = i.function
dim = i.dimension
index = i.index
weights = i.weights

if isinstance(weights, np.ndarray):
fd_order = len(weights)-1
else:
fd_order = weights.shape[-1]-1

subs = {}

indices, x0 = generate_indices(function, dim, fd_order, side=None)

# NOTE: This implementation currently assumes that indices are ordered
# according to their position in the FD stencil. This may not be the
# case in all schemes and should be changed such that the weights are
# passed as a dictionary of the form {pos: w} (or something similar).
if isinstance(weights, np.ndarray):
for j in range(len(weights)):
subs.update({function._coeff_symbol
(indices[j], deriv_order,
function, index): weights[j]})
else:
shape = weights.shape
x = weights.dimensions
for j in range(shape[-1]):
idx = list(x)
idx[-1] = j
subs.update({function._coeff_symbol
(indices[j], deriv_order, function, index):
weights[as_tuple(idx)]})

return subs

# Figure out when symbolic coefficients can be replaced
# with user provided coefficients and, if possible, generate
# replacement rules
rules = {}
for i in self.coefficients:
rules.update(generate_subs(i))

return rules


def default_rules(obj, functions):

from devito.symbolics.search import retrieve_dimensions

def generate_subs(deriv_order, function, index, indices):
dim = retrieve_dimensions(index)[0]

if dim.is_Time:
fd_order = function.time_order
elif dim.is_Space:
fd_order = function.space_order
else:
# Shouldn't arrive here
raise TypeError("Dimension type not recognised")

subs = {}

# Actual FD used indices and weights
if deriv_order == 1 and fd_order == 2:
fd_order = 1

coeffs = numeric_weights(function, deriv_order, indices, index)

for (c, i) in zip(coeffs, indices):
subs.update({function._coeff_symbol(i, deriv_order, function, index): c})

return subs

# Determine which 'rules' are missing
sym = get_sym(functions)
terms = obj.find(sym)
for i in obj.find(Weights):
for w in i.weights:
terms.update(w.find(sym))

args_present = {}
for term in terms:
key = term.args[1:]
idx = term.args[0]
args_present.setdefault(key, []).append(idx)

subs = obj.substitutions
if subs:
args_provided = [(i.deriv_order, i.function, i.index)
for i in subs.coefficients]
else:
args_provided = []

# NOTE: Do we want to throw a warning if the same arg has
# been provided twice?
args_provided = list(set(args_provided))

rules = {}
not_provided = {}
for i0 in args_present:
if any(i0 == i1 for i1 in args_provided):
# Perfect match, as expected by the legacy custom coeffs API
continue

# TODO: to make cross-derivs work, we must relax `not_provided` by
# checking not for equality, but rather for inclusion. This is ugly,
# but basically a major revamp is the only alternative... and for now,
# it does the trick
mapper = {}
deriv_order, expr, dim = i0
try:
for k, v in subs.rules.items():
ofs, do, f, d = k.args
is_dim = dim == expr.indices_ref[d]
if deriv_order == do and is_dim and f in expr._functions:
mapper[k.func(ofs, do, expr, expr.indices_ref[d])] = v
except AttributeError:
# assert subs is None
pass

if mapper:
rules.update(mapper)
else:
not_provided.update({i0: args_present[i0]})

for i, indices in not_provided.items():
rules = {**rules, **generate_subs(*i, indices)}

return rules


def get_sym(functions):
for f in functions:
try:
sym = f._coeff_symbol
return sym
except AttributeError:
pass
# Shouldn't arrive here
raise TypeError("Failed to retreive symbol")
Loading

0 comments on commit 4c8be71

Please sign in to comment.