get_figures.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d99aa2b1-0efc-4b2e-874a-cbcefc077a93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Jupyter environment detected. Enabling Open3D WebVisualizer.\n",
      "[Open3D INFO] WebRTC GUI backend enabled.\n",
      "[Open3D INFO] WebRTCWindowSystem: HTTP handshake server disabled.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import e3nn\n",
    "from e3nn import o3, io\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import open3d as o3d\n",
    "import plotly\n",
    "import plotly.graph_objects as go\n",
    "from plotly.subplots import make_subplots\n",
    "import plotly.io\n",
    "\n",
    "# plotly.io.renderers.default = \"notebook\"\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "from functools import partial, reduce\n",
    "\n",
    "\n",
    "from utils import load_model, save_model, CustomLRScheduler\n",
    "from data_generation import SimpleShapeDataset, SimpleShapeUniformRayDataset, BoxesDataset, SimpleShapeGridDataset\n",
    "from model import EncoderDecoder\n",
    "from convolution import S2ConvNet_Autoencoder\n",
    "from visualize import visualize_points, visualize_signal\n",
    "from losses import GridLoss, WeightedGridLoss, WeightedPointLoss, WeightedGridLossWithRotation\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "# device = 'cpu' # for now batch=1 so no need for gpu\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "af5ac4d2-baa4-449f-b284-0b8bc269b5aa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "LMAX = 4\n",
    "dataset = BoxesDataset(lmax=LMAX, n_samples=2)\n",
    "sphten = io.SphericalTensor(lmax=LMAX, p_arg=1, p_val=1)\n",
    "len(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "68ed111a-b3ec-4361-8220-30c4db511f44",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n",
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "S2ConvNet_Autoencoder(\n",
      "  (encoder): Sequential(\n",
      "    (0): Linear(1x0e+1x1e+1x2e+1x3e+1x4e -> 1x0e+3x1e+5x2e+7x3e+9x4e | 25 weights)\n",
      "    (1): SO3Activation (4 -> 4)\n",
      "    (2): BatchNorm (1x0e+3x1e+5x2e+7x3e+9x4e, eps=1e-05, momentum=0.1)\n",
      "    (3): Linear(1x0e+3x1e+5x2e+7x3e+9x4e -> 1x0e+3x1e+5x2e+7x3e+9x4e | 165 weights)\n",
      "    (4): BatchNorm (1x0e+3x1e+5x2e+7x3e+9x4e, eps=1e-05, momentum=0.1)\n",
      "    (5): SO3Activation (4 -> 3)\n",
      "    (6): Linear(1x0e+3x1e+5x2e+7x3e -> 1x0e+3x1e+5x2e+7x3e | 84 weights)\n",
      "    (7): BatchNorm (1x0e+3x1e+5x2e+7x3e, eps=1e-05, momentum=0.1)\n",
      "    (8): SO3Activation (3 -> 3)\n",
      "    (9): Linear(1x0e+3x1e+5x2e+7x3e -> 1x0e+3x1e+5x2e+7x3e | 84 weights)\n",
      "    (10): BatchNorm (1x0e+3x1e+5x2e+7x3e, eps=1e-05, momentum=0.1)\n",
      "    (11): SO3Activation (3 -> 3)\n",
      "    (12): Linear(1x0e+3x1e+5x2e+7x3e -> 1x0e+3x1e+5x2e+7x3e | 84 weights)\n",
      "    (13): BatchNorm (1x0e+3x1e+5x2e+7x3e, eps=1e-05, momentum=0.1)\n",
      "    (14): SO3Activation (3 -> 2)\n",
      "    (15): Linear(1x0e+3x1e+5x2e -> 1x0e+3x1e+5x2e | 35 weights)\n",
      "    (16): BatchNorm (1x0e+3x1e+5x2e, eps=1e-05, momentum=0.1)\n",
      "    (17): SO3Activation (2 -> 2)\n",
      "    (18): Linear(1x0e+3x1e+5x2e -> 1x0e+3x1e+5x2e | 35 weights)\n",
      "    (19): BatchNorm (1x0e+3x1e+5x2e, eps=1e-05, momentum=0.1)\n",
      "    (20): SO3Activation (2 -> 1)\n",
      "  )\n",
      "  (decoder): Sequential(\n",
      "    (0): Linear(1x0e+3x1e -> 1x0e+3x1e | 10 weights)\n",
      "    (1): BatchNorm (1x0e+3x1e, eps=1e-05, momentum=0.1)\n",
      "    (2): SO3Activation (1 -> 2)\n",
      "    (3): Linear(1x0e+3x1e+5x2e -> 1x0e+3x1e+5x2e | 35 weights)\n",
      "    (4): BatchNorm (1x0e+3x1e+5x2e, eps=1e-05, momentum=0.1)\n",
      "    (5): SO3Activation (2 -> 2)\n",
      "    (6): Linear(1x0e+3x1e+5x2e -> 1x0e+3x1e+5x2e | 35 weights)\n",
      "    (7): BatchNorm (1x0e+3x1e+5x2e, eps=1e-05, momentum=0.1)\n",
      "    (8): SO3Activation (2 -> 3)\n",
      "    (9): Linear(1x0e+3x1e+5x2e+7x3e -> 1x0e+3x1e+5x2e+7x3e | 84 weights)\n",
      "    (10): BatchNorm (1x0e+3x1e+5x2e+7x3e, eps=1e-05, momentum=0.1)\n",
      "    (11): SO3Activation (3 -> 3)\n",
      "    (12): Linear(1x0e+3x1e+5x2e+7x3e -> 1x0e+3x1e+5x2e+7x3e | 84 weights)\n",
      "    (13): BatchNorm (1x0e+3x1e+5x2e+7x3e, eps=1e-05, momentum=0.1)\n",
      "    (14): SO3Activation (3 -> 3)\n",
      "    (15): Linear(1x0e+3x1e+5x2e+7x3e -> 1x0e+3x1e+5x2e+7x3e | 84 weights)\n",
      "    (16): BatchNorm (1x0e+3x1e+5x2e+7x3e, eps=1e-05, momentum=0.1)\n",
      "    (17): SO3Activation (3 -> 4)\n",
      "    (18): Linear(1x0e+3x1e+5x2e+7x3e+9x4e -> 1x0e+1x1e+1x2e+1x3e+1x4e | 25 weights)\n",
      "  )\n",
      ")\n",
      "number of parameters =  86063\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/shayan/miniconda3/envs/equiv-net/lib/python3.10/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "model = S2ConvNet_Autoencoder(\n",
    "    LMAX,\n",
    "    l_list=[LMAX, 3, 3, 3, 2, 2, 1],\n",
    "    channels=[4,   8, 8, 8, 16, 16, 32]\n",
    ").to(device)\n",
    "print(model)\n",
    "print(\"number of parameters = \", sum([np.prod(x.shape) for x in model.parameters()]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9320e868-f481-4af3-a18d-5a8380efa8ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "# from utils import get_checkpoints_dir\n",
    "\n",
    "# checkpoint_path = f'{get_checkpoints_dir()}/old/checkpoint_final_state.pt'\n",
    "# checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))\n",
    "# model.load_state_dict(checkpoint)\n",
    "\n",
    "# # load_model(model, 'l=1_final')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a5497d97-e01a-46a0-ba6a-e6c143241504",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dot_almost(a, b, c, eps=1e-5, put_assert=True):\n",
    "    dt = (a * b).sum(dim=-1)\n",
    "    if (dt.max() > (c + eps)) or (dt.min() < (c - eps)):\n",
    "        print(\"dot product is \", dt, \"expected\", c)\n",
    "        assert(not put_assert)\n",
    "        return False\n",
    "    return True\n",
    "\n",
    "\n",
    "def orthogonal_vector(a, b):\n",
    "    batch, dim = a.shape\n",
    "    A = torch.stack([a, b], dim=1)\n",
    "    A = torch.cat([A, torch.eye(dim).repeat(batch, 1, 1)], dim=1).permute([0, 2, 1]) # transpose\n",
    "    q, r = torch.qr(A)\n",
    "    axis = q[:, :, -1]\n",
    "\n",
    "    dot_almost(a, axis, 0)\n",
    "    dot_almost(b, axis, 0)\n",
    "    dot_almost(axis, axis, 1)\n",
    "    return axis\n",
    "\n",
    "def rotate(R, vecs):\n",
    "    return torch.einsum('nij,nj->ni',R, vecs)\n",
    "    \n",
    "def interpolate_in_1D(v1, v2, s):\n",
    "    # todo we do linear interpolation because the signs might be different... why do we not have to care about this in other irreps?\n",
    "    return v1 + (v2 - v1) * s\n",
    "\n",
    "def interpolate_in_3D(v1, v2, s):\n",
    "    norm1 = torch.norm(v1)\n",
    "    norm2 = torch.norm(v2)\n",
    "    \n",
    "    dot = (v1 * v2).sum(dim=-1)\n",
    "    cos_angle = dot / (norm1 * norm2)\n",
    "    angle = torch.acos(cos_angle)\n",
    "    axis = torch.cross(v1, v2, dim=-1)\n",
    "    axis = axis / torch.norm(axis, dim=-1, keepdim=True)\n",
    "\n",
    "    R = o3.axis_angle_to_matrix(axis=axis, angle=angle * s)\n",
    "    \n",
    "    R_full = o3.axis_angle_to_matrix(axis=axis, angle=angle)\n",
    "    assert(dot_almost(v2, rotate(R_full, v1), norm1 * norm2))\n",
    "    \n",
    "    return rotate(R, v1) * torch.pow(norm2 / norm1, s)\n",
    "\n",
    "\n",
    "def interpolate_in_highD(irrep, v1, v2, s):\n",
    "    batch, dim = v1.shape\n",
    "    A = torch.stack([v1, v2], dim=1)\n",
    "    A = torch.cat([A, torch.eye(dim).repeat(batch, 1, 1)], dim=1).permute([0, 2, 1]) # transpose\n",
    "    q, r = torch.qr(A)\n",
    "    axis = q[:, :, :3] # this is our basis\n",
    "    v1_3d = torch.einsum('ndi,nd->ni', axis, v1)\n",
    "    v2_3d = torch.einsum('ndi,nd->ni', axis, v2)\n",
    "    v_interp_3d = interpolate_in_3D(v1_3d, v2_3d, s)\n",
    "    v_interp = torch.einsum('ndi,ni->nd', axis, v_interp_3d)\n",
    "    return v_interp\n",
    "    \n",
    "    \n",
    "    # norm1 = torch.linalg.norm(vec1, dim=-1, keepdim=True)\n",
    "    # norm2 = torch.linalg.norm(vec2, dim=-1, keepdim=True)\n",
    "    # vec1_norm = vec1 / norm1\n",
    "    # vec2_norm = vec2 / norm2\n",
    "    # axis = orthogonal_vector(vec1_norm, vec2_norm)\n",
    "    # axis_norm = torch.norm(axis, dim=-1)\n",
    "    # axis /= axis_norm.unsqueeze(-1)\n",
    "    # # theta = torch.asin(axis_norm)\n",
    "\n",
    "    # x_axis = vec1_norm\n",
    "    # y_axis = vec2_norm - (vec2_norm * x_axis).sum(dim=-1).unsqueeze(-1) * x_axis\n",
    "    # y_axis /= torch.linalg.norm(y_axis, dim=-1, keepdim=True)\n",
    "\n",
    "    # # handle degenerate cases\n",
    "    # cos = (vec1_norm * vec2_norm).sum(dim=-1)\n",
    "    # y_axis = torch.where(\n",
    "    #     (torch.abs(cos) > 1-(1e-4)).unsqueeze(-1),\n",
    "    #     orthogonal_vector(x_axis, axis),\n",
    "    #     y_axis\n",
    "    # )\n",
    "    \n",
    "    # dot_almost(x_axis, x_axis, 1)\n",
    "    # dot_almost(axis, axis, 1)\n",
    "    # dot_almost(y_axis, y_axis, 1)\n",
    "\n",
    "    # dot_almost(vec1_norm, axis, 0)\n",
    "    # dot_almost(vec2_norm, axis, 0)\n",
    "\n",
    "    # dot_almost(x_axis, y_axis, 0)\n",
    "    # dot_almost(x_axis, axis, 0)\n",
    "    # dot_almost(y_axis, axis, 0)\n",
    "\n",
    "    # x_coord = (x_axis * vec2_norm).sum(dim=-1)\n",
    "    # y_coord = (y_axis * vec2_norm).sum(dim=-1)\n",
    "\n",
    "    # theta = torch.atan2(y_coord, x_coord)\n",
    "\n",
    "    # R = o3.Irreps(f'1x{l}e').D_from_axis_angle(axis=axis, angle=theta * s)\n",
    "\n",
    "    # should_be_vec2_norm1 = torch.einsum('nij,nj->ni', o3.Irreps(f'1x{l}e').D_from_axis_angle(axis=axis, angle=+theta), vec1_norm)\n",
    "    # should_be_vec2_norm2 = torch.einsum('nij,nj->ni', o3.Irreps(f'1x{l}e').D_from_axis_angle(axis=axis, angle=-theta), vec1_norm)\n",
    "    # print('---')\n",
    "    # print(vec1)\n",
    "    # print(vec2)\n",
    "    # assert(dot_almost(should_be_vec2_norm1, vec2_norm, 1, eps=1e-2, put_assert=False) or dot_almost(should_be_vec2_norm2, vec2_norm, 1, eps=1e-2, put_assert=False))\n",
    "\n",
    "    # vec_interp_size = vec1 * torch.pow(norm2 / norm1, s)\n",
    "    # vec_interp_rotated = torch.einsum('nij,nj->ni', R, vec_interp_size)\n",
    "    # vec_interp = vec_interp_rotated\n",
    "    # res[..., ind:ind+sz] = vec_interp\n",
    "\n",
    "\n",
    "def interpolate_in_latent_space(repr, latent1, latent2, s):\n",
    "    res = torch.empty_like(latent1)\n",
    "    assert(latent1.shape == latent2.shape)\n",
    "\n",
    "    ind = 0\n",
    "    for l in repr.ls:\n",
    "        sz = 2 * l + 1\n",
    "        vec1 = latent1[..., ind:ind+sz]\n",
    "        vec2 = latent2[..., ind:ind+sz]\n",
    "\n",
    "        if sz == 1:\n",
    "            res[..., ind:ind+sz] = interpolate_in_1D(vec1, vec2, s)\n",
    "        elif sz == 3:\n",
    "            res[..., ind:ind+sz] = interpolate_in_3D(vec1, vec2, s)\n",
    "        else:\n",
    "            # raise Exception(\"l > 1 is not supported in interpolation sorry :))\")\n",
    "            res[..., ind:ind+sz] = interpolate_in_highD(o3.Irreps(f'{l}e'), vec1, vec2, s)\n",
    "        ind += sz\n",
    "    assert(ind == res.shape[-1])\n",
    "    return res\n",
    "\n",
    "\n",
    "def linear_interpolate_in_latent_space(latent1, latent2, s):\n",
    "    return latent1 + s * (latent2 - latent1)\n",
    "\n",
    "\n",
    "\n",
    "# this is just for debugging\n",
    "def interpolate_in_latent_space_hint(axis, angle, repr, latent1, latent2, s):\n",
    "    res = torch.empty_like(latent1)\n",
    "    assert(latent1.shape == latent2.shape)\n",
    "\n",
    "    ind = 0\n",
    "    for l in repr.ls:\n",
    "        sz = 2 * l + 1\n",
    "        vec1 = latent1[..., ind:ind+sz]\n",
    "        vec2 = latent2[..., ind:ind+sz]\n",
    "\n",
    "        irrep = o3.Irreps(f'{l}e')\n",
    "        D = irrep.D_from_axis_angle(axis=axis, angle=angle * s)\n",
    "        res[..., ind:ind+sz] = torch.einsum('ij,nj->ni', D, vec1)\n",
    "        ind += sz\n",
    "    assert(ind == res.shape[-1])\n",
    "    return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "a11105f9-4798-4db3-9b1c-2fdea572c21a",
   "metadata": {},
   "outputs": [],
   "source": [
    "latent_repr = model.latent_repr\n",
    "model_sphten_repr = model.model_sphten_repr\n",
    "\n",
    "sphten = model.model_sphten_repr\n",
    "\n",
    "lmax = 4\n",
    "# sh1 = dataset[0].squeeze()\n",
    "# sh2 = dataset[1].squeeze()\n",
    "\n",
    "axis = torch.tensor([0, 0, 1]).float()\n",
    "angle = torch.tensor(torch.pi/2)\n",
    "R = o3.axis_angle_to_matrix(axis, angle)\n",
    "sphten = model.model_sphten_repr\n",
    "D = sphten.D_from_matrix(R)\n",
    "\n",
    "sh1 = dataset[0].squeeze()\n",
    "sh2 = dataset[1].squeeze()\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    sh1 = sh1.unsqueeze(0).to(device)\n",
    "    sh2 = sh2.unsqueeze(0).to(device)\n",
    "    inp1, latent1, out1 = model(sh1)\n",
    "    inp2, latent2, out2 = model(sh2)\n",
    "\n",
    "init_shape = latent1.shape\n",
    "latent1 = latent1.reshape(-1).unsqueeze(0)\n",
    "latent2 = latent2.reshape(-1).unsqueeze(0)\n",
    "\n",
    "def save_sh(met, idx, sh):\n",
    "    layout = go.Layout(scene=dict(aspectmode=\"cube\",\n",
    "                               xaxis=dict(title='X', showgrid=False, visible=False),\n",
    "                               yaxis=dict(title='Y', showgrid=False, visible=False),\n",
    "                               zaxis=dict(title='Z', showgrid=False, visible=False),\n",
    "                               # Set background color to transparent\n",
    "                               bgcolor='rgba(0, 0, 0, 0)'))\n",
    "    p_val = 1\n",
    "    p_arg = 1\n",
    "    sphten = e3nn.io.SphericalTensor(lmax, p_val, p_arg)\n",
    "    fig = go.Figure([go.Surface(sphten.plotly_surface(sh, radius=True)[0])], layout=layout)\n",
    "    fig.update(layout_coloraxis_showscale=False)\n",
    "    fig.update(layout_showlegend=False)\n",
    "    fig.update_coloraxes(showscale=False)\n",
    "    fig.write_image(f'{met}-{idx}.png')\n",
    "\n",
    "    \n",
    "N = 8\n",
    "rows = 2\n",
    "columns = N\n",
    "\n",
    "i = 0\n",
    "for j in range(columns):\n",
    "    s = j/(N-1)\n",
    "    latent = interpolate_in_latent_space(latent_repr, latent1=latent1.cpu(), latent2=latent2.cpu(), s=j/(N-1)).to(device)\n",
    "    latent = latent.to(device)\n",
    "    # latent = interpolate_in_latent_space_hint(axis, angle, latent_repr, latent1=latent1.cpu(), latent2=latent2.cpu(), s=j/(N-1)).to(device)\n",
    "    latent = latent.reshape(init_shape)\n",
    "    with torch.no_grad():\n",
    "        out = model.decoder(latent)\n",
    "    save_sh('equiv_rot', j, out.cpu())\n",
    "    \n",
    "        \n",
    "\n",
    "i = 1\n",
    "for j in range(columns):\n",
    "    s = j/(N-1)\n",
    "    latent = linear_interpolate_in_latent_space(latent1=latent1.cpu(), latent2=latent2.cpu(), s=s).to(device)\n",
    "    latent = latent.to(device)\n",
    "    latent = latent.reshape(init_shape)\n",
    "    with torch.no_grad():\n",
    "        out = model.decoder(latent)\n",
    "    save_sh('lin_rot', j, out.cpu())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "2cf9201c-4110-4e06-b815-6f727434d6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "\n",
    "def get_name(name):\n",
    "    # List of image file paths\n",
    "    image_files = [f\"{name}-{i}.png\" for i in range(N)]\n",
    "    \n",
    "    # Load images and extract centers\n",
    "    images = []\n",
    "    for file in image_files:\n",
    "        img = Image.open(file)\n",
    "        width, height = img.size\n",
    "        # Calculate the center coordinates\n",
    "        D = min(width, height) * 0.5\n",
    "        left = (width - D) / 2\n",
    "        top = (height - D) / 2\n",
    "        right = (width + D) / 2\n",
    "        bottom = (height + D) / 2\n",
    "        # Crop the image to extract the center\n",
    "        cropped_img = img.crop((left, top, right, bottom))\n",
    "        images.append(cropped_img)\n",
    "    \n",
    "    # Convert images to NumPy arrays\n",
    "    image_arrays = [np.array(img) for img in images]\n",
    "    \n",
    "    # Stack the images along the vertical axis\n",
    "    stacked_image = np.hstack(image_arrays)\n",
    "    # # Convert the stacked image array back to a PIL image\n",
    "    stacked_pil_image = Image.fromarray(stacked_image)\n",
    "    stacked_pil_image.save(f'{name}-stacked.png')\n",
    "    \n",
    "    return stacked_image        \n",
    "\n",
    "\n",
    "stacked = np.vstack([get_name('equiv_rot'), get_name('lin_rot')])\n",
    "stacked_pil_image = Image.fromarray(stacked)\n",
    "stacked_pil_image.save(f'all-stacked.png')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}