neural_operators.mcmc#

class neural_operators.mcmc.mcmc.MCMC(model, prior, data, sigma_noise, pcn_beta=0.2, surrogate_to_use=None, surrogate_models=None, seed=0)[source]#

Bases: object

Preconditioned Crank-Nicolson (pCN) MCMC for Bayesian inversion of a coefficient field from noisy observations of the PDE state.

pCN’s proposal is reversible with respect to the prior, so the Metropolis acceptance ratio (see sample) depends only on the likelihood ratio, not the full posterior. model provides the true forward solve; pass surrogate_to_use/surrogate_models to run against a trained surrogate instead.

Note: run() resets pcn_beta/sigma_noise from its own arguments, overriding whatever __init__ was given – pass them to run() too if you set them there.

get_mcmc_params_for_tracer()[source]#
init_logger()[source]#
logLikelihood(current)[source]#
logPosterior(current)[source]#
logger(s)[source]#
plot(i)[source]#
process_and_print(i)[source]#
proposal(current, proposed)[source]#
run(init_m=None, n_samples=1000, n_burnin=100, pcn_beta=0.2, sigma_noise=0.01, savepath='./', savefilename='tracer', save_every=100, print_every=100, print_lvl=1, display_plot_every=100, init_tracer=False)[source]#

Run n_burnin + n_samples pCN steps.

Note: display_plot_every pops up real matplotlib windows on an interactive backend – set it past n_samples + n_burnin to disable.

sample(current)[source]#
save(i, current, accept)[source]#
solveFwd(current)[source]#
state_to_obs(u)[source]#
class neural_operators.mcmc.state.State(m_dim, u_dim, u_obs_dim)[source]#

Bases: object

A single MCMC state: m/u plus the quantities derived from them (observations, errors, log-likelihood/prior/posterior, cost).

set(a)[source]#
class neural_operators.mcmc.tracer.Tracer(mcmc)[source]#

Bases: object

Accumulates and periodically saves an MCMC run’s post-burn-in chain and running statistics. Note: accepted_samples_m etc. hold every post-burn-in sample, not just accepted ones – a rejected step just repeats the current value.

append(i, current, accept, mcmc_params=None, force_save=False)[source]#
init(current, init_sample)[source]#
save(mcmc_params)[source]#
stats()[source]#
class neural_operators.mcmc.surrogateModel.SurrogateModel(true_model, model, data)[source]#

Bases: object

Wraps a trained DeepONet/PCANet plus its DataProcessor so solveFwd(w) looks like PDEModel.solveFwd: latent field in, state out.

predict_from_m(m)[source]#

Encode/predict/decode given m directly, skipping the w->m transform. Used by callers that already parameterize m themselves (e.g. material-field/topology optimization).

solveFwd(w)[source]#
class neural_operators.mcmc.surrogateModel.SurrogateModelFNO(true_model, model, data, nodes, grid_x, grid_y, u_comps=1)[source]#

Bases: SurrogateModel

SurrogateModel for FNO: interpolates the vertex-ordered field onto the regular grid FNO trained on, predicts, interpolates back onto mesh vertices.

grid_to_nodes_u(u)[source]#
nodes_to_grid_m(m)[source]#
solveFwd(w)[source]#
class neural_operators.mcmc.correctedSurrogateModel.CorrectedSurrogateModel(surrogate, collect_diagnostics=False)[source]#

Bases: object

Wraps a SurrogateModel (or SurrogateModelFNO) so solveFwd(w) returns the residual-corrected state instead of the raw prediction. Requires the wrapped surrogate’s true_model to implement residual_correct(m, u_tilde) -> u_c.

solveFwd(w)[source]#
neural_operators.mcmc.compute_sample_errors.compute_sample_errors(mcmc)[source]#

Relative error (%) of the posterior mean and the last sample vs. ground truth, for w/m/u/u_obs. Returns a formatted string.

neural_operators.mcmc.mcmc_plot_fields.mcmc_plot_fields(mcmc, savefilename=None, params=None, use_surrogate_F_for_u=False)[source]#

mcmc_plot_fields_base, using the posterior mean and the last tracer sample.

neural_operators.mcmc.mcmc_plot_fields.mcmc_plot_fields_base(w_mean, w_sample, w_sample_i, mcmc, savefilename=None, params=None, surrogate_to_use=None, use_surrogate_F_for_u=False)[source]#

3x4 grid: ground truth / a given sample / posterior mean, for (w, m, u, u_obs), with relative errors against the ground truth row.

Post-process saved MCMC tracer results: curves and posterior field plots.

neural_operators.mcmc.mcmc_postprocess.plot_mcmc_tracer_results(savepath, mcmc, pp_params, field_suptitle, sample_i=-1, tag='Final', plot_surrogate_forward=True, save_figures=True, tracer_filename='tracer.pkl')[source]#

Load tracer from savepath and plot cost, acceptance rate, and field panels.

Parameters:

  • savepath (str) – MCMC results directory containing the tracer pickle.

  • mcmc (MCMC) – MCMC object used for the run (model, surrogate, ground truth data).

  • pp_params (dict) – Plot styling with keys curve_plot and field_plot.

  • field_suptitle (str) – Figure suptitle for the field plot panels.

  • sample_i (int) – Posterior sample index to plot. Default -1 uses the last sample.

  • tag (str) – Suffix in saved figure filenames (e.g. Final).

  • plot_surrogate_forward (bool) – If True and a surrogate is active, also save the surrogate-forward panel.

  • save_figures (bool) – If True, save PNGs to savepath; if False, display only.

  • tracer_filename (str) – Tracer pickle filename inside savepath.

Returns:

Loaded tracer object.

Return type:

tracer