{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1.2 CoefficientFunctions\n",
    "====\n",
    "\n",
    "In NGSolve, `CoefficientFunction`s are representations of functions defined on the computational domain. Examples are expressions of coordinate variables $x, y, z$ and functions that are  constant on subdomains. Much of the magic behind the seamless integration of NGSolve with python lies in `CoefficientFunction`s. This tutorial introduces you to them.\n",
    "\n",
    "After this tutorial you will know how to \n",
    "\n",
    "- **define** a `CoefficientFunction`,\n",
    "- **visualize**  a `CoefficientFunction`,\n",
    "- **evaluate** `CoefficientFunction`s at points,\n",
    "- print the **expression tree** of `CoefficientFunction`,\n",
    "- **integrate** a `CoefficientFunction`, \n",
    "- **differentiate** a `CoefficientFunction`, \n",
    "- include **parameter** in `CoefficientFunction`s,\n",
    "- **interpolate** a `CoefficientFunction` into a finite element space, \n",
    "- define **vector-valued** `CoefficientFunction`s, and \n",
    "- **compile** `CoefficientFunction`s."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ngsolve import *\n",
    "from ngsolve.webgui import Draw\n",
    "import matplotlib.pyplot as plt\n",
    "mesh = Mesh (unit_square.GenerateMesh(maxh=0.2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myfunc = x*(1-x)\n",
    "myfunc   # You have just created a CoefficientFunction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x        # This is a built-in CoefficientFunction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Visualize the function\n",
    "\n",
    "Use the `mesh` to visualize the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw(myfunc, mesh);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mip = mesh(0.2, 0.2)\n",
    "myfunc(mip)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `myfunc(0.2,0.3)` will not evaluate the function: one must give points in the form of `MappedIntegrationPoint`s like `mip` above. The `mesh` knows how to produce them."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Examining functions on sets of points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pts = [(0.1*i, 0.2) for i in range(11)]\n",
    "vals = [myfunc(mesh(*p)) for p in pts] \n",
    "for p,v in zip(pts, vals):\n",
    "    print(\"point=(%3.2f,%3.2f), value=%6.5f\"\n",
    "         %(p[0], p[1], v))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the vectorized version of the same code using `numpy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "X = np.linspace(0, 1, num=11)\n",
    "Y = np.ones_like(X) * 0.2\n",
    "myfunc(mesh(X, Y))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We may plot the restriction of the `CoefficientFunction` on a line using matplotlib. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X = np.linspace(0, 1, num=100)\n",
    "Y = np.ones_like(X) * 0.2\n",
    "plt.plot(X, myfunc(mesh(X, Y)))\n",
    "plt.xlabel('x')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Expression tree of a function\n",
    "\n",
    "Internally, coefficient functions are implemented as an expression tree made from building blocks like `x`, `y`, `sin`, etc., and arithmetic operations. (Knowing this will provide context for the later discussion of speeding up their evaluations.) For example, the expression tree for `myfunc = x*(1-x)` looks like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(myfunc) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Integrate a function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can numerically integrate the function using the mesh:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Integrate(myfunc, mesh, order=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can change the precision of the quadrature rule used for the integration using the key word argument `order`. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Differentiate a function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Automatic differentiation of a `CoefficientFunction` is possible through the `Diff` method. Here is how you get $\\partial / \\partial x$ of `myfunc`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff_myfunc = myfunc.Diff(x)\n",
    "Draw(diff_myfunc, mesh, \"derivative\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See if you can recognize an implementation of the product rule \n",
    "\n",
    "$$\n",
    "\\frac{\\partial}{\\partial x} x (1-x)\n",
    "= \n",
    "\\frac{\\partial x}{\\partial x} \n",
    "(1-x) + \n",
    "x\\frac{\\partial (1-x)}{\\partial x} \n",
    "$$\n",
    "\n",
    "in the tree-representation of the differentiated coefficient function, printed below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(diff_myfunc)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Parameters in functions\n",
    "\n",
    "When building complex coefficient functions from simple ones like `x` and `y`, you may often want to introduce `Parameter`s, which are constants whose value may be changed later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = Parameter(1.0)\n",
    "f = sin(k*y)\n",
    "Draw(f, mesh);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The same `f` may be given a different value of `k` later:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k.Set(10)\n",
    "Draw(f, mesh, order=3);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look at the expression tree of `f`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how the parameter\n",
    "is now a **node** in the expression tree.  You can differentiate a coefficient function with respect to such quantities by passing it as argument to `Diff`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print (f.Diff(k))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Integrate((f.Diff(k) - y*cos(k*y))**2, mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Interpolate a function\n",
    "\n",
    "We may `Set` a `GridFunction` using a `CoefficientFunction`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fes = H1(mesh, order=1)\n",
    "u = GridFunction(fes)\n",
    "u.Set(myfunc)\n",
    "Draw(u); "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* The `Set` method interpolates `myfunc` to an element `u` in the finite element space.\n",
    "\n",
    "* `Set` does an *Oswald-type interpolation* as follows:\n",
    "    - It first zeros the grid function;\n",
    "    - It then projects `myfunc` in $L^2$ on each mesh element;\n",
    "    - It then averages dofs on element interfaces for conformity.\n",
    "    \n",
    "* We will see other ways to interpolate in [Unit 2.10](../unit-2.10-dualbasis/dualbasis.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Vector-valued `CoefficientFunction`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an example of a vector-valued coefficient function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecfun = CoefficientFunction((-y, sin(x)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Instead of writing the long name, you can abbreviate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "vecfun = CF((-y, sin(x)))   # CF = CoefficientFunction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To draw the vector field in the `webgui` using arrows, use the `vectors` keyword argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Draw(vecfun, mesh, vectors=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another example of a vector-valued coefficient function is the gradient of a scalar `GridFunction`. Here is an example using the above-set `GridFunction` called `u`. The colors represent the magnitude of the vector field and the arrows give the direction. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u.Set(myfunc)\n",
    "gradu = grad(u)\n",
    "Draw(gradu, mesh, vectors={\"grid_size\":30});"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compiled `CoefficientFunction`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation of a `CoefficientFunction` at a point is usually done \n",
    "by traversing its expression tree and evaluating each node of the tree. When the tree has repeated nodes, this is likely wasteful. NGSolve allows you to **\"compile\"** a `CoefficientFunction` to increase the efficiency of its evaluation. The compilation translates the expression tree into a sequence of linear steps. \n",
    "\n",
    "Continuing with our simple `myfunc` example, here is how to use the `Compile` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "myfunc_compiled = myfunc.Compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now look at the differences between the compiled and non-compiled `CoefficientFunction`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(myfunc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(myfunc_compiled)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Evaluation of the compiled function is now a linear sequence of Steps 0, 1, 2, and 3 above. The output description of Steps 2 and 3 above means the following:\n",
    "\n",
    "*  Step 2: (Output of Step 1) - (Output of Step 0) \n",
    "*  Step 3: (Output of Step 0) * (Output of Step 2) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is another example, along with differences in timings for integrating the same coefficient function, in three different forms, on the same mesh. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "f0 = myfunc\n",
    "f1 = f0*y \n",
    "f2 = f1*f1 + f1*f0 + f0*f0 \n",
    "f3 = f2*f2*f2*f0**2 + f0*f2**2 + f0**2 + f1**2 + f2**2\n",
    "final = f3 + f3 + f3 \n",
    "finalc = final.Compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit Integrate(final, mesh, order=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit Integrate(finalc, mesh, order=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your NGSolve installation has a compiler script (you likely do if you built NGSolve from source using a compiler/toolchain), then you additionally have the option of letting that script optimize your coefficient function further. Here is an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finalcc = final.Compile(realcompile=True, wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit Integrate(finalcc, mesh, order=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Clearly these experiments show that the evaluation of an expression tree can be made faster by compiling the tree into a set of linear instructions. This conclusion is also supported by looking at the long outputs of the expression tree with and without the compile option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(finalc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(final)"
   ]
  }
 ],
 "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": 4
}