diff --git a/dipy/direction/__init__.py b/dipy/direction/__init__.py index 5374a1db45..a3b247841a 100644 --- a/dipy/direction/__init__.py +++ b/dipy/direction/__init__.py @@ -11,6 +11,7 @@ from .probabilistic_direction_getter import ( DeterministicMaximumDirectionGetter, ProbabilisticDirectionGetter, + FlockingDirectionGetter, ) from .ptt_direction_getter import PTTDirectionGetter @@ -19,6 +20,7 @@ "ClosestPeakDirectionGetter", "DeterministicMaximumDirectionGetter", "ProbabilisticDirectionGetter", + "FlockingDirectionGetter", "PTTDirectionGetter", "PeaksAndMetrics", "peak_directions", diff --git a/dipy/direction/probabilistic_direction_getter.pxd b/dipy/direction/probabilistic_direction_getter.pxd index 5b1469b082..9c1d99e8b0 100644 --- a/dipy/direction/probabilistic_direction_getter.pxd +++ b/dipy/direction/probabilistic_direction_getter.pxd @@ -8,3 +8,12 @@ cdef class ProbabilisticDirectionGetter(PmfGenDirectionGetter): cdef class DeterministicMaximumDirectionGetter(ProbabilisticDirectionGetter): pass + +cdef class FlockingDirectionGetter(PmfGenDirectionGetter): + cdef: + double[:, :] vertices + int particle_count + double r_min + double r_max + double delta + double alpha diff --git a/dipy/direction/probabilistic_direction_getter.pyx b/dipy/direction/probabilistic_direction_getter.pyx index 1f74f7c727..67568f370b 100644 --- a/dipy/direction/probabilistic_direction_getter.pyx +++ b/dipy/direction/probabilistic_direction_getter.pyx @@ -16,6 +16,12 @@ from dipy.direction.closest_peak_direction_getter cimport PmfGenDirectionGetter from dipy.utils.fast_numpy cimport (copy_point, cumsum, norm, normalize, where_to_insert) +from dipy.tracking.stopping_criterion cimport (StreamlineStatus, + StoppingCriterion, + TRACKPOINT, + ENDPOINT, + OUTSIDEIMAGE, + INVALIDPOINT,) cdef class ProbabilisticDirectionGetter(PmfGenDirectionGetter): """Randomly samples direction of a sphere based on probability mass @@ -187,3 +193,233 @@ cdef class DeterministicMaximumDirectionGetter(ProbabilisticDirectionGetter): newdir[2] = newdir[2] * -1 copy_point(&newdir[0], &direction[0]) return 0 + +cdef class FlockingDirectionGetter(PmfGenDirectionGetter): + """ + """ + + def __init__(self, pmf_gen, max_angle, sphere, pmf_threshold=.1, particle_count = 16, r_min=0.1, r_max=1.0, delta=0.995, alpha=1.0, **kwargs): + """Direction getter from a pmf generator. + + Parameters + ---------- + pmf_gen : PmfGen + Used to get probability mass function for selecting tracking + directions. + max_angle : float, [0, 90] + The maximum allowed angle between incoming direction and new + direction. + sphere : Sphere + The set of directions to be used for tracking. + pmf_threshold : float [0., 1.] + Used to remove direction from the probability mass function for + selecting the tracking direction. + particle_count + Number of particles in a flock + r_min + Minimum distance for gravitational attraction + r_max + Maximum distance for gravitational attraction + delta + Weight parameter between probabilistic direction and gravitational attraction + alpha + Radius of the flock seeding ball at the start + relative_peak_threshold : float in [0., 1.] + Used for extracting initial tracking directions. Passed to + peak_directions. + min_separation_angle : float in [0, 90] + Used for extracting initial tracking directions. Passed to + peak_directions. + + See Also + -------- + dipy.direction.peaks.peak_directions + + """ + PmfGenDirectionGetter.__init__(self, pmf_gen, max_angle, sphere, + pmf_threshold, **kwargs) + self.vertices = self.sphere.vertices.copy() + self.particle_count = particle_count + self.r_min = r_min + self.r_max = r_max + self.delta = delta + self.alpha= alpha + + + + cdef int get_direction_c(self, double[::1] point, double[::1] direction): + """Samples a pmf to updates ``direction`` array with a new direction. + + Parameters + ---------- + point : memory-view (or ndarray), shape (3,) + The point in an image at which to lookup tracking directions. + direction : memory-view (or ndarray), shape (3,) + Previous tracking direction. + + Returns + ------- + status : int + Returns 0 `direction` was updated with a new tracking direction, or + 1 otherwise. + + """ + cdef: + cnp.npy_intp i, idx, _len + double[:] newdir + double* pmf + double last_cdf, cos_sim + + _len = self.len_pmf + pmf = self._get_pmf(point) + + + if norm(&direction[0]) == 0: + return 1 + normalize(&direction[0]) + + with nogil: + for i in range(_len): + cos_sim = self.vertices[i][0] * direction[0] \ + + self.vertices[i][1] * direction[1] \ + + self.vertices[i][2] * direction[2] + if cos_sim < 0: + cos_sim = cos_sim * -1 + if cos_sim < self.cos_similarity: + pmf[i] = 0 + + cumsum(pmf, pmf, _len) + last_cdf = pmf[_len - 1] + if last_cdf == 0: + return 1 + + idx = where_to_insert(pmf, random() * last_cdf, _len) + + newdir = self.vertices[idx] + # Update direction and return 0 for error + if (direction[0] * newdir[0] + + direction[1] * newdir[1] + + direction[2] * newdir[2] > 0): + copy_point(&newdir[0], &direction[0]) + else: + newdir[0] = newdir[0] * -1 + newdir[1] = newdir[1] * -1 + newdir[2] = newdir[2] * -1 + copy_point(&newdir[0], &direction[0]) + return 0 + + cpdef tuple generate_streamline(self, + double[::1] seed, + double[::1] direction, + double[::1] voxel_size, + double step_size, + StoppingCriterion stopping_criterion, + cnp.float_t[:, :] streamline, + StreamlineStatus stream_status, + int fixedstep + ): + cdef: + cnp.npy_intp i + cnp.npy_intp len_streamlines = streamline.shape[0] + cnp.npy_intp particle_count = self.particle_count + double point[3] + double point_par[3] + double voxdir[3] + double dir[3] + double N, G, r_min, r_max, delta_1, delta_2, alpha, norma_vec + cnp.float_t[:, :] particles = np.empty((particle_count, 3), dtype=np.float64) + cnp.float_t[:, :] particles_dir = np.empty((particle_count, 3), dtype=np.float64) + cnp.float_t[:, :] u_m = np.empty((particle_count, 3), dtype=np.float64) + cnp.uint8_t[:] particle_values = np.ones(particle_count, dtype=np.uint8) + cnp.float_t[:, :] F_m = np.empty((particle_count, 3), dtype=np.float64) + double voxel_min[3] + double voxel_max[3] + + copy_point(&seed[0], point) + + terminate = False + r_min = self.r_min + r_max = self.r_max + delta_1 = self.delta + delta_2 = 1-delta_1 + alpha = self.alpha + + + for k in range(1, particle_count): + copy_point(&direction[0], &particles_dir[k, 0]) + vec = np.random.normal(0, 1, 3) + norma_vec = np.linalg.norm(vec) + while norma_vec == 0: + vec = np.random.normal(0, 1, 3) + norma_vec = np.linalg.norm(vec) + vec /= norma_vec + r = np.cbrt(np.random.uniform(0, 1)) * alpha + for j in range(3): + particles[k, j] = seed[j] + r*vec[j] + + copy_point(&direction[0], &particles_dir[0, 0]) + copy_point(&seed[0], &particles[0, 0]) + copy_point(&seed[0], &streamline[0,0]) + + stream_status = TRACKPOINT + for i in range(1, len_streamlines): + if terminate: + break + # Particles + for m in range(particle_count): + copy_point(&particles[m, 0], point_par) + copy_point(&particles_dir[m, 0], dir) + if self.get_direction_c(point_par, dir): + particle_values[m] = 0 + if m==0: + terminate = True + continue + copy_point(dir, &particles_dir[m, 0]) + + for m in range(particle_count): + for j in range(3): + F_m[m, j] = 0 + if particle_values[m] == 0: + continue + for k in range(particle_count): + if k != m: + for j in range(3): + u_m[k, j] = particles[m, j] - particles[k, j] + N = norm(&u_m[k,0]) + N = N*N + if N < r_min: + for j in range(3): + u_m[k, j] /= r_min + F_m[m, j] += -u_m[k, j] + elif r_min <= N and N