Merge pull request #413 from Garyfallidis/faster_sph_harm

Faster spherical harmonics

dipy/reconst/shm.py

@@ -29,14 +29,22 @@
 from numpy.linalg import pinv, svd
 from numpy.random import randint
 from dipy.reconst.odf import OdfModel, OdfFit
-from scipy.special import sph_harm, lpn
+from scipy.special import sph_harm, lpn, lpmv, gammaln
 from dipy.core.sphere import Sphere
 import dipy.core.gradients as grad
 from dipy.sims.voxel import single_tensor, all_tensor_evecs
 from dipy.core.geometry import cart2sphere
 from dipy.core.onetime import auto_attr
 from dipy.reconst.cache import Cache
 
+from distutils.version import StrictVersion
+import scipy
+
+if StrictVersion(scipy.__version__) >= StrictVersion('0.15.0'):
+    SCIPY_15_PLUS = True
+else:
+    SCIPY_15_PLUS = False
+
 
 def _copydoc(obj):
     def bandit(f):
@@ -149,10 +157,48 @@ def gen_dirac(m, n, theta, phi):
     return real_sph_harm(m, n, theta, phi)
 
 
-def real_sph_harm(m, n, theta, phi):
+def spherical_harmonics(m, n, theta, phi):
+    r""" Compute spherical harmonics
+
+    This may take scalar or array arguments. The inputs will be broadcasted
+    against each other.
+
+    Parameters
+    ----------
+    m : int ``|m| <= n``
+        The order of the harmonic.
+    n : int ``>= 0``
+        The degree of the harmonic.
+    theta : float [0, 2*pi]
+        The azimuthal (longitudinal) coordinate.
+    phi : float [0, pi]
+        The polar (colatitudinal) coordinate.
+
+    Returns
+    -------
+    y_mn : complex float
+        The harmonic $Y^m_n$ sampled at `theta` and `phi`.
+
+    Notes
+    -----
+    This is a faster implementation of scipy.special.sph_harm for
+    scipy version < 0.15.0.
+
     """
-    Compute real spherical harmonics, where the real harmonic $Y^m_n$ is
-    defined to be:
+    if SCIPY_15_PLUS:
+        return sph_harm(m, n, theta, phi)
+    x = np.cos(phi)
+    val = lpmv(m, n, x).astype(complex)
+    val *= np.sqrt((2 * n + 1) / 4.0 / np.pi)
+    val *= np.exp(0.5 * (gammaln(n - m + 1) - gammaln(n + m + 1)))
+    val = val * np.exp(1j * m * theta)
+    return val
+
+
+def real_sph_harm(m, n, theta, phi):
+    r""" Compute real spherical harmonics.
+
+    Where the real harmonic $Y^m_n$ is defined to be:
 
         Real($Y^m_n$) * sqrt(2) if m > 0
         $Y^m_n$                 if m == 0
@@ -183,7 +229,8 @@ def real_sph_harm(m, n, theta, phi):
     """
     # dipy uses a convention for theta and phi that is reversed with respect to
     # function signature of scipy.special.sph_harm
-    sh = sph_harm(np.abs(m), n, phi, theta)
+    sh = spherical_harmonics(np.abs(m), n, phi, theta)
+
     real_sh = np.where(m > 0, sh.imag, sh.real)
     real_sh *= np.where(m == 0, 1., np.sqrt(2))
     return real_sh

dipy/reconst/tests/test_shm.py

@@ -5,6 +5,7 @@
 from nose.tools import assert_equal, assert_raises, assert_true
 from numpy.testing import assert_array_equal, assert_array_almost_equal
 import numpy.testing as npt
+from scipy.special import sph_harm as sph_harm_sp
 
 from dipy.core.sphere import hemi_icosahedron
 from dipy.core.gradients import gradient_table
@@ -23,7 +24,8 @@
                               OpdtModel, normalize_data, hat, lcr_matrix,
                               smooth_pinv, bootstrap_data_array,
                               bootstrap_data_voxel, ResidualBootstrapWrapper,
-                              CsaOdfModel, QballModel, SphHarmFit)
+                              CsaOdfModel, QballModel, SphHarmFit,
+                              spherical_harmonics)
 
 def test_order_from_ncoeff():
     """
@@ -373,6 +375,46 @@ def test_sf_to_sh():
     assert_array_almost_equal(odf2d, odf2d_sf, 2)
 
 
+def test_faster_sph_harm():
+
+    sh_order = 8
+
+    m, n = sph_harm_ind_list(sh_order)
+
+    theta = np.array([1.61491146,  0.76661665,  0.11976141,  1.20198246,  1.74066314,
+                      1.5925956 ,  2.13022055,  0.50332859,  1.19868988,  0.78440679,
+                      0.50686938,  0.51739718,  1.80342999,  0.73778957,  2.28559395,
+                      1.29569064,  1.86877091,  0.39239191,  0.54043037,  1.61263047,
+                      0.72695314,  1.90527318,  1.58186125,  0.23130073,  2.51695237,
+                      0.99835604,  1.2883426 ,  0.48114057,  1.50079318,  1.07978624,
+                      1.9798903 ,  2.36616966,  2.49233299,  2.13116602,  1.36801518,
+                      1.32932608,  0.95926683,  1.070349  ,  0.76355762,  2.07148422,
+                      1.50113501,  1.49823314,  0.89248164,  0.22187079,  1.53805373,
+                      1.9765295 ,  1.13361568,  1.04908355,  1.68737368,  1.91732452,
+                      1.01937457,  1.45839   ,  0.49641525,  0.29087155,  0.52824641,
+                      1.29875871,  1.81023541,  1.17030475,  2.24953206,  1.20280498,
+                      0.76399964,  2.16109722,  0.79780421,  0.87154509])
+
+    phi = np.array([-1.5889514 , -3.11092733, -0.61328674, -2.4485381 ,  2.88058822,
+                    2.02165946, -1.99783366,  2.71235211,  1.41577992, -2.29413676,
+                    -2.24565773, -1.55548635,  2.59318232, -1.84672472, -2.33710739,
+                    2.12111948,  1.87523722, -1.05206575, -2.85381987, -2.22808984,
+                    2.3202034 , -2.19004474, -1.90358372,  2.14818373,  3.1030696 ,
+                    -2.86620183, -2.19860123, -0.45468447, -3.0034923 ,  1.73345011,
+                    -2.51716288,  2.49961525, -2.68782986,  2.69699056,  1.78566133,
+                    -1.59119705, -2.53378963, -2.02476738,  1.36924987,  2.17600517,
+                    2.38117241,  2.99021511, -1.4218007 , -2.44016802, -2.52868164,
+                    3.01531658,  2.50093627, -1.70745826, -2.7863931 , -2.97359741,
+                    2.17039906,  2.68424643,  1.77896086,  0.45476215,  0.99734418,
+                    -2.73107896,  2.28815009,  2.86276506,  3.09450274, -3.09857384,
+                    -1.06955885, -2.83826831,  1.81932195,  2.81296654])
+
+    sh = spherical_harmonics(m, n, theta[:, None], phi[:, None])
+    sh2 = sph_harm_sp(m, n, theta[:, None], phi[:, None])
+
+    assert_array_almost_equal(sh, sh2, 8)
+
+
 if __name__ == "__main__":
     import nose
     nose.runmodule()