demo_local_frames_construction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fcda3dbf",
   "metadata": {},
   "source": [
    "# LoCS Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd6bc7be",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import matplotlib_inline\n",
    "matplotlib_inline.backend_inline.set_matplotlib_formats('svg', 'pdf') # For export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80affc4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rotmat(angles):\n",
    "    cos_a = np.cos(angles)\n",
    "    sin_a = np.sin(angles)\n",
    "    return np.stack([np.stack([cos_a, -sin_a], -1),\n",
    "                     np.stack([sin_a, cos_a], -1)], -2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60dfade0",
   "metadata": {},
   "outputs": [],
   "source": [
    "colors = ['orangered', 'orange', 'cyan', 'royalblue']\n",
    "def rot_markers(i, theta=0):\n",
    "    # The markers below indicate a rectangle, a pentagon, a triangle, and an asterisk with 5 sides\n",
    "    func_markers = [(4, 0), (5, 0), (3, 0), (5, 1)]\n",
    "    prior_angles = np.array([0.0, -90.0, -30.0, -90.0])\n",
    "    return func_markers[i] + tuple([prior_angles[i] + theta])\n",
    "\n",
    "def get_angle_degrees(x):\n",
    "    return np.degrees(np.arctan2(x[:, 1], x[:, 0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71c400b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_with_origin(rows=1, cols=1, basesize=4):\n",
    "    fig, ax = plt.subplots(rows, cols, squeeze=False, sharex=True, sharey=True,\n",
    "                          figsize=(basesize * cols, basesize * rows))\n",
    "    for r in range(rows):\n",
    "        for c in range(cols):\n",
    "            ax[r, c].axvline(x=0, c=\"gray\", lw=2, alpha=0.5)\n",
    "            ax[r, c].axhline(y=0, c=\"gray\", lw=2, alpha=0.5)\n",
    "            ax[r, c].plot((1), (0), ls=\"\", marker=\">\", ms=10, color=\"gray\", alpha=0.5,\n",
    "                          transform=ax[r, c].get_yaxis_transform(), clip_on=False)\n",
    "            ax[r, c].plot((0), (1), ls=\"\", marker=\"^\", ms=10, color=\"gray\", alpha=0.5,\n",
    "                          transform=ax[r, c].get_xaxis_transform(), clip_on=False)\n",
    "            ax[r, c].set_axis_off()\n",
    "        ax[r, c].set_aspect('equal')\n",
    "        ax[r, c].set_xlim([-4.0, 4.0])\n",
    "        ax[r, c].set_ylim([-4.0, 4.0])\n",
    "    return fig, ax"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c41bdf4f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pos = np.array(\n",
    "    [[0.7,   0.3],\n",
    "     [1.2,   2.1],\n",
    "     [-0.9, -1.4],\n",
    "     [-1.4,  1.0]]\n",
    ")\n",
    "\n",
    "vel = np.array(\n",
    "    [[0.48, -0.13],\n",
    "     [-0.03,  0.5],\n",
    "     [0.43,  0.25],\n",
    "     [0.45, -0.21]]\n",
    ")\n",
    "num_objects = pos.shape[0]\n",
    "yaw = np.arctan2(vel[:, 1], vel[:, 0])\n",
    "yaw_angle = np.degrees(yaw)\n",
    "\n",
    "# In 2D, we could directly do the following:\n",
    "# R_inv = rotmat(-yaw)\n",
    "R = rotmat(yaw)\n",
    "R_inv = R.swapaxes(2, 1)\n",
    "\n",
    "rot_vel = np.einsum('nmij,nmjk->nmik', R_inv[:, np.newaxis], vel[np.newaxis, :, :, np.newaxis])\n",
    "\n",
    "rel_pos = pos[None, ...] - pos[:, None, :]\n",
    "rot_rel_pos = np.einsum('nmij,nmjk->nmik', R_inv[:, np.newaxis], rel_pos[..., np.newaxis])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeca3b7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_with_origin()\n",
    "for i in range(pos.shape[0]):\n",
    "    ax[0, 0].plot(pos[i, 0], pos[i, 1], marker=rot_markers(i, yaw_angle[i]),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20)\n",
    "    ax[0, 0].quiver(pos[i, 0], pos[i, 1], vel[i, 0], vel[i, 1], scale=4)\n",
    "plt.savefig('assets/img/demo_global_frame.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d58238f",
   "metadata": {},
   "source": [
    "![Global frame](assets/img/demo_global_frame.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9548681a",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_with_origin(2, 2)\n",
    "for i in range(num_objects):\n",
    "    r, c = divmod(i, 2)\n",
    "    for j in range(num_objects):\n",
    "        ax[r, c].plot(rot_rel_pos[i, j, 0], rot_rel_pos[i, j, 1],\n",
    "                      marker=rot_markers(j, yaw_angle[j]-yaw_angle[i]),\n",
    "                      color=colors[j], markeredgecolor='k', markersize=20)\n",
    "        ax[r, c].quiver(rot_rel_pos[i, j, 0], rot_rel_pos[i, j, 1], rot_vel[i, j, 0], rot_vel[i, j, 1], scale=4)\n",
    "plt.savefig('assets/img/demo_local_frames.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9fbce16",
   "metadata": {},
   "source": [
    "![Local frames](assets/img/demo_local_frames.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38c7e3ad",
   "metadata": {},
   "source": [
    "## Inverse Transformation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1362fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "local_pos_pred = np.array(\n",
    "    [[1.3,  0.2],\n",
    "     [0.9, -0.1],\n",
    "     [1.5, -0.2],\n",
    "     [1.4,  0.0]]\n",
    ")\n",
    "\n",
    "local_vel_pred = np.array(\n",
    "    [[0.05,  0.05],\n",
    "     [0.0,  -0.05],\n",
    "     [0.05, -0.05],\n",
    "     [-0.1,   0.0]]\n",
    ")\n",
    "\n",
    "local_vel = rot_vel[np.arange(num_objects), np.arange(num_objects)].squeeze(-1) + local_vel_pred\n",
    "local_vel_angles = get_angle_degrees(local_vel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb537aab",
   "metadata": {},
   "outputs": [],
   "source": [
    "rot_local_pos_pred = np.einsum('nij,njk->nik', R, local_pos_pred[:, :, np.newaxis]).squeeze(-1)\n",
    "rot_local_vel_pred = np.einsum('nij,njk->nik', R, local_vel_pred[:, :, np.newaxis]).squeeze(-1)\n",
    "\n",
    "global_pos_pred = pos + rot_local_pos_pred\n",
    "global_vel_pred = vel + rot_local_vel_pred\n",
    "\n",
    "global_vel_pred_angles = get_angle_degrees(global_vel_pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7509831",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_with_origin(2, 2)\n",
    "for i in range(num_objects):\n",
    "    r, c = divmod(i, 2)\n",
    "    ax[r, c].plot(rot_rel_pos[i, i, 0], rot_rel_pos[i, i, 1],\n",
    "                  marker=rot_markers(i, 0.0),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20, alpha=0.5)\n",
    "    ax[r, c].quiver(rot_rel_pos[i, i, 0], rot_rel_pos[i, i, 1], rot_vel[i, i, 0], rot_vel[i, i, 1],\n",
    "                    scale=4, alpha=0.8)\n",
    "    ax[r, c].plot(local_pos_pred[i, 0], local_pos_pred[i, 1],\n",
    "                  marker=rot_markers(i, local_vel_angles[i]),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20)\n",
    "    ax[r, c].quiver(local_pos_pred[i, 0], local_pos_pred[i, 1],\n",
    "                    local_vel[i, 0], local_vel[i, 1], scale=4)\n",
    "plt.savefig('assets/img/demo_local_pred.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ab40eba",
   "metadata": {},
   "source": [
    "![Local Predictions](assets/img/demo_local_pred.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dca41c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_with_origin(2, 2)\n",
    "for i in range(num_objects):\n",
    "    r, c = divmod(i, 2)\n",
    "    ax[r, c].plot(pos[i, 0], pos[i, 1], marker=rot_markers(i, yaw_angle[i]),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20, alpha=0.5)\n",
    "    ax[r, c].quiver(pos[i, 0], pos[i, 1], vel[i, 0], vel[i, 1],\n",
    "                    scale=4, alpha=0.5)\n",
    "    ax[r, c].plot(global_pos_pred[i, 0], global_pos_pred[i, 1],\n",
    "                  marker=rot_markers(i, global_vel_pred_angles[i]),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20)\n",
    "    ax[r, c].quiver(global_pos_pred[i, 0], global_pos_pred[i, 1],\n",
    "                    global_vel_pred[i, 0], global_vel_pred[i, 1], scale=4)\n",
    "plt.savefig('assets/img/demo_global_pred.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ae7881b",
   "metadata": {},
   "source": [
    "![Global Predictions Separate](assets/img/demo_global_pred.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b24b5f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_with_origin(1, 1)\n",
    "for i in range(num_objects):\n",
    "    ax[0, 0].plot(pos[i, 0], pos[i, 1], marker=rot_markers(i, yaw_angle[i]),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20, alpha=0.5)\n",
    "    ax[0, 0].quiver(pos[i, 0], pos[i, 1], vel[i, 0], vel[i, 1],\n",
    "                    scale=4, alpha=0.5)\n",
    "    ax[0, 0].plot(global_pos_pred[i, 0], global_pos_pred[i, 1],\n",
    "                  marker=rot_markers(i, global_vel_pred_angles[i]),\n",
    "                  color=colors[i], markeredgecolor='k', markersize=20)\n",
    "    ax[0, 0].quiver(global_pos_pred[i, 0], global_pos_pred[i, 1],\n",
    "                    global_vel_pred[i, 0], global_vel_pred[i, 1], scale=4)\n",
    "plt.savefig('assets/img/demo_global_pred_unified.png', dpi=300, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84b75074",
   "metadata": {},
   "source": [
    "![Global Predictions Unified](assets/img/demo_global_pred_unified.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8a76b69",
   "metadata": {},
   "source": [
    "Hopefully you now have an intuition on how we can construct local coordinate frames and transform object states (positions and velocities$\\dagger$) from the global frame to local frames, and how to transform predictions from local frames back to the global frame.\n",
    "\n",
    "$\\dagger$: We can include other states as well, e.g. accelerations, and rotate/translate them accordingly."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d76e6a88",
   "metadata": {},
   "source": [
    "In the next demo, `demo.ipynb`, we are applying these transformations by building a standalone LoCS network from scratch!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4daeeb55",
   "metadata": {},
   "source": [
    "## BONUS: Visualize local coordinate frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd79634c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from locs.models.standalone_fc_locs import Localizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9282995b",
   "metadata": {},
   "outputs": [],
   "source": [
    "data = torch.load('data/synth/train_feats').transpose(1, 2)\n",
    "x0 = data[[0]]\n",
    "num_objects = x0.size(1)\n",
    "num_dims = 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd90d5b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "localizer = Localizer(num_objects, num_dims)\n",
    "rel_feat, R, edge_attr = localizer(x0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25fec253",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plot_with_origin(1, 2, basesize=6)\n",
    "colors = ['r', 'g', 'b']\n",
    "\n",
    "for i in range(num_objects):\n",
    "    ax[0, 0].plot(x0[0, i, :, 0], x0[0, i, :, 1], color=colors[i])\n",
    "\n",
    "    ax[0, 1].plot(x0[0, i, :, 0], x0[0, i, :, 1], color=colors[i])\n",
    "    ax[0, 1].quiver(x0[0, i, :, 0], x0[0, i, :, 1], R[0, i, :, 0, 0], R[0, i, :, 1, 0], angles='xy', color='k')\n",
    "    ax[0, 1].quiver(x0[0, i, :, 0], x0[0, i, :, 1], R[0, i, :, 0, 1], R[0, i, :, 1, 1], angles='xy', color='grey')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d1df945",
   "metadata": {},
   "source": [
    "In 2D settings, edge features $\\mathbf{e}_{ji} = [\\mathbf{v}_{j|i}, \\mathbf{v}_{i|i}]$ contain 11 features, namely:\n",
    "\n",
    "The state of node $j$, expressed at the local frame of node $i$, $\\mathbf{v}_{j|i}$:\n",
    "- Rotated relative positions (x2)\n",
    "- Rotated euler angles (x1)\n",
    "- Rotated relative positions in spherical coordinates (x2) $\\star$\n",
    "- Rotated velocities (x2)\n",
    "\n",
    "The state of node $i$, expressed at the local frame of node $i$, $\\mathbf{v}_{j|i}$:\n",
    "- Rotated relative positions (x2)\n",
    "- Rotated velocities (x2)\n",
    "\n",
    "The rotated relative positions for node $i$ are always $0$, we include them only for consistency. When using velocities as a proxy for orientations, the rotated velocities are colinear with the $x$-axis and contain the magnitude of the velocity.\n",
    "\n",
    "$\\star$: Rotated relative positions in spherical coordinates include the distance between nodes, and the azimuth and polar angles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdb19f3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "figs, axs = plot_with_origin(1, num_objects)\n",
    "for i in range(num_objects):    \n",
    "    mask = torch.where(localizer.recv_edges == i)[0]\n",
    "    for mi in mask:\n",
    "        axs[0, i].plot(edge_attr[0, mi, :, 0], edge_attr[0, mi, :, 1], color=colors[localizer.send_edges[mi]])\n",
    "        axs[0, i].quiver(edge_attr[0, mi, :, 0], edge_attr[0, mi, :, 1],\n",
    "                         edge_attr[0, mi, :, 5], edge_attr[0, mi, :, 6], angles='xy', color='k')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d5df884",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_timesteps = x0.size(2)\n",
    "figs, axs = plot_with_origin(num_objects, num_timesteps, basesize=3)\n",
    "\n",
    "for i in range(num_objects):    \n",
    "    mask = torch.where(localizer.recv_edges == i)[0]\n",
    "    for t in range(num_timesteps):\n",
    "        for mi in mask:\n",
    "            axs[i, t].plot(edge_attr[0, mi, t, 0], edge_attr[0, mi, t, 1], 'o', color=colors[localizer.send_edges[mi]])\n",
    "            axs[i, t].quiver(edge_attr[0, mi, t, 0], edge_attr[0, mi, t, 1],\n",
    "                         edge_attr[0, mi, t, 5], edge_attr[0, mi, t, 6], angles='xy', color='k')"
   ]
  }
 ],
 "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.9.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}