diff --git a/3rd_party/ensembleLauncher/README.md b/3rd_party/ensembleLauncher/README.md index 8b9f57fbb1..bdaf47b943 100644 --- a/3rd_party/ensembleLauncher/README.md +++ b/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. diff --git a/3rd_party/ensembleLauncher/nekrs_ensemble_utils.py b/3rd_party/ensembleLauncher/nekrs_ensemble_utils.py index c2a2975143..6d996df1b6 100644 --- a/3rd_party/ensembleLauncher/nekrs_ensemble_utils.py +++ b/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) diff --git a/README.md b/README.md index a6dc473ee7..f66c16dd7d 100644 --- a/README.md +++ b/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 diff --git a/examples/README.md b/examples/README.md index 941d6a4048..a1feb6f1de 100644 --- a/examples/README.md +++ b/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 diff --git a/examples/turbChannel_srgnn/README.md b/examples/turbChannel_srgnn/README.md index bfeddde503..57b2e32434 100644 --- a/examples/turbChannel_srgnn/README.md +++ b/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: diff --git a/examples/turbChannel_srgnn/nrsrun_aurora b/examples/turbChannel_srgnn/nrsrun_aurora index 79c08502c2..2ba0118bf2 100755 --- a/examples/turbChannel_srgnn/nrsrun_aurora +++ b/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 " >> $RFILE diff --git a/examples/turbChannel_srgnn/nrsrun_polaris b/examples/turbChannel_srgnn/nrsrun_polaris index b3471b7c61..d69f404be6 100755 --- a/examples/turbChannel_srgnn/nrsrun_polaris +++ b/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 " >> $RFILE diff --git a/examples/turbChannel_srgnn/turbChannel.udf b/examples/turbChannel_srgnn/turbChannel.udf index 5ad93f53f3..d0c7324894 100644 --- a/examples/turbChannel_srgnn/turbChannel.udf +++ b/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) diff --git a/examples/turbChannel_srgnn_workflow/README.md b/examples/turbChannel_srgnn_workflow/README.md new file mode 100644 index 0000000000..5a34c7fe2d --- /dev/null +++ b/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 +``` + +For more information on how to use `gen_run_script`, use `--help` + +```sh +./gen_run_script --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 --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. diff --git a/examples/turbChannel_srgnn_workflow/clean b/examples/turbChannel_srgnn_workflow/clean new file mode 100755 index 0000000000..d5aa8c3070 --- /dev/null +++ b/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 diff --git a/examples/turbChannel_srgnn_workflow/gen_ensemble_inputs.py b/examples/turbChannel_srgnn_workflow/gen_ensemble_inputs.py new file mode 100644 index 0000000000..ae664d0f25 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/gen_ensemble_inputs.py @@ -0,0 +1,330 @@ +""" +Stage EnsembleLauncher directories for turbChannel_srgnn_workflow. + +After the initial nekRS run has written high and low order checkpoints +(e.g. ``p20`` → polynomial order 2, ``p70`` → order 7), this script creates +a run directory and EnsembleLauncher JSON configs to run additional simulations +starting from each high and low order checkpoints. + +1. Finds all requested ``_p*.f*`` checkpoints under this example directory. +2. Creates one run directory per checkpoint under ``./run_dir//``. +3. Writes a per-member ``.par`` file. +4. Symlinks each checkpoint into its member dir and copies ``.udf`` / ``.usr``. +5. Writes the three EnsembleLauncher JSON configs. +""" + +from __future__ import annotations + +import argparse +import os +import re +import shutil +import sys +from pathlib import Path +from typing import List, Sequence + +HERE = Path(__file__).resolve().parent + +sys.path.append(os.path.join(os.environ["NEKRS_HOME"], "3rd_party", "ensembleLauncher")) +from nekrs_ensemble_utils import ( + setup_ensemble_dirs, + write_ensemble_configs, +) + + +def _field_frame_index(path: Path) -> int: + m = re.search(r"\.f(\d+)$", path.name, re.IGNORECASE) + return int(m.group(1)) if m else -1 + + +def _pp_tag(path: Path, case_name: str) -> int: + """Return the ``PP`` in ``_p.f#####``.""" + m = re.search(rf"{re.escape(case_name)}_p(\d+)\.", path.name, re.IGNORECASE) + if not m: + raise ValueError( + f"Checkpoint name {path.name!r} does not match " + f"{case_name}_p.f##### (cannot read p-tag)." + ) + return int(m.group(1)) + + +def polynomial_order_from_checkpoint(path: Path, case_name: str) -> int: + """Map nek field suffix ``p`` to ``[GENERAL] polynomialOrder`` (``PP / 10``).""" + pp = _pp_tag(path, case_name) + if pp % 10 != 0: + raise ValueError( + f"Expected p-tag in {path.name!r} to be a multiple of 10 (got p{pp}); " + "nekRS multiscale naming here is p(10*polynomialOrder)." + ) + return pp // 10 + + +def p_file_suffixes_from_p_orders_arg(values: List[int]) -> List[int]: + """Map ``--p-orders`` to filename digits ``PP`` in ``_pPP``. + + * If **every** value is in ``1..9``, treat as polynomial orders → ``PP = 10*N``. + * If **every** value is ``>= 10`` and ``% 10 == 0``, treat as literal p-tags. + """ + if not values: + raise ValueError("--p-orders produced an empty list") + if any(v < 1 for v in values): + raise ValueError(f"--p-orders values must be >= 1, got {values!r}") + + all_poly_small = all(1 <= v <= 9 for v in values) + all_literal_tags = all(v >= 10 and v % 10 == 0 for v in values) + + if all_poly_small: + return [10 * v for v in values] + if all_literal_tags: + return values + + raise ValueError( + "--p-orders must be either (a) all polynomial orders in 1..9, e.g. 7,2 " + "→ …_p70.f*, …_p20.f*, or (b) all literal nek p-tags (each ≥10 and " + f"divisible by 10), e.g. 70,20. Got: {values!r}" + ) + + +def discover_snapshots_for_case( + base: Path, case_name: str, patterns: Sequence[str] +) -> List[Path]: + seen: set[Path] = set() + out: List[Path] = [] + for pattern in patterns: + for p in base.glob(pattern): + if not p.is_file(): + continue + try: + _pp_tag(p, case_name) + except ValueError: + continue + rp = p.resolve() + if rp in seen: + continue + seen.add(rp) + out.append(p) + + def sort_key(path: Path) -> tuple: + try: + pp = _pp_tag(path, case_name) + except ValueError: + pp = 0 + return (pp, _field_frame_index(path), path.name) + + out.sort(key=sort_key) + return out + + +def member_name_for_snapshot(path: Path, case_name: str) -> str: + """Filesystem-safe directory name; includes p-tag and frame index.""" + idx = _field_frame_index(path) + pp = _pp_tag(path, case_name) + if idx >= 0: + return f"from_p{pp}_f{idx:05d}" + safe = re.sub(r"[^\w.\-]+", "_", path.name) + return f"cp_p{pp}_{safe}" + + +def parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "case", + help="nekRS case basename (e.g. turbChannel for turbChannel.par / .re2).", + ) + p.add_argument( + "--p-orders", + required=True, + help="Comma-separated polynomial orders for each checkpoint (e.g. 7,2).", + ) + p.add_argument( + "--num-steps", + type=int, + required=True, + help="Number of steps to run for each nekRS simulation.", + ) + p.add_argument( + "--outdir", + default=str(HERE / "run_dir"), + help="Output directory for per-member run dirs and EL JSON configs.", + ) + p.add_argument( + "--ppn", + type=int, + default=12, + help="MPI ranks per node per member (Aurora: 12).", + ) + p.add_argument( + "--nodes-per-member", + type=int, + default=1, + help="Nodes assigned to each ensemble member.", + ) + p.add_argument( + "--ngpus-per-process", + type=int, + default=1, + help="GPUs per MPI rank.", + ) + p.add_argument( + "--system", + default="aurora", + help="System name written into system_config.json.", + ) + p.add_argument( + "--backend", + default="dpcpp", + help="OCCA backend passed to nekrs --backend.", + ) + p.add_argument( + "--cpu-bind", + default="", + help="CPU bind string for EnsembleLauncher (comma-separated IDs for EL).", + ) + p.add_argument( + "--ensemble-name", + default="turbChannel_checkpoint_ensemble", + help="Name of the ensemble inside config.json.", + ) + return p.parse_args() + + +def main() -> None: + args = parse_args() + case_name = args.case + + # Check that the cache directory exists + cache_dir = HERE / ".cache" + if not cache_dir.is_dir(): + raise FileNotFoundError( + f"{cache_dir} not found; run nekrs --build-only first so all " + "members can share the same .cache" + ) + + # Parse the polynomial orders and generate the file patterns + raw_orders = [int(x.strip()) for x in args.p_orders.split(",") if x.strip()] + pp_tags = p_file_suffixes_from_p_orders_arg(raw_orders) + patterns = [f"{case_name}_p{pp}.f*" for pp in pp_tags] + pattern_desc = ", ".join(patterns) + + # Discover the snapshots for each polynomial order + snapshots = discover_snapshots_for_case(HERE, case_name, patterns) + if not snapshots: + raise FileNotFoundError( + f"No checkpoints matched under {HERE} for {pattern_desc!r}." + ) + + # Define the ensemble members + members = [] + for snap in snapshots: + poly = polynomial_order_from_checkpoint(snap, case_name) + name = member_name_for_snapshot(snap, case_name) + base_name = snap.name + rel = snap.name + members.append( + { + "name": name, + "par_overrides": { + "GENERAL": { + "startFrom": base_name, + "numSteps": float(args.num_steps), + "polynomialOrder": int(poly), + } + }, + "symlinks": {base_name: rel}, + } + ) + + print( + f"[gen_ensemble_inputs] {len(members)} members from {pattern_desc!r}: " + f"{[m['name'] for m in members]}" + ) + + # Link the re2 and cache files + re2 = HERE / f"{case_name}.re2" + symlink_files: List[str] = [".cache"] + if re2.is_file(): + symlink_files.insert(0, f"{case_name}.re2") + + # Check that the par file exists + par_path = HERE / f"{case_name}_simple.par" + if not par_path.is_file(): + raise FileNotFoundError(f"missing {par_path} file") + + # Add the udf, usr, box files to the copy list + copy_files: List[str] = [] + udf_path = HERE / f"{case_name}_simple.udf" + if not udf_path.is_file(): + raise FileNotFoundError(f"missing {udf_path} file") + copy_files.append(udf_path.name) + + usr = HERE / f"{case_name}.usr" + if usr.is_file(): + copy_files.append(usr.name) + + box = HERE / f"{case_name}.box" + if box.is_file(): + copy_files.append(box.name) + + # Create the run directories + member_dirs = setup_ensemble_dirs( + case_name=case_name, + members=members, + base_dir=str(HERE), + output_dir=args.outdir, + copy_files=copy_files, + symlink_files=symlink_files, + par_template=str(par_path.resolve()), + ) + + nek_udf_name = f"{case_name}.udf" + for d in member_dirs: + src = d / udf_path.name + dst = d / nek_udf_name + if not src.is_file(): + raise FileNotFoundError(f"expected copied {udf_path.name} in {d}") + if dst.exists() or dst.is_symlink(): + dst.unlink() + shutil.move(str(src), str(dst)) + + # Check that the checkpoint symlinks are present + for m, d in zip(members, member_dirs): + for dst_name in m.get("symlinks") or {}: + link = d / Path(dst_name).name + if not link.exists(): + raise RuntimeError( + f"Expected checkpoint symlink missing: {link} " + f"(member {m['name']!r})." + ) + + # Write the EnsembleLauncher JSON configs + paths = write_ensemble_configs( + out_dir=args.outdir, + member_dirs=member_dirs, + case_name=case_name, + nekrs_home=os.environ["NEKRS_HOME"], + system_name=args.system, + nodes_per_member=args.nodes_per_member, + ppn=args.ppn, + ngpus_per_process=args.ngpus_per_process, + backend=args.backend, + ensemble_name=args.ensemble_name, + cpu_bind=args.cpu_bind or None, + ) + + print( + f"[gen_ensemble_inputs] {len(member_dirs)} run directories under {args.outdir}" + ) + for kind, path in paths.items(): + print(f"[gen_ensemble_inputs] wrote {kind:<8} -> {path}") + print( + "[gen_ensemble_inputs] launch with: " + f"el start {paths['config']} " + f"--system-config-file {paths['system']} " + f"--launcher-config-file {paths['launcher']}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/turbChannel_srgnn_workflow/gen_run_script b/examples/turbChannel_srgnn_workflow/gen_run_script new file mode 100755 index 0000000000..8bd8fa9704 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/gen_run_script @@ -0,0 +1,10 @@ +#!/bin/bash + +SYSTEM=$1 +NEKRS_HOME=$2 + +if [ $# -lt 2 ]; then + $NEKRS_HOME/bin/ml/setup_case --help +else + $NEKRS_HOME/bin/ml/setup_case $SYSTEM $NEKRS_HOME --model "sr-gnn" --ensemble "el" --nodes 4 ${@:3} +fi diff --git a/examples/turbChannel_srgnn_workflow/nrsrun_aurora b/examples/turbChannel_srgnn_workflow/nrsrun_aurora new file mode 100755 index 0000000000..be6df7d031 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/nrsrun_aurora @@ -0,0 +1,127 @@ +#!/bin/bash +set -e + +#-------------------------------------- +: ${QUEUE:="debug-scaling"} +: ${NEKRS_GPU_MPI:=0} +: ${NEKRS_BACKEND:="dpcpp"} +: ${RANKS_PER_NODE:=12} +: ${NODES_PER_NEKRS:=1} +: ${CPU_BIND_LIST:="1:8:16:24:32:40:53:60:68:76:84:92"} +: ${CPU_BIND_LIST_EL:="1,8,16,24,32,40,53,60,68,76,84,92"} +: ${OCCA_DPCPP_COMPILER_FLAGS:="-O3 -fsycl -fsycl-targets=intel_gpu_pvc -ftarget-register-alloc-mode=pvc:auto -fma"} +: ${ONEAPI_SDK:=""} +: ${FRAMEWORKS_MODULE:="frameworks"} +: ${VENV_PATH:=""} +#-------------------------------------- + +source $NEKRS_HOME/bin/nrsqsub_utils +setup $# 1 + +#-------------------------------------- +# Generate the run script +RFILE=run.sh +echo "#!/bin/bash -l" > $RFILE +echo "#PBS -S /bin/bash" >> $RFILE +echo "#PBS -N nekrs-ml" >> $RFILE +echo "#PBS -l select=${nodes}" >> $RFILE +echo "#PBS -l walltime=01:00:00" >> $RFILE +echo "#PBS -l filesystems=home:flare" >> $RFILE +echo "#PBS -A " >> $RFILE +echo "#PBS -q ${QUEUE}" >> $RFILE +echo "#PBS -k doe" >> $RFILE +echo "#PBS -j oe" >> $RFILE +echo "cd \$PBS_O_WORKDIR" >> $RFILE + +echo -e "\nexport TZ='/usr/share/zoneinfo/US/Central'" >> $RFILE + +echo -e "\necho Jobid: \$PBS_JOBID" >>$RFILE +echo "echo Running on host \`hostname\`" >>$RFILE +echo "echo Running on nodes \`cat \$PBS_NODEFILE\`" >>$RFILE + +echo "module restore" >> $RFILE +echo "module load ${FRAMEWORKS_MODULE}" >> $RFILE +echo "source ${VENV_PATH}" >> $RFILE +echo "module list" >> $RFILE + +echo "export ZE_FLAT_DEVICE_HIERARCHY=FLAT" >> $RFILE +echo -e "\nexport NEKRS_HOME=$NEKRS_HOME" >>$RFILE +echo "export OCCA_DPCPP_COMPILER_FLAGS=\"$OCCA_DPCPP_COMPILER_FLAGS\"" >> $RFILE +echo "export FI_CXI_RX_MATCH_MODE=hybrid" >> $RFILE # required by parRSB + +# Temporary workaround while waiting on bugfix in runtime +echo "export UR_L0_USE_COPY_ENGINE=0" >> $RFILE + +echo -e "\n# Calculate the number of nodes and ranks for the nekRS simulation" >>$RFILE +echo "RANKS_PER_NODE=${RANKS_PER_NODE}" >> $RFILE +echo "NODES_PER_NEKRS=${NODES_PER_NEKRS}" >> $RFILE +echo "NEKRS_RANKS=\$(( RANKS_PER_NODE * NODES_PER_NEKRS ))" >> $RFILE +echo "TOTAL_NODES=${nodes}" >> $RFILE +echo "TOTAL_RANKS=\$(( TOTAL_NODES * RANKS_PER_NODE ))" >> $RFILE + +echo -e "\n# Set the polynomial orders" >> $RFILE +echo "INPUT_P=2" >> $RFILE +echo "TARGET_P=7" >> $RFILE + +echo -e "\n# Build nekRS cache" >>$RFILE +echo "echo \"Building nekRS cache\"" >> $RFILE +echo "mpiexec -n \${RANKS_PER_NODE} -ppn \${RANKS_PER_NODE} --cpu-bind=list:${CPU_BIND_LIST} -- $NEKRS_HOME/bin/nekrs --setup ${case} --backend ${NEKRS_BACKEND} --build-only \${RANKS_PER_NODE}" >> $RFILE + +echo -e "\n# Run initial nekRS simulation to generate high and low p order checkpoint files" >>$RFILE +echo "echo \"Running initial nekRS simulation to generate high and low p order checkpoint files\"" >> $RFILE +echo "mpiexec -n \${NEKRS_RANKS} -ppn \${RANKS_PER_NODE} --cpu-bind=list:${CPU_BIND_LIST} -- $NEKRS_HOME/bin/nekrs --setup ${case} --backend ${NEKRS_BACKEND} 2>&1 | tee nekrs.log" >> $RFILE + +echo -e "\n# Run ensemble of nekRS simulations starting from each high- and low-p checkpoint .f file" >>$RFILE +echo "echo \"Preparing nekRS run directories and EnsembleLauncher configs\"" >> $RFILE +echo "python gen_ensemble_inputs.py ${case} \\" >> $RFILE +echo " --p-orders \${INPUT_P},\${TARGET_P} \\" >> $RFILE +echo " --num-steps 200 \\" >> $RFILE +echo " --nodes-per-member \${NODES_PER_NEKRS} \\" >> $RFILE +echo " --ppn \${RANKS_PER_NODE} \\" >> $RFILE +echo " --cpu-bind \"${CPU_BIND_LIST_EL}\" \\" >> $RFILE +echo " --backend ${NEKRS_BACKEND} \\" >> $RFILE +echo " --system aurora" >> $RFILE +echo "echo \"Launching ensemble of nekRS simulations\"" >> $RFILE +echo "el start ./run_dir/config.json \\" >> $RFILE +echo " --system-config-file ./run_dir/system_config.json \\" >> $RFILE +echo " --launcher-config-file ./run_dir/launcher_config.json" >> $RFILE + +echo -e "\n# Generate training data from nekRS checkpoint files" >>$RFILE +echo "echo \"Generating training data from nekRS checkpoint files\"" >> $RFILE +echo "TARGET_SNAP_LIST=\`ls run_dir/from_turbChannel_p\${TARGET_P}*.f0*\`" >> $RFILE +echo "INPUT_SNAP_LIST=\`ls run_dir/from_turbChannel_p\${INPUT_P}*.f0*\`" >> $RFILE +echo "python ${NEKRS_HOME}/3rd_party/gnn/sr-gnn/nek_to_pt.py \\" >> $RFILE +echo " --case_path $PWD \\" >> $RFILE +echo " --target_snap_list \${TARGET_SNAP_LIST} \\" >> $RFILE +echo " --input_snap_list \${INPUT_SNAP_LIST} \\" >> $RFILE +echo " --target_poly_order \${TARGET_P} \\" >> $RFILE +echo " --input_poly_order \${INPUT_P} \\" >> $RFILE +echo " --n_element_neighbors 12" >> $RFILE + +echo -e "\n# Train model" >>$RFILE +echo "echo \"Training model\"" >> $RFILE +echo "mpiexec -n \${TOTAL_RANKS} -ppn \${RANKS_PER_NODE} --cpu-bind=list:${CPU_BIND_LIST} \\" >> $RFILE +echo " python ${NEKRS_HOME}/3rd_party/gnn/sr-gnn/main.py \\" >> $RFILE +echo " epochs=5 \\" >> $RFILE +echo " n_element_neighbors=12 \\" >> $RFILE +echo " n_messagePassing_layers=6 \\" >> $RFILE +echo " data_dir=$PWD/pt_datasets \\" >> $RFILE +echo " model_dir=$PWD/saved_models \\" >> $RFILE +echo " 2>&1 | tee train.log" >> $RFILE + +echo -e "\n# Perform inference and generate files for visualization" >>$RFILE +echo "echo \"Performing inference and generating files for visualization\"" >> $RFILE +echo "python ${NEKRS_HOME}/3rd_party/gnn/sr-gnn/postprocess.py \\" >> $RFILE +echo " --model_path $PWD/saved_models/gnn_3_7_132_128_3_2_6_True.tar \\" >> $RFILE +echo " --case_path $PWD \\" >> $RFILE +echo " --output_name ${case} \\" >> $RFILE +echo " --target_snap_list turbChannel_p70.f00000 \\" >> $RFILE +echo " --input_snap_list turbChannel_p10.f00000 \\" >> $RFILE +echo " --target_poly_order \${TARGET_P} \\" >> $RFILE +echo " --input_poly_order \${INPUT_P} \\" >> $RFILE +echo "cd predictions/gnn_3_7_132_128_3_2_6_True" >> $RFILE +echo "$NEKRS_HOME/bin/nrsvis turbChannel_pred" >> $RFILE +echo "$NEKRS_HOME/bin/nrsvis turbChannel_error" >> $RFILE +echo "cd ../../" >> $RFILE +chmod u+x $RFILE + diff --git a/examples/turbChannel_srgnn_workflow/restart.fld b/examples/turbChannel_srgnn_workflow/restart.fld new file mode 100644 index 0000000000..c5049a1b73 Binary files /dev/null and b/examples/turbChannel_srgnn_workflow/restart.fld differ diff --git a/examples/turbChannel_srgnn_workflow/turbChannel.box b/examples/turbChannel_srgnn_workflow/turbChannel.box new file mode 100644 index 0000000000..ad5a2e3cd8 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/turbChannel.box @@ -0,0 +1,9 @@ +-3 spatial dimension ( < 0 --> generate .rea/.re2 pair) +1 number of fields +#======================================================================= +Box +-32 -16 -16 nelx,nely,nelz for Box +0 1 1. x0,x1,gain +0 1 1. y0,y1,gain +0 1 1. z0,z1,gain +P ,P ,W ,W ,P ,P bc's (3 chars each!) diff --git a/examples/turbChannel_srgnn_workflow/turbChannel.par b/examples/turbChannel_srgnn_workflow/turbChannel.par new file mode 100644 index 0000000000..ef4ee851d1 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/turbChannel.par @@ -0,0 +1,34 @@ +[GENERAL] +#verbose = true +polynomialOrder = 7 +startFrom = "restart.fld" # time 160 +stopAt = numSteps +numSteps = 1000 + +dt = targetCFL=2 + max=2e-2 + initial=5e-3 +timeStepper = tombo2 + +checkpointControl = steps +checkpointInterval = 200 + +regularization = hpfrt + nModes=1 + scalingCoeff=5 +#constFlowRate = meanVelocity=1.0 + direction=X + +[PRESSURE] +residualTol = 1e-04 + +[VELOCITY] +boundaryTypeMap = zeroValue +viscosity = 1/10000 +rho = 1.0 +residualTol = 1e-06 + +[CASEDATA] +ReTau = 550 +xLength = 6.283185307 +zLength = 3.141592653 +betaY = 2.2 + +[ML] +gnnPolynomialOrder = 2 +srGNNMultiscale = true diff --git a/examples/turbChannel_srgnn_workflow/turbChannel.re2 b/examples/turbChannel_srgnn_workflow/turbChannel.re2 new file mode 100644 index 0000000000..bb2a503e8d Binary files /dev/null and b/examples/turbChannel_srgnn_workflow/turbChannel.re2 differ diff --git a/examples/turbChannel_srgnn_workflow/turbChannel.udf b/examples/turbChannel_srgnn_workflow/turbChannel.udf new file mode 100644 index 0000000000..9df1662f65 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/turbChannel.udf @@ -0,0 +1,169 @@ +#include "gnn.hpp" + +static dfloat ReTau; +static dfloat zLength; +static dfloat xLength; +static dfloat betaY; + +#ifdef __okl__ + +#endif + +/* User Functions */ + +void userf(double time) +{ + auto mesh = nrs->mesh; + dfloat mue, rho; + platform->options.getArgs("VISCOSITY", mue); + platform->options.getArgs("DENSITY", rho); + const dfloat RE_B = rho / mue; + const dfloat DPDX = (ReTau / RE_B) * (ReTau / RE_B); + + auto o_FUx = nrs->o_NLT + 0 * nrs->fieldOffset; + platform->linAlg->fill(mesh->Nlocal, DPDX, o_FUx); +} + +void useric(nrs_t *nrs) +{ + auto mesh = nrs->mesh; + + if (platform->options.getArgs("RESTART FILE NAME").empty()) { + const auto C = 5.17; + const auto k = 0.41; + const auto eps = 1e-2; + const auto kx = 23.0; + const auto kz = 13.0; + const auto alpha = kx * 2 * M_PI / xLength; + const auto beta = kz * 2 * M_PI / zLength; + dfloat mue; + platform->options.getArgs("VISCOSITY", mue); + + auto [x, y, z] = mesh->xyzHost(); + + std::vector U(mesh->dim * nrs->fieldOffset, 0.0); + for (int i = 0; i < mesh->Nlocal; i++) { + const auto yp = (y[i] < 0) ? (1 + y[i]) * ReTau : (1 - y[i]) * ReTau; + + dfloat ux = + 1 / k * log(1 + k * yp) + (C - (1 / k) * log(k)) * (1 - exp(-yp / 11) - yp / 11 * exp(-yp / 3)); + ux *= ReTau * mue; + + U[i + 0 * nrs->fieldOffset] = ux + eps * beta * sin(alpha * x[i]) * cos(beta * z[i]); + U[i + 1 * nrs->fieldOffset] = eps * sin(alpha * x[i]) * sin(beta * z[i]); + U[i + 2 * nrs->fieldOffset] = -eps * alpha * cos(alpha * x[i]) * sin(beta * z[i]); + } + nrs->o_U.copyFrom(U.data(), U.size()); + + } +} + +void outfld_wrapper(nrs_t *nrs, std::unique_ptr &checkpointWriter, const int N, double time, int tstep, std::string fileName) +{ + if (!checkpointWriter) { + checkpointWriter = iofldFactory::create("nek"); // or "adios" + if (platform->comm.mpiRank == 0) { + printf("create a new iofldFactory... %s\n", fileName.c_str()); + } + } + + if (!checkpointWriter->isInitialized()) { + auto visMesh = (nrs->cht) ? nrs->cds->mesh[0] : nrs->mesh; + checkpointWriter->open(visMesh, iofld::mode::write, fileName); + + if (platform->options.compareArgs("LOWMACH", "TRUE")) { + checkpointWriter->addVariable("p0th", nrs->p0th[0]); + } + + if (platform->options.compareArgs("VELOCITY CHECKPOINTING", "TRUE")) { + std::vector o_V; + for (int i = 0; i < visMesh->dim; i++) { + o_V.push_back(nrs->o_U.slice(i * nrs->fieldOffset, visMesh->Nlocal)); + } + checkpointWriter->addVariable("velocity", o_V); + } + + if (platform->options.compareArgs("PRESSURE CHECKPOINTING", "TRUE")) { + auto o_p = std::vector{nrs->o_P.slice(0, visMesh->Nlocal)}; + checkpointWriter->addVariable("pressure", o_p); + } + + for (int i = 0; i < nrs->Nscalar; i++) { + if (platform->options.compareArgs("SCALAR" + scalarDigitStr(i) + " CHECKPOINTING", "TRUE")) { + const auto temperatureExists = platform->options.compareArgs("SCALAR00 IS TEMPERATURE", "TRUE"); + std::vector o_Si = {nrs->cds->o_S.slice(nrs->cds->fieldOffsetScan[i], visMesh->Nlocal)}; + if (i == 0 && temperatureExists) { + checkpointWriter->addVariable("temperature", o_Si); + } else { + const auto is = (temperatureExists) ? i - 1 : i; + checkpointWriter->addVariable("scalar" + scalarDigitStr(is), o_Si); + } + } + } + } + + const auto outXYZ = platform->options.compareArgs("CHECKPOINT OUTPUT MESH", "TRUE"); + const auto FP64 = platform->options.compareArgs("CHECKPOINT PRECISION", "FP64"); + const auto uniform = (N < 0) ? true : false; + + checkpointWriter->writeAttribute("polynomialOrder", std::to_string(abs(N))); + checkpointWriter->writeAttribute("precision", (FP64) ? "64" : "32"); + checkpointWriter->writeAttribute("uniform", (uniform) ? "true" : "false"); + checkpointWriter->writeAttribute("outputMesh", "true"); + + checkpointWriter->addVariable("time", time); + + checkpointWriter->process(); +} + + +/* UDF Functions */ + +void UDF_Setup0(MPI_Comm comm, setupAide &options) +{ + platform->par->extract("casedata", "ReTau", ReTau); + platform->par->extract("casedata", "zLength", zLength); + platform->par->extract("casedata", "xLength", xLength); + platform->par->extract("casedata", "betaY", betaY); +} + +void UDF_Setup() +{ + if (platform->options.compareArgs("CONSTANT FLOW RATE", "FALSE")) { + nrs->userVelocitySource = &userf; + } + + useric(nrs); + + // gnn plugin + // first write high order input files + bool verbose = true; + int nekMeshPOrder; + platform->options.getArgs("POLYNOMIAL DEGREE",nekMeshPOrder); + gnn_t* graph_hi = new gnn_t(nrs,nekMeshPOrder,verbose); + graph_hi->gnnSetup(); + graph_hi->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->gnnSRGNNWrite(); +} + +void UDF_ExecuteStep(double time, int tstep) +{ + // Write interpolated checkpoint at polynomial order 2 + static std::unique_ptr iofld_N2; + static std::unique_ptr iofld_N7; + if (nrs->checkpointStep) { + outfld_wrapper(nrs, iofld_N2, 2, time, tstep, "turbChannel_p2"); + outfld_wrapper(nrs, iofld_N7, 7, time, tstep, "turbChannel_p7"); + } + + if (nrs->lastStep) { + if (iofld_N2) iofld_N2->close(); + if (iofld_N7) iofld_N7->close(); + } +} diff --git a/examples/turbChannel_srgnn_workflow/turbChannel.usr b/examples/turbChannel_srgnn_workflow/turbChannel.usr new file mode 100644 index 0000000000..1993213861 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/turbChannel.usr @@ -0,0 +1,28 @@ + subroutine usrdat2 ! This routine to modify mesh coordinates + include 'SIZE' + include 'TOTAL' + + parameter(BETAM = 2.2) + parameter(XLENGTH = 6.283185307) + parameter(ZLENGTH = 3.141592653) + + call rescale_x(xm1,0.0,XLENGTH) + call rescale_x(ym1,-1.0,1.0) + call rescale_x(zm1,0.0,ZLENGTH) + + ntot = nx1*ny1*nz1*nelt + + do i=1,ntot + ym1(i,1,1,1) = tanh(BETAM*ym1(i,1,1,1))/tanh(BETAM) + enddo + + do iel=1,nelt + do ifc=1,2*ndim + if (cbc(ifc,iel,1) .eq. 'W ') boundaryID(ifc,iel) = 1 + cbc(ifc,iel,2) = cbc(ifc,iel,1) + enddo + enddo + + return + end + diff --git a/examples/turbChannel_srgnn_workflow/turbChannel_simple.par b/examples/turbChannel_srgnn_workflow/turbChannel_simple.par new file mode 100644 index 0000000000..af194d4b0b --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/turbChannel_simple.par @@ -0,0 +1,30 @@ +[GENERAL] +#verbose = true +polynomialOrder = 7 +startFrom = "restart.fld" +stopAt = numSteps +numSteps = 1000 + +dt = targetCFL=2 + max=2e-2 + initial=5e-3 +timeStepper = tombo2 + +checkpointControl = steps +checkpointInterval = 200 + +regularization = hpfrt + nModes=1 + scalingCoeff=5 +#constFlowRate = meanVelocity=1.0 + direction=X + +[PRESSURE] +residualTol = 1e-04 + +[VELOCITY] +boundaryTypeMap = zeroValue +viscosity = 1/10000 +rho = 1.0 +residualTol = 1e-06 + +[CASEDATA] +ReTau = 550 +xLength = 6.283185307 +zLength = 3.141592653 +betaY = 2.2 diff --git a/examples/turbChannel_srgnn_workflow/turbChannel_simple.udf b/examples/turbChannel_srgnn_workflow/turbChannel_simple.udf new file mode 100644 index 0000000000..8573cdf6f0 --- /dev/null +++ b/examples/turbChannel_srgnn_workflow/turbChannel_simple.udf @@ -0,0 +1,151 @@ +static dfloat ReTau; +static dfloat zLength; +static dfloat xLength; +static dfloat betaY; + +#ifdef __okl__ + +#endif + +/* User Functions */ + +void userf(double time) +{ + auto mesh = nrs->mesh; + dfloat mue, rho; + platform->options.getArgs("VISCOSITY", mue); + platform->options.getArgs("DENSITY", rho); + const dfloat RE_B = rho / mue; + const dfloat DPDX = (ReTau / RE_B) * (ReTau / RE_B); + + auto o_FUx = nrs->o_NLT + 0 * nrs->fieldOffset; + platform->linAlg->fill(mesh->Nlocal, DPDX, o_FUx); +} + +void useric(nrs_t *nrs) +{ + auto mesh = nrs->mesh; + + if (platform->options.getArgs("RESTART FILE NAME").empty()) { + const auto C = 5.17; + const auto k = 0.41; + const auto eps = 1e-2; + const auto kx = 23.0; + const auto kz = 13.0; + const auto alpha = kx * 2 * M_PI / xLength; + const auto beta = kz * 2 * M_PI / zLength; + dfloat mue; + platform->options.getArgs("VISCOSITY", mue); + + auto [x, y, z] = mesh->xyzHost(); + + std::vector U(mesh->dim * nrs->fieldOffset, 0.0); + for (int i = 0; i < mesh->Nlocal; i++) { + const auto yp = (y[i] < 0) ? (1 + y[i]) * ReTau : (1 - y[i]) * ReTau; + + dfloat ux = + 1 / k * log(1 + k * yp) + (C - (1 / k) * log(k)) * (1 - exp(-yp / 11) - yp / 11 * exp(-yp / 3)); + ux *= ReTau * mue; + + U[i + 0 * nrs->fieldOffset] = ux + eps * beta * sin(alpha * x[i]) * cos(beta * z[i]); + U[i + 1 * nrs->fieldOffset] = eps * sin(alpha * x[i]) * sin(beta * z[i]); + U[i + 2 * nrs->fieldOffset] = -eps * alpha * cos(alpha * x[i]) * sin(beta * z[i]); + } + nrs->o_U.copyFrom(U.data(), U.size()); + + } +} + +void outfld_wrapper(nrs_t *nrs, std::unique_ptr &checkpointWriter, const int N, double time, int tstep, std::string fileName) +{ + if (!checkpointWriter) { + checkpointWriter = iofldFactory::create("nek"); // or "adios" + if (platform->comm.mpiRank == 0) { + printf("create a new iofldFactory... %s\n", fileName.c_str()); + } + } + + if (!checkpointWriter->isInitialized()) { + auto visMesh = (nrs->cht) ? nrs->cds->mesh[0] : nrs->mesh; + checkpointWriter->open(visMesh, iofld::mode::write, fileName); + + if (platform->options.compareArgs("LOWMACH", "TRUE")) { + checkpointWriter->addVariable("p0th", nrs->p0th[0]); + } + + if (platform->options.compareArgs("VELOCITY CHECKPOINTING", "TRUE")) { + std::vector o_V; + for (int i = 0; i < visMesh->dim; i++) { + o_V.push_back(nrs->o_U.slice(i * nrs->fieldOffset, visMesh->Nlocal)); + } + checkpointWriter->addVariable("velocity", o_V); + } + + if (platform->options.compareArgs("PRESSURE CHECKPOINTING", "TRUE")) { + auto o_p = std::vector{nrs->o_P.slice(0, visMesh->Nlocal)}; + checkpointWriter->addVariable("pressure", o_p); + } + + for (int i = 0; i < nrs->Nscalar; i++) { + if (platform->options.compareArgs("SCALAR" + scalarDigitStr(i) + " CHECKPOINTING", "TRUE")) { + const auto temperatureExists = platform->options.compareArgs("SCALAR00 IS TEMPERATURE", "TRUE"); + std::vector o_Si = {nrs->cds->o_S.slice(nrs->cds->fieldOffsetScan[i], visMesh->Nlocal)}; + if (i == 0 && temperatureExists) { + checkpointWriter->addVariable("temperature", o_Si); + } else { + const auto is = (temperatureExists) ? i - 1 : i; + checkpointWriter->addVariable("scalar" + scalarDigitStr(is), o_Si); + } + } + } + } + + const auto outXYZ = platform->options.compareArgs("CHECKPOINT OUTPUT MESH", "TRUE"); + const auto FP64 = platform->options.compareArgs("CHECKPOINT PRECISION", "FP64"); + const auto uniform = (N < 0) ? true : false; + + checkpointWriter->writeAttribute("polynomialOrder", std::to_string(abs(N))); + checkpointWriter->writeAttribute("precision", (FP64) ? "64" : "32"); + checkpointWriter->writeAttribute("uniform", (uniform) ? "true" : "false"); + checkpointWriter->writeAttribute("outputMesh", "true"); + + checkpointWriter->addVariable("time", time); + + checkpointWriter->process(); +} + + +/* UDF Functions */ + +void UDF_Setup0(MPI_Comm comm, setupAide &options) +{ + platform->par->extract("casedata", "ReTau", ReTau); + platform->par->extract("casedata", "zLength", zLength); + platform->par->extract("casedata", "xLength", xLength); + platform->par->extract("casedata", "betaY", betaY); +} + +void UDF_Setup() +{ + if (platform->options.compareArgs("CONSTANT FLOW RATE", "FALSE")) { + nrs->userVelocitySource = &userf; + } + + useric(nrs); +} + +void UDF_ExecuteStep(double time, int tstep) +{ + // Write interpolated checkpoint through wrapper + int nekMeshPOrder; + platform->options.getArgs("POLYNOMIAL DEGREE",nekMeshPOrder); + static std::unique_ptr iofld_srgnn; + std::string iofld_name = "turbChannel_p" + std::to_string(nekMeshPOrder); + if (nrs->checkpointStep) { + outfld_wrapper(nrs, iofld_srgnn, nekMeshPOrder, time, tstep, iofld_name); + } + + if (nrs->lastStep) { + if (iofld_N2) iofld_srgnn->close(); + } +} diff --git a/scripts/ml/setup_case b/scripts/ml/setup_case index 569159a300..8d8cb6a716 100755 --- a/scripts/ml/setup_case +++ b/scripts/ml/setup_case @@ -212,8 +212,12 @@ function setup_case() { CASE_NAME=${CASE_NAME:0:${#CASE_NAME}-14} elif [ "$DEPLOYMENT" == "offline" ]; then # offline example - CASE_NAME=$(ls *.par) - CASE_NAME=${CASE_NAME:0:${#CASE_NAME}-4} + mapfile -t par_files < <(ls *.par | grep -v '_.*\.par$') + if [ ${#par_files[@]} -ne 1 ]; then + echo "Error: expected exactly one base .par file, found: ${par_files[*]}" + exit 1 + fi + CASE_NAME=$(basename -s .par "${par_files[0]}") elif [ "${DEPLOYMENT}" == "colocated" ] || [ "${DEPLOYMENT}" == "clustered" ]; then # online example CASE_NAME=$(ls *.par.safe) diff --git a/src/plugins/gnn.cpp b/src/plugins/gnn.cpp index 16f4260727..3ab5bc790d 100644 --- a/src/plugins/gnn.cpp +++ b/src/plugins/gnn.cpp @@ -212,6 +212,32 @@ void gnn_t::gnnWrite() if (rank == 0) writeToFile(writePath + "/edge_index_element_local_vertex", edge_index_local_vertex, num_vertices_local, 2); } +void gnn_t::gnnSRGNNWrite() +{ + if (verbose) printf("[RANK %d] -- in gnnSRGNNWrite() \n", rank); + MPI_Comm &comm = platform->comm.mpiComm; + + // output directory + std::filesystem::path currentPath = std::filesystem::current_path(); + currentPath /= "gnn_outputs"; + writePath = currentPath.string(); + int poly_order = mesh->Nq - 1; + writePath = writePath + "_poly_" + std::to_string(poly_order); + //if (multiscale) writePath = writePath + "_multiscale"; // unnecessary and breaks SR-GNN pre-processing pipeline + if (rank == 0) + { + if (!std::filesystem::exists(writePath)) + { + std::filesystem::create_directory(writePath); + } + } + MPI_Barrier(comm); + + // Writing element-local edge index as text file (small) + if (rank == 0) writeToFile(writePath + "/edge_index_element_local", edge_index_local, num_edges_local, 2); + if (rank == 0) writeToFile(writePath + "/edge_index_element_local_vertex", edge_index_local_vertex, num_vertices_local, 2); +} + #ifdef NEKRS_ENABLE_SMARTREDIS void gnn_t::gnnWriteDB(smartredis_client_t* client) { diff --git a/src/plugins/gnn.hpp b/src/plugins/gnn.hpp index d98995b4c4..00a708f2ae 100644 --- a/src/plugins/gnn.hpp +++ b/src/plugins/gnn.hpp @@ -48,6 +48,7 @@ class gnn_t // member functions void gnnSetup(); void gnnWrite(); + void gnnSRGNNWrite(); void interpolateField(nrs_t* nrs, occa::memory& o_field_fine, dfloat* field_coarse, int dim); #ifdef NEKRS_ENABLE_SMARTREDIS void gnnWriteDB(smartredis_client_t* client);