Repository: jax-ml/jax
Pull request #25753 (2025-01-07 14:37:15)

Add `sph_harm_y` to `jax.scipy.special` and deprecate `sph_harm`
In scipy 1.15.0, [`scipy.special.sph_harm` has been deprecated](https://docs.scipy.org/doc/scipy/release/1.15.0-notes.html#scipy-special-improvements) in favor of [`scipy.special.sph_harm_y`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.sph_harm_y.html).

The [JEP 18137](https://jax.readthedocs.io/en/latest/jep/18137-numpy-scipy-scope.html#scipy-special) lays out some issues with the existing JAX implementation of `sph_harm`, but with those same caveats, it's not too hard to (mostly) support `sph_harm_y` as demonstrated in this PR.

The main differences between `sph_harm_y` and the deprecated `sph_harm` are:

1. the order of the `m` and `n` inputs are reversed, and
2. the definitions of `theta` and `phi` are swapped.

And these are simple enough to implement. But, `sph_harm_y` also includes a new `diff_n` parameter which adds support for returning partial derivatives of the spherical harmonics. It would be possible to add support for this, but it's not obvious that it's worth the extra work at this point.

My proposal would be to (as implemented here) add a minimal `sph_harm_y` implementation that continues to support existing use cases, but it does introduce further deviation from the scipy implementation, so I'm not 100% convinced that this is the right way to go.

@jakevdp — WDYT?

jakevdp: Seems reasonable. There are some fundamental issues with the `sph_harm` implementation, but it does seem that some users get value from it.

Patch:
diff --git a/jax/_src/deprecations.py b/jax/_src/deprecations.py
index 63b40e8a4..37f2f0264 100644
--- a/jax/_src/deprecations.py
+++ b/jax/_src/deprecations.py
@@ -135,3 +135,4 @@ register('jax-numpy-quantile-interpolation')
 register('jax-numpy-reduction-non-boolean-where')
 register('jax-numpy-trimzeros-not-1d-array')
 register('pallas-gpu-triton')
+register('jax-scipy-special-sph-harm')
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
index 2fffe6381..9c7528d90 100644
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -28,6 +28,7 @@ from jax import lax
 
 from jax._src import core
 from jax._src import custom_derivatives
+from jax._src import deprecations
 from jax._src import dtypes
 from jax._src.lax.lax import _const as _lax_const
 from jax._src.numpy.util import promote_args_inexact, promote_dtypes_inexact
@@ -1731,19 +1732,19 @@ def lpmn_values(m: int, n: int, z: Array, is_normalized: bool) -> Array:
 
 
 @partial(jit, static_argnums=(4,))
-def _sph_harm(m: Array,
-              n: Array,
+def _sph_harm(n: Array,
+              m: Array,
               theta: Array,
               phi: Array,
               n_max: int) -> Array:
   """Computes the spherical harmonics."""
 
-  cos_colatitude = jnp.cos(phi)
+  cos_colatitude = jnp.cos(theta)
 
   legendre = _gen_associated_legendre(n_max, cos_colatitude, True)
   legendre_val = legendre.at[abs(m), n, jnp.arange(len(n))].get(mode="clip")
 
-  angle = abs(m) * theta
+  angle = abs(m) * phi
   vandermonde = lax.complex(jnp.cos(angle), jnp.sin(angle))
   harmonics = lax.complex(legendre_val * jnp.real(vandermonde),
                           legendre_val * jnp.imag(vandermonde))
@@ -1756,6 +1757,58 @@ def _sph_harm(m: Array,
   return harmonics
 
 
+def sph_harm_y(n: Array,
+               m: Array,
+               theta: Array,
+               phi: Array,
+               diff_n: int | None = None,
+               n_max: int | None = None) -> Array:
+  r"""Computes the spherical harmonics.
+
+  The JAX version has one extra argument `n_max`, the maximum value in `n`.
+
+  The spherical harmonic of degree `n` and order `m` can be written as
+  :math:`Y_n^m(\theta, \phi) = N_n^m * P_n^m(\cos \theta) * \exp(i m \phi)`,
+  where :math:`N_n^m = \sqrt{\frac{\left(2n+1\right) \left(n-m\right)!}
+  {4 \pi \left(n+m\right)!}}` is the normalization factor and :math:`\theta` and
+  :math:`\phi` are the colatitude and longitude, respectively. :math:`N_n^m` is
+  chosen in the way that the spherical harmonics form a set of orthonormal basis
+  functions of :math:`L^2(S^2)`.
+
+  Args:
+    n: The degree of the harmonic; must have `n >= 0`. The standard notation for
+      degree in descriptions of spherical harmonics is `l (lower case L)`. We
+      use `n` here to be consistent with `scipy.special.sph_harm_y`. Return
+      values for `n < 0` are undefined.
+    m: The order of the harmonic; must have `|m| <= n`. Return values for
+      `|m| > n` are undefined.
+    theta: The polar (colatitudinal) coordinate; must be in [0, pi].
+    phi: The azimuthal (longitudinal) coordinate; must be in [0, 2*pi].
+    diff_n: Unsupported by JAX.
+    n_max: The maximum degree `max(n)`. If the supplied `n_max` is not the true
+      maximum value of `n`, the results are clipped to `n_max`. For example,
+      `sph_harm(m=jnp.array([2]), n=jnp.array([10]), theta, phi, n_max=6)`
+      actually returns
+      `sph_harm(m=jnp.array([2]), n=jnp.array([6]), theta, phi, n_max=6)`
+  Returns:
+    A 1D array containing the spherical harmonics at (m, n, theta, phi).
+  """
+  if diff_n is not None:
+    raise NotImplementedError(
+        "The 'diff_n' argument to jax.scipy.special.sph_harm_y is not supported.")
+
+  if jnp.isscalar(theta):
+    theta = jnp.array([theta])
+
+  if n_max is None:
+    n_max = np.max(n)
+  n_max = core.concrete_or_error(
+      int, n_max, 'The `n_max` argument of `jnp.scipy.special.sph_harm` must '
+      'be statically specified to use `sph_harm` within JAX transformations.')
+
+  return _sph_harm(n, m, theta, phi, n_max)
+
+
 def sph_harm(m: Array,
              n: Array,
              theta: Array,
@@ -1763,6 +1816,11 @@ def sph_harm(m: Array,
              n_max: int | None = None) -> Array:
   r"""Computes the spherical harmonics.
 
+  Note:
+    This function is deprecated, and :func:`~jax.scipy.special.sph_harm_y`
+    should be used instead, noting that the order of ``m`` and ``n`` are
+    reversed, and definitions of ``theta`` and ``phi`` are swapped.
+
   The JAX version has one extra argument `n_max`, the maximum value in `n`.
 
   The spherical harmonic of degree `n` and order `m` can be written as
@@ -1790,17 +1848,16 @@ def sph_harm(m: Array,
   Returns:
     A 1D array containing the spherical harmonics at (m, n, theta, phi).
   """
-
-  if jnp.isscalar(phi):
-    phi = jnp.array([phi])
-
-  if n_max is None:
-    n_max = np.max(n)
-  n_max = core.concrete_or_error(
-      int, n_max, 'The `n_max` argument of `jnp.scipy.special.sph_harm` must '
-      'be statically specified to use `sph_harm` within JAX transformations.')
-
-  return _sph_harm(m, n, theta, phi, n_max)
+  # Added 2025-01-06.
+  # TODO(dfm): Remove after deprecation period.
+  deprecations.warn(
+      "jax-scipy-special-sph-harm",
+      ("jax.scipy.special.sph_harm is deprecated. Please use "
+       "jax.scipy.special.sph_harm_y instead, noting that the order of `m` and "
+       "`n` are reversed, and definitions of `theta` and `phi` are swapped."),
+      stacklevel=2,
+  )
+  return sph_harm_y(n, m, phi, theta, n_max=n_max)
 
 
 # exponential integrals
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
index c83ead11a..2ffc65a1a 100644
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -57,6 +57,7 @@ from jax._src.scipy.special import (
   softmax as softmax,
   spence as spence,
   sph_harm as sph_harm,
+  sph_harm_y as sph_harm_y,
   xlog1py as xlog1py,
   xlogy as xlogy,
   zeta as zeta,


Test patch:
diff --git a/tests/lax_scipy_test.py b/tests/lax_scipy_test.py
index c30e8da65..a6f669e72 100644
--- a/tests/lax_scipy_test.py
+++ b/tests/lax_scipy_test.py
@@ -384,7 +384,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     self._CompileAndCheck(lax_fun, args_maker, rtol=1E-6, atol=1E-6)
 
   @jtu.ignore_warning(category=DeprecationWarning,
-                      message="`scipy.special.sph_harm` is deprecated")
+                      message=".*scipy.special.sph_harm.*")
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmAccuracy(self):
     m = jnp.arange(-3, 3)[:, None]
@@ -400,7 +400,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     self.assertAllClose(actual, expected, rtol=1e-8, atol=9e-5)
 
   @jtu.ignore_warning(category=DeprecationWarning,
-                      message="`scipy.special.sph_harm` is deprecated")
+                      message=".*scipy.special.sph_harm.*")
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmOrderZeroDegreeZero(self):
     """Tests the spherical harmonics of order zero and degree zero."""
@@ -415,7 +415,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     self.assertAllClose(actual, expected, rtol=1.1e-7, atol=3e-8)
 
   @jtu.ignore_warning(category=DeprecationWarning,
-                      message="`scipy.special.sph_harm` is deprecated")
+                      message=".*scipy.special.sph_harm.*")
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmOrderZeroDegreeOne(self):
     """Tests the spherical harmonics of order one and degree zero."""
@@ -430,7 +430,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     self.assertAllClose(actual, expected, rtol=2e-7, atol=6e-8)
 
   @jtu.ignore_warning(category=DeprecationWarning,
-                      message="`scipy.special.sph_harm` is deprecated")
+                      message=".*scipy.special.sph_harm.*")
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmOrderOneDegreeOne(self):
     """Tests the spherical harmonics of order one and degree one."""
@@ -452,7 +452,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
     dtype=jtu.dtypes.all_integer,
   )
   @jtu.ignore_warning(category=DeprecationWarning,
-                      message="`scipy.special.sph_harm` is deprecated")
+                      message=".*scipy.special.sph_harm.*")
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmForJitAndAgainstNumpy(self, l_max, num_z, dtype):
     """Tests against JIT compatibility and Numpy."""
@@ -478,7 +478,7 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
       self._CheckAgainstNumpy(osp_special.sph_harm, lsp_special_fn, args_maker)
 
   @jtu.ignore_warning(category=DeprecationWarning,
-                      message="`scipy.special.sph_harm` is deprecated")
+                      message=".*scipy.special.sph_harm.*")
   @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
   def testSphHarmCornerCaseWithWrongNmax(self):
     """Tests the corner case where `n_max` is not the maximum value of `n`."""
@@ -495,6 +495,35 @@ class LaxBackedScipyTests(jtu.JaxTestCase):
 
     self.assertAllClose(actual, expected, rtol=1e-8, atol=9e-5)
 
+  @jtu.sample_product(
+    [dict(l_max=l_max, num_z=num_z)
+      for l_max, num_z in zip([1, 3, 8, 10], [2, 6, 7, 8])
+    ],
+    dtype=jtu.dtypes.all_integer,
+  )
+  @jax.numpy_dtype_promotion('standard')  # This test explicitly exercises dtype promotion
+  def testSphHarmY(self, l_max, num_z, dtype):
+    if jtu.is_device_tpu(6, "e"):
+      self.skipTest("TODO(b/364258243): fails on TPU v6e")
+    n_max = l_max
+    shape = (num_z,)
+    rng = jtu.rand_int(self.rng(), -l_max, l_max + 1)
+
+    def args_maker():
+      m = rng(shape, dtype)
+      n = abs(m)
+      theta = np.linspace(-2.0, 1.0, num_z)
+      phi = np.linspace(-4.0, 5.0, num_z)
+      return n, m, theta, phi
+
+    lsp_special_fn = partial(lsp_special.sph_harm_y, n_max=n_max)
+    self._CompileAndCheck(lsp_special_fn, args_maker)
+    if scipy_version < (1, 15, 0):
+      osp_special_fn = lambda n, m, theta, phi: osp_special.sph_harm(m, n, phi, theta)
+    else:
+      osp_special_fn = osp_special.sph_harm_y
+    self._CheckAgainstNumpy(osp_special_fn, lsp_special_fn, args_maker)
+
   @jtu.sample_product(
     n_zero_sv=n_zero_svs,
     degeneracy=degeneracies,