From 2056025f409740544ac026617a38c0ef3df756a0 Mon Sep 17 00:00:00 2001 From: Oscar Esteban Date: Wed, 23 Oct 2024 10:44:40 +0200 Subject: [PATCH] enh: add a testing module to keep convenience tooling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Jon Haitz Legarreta Gorroño --- docs/developers.rst | 1 + pyproject.toml | 1 + src/eddymotion/testing/__init__.py | 0 src/eddymotion/testing/simulations.py | 197 ++++++++++++++++++++++++++ 4 files changed, 199 insertions(+) create mode 100644 src/eddymotion/testing/__init__.py create mode 100644 src/eddymotion/testing/simulations.py diff --git a/docs/developers.rst b/docs/developers.rst index e16c028f..fffccf0b 100644 --- a/docs/developers.rst +++ b/docs/developers.rst @@ -36,4 +36,5 @@ Information on specific functions, classes, and methods. api/eddymotion.math api/eddymotion.model api/eddymotion.registration + api/eddymotion.testing api/eddymotion.utils diff --git a/pyproject.toml b/pyproject.toml index 1aa6fa6c..b3a98409 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -177,6 +177,7 @@ branch = true concurrency = ['multiprocessing'] omit = [ '*/tests/*', + '*/testing/*', '*/__init__.py', '*/conftest.py', 'src/eddymotion/_version.py' diff --git a/src/eddymotion/testing/__init__.py b/src/eddymotion/testing/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/eddymotion/testing/simulations.py b/src/eddymotion/testing/simulations.py new file mode 100644 index 00000000..e7776c8e --- /dev/null +++ b/src/eddymotion/testing/simulations.py @@ -0,0 +1,197 @@ +# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*- +# vi: set ft=python sts=4 ts=4 sw=4 et: +# +# © 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/ +# +"""Utilities for testing purposes.""" + +from __future__ import annotations + +# import nibabel as nib +import numpy as np +from dipy.core.geometry import sphere2cart +from dipy.core.gradients import gradient_table +from dipy.core.sphere import HemiSphere, Sphere, disperse_charges +from dipy.sims.voxel import all_tensor_evecs + + +def add_b0(bvals: np.ndarray, bvecs: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """ + Insert a b=0 at the front of the gradient table (b-values and b-vectors). + + Parameters + ---------- + bvals : :obj:`~numpy.ndarray` + Array of b-values. + bvecs : :obj:`~numpy.ndarray` + Array of b-vectors. + + Returns + ------- + :obj:`tuple` + Updated gradient table (b-values, b-vectors) including a b=0. + + """ + _bvals = np.insert(bvals, 0, 0) + _bvecs = np.insert(bvecs, 0, np.array([0, 0, 0]), axis=0) + return _bvals, _bvecs + + +def create_single_fiber_evecs(theta: float = 0, phi: float = 0) -> np.ndarray: + """ + Create eigenvectors for a simulated fiber given the polar coordinates of its pricipal axis. + + Parameters + ---------- + theta : :obj:`float` + Theta coordinate + phi : :obj:`float` + Phi coordinate + + Returns + ------- + :obj:`~numpy.ndarray` + Eigenvectors for a single fiber. + + """ + # Polar coordinates (theta, phi) of the principal axis of the tensor + angles = np.array([theta, phi]) + sticks = np.array(sphere2cart(1, np.deg2rad(angles[0]), np.deg2rad(angles[1]))) + evecs = all_tensor_evecs(sticks) + return evecs + + +def create_random_polar_coordinates( + hsph_dirs: int, seed: int = 1234 +) -> tuple[np.ndarray, np.ndarray]: + """ + Create random polar coordinates. + + Parameters + ---------- + hsph_dirs : :obj:`int` + Number of hemisphere directions. + seed : :obj:`int`, optional + Seed for the random number generator, by default 1234. + + Returns + ------- + :obj:`tuple` + Theta and Phi values of polar coordinates. + + """ + rng = np.random.default_rng(seed) + theta = np.pi * rng.random(hsph_dirs) + phi = 2 * np.pi * rng.random(hsph_dirs) + return theta, phi + + +def create_diffusion_encoding_gradient_dirs( + hsph_dirs: int, iterations: int = 5000, seed: int = 1234 +) -> Sphere: + """ + Create the dMRI gradient-encoding directions. + + Parameters + ---------- + hsph_dirs : :obj:`int` + Number of hemisphere directions. + iterations : :obj:`int`, optional + Number of iterations for charge dispersion, by default 5000. + seed : :obj:`int`, optional + Seed for the random number generator, by default 1234. + + Returns + ------- + :obj:`~dipy.core.sphere.Sphere` + A sphere with diffusion-encoding gradient directions. + + """ + # Create the gradient-encoding directions placing random points on a hemisphere + theta, phi = create_random_polar_coordinates(hsph_dirs, seed=seed) + hsph_initial = HemiSphere(theta=theta, phi=phi) + + # Move the points so that the electrostatic potential energy is minimized + hsph_updated, _ = disperse_charges(hsph_initial, iterations) + + # Create a sphere + return Sphere(xyz=np.vstack((hsph_updated.vertices, -hsph_updated.vertices))) + + +def create_single_shell_gradient_table( + hsph_dirs: int, bval_shell: float, iterations: int = 5000 +) -> gradient_table: + """ + Create a single-shell gradient table. + + Parameters + ---------- + hsph_dirs : :obj:`int` + Number of hemisphere directions. + bval_shell : :obj:`float` + Shell b-value. + iterations : :obj:`int`, optional + Number of iterations for charge dispersion, by default 5000. + + Returns + ------- + :obj:`~dipy.core.gradients.gradient_table` + The gradient table for the single-shell. + + """ + # Create diffusion-encoding gradient directions + sph = create_diffusion_encoding_gradient_dirs(hsph_dirs, iterations=iterations) + + # Create the gradient bvals and bvecs + vertices = sph.vertices + values = np.ones(vertices.shape[0]) + bvecs = vertices + bvals = bval_shell * values + + # Add a b0 value to the gradient table + bvals, bvecs = add_b0(bvals, bvecs) + return gradient_table(bvals, bvecs) + + +def get_query_vectors( + gtab: gradient_table, train_mask: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """ + Get the diffusion-encoding gradient vectors for estimation from the gradient table. + + Get the diffusion-encoding gradient vectors where the signal is to be estimated from the + gradient table and the training mask: the vectors of interest are those that are masked in + the training mask. b0 values are excluded. + + Parameters + ---------- + gtab : :obj:`~dipy.core.gradients.gradient_table` + Gradient table. + train_mask : :obj:`~numpy.ndarray` + Mask for selecting training vectors. + + Returns + ------- + :obj:`tuple` + Gradient vectors and indices for estimation. + + """ + idx = np.logical_and(~train_mask, ~gtab.b0s_mask) + return gtab.bvecs[idx], np.where(idx)[0]