{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Values Tutorial"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Values](../api/symforce.values.values.html#symforce.values.values.Values) objects are ordered dict-like containers used to store multiple heterogeneous objects, normally for the purpose of function generation.\n",
    "\n",
    "Typically a Values object will first be created to defines a number of different symbolic inputs (e.g. rotations, translations, scalars, poses, cameras, etc.). Then, a second Values object will be created to describe the objects to be returned from the function, which will be composed of the symbolic elements defined by the input Values object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup\n",
    "import symforce\n",
    "\n",
    "symforce.set_symbolic_api(\"sympy\")\n",
    "symforce.set_log_level(\"warning\")\n",
    "\n",
    "import symforce.symbolic as sf\n",
    "from symforce.notebook_util import display\n",
    "from symforce.values import Values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = Values(\n",
    "    x=sf.Symbol(\"x\"),\n",
    "    y=sf.Rot2.symbolic(\"c\"),\n",
    ")\n",
    "display(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `.add()` method can add a symbol using its name as the key:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs.add(sf.Symbol(\"foo\"))\n",
    "display(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding sub-values are well encouraged:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = sf.symbols(\"x y\")\n",
    "expr = x**2 + sf.sin(y) / x**2\n",
    "inputs[\"states\"] = Values(p=expr)\n",
    "display(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A Values serializes to a depth-first traversed list. This means it implements StorageOps:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(inputs.to_storage())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also get a flattened lists of keys and values, with `.` separation for sub-values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(inputs.items_recursive())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that there is a `.keys_recursive()` and a `.values_recursive()` which return flattened lists of keys and values respectively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(inputs.keys_recursive())\n",
    "display(inputs.values_recursive())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To fully reconstruct the types in the Values from the serialized scalars, we need an index that describes which parts of the serialized list correspond to which types. The spec is `T.Dict[str, IndexEntry]` where `IndexEntry` has attributes `offset, storage_dim, datatype, shape, item_index`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "index = inputs.index()\n",
    "index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With a serialized list and an index, we can get the values back:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs2 = Values.from_storage_index(inputs.to_storage(), index)\n",
    "assert inputs == inputs2\n",
    "display(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `item_index` is a recursive structure that can contain the index for a sub-values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "item_index = inputs.index()[\"states\"].item_index\n",
    "assert item_index == inputs[\"states\"].index()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also set sub-values directly with dot notation in the keys. They get split up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs[\"states.blah\"] = 3\n",
    "display(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `.attr` field also allows attribute access rather than key access:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert inputs[\"states.p\"] is inputs[\"states\"][\"p\"] is inputs.attr.states.p\n",
    "display(inputs.attr.states.p)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, SymForce adds the concept of a name scope to namespace symbols. Within a scope block, symbol names get prefixed with the scope name:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with sf.scope(\"params\"):\n",
    "    s = sf.Symbol(\"cost\")\n",
    "display(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A common use case is to call a function that adds symbols within your own name scope to avoid name collisions. You can also chain name scopes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = Values()\n",
    "v.add(sf.Symbol(\"x\"))\n",
    "with sf.scope(\"foo\"):\n",
    "    v.add(sf.Symbol(\"x\"))\n",
    "    with sf.scope(\"bar\"):\n",
    "        v.add(sf.Symbol(\"x\"))\n",
    "display(v)\n",
    "display(v.attr.foo.bar.x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The values class also provides a `.scope()` method that not only applies the scope to symbol names but also to keys added to the Values:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v = Values()\n",
    "with v.scope(\"hello\"):\n",
    "    v[\"y\"] = x**2\n",
    "    v[\"z\"] = sf.Symbol(\"z\")\n",
    "v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This flexible set of features provided by the Values class allows conveniently building up large expressions, and acts as the interface to code generation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Lie Group Operations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "One useful feature of Values objects is that element-wise Lie group operations on can be performed on them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lie_vals = Values()\n",
    "lie_vals[\"scalar\"] = sf.Symbol(\"x\")\n",
    "lie_vals[\"rot3\"] = sf.Rot3.symbolic(\"rot\")\n",
    "\n",
    "sub_lie_vals = Values()\n",
    "sub_lie_vals[\"pose3\"] = sf.Pose3.symbolic(\"pose\")\n",
    "sub_lie_vals[\"vec\"] = sf.V3.symbolic(\"vec\")\n",
    "\n",
    "lie_vals[\"sub_vals\"] = sub_lie_vals\n",
    "\n",
    "display(lie_vals)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(lie_vals.tangent_dim())\n",
    "display(len(lie_vals.to_tangent()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Importantly, we can compute the jacobian of the storage space of the object with respect to its tangent space:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(lie_vals.storage_D_tangent())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that we can use the elements of the object to compute a residual, and then compute the jacobian of such a residual with respect to the tangent space of our values object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "residual = sf.Matrix(6, 1)\n",
    "residual[0:3, 0] = lie_vals[\"rot3\"] * lie_vals[\"sub_vals.vec\"]\n",
    "residual[3:6, 0] = lie_vals[\"sub_vals.pose3\"] * lie_vals[\"sub_vals.vec\"]\n",
    "display(residual)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "residual_D_tangent = residual.jacobian(lie_vals)\n",
    "display(residual_D_tangent.shape)\n",
    "display(residual_D_tangent)"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}