Refactored CSD estimation framework (#10)

* Refactored cvxpy csd optimizer and automatic algorithm selection based on whether volume fractions are fixed.

* Added tests for CSD framework.

* Added laplace-beltrami regularization and norm of laplacian for FOD and accompanying tests.

* Updated CSD example notebook and added generalized CSD tutorial notebook.

README.md

@@ -12,8 +12,9 @@ The Dmipy software package facilitates the **reproducible estimation of diffusio
 - Any combination of tissue models (e.g. Gaussian, Cylinder, Sphere) and axon bundle representation (e.g. orientation-dispersed/diameter-distributed) can be combined into a multi-compartment model.
 - Any appropriate model can be orientation-dispersed and/or axon diameter-distributed.
 - Any predefined or custom parameter constraints or relations can be imposed.
-- Free choice of global optimizer to fit your model to the data (Brute-Force or Stochastic)
-- Ability to also fit the spherical mean of any multi-compartment model to the spherical mean of the data.
+- Free choice of global optimizer to fit your model to the data (Brute-Force or Stochastic).
+- Fit the spherical mean of any multi-compartment model to the spherical mean of the data.
+- Generalized multi-compartment constrained spherical deconvolution.
 
 **Human Connectome Project Data Interface**
 Dmipy enables you to directly download any HCP subject data using your own credentials.
@@ -65,7 +66,7 @@ To get a feeling for how to use Dmipy, we provide a few tutorial notebooks:
 - [Brute Force to Gradient Descent (Brute2Fine)](http://nbviewer.jupyter.org/github/AthenaEPI/dmipy/blob/master/examples/example_brute_force_optimization.ipynb)
 - [Stochastic (MIX) [Farooq et al. 2016]](http://nbviewer.jupyter.org/github/AthenaEPI/dmipy/blob/master/examples/example_stochastic_mix_optimization.ipynb)
 ### Constrained Spherical Deconvolution Optimizers
-- General-Purpose Multi-Compartment CSD Optimizer
+- [Generalized Multi-Shell Multi-Compartment CSD [Tournier et al. 2007, Jeurissen et al. 2014]](https://github.com/AthenaEPI/dmipy/blob/refactor_cvxpy_csd/examples/example_generalized_csd_optimizer.ipynb)
 ## Dmipy implementations of Microstructure Models in Literature
 Dmipy uses HCP data to illustrate microstructure model examples. To reproduce these examples, dmipy provides a direct way to download HCP data (using your own AWS credentials) in the [HCP tutorial](http://nbviewer.jupyter.org/github/AthenaEPI/mipy/blob/master/examples/tutorial_human_connectome_project_aws.ipynb).
 ### Single Bundle Models

dmipy/core/fitted_modeling_framework.py

@@ -5,6 +5,7 @@
     unitsphere2cart_Nd, T1_tortuosity, fractional_parameter, cart2mu)
 from ..utils.spherical_mean import (
     estimate_spherical_mean_multi_shell)
+from dipy.reconst.shm import sph_harm_ind_list
 
 
 __all__ = [
@@ -630,10 +631,32 @@ def anisotropy_index(self):
         sh_coef = self.fitted_parameters['sh_coeff']
         sh_0 = sh_coef[..., 0] ** 2
         sh_sum_squared = np.sum(sh_coef ** 2, axis=-1)
-        AI = np.sqrt(1 - sh_0 / sh_sum_squared)
-        AI[~self.mask] = 0.
+        AI = np.zeros_like(sh_0)
+        AI[self.mask] = np.sqrt(1 - sh_0[self.mask] /
+                                sh_sum_squared[self.mask])
         return AI
 
+    def norm_of_laplacian_fod(self):
+        """
+        Estimates the squared norm of the Laplacian of the FOD. Similar to
+        the anisotropy index, a higher norm means there are larger higher-order
+        coefficients in the FOD spherical harmonics expansion. This indicates
+        the FOD is more anisotropic overall. This kind of measure was explored
+        in e.g. [1]_.
+
+        References
+        ----------
+        .. [1] Descoteaux, Maxime, et al. "Regularized, fast, and robust
+            analytical Q-ball imaging." Magnetic Resonance in Medicine: An
+            Official Journal of the International Society for Magnetic
+            Resonance in Medicine 58.3 (2007): 497-510.
+        """
+        sh_coef = self.fitted_parameters['sh_coeff']
+        sh_l = sph_harm_ind_list(self.model.sh_order)[1]
+        lb_weights = sh_l ** 2 * (sh_l + 1) ** 2  # laplace-beltrami
+        norm_laplacian = np.sum(sh_coef ** 2 * lb_weights, axis=-1)
+        return norm_laplacian
+
     def predict(self, acquisition_scheme=None, S0=None, mask=None):
         """
         simulates the dMRI signal of the fitted MultiCompartmentModel for the

dmipy/core/modeling_framework.py

@@ -22,8 +22,8 @@
     FittedMultiCompartmentSphericalHarmonicsModel)
 from ..optimizers.brute2fine import (
     GlobalBruteOptimizer, Brute2FineOptimizer)
-from ..optimizers_fod.cvxpy_fod import GeneralPurposeCSDOptimizer
 from ..optimizers_fod.csd_tournier import CsdTournierOptimizer
+from ..optimizers_fod.csd_cvxpy import CsdCvxpyOptimizer
 from ..optimizers.mix import MixOptimizer
 from dipy.utils.optpkg import optional_package
 from graphviz import Digraph
@@ -1605,10 +1605,9 @@ def _check_that_one_anisotropic_kernel_is_present(self):
         if orientation_counter > 1:
             self.multiple_anisotropic_kernels = True
 
-    def fit(self, acquisition_scheme, data, mask=None,
-            solver='tournier07', unity_constraint=True,
-            use_parallel_processing=have_pathos,
-            number_of_processors=None):
+    def fit(self, acquisition_scheme, data, mask=None, solver='csd',
+            lambda_lb=1e-5, unity_constraint='kernel_dependent',
+            use_parallel_processing=have_pathos, number_of_processors=None):
         """ The main data fitting function of a
         MultiCompartmentSphericalHarmonicsModel.
 
@@ -1639,15 +1638,24 @@ def fit(self, acquisition_scheme, data, mask=None,
         mask : (N-1)-dimensional integer/boolean array of size (N_x, N_y, ...),
             Optional mask of voxels to be included in the optimization.
         solver : string,
-            can only be 'cvxpy' at this point for general-purpose csd
-            optimization using cvxpy [1]_.
-        unity_constraint: bool,
-            whether or not to constrain the volume fractions of the FOD to
-            unity.
+            Can be 'csd', 'csd_tounier07' or 'csd_cvxpy', with the default
+            being 'csd'. Using 'csd' will make the algorithm automatically
+            use the 'tournier07' solver [1]_ if there are no volume fractions
+            to fit or they are fixed. Otherwise, the slower but more general
+            cvxpy solver [2]_ is used, which follows the formulation of [3]_.
+        lambda_lb: positive float,
+            Weight for Laplace-Beltrami regularization to impose smoothness
+            into estimated FODs, follows [4]_.
+        unity_constraint: String or bool,
+            Whether or not to constrain the volume fractions of the FOD to
+            unity. The default is set to 'kernel_dependent', meaning it will
+            enforce unity if the kernel is voxel-varying or when volume
+            fractions are estimated. Otherwise unity_constraint is set to
+            False.
         use_parallel_processing : bool,
-            whether or not to use parallel processing using pathos.
+            Whether or not to use parallel processing using pathos.
         number_of_processors : integer,
-            number of processors to use for parallel processing. Defaults to
+            Number of processors to use for parallel processing. Defaults to
             the number of processors in the computer according to cpu_count().
 
         Returns
@@ -1658,13 +1666,35 @@ def fit(self, acquisition_scheme, data, mask=None,
 
         References
         ----------
-        .. [1] Diamond, Steven, and Stephen Boyd. "CVXPY: A Python-embedded
+        .. [1] Tournier, J-Donald, Fernando Calamante, and Alan Connelly.
+            "Robust determination of the fibre orientation distribution in
+            diffusion MRI: non-negativity constrained super-resolved spherical
+            deconvolution." Neuroimage 35.4 (2007): 1459-1472.
+        .. [2] Diamond, Steven, and Stephen Boyd. "CVXPY: A Python-embedded
             modeling language for convex optimization." The Journal of Machine
             Learning Research 17.1 (2016): 2909-2913.
+        .. [3] Jeurissen, Ben, et al. "Multi-tissue constrained spherical
+            deconvolution for improved analysis of multi-shell diffusion MRI
+            data." NeuroImage 103 (2014): 411-426.
+        .. [4] Descoteaux, Maxime, et al. "Regularized, fast, and robust
+            analytical Q-ball imaging." Magnetic Resonance in Medicine: An
+            Official Journal of the International Society for Magnetic
+            Resonance in Medicine 58.3 (2007): 497-510.
         """
 
         self._check_if_kernel_parameters_are_fixed()
 
+        self.voxel_varying_kernel = False
+        if bool(self.x0_parameters):  # if the dictionary is not empty
+            self.voxel_varying_kernel = True
+
+        if unity_constraint == 'kernel_dependent':
+            self.unity_constraint = False
+            if not self.volume_fractions_fixed or self.voxel_varying_kernel:
+                self.unity_constraint = True
+        else:
+            self.unity_constraint = unity_constraint
+
         # estimate S0
         self.scheme = acquisition_scheme
         data_ = np.atleast_2d(data)
@@ -1694,25 +1724,48 @@ def fit(self, acquisition_scheme, data, mask=None,
             data_, x0_)
 
         start = time()
-        if solver == 'tournier07':
+        if solver == 'csd':
+            if self.volume_fractions_fixed:
+                fit_func = CsdTournierOptimizer(
+                    acquisition_scheme, self, x0_, self.sh_order,
+                    unity_constraint=self.unity_constraint,
+                    lambda_lb=lambda_lb)
+                if use_parallel_processing:
+                    msg = 'Parallel processing turned off for tournier07'
+                    msg += ' optimizer because it does not improve fitting '
+                    msg += 'speed.'
+                    print(msg)
+                    use_parallel_processing = False
+                print('Setup Tournier07 FOD optimizer in {} seconds'.format(
+                    time() - start))
+            else:
+                fit_func = CsdCvxpyOptimizer(
+                    acquisition_scheme, self, x0_, self.sh_order,
+                    unity_constraint=self.unity_constraint,
+                    lambda_lb=lambda_lb)
+                print('Setup CVXPY FOD optimizer in {} seconds'.format(
+                    time() - start))
+        elif solver == 'csd_tournier07':
             fit_func = CsdTournierOptimizer(
                 acquisition_scheme, self, x0_, self.sh_order,
-                unity_constraint=unity_constraint)
+                unity_constraint=self.unity_constraint, lambda_lb=lambda_lb)
             if use_parallel_processing:
                 msg = 'Parallel processing turned off for tournier07 optimizer'
                 msg += ' because it does not improve fitting speed.'
                 print(msg)
                 use_parallel_processing = False
             print('Setup Tournier07 FOD optimizer in {} seconds'.format(
                 time() - start))
-        elif solver == 'cvxpy':
-            fit_func = GeneralPurposeCSDOptimizer(
-                acquisition_scheme, self, x0_, self.sh_order, unity_constraint)
+        elif solver == 'csd_cvxpy':
+            fit_func = CsdCvxpyOptimizer(
+                acquisition_scheme, self, x0_, self.sh_order,
+                unity_constraint=self.unity_constraint, lambda_lb=lambda_lb)
+            print('Setup CVXPY FOD optimizer in {} seconds'.format(
+                time() - start))
         else:
             msg = "Unknown solver name {}".format(solver)
             raise ValueError(msg)
-            print('Setup CVXPY FOD optimizer in {} seconds'.format(
-                time() - start))
+
         self.optimizer = fit_func
 
         if use_parallel_processing and not have_pathos:

dmipy/core/tests/test_csd_equivalence_parametric_distributions.py

@@ -5,7 +5,7 @@
 from dipy.data import get_sphere
 import numpy as np
 from numpy.testing import (
-    assert_array_almost_equal, assert_almost_equal, assert_raises)
+    assert_array_almost_equal, assert_almost_equal, assert_raises, assert_)
 from dipy.utils.optpkg import optional_package
 cvxpy, have_cvxpy, _ = optional_package("cvxpy")
 
@@ -29,7 +29,7 @@ def test_equivalence_csd_and_parametric_fod_tournier07(
         [stick])
     sh_mod.set_fixed_parameter('C1Stick_1_lambda_par', lambda_par)
 
-    sh_fit = sh_mod.fit(scheme, data, solver='tournier07')
+    sh_fit = sh_mod.fit(scheme, data, solver='csd_tournier07')
     fod = sh_fit.fod(sphere.vertices)
 
     watson = distributions.SD1Watson(mu=[0., 0.], odi=0.15)
@@ -57,7 +57,7 @@ def test_equivalence_csd_and_parametric_fod(
         [stick])
     sh_mod.set_fixed_parameter('C1Stick_1_lambda_par', lambda_par)
 
-    sh_fit = sh_mod.fit(scheme, data, solver='cvxpy')
+    sh_fit = sh_mod.fit(scheme, data, solver='csd_cvxpy', lambda_lb=0.)
     fod = sh_fit.fod(sphere.vertices)
 
     watson = distributions.SD1Watson(mu=[0., 0.], odi=0.15)
@@ -92,7 +92,7 @@ def test_multi_compartment_fod_with_parametric_model(
         partial_volume_1=1 - vf_intra)
     data = mc_mod.simulate_signal(scheme, simulation_parameters)
 
-    sh_fit = sh_mod.fit(scheme, data, solver='cvxpy')
+    sh_fit = sh_mod.fit(scheme, data, solver='csd_cvxpy', lambda_lb=0.)
 
     vf_intra_estimated = sh_fit.fitted_parameters['partial_volume_0']
     assert_almost_equal(vf_intra, vf_intra_estimated)
@@ -102,6 +102,94 @@ def test_multi_compartment_fod_with_parametric_model(
     assert_array_almost_equal(data, predicted_signal[0], 4)
 
 
+def test_multi_voxel_parametric_to_sm_to_sh_fod_watson():
+    stick = cylinder_models.C1Stick()
+    zeppelin = gaussian_models.G2Zeppelin()
+    watsonstick = distribute_models.SD1WatsonDistributed(
+        [stick, zeppelin])
+
+    watsonstick.set_equal_parameter(
+        'G2Zeppelin_1_lambda_par', 'C1Stick_1_lambda_par')
+    watsonstick.set_tortuous_parameter('G2Zeppelin_1_lambda_perp',
+                                       'G2Zeppelin_1_lambda_par',
+                                       'partial_volume_0')
+    mc_mod = modeling_framework.MultiCompartmentModel([watsonstick])
+
+    parameter_dict = {
+        'SD1WatsonDistributed_1_SD1Watson_1_mu': np.random.rand(10, 2),
+        'SD1WatsonDistributed_1_partial_volume_0': np.linspace(0.1, 0.9, 10),
+        'SD1WatsonDistributed_1_G2Zeppelin_1_lambda_par':
+        np.linspace(1.5, 2.5, 10) * 1e-9,
+        'SD1WatsonDistributed_1_SD1Watson_1_odi': np.linspace(0.3, 0.7, 10)
+    }
+
+    data = mc_mod.simulate_signal(scheme, parameter_dict)
+
+    sm_mod = modeling_framework.MultiCompartmentSphericalMeanModel(
+        [stick, zeppelin])
+    sm_mod.set_equal_parameter(
+        'G2Zeppelin_1_lambda_par', 'C1Stick_1_lambda_par')
+    sm_mod.set_tortuous_parameter(
+        'G2Zeppelin_1_lambda_perp', 'G2Zeppelin_1_lambda_par',
+        'partial_volume_0', 'partial_volume_1')
+
+    sf_watson = []
+    for mu, odi in zip(
+            parameter_dict['SD1WatsonDistributed_1_SD1Watson_1_mu'],
+            parameter_dict['SD1WatsonDistributed_1_SD1Watson_1_odi']):
+        watson = distributions.SD1Watson(mu=mu, odi=odi)
+        sf_watson.append(watson(sphere.vertices))
+    sf_watson = np.array(sf_watson)
+
+    sm_fit = sm_mod.fit(scheme, data)
+    sh_mod = sm_fit.return_spherical_harmonics_fod_model()
+
+    sh_fit_auto = sh_mod.fit(scheme, data)  # will pick tournier
+    fod_tournier = sh_fit_auto.fod(sphere.vertices)
+    assert_array_almost_equal(fod_tournier, sf_watson, 1)
+
+    sh_fit_tournier = sh_mod.fit(
+        scheme, data, solver='csd_tournier07', unity_constraint=False)
+    fod_tournier = sh_fit_tournier.fod(sphere.vertices)
+    assert_array_almost_equal(fod_tournier, sf_watson, 1)
+
+    sh_fit_cvxpy = sh_mod.fit(
+        scheme, data, solver='csd_cvxpy', unity_constraint=True, lambda_lb=0.)
+    fod_cvxpy = sh_fit_cvxpy.fod(sphere.vertices)
+    assert_array_almost_equal(fod_cvxpy, sf_watson, 2)
+
+    sh_fit_cvxpy = sh_mod.fit(
+        scheme, data, solver='csd_cvxpy', unity_constraint=False, lambda_lb=0.)
+    fod_cvxpy = sh_fit_cvxpy.fod(sphere.vertices)
+    assert_array_almost_equal(fod_cvxpy, sf_watson, 2)
+
+
+@np.testing.dec.skipif(not have_cvxpy)
+def test_laplacian_and_AI_with_regularization(
+        odi=0.15, mu=[0., 0.], lambda_par=1.7e-9):
+    stick = cylinder_models.C1Stick()
+    watsonstick = distribute_models.SD1WatsonDistributed(
+        [stick])
+    params = {'SD1Watson_1_odi': odi,
+              'SD1Watson_1_mu': mu,
+              'C1Stick_1_lambda_par': lambda_par}
+    data = watsonstick(scheme, **params)
+
+    sh_mod = modeling_framework.MultiCompartmentSphericalHarmonicsModel(
+        [stick])
+    sh_mod.set_fixed_parameter('C1Stick_1_lambda_par', lambda_par)
+
+    for solver in ['csd_tournier07', 'csd_cvxpy']:
+        sh_fit = sh_mod.fit(scheme, data, solver=solver, lambda_lb=0.)
+        sh_fit_reg = sh_mod.fit(scheme, data, solver=solver, lambda_lb=1e-3)
+        ai = sh_fit.anisotropy_index()
+        lb = sh_fit.norm_of_laplacian_fod()
+        ai_reg = sh_fit_reg.anisotropy_index()
+        lb_reg = sh_fit_reg.norm_of_laplacian_fod()
+        assert_(ai > ai_reg)
+        assert_(lb > lb_reg)
+
+
 @np.testing.dec.skipif(not have_cvxpy)
 def test_spherical_harmonics_model_raises(
         odi=0.15, mu=[0., 0.], lambda_par=1.7e-9):
@@ -122,5 +210,11 @@ def test_spherical_harmonics_model_raises(
 
     sh_mod = modeling_framework.MultiCompartmentSphericalHarmonicsModel(
         [stick])
+    assert_raises(ValueError, sh_mod.fit, scheme, data, solver='csd_cvxpy')
 
-    assert_raises(ValueError, sh_mod.fit, scheme, data, solver='cvxpy')
+    sh_mod = modeling_framework.MultiCompartmentSphericalHarmonicsModel(
+        [stick, ball])
+    sh_mod.set_fixed_parameter('C1Stick_1_lambda_par', lambda_par)
+    sh_mod.set_fixed_parameter('G1Ball_1_lambda_iso', 3e-9)
+    assert_raises(
+        ValueError, sh_mod.fit, scheme, data, solver='csd_tournier07')