Merge pull request #202 from astro-informatics/feature/notebook_plots

Feature/notebook plots

Author: jasonmcewen

notebooks/JAX_HEALPix_frontend.ipynb

@@ -16,8 +16,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install s2fft\n",
-    "!pip install s2fft &> /dev/null"
+    "import sys\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "# Install s2fft and data if running on google colab.\n",
+    "if IN_COLAB:\n",
+    "    !pip install s2fft &> /dev/null"
    ]
   },
   {
@@ -42,11 +46,12 @@
     "import numpy as np\n",
     "import s2fft \n",
     "\n",
-    "L = 1024\n",
-    "nside = 512\n",
+    "L = 128\n",
+    "nside = 64\n",
     "method = \"jax_healpy\"\n",
     "sampling = \"healpix\"\n",
-    "flm = np.random.randn(L, 2*L-1) + 1j*np.random.randn(L, 2*L-1)\n",
+    "rng = np.random.default_rng(23457801234570)\n",
+    "flm = s2fft.utils.signal_generator.generate_flm(rng, L)\n",
     "f = s2fft.inverse(flm, L, nside=nside, sampling=sampling, method=method)"
    ]
   },
@@ -183,7 +188,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.10.0"
   },
   "orig_nbformat": 4,
   "vscode": {

notebooks/JAX_SSHT_frontend.ipynb

@@ -16,8 +16,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install s2fft\n",
-    "!pip install s2fft &> /dev/null"
+    "import sys\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "# Install s2fft and data if running on google colab.\n",
+    "if IN_COLAB:\n",
+    "    !pip install s2fft &> /dev/null"
    ]
   },
   {
@@ -42,9 +46,10 @@
     "import numpy as np\n",
     "import s2fft \n",
     "\n",
-    "L = 1024\n",
+    "L = 128\n",
     "method = \"jax_ssht\"\n",
-    "flm = np.random.randn(L, 2*L-1) + 1j*np.random.randn(L, 2*L-1)\n",
+    "rng = np.random.default_rng(23457801234570)\n",
+    "flm = s2fft.utils.signal_generator.generate_flm(rng, L)\n",
     "f = s2fft.inverse(flm, L, method=method)"
    ]
   },
@@ -107,7 +112,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Mean absolute error = 4.909423754134027e-11\n"
+      "Mean absolute error = 7.784372519411174e-13\n"
      ]
     }
    ],
@@ -181,7 +186,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.10.0"
   },
   "orig_nbformat": 4,
   "vscode": {

notebooks/spherical_harmonic_transform.ipynb

@@ -12,12 +12,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install s2fft\n",
-    "!pip install s2fft &> /dev/null"
+    "import sys\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "# Install a spherical plotting package.\n",
+    "!pip install cartopy &> /dev/null\n",
+    "\n",
+    "# Install s2fft and data if running on google colab.\n",
+    "if IN_COLAB:\n",
+    "    !pip install s2fft &> /dev/null\n",
+    "    !mkdir data/\n",
+    "    !wget https://github.com/astro-informatics/s2fft/raw/main/notebooks/data/Gaia_EDR3_flux.npy -P data/ &> /dev/null"
    ]
   },
   {
@@ -32,20 +41,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import jax\n",
     "jax.config.update(\"jax_enable_x64\", True)\n",
     "\n",
     "import numpy as np\n",
+    "from matplotlib import pyplot as plt \n",
+    "import cartopy.crs as ccrs \n",
     "import s2fft \n",
     "\n",
-    "L = 256\n",
     "sampling = \"mw\"\n",
-    "flm = np.random.randn(L, 2*L-1) + 1j*np.random.randn(L, 2*L-1)\n",
-    "f = s2fft.inverse_jax(flm, L)"
+    "f = np.load('data/Gaia_EDR3_flux.npy')\n",
+    "L = f.shape[0]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Lets look at the input signal"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.figure(figsize=(10,5))\n",
+    "ax = plt.axes(projection=ccrs.Mollweide())\n",
+    "im = ax.imshow(f, transform=ccrs.PlateCarree(), cmap='magma')\n",
+    "plt.axis('off')\n",
+    "plt.show()"
    ]
   },
   {
@@ -62,7 +92,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -81,7 +111,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -103,7 +133,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -122,7 +152,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -144,34 +174,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Mean absolute error = 8.478196507592078e-11\n"
-     ]
-    }
-   ],
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "print(f\"Mean absolute error = {np.nanmean(np.abs(f_recov - f))}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Mean absolute error using precomputes = 8.478196507592078e-11\n"
-     ]
-    }
-   ],
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
    "source": [
     "print(f\"Mean absolute error using precomputes = {np.nanmean(np.abs(f_recov_pre - f))}\")"
    ]
@@ -193,7 +207,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.10.0"
   },
   "orig_nbformat": 4,
   "vscode": {

notebooks/spherical_rotation.ipynb

@@ -16,8 +16,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install s2fft\n",
-    "!pip install s2fft &> /dev/null"
+    "import sys\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "# Install s2fft and data if running on google colab.\n",
+    "if IN_COLAB:\n",
+    "    !pip install s2fft &> /dev/null"
    ]
   },
   {

notebooks/torch_frontend.ipynb

@@ -16,8 +16,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install s2fft\n",
-    "!pip install s2fft &> /dev/null"
+    "import sys\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "# Install s2fft and data if running on google colab.\n",
+    "if IN_COLAB:\n",
+    "    !pip install s2fft &> /dev/null"
    ]
   },
   {

notebooks/wigner_transform.ipynb

@@ -16,8 +16,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Install s2fft\n",
-    "!pip install s2fft &> /dev/null"
+    "import sys\n",
+    "IN_COLAB = 'google.colab' in sys.modules\n",
+    "\n",
+    "# Install s2fft and data if running on google colab.\n",
+    "if IN_COLAB:\n",
+    "    !pip install s2fft &> /dev/null"
    ]
   },
   {
@@ -47,7 +51,7 @@
     "L = 128\n",
     "N = 3\n",
     "reality = True\n",
-    "rng = np.random.default_rng(0)\n",
+    "rng = np.random.default_rng(83459)\n",
     "flmn = s2fft.utils.signal_generator.generate_flmn(rng, L, N, reality=reality)"
    ]
   },
@@ -190,7 +194,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.13"
+   "version": "3.10.0"
   },
   "orig_nbformat": 4,
   "vscode": {

s2fft/utils/signal_generator.py

@@ -38,16 +38,16 @@ def generate_flm(
 
     for el in range(max(L_lower, abs(spin)), L):
         if reality:
-            flm[el, 0 + L - 1] = rng.uniform()
+            flm[el, 0 + L - 1] = rng.normal()
         else:
-            flm[el, 0 + L - 1] = rng.uniform() + 1j * rng.uniform()
+            flm[el, 0 + L - 1] = rng.normal() + 1j * rng.normal()
 
         for m in range(1, el + 1):
-            flm[el, m + L - 1] = rng.uniform() + 1j * rng.uniform()
+            flm[el, m + L - 1] = rng.normal() + 1j * rng.normal()
             if reality:
                 flm[el, -m + L - 1] = (-1) ** m * np.conj(flm[el, m + L - 1])
             else:
-                flm[el, -m + L - 1] = rng.uniform() + 1j * rng.uniform()
+                flm[el, -m + L - 1] = rng.normal() + 1j * rng.normal()
 
     return torch.from_numpy(flm) if using_torch else flm
 
@@ -86,22 +86,22 @@ def generate_flmn(
     for n in range(-N + 1, N):
         for el in range(max(L_lower, abs(n)), L):
             if reality:
-                flmn[N - 1 + n, el, 0 + L - 1] = rng.uniform()
+                flmn[N - 1 + n, el, 0 + L - 1] = rng.normal()
                 flmn[N - 1 - n, el, 0 + L - 1] = (-1) ** n * flmn[
                     N - 1 + n,
                     el,
                     0 + L - 1,
                 ]
             else:
-                flmn[N - 1 + n, el, 0 + L - 1] = rng.uniform() + 1j * rng.uniform()
+                flmn[N - 1 + n, el, 0 + L - 1] = rng.normal() + 1j * rng.normal()
 
             for m in range(1, el + 1):
-                flmn[N - 1 + n, el, m + L - 1] = rng.uniform() + 1j * rng.uniform()
+                flmn[N - 1 + n, el, m + L - 1] = rng.normal() + 1j * rng.normal()
                 if reality:
                     flmn[N - 1 - n, el, -m + L - 1] = (-1) ** (m + n) * np.conj(
                         flmn[N - 1 + n, el, m + L - 1]
                     )
                 else:
-                    flmn[N - 1 + n, el, -m + L - 1] = rng.uniform() + 1j * rng.uniform()
+                    flmn[N - 1 + n, el, -m + L - 1] = rng.normal() + 1j * rng.normal()
 
     return torch.from_numpy(flmn) if using_torch else flmn