Provide PyTorch implementations by wrapping JAX functions

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 💡
 

benchmarks/precompute_spherical.py

@@ -1,5 +1,6 @@
 """Benchmarks for precompute spherical transforms."""
 
+import jax
 import numpy as np
 from benchmarking import (
     BenchmarkSetup,
@@ -10,6 +11,7 @@
 
 import s2fft
 import s2fft.precompute_transforms
+from s2fft.utils import torch_wrapper
 
 L_VALUES = [8, 16, 32, 64, 128, 256]
 SPIN_VALUES = [0]
@@ -31,11 +33,17 @@ def setup_forward(method, L, sampling, spin, reality, recursion):
         sampling=sampling,
         reality=reality,
     )
-    kernel_function = (
-        s2fft.precompute_transforms.construct.spin_spherical_kernel_jax
-        if method == "jax"
-        else s2fft.precompute_transforms.construct.spin_spherical_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # arrays to a JAX arrays before converting to torch tensors which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flm = jax.numpy.asarray(flm)
+        f = jax.numpy.asarray(f)
+    if method.startswith("torch"):
+        flm, f = torch_wrapper.tree_map_jax_array_to_torch_tensor((flm, f))
+    kernel_function = s2fft.precompute_transforms.spherical._kernel_functions[method]
     kernel = kernel_function(
         L=L,
         spin=spin,
@@ -73,11 +81,16 @@ def setup_inverse(method, L, sampling, spin, reality, recursion):
         skip("Reality only valid for scalar fields (spin=0).")
     rng = np.random.default_rng()
     flm = s2fft.utils.signal_generator.generate_flm(rng, L, spin=spin, reality=reality)
-    kernel_function = (
-        s2fft.precompute_transforms.construct.spin_spherical_kernel_jax
-        if method == "jax"
-        else s2fft.precompute_transforms.construct.spin_spherical_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # array to a JAX array before converting to a torch tensor which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flm = jax.numpy.asarray(flm)
+    if method.startswith("torch"):
+        flm = torch_wrapper.jax_array_to_torch_tensor(flm)
+    kernel_function = s2fft.precompute_transforms.spherical._kernel_functions[method]
     kernel = kernel_function(
         L=L,
         spin=spin,

benchmarks/precompute_wigner.py

@@ -1,5 +1,6 @@
 """Benchmarks for precompute Wigner-d transforms."""
 
+import jax
 import numpy as np
 from benchmarking import (
     BenchmarkSetup,
@@ -10,6 +11,7 @@
 import s2fft
 import s2fft.precompute_transforms
 from s2fft.base_transforms import wigner as base_wigner
+from s2fft.utils import torch_wrapper
 
 L_VALUES = [16, 32, 64, 128, 256]
 N_VALUES = [2]
@@ -31,11 +33,17 @@ def setup_forward(method, L, N, L_lower, sampling, reality, mode):
         sampling=sampling,
         reality=reality,
     )
-    kernel_function = (
-        s2fft.precompute_transforms.construct.wigner_kernel_jax
-        if "jax" in method
-        else s2fft.precompute_transforms.construct.wigner_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # arrays to a JAX arrays before converting to torch tensors which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flmn = jax.numpy.asarray(flmn)
+        f = jax.numpy.asarray(f)
+    if method.startswith("torch"):
+        flmn, f = torch_wrapper.tree_map_jax_array_to_torch_tensor((flmn, f))
+    kernel_function = s2fft.precompute_transforms.wigner._kernel_functions[method]
     kernel = kernel_function(
         L=L, N=N, reality=reality, sampling=sampling, forward=True, mode=mode
     )
@@ -67,11 +75,16 @@ def forward(f, kernel, method, L, N, L_lower, sampling, reality, mode):
 def setup_inverse(method, L, N, L_lower, sampling, reality, mode):
     rng = np.random.default_rng()
     flmn = s2fft.utils.signal_generator.generate_flmn(rng, L, N, reality=reality)
-    kernel_function = (
-        s2fft.precompute_transforms.construct.wigner_kernel_jax
-        if method == "jax"
-        else s2fft.precompute_transforms.construct.wigner_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # arrays to a JAX arrays before converting to torch tensors which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flmn = jax.numpy.asarray(flmn)
+    if method.startswith("torch"):
+        flmn = torch_wrapper.jax_array_to_torch_tensor(flmn)
+    kernel_function = s2fft.precompute_transforms.wigner._kernel_functions[method]
     kernel = kernel_function(
         L=L, N=N, reality=reality, sampling=sampling, forward=False, mode=mode
     )

notebooks/torch_frontend.ipynb

@@ -28,27 +28,24 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This minimal tutorial demonstrates how to use the torch frontend for `S2FFT` to compute spherical harmonic transforms. Though `S2FFT` is primarily designed for JAX, this torch functionality is fully unit tested (including gradients) and can be used straightforwardly as a learnable layer within existing models."
+    "This minimal tutorial demonstrates how to use the torch frontend for `S2FFT` to compute spherical harmonic transforms. Though `S2FFT` is primarily designed for JAX, this torch functionality is fully unit tested (including gradients) and can be used straightforwardly as a learnable layer within existing models. As the torch functions wrap the JAX implementations we need to configure JAX to use 64-bit precision floating point types by default to ensure sufficient precision for the transforms - `S2FFT` will emit a warning if this has not been done."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "JAX is not using 64-bit precision. This will dramatically affect numerical precision at even moderate L.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
+    "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",
-    "from s2fft.precompute_transforms.construct import spin_spherical_kernel\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"
    ]
   },
@@ -65,33 +62,40 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "L = 64                                                          # Spherical harmonic bandlimit\n",
-    "rng = np.random.default_rng(1234951510)                         # Random seed for signal generator\n",
-    "flm = signal_generator.generate_flm(rng, L, using_torch=True)   # Random set of spherical harmonic coefficients"
+    "L = 64 \n",
+    "rng = np.random.default_rng(1234951510)\n",
+    "flm = torch.from_numpy(signal_generator.generate_flm(rng, L))"
    ]
   },
   {
    "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"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
+     ]
+    }
+   ],
    "source": [
-    "inverse_kernel = spin_spherical_kernel(L, using_torch=True, forward=False) \n",
-    "forward_kernel = spin_spherical_kernel(L, using_torch=True, 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"
    ]
   },
   {
@@ -100,53 +104,107 @@
    "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",
      "output_type": "stream",
      "text": [
-      "Mean absolute error = 1.1866908936078849e-14\n"
+      "Mean absolute error = 2.8472981477378884e-14\n"
      ]
     }
    ],
    "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))}\")"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.10.4 ('s2fft')",
+   "display_name": "s2fft",
    "language": "python",
    "name": "python3"
   },
@@ -160,14 +218,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.11.10"
   },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "3425e24474cbe920550266ea26b478634978cc419579f9dbcf479231067df6a3"
-   }
-  }
+  "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2

pyproject.toml

@@ -28,9 +28,8 @@ classifiers = [
 description = "Differentiable and accelerated spherical transforms with JAX"
 dependencies = [
     "numpy>=1.20",
-    "jax>=0.3.13",
+    "jax>=0.3.13,<0.6.0",
     "jaxlib",
-    "torch",
 ]
 dynamic = [
     "version",
@@ -74,6 +73,10 @@ tests = [
     "pytest-cov",
     "so3",
     "pyssht",
+    "torch",
+]
+torch = [
+    "torch",
 ]
 
 [tool.scikit-build]