rediscovering_SO3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "#import emd\n",
    "import os\n",
    "\n",
    "from torch import nn\n",
    "import torch\n",
    "from IPython.display import clear_output\n",
    "\n",
    "from makegif import make_gif\n",
    "import copy\n",
    "import time\n",
    "\n",
    "torch.set_default_dtype(torch.float64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.style.use('dark_background')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = torch.tensor(np.random.randn(300,3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Some functions for visualization:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Here are some functions for visualizing the output\n",
    "\n",
    "def get_angle(v, w):\n",
    "    '''get the angle between two vectors'''\n",
    "    return v @ w / (torch.norm(v) * torch.norm(w))\n",
    "\n",
    "\n",
    "def get_axis(M):\n",
    "    '''Finds the eigenvector corresponding to the eigenvalue with the least \n",
    "    imaginary part, if the matrix is a rotation matrix or a generator of \n",
    "    rotations then this vector is the axis of rotation'''    \n",
    "\n",
    "    eig_vals, eig_vecs = torch.linalg.eig(M)\n",
    "    axis = eig_vecs.T[torch.argmin(torch.abs(eig_vals.imag))]\n",
    "    return torch.sign(torch.sum(axis).real) * axis\n",
    "\n",
    "\n",
    "def draw_vec(ax, v, color='C0', lw=4):\n",
    "    '''Draw a vector to ax, this adds lines for the projection'''\n",
    "    ax.plot([0,v[0]], [0,v[1]], [0,v[2]], color=color, lw=lw)\n",
    "\n",
    "    ax.plot([v[0],v[0]], [v[1],v[1]], [0,v[2]], color='w', alpha=.25, ls='--')\n",
    "    ax.plot([v[0],v[0]], [0   ,v[1]], [0,0   ], color='w', alpha=.25, ls='--')\n",
    "    ax.plot([0   ,v[0]], [v[1],v[1]], [0,0   ], color='w', alpha=.25, ls='--')\n",
    "\n",
    "\n",
    "def visualize_generators(params, eps=1e-3):\n",
    "    '''Imshow each of the generators'''\n",
    "    plt.figure(figsize=[12,4])\n",
    "    \n",
    "    for i, X in enumerate(params):\n",
    "        plt.subplot(1,len(params),i+1)\n",
    "        plt.imshow(X.detach().numpy(), cmap='RdBu')\n",
    "        plt.title(f'det = {np.linalg.det(np.eye(X.shape[0]) + eps * X.detach().numpy())}')\n",
    "        plt.colorbar()\n",
    "\n",
    "\n",
    "def draw_4tych(params):\n",
    "    \n",
    "    fig, [ax1, ax2, ax3, ax4] = plt.subplots(nrows=1, ncols=4, figsize=[18,4])\n",
    "    \n",
    "    # draw each generator\n",
    "    for i, [ax, X] in enumerate(zip([ax1, ax2, ax3], params)):\n",
    "        \n",
    "        ax.imshow(X.detach().numpy(), cmap='RdBu')\n",
    "        ax.set_title(f'det = {np.linalg.det(np.eye(3) + X.detach().numpy())}')\n",
    "        ax.axis('off')\n",
    "\n",
    "    ax4 = fig.add_subplot(144, projection='3d')\n",
    "    \n",
    "    # draw coordinate axis\n",
    "    ax_lim = 1\n",
    "    ax4.plot([-ax_lim,ax_lim],[0,0],[0,0], color='w', alpha=.3)\n",
    "    ax4.plot([0,0],[-ax_lim,ax_lim],[0,0], color='w', alpha=.3)\n",
    "    ax4.plot([0,0],[0,0],[-ax_lim,ax_lim], color='w', alpha=.3)\n",
    "    \n",
    "    # draw each rotation axis\n",
    "    for X, col in zip(params, ['C0', 'C1', 'C2']):\n",
    "        draw_vec(ax4, get_axis(X.detach().numpy()+np.eye(3)), col)\n",
    "\n",
    "    ax4.set_xlim(-ax_lim,ax_lim)\n",
    "    ax4.set_ylim(-ax_lim,ax_lim)\n",
    "    ax4.set_zlim(-ax_lim,ax_lim)\n",
    "    ax4.grid(False)\n",
    "    ax4.axis('off')\n",
    "    \n",
    "    ax4.set_title(f'Epoch {i}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        \n",
    "        self.X1 = nn.Linear(3,3, bias=False)\n",
    "        self.X2 = nn.Linear(3,3, bias=False)\n",
    "        self.X3 = nn.Linear(3,3, bias=False)\n",
    "        \n",
    "    def forward(self, inputs, epsilon):\n",
    "        \n",
    "        out1 = inputs[0] + epsilon * self.X1(inputs[0])\n",
    "        out2 = inputs[1] + epsilon * self.X2(inputs[1])\n",
    "        out3 = inputs[2] + epsilon * self.X3(inputs[2])\n",
    "        \n",
    "        return [out1, out2, out3]\n",
    "    \n",
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(nn.Module):\n",
    "    \n",
    "    def __init__(self, n_generators=3, n_dim=3):\n",
    "        super(MyModel, self).__init__()\n",
    "        \n",
    "        self.n_generators = n_generators\n",
    "        self.X_list = []\n",
    "        \n",
    "        for _ in range(n_generators):\n",
    "            self.X_list.append(nn.Linear(n_dim, n_dim, bias=False))\n",
    "        \n",
    "    def forward(self, inputs, epsilon):\n",
    "        \n",
    "        outputs = []\n",
    "        for i in range(self.n_generators):\n",
    "            outputs.append(inputs[0] + epsilon * self.X_list[i](inp))\n",
    "        \n",
    "        return outputs\n",
    "    \n",
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def bracket(A, B):\n",
    "    \n",
    "    return A @ B - B @ A\n",
    "\n",
    "\n",
    "def single_loss(inp, pred, X, w_infinitesimal=1, eps=1e-3):\n",
    "    \n",
    "    dl = torch.mean((torch.norm(inp, dim=1) - torch.norm(pred, dim=1))**2)\n",
    "    \n",
    "    return dl + w_infinitesimal * (torch.sum(X**2) - eps**2)**2\n",
    "\n",
    "\n",
    "def ensemble_loss(inp, preds, params, eps=1e-3):\n",
    "    \n",
    "    loss = 0\n",
    "    for i, X in enumerate(params):\n",
    "        \n",
    "        for j, Y in enumerate(params):\n",
    "            loss += get_angle(get_axis(X).real, \n",
    "                              get_axis(Y).real)**2\n",
    "            \n",
    "        loss += single_loss(inputs[0], preds[0], X1, eps)\n",
    "     \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torch_kernel",
   "language": "python",
   "name": "torch_kernel"
  },
  "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.10.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}