{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Resampling signals on the 2-sphere\n",
    "\n",
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from torch_harmonics import ResampleS2\n",
    "from torch_harmonics.quadrature import _precompute_latitudes\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "from plotting import plot_sphere"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's start by downloading our Mars topography again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "nlat = 256\n",
    "nlon = 2*nlat\n",
    "grid = \"equiangular\"\n",
    "\n",
    "xq, wq = _precompute_latitudes(nlat, grid=grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/bbonev/.zshenv:export:2: not valid in this context: :/usr/local/cuda-11.7/lib64\n",
      "--2024-08-29 15:13:01--  https://astropedia.astrogeology.usgs.gov/download/Mars/GlobalSurveyor/MOLA/thumbs/Mars_MGS_MOLA_DEM_mosaic_global_1024.jpg\n",
      "Resolving astropedia.astrogeology.usgs.gov (astropedia.astrogeology.usgs.gov)... 137.227.239.81, 2001:49c8:c000:122d::81\n",
      "Connecting to astropedia.astrogeology.usgs.gov (astropedia.astrogeology.usgs.gov)|137.227.239.81|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 \n",
      "Length: 55192 (54K) [image/jpeg]\n",
      "Saving to: ‘./data/mola_topo.jpg’\n",
      "\n",
      "./data/mola_topo.jp 100%[===================>]  53.90K   169KB/s    in 0.3s    \n",
      "\n",
      "2024-08-29 15:13:02 (169 KB/s) - ‘./data/mola_topo.jpg’ saved [55192/55192]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "!wget https://astropedia.astrogeology.usgs.gov/download/Mars/GlobalSurveyor/MOLA/thumbs/Mars_MGS_MOLA_DEM_mosaic_global_1024.jpg -O ./data/mola_topo.jpg\n",
    "\n",
    "import imageio.v3 as iio\n",
    "\n",
    "img = iio.imread('./data/mola_topo.jpg')\n",
    "#convert to grayscale\n",
    "data = np.dot(img[...,:3]/255, [0.299, 0.587, 0.114])\n",
    "# interpolate onto 512x1024 grid:\n",
    "data = nn.functional.interpolate(torch.from_numpy(data).unsqueeze(0).unsqueeze(0), size=(nlat,nlon)).squeeze().float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "let's plot it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 91512 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from plotting import plot_sphere\n",
    "\n",
    "plot_sphere(data, cmap=\"turbo\", colorbar=True, title=\"Mars topography MOLA\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resampling\n",
    "\n",
    "let us define down- and up-sampling layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling_factor = 0.25\n",
    "\n",
    "downsample = ResampleS2(nlat, nlon, int(scaling_factor * nlat), int(scaling_factor * nlon), grid_in=grid, grid_out=grid)\n",
    "upsample = ResampleS2(int(scaling_factor * nlat), int(scaling_factor * nlon), nlat, nlon, grid_in=grid, grid_out=grid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_down = downsample(data)\n",
    "data_up = upsample(data_down)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 570208 characters]",
      "text/plain": [
       "<Figure size 2304x432 with 8 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(layout='constrained', figsize=(32, 6))\n",
    "subfigs = fig.subfigures(1, 4)\n",
    "\n",
    "# spectral fitting\n",
    "plot_sphere(data, fig=subfigs[0], cmap=\"turbo\", colorbar=True, title=\"Original\")\n",
    "plot_sphere(data_down, fig=subfigs[1], cmap=\"turbo\", colorbar=True, title=\"down-sampled\")\n",
    "plot_sphere(data_up, fig=subfigs[2], cmap=\"turbo\", colorbar=True, title=\"up-sampled\")\n",
    "plot_sphere((data_up - data), fig=subfigs[3], cmap=\"turbo\", colorbar=True, title=\"residual (after upsampling)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
