{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fa561492",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img src=\"pic/ex06.png\" width=\"100%\"> <br>\n",
    "\n",
    "<h1 class=\"jp-toc-ignore\">ICON Training - Hands-on Session</h1>\n",
    "\n",
    "# Exercise 6.1: ICON ComIn, Part I\n",
    "\n",
    "---\n",
    "The **ICON Community Interface (ComIn)** enables the integration of **plugins** which is a flexible way to extend the code with your own routines and variables at run-time.  \n",
    "ComIn offers the opportunity to write your plugins in various programming languages, including Fortran, C, and Python. In this training course, we will focus on developing **Python plugins**.\n",
    "* **Step P1**: We repeat the previous task \"Step F1\", this time using only a few lines of Python code.\n",
    "* **Step P2**: We learn how to perform data MPI gathering.\n",
    "* **Step P3**: In the final step, we learn how to use the coupling library [YAC](https://dkrz-sw.gitlab-pages.dkrz.de/yac/) in a plugin to interpolate fields from the ICON unstructured grid onto a custom lat-lon grid, and then visualize the interpolated fields.\n",
    "\n",
    "To complete the following exercises, it is important to have a basic understanding of programming in Python, particularly the correct use of indentation.\n",
    "\n",
    "<br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e61c91f4-91e2-4e24-a168-049ee2687e33",
   "metadata": {},
   "source": [
    "<div style=\"width:100%;height:120px;border-top: 2px solid darkgray;border-bottom: 2px solid darkgray;padding: 10px 0px 10px 0px\">\n",
    "<img src=\"pic/logo_comin.png\" style=\"float:right;height:120px\">\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67845954-bf34-4284-8e1e-65b72a16efea",
   "metadata": {},
   "source": [
    "## Preliminary Step: Building the ICON ComIn Python Adapter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dcff6a6-f074-464e-9eee-de0b71f5c2e2",
   "metadata": {},
   "source": [
    "We will use a prebuilt ICON binary in `ICONDIR=/pool/data/ICON/ICON_training/icon`. As you are writing a Python plugin in this exercise, there is **no** need to compile the ICON code at any stage to complete the task.  \n",
    "\n",
    "For **educational purposes**, it is demonstrated below how to build ICON in a way that allows it to run ComIn plugins.\n",
    "We outline the steps for creating an out-of-source build (recommended)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "648a0fa0-d9a1-4aa9-9336-36b84a997626",
   "metadata": {
    "tags": []
   },
   "source": [
    "```bash\n",
    "mkdir -p icon/build\n",
    "cd icon/build\n",
    "../config/dkrz/levante.gcc \\\n",
    "  --enable-comin --enable-bundled-python=comin\n",
    "make -j6"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ad1a1a4-f7bf-4835-bc65-ce1c220d49a5",
   "metadata": {},
   "source": [
    "**Remarks**\n",
    "- The `--enable-bundled-python=comin` configuration option is used to build the ComIn Python adapter. This adapter is a dynamic library and can be found under `/build/externals/comin/build/plugins/python_adapter/libpython_adapter.so`\n",
    "- This is the minimal configuration setting for ComIn. For the full, `YAC` enabled version, please refer back to the <a href=\"../exercise_programming/icon_exercise_programming.ipynb\">explanation provided for F1</a>."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65bd851e-3bee-4d3b-9e0a-d9fe9eae0224",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Step P1: Writing a ComIn Python Plugin\n",
    "\n",
    "For the necessary information on ComIn, please refer to Section 9.5 of the ICON tutorial. In addition, the ComIn project website contains more detailed documentation. The ICON documentation website contains a section on ComIn featuring the documentation for the latest release version. All these references can be found at the end of this Jupyter notebook.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a86b415-8a83-4a1e-a911-58d66d7800c4",
   "metadata": {},
   "source": [
    "### Your First ComIn Plugin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d647f9d-21e3-4bae-a95a-3ab9e5b73f59",
   "metadata": {},
   "source": [
    "In this step, rather than implementing your programming directly into the ICON source code, you will develop a **ComIn plugin** to perform a task similar to **Step F1** in the ICON programming exercise. In other words, we will visualise ICON's domain decomposition."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6ef2ec0-3e0b-41ef-a1a9-bfcd67cb0e65",
   "metadata": {},
   "source": [
    "<figure >\n",
    "<img src=\"pic/EP.svg\"  width=\"100%\">\n",
    "</figure>\n",
    "This figure provides a schematic overview of the available entry points in the ICON source code: The different loops of the time stepping scheme (dynamics, fast and slow physics) can be supplemented by attached plugin functions, with the locations marked by the red arrows.\n",
    "\n",
    "<br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06213288-62a3-442f-8ce4-a5733665b38a",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "  <b style=\"color:#2d4b9b;\">Exercise:</b> \n",
    "  We will write a Python script (ComIn plugin). To this end, start from an (almost) empty <code>scripts/comin_plugin_P1.py</code>, which you can open by clicking <a href=\"scripts/comin_plugin_P1.py\">here</a>.\n",
    "  In your plugin, add a 2D diagnostic variable named <code>comin_process_id</code> for domain 1.\n",
    "  \n",
    " As outlined in the ICON tutorial, adding a new variable in your ComIn plugin involves two steps:\n",
    " <ul>\n",
    "  <li>First you need to <b>register a new variable</b> and set the required metadata for the new variable. <br/>\n",
    "      You will need the library command <code>comin.var_request_add</code> for this.\n",
    "      For a list of available metadata, refer to the <a href=\"https://icon-comin.gitlab-pages.dkrz.de/comin/d7/de6/md__2builds_2icon-comin_2comin_2doc_2comin__python__api.html\">Python API</a>.</li>\n",
    "    \n",
    "  <li>Secondly, you should <b>access your registered variable</b> like any other ICON variable within the secondary constructor of your plugin. <br/>\n",
    "      This happens with the API function <code>comin.var_get</code>.\n",
    "      Furthermore, you will need to specify the entry point for your callback function. In this exercise, we will use the <code>EP_ATM_WRITE_OUTPUT_BEFORE</code> entry point. </li>   \n",
    "  </ul>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d99c37a-d09d-4a2f-b11a-e25d7db1e8a2",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "source": [
    "<details>\n",
    "<summary style=\"font-size: 1.1em;\">Solution</summary><br/>\n",
    "\n",
    "```python\n",
    "import comin\n",
    "\n",
    "jg = 1  # set the domain id\n",
    "\n",
    "# request to register the variable\n",
    "var_descriptor = (\"comin_process_id\", jg)\n",
    "comin.var_request_add(var_descriptor, lmodexclusive=False)\n",
    "comin.metadata_set(var_descriptor, zaxis_id = comin.COMIN_ZAXIS_2D)\n",
    "\n",
    "\n",
    "@comin.EP_SECONDARY_CONSTRUCTOR\n",
    "def simple_python_constructor():\n",
    "    global comin_process_id\n",
    "    comin_process_id = comin.var_get([comin.EP_ATM_WRITE_OUTPUT_BEFORE], \n",
    "                                     (\"comin_process_id\", jg), \n",
    "                                     flag=comin.COMIN_FLAG_WRITE\n",
    "                                    )\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ccbf1fc-e842-4b41-9bd9-45c929d3b8a1",
   "metadata": {},
   "source": [
    "### Filling the data array with the process ID "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a26cad3-513d-4d51-9cbf-ed80cf5f044d",
   "metadata": {},
   "source": [
    "<div class=\"alert alert-success\">\n",
    "  <b style=\"color:#2d4b9b;\">Exercise:</b> \n",
    "Write a callback function in your plugin <a href=\"scripts/comin_plugin_P1.py\"><code>scripts/comin_plugin_P1.py</code></a> and attach it to the <code>EP_ATM_WRITE_OUTPUT_BEFORE</code> entry point. \n",
    "This is done using the <code>comin.register_callback</code> function (a function decorator).\n",
    "<p/>\n",
    "In the callback function, fill the variable <code>comin_process_id</code> with the MPI rank of the processor. \n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "360b6741-b89a-4aea-bfd9-1bb08d122dce",
   "metadata": {},
   "source": [
    "\n",
    "**Hints:**\n",
    "* In the callback function, import the NumPy package and use `numpy.asarray` to convert a variable handle you accessed into a NumPy array.\n",
    "* Refer to the Python API description for ICON ComIn to learn how to obtain the MPI rank of the ICON model in your plugin. A link to the API description is provided in the References section at the end of this Jupyter notebook.<br/>\n",
    "  You should use ComIn's so-called **host MPI communicator** which is the MPI communicator that comprises all MPI tasks of the ICON simulation which are involved in the ComIn callbacks.\n",
    "* If you want to add a log message to your plugin script, you should use <code>print(\"...\", file=sys.stderr)</code> to avoid buffering.\n",
    "\n",
    "\n",
    "<div class=\"alert alert-info\">\n",
    "<b>Remark:</b> Actually, the use of <code>np.asarray</code> is an important point to note: The `comin.var_get` call returns a handle to an ICON variable. However, since ICON contains fields with multiple time levels, the corresponding array pointer changes during the model run. Each time the plugin needs to access the array (pointer), it must be explicitly or implicitly converted to a NumPy array.\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18f7dad6-81b7-4e84-9ad8-19eab520db13",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "hide-input"
    ]
   },
   "source": [
    "<details>\n",
    "<summary style=\"font-size: 1.1em;\">Solution</summary><br/>\n",
    "\n",
    "```python\n",
    "@comin.EP_ATM_WRITE_OUTPUT_BEFORE\n",
    "def simple_python_callbackfct():\n",
    "    import numpy as np\n",
    "    # print(\"simple_python_callbackfct called!\", file=sys.stderr)\n",
    "    comin_process_id_np = np.asarray(comin_process_id)\n",
    "    comin_process_id_np[:] = comin.parallel_get_host_mpi_rank()\n",
    "```\n",
    "\n",
    "</details>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94e45b7c-2601-402b-a593-fec2a4d16b99",
   "metadata": {},
   "source": [
    "### Running the ICON Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea40314f-1ff2-46b1-ada9-dd3a29aca4ee",
   "metadata": {},
   "source": [
    "Now that you've finished writing your plugin, the next step is to run the ICON model.\n",
    "First, execute the following cell in order to set up an example ICON job, similar to the ICON-LAM setting of the previous exercise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ffaab52-08dc-4f9b-aa21-fd9eabf99c2a",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "\n",
    "user = os.environ['USER']\n",
    "home = os.environ['HOME']\n",
    "scratchdir = f\"/scratch/{user[0]}/{user}\"\n",
    "icondir = f\"/pool/data/ICON/ICON_training/icon\"\n",
    "expdir = f\"{scratchdir}/exercise_comin/P1\"\n",
    "%env ICONDIR={icondir}\n",
    "%env EXPDIR={expdir}\n",
    "%env SCRATCHDIR={scratchdir}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "205322e2-ea71-4097-9b97-799bdb9e9f94",
   "metadata": {},
   "source": [
    "Run setup script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "784a77a1-3f62-403e-a33b-909e277e95e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash $HOME/icon-training-scripts/exercise_comin/prepared/prepare_icon_run.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ade63f3-6581-4831-b113-60e0d0e42e23",
   "metadata": {},
   "source": [
    "Plugins are enabled via a special namelist `comin_nml`.\n",
    "In order to load your plugin dynamically while ICON is running, add a `comin_nml` namelist to the ICON namelist file. To this end, execute the cell below. We set the `plugin_library = libpython_adapter.so` with its absolute path such that it can be found by ICON and the operating system.\n",
    "\n",
    "*Note: Alternatively, the search path for `libpython_adapter.so` could be set through the environment variable `LD_LIBRARY_PATH`.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9b65bda-7627-44cf-9f0f-d821163759f3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Append comin_nml block\n",
    "with open(f\"{expdir}/NAMELIST_ICON\", 'a') as f:\n",
    "    f.write(f\"\"\"&comin_nml\n",
    "   plugin_list(1)%name           = \"comin_plugin\"\n",
    "   plugin_list(1)%plugin_library = \"{icondir}/build/externals/comin/build/plugins/python_adapter/libpython_adapter.so\"\n",
    "   plugin_list(1)%options        = \"{home}/icon-training-scripts/exercise_comin/scripts/comin_plugin_P1.py\"\n",
    "/\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "007c293d-6f49-4b5f-899a-3dc3080ac11e",
   "metadata": {},
   "source": [
    "We must customize `NAMELIST_ICON` to ensure that `comin_process_id` is included in the model output files for the first domain. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c22b135-0a76-448f-89c6-9d24255ff53a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Append comin_nml block\n",
    "with open(f\"{expdir}/NAMELIST_ICON\", 'a') as f:\n",
    "    f.write(f\"\"\"&output_nml\n",
    " filetype                    = 4                     ! netcdf\n",
    " dom                         = 1\n",
    " output_bounds               = 0., 10000000., 3600.  ! start, end, increment\n",
    " steps_per_file              = 1\n",
    " mode                        = 1\n",
    " include_last                = .FALSE.\n",
    " steps_per_file_inclfirst    = .FALSE.\n",
    " output_filename             = 'NWP_LAM'\n",
    " filename_format             = '<output_filename>_DOM<physdom>_<datetime2>'\n",
    " output_grid                 = .FALSE.\n",
    " remap                       = 1                     ! 1: remap to lat-lon grid\n",
    " reg_lon_def                 = 0.8,0.1,17.2\n",
    " reg_lat_def                 = 43.9,0.1,57.7\n",
    " ml_varlist                  = \"comin_process_id\"\n",
    "/\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f43a127c-a607-41de-823c-0e83c9b4867b",
   "metadata": {},
   "source": [
    "The other parts of the Fortran namelist remain unchanged. In the run script, we set up the environment variables and create symbolic links to the necessary input data required for running the ICON experiment in a directory with plenty of space (`scratch`).\n",
    "\n",
    "Now you can run the ICON model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0248e2d1-b5da-45db-9737-627e01f234f1",
   "metadata": {
    "editable": true,
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd $EXPDIR && sbatch --account=$SLURM_JOB_ACCOUNT --export=ICONDIR $EXPDIR/icon-lam.sbatch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "807d07b4-5193-41c0-8826-2ed9c3edac4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "!squeue -u $USER"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e26e7bd0-9907-4618-acf8-dc6cc103ee28",
   "metadata": {},
   "source": [
    "<table border=\"0\" cellspacing=\"0\" cellpadding=\"0\" style=\"width:100%; border-collapse:collapse;\">\n",
    "  <tr>\n",
    "    <td style=\"vertical-align:top; padding-right:24px;\">\n",
    "      <div class=\"alert alert-success\">\n",
    "        <b style=\"color:#2d4b9b;\">Exercise:</b> \n",
    "        <ul>\n",
    "          <li>\n",
    "            Use the <a href=\"../exercise_programming/scripts/icon_exercise_programming_plot01.ipynb\">plot script</a> from Step F1 to visualize <code>comin_process_id</code>.<br/>\n",
    "            Your result should look like the plot on the right.\n",
    "          </li>\n",
    "        </ul>\n",
    "      </div>\n",
    "    </td>\n",
    "    <td style=\"vertical-align:top; text-align:right;\">\n",
    "      <figure>\n",
    "        <img src=\"reference/P1_comin_plugin/domain_decomp_p1.png\" width=\"600\" style=\"display:block;\">\n",
    "        <figcaption style=\"text-align:center\"><em>Reference solution for \"Step P1: Writing a ComIn Python Plugin\".</em></figcaption>\n",
    "      </figure>\n",
    "    </td>\n",
    "  </tr>\n",
    "</table>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3099312b-470c-4f48-93e8-f00a235f59dc",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55ab88a8-1a48-4b83-a995-b46066877388",
   "metadata": {
    "tags": []
   },
   "source": [
    "<h2 class=\"jp-toc-ignore\"> <b style=\"color:#2d4b9b;\"><b style=\"color:#2d4b9b;\"> Congratulations! You have successfully completed Exercise 6, Part I.</b></h2>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f5d5129",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23b40d12-04e0-443c-b1d2-da0950692c70",
   "metadata": {},
   "source": [
    "## Further Reading and Resources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4541c17e-5548-479c-8bcb-026741e02d07",
   "metadata": {},
   "source": [
    "<ul>\n",
    "<li>\n",
    "  ICON Tutorial, Ch. 9: <a href=\"https://www.dwd.de/DE/leistungen/nwv_icon_tutorial/nwv_icon_tutorial.html\">https://www.dwd.de/DE/leistungen/nwv_icon_tutorial/nwv_icon_tutorial.html</a>  \n",
    "    <br/>A new <b>draft version</b> of the ICON Tutorial is available here: <a href=\"https://icon-training-2025-scripts-rendering-cc74a6.gitlab-pages.dkrz.de/index.html\">https://icon-training-2025-scripts-rendering-cc74a6.gitlab-pages.dkrz.de/index.html</a>. It is currently being finalized and will be published soon.\n",
    "</li>\n",
    "<li>\n",
    "  Documentation on the ICON ComIn project website: <a href=\"https://gitlab.dkrz.de/icon-comin/comin/-/blob/master/doc/icon_comin_doc.md?ref_type=heads\">https://gitlab.dkrz.de/icon-comin/comin/-/blob/master/doc/icon_comin_doc.md?ref_type=heads</a>  \n",
    "</li>\n",
    "<li>\n",
    "  Python API description for ICON ComIn: <a href=\"https://gitlab.dkrz.de/icon-comin/comin/-/blob/master/doc/comin_python_api.md?ref_type=head\">https://gitlab.dkrz.de/icon-comin/comin/-/blob/master/doc/comin_python_api.md?ref_type=head</a>  \n",
    "</li>\n",
    "<li> \n",
    " ICON Documentaton website: Information on ComIn, including the current user guide:\n",
    " <a href=\"https://docs.icon-model.org/tools/comin/comin.html\">https://docs.icon-model.org/tools/comin/comin.html</a>\n",
    "</li>\n",
    "</ul>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7923d75d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "*Author info: Deutscher Wetterdienst (DWD) 2025 ::  icon@dwd.de. For a full list of contributors, see CONTRIBUTING in the root directory. License info: see LICENSE file.*"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "0 Python 3 (based on the module python3/unstable",
   "language": "python",
   "name": "python3_unstable"
  },
  "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.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}