atlas/plotting.py at main · apluff/atlas · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
"""Plotting streamlines in 3D with Datoviz."""

from timeit import default_timer
from math import cos, sin, pi

import numpy as np

from common import *
from streamlines import *

from datoviz import canvas, run, colormap


SHOW_ALLEN = False
MAX_PATHS = 100_000


def ibl_streamlines():
    paths = load_npy(filepath(REGION, 'streamlines'))
    paths = subset(paths, MAX_PATHS)

    # NOTE: the following line should be decommented when streamlines go from bottom to top
    # paths = paths[:, ::-1, :]

    return paths


def allen_streamlines():
    paths = load_npy(filepath(REGION, 'streamlines_allen'))
    paths = subset(paths, MAX_PATHS)
    return paths


def plot_panel(panel, paths):
    assert paths.ndim == 3
    n, l, _ = paths.shape
    assert _ == 3
    length = l * np.ones(n)  # length of each path

    color = np.tile(np.linspace(0, 1, l), n)
    color = colormap(color, vmin=0, vmax=1, cmap='viridis', alpha=1)

    # Plot lines
    v = panel.visual('line_strip', depth_test=True, transform=None)
    paths[:, :, 1] *= -1
    paths = paths.reshape((-1, 3))
    paths -= paths.mean(axis=0)
    paths *= .0035
    paths[:, 1] += .2
    v.data('pos', paths)
    v.data('length', length)
    v.data('color', color)

    # # Plot points
    # v = panel.visual('point', depth_test=True)
    # v.data('pos', paths.reshape((-1, 3)))
    # v.data('ms', np.array([1.0]))
    # v.data('color', color)


c = canvas(width=1920+20, height=1080+20,
           clear_color=(0, 0, 0, 0), show_fps=False)
s = c.scene(cols=2 if SHOW_ALLEN else 1)

if SHOW_ALLEN:
    paths_allen = allen_streamlines()
    p_allen = s.panel(col=0, controller='arcball')
    plot_panel(p_allen, paths_allen)

paths_ibl = ibl_streamlines()
p_ibl = s.panel(col=1 if SHOW_ALLEN else 0, controller='arcball')
plot_panel(p_ibl, paths_ibl)
if SHOW_ALLEN:
    p_allen.link_to(p_ibl)

# We define an event callback to implement mouse picking

t0 = default_timer()


@c.connect
def on_frame(ev):
    t = default_timer() - t0
    a = +2 * pi * t / 8
    # x = 2 * cos(a)
    # y = 2 * sin(a)
    p_ibl.arcball_rotate(0, 1, 0, a)


run()