Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

api: revamp custom coefficients API #2434

Merged
merged 9 commits into from
Sep 24, 2024
Merged

api: revamp custom coefficients API #2434

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
e586281
api: revamp custom coefficients API to be passed to derivatives directly
mloubout Jul 31, 2024
803404a
misc: flake8
mloubout Jul 31, 2024
a553d5c
builtins: fix gaussian kernel for non-square
mloubout Jul 31, 2024
d88bf3d
api: add w= for weights shortcut
mloubout Aug 7, 2024
d6006ae
examples: Update custom coefficients notebook text to reflect API cha…
EdCaunt Sep 9, 2024
20313d8
examples: avoid too tight cfl for elastic
mloubout Sep 9, 2024
7c7994c
builtins: remove small clutter
mloubout Sep 12, 2024
958720a
api: fix pickle with derivative
mloubout Sep 13, 2024
77e0021
api: avoid duplicate deprecation warnings
mloubout Sep 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 4 additions & 14 deletions devito/builtins/initializers.py
View file
Open in desktop
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
293 changes: 9 additions & 284 deletions devito/finite_differences/coefficients.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,73 +1,17 @@
from functools import cached_property
from warnings import warn

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

warn("The Coefficient API is deprecated and will be removed, coefficients should"
"be passed directly to the derivative object `u.dx(weights=...)",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is creating one warning per Coefficient, another option is to register a callback to call at atexit(), or more simply use a global file-level flag to emit the warning only once ?

DeprecationWarning, stacklevel=2)
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 +28,23 @@ 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):

warn("The Coefficient API is deprecated and will be removed, coefficients should"
"be passed directly to the derivative object `u.dx(weights=...)",
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

same as before

DeprecationWarning, stacklevel=2)
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
Loading