Variational Autoencoders for PDE Inverse Problems

This repository contains the implementation of a data-driven Variational Autoencoder (VAE) framework designed to solve ill-posed Bayesian Inverse Problems.

Specifically, the model tackles a 2D steady-state diffusion (heat conduction) problem. It estimates the unknown, spatially varying diffusion coefficient (the Parameter of Interest) from sparse point observations of the solution state. By learning to model the target posterior distribution of the PDE parameters, this approach goes beyond simple regression to provide robust uncertainty quantification.

This work reproduces and extends upon the foundational methodology presented by Goh et al. (2021) [1], leveraging a custom Finite Element Method (FEM) solver for dataset generation and TensorFlow for the deep learning architecture.

🚀 Key Features

• Probabilistic Modeling: Utilizes a VAE architecture optimized via a Jensen-Shannon Divergence loss to approximate the target posterior distribution.
• Uncertainty Quantification: Generates both mean estimates and 95% credibility intervals for the reconstructed PDE parameters.
• Custom Physics Pipeline: Includes a custom FEM solver built in MATLAB to generate synthetic ground-truth data (varying mesh refinements, observation points, and forcing terms).

📂 Project Structure

• ./fem_solver: Contains the MATLAB-based finite element solver used to generate the physical datasets (solutions to the diffusion equation).
• ./test: Contains the active experimental test suites, implemented in Jupyter Notebooks using the latest version of the TensorFlow VAE framework.
• ./old_versions: Contains deprecated codebase iterations and preliminary test definitions.

🛠️ Prerequisites

To run the data generation pipeline and train the models, you will need:

• MATLAB (for the FEM data generation)
• Python 3.x
• TensorFlow * Jupyter Notebook

💻 Usage & Reproduction

To reproduce the experimental results from scratch, follow this two-step workflow:

1. Generate the Physical Dataset (MATLAB)

• Navigate to the fem_solver/ directory.
• Run the sol_trainingset.m script in MATLAB. This will solve the forward diffusion problem and generate the physical states.
• Locate the generated uh.txt output file.

2. Train and Evaluate the VAE (Python)

• Move the generated uh.txt file into the ./test/data/ directory.
• Launch Jupyter Notebook and open the desired experimental notebook inside the ./test/ folder.
• Run the notebook cells sequentially to train the VAE and visualize the reconstructed parameters and uncertainty bounds.

📖 References

[1] H. Goh, S. Sheriffdeen, J. Wittmer, and T. Bui-Thanh. Solving Bayesian inverse problems via variational autoencoders. Proceedings of Machine Learning Research, 2021.