This repo contains the code and data used to reproduce the results in the paper:
Latent Generative Modeling of Random Fields from Limited Training Data (2026).
J. E. Warner, T. A. Shah, P. E. Leser, G. F. Bomarito, J. D. Pribe, M. C. Stanley.
The ability to accurately model random fields plays a critical role in science and engineering for problems involving uncertain, spatially-varying quantities such as heterogeneous material properties and turbulent flows. Deep generative models offer a powerful tool for sampling high- or infinite-dimensional uncertainties like random fields, but their reliance on large, dense training datasets limits their applicability in contexts where sufficient data is difficult or expensive to obtain. In this work, we propose a latent-space approach to generative modeling of random fields that incorporates domain knowledge to supplement limited training data. A constraint-aware variational autoencoder (VAE) with a function decoder is first used to learn compact latent representations of continuous functions that adhere to known physical or statistical constraints, even when training data is sparse or indirect. Generative modeling is then performed in the learned latent space, decoupling constraint enforcement from the sampling process. This decoupling enables expressive multi-step generative methods to be deployed in data-limited settings where existing constrained multi-step approaches are not directly applicable. The richer latent distributions captured by the generative model also overcome limitations of standard VAEs, which rely on simple parametric priors and struggle to represent complex, multimodal, or heavy-tailed distributions over functions. Efficacy is demonstrated on two challenging applications: wind velocity field reconstruction from sparse sensors and material property inference from indirect measurements. Results show the effectiveness of incorporating domain knowledge constraints for data-limited problems and the improved sample quality and robustness of the latent generative modeling approach versus directly sampling a constrained VAE.
There are requirements.txt and environment.yml files that list the specific versions of packages used in this work. Otherwise, the base dependencies can be installed manually as follows:
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
conda install lightning -c conda-forge
conda install h5py
conda install scipy
pip install matplotlib
pip install seaborn
pip install pandas
pip install numpy
pip install torchdyn
pip install einops
Install local source:
pip install -e .
The scripts folder has subdirectories with train and test scripts to perform C-LFM for each of the four numerical examples from the paper:
gaussian_process: Demonstration 4.1.1: Random Field Reconstruction from Sparse Datalognormal_process: Demonstration 4.1.2: Random Field Reconstruction from Sparse Datapoisson_1d: Demonstration 4.2.1: Random Field Inference from Indirect Datapoisson_1d_multimodal: Demonstration 4.2.2: Random Field Inference from Indirect Datawind: Application 4.3: Wind Velocity Field Estimationmaterials: Application 4.4: Material Property Characterization
To process for performing end-to-end training and then testing for each problem is executed by running the following scripts in order:
train_vae.pytest_vae.py(optional)train_lfm.pytest_lfm.py
Step 1 pre-trains the VAE by implementing Algorithm 1 from the paper, step 3 trains a flow matching model on the latent space of the pre-trained VAE following Algorithm 2, and step 4 generates samples via latent flow matching following Algorithm 3. Step 2 is an optional step for testing the quality of the pre-trained VAE by encoding and decoding training data samples. The test scripts create some basic plots comparing generated samples to test data.
The demonstration problems, gaussian_process, lognormal_process, poisson_1d and poisson_1d_multimodal create train and test data on the fly. Train and test data is included in the repo for the materials application in the data/materials directory. Train and test data for the wind application was not included with the code due to their size, instead the data/wind_data_generation/generate_wind_data.py script needs to be run to create the datasets prior to training.
All scripts accept command line arguments to adjust hyperparameters and other details of the problem. The defaults are set to reasonable values that were used to produce the results in the paper. The exception being that the gaussian_process example must be run with --use_cov to engage the covariance constraint (it uses a correlation constraint by default) and the wind example must be run with --sparse_sensors to match the setup used in the paper. See the individual scripts for all command line arguments.
Results and model checkpoints are written to directories that are named based on the values of the command line arguments used. Each problem's subdirectory has a <problem_name>_utils.py file defining get_logging_path_and_name that specifies the results directory. A consequence of this structure is that the same command line arguments used to run train_vae.py in step 1 need to also be used in the subsequent steps in order to load stored models correctly. For example, to run the gaussian_process example with covariance constraint and three sensors for 5000 epochs:
python scripts/gaussian_process/train_vae.py --use_cov --epochs 5000 --num_sensors 3python scripts/gaussian_process/train_lfm.py --use_cov --epochs 5000 --num_sensors 3python scripts/gaussian_process/test_lfm.py --use_cov --epochs 5000 --num_sensors 3
The clfm package contains Python files implementing the core C-LFM functionality used by the scripts above, split across three subdirectories:
Contains the neural network classes that implement C-LFM. fully_connected_nets.py contain simple encoder, branch, trunk, and flow matching networks implemented with MLPs. unet1d.py contains the encoder used by the wind problem based on the U-Net architecture. vae.py contains the FunctionalVAE that implements the proposed VAE architecture with function decoder based on DeepONets that uses additional residual loss terms to incorporate physical and statistical constraints into the learning of latent function representations for random fields.
Contains problem specific Python files for each of the four examples from the paper. Each problem defines a custom PyTorch Dataset class as well as a Loss class that implements reconstruction and residual methods to define the problem's VAE loss terms. Additionally, there is a validate method which runs any problem specific validation metrics and returns them as a dictionary for logging. To customize the source code to solve a new example, one must define a new problem with appropriate dataset and loss terms.
Contains auxiliary C-LFM utilities used by the training and testing scripts, including latent_fm.py that implements training and sampling of latent flow matching and grid.py defining a spatial/temporal grid and implementing random sampling of collocation points.
If you have issues running the code or have any questions, email james.e.warner @ nasa.gov.