import cloudpickle as pickle
import numpy as np
from devito import Eq, Function, Grid, Operator, SparseFunction, TimeFunction
from devitopro import * # noqaOperator Construction and Execution Good Practices
A good DevitoPRO workflow separates the symbolic description of an Operator from the data used to run it. Build the Operator from small placeholder objects, then pass the runtime Function, TimeFunction, SparseFunction, constants, and iteration bounds to Operator.apply.
This pattern is useful because it keeps construction modular, avoids accidental host allocation, supports runtime-dependent grids, geometry, and data, and makes it practical to pickle a compiled Operator for reuse by worker processes.
Imports
Build Once From Symbolic Objects
Function, TimeFunction, SparseFunction, and related objects are symbolic until their data is touched. In particular, the underlying array is allocated lazily when .data, .data_with_halo, .coordinates.data, or another data accessor is used.
Some memory is not a user-owned data object: compiler-generated temporaries, communication buffers, streaming buffers, and transient device allocations are created and managed by the operator runtime. Use op.estimate_memory(...) when you need a pre-run estimate that includes these arrays without forcing placeholder data allocation.
The construction function below deliberately creates a tiny grid and does not touch any data.
def make_forward_operator(shape=(4, 4), nt=4, space_order=2):
grid = Grid(shape=shape)
m = Function(name='m', grid=grid)
u = TimeFunction(name='u', grid=grid, space_order=space_order,
time_order=1)
usave = TimeFunction(name='usave', grid=grid, save=nt,
space_order=space_order, time_order=1)
eqns = [Eq(u.forward, u + m*u.laplace + 1),
Eq(usave, u.forward)]
return Operator(eqns, name='Forward'), (m, u, usave)
op, (m_symbolic, u_symbolic, usave_symbolic) = make_forward_operator()(m_symbolic._data is None,
u_symbolic._data is None,
usave_symbolic._data is None)The objects used to build the Operator define the symbolic contract: names, function kinds, derivative structure, staggering, transient status, and other metadata that affect generated code. The runtime objects may have different shape, extent, save, npoint, coordinates, or data values, as long as they remain compatible with that contract.
At apply time, the keyword name selects which symbolic parameter is being overridden. The runtime object’s Python variable name and Devito name do not need to match the symbolic object for dispatch, although matching names can make diagnostics easier to read.
Why Decouple Construction And Execution
Construction is where Devito analyzes equations, lowers the symbolic IR, and prepares generated code. Execution is where concrete arrays, sparse coordinates, constants, and iteration bounds are supplied. Keeping these phases separate has three practical benefits:
- The symbolic operator can be built once and reused with many compatible runtime objects.
- The compiled operator can be pickled and shipped to another process so that process does not repeat construction and JIT work.
- The code that defines equations stays separate from orchestration code that owns data, I/O, and scheduling.
Pickle Operators For Workers
Use cloudpickle for Operator pickling. If the operator has already been applied once, or its cfunction has been requested, the compiled shared object is included in the pickle payload and restored at load time. This is the useful case for worker processes because it avoids repeating construction and JIT compilation work.
The in-memory form is:
_ = op.cfunction # JIT-compile now so the shared object is pickled
payload = pickle.dumps(op)
reloaded_op = pickle.loads(payload)
len(payload)The file form is the same two operations around a file handle:
with open('forward.pkl', 'wb') as f:
pickle.dump(op, f)
with open('forward.pkl', 'rb') as f:
op = pickle.load(f)A worker can load the operator first, then supply local runtime values through op.apply(...).
Runtime Overrides
Runtime code creates or receives concrete data objects, then calls op.apply with overrides. The operator parameter names come from the symbolic objects used during construction. The same mechanism applies to dense functions, time-dependent functions, sparse objects, constants, and iteration bounds.
def make_runtime_state(shape, nt, m_value):
grid = Grid(shape=shape)
m = Function(name='m_runtime', grid=grid)
u = TimeFunction(name='u_runtime', grid=grid, space_order=2, time_order=1)
usave = TimeFunction(name='usave_runtime', grid=grid, save=nt,
space_order=2, time_order=1)
m.data[:] = m_value
return {'m': m, 'u': u, 'usave': usave}
runtime_a = make_runtime_state(shape=(12, 10), nt=6, m_value=0.01)
summary = reloaded_op.apply(time_M=4, **runtime_a)
runtime_a['usave'].data.shapeThe same Operator can be applied to another runtime grid. No new symbolic operator is needed just because the runtime problem size changed.
runtime_b = make_runtime_state(shape=(18, 14), nt=8, m_value=0.02)
summary = reloaded_op.apply(time_M=6, **runtime_b)
runtime_b['usave'].data.shapeSparse objects follow the same rule: construct with placeholders, then pass the runtime sparse object through the keyword associated with the symbolic placeholder.
def make_sampling_operator(shape=(4, 4), npoint=1):
grid = Grid(shape=shape, extent=(1., 1.))
f = Function(name='f', grid=grid)
points = SparseFunction(name='points', grid=grid, npoint=npoint)
return Operator(points.interpolate(expr=f), name='Sample'), (f, points)
sample_op, _ = make_sampling_operator()
runtime_grid = Grid(shape=(9, 9), extent=(1., 1.))
f = Function(name='f', grid=runtime_grid)
points = SparseFunction(name='points', grid=runtime_grid, npoint=2)
f.data[:] = np.arange(f.data.size, dtype=f.dtype).reshape(f.data.shape)
points.coordinates.data[:] = np.array([[0.25, 0.25], [0.75, 0.75]])
sample_op.apply(f=f, points=points)
points.dataApplying The Pattern In Shot Loops
In an inversion framework, the outer loop can create the runtime grid, model fields, wavefields, and acquisition objects for each shot while reusing the same constructed or unpickled operator. This keeps compilation overhead out of the per-shot path when the symbolic contract is unchanged.
shots = [((8, 16), 5, 0.03),
((10, 12), 6, 0.04)]
results = []
for shape, nt, m_value in shots:
runtime = make_runtime_state(shape=shape, nt=nt, m_value=m_value)
summary = reloaded_op.apply(time_M=nt-2, **runtime)
results.append(runtime['usave'].data.shape)
resultsData Access And External Storage
Write or inspect data through f.data for the physical domain, or through f.data_with_halo to include the halo points surrounding the physical domain. Do not worry about padding: DevitoPRO, to maximize memory-aligned accesses, will automatically allocate a few extra points (“padding”) along the contiguous data dimension, based on the user-provided shape and the underlying architecture.
When an external application needs a Devito object backed by pre-existing storage, use DataReference instead of assigning internal fields manually:
from devito.data.allocators import DataReference
f = Function(name='f', grid=grid, allocator=DataReference(array_like))See submodules/devito/tests/test_data.py::TestDataReference for concrete examples.
It is still useful to know what the lower-level accessors mean when debugging or integrating with native code:
f._data: once allocated, the full local allocation, including inner halo and padding. This is the only buffer to think of as one fully contiguous Devito-owned array.f._data_allocated: the same full local allocation asf._data, but accessed through a property that may allocate storage and first ensures halo values are valid. Under MPI, this accessor may trigger a halo exchange before returning the allocation.f.data,f.data_with_halo,f._data_with_outhalo, and related accessors are views intof._data; they are domain-oriented APIs, not separate contiguous allocations.
Treat ._data and ._data_allocated as implementation details. They explain the memory layout, but they should not be the normal user-level interface for wiring external storage into Devito objects.
halo_grid = Grid(shape=(6, 6))
a = Function(name='a', grid=halo_grid, space_order=4)
a.data.shape, a.data_with_halo.shape, a._data_allocated.shapeGPU-Only Transient Data
For device backends, is_transient=True marks a function as scratch data that does not need a persistent host allocation. DevitoPRO can then allocate it directly on the device for the duration of Operator.apply.
Use this for temporaries that are not initialized from Python and are not read back from Python. Build-time and runtime objects must agree on is_transient; DevitoPRO checks this because the generated data movement is different.
transient_grid = Grid(shape=(4, 4))
scratch = TimeFunction(name='scratch', grid=transient_grid,
is_transient=True)
scratch.is_transient, scratch._data is NoneKeep Saved Writes And Reads Separate
Saved TimeFunctions may be streamed through device, host, and disk layers. A robust pattern is to keep the operator that writes a saved field separate from the operator that reads it later:
# Forward operator: write saved states.
forward_eqns = [Eq(u.forward, ...),
Eq(usave, u.forward)]
# Backward/imaging operator: read saved states.
backward_eqns = [Eq(v.backward, ... usave ...)]This matches the data streaming model and keeps the generated data movement easier to reason about. See data_streaming.ipynb for the detailed streaming behavior and layer controls.
Checklist
- Build symbolic operators outside loops that only change runtime values.
- Do not touch
.dataon placeholder objects unless construction really needs initialized data. - Pass runtime dense and sparse objects through
op.apply(name=runtime_object, ...); the keyword is the override key. - Reuse one applied or
cfunction-loaded operator withcloudpicklewhen workers need to avoid JIT work. - Use
is_transient=Truefor GPU scratch fields that do not need host data. - Use
.dataor.data_with_halofor ordinary Python-side initialization and inspection; let DevitoPRO manage padding and alignment. - Use
DataReferencewhen a Devito object must attach to pre-existing storage. - Keep saved
TimeFunctionwrites and reads in separate operators.