File: matrix.py

package info (click to toggle)
python-suitesparse-graphblas 7.4.0.0-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 968 kB
  • sloc: ansic: 9,939; python: 3,575; sh: 52; makefile: 8
file content (202 lines) | stat: -rw-r--r-- 5,232 bytes parent folder | download | duplicates (2)
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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
from suitesparse_graphblas import check_status, ffi, lib

from .io.serialize import deserialize_matrix as deserialize  # noqa
from .io.serialize import serialize_matrix as serialize  # noqa


def free(A):
    """Free a matrix."""
    check_status(A, lib.GrB_Matrix_free(A))


def new(T, nrows=lib.GxB_INDEX_MAX, ncols=lib.GxB_INDEX_MAX, *, free=free):
    """Create a new `GrB_Matrix` of type `T` and initialize it.  The
    following example creates an eight bit unsigned 2x2 matrix:

    >>> A = new(lib.GrB_UINT8, 2, 2)
    >>> shape(A)
    (2, 2)

    The default value for `nrows` and `ncols` is `lib.GxB_INDEX_MAX`
    which creates a Matrix with maximal bounds:

    >>> A = new(lib.GrB_UINT8)
    >>> shape(A) == (lib.GxB_INDEX_MAX, lib.GxB_INDEX_MAX)
    True

    The `free` argument is called when the object is garbage
    collected, the default is `matrix.free()`.  If `free` is None then
    there is no automatic garbage collection and it is up to the user
    to free the matrix.

    """
    A = ffi.new("GrB_Matrix*")
    check_status(A, lib.GrB_Matrix_new(A, T, nrows, ncols))
    if free:
        return ffi.gc(A, free)
    return A


def type(A):
    """Return the GraphBLAS type of the vector.

    >>> A = new(lib.GrB_UINT8)
    >>> type(A) == lib.GrB_UINT8
    True

    """
    T = ffi.new("GrB_Type*")
    check_status(A, lib.GxB_Matrix_type(T, A[0]))
    return T[0]


def nrows(A):
    """Return the number of rows in the matrix.

    >>> A = new(lib.GrB_UINT8, 2, 3)
    >>> nrows(A)
    2

    """
    n = ffi.new("GrB_Index*")
    check_status(A, lib.GrB_Matrix_nrows(n, A[0]))
    return n[0]


def ncols(A):
    """Return the number of columns in the matrix.

    >>> A = new(lib.GrB_UINT8, 2, 3)
    >>> ncols(A)
    3

    """
    n = ffi.new("GrB_Index*")
    check_status(A, lib.GrB_Matrix_ncols(n, A[0]))
    return n[0]


def nvals(A):
    """Return the number of stored elements in the matrix.

    >>> A = new(lib.GrB_UINT8, 2, 3)
    >>> nvals(A)
    0

    """
    n = ffi.new("GrB_Index*")
    check_status(A, lib.GrB_Matrix_nvals(n, A[0]))
    return n[0]


def shape(A):
    """Return the shape of the matrix as a two tuple `(nrows, ncols)`

    >>> A = new(lib.GrB_UINT8, 2, 2)
    >>> shape(A)
    (2, 2)

    """
    return (nrows(A), ncols(A))


def format(A):
    """Return the format of the matrix.

    >>> A = new(lib.GrB_UINT8, 2, 2)
    >>> format(A) == lib.GxB_BY_ROW
    True

    """
    format = ffi.new("int32_t*")
    check_status(A, lib.GxB_Matrix_Option_get_INT32(A[0], lib.GxB_FORMAT, format))
    return format[0]


def set_format(A, format):
    """Set the format of the matrix.

    >>> A = new(lib.GrB_UINT8, 2, 2)
    >>> set_format(A, lib.GxB_BY_COL)
    >>> format(A) == lib.GxB_BY_COL
    True

    """
    format_val = ffi.cast("int32_t", format)
    check_status(A, lib.GxB_Matrix_Option_set_INT32(A[0], lib.GxB_FORMAT, format_val))


def sparsity_status(A):
    """Get the sparsity status of the matrix."""
    sparsity_status = ffi.new("int32_t*")
    check_status(A, lib.GxB_Matrix_Option_get_INT32(A[0], lib.GxB_SPARSITY_STATUS, sparsity_status))
    return sparsity_status[0]


def sparsity_control(A):
    """Get the sparsity control of the matrix."""
    sparsity_control = ffi.new("int32_t*")
    check_status(
        A, lib.GxB_Matrix_Option_get_INT32(A[0], lib.GxB_SPARSITY_CONTROL, sparsity_control)
    )
    return sparsity_control[0]


def set_sparsity_control(A, sparsity):
    """Set the sparsity control of the matrix."""
    sparsity_control = ffi.cast("int32_t", sparsity)
    check_status(
        A, lib.GxB_Matrix_Option_set_INT32(A[0], lib.GxB_SPARSITY_CONTROL, sparsity_control)
    )


def hyper_switch(A):
    """Get the hyper switch of the matrix."""
    hyper_switch = ffi.new("double*")
    check_status(A, lib.GxB_Matrix_Option_get_FP64(A[0], lib.GxB_HYPER_SWITCH, hyper_switch))
    return hyper_switch[0]


def set_hyper_switch(A, hyper_switch):
    """Set the hyper switch of the matrix."""
    hyper_switch = ffi.cast("double", hyper_switch)
    check_status(A, lib.GxB_Matrix_Option_set_FP64(A[0], lib.GxB_HYPER_SWITCH, hyper_switch))


def bitmap_switch(A):
    """Get the bitmap switch of the matrix."""
    bitmap_switch = ffi.new("double*")
    check_status(A, lib.GxB_Matrix_Option_get_FP64(A[0], lib.GxB_BITMAP_SWITCH, bitmap_switch))
    return bitmap_switch[0]


def set_bitmap_switch(A, bitmap_switch):
    """Set the bitmap switch of the matrix."""
    bitmap_switch = ffi.cast("double", bitmap_switch)
    check_status(A, lib.GxB_Matrix_Option_set_FP64(A[0], lib.GxB_BITMAP_SWITCH, bitmap_switch))


def set_bool(A, value, i, j):
    """Set a boolean value to the matrix at row `i` column `j`.

    >>> A = new(lib.GrB_BOOL, 3, 3)
    >>> set_bool(A, True, 2, 2)
    >>> bool(A, 2, 2) == True
    True

    """
    check_status(A, lib.GrB_Matrix_setElement_BOOL(A[0], value, i, j))


def bool(A, i, j):
    """Get a boolean value from the matrix at row `i` column `j`.

    >>> A = new(lib.GrB_BOOL, 3, 3)
    >>> set_bool(A, True, 2, 2)
    >>> bool(A, 2, 2) == True
    True

    """
    value = ffi.new("bool*")
    check_status(A, lib.GrB_Matrix_extractElement_BOOL(value, A[0], i, j))
    return value[0]