probeinterface/tests/test_plotting.py at d12265e31cde8bb633fa6d0bda96a2eb4f1a43f9 · kushaangupta/probeinterface · 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
from probeinterface import Probe, ProbeGroup
from probeinterface import generate_dummy_probe, generate_dummy_probe_group
from probeinterface.plotting import plot_probe, plot_probegroup

import matplotlib.pyplot as plt
import numpy as np

import pytest


def test_plot_probe():
    probe = generate_dummy_probe()
    plot_probe(probe)
    plot_probe(probe, with_contact_id=True)
    plot_probe(probe, with_device_index=True)
    plot_probe(probe, text_on_contact=["abcde"[i % 5] for i in range(probe.get_contact_count())])

    # with color
    n = probe.get_contact_count()
    contacts_colors = np.random.rand(n, 3)
    plot_probe(probe, contacts_colors=contacts_colors)

    # 3d
    probe_3d = probe.to_3d(axes="xz")
    plot_probe(probe_3d)

    # on click
    probe.set_device_channel_indices(np.arange(probe.get_contact_count())[::-1])
    plot_probe(probe, show_channel_on_click=True)


def test_plot_probegroup():
    probegroup = generate_dummy_probe_group()

    plot_probegroup(probegroup, same_axes=True, with_contact_id=True)
    plot_probegroup(probegroup, same_axes=False)

    # 3d
    probegroup_3d = ProbeGroup()
    for probe in probegroup.probes:
        probegroup_3d.add_probe(probe.to_3d())
    probegroup_3d.probes[-1].move([0, 150, -50])
    plot_probegroup(probegroup_3d, same_axes=True)


def test_plot_probe_two_side():
    probe = Probe()
    probe.set_contacts(
        positions=np.array(
            [
                [0, 0],
                [0, 10],
                [0, 20],
                [0, 0],
                [0, 10],
                [0, 20],
            ]
        ),
        shapes="circle",
        contact_ids=["F1", "F2", "F3", "B1", "B2", "B3"],
        contact_sides=["front", "front", "front", "back", "back", "back"],
    )

    plot_probe(probe, with_contact_id=True, side="front")
    plot_probe(probe, with_contact_id=True, side="back")


if __name__ == "__main__":
    # test_plot_probe()
    test_plot_probe_two_side()
    plt.show()