Workflow

tide.workflow composes structured Maxwell operators into shot-batched, distributed, loss, and optimization workflows. Physics and derivative execution remain owned by MaxwellTM, Maxwell3D, and their linearized sessions.

Shot-Batched Modeling

Use tide.workflow.split_shots for custom mini-batch loops, or use run_shot_batches directly:

receiver = tide.workflow.run_shot_batches(
    operator,
    model,
    batch_size=8,
)

run_shot_batches rebuilds only the experiment geometry for each batch. The model, discretization, execution policy, storage policy, and named result contract remain unchanged.

index_shots and take_shot_batch expose the underlying indexing when a custom objective needs direct control. Use shot_dim=0 for shared-shot tensors shaped [S, ...] and shot_dim=1 for per-model-shot tensors shaped [B, S, ...].

Use expand_source_amplitude for the common wavelet-to-shot-amplitude step:

wavelet = tide.ricker(freq, nt, dt, device=device)
source_amplitude = tide.workflow.expand_source_amplitude(wavelet, n_shots)

Use line_acquisition_2d when a script only needs line coordinates:

acquisition = tide.workflow.line_acquisition_2d(
    source_x=torch.arange(n_shots, device=device) + 8,
    receiver_x=torch.arange(n_shots, device=device) + 12,
    source_depth=4,
    receiver_mode="paired",
)
source_location = acquisition.source_location
receiver_location = acquisition.receiver_location

Receiver Concatenation

merge_receiver_batches concatenates receiver chunks along the TIDE shot axis. It infers:

• [nt, S, R] -> shot axis 1
• [nt, B, S, R] -> shot axis 2

This keeps shared-model and batched-model outputs aligned with the solver API.

Structured Runner

run_shot_batches(operator, model, batch_size=...) accepts MaxwellTM or Maxwell3D and returns concatenated receiver data. It preserves autograd connections to the model tensors.

For loss-specific first- and second-order derivatives, use ReceiverObjective with a linearized operator:

objective = tide.workflow.ReceiverObjective(observed)
with operator.linearize(model) as linearized:
    gradient = objective.gradient(linearized)
    hvp = objective.hvp(linearized, direction, mode="gauss_newton")

With tide.optim

Optimizer state, gradients, bounds, and preconditioners remain as Torch tensors on the model device. A structured operator can be evaluated in shot batches while retaining autograd:

def objective(x: torch.Tensor) -> tuple[float, torch.Tensor]:
    packed = x.detach().clone().requires_grad_(True)
    epsilon, sigma = unpack_model(packed)
    model = tide.EMModel(epsilon, sigma, mu)
    pred = tide.workflow.run_shot_batches(
        operator,
        model,
        batch_size=batch_size,
    )
    loss = torch.mean((pred - observed) ** 2)
    loss.backward()
    if packed.grad is None:
        raise RuntimeError("objective did not produce a gradient")
    return float(loss.detach()), packed.grad.detach()

result = tide.optim.lbfgs_minimize(
    objective,
    x0,
    options=tide.optim.LBFGSOptions(
        stopping=tide.optim.StoppingCriteria(max_iter=10),
    ),
)

Preconditioners

Use curvature_preconditioner_diagonal for the common diagonal GN-style preconditioner pattern used in examples: accumulate a non-negative curvature proxy such as squared gradients, optionally smooth it, normalize it, invert it with damping, clip the scaling, and zero inactive cells.

curvature = torch.zeros_like(epsilon)

def record_curvature(_shot_indices: torch.Tensor, _loss: torch.Tensor) -> None:
    if epsilon.grad is not None:
        grad = torch.nan_to_num(epsilon.grad.detach(), nan=0.0, posinf=0.0, neginf=0.0)
        curvature.add_(grad.square())

tide.workflow.backward_shot_batches(
    objective_batch,
    shot_batches,
    zero_grad=clear_model_grads,
    zero_each_batch=True,
    after_backward=record_curvature,
)

diag = tide.workflow.curvature_preconditioner_diagonal(
    curvature,
    inactive_mask=air_mask,
    smooth_sigma=3.0,
    damping=5e-2,
    power=0.5,
    clip_min=0.3,
    clip_max=3.0,
    blend=0.7,
)
preconditioner = tide.workflow.diagonal_preconditioner(diag)

result = tide.optim.lbfgs_minimize(
    objective,
    x0,
    preconditioner=preconditioner,
    options=tide.optim.LBFGSOptions(
        stopping=tide.optim.StoppingCriteria(max_iter=10),
    ),
)

For two coupled parameter fields, accumulate the three symmetric block proxies and use curvature_preconditioner_block:

block = tide.workflow.curvature_preconditioner_block(
    curvature_ee,
    curvature_ss,
    curvature_es,
    inactive_mask=air_mask,
    smooth_sigma=3.0,
    damping=5e-2,
    power=0.5,
    clip_min=0.3,
    clip_max=3.0,
    offdiag_correlation_max=0.8,
    blend=0.7,
)
preconditioner = tide.workflow.block_preconditioner(block)

Receiver losses

receiver_mse_loss selects the observed shots corresponding to a predicted mini-batch and applies an explicit normalization policy. Use "batch" when each mini-batch objective should be a local mean. Use "all" when accumulated batches must reproduce a full-dataset mean.

loss = tide.workflow.receiver_mse_loss(
    predicted,
    observed,
    shot_indices,
    normalization="all",
)

receiver_sinkhorn_loss and receiver_gsot_loss compare traces as point clouds in time-amplitude space. Their dt, sparse sampling, transport power, and scale parameters affect the objective itself. Validate their gradient, runtime, and trace normalization on a controlled example before replacing MSE.

Shard-aware variants accept the global shot identifiers owned by one rank and the local observed-data shard. They keep data indexing and loss normalization consistent without gathering every receiver trace for each objective call.

Distributed shot sharding

init_distributed initializes one worker per selected device and returns a DistributedContext. rank_shot_indices and split_rank_shots assign global shot identifiers deterministically:

context = tide.workflow.init_distributed(enabled=True)
shot_batches = tide.workflow.split_rank_shots(
    n_shots,
    batch_size=4,
    context=context,
)

After local backward passes, call all_reduce_gradients to sum model gradients across ranks. Scale local objectives so the reduced gradient matches the intended global sum or mean. gather_receiver_shards is only needed when rank zero requires the complete receiver tensor for output or diagnostics.

tide.workflow.all_reduce_gradients(
    [epsilon, sigma],
    context=context,
)

destroy_distributed tears down a process group owned by the context. Put cleanup in finally when a worker can fail during propagation.

ReceiverObjective

ReceiverObjective separates a scalar receiver loss from Maxwell physics:

objective = tide.workflow.ReceiverObjective(
    observed_data,
    loss=custom_receiver_loss,
)

gradient(linearized) applies the loss cotangent through the operator VJP. hvp(linearized, direction, mode="gauss_newton") computes the Gauss-Newton action J.T @ loss'' @ J (positive semi-definite for convex receiver losses). mode="full" adds the nonlinear second-VJP term. Scalar hvp defaults to full.

The observed tensor must match the linearized primal receiver layout. The loss callable owns reduction and scaling, so changing a mean to a sum scales both gradient and Hessian products.

For multiple directions, use the public batched objective API:

directions = tide.EMDirection(epsilon=v_epsilon, sigma=v_sigma)
with operator.linearize(model, targets=("epsilon", "sigma")) as linearized:
    products = objective.hvp_batch(linearized, directions, block_size=4)

Each supplied direction has shape (K, *model.epsilon.shape), with the model’s dtype and device; K must be positive and identical for every supplied field. The returned EMGradient has the same leading direction axis for each selected linearization target, and None for non-target fields. Input directions need not be limited to output targets. block_size is a positive integer, defaults to four, and bounds native direction batches. Batched hvp_batch defaults to mode="gauss_newton"; mode="full" remains available through scalar composition.

Native CUDA TM Gauss-Newton execution with a shared 2-D model, device storage, and gradient sampling interval zero or one caches one background forward and reuses its snapshots across blocks and repeated calls. It does not retain second-order scattered snapshots. Each custom loss sees the original receiver layout, separately for each direction, so shot coupling, filtering inside the loss, and mean reductions retain their scalar meaning.

Other configurations use the structured scalar derivative path rather than discarding backend, fallback, storage, or experiment options. In particular, experiment-level frequency/time tapering, padding, and CFL time resampling use that path; filtering inside a receiver loss does not disable cached batching. Existing derivative restrictions still apply: TM permeability directions and dispersive TM tangent propagation are not implemented.

Close the linearization (or use its context manager) to release cached snapshots. Model, source, acquisition, and operator configuration must remain fixed during its lifetime; changes require a new linearization. A cached objective’s observed tensor must likewise remain unchanged. Switching to another observed tensor or loss releases the old objective cache. Loss callables should be pure, deterministic functions of their two tensor inputs.

Scope

The workflow module is intentionally narrow:

• no optimizer-state, model-packing, or constraint ownership
• no file I/O, plotting, logging, or device selection policy
• no replacement for the solver’s native batched-model support

It is meant to remove repeated shot-batching boilerplate from examples while keeping experiment-specific choices in user code.