tera24_athenapk.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "277823a6-4f27-4a39-a04c-72b5035f8134",
   "metadata": {},
   "source": [
    "# AthenaPK/Parthenon tutorial\n",
    "\n",
    "This tutorial is provided as jupyter notebook (and mirrored here for easier access and initial instructions).\n",
    "\n",
    "The tutorial itself is part of the **Towards exascale-ready astrophysics** workshop https://indico3-jsc.fz-juelich.de/event/169/ taking\n",
    "place 25-27 Sep 2024 online.\n",
    "\n",
    "The tutorial has been prepared by Philipp Grete <pgrete@hs.uni-hamburg.de> and is licensed under the BSD 3-Clause License.\n",
    "\n",
    "Additional resources:\n",
    "- AthenaPK\n",
    "  - Code: https://github.com/parthenon-hpc-lab/athenapk\n",
    "  - Doc: https://github.com/parthenon-hpc-lab/athenapk/blob/main/docs/README.md\n",
    "  - Chat: https://app.element.io/#/room/#AthenaPK:matrix.org\n",
    "- Parthenon (the underlying AMR framework)\n",
    "  - Code: https://github.com/parthenon-hpc-lab/parthenon\n",
    "  - Doc: https://parthenon-hpc-lab.github.io/parthenon\n",
    "  - Chat: https://app.element.io/#/room/#parthenon-general:matrix.org\n",
    "  - Papers\n",
    "    - Overview: https://doi.org/10.1177/10943420221143775\n",
    "    - Performance characteristics at scale: https://doi.org/10.1002/cpe.8069 (or https://www.osti.gov/servlets/purl/1993688 for the OA version)\n",
    "  - Some other open codes based on Parthenon:\n",
    "      - Parthenon-hydro (a hydro mini-app with full support for AMR similar to AthenaPK in structure but just 1000 LOC): https://github.com/parthenon-hpc-lab/parthenon-hydro\n",
    "      - KHARMA (GRMHD using the HARM scheme): https://github.com/AFD-Illinois/kharma\n",
    "      - Phoebus (Phifty One Ergs Blows Up a Star - GRMHD): https://github.com/lanl/phoebus"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c16f6ef-9ac6-4256-8d04-a0aadfea7f5d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The PR for the Parthenon frontend in yt is still open, so we need to use a custom version, see https://github.com/parthenon-hpc-lab/athenapk?tab=readme-ov-file#data-analysis\n",
    "# To accelerate the process, the package is already installed for the training but we need to make it known.\n",
    "import sys\n",
    "sys.path.insert(0, \"/p/project1/training2437/tera_day1/yt\")\n",
    "import yt\n",
    "yt.set_log_level(50) # disable logging\n",
    "\n",
    "# For some reason this cannot go to the top as it lead to errors when plotting with yt...\n",
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "import os\n",
    "user = os.environ['USER']\n",
    "\n",
    "from matplotlib.animation import FuncAnimation\n",
    "from IPython.display import Video"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d01aedd-231d-46ed-932e-9bd6b67bd683",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "a3924a7d-e674-4bc4-a001-890d4fd9e4a4",
   "metadata": {},
   "source": [
    "# Step 0: Getting access\n",
    "\n",
    "Follow instructions on how to get access https://gitlab.jsc.fz-juelich.de/hedgedoc/s/J-O7yP6hD#\n",
    "\n",
    "# Step 1: Preparing the tutorial environment\n",
    "Log in to https://jupyter-jsc.fz-juelich.de/\n",
    "\n",
    "Fire up a lab environment (ideally do this during the coffee break as it might take some time) with\n",
    "- Lab Config:\n",
    "  - System: `JURECA`\n",
    "  - Project: `training2437`\n",
    "  - Partition: `dc-gpu`\n",
    "- Resources (opens once `dc-gpu` is selected)\n",
    "  - Nodes: `1`\n",
    "  - GPUs: `4`\n",
    "  - Runtime: `120` (or more if you want to continue after the offical timeslot)\n",
    "- Kernels and extensions: keep defaults\n",
    "\n",
    "Once started open a terminal (which will run directly on the compute node) and execute the following to get the tutorial material\n",
    "```bash\n",
    "$ source /p/project1/training2437/tera_day1/setup_env.sh\n",
    "```\n",
    "\n",
    "Finally, navigate to the tutorial folder (`/p/project1/training2437/$USER/tera_day1/athenapk_tutorial`) on jupyter lab and open the notebook that's in the folder."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f695e99-189d-4e8b-94de-467e6e488085",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b>\n",
    "For a short discussion later (or on Slack) it might be useful to keep rough track of the code's performance for the various tests being run.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4812a37-6011-474f-bb37-dd0723d52b35",
   "metadata": {},
   "source": [
    "# Step 2: Building (not required for tutorial - just for reference)\n",
    "\n",
    "Is machine dependent and (especially for device/GPU) might take some time.\n",
    "Binaries to be run as part of the TERA24 tutorial were already prepared as follows.\n",
    "\n",
    "```bash\n",
    "$ module load Stages/2024 NVHPC OpenMPI CMake HDF5 Python Ninja\n",
    "$ module load MPI-settings/CUDA\n",
    "\n",
    "$ git clone https://github.com/parthenon-hpc-lab/athenapk.git\n",
    "$ cd athenapk\n",
    "# this will get the Parthenon and Kokkos sources that are build automatically\n",
    "$ git submodule update --init\n",
    "\n",
    "$ cmake -Bbuild-cuda-jureca -DKokkos_ENABLE_CUDA=ON -DKokkos_ARCH_AMPERE80=ON -DKokkos_ARCH_ZEN2=ON -DPARTHENON_DISABLE_HDF5_COMPRESSION=ON -DPARTHENON_ENABLE_PYTHON_MODULE_CHECK=OFF -GNinja\n",
    "$ cmake --build build-cuda-jureca -j 8\n",
    "```\n",
    "\n",
    "Please get in touch if you run into troubles building on your own machine or cluster.\n",
    "This particularly applies to the latter given that there are always machine specific details (also for running under batch systems like slurm).\n",
    "\n",
    "Additional information:\n",
    "- https://github.com/parthenon-hpc-lab/athenapk?tab=readme-ov-file#building-athenapk\n",
    "- https://parthenon-hpc-lab.github.io/parthenon/develop/src/building.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6374e01-f275-4c83-91c2-702671d072ff",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b>\n",
    "From here on the tutorial assumes that you're running the notebook on JUPYTER-JSC, use the visualization tools within the notebook, and execute the commands on the terminal opened within jupyter lab.\n",
    "\n",
    "If you know your ways around JSC systems, you can, of course, also directly connect to the systems and/or use other visualization tools like ParaView or Visit.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b043177-453e-45af-a952-a88523738587",
   "metadata": {},
   "source": [
    "# Step 3: Running 101\n",
    "\n",
    "Note that all problem generators and solvers are always included in the binary and can be chosen at runtime through the input file.\n",
    "This might be different from other common codes (including Athena++) where a binary is specifically built for a certain setup.\n",
    "\n",
    "## First simulation (Kelvin-Helmholtz instability on a uniform grid)\n",
    "\n",
    "### Run\n",
    "\n",
    "On the terminal execute the simulation (we'll look at the input file `kh.in` and output while the simulation runs in the background).\n",
    "\n",
    "```bash\n",
    "# Loading modules is only required once\n",
    "# Skip (already done by the setup_env.sh script above)\n",
    "# module load Stages/2024 NVHPC OpenMPI CMake HDF5 Python\n",
    "# module load MPI-settings/CUDA\n",
    "\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_uniform\n",
    "$ srun -n 1 ../bin/athenaPK -i kh.in |& tee ath.out\n",
    "```\n",
    "\n",
    "Given that we run interactively, we pipe the output to `ath.out` so that it's available later (for comparison).\n",
    "\n",
    "### The input file\n",
    "\n",
    "Let's take the input file apart.\n",
    "\n",
    "In general, input files are split into blocks indicated by `<BLOCK_NAME>` where\n",
    "`BLOCK_NAME` may also contain`/` to logically separate information.\n",
    "\n",
    "The `<job>` block is mandatory for AthenaPK as it defines which problem generator to load.\n",
    "These are typically connected to their respective source files in [`athenapk/src/pgen`](https://github.com/parthenon-hpc-lab/athenapk/tree/main/src/pgen).\n",
    "```\n",
    "<job>\n",
    "problem_id = kh\n",
    "```\n",
    "\n",
    "Each problem generator typically comes with additonal parameters.\n",
    "In AthenaPK it is customary to collect these parameters in the `<problem/PROBLEMNAME>` block, like\n",
    "```\n",
    "<problem/kh>\n",
    "iprob  = 5               # AMR test setup from Athena++ method paper\n",
    "vflow  = 1.0             # relative shear velocity of the interfaces\n",
    "a      = 0.01            # thickness of shearing layer\n",
    "drat   = 2               # density contrast/ratio\n",
    "amp    = 0.01            # amplitude of initial vy perturbation\n",
    "sigma  = 0.2             # thickness of perturbed layer\n",
    "```\n",
    "\n",
    "The overall mesh is configured within the `<parthenon/mesh>` block.\n",
    "Note the `parthenon` prefix, which typically means that these inputs are controlled\n",
    "by the [Parthenon](https://github.com/parthenon-hpc-lab/parthenon) framework\n",
    "and additional information might be located in the Parthenon documentation, e.g., https://parthenon-hpc-lab.github.io/parthenon/develop/src/inputs.html\n",
    "\n",
    "Here, we setup a uniform, static 2D grid with 512x512 cells and periodic boundary conditions.\n",
    "```\n",
    "<parthenon/mesh>\n",
    "refinement = none\n",
    "nghost = 2\n",
    "\n",
    "nx1 = 512\n",
    "x1min = -0.5\n",
    "x1max = 0.5\n",
    "ix1_bc = periodic\n",
    "ox1_bc = periodic\n",
    "\n",
    "nx2 = 512\n",
    "x2min = -0.5\n",
    "x2max = 0.5\n",
    "ix2_bc = periodic\n",
    "ox2_bc = periodic\n",
    "\n",
    "nx3 = 1\n",
    "x3min = -0.5\n",
    "x3max = 0.5\n",
    "ix3_bc = periodic\n",
    "ox3_bc = periodic\n",
    "```\n",
    "\n",
    "The mesh itself it logically separate into (identical) meshblocks whose size is configured in the `<parthenon/meshblock>` block.\n",
    "The size is fixed for all meshblocks and cannot be changed once the simulation has started.\n",
    "Moreover, each active point in the simulation domain is covered by exactly one meshblock (this information is particularly relevant for simulations with mesh refinement).\n",
    "All meshblocks communicate with each other via ghost/guard/halo cells (whose number is set in the `<parthenon/mesh>` block above.\n",
    "\n",
    "Here, we pick a block size of 128x512, i.e., the mesh is split into 4 blocks total.\n",
    "Picking the right block size is highly problem dependent and also connected to performance (see [101 on performance](step-4:-Performance-101)).\n",
    "```\n",
    "<parthenon/meshblock>\n",
    "nx1=128\n",
    "nx2=512\n",
    "nx3=1\n",
    "```\n",
    "\n",
    "Parthenon also takes care of time integration, which is why the following block is also prefixed with `parthenon`.\n",
    "\n",
    "The following block will use a predictor-corrector Van-Leer type integrator with a CFL of 0.4.\n",
    "The simulation will run till time 1.2 (in code time units) and not be limited by the number of cylces performed (negative `nlim`).\n",
    "```\n",
    "<parthenon/time>\n",
    "integrator = vl2\n",
    "cfl = 0.4\n",
    "tlim = 1.2\n",
    "nlim = -1\n",
    "```\n",
    "\n",
    "The `<hydro>` block is specific to AthenaPK again, see https://github.com/parthenon-hpc-lab/athenapk/blob/main/docs/input.md for additional information.\n",
    "\n",
    "Here, we choose an adiabatic equation of state with ratio of specific heats of 1.4.\n",
    "Fluxes are calcuated using an HLLC Riemann solver with second-order, piecewise-linear reconstruction (in the primitive variables).\n",
    "```\n",
    "<hydro>\n",
    "eos = adiabatic\n",
    "riemann = hllc\n",
    "reconstruction = plm\n",
    "gamma = 1.4 # gamma = C_p/C_v\n",
    "```\n",
    "\n",
    "Finally, we define an output block (again controlled by Parthenon, see https://parthenon-hpc-lab.github.io/parthenon/develop/src/outputs.html for additional details.\n",
    "\n",
    "Here, we pick the `rst` type, which translates to `restart` and allows to restart a simulation (and can also be read by standard postprocessing tools, see in the following).\n",
    "Outputs are written every 0.1 time intervals.\n",
    "<div class=\"alert alert-block alert-warning\">\n",
    "<b>Note:</b>\n",
    "Outputs are overwritten without warning.\n",
    "</div>\n",
    "\n",
    "```\n",
    "<parthenon/output0>\n",
    "file_type = rst\n",
    "dt = 0.1\n",
    "```\n",
    "\n",
    "The simulation should be done by now, so let's have a look at the output on the terminal.\n",
    "\n",
    "### The terminal output\n",
    "\n",
    "A typical output of an AthenaPK simulation (or more generally Parthenon-based simulation) will start with general information available within the code framework.\n",
    "\n",
    "Bonus info: Parthenon provides flexible interfaces to create variables that are described through metadata and afterwards handled transparently (e.g., communicated), see https://parthenon-hpc-lab.github.io/parthenon/develop/src/interface/metadata.html\n",
    "In this simulation so-called sparse variables (that are only allocated on a subset of blocks) or \"swarms\" (i.e., particles based variables/fields)\n",
    "are not used but cell-centered variables (and their fluxes).\n",
    "\n",
    "```\n",
    "Starting up hydro driver\n",
    "# Variables in use:\n",
    "# Package: parthenon::resolved_state\n",
    "# ---------------------------------------------------\n",
    "# Variables:\n",
    "# Name  Metadata flags\n",
    "# ---------------------------------------------------\n",
    "prim                      Cell,Provides,Real,Derived,Hydro,parthenon::resolved_state\n",
    "cons                      Cell,Provides,Real,Independent,FillGhost,WithFluxes,Hydro,parthenon::resolved_state\n",
    "bnd_flux::cons            Face,Provides,Real,Derived,OneCopy,Flux,parthenon::resolved_state\n",
    "# ---------------------------------------------------\n",
    "# Sparse Variables:\n",
    "# Name  sparse id       Metadata flags\n",
    "# ---------------------------------------------------\n",
    "# ---------------------------------------------------\n",
    "# Swarms:\n",
    "# Swarm Value   metadata\n",
    "# ---------------------------------------------------\n",
    "```\n",
    "\n",
    "Afterwards, the simulation actually starts.\n",
    "For each cycle, the cycle number (`cycle`), current time (`time`), current timestep (`dt`), performance in total cell updates per\n",
    "wallclock second (across all computing resources, `zone-cycles/wsec_step`), current simulation runtime in wallclock seconds (`wsec_total`),\n",
    "and wallclock seconds for the current cycle (`wsec_step`) are printed.\n",
    "Having performance numbers per cycles has turned out useful to gauge performance in simulations with varying load (like with adapative mesh refinement)\n",
    "or to identify machine issues/variability.\n",
    "\n",
    "```\n",
    "Setup complete, executing driver...\n",
    "\n",
    "cycle=0 time=0.0000000000000000e+00 dt=3.2952612821802183e-04 zone-cycles/wsec_step=0.00e+00 wsec_total=1.44e-01 wsec_step=2.40e+00\n",
    "cycle=1 time=3.2952612821802183e-04 dt=3.2952539311580179e-04 zone-cycles/wsec_step=5.16e+06 wsec_total=1.94e-01 wsec_step=5.08e-02\n",
    "cycle=2 time=6.5905152133382367e-04 dt=3.2952465716610244e-04 zone-cycles/wsec_step=2.55e+07 wsec_total=2.05e-01 wsec_step=1.03e-02\n",
    "```\n",
    "\n",
    "Finally, once the simulation is done an overview is printed.\n",
    "Note that the total performance `zone-cycles/wallsecond` might be different from the per-cycle output as it is an aggregate number.\n",
    "```\n",
    "cycle=3865 time=1.1994642549437093e+00 dt=2.7515679713439601e-04 zone-cycles/wsec_step=3.10e+07 wsec_total=3.29e+01 wsec_step=8.46e-03\n",
    "cycle=3866 time=1.1997394117408438e+00 dt=2.6058825915620254e-04 zone-cycles/wsec_step=3.10e+07 wsec_total=3.29e+01 wsec_step=8.46e-03\n",
    "cycle=3867 time=1.2000000000000000e+00 dt=2.7515150765570855e-04 zone-cycles/wsec_step=1.29e+07 wsec_total=3.29e+01 wsec_step=2.03e-02\n",
    "\n",
    "Driver completed.\n",
    "time=1.20e+00 cycle=3867\n",
    "tlim=1.20e+00 nlim=-1\n",
    "\n",
    "walltime used = 3.29e+01\n",
    "zone-cycles/wallsecond = 3.08e+07\n",
    "```\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a60491c-f6f5-4340-b798-6267f1b993fc",
   "metadata": {},
   "source": [
    "### Plotting the results\n",
    "\n",
    "Various options exists to plot AthenaPK output (or more generally, Parthenon-based codes), see https://github.com/parthenon-hpc-lab/athenapk?tab=readme-ov-file#data-analysis\n",
    "\n",
    "We'll use yt (https://yt-project.org/) with the frontend for Partheon-based codes provided above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9042e9b-4e2c-494f-9f2e-d2e58002e70f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://yt-project.org/doc/cookbook/complex_plots.html#time-series-movie\n",
    "ts = yt.load(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/kh_2d_uniform/parthenon.out0.*.rhdf\")\n",
    "\n",
    "plot = yt.SlicePlot(ts[0], \"z\", (\"gas\", \"density\"))\n",
    "plot.set_log((\"gas\", \"density\"), False)\n",
    "plot.set_zlim((\"gas\", \"density\"), 0.8, 2.1)\n",
    "\n",
    "fig = plot.plots[(\"gas\", \"density\")].figure\n",
    "\n",
    "# animate must accept an integer frame number. We use the frame number\n",
    "# to identify which dataset in the time series we want to load\n",
    "def animate(i):\n",
    "    ds = ts[i]\n",
    "    plot._switch_ds(ds)\n",
    "\n",
    "animation = FuncAnimation(fig, animate, frames=len(ts))\n",
    "animation.save(\"kh_2d_uniform.mp4\")\n",
    "\n",
    "Video(\"kh_2d_uniform.mp4\", embed=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "190597d6-c4da-4fe4-88e0-0211b479e3b2",
   "metadata": {},
   "source": [
    "NB: yt is a powerful data analysis tool.\n",
    "\n",
    "Many (derived) fields are readily available and standard plot modifcations are straightforward."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09ef6f84-9d6f-462e-976c-1d46a042e8c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = yt.load(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/kh_2d_uniform/parthenon.out0.final.rhdf\")\n",
    "slc = yt.SlicePlot(ds, \"z\", (\"gas\", \"vorticity_z\"))\n",
    "slc.set_log((\"gas\", \"vorticity_z\"), False)\n",
    "slc.annotate_grids()\n",
    "slc.annotate_streamlines((\"gas\", \"velocity_x\"), (\"gas\", \"velocity_y\"))\n",
    "\n",
    "# see ds.derived_field_list for a full list of derived fields"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5a723f3-daa1-4c9d-9645-c08b31b87207",
   "metadata": {},
   "source": [
    "## Modifications from the command line\n",
    "\n",
    "All input parameters can be modified from the command line and take precedence over the data in the input file.\n",
    "\n",
    "The variable for a given block in an input file is combined to a single argument via an additional `/`, e.g.,\n",
    "```\n",
    "<hydro>\n",
    "reconstruction = ppm\n",
    "```\n",
    "turns into `hydro/reconstruction=ppm` on the command line.\n",
    "\n",
    "### Run\n",
    "\n",
    "Run in a terminal\n",
    "```bash\n",
    "$ mkdir /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_uniform_ppm\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_uniform_ppm\n",
    "$ srun -n 1 ../bin/athenaPK -i ../kh_2d_uniform/kh.in hydro/reconstruction=ppm |& tee ath.out\n",
    "```\n",
    "\n",
    "Note that other reconstruction methods (like `wenoz` or `ppm`) require larger stencils and thus need additionan ghost zones (3 for those two methods).\n",
    "\n",
    "Can you fix this (on the command line) so that the simulation runs?\n",
    "\n",
    "<details>\n",
    "<summary>Click here for solution</summary>\n",
    "\n",
    "```bash\n",
    "$ srun -n 1 ../bin/athenaPK -i ../kh_2d_uniform/kh.in hydro/reconstruction=ppm parthenon/mesh/nghost=3 |& tee ath.out\n",
    "```\n",
    "\n",
    "</details>\n",
    "\n",
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61c31b5e-8275-48b3-987a-afb16b358add",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://yt-project.org/doc/cookbook/complex_plots.html#time-series-movie\n",
    "ts = yt.load(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/kh_2d_uniform_ppm/parthenon.out0.*.rhdf\")\n",
    "\n",
    "plot = yt.SlicePlot(ts[0], \"z\", (\"gas\", \"density\"))\n",
    "plot.set_log((\"gas\", \"density\"), False)\n",
    "plot.set_zlim((\"gas\", \"density\"), 0.8, 2.1)\n",
    "\n",
    "fig = plot.plots[(\"gas\", \"density\")].figure\n",
    "\n",
    "# animate must accept an integer frame number. We use the frame number\n",
    "# to identify which dataset in the time series we want to load\n",
    "def animate(i):\n",
    "    ds = ts[i]\n",
    "    plot._switch_ds(ds)\n",
    "\n",
    "animation = FuncAnimation(fig, animate, frames=len(ts))\n",
    "animation.save(\"kh_2d_uniform_ppm.mp4\")\n",
    "\n",
    "Video(\"kh_2d_uniform_ppm.mp4\", embed=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9571c26-de7c-4886-944d-428e0c9d0638",
   "metadata": {},
   "source": [
    "## Restarting and more command line options\n",
    "\n",
    "As mentioned above, restart/checkpointing files allow to restart a simulation (e.g., for fault tolerance or if\n",
    "simulations are running longer than the maximum walltime for a single job).\n",
    "\n",
    "In this context, two command line options are important (which are independent of each other):\n",
    "\n",
    "- `-r RESTART_FILE` to continue a simulation, and\n",
    "- `-t HH:MM:SS` to set a wallock time limit after which the simulation gracefully stops (and writes are restart/checkpoint file)\n",
    "\n",
    "### Run\n",
    "\n",
    "To try run\n",
    "\n",
    "```bash\n",
    "$ mkdir /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_uniform_restart\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_uniform_restart\n",
    "$ srun -n 1 ../bin/athenaPK -r ../kh_2d_uniform/parthenon.out0.00005.rhdf -t 00:00:10 |& tee ath.out\n",
    "```\n",
    "\n",
    "### Output\n",
    "\n",
    "```\n",
    "Blocks assigned to rank 0: 0:3\n",
    "Var: cons:5\n",
    "Starting up hydro driver\n",
    "...\n",
    "cycle=2696 time=8.6758703747135080e-01 dt=2.9846780230757841e-04 zone-cycles/wsec_step=3.09e+07 wsec_total=9.88e+00 wsec_step=8.47e-03\n",
    "cycle=2697 time=8.6788550527365838e-01 dt=2.9842965735897364e-04 zone-cycles/wsec_step=1.25e+07 wsec_total=9.90e+00 wsec_step=2.10e-02\n",
    "\n",
    "Terminating on wall-time limit\n",
    "\n",
    "Driver timed out.  Restart to continue.\n",
    "time=8.68e-01 cycle=2697\n",
    "tlim=1.20e+00 nlim=-1\n",
    "\n",
    "walltime used = 9.90e+00\n",
    "zone-cycles/wallsecond = 3.09e+07\n",
    "```\n",
    "\n",
    "Note the `Blocks assigned to rank 0: 0:3`, which is already hinting at the possibility to restart simulations with different number of processes/GPUs.\n",
    "\n",
    "Also note the `Terminating on wall-time limit`, which now triggered the end of the simulation.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17889e26-762b-4934-a9b3-0bfcab687162",
   "metadata": {},
   "source": [
    "# Step 4: Running 201\n",
    "\n",
    "## AMR - adaptive mesh refinement\n",
    "\n",
    "### Run\n",
    "\n",
    "Again lets first start the simulation and then have a look at the input file while the simulation runs.\n",
    "\n",
    "```bash\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_amr\n",
    "$ srun -n 1 ../bin/athenaPK -i kh.in |& tee ath.out\n",
    "```\n",
    "\n",
    "### Input file\n",
    "\n",
    "What has changed compared to the input file for the uniform grid?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a05db034-d9b6-43fa-88be-d2ae81d11a14",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git diff --no-index --word-diff --color-words /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_uniform/kh.in /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/kh_2d_amr/kh.in"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68d7f1d2-79dd-43ce-b2ba-41adf25c2f01",
   "metadata": {},
   "source": [
    "In other words:\n",
    "\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "    \n",
    "- Adaptive mesh refinement was enabled limiting the hierarchy to three levels total.  \n",
    "- Mesh refinement in Parthenon based codes always happens with refinment factors of two between the levels.\n",
    "- Moreover, the simulation was run at the same effective resolution (the root grid has only 1/4 of the number of cells (512->128) but two additional levels of refinement are allowed.\n",
    "- Adapative mesh refinement is triggered given a threshold in the velocity gradient (this choice here is a little arbitrary and just set to match the setup in the Athena++ code paper).\n",
    "- In addition, the meshblock size was adjusted (reduced) to allow for localized refinement -- otherwise, the entire domain would just refine).\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d04af2f1-1166-4d9c-92ab-582c5e0c9f62",
   "metadata": {},
   "source": [
    "### Terminal output\n",
    "\n",
    "The output differs slightly:\n",
    "\n",
    "```\n",
    "cycle=0 time=0.0000000000000000e+00 dt=3.2952734694010211e-04 zone-cycles/wsec_step=0.00e+00 wsec_total=8.76e-02 wsec_step=2.98e+00 zone-cycles/wsec=0.00e+00 wsec_AMR=0.00e+00\n",
    "---------------------- Current Mesh structure ----------------------\n",
    "Number of Trees = 64\n",
    "Total number of MeshBlocks = 352\n",
    "Number of physical refinement levels = 2\n",
    "Number of logical  refinement levels = 2\n",
    "  Physical level = 0 (logical level = 0): 32 MeshBlocks, cost = 32\n",
    "  Physical level = 1 (logical level = 1): 64 MeshBlocks, cost = 64\n",
    "  Physical level = 2 (logical level = 2): 256 MeshBlocks, cost = 256\n",
    "--------------------------------------------------------------------\n",
    "cycle=1 time=3.2952734694010211e-04 dt=3.2952683932585395e-04 zone-cycles/wsec_step=3.27e+05 wsec_total=3.63e-01 wsec_step=2.75e-01 zone-cycles/wsec=3.27e+05 wsec_AMR=5.67e-06\n",
    "```\n",
    "\n",
    "First, the mesh structure is being shown (also every time it changes) and additional performance counters are given:\n",
    "- `zone-cycles/wsec` differs from `zone-cycles/wsec_step` by including time for the AMR machinery (like checking for refinement)\n",
    "- `wsec_AMR` for the wallclock time spent in remeshing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce725a21-51fb-4a81-9d36-4b59d11deaf6",
   "metadata": {},
   "source": [
    "### Plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1723c1a-ac7e-450a-8445-155bd0179fd3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://yt-project.org/doc/cookbook/complex_plots.html#time-series-movie\n",
    "ts = yt.load(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/kh_2d_amr/parthenon.out0.*.rhdf\")\n",
    "\n",
    "plot = yt.SlicePlot(ts[0], \"z\", (\"gas\", \"density\"))\n",
    "plot.set_log((\"gas\", \"density\"), False)\n",
    "plot.set_zlim((\"gas\", \"density\"), 0.8, 2.1)\n",
    "plot.annotate_grids()\n",
    "\n",
    "fig = plot.plots[(\"gas\", \"density\")].figure\n",
    "\n",
    "# animate must accept an integer frame number. We use the frame number\n",
    "# to identify which dataset in the time series we want to load\n",
    "def animate(i):\n",
    "    ds = ts[i]\n",
    "    plot._switch_ds(ds)\n",
    "\n",
    "animation = FuncAnimation(fig, animate, frames=len(ts))\n",
    "animation.save(\"kh_2d_amr.mp4\")\n",
    "\n",
    "Video(\"kh_2d_amr.mp4\", embed=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dd7b71b-d5ff-40fd-9567-ad78fe8dace2",
   "metadata": {},
   "source": [
    "Bonus: Is there anything you noted about this simulation?\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "\n",
    "- Looks practically identical to the uniform grid one.  \n",
    "- But performance was not worse than unigrid (we'll cover this later).\n",
    "- On the up side: the snapshots/outputs got smaller (for effectively the same result).\n",
    "- The final snapshot only shows two levels of refinement, i.e., the very large blocks are gone meaning that the entire root grid is refinement (and thus the refinement strategy chosen/block size is not ideal).\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55f2ac81-210a-4fc8-b83b-77d332d4c457",
   "metadata": {},
   "source": [
    "## 3D MHD turbulence with additional outputs\n",
    "\n",
    "### Run\n",
    "\n",
    "Again lets first start the simulation and then have a look at the input file while the simulation runs.\n",
    "\n",
    "```bash\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/turb_3d_mhd_uniform\n",
    "$ srun -n 1 ../bin/athenaPK -i turbulence.in |& tee ath.out\n",
    "```\n",
    "\n",
    "### The input file\n",
    "\n",
    "What's new in this input file?\n",
    "\n",
    "The problem generator changed (we're now running a turbulence simulation) and associated parameters.\n",
    "\n",
    "More details on the turbulence driver can be found here: https://gitlab.com/pgrete/kathena/-/wikis/turbulence (TODO: move to GH)\n",
    "```\n",
    "<job>\n",
    "problem_id = turbulence\n",
    "\n",
    "<problem/turbulence>\n",
    "rho0         = 1.0      # initial mean density\n",
    "p0           = 1.0      # initial mean pressure\n",
    "b0           = 0.01     # initial magnetic field (x-direction)\n",
    "b_config     = 0        # 0 - net flux; 1 - no net flux uniform B; 2 - non net flux sin B; 4 - field loop\n",
    "kpeak        = 2.0      # characteristic wavenumber\n",
    "corr_time    = 1.0      # autocorrelation time of the OU forcing process\n",
    "rseed        = 20190729 # random seed of the OU forcing process\n",
    "sol_weight   = 1.0      # solenoidal weight of the acceleration field\n",
    "accel_rms    = 0.5      # root mean square value of the acceleration field\n",
    "num_modes    = 20       # number of wavemodes\n",
    "```\n",
    "\n",
    "We have additional output blocks:\n",
    "\n",
    "First a \"history\" output.This results in a text file that contains additional information of the state of the simulation\n",
    "(e.g., mean energy densities or Mach numbers) and the user can add problem specific fields as required.\n",
    "This is often useful to keep track of a high level overview of the simulation during runtime (and typically written with higher frequency, here `dt = 0.05`, than normal outputs).\n",
    "```\n",
    "<parthenon/output1>\n",
    "file_type  = hst        # History data dump\n",
    "dt         = 0.05       # time increment between outputs\n",
    "```\n",
    "\n",
    "Second, as \"hdf5\" output. This allows to write a subset of the data (and also specify the variables to be written and their precision).\n",
    "Internally they are identical to the \"restart\" outputs though the latter enforce writing data at the precision of the simulation and includes\n",
    "all variables and information to recover the state of a simulation exactly.\n",
    "Thus, they're often also written at higher cadence compared to restart outputs/checkpoint files.\n",
    "Note, that support for writing outputs using the [openPMD](https://github.com/openPMD) standard is currently undergoing final testing and will\n",
    "soon become available (and with it support using ADIOS2 rather than HDF5 for writing data).\n",
    "```\n",
    "<parthenon/output2>\n",
    "file_type  = hdf5       # Binary data dump\n",
    "variables   = prim,acc  # variables to be output\n",
    "dt         = 0.25       # time increment between outputs\n",
    "id         = prim\n",
    "single_precision_output = true\n",
    "```\n",
    "\n",
    "The simulation is now 3D (`nx3 != 1`) and the meshblock size was also adjusted accordingly.\n",
    "```\n",
    "<parthenon/mesh>\n",
    "...\n",
    "nx3        = 128        # Number of zones in X3-direction\n",
    "x3min      = 0.0        # minimum value of X3\n",
    "x3max      = 1.0        # maximum value of X3\n",
    "ix3_bc     = periodic   # inner-X3 boundary flag\n",
    "ox3_bc     = periodic   # outer-X3 boundary flag\n",
    "\n",
    "<parthenon/meshblock>\n",
    "nx1        = 64\n",
    "nx2        = 64\n",
    "nx3        = 64\n",
    "```\n",
    "\n",
    "Finally, the simulation includes magnetic fields (using a generalized Lagrangian multiplier, GLM, divergence cleaning approach) and HLLD Riemann solver,\n",
    "see https://github.com/parthenon-hpc-lab/athenapk/blob/main/docs/input.md for more info.\n",
    "\n",
    "`gamma` is set close to unity to mimick an isothermal equation of state.\n",
    "```\n",
    "<hydro>\n",
    "fluid = glmmhd\n",
    "eos = adiabatic\n",
    "riemann = hlld\n",
    "reconstruction = plm\n",
    "gamma           = 1.0001         # gamma = C_p/C_v\n",
    "```\n",
    "\n",
    "The simulation should be done now.\n",
    "\n",
    "Did you notice sth about the performance?\n",
    "\n",
    "### Plot\n",
    "\n",
    "#### yt\n",
    "\n",
    "Note that for loading the data we now load the `prim` data (which was set by the `id` variable in the output block above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a8082b3-8c79-4e94-88c1-5d1107731299",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://yt-project.org/doc/cookbook/complex_plots.html#time-series-movie\n",
    "ts = yt.load(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/turb_3d_mhd_uniform/parthenon.prim.*.phdf\")\n",
    "\n",
    "plot = yt.SlicePlot(ts[0], \"z\", (\"gas\", \"mach_number\"))\n",
    "plot.set_log((\"gas\", \"mach_number\"), False)\n",
    "plot.set_zlim((\"gas\", \"mach_number\"), 0.0, 2.0)\n",
    "\n",
    "fig = plot.plots[(\"gas\", \"mach_number\")].figure\n",
    "\n",
    "# animate must accept an integer frame number. We use the frame number\n",
    "# to identify which dataset in the time series we want to load\n",
    "def animate(i):\n",
    "    ds = ts[i]\n",
    "    plot._switch_ds(ds)\n",
    "\n",
    "animation = FuncAnimation(fig, animate, frames=len(ts))\n",
    "animation.save(\"turb_3d_mhd_uniform.mp4\")\n",
    "\n",
    "Video(\"turb_3d_mhd_uniform.mp4\", embed=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f09a3bc-7977-422c-a3e1-ffccb02d0a96",
   "metadata": {},
   "source": [
    "#### History output\n",
    "\n",
    "The history output looks like\n",
    "```\n",
    "#  History data\n",
    "# [1]=time     [2]=dt       [3]=cycle    [4]=nbtotal  [5]=mass     [6]=1-mom    [7]=2-mom    [8]=3-mom    [9]=KE       [10]=tot-E    [11]=ME       [12]=relDivB  [13]=Ms       [14]=Ma       [15]=plasma_beta\n",
    "  0.00000e+00  2.34071e-03            0            8  1.00000e+00  0.00000e+00  0.00000e+00  0.00000e+00  0.00000e+00  1.00000e+04  1.25000e-03  0.00000e+00  0.00000e+00  0.00000e+00  8.00000e+02\n",
    "  5.04687e-02  2.25420e-03           22            8  1.00000e+00 -1.52466e-20 -1.04185e-19  1.18585e-20  2.08190e-04  1.00000e+04  1.25002e-03  2.86633e-06  1.89421e-02  3.78858e-01  7.99993e+02\n",
    "```\n",
    "\n",
    "with some default columns (like time, dt, cycle, ...) and additional problem dependent columns (like the mean Mach numbers and plasma beta for the turbulence problem generator).\n",
    "\n",
    "Note, the history file is always appended to (and never overwritten).\n",
    "So lines might be duplicated if restarting from an earlier state of the simulation in the same directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49fd513a-d384-4d63-81c3-77d79202445f",
   "metadata": {},
   "outputs": [],
   "source": [
    "hst = np.genfromtxt(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/turb_3d_mhd_uniform/parthenon.out1.hst\", names=True, skip_header=1)\n",
    "\n",
    "fig, axs = plt.subplots(2,1,sharex=True)\n",
    "\n",
    "axs[0].plot(hst[\"1time\"], hst[\"9KE\"], label=\"kinetic\")\n",
    "axs[0].plot(hst[\"1time\"], hst[\"11ME\"], label=\"magnetic\")\n",
    "axs[0].set_ylabel(\"Mean energy dens.\")\n",
    "axs[0].legend()\n",
    "\n",
    "axs[1].plot(hst[\"1time\"], hst[\"13Ms\"], label=\"sonic\")\n",
    "axs[1].plot(hst[\"1time\"], hst[\"14Ma\"], label=\"Alfvenic\")\n",
    "axs[1].set_ylabel(\"Mean Mach\")\n",
    "axs[1].legend()\n",
    "\n",
    "axs[1].set_xlabel(\"time [code_time]\")\n",
    "fig.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1872ccb-c716-4180-9690-4d9351941144",
   "metadata": {},
   "source": [
    "## Turbulence with static mesh refinement (SMR)\n",
    "\n",
    "In addition to adaptive mesh refinement, a simulation can also use a mesh that is statically refined (thus not requiring resources for checking for refinement, remeshing, load balancing, etc).\n",
    "\n",
    "### Run\n",
    "\n",
    "Again lets first start the simulation and then have a look at the input file while the simulation runs.\n",
    "\n",
    "```bash\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/turb_3d_mhd_smr\n",
    "$ srun -n 1 ../bin/athenaPK -i turbulence.in |& tee ath.out\n",
    "```\n",
    "\n",
    "### Input file\n",
    "\n",
    "What has changed compared to the input file for the uniform grid?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ff88d10-75f0-4e12-abc4-f02275fc6d34",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git diff --no-index --word-diff --color-words /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/turb_3d_mhd_uniform/turbulence.in /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/turb_3d_mhd_smr/turbulence.in"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27a48c98-017f-4e06-b592-517229f369c7",
   "metadata": {},
   "source": [
    "In order to enable static refinemeht `refinement=static` is set in the `<parthenon/mesh>` block.\n",
    "\n",
    "Afterwards, an arbitarary number of static refinement regions can be added to the input file (with a unique integer number suffix)\n",
    "using `<parthenon/static_refinement#>` blocks.\n",
    "The number `4` in the example above is arbitrary.\n",
    "\n",
    "The refinement regions themselves are specified by their extends in the three directions.\n",
    "Note that all blocks that have cells within that specified extends will be refined.\n",
    "\n",
    "Finally, the level of the refinement regions is set by the `level` parameter.\n",
    "\n",
    "Note that cascading effects are likely coming into play if high levels are chosen as the refinement approach in Parthenon enforces\n",
    "only single level jumps between any two adjacent blocks.\n",
    "\n",
    "### Plot\n",
    "\n",
    "Note that the artifacts at the level boundaries in the plot below are just plotting artifacts (given that data between levels is not interpolated for calculating derivatives).\n",
    "In the simulation itself, data is prolongated/restricted in ghost cells so that derivative are calculated correctly also at level boundaries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5682372b-b916-4ac7-86b5-a5a6a7d802bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://yt-project.org/doc/cookbook/complex_plots.html#time-series-movie\n",
    "ts = yt.load(f\"/p/project1/training2437/{user}/tera_day1/athenapk_tutorial/turb_3d_mhd_smr/parthenon.prim.*.phdf\")\n",
    "\n",
    "plot = yt.SlicePlot(ts[0], \"z\", (\"gas\", \"vorticity_magnitude\"), center=[0.25, 0.25, 0.25])\n",
    "plot.set_zlim((\"gas\", \"vorticity_magnitude\"), 0.1, 100)\n",
    "plot.annotate_grids()\n",
    "\n",
    "fig = plot.plots[(\"gas\", \"vorticity_magnitude\")].figure\n",
    "\n",
    "# animate must accept an integer frame number. We use the frame number\n",
    "# to identify which dataset in the time series we want to load\n",
    "def animate(i):\n",
    "    ds = ts[i]\n",
    "    plot._switch_ds(ds)\n",
    "\n",
    "animation = FuncAnimation(fig, animate, frames=len(ts))\n",
    "animation.save(\"turb_3d_mhd_smr.mp4\")\n",
    "\n",
    "Video(\"turb_3d_mhd_smr.mp4\", embed=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d706865-41c5-40f1-963c-e32c7d41c06a",
   "metadata": {},
   "source": [
    "# Step 5: Performance 101\n",
    "\n",
    "To get a high level overview of key performacne characteristics,\n",
    "we'll use a simple linear wave test with a second order MHD (using VL2 integrator, HLLD Riemann solver, and PLM reconstruction\n",
    "known from previous tests).\n",
    "\n",
    "For a quick turnover, we won't look at outputs and just change the mesh, meshblock size (and their logical packing, see below) and take notes on the raw performance.\n",
    "\n",
    "## Increasing mesh size covered by single block\n",
    "\n",
    "```bash\n",
    "\n",
    "$ cd /p/project1/training2437/$USER/tera_day1/athenapk_tutorial/perf\n",
    "\n",
    "$ for MX in 32 64 128 256; do MB=$MX; srun -n 1 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee single.out.$MX.$MB; done\n",
    "```\n",
    "\n",
    "Note that the 256^3 mesh crashed with an `  what():  Kokkos::Impl::ParallelFor< Cuda > insufficient shared memory` error.\n",
    "\n",
    "While the 256^3 mesh still fits on a single A100 GPU on JURECA-DC, using a single 256^3 block does not work\n",
    "because by default AthenaPK explicitly caches data in the flux calculation along the x1 direction in the\n",
    "fast \"shared memory\" (cache).\n",
    "The A100 on JURECA-DC fits approx. little over 128 cells for an MHD problem.\n",
    "\n",
    "Two options exists to remedy this:\n",
    "- (not recommended): disable explicit caching and allocate the the data in device/HBM memory\n",
    "- reduce the nx1 dimension of the block\n",
    "\n",
    "Let's try both\n",
    "\n",
    "```bash\n",
    "# cache in HBM\n",
    "$ MX=256 MB=256; srun -n 1 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB hydro/scratch_level=1 |& tee single.out.$MX.$MB.hbm\n",
    "\n",
    "# cache in shared memory\n",
    "$ MX=256 MB=256; srun -n 1 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=128 parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee single.out.$MX.$MB.cache\n",
    "```\n",
    "\n",
    "And look at all numbers\n",
    "\n",
    "```bash\n",
    "$ grep one-cycles/wallsecond single*\n",
    "```\n",
    "\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "    \n",
    "```\n",
    "single.out.32.32:zone-cycles/wallsecond = 1.48e+07\n",
    "single.out.64.64:zone-cycles/wallsecond = 3.88e+07\n",
    "single.out.128.128:zone-cycles/wallsecond = 9.75e+07\n",
    "single.out.256.256.cache:zone-cycles/wallsecond = 1.62e+08\n",
    "single.out.256.256.hbm:zone-cycles/wallsecond = 1.71e+08\n",
    "```\n",
    "\n",
    "- Performance increased by ~10x when using a larger and larger problem size\n",
    "\n",
    "</details>\n",
    "\n",
    "## Changing the meshblock size\n",
    "\n",
    "Now we stick to the 256^3 mesh and vary the block size:\n",
    "\n",
    "```\n",
    "$ for MB in 32 64 128; do MX=256; srun -n 1 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee fixed256.out.$MX.$MB; done\n",
    "```\n",
    "\n",
    "And look at all numbers\n",
    "\n",
    "```bash\n",
    "$ grep one-cycles/wallsecond fixed256* single.out.256\n",
    "```\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "    \n",
    "```\n",
    "fixed256.out.256.32:zone-cycles/wallsecond = 8.37e+07\n",
    "fixed256.out.256.64:zone-cycles/wallsecond = 1.69e+08\n",
    "fixed256.out.256.128:zone-cycles/wallsecond = 2.07e+08\n",
    "single.out.256.256.cache:zone-cycles/wallsecond = 1.62e+08\n",
    "single.out.256.256.hbm:zone-cycles/wallsecond = 1.71e+08\n",
    "```\n",
    "\n",
    "- Performance peaks for 8 blocks of 128^3\n",
    "- This is tied to the parallelization of the ghost zone packing and communication.\n",
    "  - 0th order answers: this is related to the amount of exposed parallelization. Feel free to ask for more details.\n",
    "\n",
    "</details>\n",
    "\n",
    "## Logical packing of blocks\n",
    "\n",
    "Let's stick to the 256^3 mesh and block size of 64 (so we have a total of 64 blocks).\n",
    "Parthenon transparently handles multiple blocks at the same time (in a single kernel launch)\n",
    "and the number of blocks in pack are controlled by the `parthenon/mesh/pack_size` parameter\n",
    "(which defaults to `-1`, i.e., a single pack for all blocks as it performs best on average on devices.\n",
    "\n",
    "```bash\n",
    "$ for PS in 1 2 4 8 16 32 64; do MX=256 MB=64; srun -n 1 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB parthenon/mesh/pack_size=$PS |& tee pack.out.$MX.$MB.$PS; done\n",
    "```\n",
    "\n",
    "And look at all numbers\n",
    "\n",
    "```bash\n",
    "$ grep one-cycles/wallsecond pack.out*\n",
    "```\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "    \n",
    "```\n",
    "pack.out.256.64.1:zone-cycles/wallsecond = 4.72e+07\n",
    "pack.out.256.64.2:zone-cycles/wallsecond = 7.95e+07\n",
    "pack.out.256.64.4:zone-cycles/wallsecond = 1.23e+08\n",
    "pack.out.256.64.8:zone-cycles/wallsecond = 1.44e+08\n",
    "pack.out.256.64.16:zone-cycles/wallsecond = 1.56e+08\n",
    "pack.out.256.64.32:zone-cycles/wallsecond = 1.66e+08\n",
    "pack.out.256.64.64:zone-cycles/wallsecond = 1.68e+08\n",
    "```\n",
    "\n",
    "- More blocks per pack are faster BUT...\n",
    "- ... this is not always the case -- especially for small block sizes as it also affect communication patters, see https://doi.org/10.1002/cpe.8069 (or https://www.osti.gov/servlets/purl/1993688 for the OA version)\n",
    "- Always benchmark your specific problem setup and the target machine\n",
    "\n",
    "</details>\n",
    "\n",
    "\n",
    "## Running in parallel\n",
    "\n",
    "Nothing (*) needs to be changed, just increase the number of ranks to be used.\n",
    "\n",
    "(*) this is highly machine dependent as it's sometimes (let's say often) not trivial to get the correct CPU cores assigned to the correct GPU etc. Always read the system documentation!\n",
    "\n",
    "### Simple weak scaling\n",
    "\n",
    "We keep the best performing case (256^3 mesh with 128^3 blocks) as basis and run with 1/2/4 ranks (GPUs).\n",
    "\n",
    "Reminder: for a weak scaling test we keep the problem size per compute resource constant.\n",
    "\n",
    "\n",
    "```bash\n",
    "$ MB=128 srun -n 1 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=256 parthenon/mesh/nx2=256 parthenon/mesh/nx3=256 parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee weak.out.1\n",
    "$ MB=128 srun -n 2 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=512 parthenon/mesh/nx2=256 parthenon/mesh/nx3=256 parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee weak.out.2\n",
    "$ MB=128 srun -n 4 ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=512 parthenon/mesh/nx2=512 parthenon/mesh/nx3=256 parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee weak.out.4\n",
    "```\n",
    "\n",
    "And look at all numbers\n",
    "\n",
    "```bash\n",
    "$ grep one-cycles/wallsecond weak.out*\n",
    "```\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "    \n",
    "```\n",
    "weak.out.1:zone-cycles/wallsecond = 2.07e+08\n",
    "weak.out.2:zone-cycles/wallsecond = 4.03e+08\n",
    "weak.out.4:zone-cycles/wallsecond = 7.89e+08\n",
    "```\n",
    "\n",
    "</details>\n",
    "\n",
    "### Simple strong scaling\n",
    "\n",
    "Again, we keep the best performing case (256^3 mesh with 128^3 blocks) and run with 1/2/4 ranks (GPUs).\n",
    "\n",
    "Reminder: for a strong scaling test we keep the problem size constant in increase the computing resources to solve it.\n",
    "\n",
    "```bash\n",
    "$ for NR in 1 2 4; do MX=256 MB=128 srun -n $NR ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=$MB |& tee strong.out.$NR; done\n",
    "\n",
    "# one more test case with the problem decomposed into more blocks\n",
    "$ MX=256 MB=128 NR=4 srun -n $NR ../bin/athenaPK -i linwave.in parthenon/mesh/nx1=$MX parthenon/mesh/nx2=$MX parthenon/mesh/nx3=$MX parthenon/meshblock/nx1=$MB parthenon/meshblock/nx2=$MB parthenon/meshblock/nx3=64 |& tee strong.out.$NR.moreblocks\n",
    "```\n",
    "\n",
    "And look at all numbers\n",
    "\n",
    "```bash\n",
    "$ grep one-cycles/wallsecond strong.out*\n",
    "```\n",
    "<details>    \n",
    "<summary>Click</summary>\n",
    "    \n",
    "```\n",
    "strong.out.1:zone-cycles/wallsecond = 2.07e+08\n",
    "strong.out.2:zone-cycles/wallsecond = 3.55e+08\n",
    "strong.out.4:zone-cycles/wallsecond = 5.72e+08\n",
    "strong.out.4.moreblocks:zone-cycles/wallsecond = 6.44e+08\n",
    "```\n",
    "\n",
    "- As above: always benchmark your specific problem setup and the target machine\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "785ac4ec-d303-4068-82cf-41ecd7f313d4",
   "metadata": {},
   "source": [
    "# Step 6: Code modifications (beyond the tutorial today)\n",
    "\n",
    "Basic instruction on how to add your own problem generator can be found here: https://github.com/parthenon-hpc-lab/athenapk/blob/main/docs/pgen.md\n",
    "\n",
    "Generally, the recommendation is to copy an existing (simple) problem generator (like https://github.com/parthenon-hpc-lab/athenapk/blob/main/src/pgen/orszag_tang.cpp) and start from there.\n",
    "\n",
    "Finally (and probably most importantly): Feel free to ask questions either by\n",
    "- create a discussion (https://github.com/parthenon-hpc-lab/athenapk/discussions) or open an issue (https://github.com/parthenon-hpc-lab/athenapk/issues)\n",
    "- join the chat: https://app.element.io/#/room/#AthenaPK:matrix.org"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "216bfd21-dfd0-4276-ba4c-0e08758db704",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}