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
|
import numpy as np
from .._utils import can_store, check_consistent_fill_value, normalize_axis
def concatenate(arrays, axis=0, compressed_axes=None):
from .compressed import GCXS
check_consistent_fill_value(arrays)
arrays = [arr if isinstance(arr, GCXS) else GCXS(arr, compressed_axes=(axis,)) for arr in arrays]
axis = normalize_axis(axis, arrays[0].ndim)
dim = sum(x.shape[axis] for x in arrays)
shape = list(arrays[0].shape)
shape[axis] = dim
assert all(x.shape[ax] == arrays[0].shape[ax] for x in arrays for ax in set(range(arrays[0].ndim)) - {axis})
if compressed_axes is None:
compressed_axes = (axis,)
if arrays[0].ndim == 1:
from .._coo.common import concatenate as coo_concat
arrays = [arr.tocoo() for arr in arrays]
return coo_concat(arrays, axis=axis)
# arrays may have different compressed_axes
# concatenating becomes easy when compressed_axes are the same
arrays = [arr.change_compressed_axes((axis,)) for arr in arrays]
ptr_list = []
for i, arr in enumerate(arrays):
if i == 0:
ptr_list.append(arr.indptr)
continue
ptr_list.append(arr.indptr[1:])
indptr = np.concatenate(ptr_list)
indices = np.concatenate([arr.indices for arr in arrays])
data = np.concatenate([arr.data for arr in arrays])
ptr_len = arrays[0].indptr.shape[0]
nnz = arrays[0].nnz
total_nnz = sum(int(arr.nnz) for arr in arrays)
if not can_store(indptr.dtype, total_nnz):
indptr = indptr.astype(np.min_scalar_type(total_nnz))
for i in range(1, len(arrays)):
indptr[ptr_len:] += nnz
nnz = arrays[i].nnz
ptr_len += arrays[i].indptr.shape[0] - 1
return GCXS(
(data, indices, indptr),
shape=tuple(shape),
compressed_axes=arrays[0].compressed_axes,
fill_value=arrays[0].fill_value,
).change_compressed_axes(compressed_axes)
def stack(arrays, axis=0, compressed_axes=None):
from .compressed import GCXS
check_consistent_fill_value(arrays)
arrays = [arr if isinstance(arr, GCXS) else GCXS(arr, compressed_axes=(axis,)) for arr in arrays]
axis = normalize_axis(axis, arrays[0].ndim + 1)
assert all(x.shape[ax] == arrays[0].shape[ax] for x in arrays for ax in set(range(arrays[0].ndim)) - {axis})
if compressed_axes is None:
compressed_axes = (axis,)
if arrays[0].ndim == 1:
from .._coo.common import stack as coo_stack
arrays = [arr.tocoo() for arr in arrays]
return coo_stack(arrays, axis=axis)
# arrays may have different compressed_axes
# stacking becomes easy when compressed_axes are the same
ptr_list = []
for i in range(len(arrays)):
shape = list(arrays[i].shape)
shape.insert(axis, 1)
arrays[i] = arrays[i].reshape(shape).change_compressed_axes((axis,))
if i == 0:
ptr_list.append(arrays[i].indptr)
continue
ptr_list.append(arrays[i].indptr[1:])
shape[axis] = len(arrays)
indptr = np.concatenate(ptr_list)
indices = np.concatenate([arr.indices for arr in arrays])
data = np.concatenate([arr.data for arr in arrays])
ptr_len = arrays[0].indptr.shape[0]
nnz = arrays[0].nnz
total_nnz = sum(int(arr.nnz) for arr in arrays)
if not can_store(indptr.dtype, total_nnz):
indptr = indptr.astype(np.min_scalar_type(total_nnz))
for i in range(1, len(arrays)):
indptr[ptr_len:] += nnz
nnz = arrays[i].nnz
ptr_len += arrays[i].indptr.shape[0] - 1
return GCXS(
(data, indices, indptr),
shape=tuple(shape),
compressed_axes=arrays[0].compressed_axes,
fill_value=arrays[0].fill_value,
).change_compressed_axes(compressed_axes)
|
