Add torch wrappers for Wigner OTF transforms

Author: matt-graham

s2fft/transforms/wigner.py

@@ -8,6 +8,7 @@
 import s2fft
 from s2fft.sampling import so3_samples as samples
 from s2fft.transforms import c_backend_spherical as c_sph
+from s2fft.utils import torch_wrapper
 
 
 def inverse(
@@ -84,7 +85,7 @@ def inverse(
     if method not in _inverse_functions:
         raise ValueError(f"Method {method} not recognised.")
 
-    if N >= 8 and method in ("numpy", "jax"):
+    if N >= 8 and method in ("numpy", "jax", "torch"):
         raise Warning("Recursive transform may provide lower precision beyond N ~ 8")
 
     inverse_kwargs = {
@@ -96,7 +97,7 @@ def inverse(
         "reality": reality,
     }
 
-    if method in ("jax", "numpy"):
+    if method in ("jax", "numpy", "torch"):
         inverse_kwargs.update(nside=nside, precomps=precomps)
 
     if method == "jax_ssht":
@@ -288,6 +289,9 @@ def func(flm, spin, p0, p1, p2, p3, p4):
     return f
 
 
+inverse_torch = torch_wrapper.wrap_as_torch_function(inverse_jax)
+
+
 def inverse_jax_ssht(
     flmn: jnp.ndarray,
     L: int,
@@ -413,7 +417,7 @@ def forward(
     if method not in _inverse_functions:
         raise ValueError(f"Method {method} not recognised.")
 
-    if N >= 8 and method in ("numpy", "jax"):
+    if N >= 8 and method in ("numpy", "jax", "torch"):
         raise Warning("Recursive transform may provide lower precision beyond N ~ 8")
 
     forward_kwargs = {
@@ -425,7 +429,7 @@ def forward(
         "reality": reality,
     }
 
-    if method in ("jax", "numpy"):
+    if method in ("jax", "numpy", "torch"):
         forward_kwargs.update(nside=nside, precomps=precomps)
 
     if method == "jax_ssht":
@@ -642,6 +646,9 @@ def func(fba, spin, p0, p1, p2, p3, p4):
     return flmn
 
 
+forward_torch = torch_wrapper.wrap_as_torch_function(forward_jax)
+
+
 def forward_jax_ssht(
     f: jnp.ndarray,
     L: int,
@@ -829,10 +836,12 @@ def _fban_to_f(fban: jnp.ndarray, L: int, N: int, reality: bool = False) -> jnp.
     "numpy": inverse_numpy,
     "jax": inverse_jax,
     "jax_ssht": inverse_jax_ssht,
+    "torch": inverse_torch,
 }
 
 _forward_functions = {
     "numpy": forward_numpy,
     "jax": forward_jax,
     "jax_ssht": forward_jax_ssht,
+    "torch": forward_torch,
 }

tests/test_wigner_transform.py

@@ -16,9 +16,16 @@
 N_to_test = [2]
 L_lower_to_test = [0, 2]
 sampling_to_test = ["mw", "mwss", "dh", "gl"]
-method_to_test = ["numpy", "jax"]
+method_to_test = ["numpy", "jax", "torch"]
 reality_to_test = [False, True]
 
+_generate_precomputes_functions = {
+    "jax": generate_precomputes_wigner_jax,
+    "numpy": generate_precomputes_wigner,
+    # torch method wraps jax so use jax to generate precomputess
+    "torch": generate_precomputes_wigner_jax,
+}
+
 
 @pytest.mark.parametrize("L", L_to_test)
 @pytest.mark.parametrize("N", N_to_test)
@@ -38,15 +45,9 @@ def test_inverse_wigner_transform(
 ):
     flmn = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     f_check = base_wigner.inverse(flmn, L, N, L_lower, sampling, reality)
-
-    if method.lower() == "jax":
-        precomps = generate_precomputes_wigner_jax(
-            L, N, sampling, None, False, reality, L_lower
-        )
-    else:
-        precomps = generate_precomputes_wigner(
-            L, N, sampling, None, False, reality, L_lower
-        )
+    precomps = _generate_precomputes_functions[method](
+        L, N, sampling, None, False, reality, L_lower
+    )
     f = wigner.inverse(flmn, L, N, None, sampling, method, reality, precomps, L_lower)
     np.testing.assert_allclose(f, f_check, atol=1e-14)
 
@@ -69,15 +70,9 @@ def test_forward_wigner_transform(
 ):
     flmn = flmn_generator(L=L, N=N, L_lower=L_lower, reality=reality)
     f = base_wigner.inverse(flmn, L, N, L_lower, sampling, reality)
-
-    if method.lower() == "jax":
-        precomps = generate_precomputes_wigner_jax(
-            L, N, sampling, None, True, reality, L_lower
-        )
-    else:
-        precomps = generate_precomputes_wigner(
-            L, N, sampling, None, True, reality, L_lower
-        )
+    precomps = _generate_precomputes_functions[method](
+        L, N, sampling, None, True, reality, L_lower
+    )
     flmn_check = wigner.forward(
         f, L, N, None, sampling, method, reality, precomps, L_lower
     )