Tutorial-Indexed_Expressions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<script async src=\"https://www.googletagmanager.com/gtag/js?id=UA-59152712-8\"></script>\n",
    "<script>\n",
    "  window.dataLayer = window.dataLayer || [];\n",
    "  function gtag(){dataLayer.push(arguments);}\n",
    "  gtag('js', new Date());\n",
    "\n",
    "  gtag('config', 'UA-59152712-8');\n",
    "</script>\n",
    "\n",
    "# Indexed Expressions: Representing and manipulating tensors, pseudotensors, etc. in NRPy+\n",
    "\n",
    "## Author: Zach Etienne\n",
    "### Formatting improvements courtesy Brandon Clark\n",
    "\n",
    "## This notebook demonstrates the use of NRPy+ in symbolic tensor computations and numerical code generation. It introduces techniques for manipulating indexed expressions with varying ranks, demonstrating practical applications such as tensor contractions and the conversion of C code with SIMD support.\n",
    "\n",
    "**Notebook Status:** <font color='red'><b> Not Validated </b></font>\n",
    "\n",
    "### NRPy+ Source Code for this module: [indexedexp.py](../edit/indexedexp.py)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='toc'></a>\n",
    "\n",
    "# Table of Contents\n",
    "$$\\label{toc}$$\n",
    "\n",
    "This notebook is organized as follows\n",
    "\n",
    "1. [Step 1](#initializenrpy): Initialize core Python/NRPy+ modules\n",
    "1. [Step 2](#idx1): Rank-1 Indexed Expressions\n",
    "    1. [Step 2.a](#dot): Performing a Dot Product\n",
    "1. [Step 3](#idx2): Rank-2 and Higher Indexed Expressions \n",
    "    1. [Step 3.a](#con): Creating C Code for the contraction variable \n",
    "    1. [Step 3.b](#simd): Enable SIMD support\n",
    "1. [Step 4](#exc): Exercise\n",
    "1. [Step 5](#latex_pdf_output): Output this notebook to $\\LaTeX$-formatted PDF file"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='initializenrpy'></a>\n",
    "\n",
    "# Step 1: Initialize core NRPy+ modules \\[Back to [top](#toc)\\]\n",
    "$$\\label{initializenrpy}$$\n",
    "\n",
    "Let's start by importing all the needed modules from NRPy+ for dealing with indexed expressions and ouputting C code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:53.882166Z",
     "iopub.status.busy": "2021-03-07T17:11:53.881429Z",
     "iopub.status.idle": "2021-03-07T17:11:54.205065Z",
     "shell.execute_reply": "2021-03-07T17:11:54.205516Z"
    }
   },
   "outputs": [],
   "source": [
    "# The NRPy_param_funcs module sets up global structures that manage free parameters within NRPy+\n",
    "import NRPy_param_funcs as par   # NRPy+: Parameter interface\n",
    "# The indexedexp module defines various functions for defining and managing indexed quantities like tensors and pseudotensors\n",
    "import indexedexp as ixp         # NRPy+: Symbolic indexed expression (e.g., tensors, vectors, etc.) support\n",
    "# The grid module defines various parameters related to a numerical grid or the dimensionality of indexed expressions\n",
    "# For example, it declares the parameter DIM, which specifies the dimensionality of the indexed expression\n",
    "import grid as gri               # NRPy+: Functions having to do with numerical grids\n",
    "from outputC import outputC      # NRPy+: Basic C code output functionality"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='idx1'></a>\n",
    "\n",
    "# Step 2: Rank-1 Indexed Expressions \\[Back to [top](#toc)\\]\n",
    "$$\\label{idx1}$$\n",
    "\n",
    "Indexed expressions of rank 1 are stored as [Python lists](https://www.tutorialspoint.com/python/python_lists.htm). \n",
    "\n",
    "There are two standard ways to declare indexed expressions:\n",
    "+ **Initialize indexed expression to zero:** \n",
    "    + **zerorank1(DIM=-1)** $\\leftarrow$ As we will see below, initializing to zero is useful if the indexed expression depends entirely on some other indexed or non-indexed expressions.\n",
    "        + **DIM** is an *optional* parameter that, if set to -1, will default to the dimension as set in the **grid** module: `par.parval_from_str(\"grid::DIM\")`. Otherwise, the rank-1 indexed expression will have dimension **DIM**.\n",
    "+ **Initialize indexed expression symbolically:** \n",
    "    + **declarerank1(symbol, DIM=-1)**. \n",
    "        + As in **`zerorank1()`**, **DIM** is an *optional* parameter that, if set to -1, will default to the dimension as set in the **grid** module: `par.parval_from_str(\"grid::DIM\")`. Otherwise, the rank-1 indexed expression will have dimension **DIM**.\n",
    "\n",
    "`zerorank1()` and `declarerank1()` are both wrapper functions for the more general function `declare_indexedexp()`.\n",
    "+ **declare_indexedexp(rank, symbol=None, symmetry=None, dimension=None)**.\n",
    "    + The following are optional parameters: **symbol**, **symmetry**, and **dimension**. If **symbol** is not specified, then `declare_indexedexp()` will initialize an indexed expression to zero. If **symmetry** is not specified or has value \"nosym\", then an indexed expression will not be symmetrized, which has no relevance for an indexed expression of rank 1. If **dimension** is not specified or has value -1, then **dimension** will default to the dimension as set in the **grid** module: `par.parval_from_str(\"grid::DIM\")`.\n",
    "\n",
    "For example, the 3-vector $\\beta^i$ (upper index denotes contravariant) can be initialized to zero as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.211655Z",
     "iopub.status.busy": "2021-03-07T17:11:54.210632Z",
     "iopub.status.idle": "2021-03-07T17:11:54.214370Z",
     "shell.execute_reply": "2021-03-07T17:11:54.214840Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 0, 0]\n"
     ]
    }
   ],
   "source": [
    "# Declare rank-1 contravariant (\"U\") vector\n",
    "betaU = ixp.zerorank1()\n",
    "\n",
    "# Print the result. It's a list of zeros!\n",
    "print(betaU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next set $\\beta^i = \\sum_{j=0}^i j = \\{0,1,3\\}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.220338Z",
     "iopub.status.busy": "2021-03-07T17:11:54.219693Z",
     "iopub.status.idle": "2021-03-07T17:11:54.222488Z",
     "shell.execute_reply": "2021-03-07T17:11:54.222955Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The 3-vector betaU is now set to: [0, 1, 3]\n"
     ]
    }
   ],
   "source": [
    "# Get the dimension we just set, so we know how many indices to loop over\n",
    "DIM = par.parval_from_str(\"grid::DIM\")\n",
    "\n",
    "for i in range(DIM): # sum i from 0 to DIM-1, inclusive\n",
    "    for j in range(i+1): # sum j from 0 to i, inclusive\n",
    "        betaU[i] += j\n",
    "\n",
    "print(\"The 3-vector betaU is now set to: \"+str(betaU))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, the 3-vector $\\beta^i$ can be initialized **symbolically** as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.228710Z",
     "iopub.status.busy": "2021-03-07T17:11:54.227931Z",
     "iopub.status.idle": "2021-03-07T17:11:54.231683Z",
     "shell.execute_reply": "2021-03-07T17:11:54.231083Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[betaU0, betaU1, betaU2]\n"
     ]
    }
   ],
   "source": [
    "# Set the dimension to 3\n",
    "par.set_parval_from_str(\"grid::DIM\",3)\n",
    "\n",
    "# Declare rank-1 contravariant (\"U\") vector\n",
    "betaU = ixp.declarerank1(\"betaU\")\n",
    "\n",
    "# Print the result. It's a list!\n",
    "print(betaU)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Declaring $\\beta^i$ symbolically is standard in case `betaU0`, `betaU1`, and `betaU2` are defined elsewhere (e.g., read in from main memory as a gridfunction).\n",
    "\n",
    "As can be seen, NRPy+'s standard naming convention for indexed rank-1 expressions is \n",
    "+ **\\[base variable name\\]+\\[\"U\" for contravariant (up index) or \"D\" for covariant (down index)\\]**\n",
    "\n",
    "*Caution*: after declaring the vector, `betaU0`, `betaU1`, and `betaU2` can only be accessed or manipulated through list access; i.e., via `betaU[0]`, `betaU[1]`, and `betaU[2]`, respectively. Attempts to access `betaU0` directly will fail.\n",
    "\n",
    "Knowing this, let's multiply `betaU1` by 2:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.239050Z",
     "iopub.status.busy": "2021-03-07T17:11:54.237932Z",
     "iopub.status.idle": "2021-03-07T17:11:54.241724Z",
     "shell.execute_reply": "2021-03-07T17:11:54.242357Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The 3-vector betaU is now set to [betaU0, 2*betaU1, betaU2]\n",
      "The component betaU[1] is now set to 2*betaU1\n"
     ]
    }
   ],
   "source": [
    "betaU[1] *= 2\n",
    "print(\"The 3-vector betaU is now set to \"+str(betaU))\n",
    "print(\"The component betaU[1] is now set to \"+str(betaU[1]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='dot'></a>\n",
    "\n",
    "## Step 2.a: Performing a Dot Product \\[Back to [top](#toc)\\]\n",
    "$$\\label{dot}$$\n",
    "\n",
    "Next, let's declare the variable $\\beta_j$ and perform the dot product $\\beta^i \\beta_i$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.251924Z",
     "iopub.status.busy": "2021-03-07T17:11:54.250700Z",
     "iopub.status.idle": "2021-03-07T17:11:54.287349Z",
     "shell.execute_reply": "2021-03-07T17:11:54.286718Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "betaD0*betaU0 + betaD1*betaU1 + betaD2*betaU2\n"
     ]
    }
   ],
   "source": [
    "# First set betaU back to its initial value\n",
    "betaU = ixp.declarerank1(\"betaU\")\n",
    "\n",
    "# Declare beta_j:\n",
    "betaD = ixp.declarerank1(\"betaD\")\n",
    "\n",
    "# Get the dimension we just set, so we know how many indices to loop over\n",
    "DIM = par.parval_from_str(\"grid::DIM\")\n",
    "\n",
    "# Initialize dot product to zero\n",
    "dotprod = 0\n",
    "\n",
    "# Perform dot product beta^i beta_i\n",
    "for i in range(DIM):\n",
    "    dotprod += betaU[i]*betaD[i]\n",
    "\n",
    "# Print result!\n",
    "print(dotprod)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='idx2'></a>\n",
    "\n",
    "# Step 3: Rank-2 and Higher Indexed Expressions \\[Back to [top](#toc)\\]\n",
    "$$\\label{idx2}$$\n",
    "\n",
    "Moving to higher ranks, rank-2 indexed expressions are stored as lists of lists, rank-3 indexed expressions as lists of lists of lists, etc. For example:\n",
    "\n",
    "+ the covariant rank-2 tensor $g_{ij}$ is declared as `gDD[i][j]` in NRPy+, so that e.g., `gDD[0][2]` is stored with name `gDD02` and\n",
    "+ the rank-2 tensor $T^{\\mu}{}_{\\nu}$ is declared as `TUD[m][n]` in NRPy+ (index names are of course arbitrary).\n",
    "\n",
    "*Caveat*: note that it is currently up to the user to determine whether the combination of indexed expressions makes sense; NRPy+ does not track whether up and down indices are written consistently.\n",
    "\n",
    "NRPy+ supports symmetries in indexed expressions (above rank 1), so that if $h_{ij} = h_{ji}$, then declaring `hDD[i][j]` to be symmetric in NRPy+ will result in both `hDD[0][2]` and `hDD[2][0]` mapping to the *single* SymPy variable `hDD02`.\n",
    "\n",
    "To see how this works in NRPy+, let's define in NRPy+ a symmetric, rank-2 tensor $h_{ij}$ in three dimensions, and then compute the contraction, which should be given by $$con = h^{ij}h_{ij} = h_{00} h^{00} + h_{11} h^{11} + h_{22} h^{22} + 2 (h_{01} h^{01} + h_{02} h^{02} + h_{12} h^{12}).$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.299407Z",
     "iopub.status.busy": "2021-03-07T17:11:54.298558Z",
     "iopub.status.idle": "2021-03-07T17:11:54.301860Z",
     "shell.execute_reply": "2021-03-07T17:11:54.302510Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "hDD00*hUU00 + 2*hDD01*hUU01 + 2*hDD02*hUU02 + hDD11*hUU11 + 2*hDD12*hUU12 + hDD22*hUU22\n"
     ]
    }
   ],
   "source": [
    "# Get the dimension we just set (should be set to 3).\n",
    "DIM = par.parval_from_str(\"grid::DIM\")\n",
    "\n",
    "# Declare h_{ij}=hDD[i][j] and h^{ij}=hUU[i][j]\n",
    "hUU = ixp.declarerank2(\"hUU\",\"sym01\")\n",
    "hDD = ixp.declarerank2(\"hDD\",\"sym01\")\n",
    "\n",
    "# Perform sum h^{ij} h_{ij}, initializing contraction result to zero\n",
    "con = 0\n",
    "for i in range(DIM):\n",
    "    for j in range(DIM):\n",
    "        con += hUU[i][j]*hDD[i][j]\n",
    "\n",
    "# Print result\n",
    "print(con)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='con'></a>\n",
    "\n",
    "## Step 3.a: Creating C Code for the contraction variable $\\text{con}$ \\[Back to [top](#toc)\\]\n",
    "$$\\label{con}$$\n",
    "\n",
    "Next let's create the C code for the contraction variable $\\text{con}$, without CSE (common subexpression elimination)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.311229Z",
     "iopub.status.busy": "2021-03-07T17:11:54.310497Z",
     "iopub.status.idle": "2021-03-07T17:11:54.361366Z",
     "shell.execute_reply": "2021-03-07T17:11:54.361835Z"
    },
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/*\n",
      " *  Original SymPy expression:\n",
      " *  \"con = hDD00*hUU00 + 2*hDD01*hUU01 + 2*hDD02*hUU02 + hDD11*hUU11 + 2*hDD12*hUU12 + hDD22*hUU22\"\n",
      " */\n",
      "{\n",
      "  con = hDD00*hUU00 + 2*hDD01*hUU01 + 2*hDD02*hUU02 + hDD11*hUU11 + 2*hDD12*hUU12 + hDD22*hUU22;\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "outputC(con,\"con\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='simd'></a>\n",
    "\n",
    "## Step 3.b: Enable SIMD support \\[Back to [top](#toc)\\]\n",
    "$$\\label{simd}$$\n",
    "\n",
    "Finally, let's see how it looks with SIMD support enabled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.380421Z",
     "iopub.status.busy": "2021-03-07T17:11:54.379784Z",
     "iopub.status.idle": "2021-03-07T17:11:54.382852Z",
     "shell.execute_reply": "2021-03-07T17:11:54.383317Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/*\n",
      " *  Original SymPy expression:\n",
      " *  \"con = hDD00*hUU00 + 2*hDD01*hUU01 + 2*hDD02*hUU02 + hDD11*hUU11 + 2*hDD12*hUU12 + hDD22*hUU22\"\n",
      " */\n",
      "{\n",
      "  const double tmp_Integer_2 = 2.0;\n",
      "  const REAL_SIMD_ARRAY _Integer_2 = ConstSIMD(tmp_Integer_2);\n",
      "\n",
      "  con = FusedMulAddSIMD(hDD22, hUU22, FusedMulAddSIMD(_Integer_2, MulSIMD(hDD01, hUU01), FusedMulAddSIMD(_Integer_2, MulSIMD(hDD02, hUU02), FusedMulAddSIMD(_Integer_2, MulSIMD(hDD12, hUU12), FusedMulAddSIMD(hDD00, hUU00, MulSIMD(hDD11, hUU11))))));\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "outputC(con,\"con\",params=\"enable_SIMD=True\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='exc'></a>\n",
    "\n",
    "# Step 4: Exercise \\[Back to [top](#toc)\\]\n",
    "$$\\label{exc}$$\n",
    "\n",
    "Setting $\\beta^i$ via the declarerank1(), write the NRPy+ code required to generate the needed C code for the lowering operator: $g_{ij} \\beta^i$, and set the result to C variables `betaD0out`, `betaD1out`, and `betaD2out` [solution](Tutorial-Indexed_Expressions_soln.ipynb). *Hint: You will want to use the `zerorank1()` function*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**To complete this exercise, you must first reset all variables in the notebook:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.386900Z",
     "iopub.status.busy": "2021-03-07T17:11:54.386305Z",
     "iopub.status.idle": "2021-03-07T17:11:54.388451Z",
     "shell.execute_reply": "2021-03-07T17:11:54.388978Z"
    }
   },
   "outputs": [],
   "source": [
    "# *Uncomment* the below %reset command and then press <Shift>+<Enter>.\n",
    "#    Respond with \"y\" in the dialog box to reset all variables.\n",
    "\n",
    "# %reset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Write your solution below:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='latex_pdf_output'></a>\n",
    "\n",
    "# Step 5: Output this notebook to $\\LaTeX$-formatted PDF file \\[Back to [top](#toc)\\]\n",
    "$$\\label{latex_pdf_output}$$\n",
    "\n",
    "The following code cell converts this Jupyter notebook into a proper, clickable $\\LaTeX$-formatted PDF file. After the cell is successfully run, the generated PDF may be found in the root NRPy+ tutorial directory, with filename\n",
    "[Tutorial-Indexed_Expressions.pdf](Tutorial-Indexed_Expressions.pdf). (Note that clicking on this link may not work; you may need to open the PDF file through another means.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2021-03-07T17:11:54.393323Z",
     "iopub.status.busy": "2021-03-07T17:11:54.392676Z",
     "iopub.status.idle": "2021-03-07T17:11:57.883821Z",
     "shell.execute_reply": "2021-03-07T17:11:57.884810Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created Tutorial-Indexed_Expressions.tex, and compiled LaTeX file to PDF\n",
      "    file Tutorial-Indexed_Expressions.pdf\n"
     ]
    }
   ],
   "source": [
    "import cmdline_helper as cmd    # NRPy+: Multi-platform Python command-line interface\n",
    "cmd.output_Jupyter_notebook_to_LaTeXed_PDF(\"Tutorial-Indexed_Expressions\")"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}