{
 "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 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": [
      "--2025-01-20 05:35:21--  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  --.-KB/s    in 0.07s   \n",
      "\n",
      "2025-01-20 05:35:21 (740 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, 146568 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "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": [
    "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 0x7fbcebbf0410>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 145744 characters]",
      "text/plain": [
       "<Figure size 640x480 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plot_sphere(data_res, 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, 663696 characters]",
      "text/plain": [
       "<Figure size 3200x600 with 8 Axes>"
      ]
     },
     "metadata": {},
     "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": 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, 1011192 characters]",
      "text/plain": [
       "<Figure size 3200x1200 with 12 Axes>"
      ]
     },
     "metadata": {},
     "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], cmap=\"turbo\", colorbar=True, title=\"Original\")\n",
    "plot_sphere(data_down, fig=subfigs[0, 1], cmap=\"turbo\", colorbar=True, title=\"down-sampled\")\n",
    "plot_sphere(data_up, fig=subfigs[0, 2], cmap=\"turbo\", colorbar=True, title=\"up-sampled\")\n",
    "plot_sphere(data_sht_up, fig=subfigs[0, 3], cmap=\"turbo\", colorbar=True, title=\"SHT up-sampled\")\n",
    "plot_sphere((data_up - data), fig=subfigs[1, 2], cmap=\"turbo\", colorbar=True, title=\"residual (after upsampling)\")\n",
    "plot_sphere((data_sht_up - data), fig=subfigs[1, 3], 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())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "from functools import cache\n",
    "from typing import Optional\n",
    "import torch\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "from torch_harmonics.quadrature import legendre_gauss_weights, lobatto_weights, clenshaw_curtiss_weights\n",
    "from torch_harmonics.legendre import _precompute_legpoly, _precompute_dlegpoly\n",
    "\n",
    "\n",
    "#@cache\n",
    "def legpoly_torch(mmax: int, lmax: int, x: torch.Tensor, norm: Optional[str]=\"ortho\", inverse: Optional[bool]=False, csphase: Optional[bool]=True) -> torch.Tensor:\n",
    "    r\"\"\"\n",
    "    Computes the values of (-1)^m c^l_m P^l_m(x) at the positions specified by x.\n",
    "    The resulting tensor has shape (mmax, lmax, len(x)). The Condon-Shortley Phase (-1)^m\n",
    "    can be turned off optionally.\n",
    "\n",
    "    method of computation follows\n",
    "    [1] Schaeffer, N.; Efficient spherical harmonic transforms aimed at pseudospectral numerical simulations, G3: Geochemistry, Geophysics, Geosystems.\n",
    "    [2] Rapp, R.H.; A Fortran Program for the Computation of Gravimetric Quantities from High Degree Spherical Harmonic Expansions, Ohio State University Columbus; report; 1982;\n",
    "        https://apps.dtic.mil/sti/citations/ADA123406\n",
    "    [3] Schrama, E.; Orbit integration based upon interpolated gravitational gradients\n",
    "    \"\"\"\n",
    "\n",
    "    # compute the tensor P^m_n:\n",
    "    nmax = max(mmax,lmax)\n",
    "    vdm = torch.zeros((nmax, nmax, len(x)), dtype=torch.float64)\n",
    "        \n",
    "    norm_factor = 1. if norm == \"ortho\" else np.sqrt(4 * np.pi)\n",
    "    norm_factor = 1. / norm_factor if inverse else norm_factor\n",
    "\n",
    "    # initial values to start the recursion\n",
    "    vdm[0,0,:] = norm_factor / np.sqrt(4 * np.pi)\n",
    "\n",
    "    # fill the diagonal and the lower diagonal\n",
    "    for l in range(1, nmax):\n",
    "        vdm[l-1, l, :] = np.sqrt(2*l + 1) * x * vdm[l-1, l-1, :]\n",
    "        vdm[l, l, :] = torch.sqrt( (2*l + 1) * (1 + x) * (1 - x) / 2.0 / l ) * vdm[l-1, l-1, :]\n",
    "\n",
    "    # fill the remaining values on the upper triangle and multiply b\n",
    "    for l in range(2, nmax):\n",
    "        for m in range(0, l-1):\n",
    "            vdm[m, l, :] = x * np.sqrt((2*l - 1) / (l - m) * (2*l + 1) / float(l + m)) * vdm[m, l-1, :] \\\n",
    "                            - np.sqrt((l + m - 1) / float((l - m) * (2*l + 1)) / float((2*l - 3) * (l - m - 1)) / float(l + m)) * vdm[m, l-2, :]\n",
    "\n",
    "    if norm == \"schmidt\":\n",
    "        for l in range(0, nmax):\n",
    "            if inverse:\n",
    "                vdm[:, l, : ] = vdm[:, l, : ] * np.sqrt(2*l + 1)\n",
    "            else:\n",
    "                vdm[:, l, : ] = vdm[:, l, : ] / np.sqrt(2*l + 1)\n",
    "\n",
    "    vdm = vdm[:mmax, :lmax]\n",
    "\n",
    "    if csphase:\n",
    "        for m in range(1, mmax, 2):\n",
    "            vdm[m] *= -1\n",
    "\n",
    "    return vdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([False, False, False, False, False, False, False, False, False, False])\n",
      "[[[ 0.28209479  0.28209479  0.28209479  0.28209479  0.28209479\n",
      "    0.28209479  0.28209479  0.28209479  0.28209479  0.28209479]\n",
      "  [ 0.47585318  0.42267213  0.33196122  0.21175808  0.07274038\n",
      "   -0.07274038 -0.21175808 -0.33196122 -0.42267213 -0.47585318]\n",
      "  [ 0.58204939  0.39266372  0.12136018 -0.13767009 -0.29442096\n",
      "   -0.29442096 -0.13767009  0.12136018  0.39266372  0.58204939]]\n",
      "\n",
      " [[-0.         -0.         -0.         -0.         -0.\n",
      "   -0.         -0.         -0.         -0.         -0.        ]\n",
      "  [-0.07840972 -0.1733215  -0.25350953 -0.31136067 -0.34164401\n",
      "   -0.34164401 -0.31136067 -0.25350953 -0.1733215  -0.07840972]\n",
      "  [-0.17075452 -0.33526279 -0.38513319 -0.30174011 -0.11373095\n",
      "    0.11373095  0.30174011  0.38513319  0.33526279  0.17075452]]\n",
      "\n",
      " [[ 0.          0.          0.          0.          0.\n",
      "    0.          0.          0.          0.          0.        ]\n",
      "  [ 0.          0.          0.          0.          0.\n",
      "    0.          0.          0.          0.          0.        ]\n",
      "  [ 0.01989547  0.09721184  0.20797105  0.31371971  0.37771299\n",
      "    0.37771299  0.31371971  0.20797105  0.09721184  0.01989547]]] \n",
      " tensor([[[ 0.2821,  0.2821,  0.2821,  0.2821,  0.2821,  0.2821,  0.2821,\n",
      "           0.2821,  0.2821,  0.2821],\n",
      "         [ 0.1119,  0.2568,  0.4025,  0.5485,  0.6945,  0.8405,  0.9865,\n",
      "           1.1325,  1.2782,  1.4231],\n",
      "         [-0.0135,  0.1982,  0.5791,  1.1292,  1.8485,  2.7368,  3.7941,\n",
      "           5.0198,  6.4124,  7.9638]],\n",
      "\n",
      "        [[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000, -0.0000, -0.0000,\n",
      "          -0.0000, -0.0000, -0.0000],\n",
      "         [-0.3363, -0.2939, -0.1958,     nan,     nan,     nan,     nan,\n",
      "              nan,     nan,     nan],\n",
      "         [-0.1722, -0.3454, -0.3608,     nan,     nan,     nan,     nan,\n",
      "              nan,     nan,     nan]],\n",
      "\n",
      "        [[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "           0.0000,  0.0000,  0.0000],\n",
      "         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,  0.0000,\n",
      "           0.0000,  0.0000,  0.0000],\n",
      "         [ 0.3660,  0.2796,  0.1241,     nan,     nan,     nan,     nan,\n",
      "              nan,     nan,     nan]]], dtype=torch.float64)\n"
     ]
    }
   ],
   "source": [
    "cost, w = legendre_gauss_weights(10, -1, 1)\n",
    "tq = np.flip(np.arccos(cost)).copy()\n",
    "x = torch.from_numpy(tq).clone()\n",
    "print(torch.isnan(x))\n",
    "lpt = legpoly_torch(3, 3, x)\n",
    "\n",
    "lpn = _precompute_legpoly(3, 3, tq)\n",
    "\n",
    "print(lpn, \"\\n\", lpt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dace",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
