Update README and notebook to indicate wider torch support

Author: matt-graham

README.md

@@ -195,7 +195,7 @@ f = fft.wigner.inverse_jax(flmn, L, N, method="jax")
 For further details on usage see the [documentation](https://astro-informatics.github.io/s2fft/) and associated [notebooks](https://astro-informatics.github.io/s2fft/tutorials/spherical_harmonic/spherical_harmonic_transform.html).
 
 > [!NOTE]  
-> We also provide PyTorch support for the precompute version of our transforms, as demonstrated in the [_Torch frontend_ tutorial notebook](https://astro-informatics.github.io/s2fft/tutorials/torch_frontend/torch_frontend.html).
+> We also provide PyTorch support for our transforms, as demonstrated in the [_Torch frontend_ tutorial notebook](https://astro-informatics.github.io/s2fft/tutorials/torch_frontend/torch_frontend.html). This wraps the JAX implementations so JAX will need to be installed in addition to PyTorch.
 
 ## SSHT & HEALPix wrappers 💡
 

notebooks/torch_frontend.ipynb

@@ -40,8 +40,11 @@
     "import jax\n",
     "jax.config.update(\"jax_enable_x64\", True)\n",
     "import torch \n",
-    "import numpy as np \n",
-    "from s2fft.precompute_transforms.spherical import inverse, forward\n",
+    "import numpy as np\n",
+    "from s2fft.transforms.spherical import inverse, forward\n",
+    "from s2fft.precompute_transforms.spherical import (\n",
+    "    inverse as precompute_inverse, forward as precompute_forward\n",
+    ")\n",
     "from s2fft.precompute_transforms.construct import spin_spherical_kernel_torch\n",
     "from s2fft.utils import signal_generator"
    ]
@@ -68,7 +71,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "For the fully precompute transform we must also generate the precompute kernels which we store as a torch tensors."
+    "Now lets calculate the signal on the sphere by applying the inverse spherical harmonic transform"
    ]
   },
   {
@@ -85,15 +88,14 @@
     }
    ],
    "source": [
-    "inverse_kernel = spin_spherical_kernel_torch(L, forward=False) \n",
-    "forward_kernel = spin_spherical_kernel_torch(L, forward=True) "
+    "f = inverse(flm, L, method=\"torch\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now lets calculate the signal on the sphere by applying the inverse spherical harmonic transform"
+    "To calculate the corresponding spherical harmonic representation execute"
    ]
   },
   {
@@ -102,36 +104,90 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "f = inverse(flm, L, 0, inverse_kernel, method=\"torch\")"
+    "flm_check = forward(f, L, method=\"torch\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "To calculate the corresponding spherical harmonic representation execute"
+    "Finally, lets check the error on the round trip is as expected for 64 bit machine precision floating point arithmetic"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 6,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Mean absolute error = 2.8915048238993476e-14\n"
+     ]
+    }
+   ],
    "source": [
-    "flm_check = forward(f, L, 0, forward_kernel, method=\"torch\")"
+    "print(f\"Mean absolute error = {np.nanmean(np.abs(flm_check - flm))}\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Finally, lets check the error on the roundtrip is at 64bit machine precision"
+    "For the fully precompute transform we must also generate the precompute kernels which we store as a torch tensors."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 7,
    "metadata": {},
+   "outputs": [],
+   "source": [
+    "inverse_kernel = spin_spherical_kernel_torch(L, forward=False) \n",
+    "forward_kernel = spin_spherical_kernel_torch(L, forward=True) "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We then pass the kernels as additional arguments to the transform functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'orward_kernel' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[8], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m precompute_f \u001b[38;5;241m=\u001b[39m precompute_inverse(flm, L, kernel\u001b[38;5;241m=\u001b[39minverse_kernel, method\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtorch\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m precompute_flm_check \u001b[38;5;241m=\u001b[39m precompute_forward(f, L, kernel\u001b[38;5;241m=\u001b[39m\u001b[43morward_kernel\u001b[49m, method\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtorch\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+      "\u001b[0;31mNameError\u001b[0m: name 'orward_kernel' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "precompute_f = precompute_inverse(flm, L, kernel=inverse_kernel, method=\"torch\")\n",
+    "precompute_flm_check = precompute_forward(f, L, kernel=forward_kernel, method=\"torch\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Again, we check the error on the round trip is as expected"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
    "outputs": [
     {
      "name": "stdout",
@@ -142,7 +198,7 @@
     }
    ],
    "source": [
-    "print(f\"Mean absolute error = {np.nanmean(np.abs(flm_check - flm))}\")"
+    "print(f\"Mean absolute error = {np.nanmean(np.abs(precompute_flm_check - flm))}\")"
    ]
   }
  ],