scheme_base

The class every iterative ensemble data-assimilation scheme inherits.

This is the PIPT counterpart to :mod:popt.optimization_methods.optimizer_base, and deliberately mirrors its shape: the scheme object owns its own iteration loop, its convergence checks, and its checkpoint/restart handling, while subclasses supply only the algorithm-specific analysis step.

The two packages differ in what the iteration acts on. An optimizer is handed callables (fun, jac, hess) and drives a control vector. An assimilation scheme is handed an ensemble collaborator, which owns the state realisations, the observed data, and the forward simulator.

Ensemble collaborator protocol

The scheme only relies on the following members, so anything satisfying them can be substituted (a lightweight fake is used in the unit tests):

ensemble.forecast() Run the forward simulator on the current state and refresh pred_data. ensemble.enX State ensemble matrix, shape (nx, ne). ensemble.pred_data Predicted data for the current state. ensemble.logger A :class:ensemble.logger.PetLogger, a no-op :class:ensemble.logger.NullLogger (set when the ensemble's logit option is false), or None (e.g. a test double with no logger at all). ensemble.keys_da The parsed dataassim config. Read at every hook, since which diagnostics and artifacts a run produces is a matter of configuration. ensemble.sim The forward simulator. Only input_dict is read here, to decide whether QA/QC was asked for. ensemble._saving_enabled Whether the run writes artifacts at all.

Reaching the ensemble's state

A scheme reads plenty of ensemble state -- enX, pred_data, keys_da, localization and friends -- and so do the analyses, through the scheme. Rather than forwarding unknown attributes at lookup time, each of those names is declared as an explicit :class:property on :class:AssimilationScheme (see the block of _ensemble_attr / _own_or_ensemble_attr declarations below). The scheme is therefore a façade: everything an analysis needs is reachable as scheme.<name>, whether the value lives on the scheme or on its ensemble, and an analysis never has to know which.

Reads delegate; writes do not. Assigning ensemble state goes through self.ensemble.<name> = ... explicitly, because that is the object the forecast reads back. The four names a scheme may legitimately compute for itself (cov_data, scale_data, proj, Am) are the exception and have setters.

Relationship to the legacy design

Historically a PIPT scheme inherited from pipt.loop.ensemble.Ensemble and an external pipt.loop.assimilation.Assimilate object drove the loop. That made every scheme simultaneously an algorithm and a data container, and made the analysis flavour (approx/full/subspace) part of the class name. Here the ensemble is a collaborator rather than a superclass, matching how OptimizerBase composes with its callables.

AssimilationResult

Bases: OptimizeResult

Result of an assimilation run.

A dict subclass with attribute access, mirroring :class:scipy.optimize.OptimizeResult so that PIPT and POPT results can be handled the same way. Typical fields:

nit Number of accepted iterations. success Whether the run stopped on a convergence criterion rather than by exhausting maxiter. message Human-readable reason the run stopped. why_stop Mapping of criterion name to whether it fired. data_misfit / prior_data_misfit Final and initial mean data misfit.

AssimilationScheme

Bases: AnalysisBindingMixin, RestartMixin, ABC

What every iterative ensemble data-assimilation scheme inherits.

Subclasses implement :meth:update_step, which performs one iteration and reports what it produced. Everything else is here: the loop, convergence bookkeeping, restart files, the run table, the result object, and the diagnostics and artifact saving that surround a run.

Those last two used to be a separate AssimilationWorkflowMixin that a combined class mixed in ahead of the loop. The split bought nothing -- every shipped scheme wanted both halves -- and cost the reader two classes and one load-bearing MRO order, in which listing the mixin second silently stopped a run from saving anything.

RESTART_ATTRIBUTES: tuple = ()

Attributes a scheme needs restored to resume mid-run: what its iterations change and what it drew at construction (perturbed observations, a damping parameter). The loop's own bookkeeping and the ensemble's state are covered by the base state; a subclass only names what it adds. Missing names are skipped, so a scheme that has not yet set one of them checkpoints fine.

__init__(ensemble, **options)

Parameters:

Name	Type	Description	Default
ensemble	object	Collaborator satisfying the ensemble protocol described in the module docstring. Owns the state, the observed data and the forward simulator.	required
**options		Scheme configuration. maxiter: Maximum number of accepted iterations (default: 100). misfit_tol: Relative data-misfit tolerance for convergence (default: 0.01). The assimilation counterpart of an optimizer's ftol. step_tol: Absolute tolerance on the norm of the state update (default: 1e-8). Counterpart of an optimizer's xtol. restart: Restore from a restart file on startup (default: False). restartsave: Write a restart file after the prior forecast and each accepted iteration (default: False). restart_file: Path for the restart file (default: '{scheme_name}_restart.pkl'). Config-driven schemes take these three from the [dataassim] block via :func:restart_options.	{}

after_accepted_iteration()

Persist iteration artifacts and run QA/QC after an accepted update.

after_analysis()

Between analysis and forecast.

The odd one out: it marks a point inside :meth:update_step, and this class does not dictate the shape of a step, so a scheme calls it itself. The rest of the hooks here are called by :meth:run_assimilation. Nothing runs here at present; it used to refresh QA/QC's variance after data screening, which is no longer supported.

after_forecast(state)

Between forecast and scoring: replace outlier members.

Ordering matters -- outliers are replaced before the misfit is scored, so the replacement feeds into the number the scheme sees. The resampled state is returned rather than written back, so the caller keeps ownership of what it is forecasting.

after_loop(converged)

Save the posterior and the reason the run stopped.

after_prior_forecast()

Handle the prior forecast: prior QA, saved artifacts.

Outlier replacement is not done here. The prior goes through :meth:after_forecast like every other forecast, so it has already happened by the time this runs -- and before :meth:record_prior_score computes the misfit, which is the order that matters.

assimilate(*args, **options)

Construct the scheme and run it to completion.

The assimilation counterpart of scipy.optimize.minimize: one call that builds the scheme, runs every iteration, and returns the outcome. Use it when the scheme object itself is not needed afterwards; when it is, construct the class and call :meth:run_assimilation instead.

Every argument is forwarded verbatim to the constructor, so this accepts whatever the scheme accepts rather than imposing a second signature.

Parameters:

Name	Type	Description	Default
*args		Positional arguments for the constructor. For the shipped PIPT schemes that is (keys_da, keys_en, sim) -- the parsed data-assimilation config, the parsed ensemble config, and the forward simulator -- from which the scheme builds its own ensemble. A scheme written directly against the collaborator protocol is handed its ensemble here instead.	()
**options		Keyword arguments for the constructor, such as analysis to override the flavour named in the config.	{}

Returns:

Type	Description
AssimilationResult	Outcome of the run. x is the posterior state ensemble, nit the number of accepted iterations, data_misfit and prior_data_misfit the final and initial mean misfits, and message the reason the run stopped.

Examples:

>>> keys_da, keys_sim, keys_en = read_config.read("case.toml")
>>> result = ESMDA.assimilate(keys_da, keys_en, flow(keys_sim))
>>> result.prior_data_misfit, result.data_misfit
(539.2, 70.1)

Overriding the flavour named in the config:

>>> result = ESMDA.assimilate(keys_da, keys_en, sim, analysis="subspace")

Notes

success reports whether the run stopped on a convergence criterion rather than by exhausting maxiter. Schemes with a fixed iteration schedule -- ES-MDA in particular -- therefore finish normally with success=False, which is expected rather than a failure.

See Also

run_assimilation : Run an already-constructed scheme.

build_ensemble(keys_da, keys_en, sim, ensemble=None)

The collaborator to run on: ensemble if given, else a fresh ENSEMBLE_CLASS.

Handing one in lets two schemes share a prior and its forecasts, and lets a test substitute a stand-in without the config, data files and simulator a real ensemble needs.

check_convergence()

Check scheme-specific convergence criteria.

Returns:

Type	Description
bool	True if a subclass-specific stopping criterion is satisfied. The default implementation never stops the loop.

check_misfit_convergence()

Check convergence on the relative change in mean data misfit.

check_state_convergence()

Check convergence on the norm of the state update.

The counterpart of :meth:popt.optimization_methods.optimizer_base. OptimizerBase.check_state_convergence, which compares xk against xk_old. enX_old is snapshotted by :meth:run_assimilation before each attempt, but only when step_tol > 0 -- see there for why.

Opt-in in practice: every shipped scheme passes step_tol=0.0, because ‖Δx‖₂ over a state that mixes variables on different scales (log-permeability alongside saturations, say) has no tolerance that is meaningful across cases. The base default of 1e-8 is small enough to mean "the state did not move at all" rather than being a guess at a scale.

log_columns(prior_run=False)

Trailing columns for the run table -- typically the scheme's control parameter, e.g. {"λ": self.lam}. Empty by default.

log_update(success=None, prior_run=False)

Log one row of the run table.

Called by :meth:run_assimilation -- once for the prior and once per accepted iteration -- so a scheme gets its rows without asking, and the attempts it makes inside :meth:update_step stay its own business. The row is the same for every scheme apart from its control parameter, which :meth:log_columns supplies.

propose_state(result, step_scale=1.0)

The trial state an analysis result implies.

Parameters:

Name	Type	Description	Default
result	AnalysisResult or array - like	What self.update(...) returned. A plain array is a state-space step.	required
step_scale	float	Step length applied to the step (GN-EnRML's gamma); 1 for schemes without one.	1.0

Returns:

Type	Description
ndarray	The state to forecast. Weight-space results also advance self.W from self.current_W; the scheme commits W to current_W when it accepts the step.

record_prior_score()

Score the prior forecast and record it, before any iteration.

Sets prior_data_misfit_mean, data_misfit_mean and the per-realisation ensemble_misfit, so the prior is described by the same attributes as every later iteration -- and early enough that the iteration-0 artifacts written by :meth:after_prior_forecast can capture them.

This used to be a score_prior() hook that each scheme implemented, which meant every scheme spelled out both the misfit expression and the five assignments around it. The expression is now :meth:score and the bookkeeping is here; a scheme customises the former.

Does nothing when :meth:score reports no misfit, which is how a scheme with nothing to score opts out.

run_assimilation()

Run this scheme's assimilation to completion.

Named for the job rather than the mechanism, and matching the run_forecast already on this class. The counterpart in popt is OptimizerBase.run_optimization.

Restores a checkpoint if configured, forecasts and scores the prior, then calls :meth:update_step until a convergence criterion fires or maxiter iterations have been taken. One call is one iteration: a scheme that retries -- re-damping, backtracking a step length -- does so inside :meth:update_step, so a report coming back rejected means it has run out of attempts, and the run stops rather than asking again for a step it just said it could not find.

Convergence is checked on rejected reports too, before that stop takes effect: a scheme's :meth:check_convergence can legitimately fire on a step it is about to reject (a stalled misfit that did not actually improve), and that verdict decides how the run is reported.

Returns:

Type	Description
AssimilationResult	Populated result object, also stored on self.results.

run_forecast(state)

Forecast state, then run the post-forecast step.

Returns the state to carry forward -- the same one unless :meth:after_forecast replaced members in it.

score(pred_data=None)

Per-realisation data misfit of a forecast.

Called every time a new state has been forecast and needs a number: once for the prior, by :meth:record_prior_score, and then by each scheme for every attempt it takes inside :meth:update_step. One definition per scheme, rather than the same expression repeated in a prior-scoring hook and again in the step.

Parameters:

Name	Type	Description	Default
pred_data	optional	The forecast to score -- a PETDataFrame or an (nd, ne) matrix. Defaults to self.pred_data, which is what the ensemble's most recent forecast produced, so the usual call is self.score() straight after run_forecast. Pass one explicitly to score a forecast the ensemble no longer holds.	None

Returns:

Type	Description
ndarray or None	(ne,) misfit per realisation, or None when the scheme has no observation ensemble bound -- a scheme that scores some other way overrides this, and one that reports no misfit at all (the base's own tests) leaves the loop's misfit bookkeeping alone.

Notes

The default is the objective function every shipped scheme uses,

.. math::

\Phi_j = (g(m_j) - d_j)^{\mathsf T} C_d^{-1} (g(m_j) - d_j),

against the perturbed observations enObs and the data covariance cov_data. ES-MDA overrides it to score against an un-inflated copy of the perturbations (enObs_conv); the multilevel scheme to score all fidelity levels at once.

update_step()

Perform one scheme-specific analysis step.

Implementations compute the analysis update, apply it to the ensemble state, run the resulting forecast, and score the result. How they do that is entirely theirs -- the base calls this and nothing inside it.

Returns:

Type	Description
StepReport	accepted decides whether the loop advances or gives the scheme another attempt at the same iteration number, which is how the Levenberg-Marquardt schemes back off by increasing their damping parameter. misfit is the per-realisation data misfit as of now; the loop derives data_misfit and data_misfit_std from it.

StepReport

What one attempt produced. Returned by :meth:update_step.

The base does not dictate how a scheme takes its step; this is what it needs back afterwards, to score convergence, log, and build the result. Required fields are positional, so forgetting one is a TypeError at construction rather than a None surfacing several iterations later.

accepted: bool

Keep this step? False says the scheme found no improving step and has exhausted the attempts it makes inside :meth:update_step, so the loop stops rather than asking for the same step again.

misfit: np.ndarray

Per-realisation data misfit as of now. The loop derives data_misfit and data_misfit_std from it, so the three can no longer drift apart the way separately-assigned attributes could.

"As of now" matters for a scheme that gives up: LM-EnRML restores the last accepted misfit when it backs off, and returns that, so the value the loop records and logs is the one the run actually reached.

state: Any

The state this attempt produced, committed by the loop when accepted. A scheme still writes it to ensemble.enX_temp first, because that is what the forecast predicts on -- but handing it back here is what lets the loop own the commit, rather than every scheme remembering the same two lines. Forgetting them used to give a run that iterated and logged normally while returning the prior untouched.

why_stop: dict | None = None

Criterion record, merged into result.why_stop.

restart_options(keys_da)

The checkpoint settings of a config's [dataassim] block, as scheme options.

restart (resume from the checkpoint), restartsave (write one after the prior forecast and every accepted iteration) and restart_file (default <scheme>_restart.pkl). Legacy yes/no strings are accepted. Schemes pass **restart_options(keys_da) to the base so the keys reach :class:~ensemble.checkpoint.RestartMixin; they used to stop at the ensemble, which loaded a pickle of itself and left the scheme's own state -- iteration, damping, misfit history -- at its initial values.