Add workflow to train SR-GNN model from ensemble of nekRS simulations

3rd_party/ensembleLauncher/README.md

@@ -16,6 +16,9 @@ What this directory provides:
   - **symlinks** to the large/shared paths declared by the caller (typically
     `.re2` and `.cache`, both of which are identical across members when the
     sweep is over runtime parameters rather than over the mesh).
+  - optional per-member **extra symlinks** via each member dict's ``symlinks``
+    key: ``{dest_basename: src_relpath_from_base}`` (e.g. symlink one checkpoint
+    ``.f`` file into each member directory for ``startFrom`` restarts).
 
   The split between copies and symlinks matters: `.cache` in particular is
   large and members must not race to rebuild it.

3rd_party/ensembleLauncher/nekrs_ensemble_utils.py

@@ -63,6 +63,9 @@ def setup_ensemble_dirs(
         ``"par_overrides"``: a mapping ``{section: {key: value, ...}, ...}``
         that is applied to the base ``.par`` (see ``apply_par_overrides``
         for matching/formatting rules).
+        Optional ``"symlinks"``: ``{dest_basename: src_relpath_from_base, ...}``
+        creates extra symlinks in that member directory (e.g. symlink each
+        checkpoint ``.f`` file next to the per-member ``.par`` for ``startFrom``).
     base_dir : str
         Directory containing the source case files (default: current working
         directory).
@@ -137,6 +140,29 @@ def setup_ensemble_dirs(
                 shutil.rmtree(dst)
             os.symlink(src, dst)
 
+        # Optional per-member symlinks: dest filename in member dir -> path
+        # relative to ``base_dir`` (same as ``copy_files`` / ``symlink_files``).
+        extra_links = member.get("symlinks")
+        if extra_links:
+            if not isinstance(extra_links, dict):
+                raise TypeError(
+                    "member['symlinks'] must be a dict "
+                    "{dest_basename: src_relpath_from_base}"
+                )
+            for dst_name, src_rel in extra_links.items():
+                src = (base / str(src_rel)).resolve()
+                if not src.is_file():
+                    raise FileNotFoundError(
+                        f"member['symlinks'] source not found: {src} "
+                        f"(member={member['name']!r})"
+                    )
+                dst = d / Path(dst_name).name
+                if dst.is_symlink() or dst.is_file():
+                    dst.unlink()
+                elif dst.is_dir():
+                    shutil.rmtree(dst)
+                os.symlink(src, dst)
+
         dirs.append(d)
 
     return dirs
@@ -202,6 +228,8 @@ def apply_par_overrides(
 def _fmt_par_value(v) -> str:
     if isinstance(v, bool):
         return "true" if v else "false"
+    if isinstance(v, int):
+        return str(v)
     if isinstance(v, float):
         return f"{v:.10g}"
     return str(v)

README.md

@@ -22,21 +22,22 @@ Some key functionalities of nekRS-ML are:
 * [In-memory data staging with SmartSim](./src/plugins/smartRedis.hpp): nekRS-ML can also be linked to the [SmartRedis](https://github.com/CrayLabs/SmartRedis) library, which when coupled with a [SmartSim](https://github.com/CrayLabs/SmartSim) workflow enables online training and inference with in-memory data-staging.
 * [Efficient deployment of nekRS ensembles](./examples/periodicHill_ensemble/): nekRS-ML provides utilities to setup and launch nekRS ensembles with [EnsembleLauncher](https://github.com/argonne-lcf/ensemble_launcher) (EL), which is a light-weight, scalable task launcher developed at the ALCF. This tool is useful for deploying parameter sweeps, scaling studies or gathering training data from various simulations by launching large ensembles on HPC systems.
 
-### Progression of AI-enabled examples
+### Progression of AI/ML-enabled examples
 
-nekRS-ML hosts a series of AI-enabled examples listed below in order of complexity to provide a smooth learning progression. 
+nekRS-ML hosts a series of AI/ML-enabled examples listed below in order of complexity to provide a smooth learning progression. 
 Users can find more details on each of the examples in the  README files contained within the respective directories. 
 
 * [tgv_gnn_offline](./examples/tgv_gnn_offline/): Offline training pipeline to generate data and perform time independent training of the Dist-GNN model.
 * [tgv_gnn_offline_coarse_mesh](./examples/tgv_gnn_offline_coarse_mesh/): Offline training pipeline to generate data and perform time independent training of the Dist-GNN model on a p-coarsened grid relative to the one used by the nekRS simulation.
 * [tgv_gnn_offline_traj](./examples/tgv_gnn_offline_traj/): Offline training pipeline to generate data and perform time dependent training of the Dist-GNN model.
-* [tuurbChannel_srgnn](./examples/turbChannel_srgnn/): Offline training pipeline to generate data, perform training, and evaluate the model through inference with the SR-GNN model. 
+* [turbChannel_srgnn](./examples/turbChannel_srgnn/): Offline training pipeline to generate data, perform training, and evaluate the model through inference with the SR-GNN model. 
 * [turbChannel_wallModel_ML](./examples/turbChannel_wallModel_ML/): Online training and inference workflows of a data-driven wall shear stress model for LES applied to a turbulent channel flow at a friction Reynolds number of 950. This example is an extension to [turbChannel_wallModel](./examples/turbChannel_wallModel/), which uses an algebraic equilibrium wall model (no ML).
 * [tgv_gnn_online](./examples/tgv_gnn_online/): Online training workflow using SmartSim to concurrently generate data and perform time independent training of the Dist-GNN model.
 * [tgv_gnn_online_traj](./examples/tgv_gnn_online_traj/): Online training workflow using SmartSim to concurrently generate data and perform time dependent training of the Dist-GNN model.
 * [tgv_gnn_online_traj_adios](./examples/tgv_gnn_online_traj_adios/): Online training workflow using ADIOS2 to concurrently generate data and perform time dependent training of the Dist-GNN model.
 * [shooting_workflow_smartredis](./examples/shooting_workflow_smartredis/): Online training workflow using SmartSim to shoot the nekRS solution forward in time leveraging the Dist-GNN model.
 * [shooting_workflow_adios](./examples/shooting_workflow_adios/): Online training workflow using ADIOS2 to shoot the nekRS solution forward in time leveraging the Dist-GNN model.
+* [turbChannel_srgnn_workflow](./examples/turbChannel_srgnn_workflow/): Similar to the [turbChannel_srgnn](./examples/turbChannel_srgnn/), however expanded to train the SR-GNN model using low-polynomial order nekRS data instead of projected data. After an initial run saving both high and low p-order checkpoints, EnsembleLauncher is used to run short nekRS simulations from each of the checkpoints to produce the training data. 
 
 ### Other examples
 

examples/README.md

@@ -15,6 +15,7 @@
 | shooting_workflow_smartredis |   x      |        |         |        x           |                  |              |                  |
 | shooting_workflow_adios      |   x      |        |         |                    |        x         |              |                  |
 | periodicHill_ensemble        |          |        |         |                    |                  |              |        x         |
+| turbChannel_srgnn_workflow   |          |    x   |    x    |                    |                  |        x     |        x         |
 
 
 ## Plain nekRS examples

examples/turbChannel_srgnn/README.md

@@ -1,7 +1,7 @@
 # Offline training of the SR-GNN model for mesh-based, three-dimensional super-resolution
 
 This example demonstrates the pipeline for training and deploying the SR-GNN model on nekRS field data. It builds upon the [turbulent channel example](../turbChannel/) available with nekRS by modifying the `.par` and `.udf` files and calling on the scripts available in the [SR-GNN repository](../../3rd_party/gnn/sr-gnn/).
-The example was adapted from the work of Shivam Barwey (ANL) and is based on this [paper](https://www.sciencedirect.com/science/article/abs/pii/S0045782525003445).
+The example was adapted from the work of Prof. Shivam Barwey (U. Notre Dame) and is based on this [paper](https://www.sciencedirect.com/science/article/abs/pii/S0045782525003445).
 
 The following important steps and modifications are highlighted for this example:
 

examples/turbChannel_srgnn/nrsrun_aurora

@@ -25,7 +25,7 @@ RFILE=run.sh
 echo "#!/bin/bash -l" > $RFILE
 echo "#PBS -S /bin/bash" >> $RFILE
 echo "#PBS -N nekrs-ml" >> $RFILE
-echo "#PBS -l select=1" >> $RFILE
+echo "#PBS -l select=${nodes}" >> $RFILE
 echo "#PBS -l walltime=00:30:00" >> $RFILE
 echo "#PBS -l filesystems=home:flare" >> $RFILE
 echo "#PBS -A <project_name>" >> $RFILE

examples/turbChannel_srgnn/nrsrun_polaris

@@ -22,7 +22,7 @@ RFILE=run.sh
 echo "#!/bin/bash -l" > $RFILE
 echo "#PBS -S /bin/bash" >> $RFILE
 echo "#PBS -N nekrs-ml" >> $RFILE
-echo "#PBS -l select=1" >> $RFILE
+echo "#PBS -l select=${nodes}" >> $RFILE
 echo "#PBS -l walltime=00:30:00" >> $RFILE
 echo "#PBS -l filesystems=home:eagle" >> $RFILE
 echo "#PBS -A <project_name>" >> $RFILE

examples/turbChannel_srgnn/turbChannel.udf

@@ -142,14 +142,14 @@ void UDF_Setup()
   platform->options.getArgs("POLYNOMIAL DEGREE",nekMeshPOrder);
   gnn_t* graph_ho = new gnn_t(nrs,nekMeshPOrder,verbose);
   graph_ho->gnnSetup();
-  graph_ho->gnnWrite();
+  graph_ho->gnnSRGNNWrite();
 
   // then write low order input files
   int gnnMeshPOrder;
   platform->options.getArgs("GNN POLY ORDER",gnnMeshPOrder);
   gnn_t* graph_lo = new gnn_t(nrs,gnnMeshPOrder,verbose);
   graph_lo->gnnSetup();
-  graph_lo->gnnWrite();
+  graph_lo->gnnSRGNNWrite();
 }
 
 void UDF_ExecuteStep(double time, int tstep)

examples/turbChannel_srgnn_workflow/README.md

@@ -0,0 +1,78 @@
+# SR-GNN Model Training Workflow with EnsembleLauncher
+
+This example extends the [turbChannel_srgnn](../turbChannel_srgnn/) example to train the SR-GNN model using low-polynomial order nekRS data instead of data projected from high to low p-order. After an initial nekRS run saving both high and projected low p-order checkpoints, EnsembleLauncher (EL) is used to automatically run short nekRS simulations from each of the checkpoints and produce more realistic low-p snapshots to be used as the training inputs.
+The example still uses the [turbulent channel](../turbChannel/) as the flow case.
+The SR-GNN model is based on the [work](https://www.sciencedirect.com/science/article/abs/pii/S0045782525003445) of Prof. Shivam Barwey (U. Notre Dame).
+
+This example combines many of the tools demonstrated in the [turbChannel_srgnn](../turbChannel_srgnn/) and [periodicHill_ensemble](../periodicHill_ensemble/) examples, with the following additional details:
+
+* The example uses two `.par` files. `turbChannel.par` is very similar to the one used by the regular [turbChannel_srgnn](../turbChannel_srgnn/) example and is used to run the initial nekRS simulation writing the `gnn_outputs_poly_*` directories and the `.f` files for the high p-order used by nekRS and the lower p-order specified with the `gnnPolynomialOrder` parameter (in this case set to 2). The second file called `turbChannel_simple.par` is used by the second set of nekRS simulations launched with EL which start from the aforementioned `.f` files and simply advance the simulation for a small number of time steps writing additional `.f` files to be used for training.
+* The example also uses two `.udf` files. Similarly to above, `turbChannel.udf` is used for the initial simulation and `turbChannel_simple.udf` is used for the additional simulations launched with EL.
+* The example contains a `gen_ensemble_inputs.py` script used to prepare the ensemble of nekRS simulations by creating run directorties with the appropriate files and the JSON configuration for EL all stored in `./run_dir`. The main arguments to the script are the case name, the values of the polynimial orders used to run the additional nekRS simulations, the number of time steps to run. For each additional nekRS simulation, `turbChannel_simple.par` will be copied into the run directory and renamed to `turbChannel.par` with the correct polynomial order and number of time steps. Additionally, the appropriate `.f` checkpoint file from the initial simulation will be linked and renamed to `restart.fld` so the new simulation can restart from the correct snapshot and p-order.
+* The training data from the model is collected from the checkpoint files produced by the ensemble of nekRS simulations, with the low p-order simulations contributing the coarse inputs and the high p-order simulations contributing the target outputs. 
+
+
+## Building nekRS
+
+Requirements:
+* Linux, Mac OS X (Microsoft WSL and Windows is not supported)
+* GNU/oneAPI/NVHPC/ROCm compilers (C++17/C99 compatible)
+* MPI-3.1 or later
+* CMake version 3.21 or later
+* PyTorch, PyTorch Geometric and PyTorch Cluster
+* Pymech (for reaking nekRS files from Python)
+* EnsembleLauncher for orchestrating the ensemble of nekRS runs
+
+To build nekRS and the required dependencies, first clone our GitHub repository:
+
+```sh
+https://github.com/argonne-lcf/nekRS-ML.git
+```
+
+Then, simply execute one of the build scripts contained in the repository.
+The HPC systems currently supported for this example are:
+* [Aurora](https://docs.alcf.anl.gov/aurora/) (Argonne LCF)
+
+For example, to build nekRS-ML on Aurora, from the login nodes execute 
+
+```sh
+./BuildMeOnAurora
+```
+
+## Running the example
+
+Scripts are provided to conveniently generate run scripts and config files for the workflow on the different ALCF systems.
+Note that a virtual environment with EnsembleLauncher, PyTorch Geometric and PyTorch Cluster is needed to launch the ensemble and training/inference, and by default the `gen_run_script` will create one with the required dependencies.
+
+**From a login node** execute:
+```sh
+./gen_run_script <system_name> </path/to/nekRS>
+```
+
+For more information on how to use `gen_run_script`, use `--help`
+
+```sh
+./gen_run_script <system_name> </path/to/nekRS> --help
+```
+
+The script will produce a `run.sh` script specifically tailored to the desired system and using the desired nekRS install directory. By default, the script is set up to run on 4 nodes. To change the number of nodes to run on, simply add the number of nodes to the script as follows
+
+```sh
+./gen_run_script <system_name> </path/to/nekRS> --nodes 8
+```
+
+Finally, to run the example simply submit the run script with
+
+```bash
+qsub run.sh
+```
+
+The `run.sh` script is composed of six steps:
+
+1. A precompilation step in which nekRS is run with the `--build-only` flag. This is done such that the `.cache` directory can be built beforehand.
+2. The initial nekRS simulation to generate the checkpoint `.f` files at the low and high p-orders and the `gnn_outputs_poly_*` directories. The example sets the higher p-order to 7 and the lower one to 2, however these values can be changed in the `run.sh` script.
+3. The ensemble of additional nek runs. First, `gen_ensemble_inputs.py` is executed to prepare the EL configuration and run directories, then EL is launched with the CLI command `el start`. 
+    * NOTE: The turbulent channel case can easily be run on a single node of Aurora using all PVC 6 GPUs (12 tiles). Therefore, the initial simulation and all other simulations launched by nekRS are set up to run on a single node. With 4 nodes available, this means initially in step 2, 3 nodes are not being utilized, but then all 4 nodes are used to run 4 parallel nekRS simulations during this step. Training also uses all 4 nodes. If more nodes are requested for this example, only the ensemble and training will benefit from these additional resources. Moreover, if this example is to be run on a different case requiring more nodes, simple change the value assigned to `NODES_PER_NEKRS` in the `run.sh` script and allocate sufficient number of nodes to run the workflow.
+4. Training data is collected from the ensemble of runs and the files are passed to the SR-GNN utility `nek_to_pt.py` to create the input data to the model in PyTorch format.
+5. Training of the SR-GNN model is performed in parallel using all available nodes and GPU.
+6. Inference is performed at the end to create `.f` files to be used to visualize the reconstruction of the SR-GNN model. 

examples/turbChannel_srgnn_workflow/clean

@@ -0,0 +1,9 @@
+#!/bin/bash
+
+rm -rf pytorch_cluster
+rm turbChannel0.f* turbChannel_p10.f* turbChannel_p70.f* turbChannel.nek5000 turbChannel_p1.nek5000 turbChannel_p7.nek5000
+rm -r gnn_outputs_poly*
+rm -r pt_datasets saved_models outputs ckpt predictions
+rm run.sh train.log nekrs.log
+rm -r run_dir logs
+rm main_status.json results.json