Skip to content

Commit

Permalink
api: add harmonic averaging option
Browse files Browse the repository at this point in the history
  • Loading branch information
mloubout committed Sep 11, 2024
1 parent 4698a7d commit 69fc6f2
Show file tree
Hide file tree
Showing 8 changed files with 171 additions and 134 deletions.
2 changes: 1 addition & 1 deletion devito/finite_differences/differentiable.py
Original file line number Diff line number Diff line change
Expand Up @@ -1033,7 +1033,7 @@ def _(expr, x0, **kwargs):
def _(expr, x0, **kwargs):
from devito.finite_differences.derivative import Derivative
x0_expr = {d: v for d, v in x0.items() if v is not expr.indices_ref[d]}
if x0_expr and not expr.is_parameter:
if x0_expr:

This comment has been minimized.

Copy link
@guaacoelho

guaacoelho Nov 22, 2024

Hi @mloubout,

Could you clarify the reasoning behind this change? I'm trying to understand if there are specific cases where not expr.is_parameter was no longer necessary or if this simplifies the logic in a particular context.

Thanks in advance for your insights!

dims = tuple((d, 0) for d in x0_expr)
fd_o = tuple([2]*len(dims))
return Derivative(expr, *dims, fd_order=fd_o, x0=x0_expr)._evaluate(**kwargs)
Expand Down
20 changes: 16 additions & 4 deletions devito/types/basic.py
Original file line number Diff line number Diff line change
Expand Up @@ -852,7 +852,7 @@ class AbstractFunction(sympy.Function, Basic, Pickable, Evaluable):
"""

__rkwargs__ = ('name', 'dtype', 'grid', 'halo', 'padding', 'ghost',
'alias', 'space', 'function', 'is_transient')
'alias', 'space', 'function', 'is_transient', 'avg_mode')

def __new__(cls, *args, **kwargs):
# Preprocess arguments
Expand Down Expand Up @@ -983,6 +983,11 @@ def __init_finalize__(self, *args, **kwargs):
# certain optimizations, such as avoiding memory copies
self._is_transient = kwargs.get('is_transient', False)

# Averaging mode for off the grid evaluation
self._avg_mode = kwargs.get('avg_mode', 'arithmetic')
assert self._avg_mode in ['arithmetic', 'harmonic'], "Accepted avg_mode " \
"values are 'arithmetic' or 'harmonic', invalid %s" % self._avg_mode

@classmethod
def __args_setup__(cls, *args, **kwargs):
"""
Expand Down Expand Up @@ -1147,13 +1152,16 @@ def _evaluate(self, **kwargs):
return self

# Base function
retval = self.function
retval = 1 / self.function if self._avg_mode == 'harmonic' else self.function
# Apply interpolation from inner most dim
for d, i in self._grid_map.items():
retval = retval.diff(d, 0, fd_order=2, x0={d: i})
retval = retval.diff(d, deriv_order=0, fd_order=2, x0={d: i}).evaluate
if self._avg_mode == 'harmonic':
retval = 1 / retval

# Evaluate. Since we used `self.function` it will be on the grid when evaluate
# is called again within FD
return retval.evaluate
return retval.evaluate.expand()

@property
def shape(self):
Expand Down Expand Up @@ -1255,6 +1263,10 @@ def is_const(self):
def is_transient(self):
return self._is_transient

@property
def avg_mode(self):
return self._avg_mode

@property
def alias(self):
return self._alias
Expand Down
4 changes: 2 additions & 2 deletions examples/seismic/elastic/elastic_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,8 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
def test_elastic(dtype):
_, _, _, [rec1, rec2, v, tau] = run(dtype=dtype)
assert np.isclose(norm(rec1), 19.25636, atol=1e-3, rtol=0)
assert np.isclose(norm(rec2), 0.6276, atol=1e-3, rtol=0)
assert np.isclose(norm(rec1), 19.43005, atol=1e-3, rtol=0)
assert np.isclose(norm(rec2), 0.6345, atol=1e-3, rtol=0)


@pytest.mark.parametrize('shape', [(51, 51), (16, 16, 16)])
Expand Down
10 changes: 7 additions & 3 deletions examples/seismic/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -170,12 +170,12 @@ def physical_params(self, **kwargs):
return {i.name: kwargs.get(i.name, i) or i for i in known}

def _gen_phys_param(self, field, name, space_order, is_param=True,
default_value=0, **kwargs):
default_value=0, avg_mode='arithmetic', **kwargs):
if field is None:
return default_value
if isinstance(field, np.ndarray):
function = Function(name=name, grid=self.grid, space_order=space_order,
parameter=is_param)
parameter=is_param, avg_mode=avg_mode)
initialize_function(function, field, self.padsizes)
else:
function = Constant(name=name, value=field, dtype=self.grid.dtype)
Expand Down Expand Up @@ -307,7 +307,11 @@ def _initialize_physics(self, vp, space_order, **kwargs):
vs = kwargs.pop('vs')
self.lam = self._gen_phys_param((vp**2 - 2. * vs**2)/b, 'lam', space_order,
is_param=True)
self.mu = self._gen_phys_param(vs**2 / b, 'mu', space_order, is_param=True)
# Need to add small value to avoid division by zero
if isinstance(vs, np.ndarray):
vs = vs + 1e-12
self.mu = self._gen_phys_param(vs**2 / b, 'mu', space_order, is_param=True,
avg_mode='harmonic')
else:
# All other seismic models have at least a velocity
self.vp = self._gen_phys_param(vp, 'vp', space_order)
Expand Down
236 changes: 114 additions & 122 deletions examples/seismic/tutorials/06_elastic_varying_parameters.ipynb

Large diffs are not rendered by default.

2 changes: 1 addition & 1 deletion examples/seismic/viscoelastic/viscoelastic_example.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ def run(shape=(50, 50), spacing=(20.0, 20.0), tn=1000.0,
@pytest.mark.parametrize("dtype", [np.float32, np.float64])
def test_viscoelastic(dtype):
_, _, _, [rec1, rec2, v, tau] = run(dtype=dtype)
assert np.isclose(norm(rec1), 12.28040, atol=1e-3, rtol=0)
assert np.isclose(norm(rec1), 12.30114, atol=1e-3, rtol=0)
assert np.isclose(norm(rec2), 0.312462, atol=1e-3, rtol=0)


Expand Down
2 changes: 1 addition & 1 deletion examples/userapi/03_subdomains.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -941,7 +941,7 @@
"outputs": [],
"source": [
"from devito import norm\n",
"assert np.isclose(norm(v[0]), 0.10312, rtol=1e-4)"
"assert np.isclose(norm(v[0]), 0.10301, rtol=1e-4)"
]
},
{
Expand Down
29 changes: 29 additions & 0 deletions tests/test_differentiable.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,8 @@
from itertools import product

import sympy
import pytest

from devito import Function, Grid, Differentiable, NODE
from devito.finite_differences.differentiable import Add, Mul, Pow, diffify, interp_for_fd

Expand Down Expand Up @@ -82,3 +86,28 @@ def test_interp():
sa + interp_for_fd(a, {x: x + x.spacing/2}, expand=True))
assert sp_diff(interp_for_fd(a + sa, {x: x}, expand=True),
a + interp_for_fd(sa, {x: x}, expand=True))


@pytest.mark.parametrize('ndim', [1, 2, 3])
def test_avg_mode(ndim):
grid = Grid([11]*ndim)
v = Function(name='v', grid=grid, staggered=grid.dimensions)
a0 = Function(name="a0", grid=grid)
a = Function(name="a", grid=grid, parameter=True)
b = Function(name="b", grid=grid, parameter=True, avg_mode='harmonic')

a0_avg = a0._eval_at(v)
a_avg = a._eval_at(v).evaluate
b_avg = b._eval_at(v).evaluate

assert a0_avg == a0

# Indices around the point at the center of a cell
all_shift = tuple(product(*[[0, 1] for _ in range(ndim)]))
args = [{d: d + i * d.spacing for d, i in zip(grid.dimensions, s)} for s in all_shift]

# Default is arithmetic average
assert sympy.simplify(a_avg - 0.5**ndim * sum(a.subs(arg) for arg in args)) == 0

# Harmonic average, h(a[.5]) = 1/(.5/a[0] + .5/a[1])
assert sympy.simplify(b_avg - 1/(0.5**ndim * sum(1/b.subs(arg) for arg in args))) == 0

0 comments on commit 69fc6f2

Please sign in to comment.