Merge pull request #168 from ajnonaka/ajn_pytuq

ajnonaka · web-flow · commit 2b9b2eed45c8 · 2026-01-23T13:59:19.000-08:00
AMReX+PyTUQ
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
@@ -19,7 +19,7 @@ jobs:
           .github/workflows/dependencies/documentation.sh
           echo "Installing python packages for docs..."
           python3 -m pip install --upgrade pip
-          python3 -m pip install sphinx sphinx_rtd_theme breathe sphinxcontrib.bibtex docutils
+          python3 -m pip install sphinx sphinx_rtd_theme breathe sphinxcontrib.bibtex docutils sphinx-copybutton sphinx-design
 
       - name: Install and Build
         run: |
diff --git a/Docs/Readme.sphinx b/Docs/Readme.sphinx
@@ -23,6 +23,10 @@ If you have conda, you can install the necessary packages as follows:
 
        2. Type "conda install sphinx_rtd_theme"
 
+       3. Type "conda install sphinx-design"
+
+       4. Type "conda install -c conda-forge sphinx-copybutton"
+
 If you don't have a conda Python installation, you get one by doing the following:
 
        1.  Go to https://conda.io/miniconda.html and download the Python 3.6 64-bit (bash installer), "Miniconda3-latest-Linux-x86_64.sh".  Save this script to your hard drive.
diff --git a/Docs/source/HeatEquation_UQ.rst b/Docs/source/HeatEquation_UQ.rst
@@ -0,0 +1,209 @@
+.. _tutorials_uq:
+
+.. _guided_pytuq_integration:
+
+.. _pytuq_quickstart:
+
+UQ with PyTUQ
+===========
+
+.. admonition:: **Time to Complete**: 30-60 minutes
+   :class: note
+
+   **What you will learn**:
+      - Install PyTUQ
+      - Run an AMReX+PyTUQ to perform sensitivity analysis
+
+Overview
+--------
+
+AMReX simulations deliver high-fidelity, accurate results for complex physics problems, but the effect on simulation results due to uncertain inputs can require hundreds or thousands of simulations, making comprehensive analysis computationally prohibitive.
+
+This tutorial demonstrates how to improve efficiency without sacrificing accuracy by using polynomial chaos expansions to build fast surrogate models that identify which input parameters truly matter.
+
+PyTUQ (Python interface to the UQ Toolkit) provides specialized tools for surrogate construction and global sensitivity analysis, enabling rapid parameter space exploration and dimensionality reduction for scientific applications.
+
+We demonstrate how to integrate PyTUQ with your AMReX application through a practical workflow; the AMReX-based heat equation tutorial is equipped to perform sensitivity analysis.
+
+In these examples we model the heat equation
+
+.. math:: \frac{\partial\phi}{\partial t} = D\nabla^2 \phi,
+
+with initial condition
+
+.. math:: \phi_0 = 1 + A e^{-r^2 / (2V)},
+
+where ``r`` is the distance from the center of the domain, and with uncertain parameters ``diffusion_coeff`` (:math:`D`), ``init_amplitude`` (:math:`A`), and ``init_variance`` (:math:`V`).
+The outputs of interest are the maximum temperature, mean temperature, standard deviation of temperature, and the temperature at a specified point.
+
+.. toctree::
+   :maxdepth: 1
+
+   HeatEquation_UQ_MathematicalDetails
+
+Located in ``amrex-tutorials/ExampleCodes/UQ/HeatEquation``, this example illustrates a complete forward UQ workflow from parameter sampling randomized input parameters to perform sensitivity analysis.
+By understanding this example, you will have a basis for understanding how to adapt this workflow to your own AMReX application.
+
+More specifically, you can directly compare/diff ``amrex-tutorials/ExampleCodes/UQ/HeatEquation/main.cpp`` against the original heat equation tutorial ``amrex-tutorials/GuidedTutorials/HeatEquation_Simple/main.cpp`` to see exactly what source code changes are made to the AMReX application in this example.
+
+Installation
+------------
+
+We now describe the installation and workflow process on a local workstation.
+First, install pytuq using this script (based on information provided in `pytuq/README.md <https://github.com/sandialabs/pytuq/blob/main/README.md>`_):
+
+.. code-block:: bash
+   :caption: Pytuq installation script
+
+   #!/bin/bash
+
+   # 1. Clone repositories
+   git clone https://github.com/sandialabs/pytuq.git
+
+   # 2. Create a conda environment with python (optional, you can add to an existing env)
+   conda create --name pytuq
+   conda activate pytuq
+   conda install python=3.11
+
+   # 3. Install PyTUQ and requirements
+   cd pytuq
+   pip install -r requirements.txt
+   pip install .
+   conda install dill
+
+   # 4. Verify installation
+   conda list | grep pytuq    # Should show pytuq 1.0.0
+
+Examples
+--------
+
+C++ AMReX + PyTUQ (BASH driven)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+   **Prerequisites**:
+
+   - AMReX and pytuq cloned at the same directory level as amrex-tutorials
+   - pytuq installation described above
+   - GNU Parallel for task management: ``sudo apt-get install parallel`` (mpiexec option also exists in script)
+
+   .. code-block:: bash
+      :caption: Build AMReX appplication
+
+      cd /path/to/amrex-tutorials/ExampleCodes/UQ/HeatEquation/
+      make -j4
+
+   .. code-block:: bash
+      :caption: Run sensitivity analysis with bash script
+
+      ./workflow_uqpc.x
+
+Understanding GNU Parallel Workflow Pattern
+-------------------------------------------
+
+   The ``workflow_uqpc.x`` bash script relies on the user augmenting their codes to write outputs of interest to ASCII text files.
+   In this case, the ``main.cpp`` was modified from the ``amrex-tutorials/GuidedTutorials/HeatEquation_Simple/main.cpp`` in the following ways:
+
+   First, support for parsing ``diffusion_coeff``, ``init_amplitude``, and ``init_variance`` from the input file and command line were added.
+
+   Second, support for writing outputs of interest to ASCII text files is added.  The ``datalog_filename`` input parameter is generated by the bash
+   script to give each simulation output a unique identifier.
+   The ``datalog_int`` parameter gives the user the option to write the outputs of interest at a given time step interval, but in this example,
+   the outputs of interest after the final step are those that matter, and are extracted by the bash script to create a master output file
+   containing a separate set of simulation outputs of interest in each row.
+
+   The bash script calls PyTUQ scripts that generate an input parameter files for training and testing (``ptrain.txt`` and ``ptest.txt``)
+   based on polynomial chaos settings, and then uses GNU Parallel to run multiple simulations
+   efficiently and collect outputs into results files (``ytrain.txt`` and ``ytest.txt``) that PyTUQ can use for surrogate model fitting.
+
+Understanding the Output
+------------------------
+
+| `pnames.txt`, `outnames.txt`
+| -names of input and output parameters
+
+| `param_margpc.txt`
+| -Each row contains the mean and standard deviation for an uncertain input parameter
+
+| `qtrain.txt`, `qtest.txt`
+| -each row is a separate set of normal random variables used to generate uncertain inputs
+
+| `ptrain.txt`, `ptest.txt`, `pall.txt`
+| -each row is a separate set of input parameters for each simulation
+
+| `stdoutlog_train##.txt`, `stdoutlog_test##.txt`
+| -All standard output from AMReX simulations for training and testing data.  The numbers refer to separate simulations.
+
+| `datalog_train##.txt`, `datalog_test##.txt`
+| -Specifically-chosen output from AMReX simulations for training and testing data.  In this example it reports the outputs (max, mean, standard deviation, and specific cell temperature) at user-specified intervals.
+
+| `ytrain.txt`, `ytest.txt`, `yall.txt`
+| -agglomeration of outputs of interest from all simulations
+
+| `results.pk`
+| -Python pickle file encapsulating the results
+
+| `labels.txt`
+| -list of labels of simulation types (training or testing) for diagnostic/plot generation purposes
+
+| `xx_<INVAR1>_<INVAR2>.png`
+| -scatter plot of 2 inputs for training and testing
+
+| `pcoord_1.png`, `pcoord_1_lab_Testing.png`, `pcoord_1_lab_Training.png`
+| -graphical representation of how inputs correlate to outputs for each individual simulation
+
+| `yx_<OUTVAR>.png`, `yx_<OUTVAR>_log.png`
+| -scatter plots of output as a function of each input
+
+| `yxx_#.png`
+| -scatter plots of output a function of multiple inputs
+
+| `pdf_tri_inputs.png`, `pdf_tri_output.png`
+| -PDFs of inputs and outputs
+
+| `ensemble.png`
+| -graphical display of all output values
+
+| `idm_#_training.png`, `idm_#_testing.png`
+| -graphical display of output values
+
+| `dm_#.png`
+| -”data vs model” parity plots for each output; compares the predicted values from a surrogate model or approximation against the true or actual values from the full computational model
+
+| `fit_s#_training.png`, `fit_s#_testing.png`
+| -Shows model vs PC approximation for a single simulation.
+
+| `pdf_output_#.png`, `pdf_joyplot.png`
+| -PDF of output variables
+
+| `allsens_main.txt`, `sens_main.png`
+| -raw data and plot for parameter sensitivities
+
+| `jsens_#.png`, `Jsens_ave.png`
+| -joint sensitivities of output with respect to inputs
+
+| `sensmat_main.png`
+| -sensitivity matrix
+
+| `pcslices_o#.png`
+| -polynomial chaos fits
+
+| `pccont_o#_d#_d#.png`
+| -polynomial chaos fits of output variables with respect to two input variables
+
+
+Additional Resources
+--------------------
+
+**PyTUQ Resources:**
+
+- `PyTUQ Documentation <https://sandialabs.github.io/pytuq>`_
+- `PyTUQ Examples directory <https://github.com/sandialabs/pytuq/tree/main/examples>`_
+- eebaill, ksargsyan, & Bert Debusschere. (2025). sandialabs/pytuq: v1.0.0z (v1.0.0z). Zenodo. https://doi.org/10.5281/zenodo.17110054
+
+**AMReX Resources:**
+
+- `AMReX Documentation <https://amrex-codes.github.io/amrex/docs_html/>`_
+
+**Uncertainty Quantification Theory:**
+
+- Ghanem, Roger, David Higdon, and Houman Owhadi, eds. *Handbook of Uncertainty Quantification*. Vol. 6. New York: Springer, 2017. (For workflow, plotting, and analysis specifics)
diff --git a/Docs/source/HeatEquation_UQ_MathematicalDetails.rst b/Docs/source/HeatEquation_UQ_MathematicalDetails.rst
@@ -0,0 +1,48 @@
+.. _pytuq_mathematical_details:
+
+Parameter Sensitivity via Polynomial Chaos
+-------------------------------------------
+
+In this example, PyTUQ constructs polynomial chaos expansions to analyze how uncertain physical parameters affect simulation outputs in a heat diffusion problem.
+
+The Heat Equation
+^^^^^^^^^^^^^^^^^
+
+The governing equation for this tutorial is the heat diffusion equation:
+
+.. math:: \frac{\partial\phi}{\partial t} = D\nabla^2 \phi,
+
+with initial condition
+
+.. math:: \phi_0 = 1 + A e^{-r^2 / (2V)},
+
+where ``r`` is the distance from the center of the domain, and with uncertain parameters ``diffusion_coeff`` (:math:`D`), ``init_amplitude`` (:math:`A`), and ``init_variance`` (:math:`V`).
+
+
+
+Quantities of Interest (Outputs)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The simulation extracts four statistical quantities at the final timestep:
+
+1. **max_temp**: Maximum temperature in the domain
+2. **mean_temp**: Average temperature across all cells
+3. **std_temp**: Standard deviation of temperature
+4. **cell_temp**: Temperature in a particular cell, in this case, cell (9,9,9)
+
+These outputs are computed in ``main.cpp``and written to the datalog file.
+
+PyTUQ Workflow
+^^^^^^^^^^^^^^
+
+PyTUQ uses polynomial chaos expansion to construct a surrogate model:
+
+1. **Parameter sampling**: Generate sample points in the 3D parameter space based on the specified distributions
+2. **Forward simulations**: Run the AMReX heat equation solver for each parameter set
+3. **Surrogate construction**: Fit polynomial chaos coefficients to map inputs → outputs
+4. **Sensitivity analysis**: Compute Sobol indices to identify which parameters most influence each output
+
+The connection is:
+
+- **Inputs**: ParmParse parameters (``diffusion_coeff``, ``init_amplitude``, ``init_variance``) specified in ``inputs`` file or command line
+- **Outputs**: Quantities of interest extracted from datalog files
diff --git a/Docs/source/conf.py b/Docs/source/conf.py
@@ -39,6 +39,8 @@ def get_amrex_version():
 extensions = ['sphinx.ext.mathjax',
               'sphinx.ext.githubpages',
               'sphinx.ext.viewcode',
+              'sphinx_design',
+              'sphinx_copybutton',
               'sphinx.ext.intersphinx',
               'sphinx_rtd_theme']
 
@@ -51,6 +53,9 @@ def get_amrex_version():
 # Add any paths that contain templates here, relative to this directory.
 templates_path = ['ytemplates']
 
+# sphinx-copybutton configuration
+copybutton_exclude = '.linenos, .gp, .go'
+
 # The suffix(es) of source filenames.
 # You can specify multiple suffix as a list of string:
 #
diff --git a/Docs/source/index.rst b/Docs/source/index.rst
@@ -50,6 +50,7 @@ sorted by the following categories:
 - :ref:`GPU<tutorials_gpu>`  -- Offload work to the GPUs using AMReX tools.
 - :ref:`Linear Solvers<tutorials_linearsolvers>`  -- Examples of several linear solvers.
 - :ref:`ML/PYTORCH<tutorials_ml>`  -- Use of pytorch models to replace point-wise computational kernels.
+- :ref:`UQ with pytuq<tutorials_uq>`  -- Use of pytuq methods to perform generalized sensitivity analysis.
 - :ref:`MPMD<tutorials_mpmd>` -- Usage of AMReX-MPMD (Multiple Program Multiple Data) framework.
 - :ref:`MUI<tutorials_mui>`  -- Incorporates the MxUI/MUI (Multiscale Universal interface) frame into AMReX.
 - :ref:`Particles<tutorials_particles>`  -- Basic usage of AMReX's particle data structures.
@@ -72,6 +73,7 @@ sorted by the following categories:
    GPU_Tutorial
    LinearSolvers_Tutorial
    ML_Tutorial
+   HeatEquation_UQ
    MPMD_Tutorials
    MUI_Tutorial
    Particles_Tutorial
@@ -97,6 +99,8 @@ sorted by the following categories:
 
 .. _`Linear Solvers`:  LinearSolvers_Tutorial.html
 
+.. _`UQ`:  UQ_Tutorial.html
+
 .. _`MPMD`:  MPMD_Tutorials.html
 
 .. _`MUI`: MUI_Tutorial.html
diff --git a/ExampleCodes/UQ/HeatEquation/GNUmakefile b/ExampleCodes/UQ/HeatEquation/GNUmakefile
@@ -0,0 +1,16 @@
+# AMREX_HOME defines the directory in which we will find all the AMReX code.
+AMREX_HOME ?= ../../../../amrex
+
+DEBUG        = FALSE
+USE_MPI      = TRUE
+USE_OMP      = FALSE
+COMP         = gnu
+DIM          = 3
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.defs
+
+include ./Make.package
+
+include $(AMREX_HOME)/Src/Base/Make.package
+
+include $(AMREX_HOME)/Tools/GNUMake/Make.rules
diff --git a/ExampleCodes/UQ/HeatEquation/Make.package b/ExampleCodes/UQ/HeatEquation/Make.package
@@ -0,0 +1,2 @@
+CEXE_sources += main.cpp
+
diff --git a/ExampleCodes/UQ/HeatEquation/inputs b/ExampleCodes/UQ/HeatEquation/inputs
@@ -0,0 +1,16 @@
+# Grid/Domain parameters
+n_cell = 32                 # number of cells on each side of domain
+max_grid_size = 16          # size of each box (or grid)
+
+# Time stepping parameters  
+nsteps = 100                # total steps in simulation
+dt = 5.e-8                  # time step
+
+# Output control
+plot_int = -1                # how often to write plotfile (-1 = no plots)
+datalog_int = -1             # how often to write datalog (-1 = no regular output)
+   
+# Physics parameters (these are what we vary for UQ and values here will be overwritten
+diffusion_coeff = 100.      # diffusion coefficient for heat equation
+init_amplitude = 1.         # amplitude of initial temperature profile
+init_variance = 0.01        # variance of initial temperature profile
diff --git a/ExampleCodes/UQ/HeatEquation/main.cpp b/ExampleCodes/UQ/HeatEquation/main.cpp
diff --git a/ExampleCodes/UQ/HeatEquation/workflow_uqpc.x b/ExampleCodes/UQ/HeatEquation/workflow_uqpc.x
diff --git a/GuidedTutorials/HeatEquation_Simple/main.cpp b/GuidedTutorials/HeatEquation_Simple/main.cpp
