Recent progresses in generative models for text, images and video have shown that neural networks can learn complex data structures. In this example notebook, we train the generative model to learn the probability distribution of a large ensemble of the shear velocity generated by a Monte Carlo study, use the trained model on Gadi to reproduce the ensemble and evaluate the inference result by comparing its statistics with that of the input dataset.

The example is mainly a replication of the work reported in section 4.2 in the paper "Upscaling and downscaling Monte Carlo ensembles with generative models" https://doi.org/10.1093/gji/ggac100 [1] using the ensemble generated by Dr. Sima Mousavi in this paper "Lowermost mantle shear-velocity structure from hierarchical trans-dimensional Bayesian tomography" https://doi.org/10.1029/2020JB021557 [2]

Similar method has great potential in surrogates of complex data generated by big physics facilities, such as radiograph images, temporally and spatially-resolved spectroscopy, that are currently converted to scalar or low-dimensional vector summaries to suit the storage limitation.

Run the training script

Copy all the relevant scripts to your local directory and then submit the job to train the model.

$ module use /g/data/dk92/apps/Modules/modulefiles/
$ module load NCI-ai-ml/23.05
$ wdir=/scratch/$PROJECT/$USER
$ cp /g/data/dk92/notebooks/examples-aiml/gan_mcmc/job.sh $wdir
$ cp /g/data/dk92/notebooks/examples-aiml/gan_mcmc/neural_networks.py $wdir
$ cp /g/data/dk92/notebooks/examples-aiml/gan_mcmc/generative_model.py $wdir
$ cp /g/data/dk92/notebooks/examples-aiml/gan_mcmc/train_ndcgan.py $wdir
$ cd $wdir
$ qsub -P $PROJECT job.sh

Run the tested notebook

  1. Go to ARE site: are.nci.org.au
  2. Fill out the JupyterLab request form
    1. Walltime (hours): 2
    2. Queue: normalsr
    3. Compute Size: small
    4. Project:  <xy01>
    5. Storage: gdata/dk92+scratch/<xy01>
  3. Click Advanced options and fill in the following fields
    1. Module directories: /g/data/dk92/apps/Modules/modulefiles/
    2. Modules: NCI-ai-ml/23.05
    3. Jobfs size: 10GB
  4. Launch the session to run the tested notebook

 

Note on the tested notebook

  • copy the tested notebook from "/g/data/dk92/notebooks/examples-aiml/gan_mcmc/gan_mcmc.ipynb" to your own working directory
  • the notebook gan_mcmc.ipynb contains an example to train a DCGAN model with the ensemble generated by a Monte Carlo study, reproduce the ensemble and evaluate the inference result.
  • if doing inference only, in order to load the checkpoint file, copy two more python scripts, namely, "/g/data/dk92/notebooks/examples-aiml/gan_mcmc/neural_networks.py" and  "/g/data/dk92/notebooks/examples-aiml/gan_mcmc/generative_model" to your working directory.
  • if necessary, modify the checkpoint file in the notebook to the one you trained by yourself. The current one used in the notebook is trained on Gadi for 10 epochs using the reduced-sized ensemble (50%).

Statistics Comparison

The trained GAN reproduces the ensemble in great details, not only the distribution at individual locations, see EMD map and the distribution comparison, but also the spatial relations across different locations, see the covariance matrix across all the 529 sites.


Citations

  1. Scheiter, M., Valentine, A., Sambridge, M., 2022. Upscaling and downscaling Monte Carlo ensembles with generative models, Geophys. J. Int., 230(2):916--931.

  2. Mousavi, S., Tkalcic, H., Hawkins, R., & Sambridge, M., 2021. Lowermost mantle shear-velocity structure from hierarchical trans-dimensional Bayesian tomography, J. Geophys. Res. Solid Earth, 126, e2020JB021557

  • No labels