test_ts2kit.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test TS2Kit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch \n",
    "import numpy as np\n",
    "from ts2kit import FTSHT, ITSHT, gridDH, clearTS2KitCache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to clear cache (assuming you've set the cache path in ts2kit.py)\n",
    "#clearTS2KitCache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inverse SHT, followed by forward SHT\n",
    "class tshtTest(torch.nn.Module):\n",
    "    def __init__(self, B):\n",
    "        super(tshtTest, self).__init__()\n",
    "\n",
    "        self.fsht = FTSHT(B)\n",
    "        self.isht = ITSHT(B)\n",
    "\n",
    "    def forward(self, x):\n",
    "        \n",
    "        return self.fsht(self.isht(x));\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Parameters - change me to try different things\n",
    "\n",
    "## Bandlimit\n",
    "B = 64\n",
    "\n",
    "## Number of batch dimensions\n",
    "b = 4096;\n",
    "\n",
    "## Torch device (GPU)\n",
    "device = torch.device('cuda')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Initalize the module at double and floating point precision\n",
    "\n",
    "# On first run at a given bandwidth, you should see several statements printed in the console confirming\n",
    "# that various tensors have been pre-computed sucesssfully. These will be saved to the cache directory and \n",
    "# automatically loaded during subsequent initalizations\n",
    "\n",
    "test_double = tshtTest(B).to(device)\n",
    "test_float = tshtTest(B).to(device).float()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Generate random SH coefficents\n",
    "Psi_double = torch.view_as_complex(2*(torch.rand(b, 2*B -1, B, 2).double() - 0.5)).to(device)\n",
    "\n",
    "for m in range(-(B-1), B):\n",
    "    for l in range(0, B):\n",
    "        if (l < np.abs(m)):\n",
    "            Psi_double[:, m + (B-1), l] = 0.0;\n",
    "\n",
    "Psi_float = Psi_double.clone().cfloat()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Run time and reconstruction error at double precision\n",
    "start = torch.cuda.Event(enable_timing=True)\n",
    "end = torch.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "PsiR_double = test_double(Psi_double);\n",
    "end.record()\n",
    "\n",
    "torch.cuda.synchronize()\n",
    "\n",
    "print('Run time: {:06.4f} ms'.format(start.elapsed_time(end)), flush=True);\n",
    "print('error = {}'.format(torch.sum(torch.abs(Psi_double - PsiR_double))/torch.sum(torch.abs(Psi_double))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Run time and reconstruction error at floating precision\n",
    "start = torch.cuda.Event(enable_timing=True)\n",
    "end = torch.cuda.Event(enable_timing=True)\n",
    "\n",
    "start.record()\n",
    "PsiR_float = test_float(Psi_float);\n",
    "end.record()\n",
    "\n",
    "torch.cuda.synchronize()\n",
    "\n",
    "print('Run time: {:06.4f} ms'.format(start.elapsed_time(end)), flush=True);\n",
    "print('error = {}'.format(torch.sum(torch.abs(Psi_float - PsiR_float))/torch.sum(torch.abs(Psi_float))))"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}