metaarray/test.py at main · outofculture/metaarray · 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
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
import os

import numpy as np

from MetaArray import axis, MetaArray


def test_metaarray():
    # Create an array with every option possible
    arr = np.zeros((2, 5, 3, 5), dtype=int)
    for i in range(arr.shape[0]):
        for j in range(arr.shape[1]):
            for k in range(arr.shape[2]):
                for w in range(arr.shape[3]):
                    arr[i, j, k, w] = (i + 1) * 1000 + (j + 1) * 100 + (k + 1) * 10 + (w + 1)

    info = [
        axis("Axis1"),
        axis("Axis2", values=[1, 2, 3, 4, 5]),
        axis("Axis3", cols=["Ax3Col1", ("Ax3Col2", "mV", "Axis3 Column2"), (("Ax3", "Col3"), "A", "Axis3 Column3")]),
        {"name": "Axis4", "values": np.array([1.1, 1.2, 1.3, 1.4, 1.5]), "units": "s"},
        {"extra": "info"},
    ]

    ma = MetaArray(arr, info=info)

    print("====  Original Array =======")
    print(ma)
    print("\n\n")

    # Index/slice tests: check that all values and meta info are correct after slice
    print("\n -- normal integer indexing\n")

    print("\n  ma[1]")
    print(ma[1])

    print("\n  ma[1, 2:4]")
    print(ma[1, 2:4])

    print("\n  ma[1, 1:5:2]")
    print(ma[1, 1:5:2])

    print("\n -- named axis indexing\n")

    print("\n  ma['Axis2':3]")
    print(ma["Axis2":3])

    print("\n  ma['Axis2':3:5]")
    print(ma["Axis2":3:5])

    print("\n  ma[1, 'Axis2':3]")
    print(ma[1, "Axis2":3])

    print("\n  ma[:, 'Axis2':3]")
    print(ma[:, "Axis2":3])

    print("\n  ma['Axis2':3, 'Axis4':0:2]")
    print(ma["Axis2":3, "Axis4":0:2])

    print("\n -- column name indexing\n")

    print("\n  ma['Axis3':'Ax3Col1']")
    print(ma["Axis3":"Ax3Col1"])

    print("\n  ma['Axis3':('Ax3','Col3')]")
    print(ma["Axis3":("Ax3", "Col3")])

    print("\n  ma[:, :, 'Ax3Col2']")
    print(ma[:, :, "Ax3Col2"])

    print("\n  ma[:, :, ('Ax3','Col3')]")
    print(ma[:, :, ("Ax3", "Col3")])

    print("\n -- axis value range indexing\n")

    print("\n  ma['Axis2':1.5:4.5]")
    print(ma["Axis2":1.5:4.5])

    print("\n  ma['Axis4':1.15:1.45]")
    print(ma["Axis4":1.15:1.45])

    print("\n  ma['Axis4':1.15:1.25]")
    print(ma["Axis4":1.15:1.25])

    print("\n -- list indexing\n")

    print("\n  ma[:, [0,2,4]]")
    print(ma[:, [0, 2, 4]])

    print("\n  ma['Axis4':[0,2,4]]")
    print(ma["Axis4":[0, 2, 4]])

    print("\n  ma['Axis3':[0, ('Ax3','Col3')]]")
    print(ma["Axis3":[0, ("Ax3", "Col3")]])

    print("\n -- boolean indexing\n")

    print("\n  ma[:, array([True, True, False, True, False])]")
    print(ma[:, np.array([True, True, False, True, False])])

    print("\n  ma['Axis4':array([True, False, False, False])]")
    print(ma["Axis4": np.array([True, False, False, False, False])])

    # Array operations:
    #  - Concatenate
    #  - Append
    #  - Extend
    #  - Rowsort

    # File I/O tests

    print("\n================  File I/O Tests  ===================\n")
    tf = "test.ma"
    # write whole array

    print("\n  -- write/read test")
    ma.write(tf)
    ma2 = MetaArray(file=tf)

    print("\nArrays are equivalent:", np.all(ma == ma2))

    del ma
    del ma2
    ma = MetaArray(file=tf, writable=True)
    before = ma[0][0][0].mean()
    ma[0][0] += 1
    after = ma[0][0][0].mean()
    assert before + 1 == after, (before, after)

    # CSV write

    # append mode

    print("\n================append test (%s)===============" % tf)
    ma = MetaArray(file=tf)
    ma["Axis2":0:2].write(tf, appendAxis="Axis2")
    for i in range(2, ma.shape[1]):
        ma["Axis2":[i]].write(tf, appendAxis="Axis2")

    ma2 = MetaArray(file=tf)

    print("\nArrays are equivalent:", (ma == ma2).all())

    os.remove(tf)

    # memmap test
    print("\n==========Memmap test============")
    ma.write(tf, mappable=True)
    ma2 = MetaArray(file=tf, mmap=True)
    print("\nArrays are equivalent:", (ma == ma2).all())
    os.remove(tf)


if __name__ == "__main__":
    test_metaarray()