{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Codegen Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One of the most important features of symforce is the ability to generate computationally efficient code from symbolic expressions. Before progressing, first make sure you are familiar with the other symforce tutorials, especially the [Values tutorial](../tutorials/values_tutorial.html).\n",
    "\n",
    "The typical workflow for generating a function is to define a Python function that operates on symbolic inputs to return the symbolic result.  Typically this will look like:\n",
    "\n",
    "1. Define a Python function that operates on symbolic inputs\n",
    "2. Create a Codegen object using `Codegen.function`.  Various properties of the function will be deduced automatically; for instance, the name of the generated function is generated from the name of the Python function, and the argument names and types are deduced from the Python function argument names and type annotations.\n",
    "3. Generate the code in your desired language\n",
    "\n",
    "Alternately, you may want to define the input and output symbolic `Values` explicitly, with the following steps:\n",
    "\n",
    "1. Build an input Values object that defines a symbolic representation of each input to the function. Note that inputs and outputs can be Values objects themselves, which symforce will automatically generate into custom types.\n",
    "2. Build an output Values object that defines the outputs of the function in terms of the objects in the input Values.\n",
    "3. Generate the code in your desired language"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup\n",
    "import numpy as np\n",
    "\n",
    "import symforce\n",
    "\n",
    "symforce.set_symbolic_api(\"symengine\")\n",
    "symforce.set_log_level(\"warning\")\n",
    "\n",
    "# Set epsilon to a symbol for safe code generation.  For more information, see the Epsilon tutorial:\n",
    "# https://symforce.org/tutorials/epsilon_tutorial.html\n",
    "symforce.set_epsilon_to_symbol()\n",
    "\n",
    "import symforce.symbolic as sf\n",
    "from symforce import codegen\n",
    "from symforce.codegen import codegen_util\n",
    "from symforce.notebook_util import display\n",
    "from symforce.notebook_util import display_code_file\n",
    "from symforce.values import Values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating from a Python function"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we look at using existing python functions to generate an equivalent function using the codegen package. The inputs to the function are automatically deduced from the signature and type annotations.  Additionally, we can change how the generated function is declared (e.g. whether to return an object using a return statement or a pointer passed as an argument to the function)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def az_el_from_point(\n",
    "    nav_T_cam: sf.Pose3, nav_t_point: sf.Vector3, epsilon: sf.Scalar = 0\n",
    ") -> sf.Vector2:\n",
    "    \"\"\"\n",
    "    Transform a nav point into azimuth / elevation angles in the\n",
    "    camera frame.\n",
    "\n",
    "    Args:\n",
    "        nav_T_cam (sf.Pose3): camera pose in the world\n",
    "        nav_t_point (sf.Matrix): nav point\n",
    "        epsilon (Scalar): small number to avoid singularities\n",
    "\n",
    "    Returns:\n",
    "        sf.Matrix: (azimuth, elevation)\n",
    "    \"\"\"\n",
    "    cam_t_point = nav_T_cam.inverse() * nav_t_point\n",
    "    x, y, z = cam_t_point\n",
    "    theta = sf.atan2(y, x + epsilon)\n",
    "    phi = sf.pi / 2 - sf.acos(z / (cam_t_point.norm() + epsilon))\n",
    "    return sf.V2(theta, phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az_el_codegen = codegen.Codegen.function(\n",
    "    func=az_el_from_point,\n",
    "    config=codegen.CppConfig(),\n",
    ")\n",
    "az_el_codegen_data = az_el_codegen.generate_function()\n",
    "\n",
    "print(\"Files generated in {}:\\n\".format(az_el_codegen_data.output_dir))\n",
    "for f in az_el_codegen_data.generated_files:\n",
    "    print(\"  |- {}\".format(f.relative_to(az_el_codegen_data.output_dir)))\n",
    "\n",
    "display_code_file(az_el_codegen_data.generated_files[0], \"C++\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating function jacobians"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "codegen_with_jacobians = az_el_codegen.with_jacobians(\n",
    "    # Just compute wrt the pose and point, not epsilon\n",
    "    which_args=[\"nav_T_cam\", \"nav_t_point\"],\n",
    "    # Include value, not just jacobians\n",
    "    include_results=True,\n",
    ")\n",
    "\n",
    "data = codegen_with_jacobians.generate_function()\n",
    "\n",
    "display_code_file(data.generated_files[0], \"C++\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code generation using implicit functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we look at generating functions using a list of input variables and output expressions that are a function of those variables. In this case we don't need to explicitly define a function in python, but can instead generate one directly using the codegen package."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's set up an example for the [double pendulum](https://www.myphysicslab.com/pendulum/double-pendulum-en.html). We'll skip the derivation and just define the equations of motion for the angular acceleration of the two links:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define symbols\n",
    "L = sf.V2.symbolic(\"L\").T  # Length of the two links\n",
    "m = sf.V2.symbolic(\"m\").T  # Mass of the two links\n",
    "ang = sf.V2.symbolic(\"a\").T  # Angle of the two links\n",
    "dang = sf.V2.symbolic(\"da\").T  # Angular velocity of the two links\n",
    "g = sf.Symbol(\"g\")  # Gravity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Angular acceleration of the first link\n",
    "ddang_0 = (\n",
    "    -g * (2 * m[0] + m[1]) * sf.sin(ang[0])\n",
    "    - m[1] * g * sf.sin(ang[0] - 2 * ang[1])\n",
    "    - 2\n",
    "    * sf.sin(ang[0] - ang[1])\n",
    "    * m[1]\n",
    "    * (dang[1] * 2 * L[1] + dang[0] * 2 * L[0] * sf.cos(ang[0] - ang[1]))\n",
    ") / (L[0] * (2 * m[0] + m[1] - m[1] * sf.cos(2 * ang[0] - 2 * ang[1])))\n",
    "display(ddang_0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Angular acceleration of the second link\n",
    "ddang_1 = (\n",
    "    2\n",
    "    * sf.sin(ang[0] - ang[1])\n",
    "    * (\n",
    "        dang[0] ** 2 * L[0] * (m[0] + m[1])\n",
    "        + g * (m[0] + m[1]) * sf.cos(ang[0])\n",
    "        + dang[1] ** 2 * L[1] * m[1] * sf.cos(ang[0] - ang[1])\n",
    "    )\n",
    ") / (L[1] * (2 * m[0] + m[1] - m[1] * sf.cos(2 * ang[0] - 2 * ang[1])))\n",
    "display(ddang_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's organize the input symbols into a Values hierarchy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = Values()\n",
    "\n",
    "inputs[\"ang\"] = ang\n",
    "inputs[\"dang\"] = dang\n",
    "\n",
    "with inputs.scope(\"constants\"):\n",
    "    inputs[\"g\"] = g\n",
    "\n",
    "with inputs.scope(\"params\"):\n",
    "    inputs[\"L\"] = L\n",
    "    inputs[\"m\"] = m\n",
    "\n",
    "display(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output will simply be a 2-vector of the angular accelerations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "outputs = Values(ddang=sf.V2(ddang_0, ddang_1))\n",
    "\n",
    "display(outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now run code generation to produce an executable module (in a temp directory if none provided):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "double_pendulum = codegen.Codegen(\n",
    "    inputs=inputs,\n",
    "    outputs=outputs,\n",
    "    config=codegen.CppConfig(),\n",
    "    name=\"double_pendulum\",\n",
    "    return_key=\"ddang\",\n",
    ")\n",
    "double_pendulum_data = double_pendulum.generate_function()\n",
    "\n",
    "# Print what we generated\n",
    "print(\"Files generated in {}:\\n\".format(double_pendulum_data.output_dir))\n",
    "for f in double_pendulum_data.generated_files:\n",
    "    print(\"  |- {}\".format(f.relative_to(double_pendulum_data.output_dir)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_code_file(double_pendulum_data.function_dir / \"double_pendulum.h\", \"C++\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also generate functions with different function declarations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function using structs as inputs and outputs (returned as pointer arg)\n",
    "input_values = Values(inputs=inputs)\n",
    "output_values = Values(outputs=outputs)\n",
    "namespace = \"double_pendulum\"\n",
    "double_pendulum_values = codegen.Codegen(\n",
    "    inputs=input_values,\n",
    "    outputs=output_values,\n",
    "    config=codegen.CppConfig(),\n",
    "    name=\"double_pendulum\",\n",
    ")\n",
    "double_pendulum_values_data = double_pendulum_values.generate_function(\n",
    "    namespace=namespace,\n",
    ")\n",
    "\n",
    "# Print what we generated. Note the nested structs that were automatically\n",
    "# generated.\n",
    "print(\"Files generated in {}:\\n\".format(double_pendulum_values_data.output_dir))\n",
    "for f in double_pendulum_values_data.generated_files:\n",
    "    print(\"  |- {}\".format(f.relative_to(double_pendulum_values_data.output_dir)))\n",
    "\n",
    "display_code_file(\n",
    "    double_pendulum_values_data.function_dir / \"double_pendulum.h\",\n",
    "    \"C++\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we can generate the same function in other languages as well:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "namespace = \"double_pendulum\"\n",
    "double_pendulum_python = codegen.Codegen(\n",
    "    inputs=inputs,\n",
    "    outputs=outputs,\n",
    "    config=codegen.PythonConfig(use_eigen_types=False),\n",
    "    name=\"double_pendulum\",\n",
    "    return_key=\"ddang\",\n",
    ")\n",
    "double_pendulum_python_data = double_pendulum_python.generate_function(\n",
    "    namespace=namespace,\n",
    ")\n",
    "\n",
    "print(\"Files generated in {}:\\n\".format(double_pendulum_python_data.output_dir))\n",
    "for f in double_pendulum_python_data.generated_files:\n",
    "    print(\"  |- {}\".format(f.relative_to(double_pendulum_python_data.output_dir)))\n",
    "\n",
    "display_code_file(\n",
    "    double_pendulum_python_data.function_dir / \"double_pendulum.py\",\n",
    "    \"python\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "constants_t = codegen_util.load_generated_lcmtype(\n",
    "    namespace, \"constants_t\", double_pendulum_python_data.python_types_dir\n",
    ")\n",
    "\n",
    "params_t = codegen_util.load_generated_lcmtype(\n",
    "    namespace, \"params_t\", double_pendulum_python_data.python_types_dir\n",
    ")\n",
    "\n",
    "ang = np.array([[0.0, 0.5]])\n",
    "dang = np.array([[0.0, 0.0]])\n",
    "consts = constants_t()\n",
    "consts.g = 9.81\n",
    "params = params_t()\n",
    "params.L = [0.5, 0.3]\n",
    "params.m = [0.3, 0.2]\n",
    "\n",
    "gen_module = codegen_util.load_generated_package(\n",
    "    namespace, double_pendulum_python_data.function_dir\n",
    ")\n",
    "gen_module.double_pendulum(ang, dang, consts, params)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}