Measure length of streamline passing trough each voxel. Plus some utility functions to test, bench and vizualize.

length_per_voxel.py

from dipy.viz import fvtk

import numpy as np

from numpy.testing import measure
from nose.tools import assert_equal, assert_almost_equal
from numpy.testing import assert_array_equal, assert_array_almost_equal

norm = lambda x: np.sqrt(np.sum(x**2))


def get_closest_edge(p, d, eps=1):
    edge = np.zeros_like(p)
    edge[0] = np.floor(p[0]+eps) if d[0] >= 0 else np.ceil(p[0]-eps)
    edge[1] = np.floor(p[1]+eps) if d[1] >= 0 else np.ceil(p[1]-eps)
    edge[2] = np.floor(p[2]+eps) if d[2] >= 0 else np.ceil(p[2]-eps)

    return edge


def show(streamline, points):
    r = fvtk.ren()

    fvtk.add(r, fvtk.line(streamline, (0, 0, 1), opacity=1, linewidth=1))

    min_pts = np.floor(np.min(streamline, axis=0)).astype('int')
    max_pts = np.ceil(np.max(streamline, axis=0)).astype('int')

    grid = []
    for x in range(min_pts[0], max_pts[0]+1):
        for y in range(min_pts[1], max_pts[1]+1):
            for z in range(min_pts[2], max_pts[2]+1):
                grid.append([(min_pts[0], y, z), (max_pts[0], y, z)])
                grid.append([(x, min_pts[1], z), (x, max_pts[1], z)])
                grid.append([(x, y, min_pts[2]), (x, y, max_pts[2])])

    fvtk.add(r, fvtk.line(grid, (0.2, 0.2, 0.2), opacity=1, linewidth=0.5))

    fvtk.add(r, fvtk.point([(0, 0, 0)], (0, 1, 0), opacity=1, point_radius=0.01))
    fvtk.add(r, fvtk.line([[(0, 0, 0), (.1, 0, 0)]], (1, 0, 0), opacity=1, linewidth=0.5))
    fvtk.add(r, fvtk.line([[(0, 0, 0), (0, .1, 0)]], (0, 1, 0), opacity=1, linewidth=0.5))
    fvtk.add(r, fvtk.line([[(0, 0, 0), (0, 0, .1)]], (0, 0, 1), opacity=1, linewidth=0.5))

    fvtk.add(r, fvtk.point(points, (1, 0, 0), opacity=1, point_radius=0.01))

    fvtk.show(r, title='Streamlines')


# CAUTION: Modify array x
def insert_unique_neighbor(x, k, v):
    if len(x) == 0 or x[-1][0] != k:
        x.append([k, v])
    else:
        x[-1][1] += v


def get_lengths_per_voxel(streamline):
    flags = np.seterr(divide="ignore", under="ignore")

    lengths = []
    pts = []
    for a, b in zip(streamline[:-1], streamline[1:]):
        d = b-a
        norm_d = norm(d)
        dist = norm_d

        # Check if first point is lying on an edge.
        edge = a
        if not np.any(np.floor(a) == a):
            edge = get_closest_edge(a, d)

        while True:
            ratio = np.nanmin(np.abs((edge-a)/d))
            dist -= ratio*norm_d

            # Use np.allclose because last point could be already lying on an edge.
            if dist < 0 and not np.allclose(dist, 0):
                break

            k = tuple(np.floor(a + 0.5*ratio*d))
            v = ratio*norm_d
            insert_unique_neighbor(lengths, k, v)

            a = ratio*d + a
            a[np.abs(a) <= 1e-16] = 0.0  # Snap values near zero
            pts.append(a)
            edge = get_closest_edge(a, d)

        k = tuple(np.floor(a + 0.5*(b-a)))
        v = norm(b-a)
        insert_unique_neighbor(lengths, k, v)

    np.seterr(**flags)
    return lengths, pts


def test():
    # Superior-Inferior
    streamline = np.array([(.7, 1.2, 0.5),
                           (.7, -0.2, 0.5)])

    lengths, pts = get_lengths_per_voxel(streamline)

    assert_equal(len(pts), 2)
    assert_array_equal(pts[0], [0.7, 1., 0.5])
    assert_array_equal(pts[1], [0.7, 0., 0.5])

    assert_equal(len(lengths), 3)
    assert_equal(lengths[0][0], (0, 1, 0))
    assert_equal(lengths[1][0], (0, 0, 0))
    assert_equal(lengths[2][0], (0, -1, 0))

    assert_almost_equal(lengths[0][1], 0.2)
    assert_almost_equal(lengths[1][1], 1.)
    assert_almost_equal(lengths[2][1], 0.2)

    # Superior-Inferior and Left-Right
    streamline = np.array([(1.2, -0.2, 0.5),
                           (-0.2, 1.2, 0.5)])

    lengths, pts = get_lengths_per_voxel(streamline)

    assert_equal(len(pts), 2)
    assert_array_equal(pts[0], [1., 0., 0.5])
    assert_array_equal(pts[1], [0., 1., 0.5])

    assert_equal(len(lengths), 3)
    assert_equal(lengths[0][0], (1, -1, 0))
    assert_equal(lengths[1][0], (0, 0, 0))
    assert_equal(lengths[2][0], (-1, 1, 0))

    assert_almost_equal(lengths[0][1], np.sqrt(2 * 0.2**2))
    assert_almost_equal(lengths[1][1], np.sqrt(2 * 1.0**2))
    assert_almost_equal(lengths[2][1], np.sqrt(2 * 0.2**2))

    # # Along an edge
    # streamline = np.array([(1, 1.2, 0.5),
    #                        (1, -0.2, 0.5)])

    # *Not working* Along an edge, starting inside
    # streamline = np.array([(1, 0.8, 0.5),
    #                        (1, -0.2, 0.5)])

    # More complex test
    streamline = np.array([(0, 1.2, 0),
                           (0.1, 0.7, 0.1),
                           (0.5, 0.4, 1.3),
                           (1.2, -0.2, 1)])

    lengths, pts = get_lengths_per_voxel(streamline)

    assert_equal(len(pts), 6)
    assert_array_equal(pts[0], [0., 1.2, 0.])  # Exact value since it is the streamline's first point.
    assert_array_almost_equal(pts[1], [0.04, 1., 0.04])
    assert_array_almost_equal(pts[2], [0.4, 0.475, 1.])
    assert_array_almost_equal(pts[3], [0.96666667, 0., 1.1])
    assert_array_almost_equal(pts[4], [1. , -0.02857143, 1.08571429])
    assert_array_almost_equal(pts[5], [1.2, -0.2, 1.])

    assert_equal(len(lengths), 5)
    assert_equal(lengths[0][0], (0, 1, 0))
    assert_equal(lengths[1][0], (0, 0, 0))
    assert_equal(lengths[2][0], (0, 0, 1))
    assert_equal(lengths[3][0], (0, -1, 1))
    assert_equal(lengths[4][0], (1, -1, 1))


def bench():
    NB_POINTS = 100
    repeat = 10
    streamline = np.random.rand(NB_POINTS, 3) * 15
    print("Timing get_lengths_per_voxel() in Python")
    python_time = measure("get_lengths_per_voxel(streamline)", repeat)
    print("Python time: {0:.2}sec".format(python_time))


if __name__ == '__main__':
    test()
    bench()

    streamline = np.array([(0, 1.2, 0),
                           (0.1, 0.7, 0.1),
                           (0.5, 0.4, 1.3),
                           (1.2, -0.2, 1),
                           (0.3, 1.1, 0.3)])

    lengths, pts = get_lengths_per_voxel(streamline)
    print "\n".join(map(str, lengths))
    print "\n".join(map(str, pts))
    show(streamline, pts)