Extending the Framework

Every built-in solver is a thin subclass of Propagator (recipes/solver.py). The base class owns operator assembly, wavefield management, and the run loop; subclasses supply only the PDE-specific pieces. This page explains what those pieces are and how to override them.

---

1. Architecture overview

Propagator                         # recipes/solver.py
├── state / saved_state            # wavefield TimeFunctions (cached)
├── src / rec                      # source & receiver SparseTimeFunctions
├── fwd_op()                       # Devito Operator for forward pass
├── adj_op()                       # Devito Operator for adjoint pass
├── jacobian_adj_op()              # Devito Operator for gradient pass
├── forward() / adjoint()          # high-level run methods
└── jacobian_adjoint()

The equation list assembled inside fwd_op is:

_make_stencils(fw=True)   ➜  PDE update equations
acquisition_eq()          ➜  source injection + receiver interpolation
eq_save()                 ➜  snapshot equations  (when save=True)
extra                     ➜  any additional per-step equations

---

2. Minimum implementation

A concrete solver must set four class attributes and implement two methods.

from devito import Eq, solve, TimeFunction
from devito.tools import as_tuple
from recipes.solver import Propagator


class MySolver(Propagator):

    # -----------------------------------------------------------------
    # Class attributes
    # -----------------------------------------------------------------

    # Time discretization order (1 or 2)
    _time_order = 2

    # Names of physical parameters read from self.model.
    # Each name must correspond to a Function or Constant on the model.
    _parameters = ['m', 'b']

    # Wavefield definitions: list of (name, TimeFunction-like type) pairs.
    # Use devitopro.TimeFunction so dtype and staggering options are available.
    _fields = [('u', TimeFunction)]

    # Absorbing boundary convention: 'damp' (additive) or 'mask' (multiplicative)
    _abc = 'damp'

    # -----------------------------------------------------------------
    # Required methods
    # -----------------------------------------------------------------

    def _make_stencils(self, fw=False, **kwargs):
        """
        Return (equations, temporaries).

        equations   -- list of Devito Eq for the wavefield time update.
        temporaries -- list of Eq for auxiliary quantities computed before
                       the update (e.g. memory variables, precomputed products).

        Use self._physical to restrict equations to the interior subdomain
        and self._boundary for the absorbing layer region.
        The argument fw=True selects the forward direction; fw=False selects
        the adjoint direction.
        """
        u = self.state['u']
        unext = u.forward if fw else u.backward
        m, b, damp = self.model.m, self.model.b, self.model.damp

        pde = m * u.dt2 - b * u.laplace + damp * u.dt
        eq = Eq(unext, solve(pde, unext), subdomain=self._physical)
        return [eq], []

    @property
    def pressure(self):
        """
        Symbolic expression delivered to receivers and exposed via
        solver.pressure / solver.pressure_data.
        """
        return self.state['u']

    def sensitivity(self, params='m'):
        """
        Gradient equations accumulated during the adjoint pass.

        Access saved forward snapshots via self.saved_state (dict keyed by
        '<name>_save'), adjoint wavefields via self.state, and the gradient
        accumulator via self.perturbation(param).
        """
        params = as_tuple(params)
        param  = params[0]

        u_fwd = self.saved_state['u_save']
        v     = self.state['u']
        dm    = self.perturbation(param=param)
        dt    = u_fwd.indices[0].spacing

        return [Eq(dm, dm - dt * u_fwd * v.dt)]

Register the solver so it is available via recipes_registry and on the CLI:

from recipes import recipes_registry
recipes_registry['my-solver'] = MySolver

---

3. Customisation hooks

3.1 Wavefield definition

Hook	Kind	Purpose
_fields	class attribute	List of (name, FunctionType) pairs that define the wavefields
_time_order	class attribute	Sets the rolling buffer depth and the time derivative order
save_fields	property	Subset of _fields to snapshot — return only what the adjoint needs
save_expr(u)	method	Expression to write into the snapshot (default: u; use u.dt for velocity)

@property
def save_fields(self):
    """Save only velocity; recompute stress on the fly during adjoint."""
    return [self._fields[0]]   # first field only

def save_expr(self, u):
    """Store particle velocity instead of displacement."""
    return u.dt

3.2 Source injection

Hook	Kind	Purpose
_src_fields	property	Fields to inject the source into
src_scale(fw)	method	Scaling factor applied to the source expression
acquisition_eq(fw, src_field, measurement)	method	Full src/rec equation builder

@property
def _src_fields(self):
    """Inject into diagonal stress components (elastic convention)."""
    return self.state['tau'].diagonal()

def src_scale(self, fw=True):
    """Scale by dt² / (m·b) so the source amplitude is correct."""
    dt = self.grid.time_dim.spacing
    return dt**2 / (self.model.m * self.model.b)

3.3 Receiver interpolation

Hook	Kind	Purpose
_rec_expr	property	Expression interpolated at receiver positions (default: self.pressure)

@property
def _rec_expr(self):
    """Record the hydrostatic pressure (trace of stress tensor)."""
    return self.state['tau'].trace()

3.4 Free-surface mirroring

Override acquisition_eq and _make_stencils using the utilities from recipes.utils:

from recipes.utils import mirror, mirror_source

def acquisition_eq(self, fw=True, measurement=True, src_field=None):
    s_eq, r_eq = super().acquisition_eq(
        fw=fw, measurement=measurement, src_field=src_field)
    if self.model.fs:
        s_eq = mirror_source(self.model, s_eq)
    return s_eq, r_eq

def _make_stencils(self, fw=False, **kwargs):
    # ... build update_u ...
    if self.model.fs:
        update_u += mirror(self.model, u.forward if fw else u.backward)
    return update_u, []

3.5 Extra per-step equations

The extra property appends equations to every time-step iteration of the operator. Use it for quantities that must accumulate over the entire propagation (illumination, energy, diagnostics):

from devito import Function, Inc

@property
def extra(self):
    u   = self.state['u']
    ill = Function(name='illumination', grid=self.grid)
    return [Inc(ill, u * u)]

After solver.forward(), ill.data contains the per-point energy integral.

3.6 Runtime keyword preparation

Override _prepare_kw to pass extra objects to op.apply — for example, memory variables that must be allocated before execution:

def _prepare_kw(self, **kwargs):
    kw = super()._prepare_kw(**kwargs)
    r = memory_variable(self.state['u'])
    kw[r.name] = r
    return kw

3.7 Optimisation defaults

Override _default_opt to set physics-specific compiler flags:

@property
def _default_opt(self):
    return ('advanced', {'openmp': True, 'errctl': 'basic'})

---

4. Custom source and receiver geometry

Replace acquisition_eq entirely to implement multi-source experiments, marine streamer geometries, or any other non-standard acquisition:

import numpy as np
from examples.seismic import RickerSource, Receiver, TimeAxis


class MultiSrcSolver(AcousticIsotropic):
    """Forward pass with two simultaneous sources."""

    def acquisition_eq(self, fw=True, src_field=None, measurement=True):
        nt   = self.options['nt']
        time = TimeAxis(start=0, step=self.model.critical_dt, num=nt)

        src = RickerSource(name='src', grid=self.grid,
                           f0=self.options['f0'], time_range=time,
                           npoint=2)
        src.coordinates.data[0] = [500.,  20.]
        src.coordinates.data[1] = [1000., 20.]

        scale  = self.src_scale(fw=fw)
        fields = self.src_fields(fw=fw, src_field=src_field)
        s_eq   = src.inject(field=fields, expr=scale * src)

        r_eq = self.rec.interpolate(expr=self._rec_expr) if measurement else []
        return s_eq, r_eq

---

5. Custom imaging condition

Override sensitivity to implement any imaging condition — e.g. a directional gradient, an elastic imaging condition, or one weighted by a precomputed operator.

def sensitivity(self, params='m'):
    """Horizontal-derivative imaging condition: dm ~ ∫ u_fwd.dx · v_adj dt"""
    from devito.tools import as_tuple
    params = as_tuple(params)
    param  = params[0]

    u_fwd = self.saved_state['u_save']
    v     = self.state['u']
    dm    = self.perturbation(param=param)
    dt    = u_fwd.indices[0].spacing

    return [Eq(dm, dm - dt * u_fwd.dx * v.dx)]

---

6. Complete example — illumination-weighted gradient

This self-contained example wires together several of the hooks above.

import numpy as np
from devito import Eq, Inc, Function
from devito.tools import as_tuple

from recipes import recipes_registry, layered_model
from recipes.isotropic_acoustic import AcousticIsotropic


class IllumWeightedAcoustic(AcousticIsotropic):
    """
    Acoustic solver that weights the gradient by source illumination.

    Imaging condition:
        dm(x) = -∫ u_fwd(x,t) · v_adj.dt(x,t) / I(x) dt

    where the illumination I(x) = ∫ u_fwd²(x,t) dt is accumulated as an
    extra equation during the forward pass.
    """

    @property
    def _illumination(self):
        """Persistent illumination Function (allocated once per solver)."""
        if not hasattr(self, '_ill'):
            self._ill = Function(name='ill', grid=self.grid)
        return self._ill

    @property
    def extra(self):
        u = self.state['u']
        return [Inc(self._illumination, u * u)]

    def sensitivity(self, params='m'):
        params = as_tuple(params)
        param  = params[0]

        u_fwd    = self.saved_state['u_save']
        v        = self.state['u']
        ill_safe = self._illumination + 1e-10   # avoid division by zero
        dm       = self.perturbation(param=param)
        dt       = u_fwd.indices[0].spacing

        return [Eq(dm, dm - dt * u_fwd * v.dt / ill_safe)]


recipes_registry['illum-acoustic'] = IllumWeightedAcoustic

# ── usage ────────────────────────────────────────────────────────────
model  = layered_model('iso-acoustic', shape=(101, 101), spacing=(10., 10.))
solver = IllumWeightedAcoustic(model, {'nt': 500, 'space_order': 8})

solver.forward(save=True)

# Inject residuals
solver.rec.data[:] -= np.random.randn(*solver.rec.data.shape) * 1e-6

solver.jacobian_adjoint(params='m')
grad = solver.perturbation(param='m').data
ill  = solver._illumination.data

---

7. Hook quick-reference

Hook	Type	Required?	Override to…
_time_order	class attr	yes	Change time discretization order
_parameters	class attr	yes	Declare model parameters used by the PDE
_fields	class attr	yes	Define wavefield names and types
_abc	class attr	yes	Select absorbing boundary convention ('damp'/'mask')
_make_stencils(fw)	method	yes	Define the PDE update equations
pressure	property	yes	Set which expression is delivered to receivers
sensitivity(params)	method	yes	Define the gradient imaging condition
save_fields	property	no	Save only a subset of wavefields
save_expr(u)	method	no	Change what is written into the snapshot
_src_fields	property	no	Choose which fields receive source injection
_rec_expr	property	no	Choose which expression is recorded
src_scale(fw)	method	no	Apply PDE-normalisation scaling to source
acquisition_eq(...)	method	no	Full custom src/rec injection
extra	property	no	Add per-step accumulation equations
_prepare_kw(**kw)	method	no	Pass extra runtime arguments to op.apply
_default_opt	property	no	Override default compiler optimisation