utils.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:22.860018Z",
     "start_time": "2020-08-06T18:00:22.858105Z"
    }
   },
   "outputs": [],
   "source": [
    "# default_exp utils"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This contains general utilities used in different modules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.919997Z",
     "start_time": "2020-08-06T18:00:22.863204Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "import hashlib\n",
    "import json\n",
    "import math\n",
    "import os\n",
    "import re\n",
    "import time\n",
    "import warnings\n",
    "from pathlib import Path\n",
    "\n",
    "import ipykernel\n",
    "import nbdev.export\n",
    "import numpy as np\n",
    "import requests\n",
    "import seaborn as sns\n",
    "import torch\n",
    "from IPython.display import Javascript, display\n",
    "from notebook.notebookapp import list_running_servers\n",
    "from PIL import Image\n",
    "from shapely.geometry import Point\n",
    "from shapely.geometry.polygon import Polygon\n",
    "from torchvision import transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.924515Z",
     "start_time": "2020-08-06T18:00:23.921552Z"
    }
   },
   "outputs": [],
   "source": [
    "from IPython.core.debugger import set_trace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# numpy/torch conversion stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In general I'd like to have functions which work for both numpy and torch since the APIs aren't exactly the same. The approach I've taken is to write the function in `torch` (if possible) then add a function decorator `numpyify` to allow it to work with `numpy` arrays. This approach is far from perfect but it seems to work for most my cases."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`args_loop` assumes arguments are grouped in tuples or dictionaries, which is mostly true as `*args` and `**kwargs` are tuples and dictionaries, respectively. This will loop over and apply `callback` to each argument. Again, not perfect, but seems to work for most cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.937859Z",
     "start_time": "2020-08-06T18:00:23.926323Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def args_loop(args, callback):\n",
    "    if isinstance(args, tuple):  return tuple(args_loop(arg, callback) for arg in args)\n",
    "    elif isinstance(args, dict): return {key: args_loop(arg, callback) for key,arg in args.items()}\n",
    "    else:                        return callback(args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`Formatter` is a base class which will format input arguments; it also keeps track if any arguments were successfully formatted. Again, not perfect but works in most cases."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.945410Z",
     "start_time": "2020-08-06T18:00:23.939336Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Formatter():\n",
    "    def __init__(self): self.formatted = False\n",
    "    \n",
    "    def predicate(self, arg): raise NotImplementedError('Please implement predicate')\n",
    "    def formatter(self, arg): raise NotImplementedError('Please implement formatter')\n",
    "        \n",
    "    def callback(self, arg):\n",
    "        if self.predicate(arg): \n",
    "            self.formatted = True\n",
    "            return self.formatter(arg)\n",
    "        else:                   \n",
    "            return arg\n",
    "        \n",
    "    def __call__(self, args):\n",
    "        self.formatted = False\n",
    "        return args_loop(args, self.callback)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`torch2np` will convert arguments from torch to numpy arrays. It should work for all cases (tensors which require gradients and cuda tensors) and will use the same underlying data and `dtype`.\n",
    "\n",
    "NOTE: The `formatted` flag for `torch2np` will not be thread safe, calling `torch2np = Torch2np()` per invocation should make it safe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.951452Z",
     "start_time": "2020-08-06T18:00:23.946295Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Torch2np(Formatter):\n",
    "    def predicate(self, arg): return isinstance(arg, torch.Tensor)\n",
    "    def formatter(self, arg): return arg.detach().cpu().numpy()\n",
    "torch2np = Torch2np()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`np2torch` will convert to tensor with the same `dtype`. It's not as general as `torch2np` since output tensor might need to be on the gpu.\n",
    "\n",
    "NOTE: The `formatted` flag for `np2torch` will not be thread safe, calling `np2torch = Np2Torch()` per invocation should make it safe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.957208Z",
     "start_time": "2020-08-06T18:00:23.952309Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "class Np2torch(Formatter):\n",
    "    def predicate(self, arg): return isinstance(arg, np.ndarray)\n",
    "    def formatter(self, arg): return torch.from_numpy(arg)\n",
    "np2torch = Np2torch()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`numpyify` decorator will allow functions designed for torch tensors to also work for numpy tensors. \n",
    "\n",
    "NOTE: this will fail if input does not contain a `torch.tensor` (i.e. like the `eye` function, which takes an integer and returns a tensor). To ensure this works, one of the inputs must be a `torch.tensor`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.962877Z",
     "start_time": "2020-08-06T18:00:23.958180Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def numpyify(f):\n",
    "    def _numpyify(*args, **kwargs):\n",
    "        np2torch = Np2torch()                      # For thread safety, make a local instantiation\n",
    "        args, kwargs = np2torch((args, kwargs))\n",
    "        out = f(*args, **kwargs)\n",
    "        if np2torch.formatted: out = torch2np(out) # Convert back\n",
    "        return out\n",
    "    return _numpyify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tests"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`assert_allclose` checks if two things, `A` and `B`, are close to each other. I use `np.allclose` because its more robust than `torch.allclose` (i.e. numpys version will work with datatypes other than `np.array`s, like `bool`s and `int`s)\n",
    "\n",
    "NOTE: I'm assuming the format of the inputs is the same; if not I'm assuming this is programmer error."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.971599Z",
     "start_time": "2020-08-06T18:00:23.964517Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def _assert_allclose(A, B, **kwargs):\n",
    "    if isinstance(A, tuple):            \n",
    "        for a,b in zip(A,B): _assert_allclose(a, b, **kwargs) # Possibly add \"strict\" keyword here\n",
    "    elif isinstance(A, dict):           \n",
    "        for key in A.keys() | B.keys(): _assert_allclose(A[key], B[key], **kwargs)\n",
    "    else:\n",
    "        try:    assert(np.allclose(A, B, **kwargs))\n",
    "        except: assert(np.all(A == B))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.977258Z",
     "start_time": "2020-08-06T18:00:23.973549Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def assert_allclose(A, B, **kwargs): _assert_allclose(*torch2np((A, B)), **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.985992Z",
     "start_time": "2020-08-06T18:00:23.978256Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[0.6037, 0.6749, 0.5302],\n",
       "         [0.1848, 0.6734, 0.6051],\n",
       "         [0.8319, 0.7146, 0.7991],\n",
       "         [0.7886, 0.8095, 0.5742]]),\n",
       " array([[-0.59408999,  2.28263673,  0.90287216],\n",
       "        [ 1.79793451,  1.01680026, -0.49833506],\n",
       "        [ 0.11928329, -1.87143434, -0.70390851],\n",
       "        [-1.42716622, -0.09061538,  0.28871047]]))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = torch.rand((4,3))\n",
    "B = np.random.normal(size=(4,3))\n",
    "A, B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.990677Z",
     "start_time": "2020-08-06T18:00:23.986978Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(A, A+1e-5, atol=1e-5)\n",
    "assert_allclose((A, (B, {'test': 1.})), (A+1e-5, (B+1e-5, {'test': 1 + 1e-5})), atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we have multiple functions that should work for torch and numpy, we should have a way to test for both without having to write duplicate tests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:23.996311Z",
     "start_time": "2020-08-06T18:00:23.991874Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def assert_allclose_f(f, x, y, **kwargs):\n",
    "    if not isinstance(x, tuple): x = (x,)\n",
    "    assert_allclose(f(*x), y, **kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.002307Z",
     "start_time": "2020-08-06T18:00:23.997449Z"
    }
   },
   "outputs": [],
   "source": [
    "# export \n",
    "def assert_allclose_f_ttn(f, x, y, **kwargs): # ttn == \"torch, then numpy\"\n",
    "    torch2np = Torch2np()\n",
    "    assert_allclose_f(f, x, y, **kwargs) # Torch test\n",
    "    x, y = torch2np((x, y))\n",
    "    assert(torch2np.formatted)           # Make sure something was converted\n",
    "    assert_allclose_f(f, x, y, **kwargs) # Numpy test"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For some reason I can't find a built-in that will reverse and return a list without doing some iterable thing.\n",
    "\n",
    "TODO: get this to work for torch tensors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.008591Z",
     "start_time": "2020-08-06T18:00:24.003472Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def reverse(A): return A[::-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.014376Z",
     "start_time": "2020-08-06T18:00:24.009712Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(reverse(['a', 'b', 'c']), ['c', 'b', 'a'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`shape` returns the tensor shape as a tensor the same type of the tensor. Convenient if you need to do arithmetic based on the size of the tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.020350Z",
     "start_time": "2020-08-06T18:00:24.015442Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def shape(A, dtype=None):\n",
    "    if dtype is None: dtype = A.dtype\n",
    "    return A.new_tensor(A.shape, dtype=dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.026573Z",
     "start_time": "2020-08-06T18:00:24.021525Z"
    }
   },
   "outputs": [],
   "source": [
    "A = torch.rand(3,4)\n",
    "assert_allclose_f_ttn(shape, A, torch.FloatTensor([3, 4]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.032476Z",
     "start_time": "2020-08-06T18:00:24.028006Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def stackify(A, dim=0):\n",
    "    if isinstance(A, tuple): return torch.stack([stackify(a, dim) for a in A], dim)\n",
    "    return A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.038754Z",
     "start_time": "2020-08-06T18:00:24.033599Z"
    }
   },
   "outputs": [],
   "source": [
    "x = (tuple(torch.FloatTensor([1,2])), tuple(torch.FloatTensor([1,2])))\n",
    "assert_allclose_f_ttn(stackify, (x,), torch.FloatTensor([[1, 2],[1, 2]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "pytorch doesnt have an equivalent of `np.delete`\n",
    "\n",
    "NOTE: temporarily not `numpyified` because I need `np.array(,dtype=np.object)` to work since torch does not support jagged/nested tensors yet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.046789Z",
     "start_time": "2020-08-06T18:00:24.039922Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def delete(A, idx_delete):\n",
    "    idx = torch.ones(len(A), dtype=torch.bool)\n",
    "    idx[idx_delete] = False\n",
    "    return A[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.053423Z",
     "start_time": "2020-08-06T18:00:24.047849Z"
    }
   },
   "outputs": [],
   "source": [
    "A = torch.FloatTensor([[1,2,3],\n",
    "                       [4,5,6],\n",
    "                       [7,8,9]])\n",
    "assert_allclose_f_ttn(delete, (A, 1), torch.FloatTensor([[1, 2, 3],\n",
    "                                                         [7, 8, 9]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.059325Z",
     "start_time": "2020-08-06T18:00:24.054377Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def rescale(A, r1, r2): return (A-r1[0])/(r1[1]-r1[0])*(r2[1]-r2[0])+r2[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.065729Z",
     "start_time": "2020-08-06T18:00:24.061846Z"
    }
   },
   "outputs": [],
   "source": [
    "A = torch.FloatTensor([1,2])\n",
    "assert_allclose_f_ttn(rescale, (A, torch.FloatTensor([1,2]), torch.FloatTensor([2,4])), \n",
    "                      torch.FloatTensor([2, 4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Points/Vectors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-19T19:46:17.283758Z",
     "start_time": "2020-07-19T19:46:17.280117Z"
    }
   },
   "source": [
    "Kind of hard to keep separate distinction between points and vectors even though they technically aren't the same thing."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General point stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`singlify` decorator will allow functions designed for multiple points to work for single point inputs. points should have a shape of `(N, [2,3])` whereas a single point will have shape of `([2,3])`, so its convenient to just define the function for multiple points and use the decorator so it will work for a single point.\n",
    "\n",
    "NOTE: first argument must be `ps`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.071507Z",
     "start_time": "2020-08-06T18:00:24.067364Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def singlify(f):\n",
    "    def _singlify(ps, *args, **kwargs):\n",
    "        single = len(ps.shape) == 1\n",
    "        if single: ps = ps[None]\n",
    "        ps = f(ps, *args, **kwargs)\n",
    "        if single: ps = ps[0]\n",
    "        return ps\n",
    "    return _singlify"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`augment` will add ones to points; useful for affine and homography xforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.077435Z",
     "start_time": "2020-08-06T18:00:24.072397Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "@singlify\n",
    "def augment(ps): return torch.cat([ps, ps.new_ones((len(ps), 1))], dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.084110Z",
     "start_time": "2020-08-06T18:00:24.078476Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1940, 0.2536],\n",
    "                        [0.2172, 0.1626],\n",
    "                        [0.9834, 0.2700],\n",
    "                        [0.5324, 0.7137]])\n",
    "assert_allclose_f_ttn(augment, ps, torch.FloatTensor([[0.1940, 0.2536, 1.0000],\n",
    "                                                      [0.2172, 0.1626, 1.0000],\n",
    "                                                      [0.9834, 0.2700, 1.0000],\n",
    "                                                      [0.5324, 0.7137, 1.0000]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`deaugment` will remove last column; might wanna add check to make sure column contains ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.090251Z",
     "start_time": "2020-08-06T18:00:24.085019Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@singlify\n",
    "def deaugment(ps): return ps[:, 0:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.096870Z",
     "start_time": "2020-08-06T18:00:24.091365Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1940, 0.2536, 1.0000],\n",
    "                        [0.2172, 0.1626, 1.0000],\n",
    "                        [0.9834, 0.2700, 1.0000],\n",
    "                        [0.5324, 0.7137, 1.0000]])\n",
    "assert_allclose_f_ttn(deaugment, ps, torch.FloatTensor([[0.1940, 0.2536],\n",
    "                                                        [0.2172, 0.1626],\n",
    "                                                        [0.9834, 0.2700],\n",
    "                                                        [0.5324, 0.7137]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`normalize` will divide by last column and remove it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.107605Z",
     "start_time": "2020-08-06T18:00:24.097990Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@singlify\n",
    "def normalize(ps): return deaugment(ps/ps[:, [-1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.117059Z",
     "start_time": "2020-08-06T18:00:24.108607Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1940, 0.2536, 2.0000],\n",
    "                        [0.2172, 0.1626, 3.0000],\n",
    "                        [0.9834, 0.2700, 4.0000],\n",
    "                        [0.5324, 0.7137, 5.0000]])\n",
    "assert_allclose_f_ttn(normalize, ps, torch.FloatTensor([[0.0970, 0.1268],\n",
    "                                                        [0.0724, 0.0542],\n",
    "                                                        [0.2458, 0.0675],\n",
    "                                                        [0.1065, 0.1427]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-19T18:07:45.252228Z",
     "start_time": "2020-07-19T18:07:45.249778Z"
    }
   },
   "source": [
    "## Bounding box stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-26T00:03:08.202947Z",
     "start_time": "2020-05-26T00:03:08.198828Z"
    }
   },
   "source": [
    "`ps_bb` is points bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.128856Z",
     "start_time": "2020-08-06T18:00:24.118085Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def ps_bb(ps): return stackify((torch.min(ps, dim=0).values, torch.max(ps, dim=0).values))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.135121Z",
     "start_time": "2020-08-06T18:00:24.129795Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1940, 0.2536],\n",
    "                        [0.2172, 0.1626],\n",
    "                        [0.9834, 0.2700],\n",
    "                        [0.5324, 0.7137]])\n",
    "assert_allclose_f_ttn(ps_bb, ps, torch.FloatTensor([[0.194 , 0.1626],\n",
    "                                                    [0.9834, 0.7137]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`array_bb` is array bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.140438Z",
     "start_time": "2020-08-06T18:00:24.136027Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def array_bb(arr, dtype=None):\n",
    "    if dtype is None: dtype = arr.dtype\n",
    "    return arr.new_tensor([[0,0], [arr.shape[1]-1, arr.shape[0]-1]], dtype=dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.147195Z",
     "start_time": "2020-08-06T18:00:24.141913Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.rand(5,4)\n",
    "assert_allclose_f_ttn(array_bb, arr, torch.FloatTensor([[0, 0],\n",
    "                                                        [3, 4]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bb_sz` returns the size of a bounding box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.153045Z",
     "start_time": "2020-08-06T18:00:24.148504Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def bb_sz(bb): return stackify((bb[1,1]-bb[0,1]+1, bb[1,0]-bb[0,0]+1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.159626Z",
     "start_time": "2020-08-06T18:00:24.153917Z"
    }
   },
   "outputs": [],
   "source": [
    "bb = torch.FloatTensor([[0,0],[5,4]])\n",
    "assert_allclose_f_ttn(bb_sz, bb, torch.FloatTensor([5, 6]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bb_grid` is bounding box grid; i,j is swapped to x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.167158Z",
     "start_time": "2020-08-06T18:00:24.160660Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def bb_grid(bb, dtype=None):\n",
    "    if dtype is None: dtype = bb.dtype\n",
    "    return stackify(reverse(torch.meshgrid(torch.arange(bb[0,1], bb[1,1]+1, dtype=dtype, device=bb.device), \n",
    "                                           torch.arange(bb[0,0], bb[1,0]+1, dtype=dtype, device=bb.device))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.174178Z",
     "start_time": "2020-08-06T18:00:24.168169Z"
    }
   },
   "outputs": [],
   "source": [
    "bb = torch.FloatTensor([[0,0],[2,1]])\n",
    "assert_allclose_f_ttn(bb_grid, bb, torch.FloatTensor([[[0, 1, 2],\n",
    "                                                       [0, 1, 2]],\n",
    "                                                      [[0, 0, 0],\n",
    "                                                       [1, 1, 1]]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`bb_array` applies bounding box to array and returns the sub array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.179461Z",
     "start_time": "2020-08-06T18:00:24.175497Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def bb_array(arr, bb): \n",
    "    bb = bb.long() # Must be long for indexing; cannot round either incase input bb is type long\n",
    "    return arr[bb[0,1]:bb[1,1]+1, bb[0,0]:bb[1,0]+1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.185927Z",
     "start_time": "2020-08-06T18:00:24.180889Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[1,2,3],[4,5,6],[7,8,9]])\n",
    "bb = torch.FloatTensor([[1,1],[2,2]])\n",
    "assert_allclose_f_ttn(bb_array, (arr, bb), torch.FloatTensor([[5, 6],[8, 9]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.191438Z",
     "start_time": "2020-08-06T18:00:24.186994Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def is_p_in_bb(p, bb): return p[0] >= bb[0,0] and p[1] >= bb[0,1] and p[0] <= bb[1,0] and p[1] <= bb[1,1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.198305Z",
     "start_time": "2020-08-06T18:00:24.192764Z"
    }
   },
   "outputs": [],
   "source": [
    "p1 = torch.FloatTensor([0.0, 0.0])\n",
    "p2 = torch.FloatTensor([1.5, 1.5])\n",
    "bb = torch.FloatTensor([[1,1], [2,2]])\n",
    "assert_allclose_f_ttn(is_p_in_bb, (p1, bb), False)\n",
    "assert_allclose_f_ttn(is_p_in_bb, (p2, bb), True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.203666Z",
     "start_time": "2020-08-06T18:00:24.199494Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def is_bb_in_bb(bb1, bb2): return is_p_in_bb(bb1[0], bb2) and is_p_in_bb(bb1[1], bb2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.210348Z",
     "start_time": "2020-08-06T18:00:24.204648Z"
    }
   },
   "outputs": [],
   "source": [
    "bb1 = torch.FloatTensor([[1,1],[5,5]])\n",
    "bb2 = torch.FloatTensor([[2,2],[4,4]])\n",
    "bb3 = torch.FloatTensor([[8,8],[9,9]])\n",
    "assert_allclose_f_ttn(is_bb_in_bb, (bb2, bb1), True)\n",
    "assert_allclose_f_ttn(is_bb_in_bb, (bb3, bb1), False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Boundary stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.215451Z",
     "start_time": "2020-08-06T18:00:24.211353Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def is_p_in_b(p, b): return Polygon(b).contains(Point(*p))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.222281Z",
     "start_time": "2020-08-06T18:00:24.216563Z"
    }
   },
   "outputs": [],
   "source": [
    "b  = torch.FloatTensor([[0,0],[0,1],[1,1],[1,0]])\n",
    "p1 = torch.FloatTensor([0.5, 0.5])\n",
    "p2 = torch.FloatTensor([1.5, 1.5])\n",
    "assert_allclose_f_ttn(is_p_in_b, (p1, b), True)\n",
    "assert_allclose_f_ttn(is_p_in_b, (p2, b), False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.229505Z",
     "start_time": "2020-08-06T18:00:24.223312Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def bb2b(bb): return bb[[[0,0],[0,1],[1,1],[1,0]],\n",
    "                        [[0,0],[0,1],[0,1],[0,0]]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.235468Z",
     "start_time": "2020-08-06T18:00:24.230536Z"
    }
   },
   "outputs": [],
   "source": [
    "bb = torch.FloatTensor([[1,1],[5,5]])\n",
    "assert_allclose_f_ttn(bb2b, bb, torch.FloatTensor([[1., 1.],\n",
    "                                                   [1., 5.],\n",
    "                                                   [5., 5.],\n",
    "                                                   [5., 1.]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Point geometries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`grid2ps` converts grid to points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.241614Z",
     "start_time": "2020-08-06T18:00:24.236535Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def grid2ps(X, Y, order='row'):\n",
    "    if   order == 'row': return stackify((X.flatten(), Y.flatten()), dim=1)\n",
    "    elif order == 'col': return grid2ps(X.T, Y.T, order='row')\n",
    "    else: raise RuntimeError(f'Unrecognized option: {order}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.248222Z",
     "start_time": "2020-08-06T18:00:24.242878Z"
    }
   },
   "outputs": [],
   "source": [
    "X = torch.FloatTensor([[1,2],[1,2]])\n",
    "Y = torch.FloatTensor([[1,1],[2,2]])\n",
    "assert_allclose_f_ttn(grid2ps, (X, Y, 'row'), torch.FloatTensor([[1, 1],\n",
    "                                                                 [2, 1],\n",
    "                                                                 [1, 2],\n",
    "                                                                 [2, 2]]))\n",
    "assert_allclose_f_ttn(grid2ps, (X, Y, 'col'), torch.FloatTensor([[1, 1],\n",
    "                                                                 [1, 2],\n",
    "                                                                 [2, 1],\n",
    "                                                                 [2, 2]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.253173Z",
     "start_time": "2020-08-06T18:00:24.249117Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def array_ps(arr, dtype=None):\n",
    "    if dtype is None: dtype = arr.dtype\n",
    "    return grid2ps(*bb_grid(array_bb(arr), dtype))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.260097Z",
     "start_time": "2020-08-06T18:00:24.254149Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.rand(3,2)\n",
    "assert_allclose_f_ttn(array_ps, arr, torch.FloatTensor([[0, 0],\n",
    "                                                        [1, 0],\n",
    "                                                        [0, 1],\n",
    "                                                        [1, 1],\n",
    "                                                        [0, 2],\n",
    "                                                        [1, 2]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`crrgrid` is centered rectanglular rectangle grid\n",
    "\n",
    "NOTE: not yet numpyified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.267537Z",
     "start_time": "2020-08-06T18:00:24.261020Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def crrgrid(num_h, num_w, spacing_h, spacing_w, dtype, device=None):\n",
    "    h, w = spacing_h*(num_h-1), spacing_w*(num_w-1)\n",
    "    y = torch.linspace(-h/2, h/2, int(num_h), dtype=dtype, device=device)\n",
    "    x = torch.linspace(-w/2, w/2, int(num_w), dtype=dtype, device=device)\n",
    "    return grid2ps(*reverse(torch.meshgrid(y, x)), 'col')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.274138Z",
     "start_time": "2020-08-06T18:00:24.268580Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose_f_ttn(crrgrid, (2, 2, 1, 1, torch.float), torch.FloatTensor([[-0.5000, -0.5000],\n",
    "                                                                             [-0.5000,  0.5000],\n",
    "                                                                             [ 0.5000, -0.5000],\n",
    "                                                                             [ 0.5000,  0.5000]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`csrgrid` is centered square rectangle grid\n",
    "\n",
    "NOTE: not yet numpyified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.279076Z",
     "start_time": "2020-08-06T18:00:24.275506Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def csrgrid(num_h, num_w, spacing, dtype, device=None):\n",
    "    return crrgrid(num_h, num_w, spacing, spacing, dtype, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.285953Z",
     "start_time": "2020-08-06T18:00:24.280363Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose_f_ttn(csrgrid, (2, 2, 1, torch.float), torch.FloatTensor([[-0.5, -0.5],\n",
    "                                                                          [-0.5,  0.5],\n",
    "                                                                          [ 0.5, -0.5],\n",
    "                                                                          [ 0.5,  0.5]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-08T21:15:26.373537Z",
     "start_time": "2020-07-08T21:15:26.364762Z"
    }
   },
   "source": [
    "`csdgrid` is centered square diamond grid\n",
    "\n",
    "NOTE: not yet numpyified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.295482Z",
     "start_time": "2020-08-06T18:00:24.287224Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def csdgrid(num_h, num_w, spacing, fo, dtype, device=None):\n",
    "    h, w = spacing*(num_h-1), spacing*(num_w-1)\n",
    "    xs_grid = torch.linspace(-w/2, w/2, int(num_w), dtype=dtype, device=device)\n",
    "    ys_grid = torch.linspace(-h/2, h/2, int(num_h), dtype=dtype, device=device)\n",
    "    ps = []\n",
    "    for x_grid in xs_grid:\n",
    "        if fo: ys, fo = ys_grid[0::2], False\n",
    "        else:  ys, fo = ys_grid[1::2], True\n",
    "        xs = torch.full((len(ys),), x_grid, dtype=dtype, device=device)\n",
    "        ps.append(stackify((xs, ys), dim=1))\n",
    "    return torch.cat(ps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.304204Z",
     "start_time": "2020-08-06T18:00:24.296860Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose_f_ttn(csdgrid, (5, 4, 1, True, torch.float), torch.FloatTensor([[-1.5, -2],\n",
    "                                                                                [-1.5,  0],\n",
    "                                                                                [-1.5,  2],\n",
    "                                                                                [-0.5, -1],\n",
    "                                                                                [-0.5,  1],\n",
    "                                                                                [ 0.5, -2],\n",
    "                                                                                [ 0.5,  0],\n",
    "                                                                                [ 0.5,  2],\n",
    "                                                                                [ 1.5, -1],\n",
    "                                                                                [ 1.5,  1]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`cfpgrid` is centered four point grid\n",
    "\n",
    "NOTE: not yet numpyified"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.309486Z",
     "start_time": "2020-08-06T18:00:24.305152Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def cfpgrid(h, w, dtype, device=None): return crrgrid(2, 2, h, w, dtype, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.315717Z",
     "start_time": "2020-08-06T18:00:24.310527Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose_f_ttn(cfpgrid, (2 ,2, torch.float), torch.FloatTensor([[-1, -1],\n",
    "                                                                       [-1,  1],\n",
    "                                                                       [ 1, -1],\n",
    "                                                                       [ 1,  1]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## General vector stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-26T02:00:08.716743Z",
     "start_time": "2020-05-26T02:00:08.710841Z"
    }
   },
   "source": [
    "`unitize` will make norm of each vector 1\n",
    "\n",
    "NOTE: Handling of zero vector might be a little tricky"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.320779Z",
     "start_time": "2020-08-06T18:00:24.316618Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "@singlify\n",
    "def unitize(vs): return vs/torch.norm(vs, dim=1, keepdim=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.330276Z",
     "start_time": "2020-08-06T18:00:24.321690Z"
    }
   },
   "outputs": [],
   "source": [
    "vs = torch.FloatTensor([[0.1940, 0.2536],\n",
    "                        [0.2172, 0.1626],\n",
    "                        [0.9834, 0.2700],\n",
    "                        [0.5324, 0.7137]])\n",
    "assert_allclose_f_ttn(unitize, vs, torch.FloatTensor([[0.6075896, 0.7942511],\n",
    "                                                      [0.8005304, 0.5992921],\n",
    "                                                      [0.9643143, 0.2647598],\n",
    "                                                      [0.5979315, 0.8015471]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.337695Z",
     "start_time": "2020-08-06T18:00:24.331304Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def cross_mat(v):\n",
    "    zero = v.new_tensor(0)\n",
    "    \n",
    "    return stackify((( zero, -v[2],  v[1]),\n",
    "                     ( v[2],  zero, -v[0]),\n",
    "                     (-v[1],  v[0],  zero)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.344685Z",
     "start_time": "2020-08-06T18:00:24.338652Z"
    }
   },
   "outputs": [],
   "source": [
    "v = torch.FloatTensor([1,2,3])\n",
    "assert_allclose_f_ttn(cross_mat, v, torch.FloatTensor([[ 0,-3, 2],\n",
    "                                                       [ 3, 0,-1],\n",
    "                                                       [-2, 1, 0]]))\n",
    "assert_allclose(cross_mat(v)@v, 0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transforms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pmm` is point matrix multiplication"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.349276Z",
     "start_time": "2020-08-06T18:00:24.345884Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@singlify\n",
    "def pmm(ps, A, aug=False):\n",
    "    if aug: ps = augment(ps)\n",
    "    ps = (A@ps.T).T\n",
    "    if aug: ps = normalize(ps) # works for both affine and homography transforms\n",
    "    return ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.355760Z",
     "start_time": "2020-08-06T18:00:24.350319Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.1940, 0.2536],\n",
    "                        [0.2172, 0.1626],\n",
    "                        [0.9834, 0.2700],\n",
    "                        [0.5324, 0.7137]])\n",
    "A = torch.FloatTensor([[0.9571, 0.5551],\n",
    "                       [0.8914, 0.2626]])\n",
    "assert_allclose_f_ttn(pmm, (ps, A), torch.FloatTensor([[0.3265, 0.2395],\n",
    "                                                       [0.2981, 0.2363],\n",
    "                                                       [1.0911, 0.9475],\n",
    "                                                       [0.9057, 0.6620]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Spherical coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.361685Z",
     "start_time": "2020-08-06T18:00:24.356789Z"
    }
   },
   "outputs": [],
   "source": [
    "@numpyify\n",
    "@singlify\n",
    "def cart2spherical(ps):\n",
    "    x, y, z = ps[:,0], ps[:,1], ps[:,2]\n",
    "    r = torch.sqrt(x**2 + y**2 + z**2)\n",
    "    theta = torch.atan2(torch.sqrt(x**2 + y**2), z)\n",
    "    phi = torch.atan2(y, x)\n",
    "    return stackify((r, theta, phi), 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.371003Z",
     "start_time": "2020-08-06T18:00:24.367751Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.4082785 , 0.1256695 , 0.54343185],\n",
    "                        [0.93478803, 0.40557636, 0.40224384],\n",
    "                        [0.48831708, 0.82743735, 0.68537884]])\n",
    "assert_allclose_f_ttn(cart2spherical, ps, torch.FloatTensor([[0.69123247, 0.66619615, 0.29860039],\n",
    "                                                             [1.09550032, 1.19482283, 0.40935945],\n",
    "                                                             [1.18019079, 0.95116431, 1.03764654]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.378683Z",
     "start_time": "2020-08-06T18:00:24.374648Z"
    }
   },
   "outputs": [],
   "source": [
    "@numpyify\n",
    "@singlify\n",
    "def spherical2cart(ps):\n",
    "    r, theta, phi = ps[:,0], ps[:,1], ps[:,2]\n",
    "    x = r*torch.sin(theta)*torch.cos(phi)\n",
    "    y = r*torch.sin(theta)*torch.sin(phi)\n",
    "    z = r*torch.cos(theta)\n",
    "    return stackify((x, y, z), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.385196Z",
     "start_time": "2020-08-06T18:00:24.379733Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.69123247, 0.66619615, 0.29860039],\n",
    "                        [1.09550032, 1.19482283, 0.40935945],\n",
    "                        [1.18019079, 0.95116431, 1.03764654]])\n",
    "assert_allclose_f_ttn(spherical2cart, ps, torch.FloatTensor([[0.4082785 , 0.1256695 , 0.5434318 ],\n",
    "                                                             [0.9347881 , 0.4055764 , 0.40224388],\n",
    "                                                             [0.4883171 , 0.82743734, 0.68537885]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.391417Z",
     "start_time": "2020-08-06T18:00:24.386317Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0.6933, 0.7242, 0.6667],\n",
       "        [0.9893, 0.9539, 0.8397],\n",
       "        [0.0376, 0.9858, 0.5857],\n",
       "        [0.7251, 0.7630, 0.4017]])"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ps = torch.rand(4,3)\n",
    "ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.396432Z",
     "start_time": "2020-08-06T18:00:24.392435Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(ps, spherical2cart(cart2spherical(ps)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Point conditioning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`condition_mat` is typically used to \"condition\" points to improve conditioning; its inverse is usually applied afterwards. It sets the mean of the points to zero and the average distance to `sqrt(2)`. I use the term \"condition\" here so I don't get confused with \"normalization\" which is used above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.402158Z",
     "start_time": "2020-08-06T18:00:24.397933Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def condition_mat(ps):\n",
    "    zero, one = ps.new_tensor(0), ps.new_tensor(1)\n",
    "\n",
    "    xs, ys = ps[:, 0], ps[:, 1]\n",
    "    mean_x, mean_y = xs.mean(), ys.mean()\n",
    "    s_m = math.sqrt(2)*len(ps)/(torch.sqrt((xs-mean_x)**2+(ys-mean_y)**2)).sum()\n",
    "    return stackify((( s_m, zero, -mean_x*s_m),\n",
    "                     (zero,  s_m, -mean_y*s_m),\n",
    "                     (zero, zero,         one)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.409421Z",
     "start_time": "2020-08-06T18:00:24.403416Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.5466, 0.0889],\n",
    "                        [0.1493, 0.6591],\n",
    "                        [0.5600, 0.0352],\n",
    "                        [0.7287, 0.5892]])\n",
    "assert_allclose_f_ttn(condition_mat, ps, torch.FloatTensor([[ 4.0950,  0.0000, -2.0317],\n",
    "                                                            [ 0.0000,  4.0950, -1.4050],\n",
    "                                                            [ 0.0000,  0.0000,  1.0000]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.416318Z",
     "start_time": "2020-08-06T18:00:24.410449Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def condition(ps):\n",
    "    T = condition_mat(ps)\n",
    "    return pmm(ps, T, aug=True), T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.423251Z",
     "start_time": "2020-08-06T18:00:24.417337Z"
    }
   },
   "outputs": [],
   "source": [
    "ps = torch.FloatTensor([[0.5466, 0.0889],\n",
    "                        [0.1493, 0.6591],\n",
    "                        [0.5600, 0.0352],\n",
    "                        [0.7287, 0.5892]])\n",
    "ps_cond, T = condition(ps)\n",
    "assert_allclose(ps_cond.mean(), 0, atol=1e-6)\n",
    "assert_allclose(ps_cond.norm(dim=1).mean(), math.sqrt(2), atol=1e-6)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Homographies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`homography` estimates a homography between two sets of points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.428525Z",
     "start_time": "2020-08-06T18:00:24.424199Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def homography(ps1, ps2):    \n",
    "    # Condition and augment points\n",
    "    (ps1_cond, T1), (ps2_cond, T2) = map(condition, [ps1, ps2])\n",
    "    ps1_cond, ps2_cond = map(augment, [ps1_cond, ps2_cond])\n",
    "\n",
    "    # Form homogeneous system\n",
    "    L = torch.cat([torch.cat([ps1_cond, torch.zeros_like(ps1_cond), -ps2_cond[:, 0:1]*ps1_cond], dim=1),\n",
    "                   torch.cat([torch.zeros_like(ps1_cond), ps1_cond, -ps2_cond[:, 1:2]*ps1_cond], dim=1)])\n",
    "\n",
    "    # Solution is the last column of V\n",
    "    H12_cond = torch.svd(L, some=False).V[:,-1].reshape(3,3)\n",
    "\n",
    "    # Undo conditioning\n",
    "    H12 = torch.inverse(T2)@H12_cond@T1\n",
    "    H12 = H12/H12[2,2] # Sets H12[2,2] to 1\n",
    "    return H12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.436341Z",
     "start_time": "2020-08-06T18:00:24.429515Z"
    }
   },
   "outputs": [],
   "source": [
    "ps1 = torch.FloatTensor([[-350, -350],\n",
    "                         [-350,  350],\n",
    "                         [ 350, -350],\n",
    "                         [ 350,  350]])\n",
    "ps2 = torch.FloatTensor([[ 970,  517],\n",
    "                         [ 156,  498],\n",
    "                         [ 973, 1317],\n",
    "                         [ 192, 1279]])\n",
    "assert_allclose_f_ttn(homography, (ps1, ps2), \n",
    "                      torch.FloatTensor([[ 6.252e-02, -1.117e+00,  5.6786e+02],\n",
    "                                         [ 1.183e+00, -8.049e-03,  9.1087e+02],\n",
    "                                         [ 6.000e-05,  3.650e-05,  1.0000e+00]]), atol=1e-3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rotations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`approx_R` gives the nearest rotational approximation to the input matrix (frobenius norm?). Note that for a proper rotation determinant must be +1, which is checked after."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.442410Z",
     "start_time": "2020-08-06T18:00:24.437339Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def approx_R(R):\n",
    "    one = R.new_tensor(1)\n",
    "    \n",
    "    [U,_,V] = torch.svd(R)\n",
    "    R = U@V.T\n",
    "    if not torch.isclose(torch.det(R), one):\n",
    "        R = R.new_full((3,3), math.nan)\n",
    "    return R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.449384Z",
     "start_time": "2020-08-06T18:00:24.443600Z"
    }
   },
   "outputs": [],
   "source": [
    "R = torch.FloatTensor([[0.0958, 0.8441, 0.2009],\n",
    "                       [0.7877, 0.9110, 0.9277],\n",
    "                       [0.3727, 0.7262, 0.1417]])\n",
    "assert_allclose_f_ttn(approx_R, R, torch.FloatTensor([[-0.5659,  0.8152,  0.1237],\n",
    "                                                      [ 0.5155,  0.2328,  0.8247],\n",
    "                                                      [ 0.6434,  0.5304, -0.5520]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-19T13:31:59.556732Z",
     "start_time": "2020-07-19T13:31:59.553682Z"
    }
   },
   "source": [
    "#### Euler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.454951Z",
     "start_time": "2020-08-06T18:00:24.450421Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def euler2R(euler):\n",
    "    s, c = torch.sin, torch.cos\n",
    "\n",
    "    e_x, e_y, e_z = euler\n",
    "    return stackify((\n",
    "        (c(e_y)*c(e_z), c(e_z)*s(e_x)*s(e_y) - c(e_x)*s(e_z), s(e_x)*s(e_z) + c(e_x)*c(e_z)*s(e_y)),\n",
    "        (c(e_y)*s(e_z), c(e_x)*c(e_z) + s(e_x)*s(e_y)*s(e_z), c(e_x)*s(e_y)*s(e_z) - c(e_z)*s(e_x)),\n",
    "        (      -s(e_y),                        c(e_y)*s(e_x),                        c(e_x)*c(e_y))\n",
    "    ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.461584Z",
     "start_time": "2020-08-06T18:00:24.456215Z"
    }
   },
   "outputs": [],
   "source": [
    "euler = torch.FloatTensor([0.2748, 0.0352, 0.4496])\n",
    "assert_allclose_f_ttn(euler2R, euler, torch.FloatTensor([[ 0.9001, -0.4097,  0.1484],\n",
    "                                                         [ 0.4343,  0.8710, -0.2297],\n",
    "                                                         [-0.0352,  0.2712,  0.9619]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.466750Z",
     "start_time": "2020-08-06T18:00:24.462624Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def R2euler(R):\n",
    "    return stackify((torch.atan2( R[2, 1], R[2, 2]),\n",
    "                     torch.atan2(-R[2, 0], torch.sqrt(R[0, 0]**2+R[1, 0]**2)),\n",
    "                     torch.atan2( R[1, 0], R[0, 0])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.475888Z",
     "start_time": "2020-08-06T18:00:24.467848Z"
    }
   },
   "outputs": [],
   "source": [
    "R = torch.FloatTensor([[ 0.9001, -0.4097,  0.1484],\n",
    "                       [ 0.4343,  0.8710, -0.2297],\n",
    "                       [-0.0352,  0.2712,  0.9619]])\n",
    "assert_allclose_f_ttn(R2euler, R, torch.FloatTensor([0.2748, 0.0352, 0.4496]), atol=1e-4) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ensure euler => R => euler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.481599Z",
     "start_time": "2020-08-06T18:00:24.476885Z"
    }
   },
   "outputs": [],
   "source": [
    "euler = torch.FloatTensor([0.2748, 0.0352, 0.4496])\n",
    "assert_allclose(R2euler(euler2R(euler)), euler)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Rodrigues"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "rodrigues conversion from:\n",
    "* https://www2.cs.duke.edu/courses/fall13/compsci527/notes/rodrigues.pdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.490282Z",
     "start_time": "2020-08-06T18:00:24.482707Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def rodrigues2R(r):\n",
    "    zero = r.new_tensor(0)\n",
    "    \n",
    "    theta = torch.norm(r)\n",
    "    if theta > math.pi: warnings.warn('Theta greater than pi')\n",
    "    if torch.isclose(theta, zero): return torch.eye(3, dtype=r.dtype, device=r.device)\n",
    "    u = r/theta\n",
    "    return torch.eye(3, dtype=r.dtype, device=r.device)*torch.cos(theta) + \\\n",
    "           (1-torch.cos(theta))*u[:,None]@u[:,None].T + \\\n",
    "           cross_mat(u)*torch.sin(theta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.497637Z",
     "start_time": "2020-08-06T18:00:24.491346Z"
    }
   },
   "outputs": [],
   "source": [
    "theta = math.pi/4\n",
    "k = torch.FloatTensor([ 0.8155,  0.0937, -0.5711])\n",
    "r = theta*k\n",
    "assert_allclose_f_ttn(rodrigues2R, r, torch.FloatTensor([[ 0.9019,  0.4262, -0.0702],\n",
    "                                                         [-0.3814,  0.7097, -0.5923],\n",
    "                                                         [-0.2027,  0.5610,  0.8026]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 89,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.504738Z",
     "start_time": "2020-08-06T18:00:24.498580Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def R2rodrigues(R):\n",
    "    zero, one, pi = R.new_tensor(0), R.new_tensor(1), R.new_tensor(math.pi)\n",
    "    \n",
    "    A = (R-R.T)/2\n",
    "    rho = A[[2,0,1],[1,2,0]]\n",
    "    s = torch.norm(rho)\n",
    "    c = (R.trace()-1)/2\n",
    "    if torch.isclose(s, zero) and torch.isclose(c, one): \n",
    "        r = R.new_zeros(3)\n",
    "    elif torch.isclose(s, zero) and torch.isclose(c, -one):\n",
    "        V = R + torch.eye(3, dtype=R.dtype, device=R.device)\n",
    "        v = V[:, torch.where(~torch.isclose(torch.norm(V, dim=0), zero))[0][0]] # Just get first non-zero\n",
    "        u = unitize(v)\n",
    "        def S_half(r):\n",
    "            if torch.isclose(torch.norm(r), pi) and \\\n",
    "               ((torch.isclose(r[0], zero) and torch.isclose(r[1], zero) and r[2] < 0) or \\\n",
    "                (torch.isclose(r[0], zero) and r[1] < 0) or \\\n",
    "                (r[0] < 0)):\n",
    "                return -r\n",
    "            else: \n",
    "                return r\n",
    "        r = S_half(u*math.pi)\n",
    "    elif not torch.isclose(s, zero):\n",
    "        u = rho/s\n",
    "        theta = torch.atan2(s,c)\n",
    "        r = u*theta\n",
    "    else: raise RuntimeError('This shouldnt happen; please debug')\n",
    "    return r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.513979Z",
     "start_time": "2020-08-06T18:00:24.505776Z"
    }
   },
   "outputs": [],
   "source": [
    "R = torch.FloatTensor([[ 0.9019,  0.4262, -0.0702],\n",
    "                       [-0.3814,  0.7097, -0.5923],\n",
    "                       [-0.2027,  0.5610,  0.8026]])\n",
    "assert_allclose_f_ttn(R2rodrigues, R, torch.FloatTensor([ 0.6405,  0.0736, -0.4485]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.519645Z",
     "start_time": "2020-08-06T18:00:24.514980Z"
    }
   },
   "outputs": [],
   "source": [
    "r = torch.FloatTensor([.1,.2,.3])\n",
    "assert_allclose(R2rodrigues(rodrigues2R(r)), r)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Other rotation stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`approx_R` gives the nearest rotational approximation to the input matrix. Note that for a proper rotation determinant must be +1, which is checked after."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.525111Z",
     "start_time": "2020-08-06T18:00:24.520566Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def approx_R(R):\n",
    "    one = R.new_tensor(1)\n",
    "    \n",
    "    [U,_,V] = torch.svd(R)\n",
    "    R = U@V.T\n",
    "    if not torch.isclose(torch.det(R), one):\n",
    "        R = R.new_full((3,3), math.nan)\n",
    "    return R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.531963Z",
     "start_time": "2020-08-06T18:00:24.526244Z"
    }
   },
   "outputs": [],
   "source": [
    "R = torch.FloatTensor([[0.0958, 0.8441, 0.2009],\n",
    "                       [0.7877, 0.9110, 0.9277],\n",
    "                       [0.3727, 0.7262, 0.1417]])\n",
    "assert_allclose_f_ttn(approx_R, R, torch.FloatTensor([[-0.5659,  0.8152,  0.1237],\n",
    "                                                      [ 0.5155,  0.2328,  0.8247],\n",
    "                                                      [ 0.6434,  0.5304, -0.5520]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rigid"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Rigid transforms are a rotation followed by translation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.539849Z",
     "start_time": "2020-08-06T18:00:24.532886Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def Rt2M(R, t):\n",
    "    M = torch.cat([R, t[:,None]], dim=1)\n",
    "    M = torch.cat([M, M.new_tensor([[0,0,0,1]])])\n",
    "    return M"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 95,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.546415Z",
     "start_time": "2020-08-06T18:00:24.541073Z"
    }
   },
   "outputs": [],
   "source": [
    "R = torch.FloatTensor([[-0.5659,  0.8152,  0.1237],\n",
    "                       [ 0.5155,  0.2328,  0.8247],\n",
    "                       [ 0.6434,  0.5304, -0.5520]])\n",
    "t = torch.FloatTensor([1,2,3])\n",
    "assert_allclose_f_ttn(Rt2M, (R,t), torch.FloatTensor([[-0.5659,  0.8152,  0.1237,  1.0000],\n",
    "                                                      [ 0.5155,  0.2328,  0.8247,  2.0000],\n",
    "                                                      [ 0.6434,  0.5304, -0.5520,  3.0000],\n",
    "                                                      [ 0.0000,  0.0000,  0.0000,  1.0000]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 96,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.551967Z",
     "start_time": "2020-08-06T18:00:24.547514Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def M2Rt(M): return M[0:3,0:3], M[0:3,3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.558740Z",
     "start_time": "2020-08-06T18:00:24.553053Z"
    }
   },
   "outputs": [],
   "source": [
    "M = torch.FloatTensor([[-0.5659,  0.8152,  0.1237,  1.0000],\n",
    "                       [ 0.5155,  0.2328,  0.8247,  2.0000],\n",
    "                       [ 0.6434,  0.5304, -0.5520,  3.0000],\n",
    "                       [ 0.0000,  0.0000,  0.0000,  1.0000]])\n",
    "assert_allclose_f_ttn(M2Rt, M, (torch.FloatTensor([[-0.5659,  0.8152,  0.1237],\n",
    "                                                   [ 0.5155,  0.2328,  0.8247],\n",
    "                                                   [ 0.6434,  0.5304, -0.5520]]), torch.FloatTensor([1,2,3])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 98,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.564260Z",
     "start_time": "2020-08-06T18:00:24.559724Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def invert_rigid(M):\n",
    "    R, t = M2Rt(M)\n",
    "    return Rt2M(R.T, -R.T@t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 99,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.570799Z",
     "start_time": "2020-08-06T18:00:24.565145Z"
    }
   },
   "outputs": [],
   "source": [
    "M = torch.FloatTensor([[-0.5659,  0.8152,  0.1237,  1.0000],\n",
    "                       [ 0.5155,  0.2328,  0.8247,  2.0000],\n",
    "                       [ 0.6434,  0.5304, -0.5520,  3.0000],\n",
    "                       [ 0.0000,  0.0000,  0.0000,  1.0000]])\n",
    "assert_allclose_f_ttn(invert_rigid, M, torch.inverse(M), atol=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 100,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.578414Z",
     "start_time": "2020-08-06T18:00:24.571753Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def mult_rigid(M1, M2):\n",
    "    R1, t1 = M2Rt(M1)\n",
    "    R2, t2 = M2Rt(M2)\n",
    "    return Rt2M(R1@R2, R1@t2+t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 101,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.587290Z",
     "start_time": "2020-08-06T18:00:24.579354Z"
    }
   },
   "outputs": [],
   "source": [
    "M = torch.FloatTensor([[-0.5659,  0.8152,  0.1237,  1.0000],\n",
    "                       [ 0.5155,  0.2328,  0.8247,  2.0000],\n",
    "                       [ 0.6434,  0.5304, -0.5520,  3.0000],\n",
    "                       [ 0.0000,  0.0000,  0.0000,  1.0000]])\n",
    "assert_allclose_f_ttn(mult_rigid, (M,M), M@M)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## More Vector stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Stuff placed here mainly because it requires stuff from transforms section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a random unit vector uniformly sampled over a sphere, if you randomly sample `theta` and `phi` directly, you will over sample near the poles. Instead, you can do an \"equal area\" projection of the sphere onto a cylinder (i.e. rectangle), uniformly sample the cylinder/rectangle, and then project the point back. More info here:\n",
    "* https://math.stackexchange.com/a/44691\n",
    "\n",
    "NOTE: `torch.rand` is random uniform between `[0,1)`, so it is rescaled as needed. Also, since `random_unit` doesnt take in a `tensor`, it cannot be `numpyified`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 102,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.592975Z",
     "start_time": "2020-08-06T18:00:24.588261Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def random_unit(dtype=None, device=None):\n",
    "    r = torch.ones(1, dtype=dtype, device=device)[0]\n",
    "    theta = torch.acos(rescale(torch.rand(1, dtype=dtype, device=device)[0], [0,1], [-1,1]))\n",
    "    phi = rescale(torch.rand(1, dtype=dtype, device=device)[0], [0,1], [0, 2*math.pi])\n",
    "    return spherical2cart(stackify((r, theta, phi)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.599247Z",
     "start_time": "2020-08-06T18:00:24.594201Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(torch.norm(random_unit()), 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`v_v_angle` is vector-vector angle and returns the angle between two vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.604758Z",
     "start_time": "2020-08-06T18:00:24.600205Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def v_v_angle(a, b):   \n",
    "    x = torch.dot(a,b)/(torch.norm(a)*torch.norm(b))\n",
    "    return torch.acos(x.clamp(-1, 1)) # Precision errors can make this go outside domain [-1,1] => NaNs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.611411Z",
     "start_time": "2020-08-06T18:00:24.605812Z"
    }
   },
   "outputs": [],
   "source": [
    "a = torch.FloatTensor([ 0.5568, -0.4851, -0.6743])\n",
    "b = torch.FloatTensor([-0.8482, -0.4175, -0.3260])\n",
    "assert_allclose_f_ttn(v_v_angle, (a, b), 1.6207424465344742, atol=1e-5)\n",
    "assert_allclose_f_ttn(v_v_angle, (a, a), 0, atol=1e-5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-19T17:45:09.147783Z",
     "start_time": "2020-07-19T17:45:09.144252Z"
    }
   },
   "source": [
    "`v_v_R` is vector vector rotation matrix and returns the rotation matrix between the two vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.616631Z",
     "start_time": "2020-08-06T18:00:24.612441Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def v_v_R(v1, v2):\n",
    "    zero = v1.new_tensor(0)\n",
    "\n",
    "    theta = v_v_angle(v1, v2)\n",
    "    if torch.isclose(theta, zero): return torch.eye(3, dtype=v1.dtype, device=v1.device)\n",
    "    v3 = unitize(torch.cross(v1, v2))\n",
    "    return rodrigues2R(theta*v3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 107,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.625182Z",
     "start_time": "2020-08-06T18:00:24.617648Z"
    }
   },
   "outputs": [],
   "source": [
    "v1 = torch.FloatTensor([ 0.5568, -0.4851, -0.6743])\n",
    "v2 = torch.FloatTensor([-0.8482, -0.4175, -0.3260])\n",
    "assert_allclose_f_ttn(v_v_R, (v1,v2), torch.FloatTensor([[-0.03390428,  0.54606884,  0.83705395],\n",
    "                                                         [-0.74174788,  0.54757324, -0.3872643 ],\n",
    "                                                         [-0.66982131, -0.63401291,  0.38648033]]))\n",
    "assert_allclose_f_ttn(v_v_R, (v1,v1), torch.FloatTensor([[1, 0, 0],\n",
    "                                                         [0, 1, 0],\n",
    "                                                         [0, 0, 1]]))\n",
    "assert_allclose(unitize(pmm(v1, v_v_R(v1,v2))), unitize(v2), atol=1e-4) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Line stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.630689Z",
     "start_time": "2020-08-06T18:00:24.626204Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def pm2l(p, m):\n",
    "    zero, one = p.new_tensor(0), p.new_tensor(1)\n",
    "    \n",
    "    x, y = p\n",
    "    if not torch.isfinite(m): a, b, c = one, zero,    -x\n",
    "    else:                     a, b, c =   m, -one, y-m*x\n",
    "    return stackify((a,b,c))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.637320Z",
     "start_time": "2020-08-06T18:00:24.631570Z"
    }
   },
   "outputs": [],
   "source": [
    "p = torch.FloatTensor([1.5, 2.5])\n",
    "m = torch.FloatTensor([0.5])[0]\n",
    "assert_allclose_f_ttn(pm2l, (p, m), torch.FloatTensor([ 0.5000, -1.0000,  1.7500]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 110,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.642966Z",
     "start_time": "2020-08-06T18:00:24.638193Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def ps2l(p1, p2):\n",
    "    zero = p1.new_tensor(0)\n",
    "\n",
    "    x1, y1 = p1\n",
    "    x2, y2 = p2\n",
    "    if not torch.isclose(x2-x1, zero): m = (y2-y1)/(x2-x1)\n",
    "    else:                              m = p1.new_tensor(math.inf)\n",
    "    return pm2l(p1, m)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.649989Z",
     "start_time": "2020-08-06T18:00:24.644023Z"
    }
   },
   "outputs": [],
   "source": [
    "p1 = torch.FloatTensor([1.5, 2.5])\n",
    "p2 = torch.FloatTensor([2.5, 3.5])\n",
    "assert_allclose_f_ttn(ps2l, (p1, p2), torch.FloatTensor([ 1.0000,  -1.0000, 1.000]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.655264Z",
     "start_time": "2020-08-06T18:00:24.651055Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def pld(p, l):\n",
    "    x, y = p\n",
    "    a, b, c = l\n",
    "    return torch.abs(a*x + b*y + c)/torch.sqrt(a**2 + b**2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.662207Z",
     "start_time": "2020-08-06T18:00:24.656210Z"
    }
   },
   "outputs": [],
   "source": [
    "p = torch.FloatTensor([0.5, 1.5])\n",
    "l = torch.FloatTensor([1.0000,  0.0000, -1.5000])\n",
    "assert_allclose_f_ttn(pld, (p, l), 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.667580Z",
     "start_time": "2020-08-06T18:00:24.663271Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def l_l_intersect(l1, l2):\n",
    "    a1, b1, c1 = l1\n",
    "    a2, b2, c2 = l2\n",
    "    return stackify(((-c1*b2 + b1*c2)/(a1*b2 - b1*a2), (-a1*c2 + c1*a2)/(a1*b2 - b1*a2)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.674412Z",
     "start_time": "2020-08-06T18:00:24.668550Z"
    }
   },
   "outputs": [],
   "source": [
    "l1 = torch.FloatTensor([.1, .2, .3])\n",
    "l2 = torch.FloatTensor([.4, .2, .2])\n",
    "assert_allclose_f_ttn(l_l_intersect, (l1, l2), torch.FloatTensor([ 0.3333, -1.6667]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`b_ls` gets lines of boundary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.682060Z",
     "start_time": "2020-08-06T18:00:24.675589Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def b_ls(b): \n",
    "    return stackify(tuple(ps2l(b[idx], b[torch.remainder(idx+1, len(b))]) for idx in torch.arange(len(b))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.690130Z",
     "start_time": "2020-08-06T18:00:24.683249Z"
    }
   },
   "outputs": [],
   "source": [
    "b = torch.FloatTensor([[1., 1.],\n",
    "                       [1., 4.],\n",
    "                       [5., 4.],\n",
    "                       [5., 1.]])\n",
    "assert_allclose_f_ttn(b_ls, b, torch.FloatTensor([[ 1.,  0., -1.],\n",
    "                                                  [ 0., -1.,  4.],\n",
    "                                                  [ 1.,  0., -5.],\n",
    "                                                  [-0., -1.,  1.]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "TODO: handle edge cases"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.695108Z",
     "start_time": "2020-08-06T18:00:24.691410Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def b_l_intersect(b, l):\n",
    "    ps = []\n",
    "    for idx in torch.arange(len(b)):\n",
    "        p1, p2 = b[idx], b[torch.remainder(idx+1, len(b))]\n",
    "        p = l_l_intersect(ps2l(p1, p2), l)\n",
    "        d, d1, d2 = torch.norm(p2-p1), torch.norm(p1-p), torch.norm(p2-p)\n",
    "        if torch.isclose(d, d1+d2): ps.append(p)\n",
    "    return stackify(tuple(ps))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.704085Z",
     "start_time": "2020-08-06T18:00:24.696146Z"
    }
   },
   "outputs": [],
   "source": [
    "b = torch.FloatTensor([[-100., -100.],\n",
    "                       [-100.,  200.],\n",
    "                       [ 200.,  200.],\n",
    "                       [ 200., -100.]])\n",
    "l = torch.FloatTensor([.1, .2, .3])\n",
    "assert_allclose_f_ttn(b_l_intersect, (b, l), torch.FloatTensor([[-100.0000,   48.5000],\n",
    "                                                                [ 197.0000, -100.0000]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Ellipse stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`sample_2pi` prevents accidentally resampling 2pi twice by linspacing with an additional sample and then removing the last sample\n",
    "\n",
    "NOTE: Get numpy version working"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.707491Z",
     "start_time": "2020-08-06T18:00:24.705220Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def sample_2pi(num_samples, dtype=None, device=None): \n",
    "    return torch.linspace(0, 2*math.pi, int(num_samples)+1, dtype=dtype, device=device)[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.713775Z",
     "start_time": "2020-08-06T18:00:24.708372Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose_f_ttn(sample_2pi, 3, torch.FloatTensor([0.0000, 2.0944, 4.1888]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.720661Z",
     "start_time": "2020-08-06T18:00:24.714805Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def sample_ellipse(e, num_samples):\n",
    "    h, k, a, b, alpha = e\n",
    "    thetas = sample_2pi(num_samples, e.dtype, e.device)\n",
    "    return stackify((a*torch.cos(alpha)*torch.cos(thetas) - b*torch.sin(alpha)*torch.sin(thetas) + h,\n",
    "                     a*torch.sin(alpha)*torch.cos(thetas) + b*torch.cos(alpha)*torch.sin(thetas) + k), dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.729454Z",
     "start_time": "2020-08-06T18:00:24.721759Z"
    }
   },
   "outputs": [],
   "source": [
    "e = torch.FloatTensor([1, 2, 3, 4, math.pi/4])\n",
    "assert_allclose_f_ttn(sample_ellipse, (e, 3), torch.FloatTensor([[ 3.1213,  4.1213],\n",
    "                                                                 [-2.5101,  3.3888],\n",
    "                                                                 [ 2.3888, -1.5101]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.734986Z",
     "start_time": "2020-08-06T18:00:24.730616Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def ellipse2conic(e):\n",
    "    h, k, a, b, alpha = e\n",
    "    A = a**2*torch.sin(alpha)**2 + b**2*torch.cos(alpha)**2\n",
    "    B = 2*(b**2 - a**2)*torch.sin(alpha)*torch.cos(alpha)\n",
    "    C = a**2*torch.cos(alpha)**2 + b**2*torch.sin(alpha)**2\n",
    "    D = -2*A*h - B*k\n",
    "    E = -B*h - 2*C*k\n",
    "    F = A*h**2 + B*h*k + C*k**2 - a**2*b**2\n",
    "    return stackify(((  A, B/2, D/2),\n",
    "                     (B/2,   C, E/2),\n",
    "                     (D/2, E/2,   F)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 125,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.742342Z",
     "start_time": "2020-08-06T18:00:24.736023Z"
    }
   },
   "outputs": [],
   "source": [
    "e = torch.FloatTensor([1, 2, 3, 4, math.pi/4])\n",
    "assert_allclose_f_ttn(ellipse2conic, e, torch.FloatTensor([[ 12.5000,   3.5000, -19.5000],\n",
    "                                                           [  3.5000,  12.5000, -28.5000],\n",
    "                                                           [-19.5000, -28.5000, -67.5000]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 126,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.751984Z",
     "start_time": "2020-08-06T18:00:24.743278Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def conic2ellipse(Aq):\n",
    "    zero, pi = Aq.new_tensor(0), Aq.new_tensor(math.pi)\n",
    "    \n",
    "    A = Aq[0, 0]\n",
    "    B = 2*Aq[0, 1]\n",
    "    C = Aq[1, 1]\n",
    "    D = 2*Aq[0, 2]\n",
    "    E = 2*Aq[1, 2]\n",
    "    F = Aq[2, 2]\n",
    "\n",
    "    # Return nans if input conic is not ellipse\n",
    "    if torch.any(~torch.isfinite(Aq)) or torch.isclose(B**2-4*A*C, zero) or B**2-4*A*C > 0:\n",
    "        return Aq.new_full((5,), math.nan)\n",
    "\n",
    "    # Equations below are from https://math.stackexchange.com/a/820896/39581\n",
    "\n",
    "    # \"coefficient of normalizing factor\"\n",
    "    q = 64*(F*(4*A*C-B**2)-A*E**2+B*D*E-C*D**2)/(4*A*C-B**2)**2\n",
    "\n",
    "    # distance between center and focal point\n",
    "    s = 1/4*torch.sqrt(torch.abs(q)*torch.sqrt(B**2+(A-C)**2))\n",
    "\n",
    "    # ellipse parameters\n",
    "    h = (B*E-2*C*D)/(4*A*C-B**2)\n",
    "    k = (B*D-2*A*E)/(4*A*C-B**2)\n",
    "    a = 1/8*torch.sqrt(2*torch.abs(q)*torch.sqrt(B**2+(A-C)**2)-2*q*(A+C))\n",
    "    b = torch.sqrt(a**2-s**2)\n",
    "    # Get alpha; note that range of alpha is [0, pi)\n",
    "    if torch.isclose(q*A-q*C, zero) and torch.isclose(q*B, zero): alpha = zero # Circle\n",
    "    elif torch.isclose(q*A-q*C, zero) and q*B > 0:                alpha = 1/4*pi\n",
    "    elif torch.isclose(q*A-q*C, zero) and q*B < 0:                alpha = 3/4*pi\n",
    "    elif q*A-q*C > 0 and (torch.isclose(q*B, zero) or q*B > 0):   alpha = 1/2*torch.atan(B/(A-C))\n",
    "    elif q*A-q*C > 0 and q*B < 0:                                 alpha = 1/2*torch.atan(B/(A-C)) + pi\n",
    "    elif q*A-q*C < 0:                                             alpha = 1/2*torch.atan(B/(A-C)) + 1/2*pi\n",
    "    else: raise RuntimeError('\"Impossible\" condition reached; please debug')\n",
    "\n",
    "    return stackify((h, k, a, b, alpha))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 127,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.764046Z",
     "start_time": "2020-08-06T18:00:24.753075Z"
    }
   },
   "outputs": [],
   "source": [
    "Aq = torch.FloatTensor([[  9.56324965,  -1.90407389,  -5.75510187],\n",
    "                        [ -1.90407389,  15.43675035, -28.96942682],\n",
    "                        [ -5.75510187, -28.96942682, -80.3060445 ]])\n",
    "assert_allclose_f_ttn(conic2ellipse, Aq, \n",
    "                      torch.FloatTensor([1.0000, 2.0000, 4.0000, 3.0000, 0.2876]), atol=1e-4)\n",
    "assert_allclose(ellipse2conic(conic2ellipse(Aq)), Aq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# General image processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 128,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.767491Z",
     "start_time": "2020-08-06T18:00:24.765103Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def rgb2gray(arr): # From Pillow documentation\n",
    "    return arr[:,:,0]*(299/1000) + arr[:,:,1]*(587/1000) + arr[:,:,2]*(114/1000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 129,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.776144Z",
     "start_time": "2020-08-06T18:00:24.768403Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[[.1,.2,.3],[.2,.2,.2]]])\n",
    "assert_allclose_f_ttn(rgb2gray, arr, torch.FloatTensor([[0.1815, 0.2000]]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.782251Z",
     "start_time": "2020-08-06T18:00:24.777290Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def imresize(arr, sz, mode='bilinear', align_corners=True):\n",
    "    if not isinstance(sz, tuple): sz = tuple((shape(arr)//(shape(arr)/sz).min()).long())\n",
    "    return torch.nn.functional.interpolate(arr[None, None, :, :], \n",
    "                                           size=sz, \n",
    "                                           mode=mode, \n",
    "                                           align_corners=align_corners).squeeze(0).squeeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 131,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.789557Z",
     "start_time": "2020-08-06T18:00:24.783172Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[1,1,1,1],\n",
    "                         [2,2,2,2]])\n",
    "assert_allclose_f_ttn(imresize, (arr, 3), \n",
    "                      torch.FloatTensor([[1.0000, 1.0000, 1.0000, 1.0000, 1.0000, 1.0000],\n",
    "                                         [1.5000, 1.5000, 1.5000, 1.5000, 1.5000, 1.5000],\n",
    "                                         [2.0000, 2.0000, 2.0000, 2.0000, 2.0000, 2.0000]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`conv2d` is actually cross correlation, but you can transpose and do the same operation. This is mainly just a helper function to do 2d convolutions easily."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.794564Z",
     "start_time": "2020-08-06T18:00:24.790677Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def conv2d(arr, kernel, **kwargs):\n",
    "    return torch.nn.functional.conv2d(arr[None,None], kernel[None, None], **kwargs).squeeze(0).squeeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 133,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.801804Z",
     "start_time": "2020-08-06T18:00:24.795687Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[1,2,3],\n",
    "                         [3,2,1],\n",
    "                         [1,1,1]])\n",
    "kernel = torch.FloatTensor([[-0.25, 0.25],\n",
    "                            [-0.25, 0.25]])\n",
    "assert_allclose_f_ttn(conv2d, (arr, kernel), torch.FloatTensor([[ 0.0000,  0.0000],\n",
    "                                                                [-0.2500, -0.2500]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`pad` is just a helper function to do 2d convolutions easily"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 134,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.806388Z",
     "start_time": "2020-08-06T18:00:24.802970Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def pad(arr, pad, mode):\n",
    "    return torch.nn.functional.pad(arr[None,None], pad, mode=mode).squeeze(0).squeeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.813347Z",
     "start_time": "2020-08-06T18:00:24.807642Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[1,2,3],\n",
    "                         [3,2,1],\n",
    "                         [1,1,1]])\n",
    "assert_allclose_f_ttn(pad, (arr, (1,1,1,1), 'replicate'), torch.FloatTensor([[1,1,2,3,3],\n",
    "                                                                             [1,1,2,3,3],\n",
    "                                                                             [3,3,2,1,1],\n",
    "                                                                             [1,1,1,1,1],\n",
    "                                                                             [1,1,1,1,1]]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-07-12T19:14:18.162748Z",
     "start_time": "2020-07-12T19:14:18.158796Z"
    }
   },
   "source": [
    "`grad_array` by default uses replication to help with gradients along the edges. The default padding in `nn.functional.conv2d` is zero padding, which results in bad gradients along the edges."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 136,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.820990Z",
     "start_time": "2020-08-06T18:00:24.814570Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def grad_array(arr):\n",
    "    kernel_sobel = arr.new_tensor([[-0.1250, 0, 0.1250],\n",
    "                                   [-0.2500, 0, 0.2500],\n",
    "                                   [-0.1250, 0, 0.1250]])\n",
    "    arr = pad(arr, pad=(1,1,1,1), mode='replicate')\n",
    "    return tuple(conv2d(arr, kernel) for kernel in (kernel_sobel, kernel_sobel.T))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 137,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.829145Z",
     "start_time": "2020-08-06T18:00:24.822618Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[1,2,3],\n",
    "                         [3,2,1],\n",
    "                         [1,1,1]])\n",
    "assert_allclose_f_ttn(grad_array, arr, (torch.FloatTensor([[ 0.2500,  0.5000,  0.2500],\n",
    "                                                           [-0.1250, -0.2500, -0.1250],\n",
    "                                                           [-0.1250, -0.2500, -0.1250]]), \n",
    "                                        torch.FloatTensor([[ 0.7500,  0.0000, -0.7500],\n",
    "                                                           [-0.1250, -0.5000, -0.8750],\n",
    "                                                           [-0.8750, -0.5000, -0.1250]])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 138,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.833920Z",
     "start_time": "2020-08-06T18:00:24.830375Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def interp_array(arr, ps, align_corners=True, **kwargs):    \n",
    "    ps = stackify((rescale(ps[:, 0], [0, arr.shape[1]-1], [-1, 1]),\n",
    "                   rescale(ps[:, 1], [0, arr.shape[0]-1], [-1, 1])), dim=1) # ps must be rescaled to [-1,1]\n",
    "    return torch.nn.functional.grid_sample(arr[None, None], \n",
    "                                           ps[None, None], \n",
    "                                           align_corners=align_corners,\n",
    "                                           **kwargs).squeeze()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 139,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.843507Z",
     "start_time": "2020-08-06T18:00:24.835012Z"
    }
   },
   "outputs": [],
   "source": [
    "arr = torch.FloatTensor([[1,2,3],\n",
    "                         [4,5,6],\n",
    "                         [7,8,9]])\n",
    "ps = array_ps(arr)*0.8\n",
    "assert_allclose_f_ttn(interp_array, (arr, ps), torch.FloatTensor([1.,1.8,2.6,3.4,4.2,5.,5.8,6.6,7.4]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optimization stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`wlstsq` is weighted least squares"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 140,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.847926Z",
     "start_time": "2020-08-06T18:00:24.844565Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "@numpyify\n",
    "def wlstsq(A, b, W=None):\n",
    "    single = len(b.shape) == 1\n",
    "    if single: b = b[:, None]\n",
    "    if W is not None: # Weight matrix is a diagonal matrix with sqrt of the input weights\n",
    "        W = torch.sqrt(W.reshape(-1,1))\n",
    "        A, b = A*W, b*W\n",
    "    x = torch.lstsq(b, A).solution[:A.shape[1],:] # first n rows contains solution\n",
    "    if single: x = x.squeeze(1)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 141,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.855038Z",
     "start_time": "2020-08-06T18:00:24.848793Z"
    }
   },
   "outputs": [],
   "source": [
    "A = torch.FloatTensor([[1, 2, 3],\n",
    "                       [2, 3, 4],\n",
    "                       [4, 2, 5],\n",
    "                       [3, 3, 2],\n",
    "                       [1, 6, 7]])\n",
    "b = torch.FloatTensor([[1, 2],\n",
    "                       [1, 2],\n",
    "                       [1, 2],\n",
    "                       [2, 3],\n",
    "                       [7, 3]])\n",
    "W = torch.FloatTensor([1, 2, 3, 4, 5])\n",
    "assert_allclose_f_ttn(wlstsq, (A, b, W), torch.FloatTensor([[-0.5300,  0.4480],\n",
    "                                                            [ 1.0744,  0.6981],\n",
    "                                                            [ 0.1175, -0.2283]]), atol=1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Plotting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.860421Z",
     "start_time": "2020-08-06T18:00:24.856304Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def get_colors(n): return sns.color_palette(None, n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook stuff"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are kind of hacky, but I like being able to rerun a notebook and have it auto save/build/convert at the end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.866422Z",
     "start_time": "2020-08-06T18:00:24.861413Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def get_notebook_file():\n",
    "    id_kernel = re.search('kernel-(.*).json', ipykernel.connect.get_connection_file()).group(1)\n",
    "    for server in list_running_servers():\n",
    "        response = requests.get(requests.compat.urljoin(server['url'], 'api/sessions'),\n",
    "                                params={'token': server.get('token', '')})\n",
    "        for r in json.loads(response.text):\n",
    "            if 'kernel' in r and r['kernel']['id'] == id_kernel:\n",
    "                return Path(r['notebook']['path'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 144,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.881007Z",
     "start_time": "2020-08-06T18:00:24.867319Z"
    }
   },
   "outputs": [],
   "source": [
    "assert_allclose(get_notebook_file().as_posix(), 'utils.ipynb')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 145,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.884629Z",
     "start_time": "2020-08-06T18:00:24.882102Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def save_notebook():\n",
    "    file_notebook = get_notebook_file()\n",
    "    _get_md5 = lambda : hashlib.md5(file_notebook.read_bytes()).hexdigest() \n",
    "    md5_start = _get_md5()\n",
    "    display(Javascript('IPython.notebook.save_checkpoint();')) # Asynchronous\n",
    "    while md5_start == _get_md5(): time.sleep(1e-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 146,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.890544Z",
     "start_time": "2020-08-06T18:00:24.885909Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def build_notebook(save=True):\n",
    "    if save: save_notebook()\n",
    "    nbdev.export.notebook2script(fname=get_notebook_file().as_posix())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 147,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:24.896565Z",
     "start_time": "2020-08-06T18:00:24.891651Z"
    }
   },
   "outputs": [],
   "source": [
    "# export\n",
    "def convert_notebook(save=True, t='markdown'):\n",
    "    if save: save_notebook()\n",
    "    os.system(f'jupyter nbconvert --to {t} {get_notebook_file().as_posix()}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-05-03T20:36:21.604772Z",
     "start_time": "2020-05-03T20:36:21.600609Z"
    }
   },
   "source": [
    "# Build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 148,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-08-06T18:00:25.170751Z",
     "start_time": "2020-08-06T18:00:24.897478Z"
    },
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "application/javascript": [
       "IPython.notebook.save_checkpoint();"
      ],
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converted utils.ipynb.\n"
     ]
    }
   ],
   "source": [
    "build_notebook()"
   ]
  }
 ],
 "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.6.9"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}