Use correct kernel construction functions

matt-graham · matt-graham · commit 8919b2f7baa1 · 2024-12-11T13:15:18.000Z
diff --git a/s2fft/precompute_transforms/spherical.py b/s2fft/precompute_transforms/spherical.py
@@ -624,7 +624,7 @@ def forward_transform_torch(
 }
 
 _kernel_functions = {
-    "numpy": partial(construct.fourier_wigner_kernel, using_torch=False),
-    "jax": construct.fourier_wigner_kernel_jax,
-    "torch": partial(construct.fourier_wigner_kernel, using_torch=True),
+    "numpy": partial(construct.spin_spherical_kernel, using_torch=False),
+    "jax": construct.spin_spherical_kernel_jax,
+    "torch": partial(construct.spin_spherical_kernel, using_torch=True),
 }

Original file line number	Diff line number	Diff line change
`@@ -624,7 +624,7 @@ def forward_transform_torch(`
`624`	`624`	`}`
`625`	`625`
`626`	`626`	`_kernel_functions = {`
`627`		`- "numpy": partial(construct.fourier_wigner_kernel, using_torch=False),`
`628`		`- "jax": construct.fourier_wigner_kernel_jax,`
`629`		`- "torch": partial(construct.fourier_wigner_kernel, using_torch=True),`
	`627`	`+ "numpy": partial(construct.spin_spherical_kernel, using_torch=False),`
	`628`	`+ "jax": construct.spin_spherical_kernel_jax,`
	`629`	`+ "torch": partial(construct.spin_spherical_kernel, using_torch=True),`
`630`	`630`	`}`