diff --git a/dipy/reconst/shm.py b/dipy/reconst/shm.py index c93acded94..ded51fb993 100644 --- a/dipy/reconst/shm.py +++ b/dipy/reconst/shm.py @@ -29,7 +29,7 @@ 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 @@ -37,6 +37,14 @@ 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 diff --git a/dipy/reconst/tests/test_shm.py b/dipy/reconst/tests/test_shm.py index bc78566a75..b4a237670b 100644 --- a/dipy/reconst/tests/test_shm.py +++ b/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()