From 2234d0651c99f9379c15417556ce27853e4e19e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jon=20Haitz=20Legarreta=20Gorro=C3=B1o?= Date: Sun, 7 Apr 2024 10:36:04 -0400 Subject: [PATCH] ENH: Implement Gaussian Process Implement Gaussian Process. --- pyproject.toml | 1 + src/eddymotion/estimator.py | 1 + src/eddymotion/model/__init__.py | 2 + src/eddymotion/model/base.py | 105 +++++ src/eddymotion/model/tests/__init__.py | 0 src/eddymotion/model/tests/test_utils.py | 170 ++++++++ src/eddymotion/model/utils.py | 481 +++++++++++++++++++++++ test/test_model.py | 41 ++ 8 files changed, 801 insertions(+) create mode 100644 src/eddymotion/model/tests/__init__.py create mode 100644 src/eddymotion/model/tests/test_utils.py create mode 100644 src/eddymotion/model/utils.py diff --git a/pyproject.toml b/pyproject.toml index 6faf5185..8d876958 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "numpy>=1.17.3", "nest-asyncio>=1.5.1", "scikit-image>=0.14.2", + "scikit_learn", "scipy>=1.8.0", ] dynamic = ["version"] diff --git a/src/eddymotion/estimator.py b/src/eddymotion/estimator.py index b11bab48..ee4a6316 100644 --- a/src/eddymotion/estimator.py +++ b/src/eddymotion/estimator.py @@ -110,6 +110,7 @@ def estimate( "avg", "average", "mean", + "gp", ) or model.lower().startswith("full") dwmodel = None diff --git a/src/eddymotion/model/__init__.py b/src/eddymotion/model/__init__.py index 3c44e2ad..b64712ff 100644 --- a/src/eddymotion/model/__init__.py +++ b/src/eddymotion/model/__init__.py @@ -26,6 +26,7 @@ AverageDWModel, DKIModel, DTIModel, + GaussianProcessModel, ModelFactory, PETModel, TrivialB0Model, @@ -36,6 +37,7 @@ "AverageDWModel", "DKIModel", "DTIModel", + "GaussianProcessModel", "TrivialB0Model", "PETModel", ) diff --git a/src/eddymotion/model/base.py b/src/eddymotion/model/base.py index 38207741..d06a2c0d 100644 --- a/src/eddymotion/model/base.py +++ b/src/eddymotion/model/base.py @@ -25,6 +25,7 @@ import warnings import numpy as np +import sklearn from dipy.core.gradients import gradient_table from joblib import Parallel, delayed @@ -127,6 +128,8 @@ def __init__(self, gtab, S0=None, mask=None, b_max=None, **kwargs): if not model_str: raise TypeError("No model defined") + # ToDo + # Use lazy loading ? from importlib import import_module module_name, class_name = model_str.rsplit(".", 1) @@ -409,6 +412,108 @@ class DKIModel(BaseModel): _model_class = "dipy.reconst.dki.DiffusionKurtosisModel" +class GaussianProcessModel: + """A Gaussian Process model based on [Andersson16a]_ (fig 1). + DWIs need to be transformed to a single ref space (fig 2 [Andersson16b]_ ?) + + Definitions: + s: reference/undistorted space: used to denote the space or any image in + that space + f: observed/distorted image: used to denote any image in acquisition space + a: acquisition parameters: PE‐direction and bandwidth in PE‐direction + r: rigid body (subject movement) parameters + \beta: Eddy current parameters + e(\beta): Eddy current‐induced off resonance field (Hz) + h: Susceptibility induced off‐resonance field (Hz) + + (fig 1) and algorithm: + 1. Input: N DWI volumes f_{i} with acq parameters a_{i}; susceptibility + field h + 2. Initialize: set all beta_{i} and r_i{i} = 0 + 3. Compute for M iterations + - Load GP prediction maker + - For all i in N (DWIs) + - Compute \\hat{s}_{i} (f_{i}, h, \beta_{i}, r_{i}, a_{i}) eqs 2 and 4 + - Load \\hat{s}_{i} (f_{i}, h, \beta_{i}, r_{i}, a_{i}) as training + data for GP + - Estimate hyperparameters for the GP used to predict the signal shape + for every voxel + - Update EC and movement parameters + - For all i in N (DWIs) + - Draw a prediction s_{i} from the GP + - Compute \\hat{f}_{i} (s_{i}, h, \beta_{i}, r_{i}, a_{i}) + - Use \\hat{f}_{i} - f_{i} to update \beta_{i} and r_{i} (eq 6) + + a: direction of the PE and the total readout time (here defined as the time + between the acquisition of the center of the first and last echoes). + Internally, a is divided into a = [p t] where p is a unity length 1 x 3 + vector defining the PE direction (such that for example [1 0 0], [−1 0 0], + [0 1 0] and [0 −1 0] denote R → L, L → R, P → A and A P PE direction + respectively) and where t denotes the readout time (in seconds). + r: 1x6 vector: 3 translations, 3 rotations + \beta: four for linear; ten for quadratic, twenty for cubic. + h: assumed to be in the same space as the first b = 0 image supplied to + eddy, which will be automatically fulfilled if it was estimated by topup + and that same b = 0 image was the first of those supplied to topup. Hence, + it can be said to help define the reference/undistorted space as the first + b = 0 image after distortion correction by h. + + See Appendix A for further details. + + Add the outlier detection part in [Andersson16b]? + + References + ---------- + .. [Andersson16a] J. L. R. Andersson. et al., An integrated approach to + correction for off-resonance effects and subject movement in diffusion MR + imaging, NeuroImage 125 (2016) 1063–1078 + .. [Andersson16b] J. L. R. Andersson. et al., Incorporating outlier + detection and replacement into a non-parametric framework for movement and + distortion correction of diffusion MR images, NeuroImage 141 (2016) 556–572 + """ + + __slots__ = ("_dwi", "_a", "_h", "_kernel", "_num_iterations", "_betas", "_r", "_gpr") + + def __init__(self, dwi, a, h, kernel, num_iterations=5, **kwargs): + """Implement object initialization.""" + + # ToDo + # This should be he HDF5 file dwi object, so that we avoid having the + # entire 4D volume in memory + self._dwi = dwi + self._a = a + self._h = h + self._num_iterations = num_iterations + + # Initialize + self._betas = 0 + self._r = 0 + + # ToDo + # Build the GP kernel here or in fit ? + # self._gpr = None + # Does the kernel depend on which data we use as the training data (i.e. + # varies with the index we choose to predict)? + self._kernel = kernel + + def fit(self, *args, **kwargs): + """The x are our gradient directions; the observations are our diffusion + volumes. + X_train: array-like of shape (n_samples, n_features), n_samples being + the number of gradients, and the n_features the number of shells ? + y_train_array-like of shape (n_samples,) or (n_samples, n_targets)""" + + self._gpr = sklearn.gaussian_process.GaussianProcessRegressor(kernel=self._kernel) + self._gpr.fit(X_train, y_train) + + def predict(self, gradient, **kwargs): + """Return the Gaussian Process prediction according to [Andersson16]_""" + # ToDo + # Call self._gprlog_marginal_likelihood for eq. 12 in Andersson 15 ? + y_mean, y_std = self._gpr.predict(X, return_std=True) + return y_mean, y_std + + def _rasb2dipy(gradient): gradient = np.asanyarray(gradient) if gradient.ndim == 1: diff --git a/src/eddymotion/model/tests/__init__.py b/src/eddymotion/model/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/model/tests/test_utils.py b/src/eddymotion/model/tests/test_utils.py new file mode 100644 index 00000000..0fd3dcdc --- /dev/null +++ b/src/eddymotion/model/tests/test_utils.py @@ -0,0 +1,170 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2022 The NiPreps Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# We support and encourage derived works from this project, please read +# about our expectations at +# +# https://www.nipreps.org/community/licensing/ +# +import nibabel as nib +import numpy as np +from eddymotion.model.utils import ( + extract_dmri_shell, + find_shelling_scheme, + is_positive_definite, + # update_covariance1, + # update_covariance2, +) + + +def test_is_positive_definite(): + + matrix = np.array([[4, 1, 2], [1, 3, 1], [2, 1, 5]]) + assert is_positive_definite(matrix) + + matrix = np.array([[4, 1, 2], [1, -3, 1], [2, 1, 5]]) + assert not is_positive_definite(matrix) + + +def test_update_covariance(): + + _K = np.random.rand(5, 5) + _thpar = [0.5, 1.0, 2.0] + update_covariance1(_K, _thpar) + print(_K) # Updated covariance matrix + + +def test_extract_dmri_shell(): + + # dMRI volume with 5 gradients + bvals = np.asarray([0, 1980, 12, 990, 2000]) + bval_count = len(bvals) + vols_size = (10, 15, 20) + dwi = np.ones((*vols_size, bval_count)) + bvecs = np.ones((bval_count, 3)) + # Set all i-th gradient dMRI volume data and bvecs values to i + for i in range(bval_count): + dwi[..., i] = i + bvecs[i, :] = i + dwi_img = nib.Nifti1Image(dwi, affine=np.eye(4)) + + bvals_to_extract = [0, 2000] + tol = 15 + + expected_indices = np.asarray([0, 2, 4]) + expected_shell_data = np.stack([i*np.ones(vols_size) for i in expected_indices], axis=-1) + expected_shell_bvals = np.asarray([0, 12, 2000]) + expected_shell_bvecs = np.asarray([[i]*3 for i in expected_indices]) + + ( + obtained_indices, + obtained_shell_data, + obtained_shell_bvals, + obtained_shell_bvecs + ) = extract_dmri_shell( + dwi_img, bvals, bvecs, bvals_to_extract=bvals_to_extract, tol=tol) + + assert np.array_equal(obtained_indices, expected_indices) + assert np.array_equal(obtained_shell_data, expected_shell_data) + assert np.array_equal(obtained_shell_bvals, expected_shell_bvals) + assert np.array_equal(obtained_shell_bvecs, expected_shell_bvecs) + + bvals = np.asarray([0, 1010, 12, 990, 2000]) + bval_count = len(bvals) + vols_size = (10, 15, 20) + dwi = np.ones((*vols_size, bval_count)) + bvecs = np.ones((bval_count, 3)) + # Set all i-th gradient dMRI volume data and bvecs values to i + for i in range(bval_count): + dwi[..., i] = i + bvecs[i, :] = i + dwi_img = nib.Nifti1Image(dwi, affine=np.eye(4)) + + bvals_to_extract = [0, 1000] + tol = 20 + + expected_indices = np.asarray([0, 1, 2, 3]) + expected_shell_data = np.stack([i*np.ones(vols_size) for i in expected_indices], axis=-1) + expected_shell_bvals = np.asarray([0, 1010, 12, 990]) + expected_shell_bvecs = np.asarray([[i]*3 for i in expected_indices]) + + ( + obtained_indices, + obtained_shell_data, + obtained_shell_bvals, + obtained_shell_bvecs + ) = extract_dmri_shell( + dwi_img, bvals, bvecs, bvals_to_extract=bvals_to_extract, tol=tol) + + assert np.array_equal(obtained_indices, expected_indices) + assert np.array_equal(obtained_shell_data, expected_shell_data) + assert np.array_equal(obtained_shell_bvals, expected_shell_bvals) + assert np.array_equal(obtained_shell_bvecs, expected_shell_bvecs) + + +def test_find_shelling_scheme(): + + tol = 20 + bvals = np.asarray([0, 0]) + expected_shells = np.asarray([0]) + expected_bval_centroids = np.asarray([0, 0]) + obtained_shells, obtained_bval_centroids = find_shelling_scheme( + bvals, tol=tol) + + assert np.array_equal(obtained_shells, expected_shells) + assert np.array_equal(obtained_bval_centroids, expected_bval_centroids) + + bvals = np.asarray([ + 5, 300, 300, 300, 300, 300, 305, 1005, 995, 1000, 1000, 1005, 1000, + 1000, 1005, 995, 1000, 1005, 5, 995, 1000, 1000, 995, 1005, 995, 1000, + 995, 995, 2005, 2000, 2005, 2005, 1995, 2000, 2005, 2000, 1995, 2005, 5, + 1995, 2005, 1995, 1995, 2005, 2005, 1995, 2000, 2000, 2000, 1995, 2000, 2000, + 2005, 2005, 1995, 2005, 2005, 1990, 1995, 1995, 1995, 2005, 2000, 1990, 2010, 5 + ]) + expected_shells = np.asarray([5., 300.83333333, 999.5, 2000.]) + expected_bval_centroids = ([ + 5., 300.83333333, 300.83333333, 300.83333333, 300.83333333, 300.83333333, 300.83333333, 999.5, 999.5, 999.5, 999.5, 999.5, 999.5, + 999.5, 999.5, 999.5, 999.5, 999.5, 5., 999.5, 999.5, 999.5, 999.5, 999.5, 999.5, 999.5, + 999.5, 999.5, 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 5., + 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., + 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 2000., 5. + ]) + obtained_shells, obtained_bval_centroids = find_shelling_scheme( + bvals, tol=tol) + + # ToDo + # Giving a tolerance of 15 this fails because it finds 5 clusters + assert np.allclose(obtained_shells, expected_shells) + assert np.allclose(obtained_bval_centroids, expected_bval_centroids) + + bvals = np.asarray([0, 1980, 12, 990, 2000]) + expected_shells = np.asarray([6, 990, 1980, 2000]) + expected_bval_centroids = np.asarray([6, 1980, 6, 990, 2000]) + obtained_shells, obtained_bval_centroids = find_shelling_scheme( + bvals, tol=tol) + + assert np.allclose(obtained_shells, expected_shells) + assert np.allclose(obtained_bval_centroids, expected_bval_centroids) + + bvals = np.asarray([0, 1010, 12, 990, 2000]) + tol = 60 + expected_shells = np.asarray([6, 1000, 2000]) + expected_bval_centroids = np.asarray([6, 1000, 6, 1000, 2000]) + obtained_shells, obtained_bval_centroids = find_shelling_scheme(bvals, tol) + + assert np.allclose(obtained_shells, expected_shells) + assert np.allclose(obtained_bval_centroids, expected_bval_centroids) diff --git a/src/eddymotion/model/utils.py b/src/eddymotion/model/utils.py new file mode 100644 index 00000000..68137659 --- /dev/null +++ b/src/eddymotion/model/utils.py @@ -0,0 +1,481 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# Copyright 2022 The NiPreps Developers +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY kIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# We support and encourage derived works from this project, please read +# about our expectations at +# +# https://www.nipreps.org/community/licensing/ +# +import numpy as np +from dipy.core.gradients import get_bval_indices +from sklearn.cluster import KMeans + +B0_THRESHOLD = 50 # from dmriprep +SHELL_DIFF_THRES = 20 # 150 in dmriprep + + +def is_positive_definite(matrix): + """Check whether the given matrix is positive definite. Any positive + definite matrix can be decomposed as the product of a lower triangular + matrix and its conjugate transpose by performing the Cholesky decomposition. + + Parameters + ---------- + matrix : np.ndarray + The matrix to check. + + Returns + ------- + True is the matrix is positive definite; False otherwise + """ + + try: + # Attempt Cholesky decomposition + np.linalg.cholesky(matrix) + return True + except np.linalg.LinAlgError: + # Matrix is not positive definite + return False + + +def variance(data, indx): + m = np.mean(data[indx]) + v = np.sum(np.square(data[indx] - m)) + return v / (len(indx) - 1) + + +def compute_hyperparams_spherical_matrix(data): + """Used for single-shell data. SphericalKMatrix::GetHyperParGuess""" + + num_groups = len(grps) + + # Calculate group-wise variances (across directions) averaged over voxels + vars = np.zeros(num_groups) + for i in range(num_groups): + for j in range(len(data)): + vars[i] += variance(data[j], grps[i]) + vars[i] /= len(data) + + # Make (semi-educated) guesses for hyperparameters + hpar = [0.0] * n_par(num_groups) + sn_wgt = 0.25 + cntr = 0 + + # ToDo + # Vectorize this + for cntr in range(100): + for j in range(num_groups): + for i in range(j, num_groups): + pi = ij_to_parameter_index(i, j, num_groups) + if i == j: + hpar[pi] = np.log(vars[i] / 2.7) + hpar[pi + 1] = 1.5 + hpar[pi + 2] = np.log(sn_wgt * vars[i]) + else: + hpar[pi] = np.log(0.8 * min(vars[i], vars[j]) / 2.7) + hpar[pi + 1] = 1.5 + + if not valid_hpars(hpar): + sn_wgt *= 1.1 + else: + break + + if cntr == 99: + raise ValueError("Unable to find valid hyperparameters") + + return hpar + + +def compute_hyperparams_newspherical_matrix(data): + """Used for multi-shell data. NewSphericalKMatrix::GetHyperParGuess""" + + num_groups = len(grps) + vars = np.zeros(num_groups) + var = 0.0 + + # ToDo + # Vectorize this + for i in range(num_groups): + for j in range(len(data)): + vars[i] += variance(data[j], grps[i]) + vars[i] /= len(data) + var += vars[i] + + var /= num_groups + + # Make (semi-educated) guesses for hyperparameters based on findings in GP + # paper + hpar = [0.0] * n_par(num_groups) + hpar[0] = 0.9 * np.log(var / 3.0) + hpar[1] = 0.45 + + if num_groups == 1: + hpar[2] = np.log(var / 3.0) + else: + hpar[2] = 0.0 + delta = 0.2 / (num_groups - 1) + # ToDo + # Vectorize this + for i in range(num_groups): + hpar[3 + i] = (1.1 - i * delta) * np.log(var / 3.0) + + if not valid_hpars(hpar): + raise ValueError("Unable to find valid hyperparameters") + + return hpar + + +def calculate_angle_matrix(bvecs): + """Matrix of angles between gradients. + FSL crew takes as input dpars, which are diffusion parameters""" + + # ToDo + # Does this deal with multi-shell? What does the matrix look like cols vs rows? + # Should the b0 be included? + # What if whe have different bvals ? + # Note that the double loops assume bvecs is a 2D array + + n = len(bvecs) + angle_mat = np.zeros((n, n)) + + for j in range(n): + for i in range(j, n): + angle_mat[i, j] = np.arccos(min(1.0, abs(np.dot(bvecs[i], bvecs[j])))) + + return angle_mat + + +def compute_exponential_function(th, a): + """Compute the exponential function according to eq. 9 in [Andersson15]_. + + .. math:: + + C(\theta) = \exp(- \frac{\theta}{a}) + + Parameters + ---------- + th : np.ndarray + . + a : float > 0 + Positive scale parameter that here determines the "distance" at which θ + the covariance goes to zero. + + Returns + ------- + The spherical function. + + References + ---------- + .. [Andersson15] J. L. R. Andersson. et al., Non-parametric representation + and prediction of single- and multi-shell diffusion-weighted MRI data using + Gaussian processes, NeuroImage 122 (2015) 166–176 + """ + + assert a > 0 + return np.exp(-th / a) + + +def compute_spherical_function(theta, a): + """Compute the spherical function according to eq. 10 in [Andersson15]_. + + .. math:: + + C(\theta) = \begin{cases} + 1 - \frac{3 \theta}{2 a} + \frac{\theta^3}{2 a^3} & \textnormal{if} \; \theta \leq a \\ + 0 & \textnormal{if} \; \theta > a + \end{cases} + + Parameters + ---------- + theta : np.ndarray + . + a : float > 0 + Positive scale parameter that here determines the "distance" at which θ + the covariance goes to zero. + + Returns + ------- + The spherical function. + + References + ---------- + .. [Andersson15] J. L. R. Andersson. et al., Non-parametric representation + and prediction of single- and multi-shell diffusion-weighted MRI data using + Gaussian processes, NeuroImage 122 (2015) 166–176 + """ + + assert a > 0 + return 1.0 - 1.5 * theta / a + 0.5 * (theta ** 3) / (a ** 3) + + +def compute_squared_exponential_function(grpi, grpb, l): + """Compute the squared exponential smooth function describing how the + covariance changes along the b direction. + + It uses the log of the b-values as the measure of distance along the + b-direction according to eq. 15 in [Andersson15]_. + + .. math:: + + C_{b}(b, b'; \ell) = \exp\left( - \frac{(\log b - \log b')^2}{2 \ell^2} \right) + + Parameters + ---------- + grpi : np.ndarray + Group of indices. + grpb : np.ndarray + Groups of b-values. + l : float + + Returns + ------- + The squared exponential function. + + References + ---------- + .. [Andersson15] J. L. R. Andersson. et al., Non-parametric representation + and prediction of single- and multi-shell diffusion-weighted MRI data using + Gaussian processes, NeuroImage 122 (2015) 166–176 + """ + + # Compute log probability of b-values + log_grpb = np.log(grpb) + bv_diff = log_grpb[grpi[:, None]] - log_grpb[grpi] + return np.exp(-(bv_diff ** 2) / (2 * l ** 2)) + +# ToDo +# Long-term: generalize this so that the loop block can be put into a function +# and be called from the single-shell and multi-shell +def compute_single_shell_covariance_matrix(k, angle_mat, grpi, ngrp, thpar): + """Compute single-shell covariance. + SphericalKMatrix::calculate_K_matrix""" + + if k.shape[0] != angle_mat.shape[0]: + k = np.zeros_like(angle_mat) + + # Compute angular covariance + # ToDo + # Vectorize this + for j in range(k.shape[1]): + for i in range(j, k.shape[0]): + pindx = ij_to_parameter_index(grpi[i], grpi[j], ngrp) + sm = thpar[pindx] + a = thpar[pindx + 1] + theta = angle_mat[i, j] + + # eq. 10 in [Andersson15]_ + if a > theta: + k[i+1, j+1] = sm * (1.0 - 1.5 * theta / a + 0.5 * (theta ** 3) / (a ** 3)) + else: + k[i+1, j+1] = 0.0 + if i == j: + k[i+1, j+1] += thpar[pindx + 2] + + +def compute_multi_shell_covariance_matrix1(k, thpar, grpb, grpi): + """Compute multi-shell covariance. + Indicies2KMatrix::common_construction + hpar are hyperparameters; thpar are "transformed" (for example + exponentiated) hyperparameters + """ + + sm = thpar[0] + a = thpar[1] + l = thpar[2] + + # Make angle matrix (eq. 11 in [Andersson15]_) + # ToDo + # Vectorize this + # _k.resize(nval, nval); missing + for j in range(nval): + for i in range(j, nval): + if i == j: + k[i, j] = 0.0 # set diagonal elements to 0 + else: + k[i, j] = np.arccos(min(1.0, abs(np.dot(bvecs[i], bvecs[j])))) + k[j, i] = k[i, j] # make it symmetric + + # Compute angular covariance + # ToDo + # Vectorize this + for j in range(k.shape[1]): + for i in range(j, k.shape[0]): + theta = k[i+1, j+1] + # eq. 10 in [Andersson15]_ + if a > theta: + k[i+1, j+1] = sm * (1.0 - 1.5 * theta / a + 0.5 * (theta ** 3) / (a ** 3)) + else: + k[i+1, j+1] = 0.0 + if i != j: + k[j+1, i+1] = k[i+1, j+1] # make it symmetric + + # Compute b-value covariance + # ToDo + # Vectorize this + if len(grpb) > 1: + log_grpb = np.log(grpb) + # Takes upper triangular elements + for j in range(k.shape[1]): + for i in range(j + 1, k.shape[0]): + if k[i+1, j+1] != 0.0: + bvdiff = log_grpb[grpi[i]] - log_grpb[grpi[j]] + if bvdiff: + # eq. 15 in [Andersson15]_ + k[i+1, j+1] *= np.exp(-(bvdiff ** 2) / (2 * l ** 2)) + k[j+1, i] = k[i+1, j+1] # make it symmetric + + +def compute_multi_shell_covariance_matrix2(k, angle_mat, thpar, grpb, grpi): + """Compute multi-shell covariance. + NewSphericalKMatrix::calculate_K_matrix + """ + + sm = thpar[0] + a = thpar[1] + l = thpar[2] + + # Compute angular covariance + # ToDo + # Vectorize this + for j in range(K.shape[1]): + for i in range(j, K.shape[0]): + theta = angle_mat[i+1, j+1] + if a > theta: + K[i+1, j+1] = sm * (1.0 - 1.5 * theta / a + 0.5 * (theta ** 3) / (a ** 3)) + else: + K[i+1, j+1] = 0.0 + + # Compute b-value covariance + # ToDo + # Vectorize this + if ngrp > 1: + log_grpb = np.log(grpb()) + for j in range(K.shape[1]): + for i in range(j + 1, K.shape[0]): + bvdiff = log_grpb[grpi[i]] - log_grpb[grpi[j]] + if bvdiff: + K[i+1, j+1] *= np.exp(-(bvdiff ** 2) / (2 * l ** 2)) + + +# ToDo +# Naming: DWI vs dMRI +def extract_dmri_shell(dwi, bvals, bvecs, bvals_to_extract, tol=SHELL_DIFF_THRES): + """Extract the DWI volumes that are on the given b-value shells. Multiple + shells can be extracted at once by specifying multiple b-values. The + extracted volumes will be in the same order as in the original file. + + Parameters + ---------- + dwi : nib.Nifti1Image + Original dMRI multi-shell volume. + bvals : ndarray + b-values in FSL format. + bvecs : ndarray + b-vectors in FSL format. + bvals_to_extract : list of int + List of b-values to extract. + tol : int, optional + Tolerance between the b-values to extract and the actual b-values. + + Returns + ------- + indices : ndarray + Indices of the volumes corresponding to the given ``bvals``. + shell_data : ndarray + Volumes corresponding to the given ``bvals``. + output_bvals : ndarray + Selected b-values (as extracted from ``bvals``). + output_bvecs : ndarray + Selected b-vectors. + """ + + indices = [ + get_bval_indices(bvals, shell, tol=tol) for shell in bvals_to_extract + ] + indices = np.unique(np.sort(np.hstack(indices))) + + if len(indices) == 0: + raise ValueError( + f"No dMRI volumes found corresponding to the given b-values: {bvals_to_extract}" + ) + + shell_data = dwi.get_fdata()[..., indices] + output_bvals = bvals[indices].astype(int) + output_bvecs = bvecs[indices, :] + + return indices, shell_data, output_bvals, output_bvecs + + +# ToDo +# Long term: use DiffusionGradientTable from dmriprep, normalizing gradients, +# etc. ? +def find_shelling_scheme(bvals, tol=SHELL_DIFF_THRES): + """Find the shelling scheme on the given b-values: extract the b-value + shells as the b-values centroids using k-means clustering. + + Parameters + ---------- + bvals : ndarray + b-values in FSL format. + tol : int, optional + Tolerance between the b-values and the centroids in the average squared + distance sense. + + Returns + ------- + shells : ndarray + b-value shells. + bval_centroids : ndarray + Shell value corresponding to each value in ``bvals``. + """ + + # Use kmeans to find the shelling scheme + for k in range(1, len(np.unique(bvals)) + 1): + kmeans_res = KMeans(n_clusters=k).fit(bvals.reshape(-1, 1)) + # ToDo + # The tolerance is not a very intuitive value, as it has to do with the + # sum of squared distances across all samples to the centroids + # (_inertia) + # Alternatives: + # - We could accept the number of clusters as a parameter and do + # kmeans_res = KMeans(n_clusters=n_clusters) + # Setting that to 3 in the last testing case, where tol = 60 is not + # intuitive would give the expected 6, 1000, 2000 clusters. + # Passes all tests. But maybe not tested corner cases + # We could have both k and tol as optional parameters, set to None by + # default to force the user set one + # - Use get_bval_indices to get the cluster centroids and then + # substitute the values in bvals with the corresponding values + # indices = [get_bval_indices(bvals, shell, tol=tol) for shell in bvals_to_extract] + # result = np.zeros_like(bvals) + # for i, idx in enumerate(indices): + # result[idx] = bvals_to_extract[i] + + if kmeans_res.inertia_ / len(bvals) < tol: + break + else: + raise ValueError( + f"bvals parsing failed: no shells found more than {tol} apart" + ) + + # Convert the kclust labels to an array + shells = kmeans_res.cluster_centers_ + bval_centroids = np.zeros(bvals.shape) + for i in range(shells.size): + bval_centroids[kmeans_res.labels_ == i] = shells[i][0] + + return np.sort(np.squeeze(shells, axis=-1)), bval_centroids diff --git a/test/test_model.py b/test/test_model.py index 5f7fa4b1..c5e03a3b 100644 --- a/test/test_model.py +++ b/test/test_model.py @@ -91,6 +91,47 @@ def test_average_model(): assert np.all(tmodel_2000.predict([0, 0, 0]) == 1100) +def test_gp_model(): + + import nibabel as nib + # ToDo + # Transition to the HDF5 format eventually + root_path = "nireports/nireports/tests/data" + fname_dwi = root_path + "/ds000114_sub-01_ses-test_desc-trunc_dwi.nii.gz" + fname_bval = root_path + "/ds000114_sub-01_ses-test_desc-trunc_dwi.bval" + fname_bvec = root_path + "/ds000114_sub-01_ses-test_desc-trunc_dwi.bvec" + dwi = nib.load(fname_dwi) + bvecs = np.loadtxt(fname_bvec).T + bvals = np.loadtxt(fname_bval) + b0s_mask = bvals < 50 + gradients = np.hstack([bvecs[~b0s_mask], bvals[~b0s_mask, None]]) + # ToDo: what if we are in the multi-shell case ? + # Assume single shell case for now + num_gradients = 1 + # a = np.zeros(num_gradients) # acquisition parameters + # h = nib.load() # Susceptibility induced off‐resonance field (Hz) + # ToDo + # Provide proper values/estimates for these + a = 1 + h = 1 # should be a nifti image + + # ToDo + # Build kernel properly + kernel = np.ones(20) + + num_iterations = 5 + gp = model.GaussianProcessModel(dwi=dwi, a=a, h=h, kernel=kernel, num_iterations=num_iterations) + + _dwi = dwi[..., ~b0s_mask] + + lovo_index = 5 + train_idx = range(len(bvals)) != lovo_index + data_train = _dwi[..., train_idx] + data_train = _dwi + gp.fit(data_train[0], data_train[1]) + predicted1 = gp.predict(data_test[1]) + + def test_two_initialisations(datadir): """Check that the two different initialisations result in the same models"""