{
 "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 os\n",
    "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 torch_harmonics.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 = 257\n",
    "nlon = 2*(nlat-1)\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",
      "--2025-03-20 22:41:46--  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... "
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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   174KB/s    in 0.3s    \n",
      "\n",
      "2025-03-20 22:41:47 (174 KB/s) - ‘./data/mola_topo.jpg’ saved [55192/55192]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "os.makedirs(\"./data\", exist_ok=True)\n",
    "!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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88412 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 torch_harmonics.plotting import plot_sphere\n",
    "\n",
    "plot_sphere(data, projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"Mars topography MOLA\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's bring the data to our desired input grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid_up = \"equiangular\"\n",
    "grid_down = \"legendre-gauss\"\n",
    "\n",
    "input_resample = ResampleS2(nlat, nlon, nlat, nlon, grid_in=\"equiangular\", grid_out=grid_up)\n",
    "\n",
    "data_res = input_resample(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<cartopy.mpl.geocollection.GeoQuadMesh at 0x762e773c0520>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 88204 characters]",
      "text/plain": [
       "<Figure size 432x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_sphere(data_res, projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"interpolated input data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resampling\n",
    "\n",
    "let us define down- and up-sampling layers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "scale_factor = 4\n",
    "\n",
    "nlat_down = math.ceil(nlat / scale_factor)\n",
    "nlon_down = math.ceil(nlon / scale_factor)\n",
    "\n",
    "downsample = ResampleS2(nlat, nlon, nlat_down, nlon_down, grid_in=grid_up)\n",
    "upsample = ResampleS2(nlat_down, nlon_down, nlat, nlon, grid_in=grid_down)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_down = downsample(data)\n",
    "data_up = upsample(data_down)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 471848 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], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"Original\")\n",
    "plot_sphere(data_down, fig=subfigs[1], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"down-sampled\")\n",
    "plot_sphere(data_up, fig=subfigs[2], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"up-sampled\")\n",
    "plot_sphere((data_up - data), fig=subfigs[3], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"residual (after upsampling)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5654)\n"
     ]
    }
   ],
   "source": [
    "print((data_up - data).abs().max())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Upsampling error\n",
    "\n",
    "In order to estimate the accuracy of up and down-sampling, let us compare it to idealized SHT upsampling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch_harmonics import RealSHT, InverseRealSHT\n",
    "\n",
    "sht = RealSHT(nlat_down, nlon_down, grid=grid_down)\n",
    "isht = InverseRealSHT(nlat, nlon, lmax=sht.lmax, mmax=sht.mmax, grid=grid_up)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_up = upsample(data_down)\n",
    "data_sht_up = isht(sht(data_down))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 731872 characters]",
      "text/plain": [
       "<Figure size 2304x864 with 12 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure(layout='constrained', figsize=(32, 12))\n",
    "subfigs = fig.subfigures(2, 4)\n",
    "\n",
    "# spectral fitting\n",
    "plot_sphere(data, fig=subfigs[0, 0], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"Original\")\n",
    "plot_sphere(data_down, fig=subfigs[0, 1], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"down-sampled\")\n",
    "plot_sphere(data_up, fig=subfigs[0, 2], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"up-sampled\")\n",
    "plot_sphere(data_sht_up, fig=subfigs[0, 3], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"SHT up-sampled\")\n",
    "plot_sphere((data_up - data), fig=subfigs[1, 2], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"residual (after upsampling)\")\n",
    "plot_sphere((data_sht_up - data), fig=subfigs[1, 3], projection=\"orthographic\", cmap=\"turbo\", colorbar=True, title=\"residual (after SHT upsampling)\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that SHT upsampling suffers from ripples. This can be mitigated by choosing a grid which uses exact quadrature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0.5654)\n",
      "tensor(0.6636)\n"
     ]
    }
   ],
   "source": [
    "print((data_up - data).abs().max())\n",
    "print((data_sht_up - data).abs().max())"
   ]
  }
 ],
 "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
}
