import numpy as np
from scipy.special import lpmv, gammaln
from tqdm import tqdm
from dipy.align import Bunch
from dipy.tracking.local_tracking import LocalTracking
import random
import sys
import math
def spherical_harmonics(m, n, theta, phi):
"""
An implementation of spherical harmonics that overcomes conda compilation
issues. See: https://github.com/nipy/dipy/issues/852
"""
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
TissueTypes = Bunch(OUTSIDEIMAGE=-1, INVALIDPOINT=0, TRACKPOINT=1, ENDPOINT=2)
class VerboseLocalTracking(LocalTracking):
def __init__(self, *args, min_length=10, max_length=1000, **kwargs):
super().__init__(*args, **kwargs)
self.min_length = min_length
self.max_length = max_length
def _generate_streamlines(self):
"""A streamline generator"""
# Get inverse transform (lin/offset) for seeds
inv_A = np.linalg.inv(self.affine)
lin = inv_A[:3, :3]
offset = inv_A[:3, 3]
F = np.empty((self.max_length + 1, 3), dtype=float)
B = F.copy()
for s in tqdm(self.seeds):
s = np.dot(lin, s) + offset
# Set the random seed in numpy and random
if self.random_seed is not None:
s_random_seed = hash(np.abs((np.sum(s)) + self.random_seed)) \
% (2**32 - 1)
random.seed(s_random_seed)
np.random.seed(s_random_seed)
directions = self.direction_getter.initial_direction(s)
if directions.size == 0 and self.return_all:
# only the seed position
yield [s]
directions = directions[:self.max_cross]
for first_step in directions:
stepsF, tissue_class = self._tracker(s, first_step, F)
if not (self.return_all
or tissue_class == TissueTypes.ENDPOINT
or tissue_class == TissueTypes.OUTSIDEIMAGE):
continue
first_step = -first_step
stepsB, tissue_class = self._tracker(s, first_step, B)
if not (self.return_all
or tissue_class == TissueTypes.ENDPOINT
or tissue_class == TissueTypes.OUTSIDEIMAGE):
continue
if stepsB == 1:
streamline = F[:stepsF].copy()
else:
parts = (B[stepsB - 1:0:-1], F[:stepsF])
streamline = np.concatenate(parts, axis=0)
len_sl = len(streamline)
if len_sl < self.min_length * self.step_size \
or len_sl > self.max_length * self.step_size:
continue
else:
yield streamline
def in_place_norm(vec, axis=-1, keepdims=False, delvec=True):
""" Return Vectors with Euclidean (L2) norm
See :term:`unit vector` and :term:`Euclidean norm`
Parameters
-------------
vec : array_like
Vectors to norm. Squared in the process of calculating the norm.
axis : int, optional
Axis over which to norm. By default norm over last axis. If `axis` is
None, `vec` is flattened then normed. Default is -1.
keepdims : bool, optional
If True, the output will have the same number of dimensions as `vec`,
with shape 1 on `axis`. Default is False.
delvec : bool, optional
If True, vec is deleted as soon as possible.
If False, vec is not deleted, but still squared. Default is True.
Returns
---------
norm : array
Euclidean norms of vectors.
Examples
--------
>>> vec = [[8, 15, 0], [0, 36, 77]]
>>> in_place_norm(vec)
array([ 17., 85.])
>>> vec = [[8, 15, 0], [0, 36, 77]]
>>> in_place_norm(vec, keepdims=True)
array([[ 17.],
[ 85.]])
>>> vec = [[8, 15, 0], [0, 36, 77]]
>>> in_place_norm(vec, axis=0)
array([ 8., 39., 77.])
"""
vec = np.asarray(vec)
if keepdims:
ndim = vec.ndim
shape = vec.shape
np.square(vec, out=vec)
vec_norm = vec.sum(axis)
if delvec:
del vec
try:
np.sqrt(vec_norm, out=vec_norm)
except TypeError:
vec_norm = vec_norm.astype(float)
np.sqrt(vec_norm, out=vec_norm)
if keepdims:
if axis is None:
shape = [1] * ndim
else:
shape = list(shape)
shape[axis] = 1
vec_norm = vec_norm.reshape(shape)
return vec_norm
def tensor_odf(evals, evecs, sphere, num_batches=100):
"""
Calculate the tensor Orientation Distribution Function
Parameters
----------
evals : array (4D)
Eigenvalues of a tensor. Shape (x, y, z, 3).
evecs : array (5D)
Eigenvectors of a tensor. Shape (x, y, z, 3, 3)
sphere : sphere object
The ODF will be calculated in each vertex of this sphere.
num_batches : int
Split the calculation into batches. This reduces memory usage.
If memory use is not an issue, set to 1.
If set to -1, there will be 1 batch per vertex in the sphere.
Default: 100
"""
num_vertices = sphere.vertices.shape[0]
if num_batches == -1:
num_batches = num_vertices
batch_size = math.ceil(num_vertices / num_batches)
batches = range(num_batches)
mask = np.where((evals[..., 0] > 0)
& (evals[..., 1] > 0)
& (evals[..., 2] > 0))
evecs = evecs[mask]
proj_norm = np.zeros((num_vertices, evecs.shape[0]))
it = tqdm(batches) if num_batches != 1 else batches
for i in it:
start = i * batch_size
end = (i + 1) * batch_size
if end > num_vertices:
end = num_vertices
proj = np.dot(sphere.vertices[start:end], evecs)
proj /= np.sqrt(evals[mask])
proj_norm[start:end, :] = in_place_norm(proj)
proj_norm **= -3
proj_norm /= 4 * np.pi * np.sqrt(np.prod(evals[mask], -1))
odf = np.zeros((evals.shape[:3] + (sphere.vertices.shape[0],)))
odf[mask] = proj_norm.T
return odf
