example/01_reading_and_writing_files.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "40786f36",
   "metadata": {},
   "source": [
    "# Reading and writing files with Loki\n",
    "\n",
    "This is the first introductory notebook on how to work with Loki. The intention is to give an overview of how Fortran files can be read into Loki's internal representation to be able to work on their content and apply transformations, and how we can generate Fortran source code again. It includes a short peak at the control flow representation but details will be discussed in other notebooks.\n",
    "\n",
    "Let's start by parsing the file `src/phys_mod.F90` from the `example` directory.\n",
    "Loki uses a [_Sourcefile_](https://sites.ecmwf.int/docs/loki/main/loki.sourcefile.html#module-loki.sourcefile) object to represent an entire source file, which can contain modules or subroutines. To initialize the object with the content of a file on disc, we use the `from_file` class method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "fa7c571e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Loki::Sourcefile] Constructed from src/phys_mod.F90 in 7.46s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<loki.sourcefile.Sourcefile at 0x14c39b54a3d0>"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from loki import Sourcefile\n",
    "source = Sourcefile.from_file('src/phys_mod.F90', preprocess=True)\n",
    "source"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "836efe96",
   "metadata": {},
   "source": [
    "Let's examine the content of the source file by looking at the modules and subroutines contained in that file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c262254",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Modules: (Module:: phys_mod,)\n",
      "Subroutines: ()\n"
     ]
    }
   ],
   "source": [
    "print(f\"Modules: {source.modules}\")\n",
    "print(f\"Subroutines: {source.subroutines}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0a6d4aa",
   "metadata": {},
   "source": [
    "We can see from the above that `source` contains one module by the name \"phys_mod\" and no free subroutines.\n",
    "We can access modules either via their index in the `modules` property (i.e., `source.modules[0]`) or using a subscript operator with their name directly on the `Sourcefile` object:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5c7ba4af",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Module:: phys_mod"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "phys_mod = source['phys_mod']\n",
    "phys_mod"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c4518e",
   "metadata": {},
   "source": [
    "Fortran modules are represented as [_Module_](https://sites.ecmwf.int/docs/loki/main/loki.module.html#loki.module.Module) objects in Loki. They consist of a specification part and may contain, e.g., subroutines. Let's examine this object further:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6b98dc7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Spec: Section::\n",
      "Subroutines: (Subroutine:: phys_kernel_LITE_LOOP, Subroutine:: phys_kernel_VERT_SEARCH, Subroutine:: phys_kernel_NASTY_EXPS, Subroutine:: phys_kernel_LU_SOLVER, Subroutine:: phys_kernel_LU_SOLVER_COMPACT)\n"
     ]
    }
   ],
   "source": [
    "print(f\"Spec: {phys_mod.spec}\")\n",
    "print(f\"Subroutines: {phys_mod.subroutines}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10c81ec9",
   "metadata": {},
   "source": [
    "The specification part consists of a [_Section_](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Section) node, which acts as the root node of Loki's control flow tree. At this point, it may be useful to learn more about Loki's internal representation by reading the [relevant part](https://sites.ecmwf.int/docs/loki/main/internal_representation.html) of the documentation. But for the objectives of this notebook we can also carry on and treat them as a black box for now.\n",
    "\n",
    "To get an impression of what the IR of the specification part looks like, we can call `view()` on any of the nodes to print a representation of this node and the tree below it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a6159d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Section::>\n",
      "  <Import:: iso_fortran_env => ()>\n",
      "  <Comment:: >\n",
      "  <Import:: omp_lib => ()>\n",
      "  <Comment:: >\n",
      "  <Intrinsic:: IMPLICIT NONE>\n",
      "  <Comment:: >\n",
      "  <VariableDeclaration:: sp>\n",
      "  <VariableDeclaration:: dp>\n",
      "  <CommentBlock:: >\n",
      "  <VariableDeclaration:: lp>\n",
      "  <CommentBlock:: >\n",
      "  <VariableDeclaration:: ip>\n",
      "  <Comment:: >\n",
      "  <VariableDeclaration:: cst1, cst2>\n",
      "  <VariableDeclaration:: nspecies>\n",
      "  <Comment:: >"
     ]
    }
   ],
   "source": [
    "phys_mod.spec.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7fa4ba7-a0f6-4f10-a47f-d508a121653d",
   "metadata": {},
   "source": [
    "Or alternativly, if `graphviz` is available, we can call `ir_graph()` on any of the nodes to view a graph representation of this node and the tree below it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6942fbb4-113c-466d-be0e-4fce35d837ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Loki::Graph Visualization] Created graph visualization in 0.01s\n"
     ]
    },
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 2.43.0 (0)\n",
       " -->\n",
       "<!-- Title: %3 Pages: 1 -->\n",
       "<svg width=\"521pt\" height=\"476pt\"\n",
       " viewBox=\"0.00 0.00 521.06 476.00\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 472)\">\n",
       "<title>%3</title>\n",
       "<polygon fill=\"white\" stroke=\"transparent\" points=\"-4,4 -4,-472 517.06,-472 517.06,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"66.94\" cy=\"-234\" rx=\"66.89\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"66.94\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;Section::&gt;</text>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>1</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-450\" rx=\"166.27\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-446.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;Import:: iso_fortran_env =&gt; ()&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;1 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>0&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M72,-252.14C81.61,-290.6 109.66,-380.32 169.89,-423 177.18,-428.16 185.08,-432.49 193.35,-436.11\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"192.25,-439.44 202.84,-439.9 194.85,-432.94 192.25,-439.44\"/>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>2</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-396\" rx=\"131.08\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-392.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;Import:: omp_lib =&gt; ()&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;2 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>0&#45;&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M75.91,-251.98C90.11,-281.73 122.8,-340.74 169.89,-369 182.32,-376.46 196.14,-382.06 210.37,-386.25\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"209.72,-389.7 220.29,-388.92 211.55,-382.94 209.72,-389.7\"/>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>3</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-342\" rx=\"151.37\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-338.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;Intrinsic:: IMPLICIT NONE&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;3 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>0&#45;&gt;3</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M84.49,-251.74C103.26,-270.68 135.74,-299.99 169.89,-315 182.82,-320.69 196.78,-325.19 210.94,-328.77\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"210.22,-332.19 220.76,-331.09 211.83,-325.38 210.22,-332.19\"/>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-288\" rx=\"137.58\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-284.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;VariableDeclaration:: sp&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;4 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>0&#45;&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M113.38,-247.09C130.94,-251.88 151.26,-257.07 169.89,-261 191.39,-265.54 214.55,-269.69 236.67,-273.3\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"236.32,-276.79 246.75,-274.93 237.44,-269.88 236.32,-276.79\"/>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>5</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-234\" rx=\"138.38\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-230.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;VariableDeclaration:: dp&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;5 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>0&#45;&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M133.92,-234C151.92,-234 172.18,-234 192.79,-234\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"192.93,-237.5 202.93,-234 192.93,-230.5 192.93,-237.5\"/>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>6</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-180\" rx=\"135.68\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-176.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;VariableDeclaration:: lp&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;6 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>0&#45;&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M113.38,-220.91C130.94,-216.12 151.26,-210.93 169.89,-207 191.51,-202.44 214.83,-198.26 237.07,-194.63\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"237.88,-198.05 247.2,-193 236.77,-191.14 237.88,-198.05\"/>\n",
       "</g>\n",
       "<!-- 7 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>7</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-126\" rx=\"135.68\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-122.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;VariableDeclaration:: ip&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;7 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>0&#45;&gt;7</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M84.49,-216.26C103.26,-197.32 135.74,-168.01 169.89,-153 184.55,-146.55 200.52,-141.62 216.61,-137.85\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"217.62,-141.22 226.64,-135.66 216.12,-134.38 217.62,-141.22\"/>\n",
       "</g>\n",
       "<!-- 8 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>8</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-72\" rx=\"171.67\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-68.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;VariableDeclaration:: cst1, cst2&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;8 -->\n",
       "<g id=\"edge8\" class=\"edge\">\n",
       "<title>0&#45;&gt;8</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M75.91,-216.02C90.11,-186.27 122.8,-127.26 169.89,-99 178.09,-94.08 186.9,-89.96 196.04,-86.53\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"197.37,-89.77 205.67,-83.19 195.07,-83.16 197.37,-89.77\"/>\n",
       "</g>\n",
       "<!-- 9 -->\n",
       "<g id=\"node10\" class=\"node\">\n",
       "<title>9</title>\n",
       "<ellipse fill=\"none\" stroke=\"black\" cx=\"341.47\" cy=\"-18\" rx=\"167.07\" ry=\"18\"/>\n",
       "<text text-anchor=\"middle\" x=\"341.47\" y=\"-14.3\" font-family=\"Times,serif\" font-size=\"14.00\">&lt;VariableDeclaration:: nspecies&gt;</text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;9 -->\n",
       "<g id=\"edge9\" class=\"edge\">\n",
       "<title>0&#45;&gt;9</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M72,-215.86C81.61,-177.4 109.66,-87.68 169.89,-45 177.18,-39.84 185.08,-35.51 193.35,-31.89\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"194.85,-35.06 202.84,-28.1 192.25,-28.56 194.85,-35.06\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x7fc53b1c0550>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "graph = None\n",
    "try:\n",
    "    graph = phys_mod.spec.ir_graph()\n",
    "except ImportError:\n",
    "    print(\"Install graphviz if you want to view the graph representation!\")\n",
    "graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80f46c51",
   "metadata": {},
   "source": [
    "We can see a number of (empty) [comments](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Comment) - which are simply empty lines and retained to be able to produce Fortran code with a formatting similar to the original source. Since [comments](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Comment) might introduce additional noise, they are ignored by default in the graph representation. Other than that, we also have some [_Import_](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Import) statements, [preprocessor directives](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.PreprocessorDirective) and [declarations](https://sites.ecmwf.int/docs/loki/main/loki.ir.html#loki.ir.Declaration).\n",
    "\n",
    "We can also convert this representation of the specification part back into a Fortran representation using the Fortran backend via [_fgen_](https://sites.ecmwf.int/docs/loki/main/loki.backend.fgen.html):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d0ccd9c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "USE iso_fortran_env\n",
      "\n",
      "USE omp_lib\n",
      "\n",
      "IMPLICIT NONE\n",
      "\n",
      "INTEGER, PARAMETER :: sp = REAL32\n",
      "INTEGER, PARAMETER :: dp = REAL64\n",
      "\n",
      "\n",
      "\n",
      "INTEGER, PARAMETER :: lp = dp!! lp : \"local\" precision\n",
      "\n",
      "\n",
      "INTEGER, PARAMETER :: ip = INT64\n",
      "\n",
      "REAL(KIND=lp) :: cst1 = 2.5, cst2 = 3.14\n",
      "INTEGER, PARAMETER :: nspecies = 5\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from loki import fgen\n",
    "print(fgen(phys_mod.spec))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e00f560e",
   "metadata": {},
   "source": [
    "When comparing the Fortran code to the above internal representation makes it easy to identify the one-to-one correlation between IR nodes and statements in the original source code.\n",
    "\n",
    "Let's pick out one of the kernel loops next:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d06624b9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Subroutine:: phys_kernel_LITE_LOOP"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lite_loop = phys_mod['phys_kernel_LITE_LOOP']\n",
    "lite_loop"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7aac3d17",
   "metadata": {},
   "source": [
    "Subroutines and functions are represented as a [_Subroutine_](https://sites.ecmwf.int/docs/loki/main/loki.subroutine.html#loki.subroutine.Subroutine) object. This allows, for example, to inspect the names of the dummy arguments expected by this routine:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7e66ac24",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['dim1',\n",
       " 'dim2',\n",
       " 'i1',\n",
       " 'i2',\n",
       " 'in1',\n",
       " 'in2',\n",
       " 'in3',\n",
       " 'in4',\n",
       " 'in5',\n",
       " 'in6',\n",
       " 'in7',\n",
       " 'in8',\n",
       " 'in9',\n",
       " 'in10',\n",
       " 'out1']"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "lite_loop.argnames"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb0b6061",
   "metadata": {},
   "source": [
    "Furthermore, all subroutines contain a specification and body part (either of which can of course be empty in principal):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4a1850b5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Section::>\n",
      "  <VariableDeclaration:: dim1, dim2, i1, i2>\n",
      "  <VariableDeclaration:: in1(1:dim1, 1:dim2), in2(1:dim1, 1:dim2), in3(1:dim1, 1:dim2), \n",
      "  in4(1:dim1, 1:dim2), in5(1:dim1, 1:dim2), in6(1:dim1, 1:dim2), in7(1:dim1, 1:dim2), \n",
      "  in8(1:dim1, 1:dim2), in9(1:dim1, 1:dim2), in10(1:dim1, 1:dim2)>\n",
      "  <VariableDeclaration:: out1(1:dim1, 1:dim2)>\n",
      "  <Comment:: >\n",
      "  <VariableDeclaration:: i, k>"
     ]
    }
   ],
   "source": [
    "lite_loop.spec.view()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "6dad7303",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Section::>\n",
      "  <Loop:: k=1:dim2>\n",
      "    <Loop:: i=i1:i2>\n",
      "      <Assignment:: out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, \n",
      "      k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) + in10(i, k))*0.1>\n",
      "      <Assignment:: in1(i, k) = out1(i, k)>"
     ]
    }
   ],
   "source": [
    "lite_loop.body.view()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67ad8799",
   "metadata": {},
   "source": [
    "As we can see from the above, this kernel accepts a large number of arguments and consists essentially of two nested loops. Instead of viewing abstract representation, we can also produce Fortran source code again, either by calling `fgen` for individual parts or the entire `Subroutine` object, or, in this case, we can also use a convenience API offered by the object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e5aba927",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SUBROUTINE phys_kernel_LITE_LOOP (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "  INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "  REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "  \n",
      "  INTEGER(KIND=ip) :: i, k\n",
      "  DO k=1,dim2\n",
      "    DO i=i1,i2\n",
      "      out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) +  &\n",
      "      & in10(i, k))*0.1\n",
      "      in1(i, k) = out1(i, k)\n",
      "    END DO\n",
      "  END DO\n",
      "END SUBROUTINE phys_kernel_LITE_LOOP\n"
     ]
    }
   ],
   "source": [
    "print(lite_loop.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "773a5f1f",
   "metadata": {},
   "source": [
    "In this notebook, we will not go into detail on how to actually modify the control flow tree of this routine. But we will extract this routine from the module and put it into a standalone module.\n",
    "\n",
    "Let's start by creating a clone of this routine with a new name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c3343bd9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Subroutine:: my_routine"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "my_routine = lite_loop.clone(name='my_routine')\n",
    "my_routine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f00603d",
   "metadata": {},
   "source": [
    "Next, we create a new module and insert `my_routine` as a subroutine. To make sure the relevant declarations from the original module are available, we create a copy of the relevant spec:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "59e719b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Module:: my_module"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from loki import Module\n",
    "my_module = Module(name='my_module', spec=phys_mod.spec.clone(), contains=(my_routine,))\n",
    "my_module"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b368f8b2",
   "metadata": {},
   "source": [
    "Let's ensure the new module contains `my_routine`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "85448ed7",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert len(my_module.subroutines) == 1\n",
    "assert my_module.subroutines[0] is my_routine"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "571c62a4",
   "metadata": {},
   "source": [
    "We can also take a look at the Fortran code of this new module:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0f14e7ba",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MODULE my_module\n",
      "  USE iso_fortran_env\n",
      "  \n",
      "  USE omp_lib\n",
      "  \n",
      "  IMPLICIT NONE\n",
      "  \n",
      "  INTEGER, PARAMETER :: sp = REAL32\n",
      "  INTEGER, PARAMETER :: dp = REAL64\n",
      "  \n",
      "  \n",
      "  \n",
      "  INTEGER, PARAMETER :: lp = dp  !! lp : \"local\" precision\n",
      "  \n",
      "  \n",
      "  INTEGER, PARAMETER :: ip = INT64\n",
      "  \n",
      "  REAL(KIND=lp) :: cst1 = 2.5, cst2 = 3.14\n",
      "  INTEGER, PARAMETER :: nspecies = 5\n",
      "  \n",
      "  CONTAINS\n",
      "  SUBROUTINE my_routine (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "    INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "    REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "    REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "    \n",
      "    INTEGER(KIND=ip) :: i, k\n",
      "    DO k=1,dim2\n",
      "      DO i=i1,i2\n",
      "        out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) + &\n",
      "        &  in10(i, k))*0.1\n",
      "        in1(i, k) = out1(i, k)\n",
      "      END DO\n",
      "    END DO\n",
      "  END SUBROUTINE my_routine\n",
      "END MODULE my_module\n"
     ]
    }
   ],
   "source": [
    "print(my_module.to_fortran())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6fe01e4c",
   "metadata": {},
   "source": [
    "And, ultimately, we can write this to a separate source file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3ff6dcd8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Loki::Sourcefile] Writing to my_module.F90\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "Sourcefile.to_file(fgen(my_module), Path('my_module.F90'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34e235dc",
   "metadata": {},
   "source": [
    "Finally, let's take a peek at the generated file (disregard the pylint comment, which is there only for technical reasons related to our automated testing):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "acf60783",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MODULE my_module\n",
      "  USE iso_fortran_env\n",
      "  \n",
      "  USE omp_lib\n",
      "  \n",
      "  IMPLICIT NONE\n",
      "  \n",
      "  INTEGER, PARAMETER :: sp = REAL32\n",
      "  INTEGER, PARAMETER :: dp = REAL64\n",
      "  \n",
      "  \n",
      "  \n",
      "  INTEGER, PARAMETER :: lp = dp  !! lp : \"local\" precision\n",
      "  \n",
      "  \n",
      "  INTEGER, PARAMETER :: ip = INT64\n",
      "  \n",
      "  REAL(KIND=lp) :: cst1 = 2.5, cst2 = 3.14\n",
      "  INTEGER, PARAMETER :: nspecies = 5\n",
      "  \n",
      "  CONTAINS\n",
      "  SUBROUTINE my_routine (dim1, dim2, i1, i2, in1, in2, in3, in4, in5, in6, in7, in8, in9, in10, out1)\n",
      "    INTEGER(KIND=ip), INTENT(IN) :: dim1, dim2, i1, i2\n",
      "    REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: in1, in2, in3, in4, in5, in6, in7, in8, in9, in10\n",
      "    REAL(KIND=lp), INTENT(INOUT), DIMENSION(1:dim1, 1:dim2) :: out1\n",
      "    \n",
      "    INTEGER(KIND=ip) :: i, k\n",
      "    DO k=1,dim2\n",
      "      DO i=i1,i2\n",
      "        out1(i, k) = (in1(i, k) + in2(i, k) + in3(i, k) + in4(i, k) + in5(i, k) + in6(i, k) + in7(i, k) + in8(i, k) + in9(i, k) + &\n",
      "        &  in10(i, k))*0.1\n",
      "        in1(i, k) = out1(i, k)\n",
      "      END DO\n",
      "    END DO\n",
      "  END SUBROUTINE my_routine\n",
      "END MODULE my_module\n"
     ]
    }
   ],
   "source": [
    "# pylint: disable=undefined-variable\n",
    "%cat my_module.F90"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf2467a7",
   "metadata": {},
   "source": [
    "Loki's documentation holds further details on [how to read files](https://sites.ecmwf.int/docs/loki/main/frontends.html) and additional options (choice of frontends, preprocessing) for that as well as the [different backends](https://sites.ecmwf.int/docs/loki/main/backends.html) that are available to generate code."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.8 ('loki_env': venv)",
   "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.8.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "5b6429b76fde06fc4400bf3c27b3ae893ffb7a047f8b8ee9418a3bc77878d107"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}