storage.py 9.33 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import warnings
rusty1s's avatar
rusty1s committed
2

rusty1s's avatar
rusty1s committed
3
import torch
rusty1s's avatar
rusty1s committed
4
from torch_scatter import segment_csr
rusty1s's avatar
rusty1s committed
5

rusty1s's avatar
typo  
rusty1s committed
6
__cache__ = {'enabled': True}
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
9

def is_cache_enabled():
rusty1s's avatar
typo  
rusty1s committed
10
    return __cache__['enabled']
rusty1s's avatar
rusty1s committed
11
12
13


def set_cache_enabled(mode):
rusty1s's avatar
typo  
rusty1s committed
14
    __cache__['enabled'] = mode
rusty1s's avatar
rusty1s committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31


class no_cache(object):
    def __enter__(self):
        self.prev = is_cache_enabled()
        set_cache_enabled(False)

    def __exit__(self, *args):
        set_cache_enabled(self.prev)
        return False

    def __call__(self, func):
        def decorate_no_cache(*args, **kwargs):
            with self:
                return func(*args, **kwargs)

        return decorate_no_cache
rusty1s's avatar
rusty1s committed
32
33


rusty1s's avatar
rusty1s committed
34
35
36
class cached_property(object):
    def __init__(self, func):
        self.func = func
rusty1s's avatar
sorting  
rusty1s committed
37

rusty1s's avatar
rusty1s committed
38
39
40
41
    def __get__(self, obj, cls):
        value = getattr(obj, f'_{self.func.__name__}', None)
        if value is None:
            value = self.func(obj)
rusty1s's avatar
typo  
rusty1s committed
42
            if is_cache_enabled():
rusty1s's avatar
rusty1s committed
43
                setattr(obj, f'_{self.func.__name__}', value)
rusty1s's avatar
rusty1s committed
44
45
46
        return value


rusty1s's avatar
rusty1s committed
47
48
49
50
def optional(func, src):
    return func(src) if src is not None else src


rusty1s's avatar
rusty1s committed
51
52
53
54
55
56
57
58
59
60
61
62
layouts = ['coo', 'csr', 'csc']


def get_layout(layout=None):
    if layout is None:
        layout = 'coo'
        warnings.warn('`layout` argument unset, using default layout '
                      '"coo". This may lead to unexpected behaviour.')
    assert layout in layouts
    return layout


rusty1s's avatar
rusty1s committed
63
class SparseStorage(object):
rusty1s's avatar
rusty1s committed
64
65
66
67
    cache_keys = [
        'rowcount', 'rowptr', 'colcount', 'colptr', 'csr2csc', 'csc2csr'
    ]

rusty1s's avatar
rusty1s committed
68
69
70
71
72
73
74
75
76
77
78
    def __init__(self,
                 index,
                 value=None,
                 sparse_size=None,
                 rowcount=None,
                 rowptr=None,
                 colcount=None,
                 colptr=None,
                 csr2csc=None,
                 csc2csr=None,
                 is_sorted=False):
rusty1s's avatar
rusty1s committed
79
80
81

        assert index.dtype == torch.long
        assert index.dim() == 2 and index.size(0) == 2
rusty1s's avatar
rusty1s committed
82
        index = index.contiguous()
rusty1s's avatar
rusty1s committed
83
84

        if value is not None:
rusty1s's avatar
rusty1s committed
85
86
            assert value.device == index.device
            assert value.size(0) == index.size(1)
rusty1s's avatar
rusty1s committed
87
88
            value = value.contiguous()

rusty1s's avatar
rusty1s committed
89
90
91
        if sparse_size is None:
            sparse_size = torch.Size((index.max(dim=-1)[0] + 1).tolist())

rusty1s's avatar
rusty1s committed
92
93
94
95
96
        if rowcount is not None:
            assert rowcount.dtype == torch.long
            assert rowcount.device == index.device
            assert rowcount.dim() == 1 and rowcount.numel() == sparse_size[0]

rusty1s's avatar
rusty1s committed
97
        if rowptr is not None:
rusty1s's avatar
rusty1s committed
98
99
            assert rowptr.dtype == torch.long
            assert rowptr.device == index.device
rusty1s's avatar
rusty1s committed
100
            assert rowptr.dim() == 1 and rowptr.numel() - 1 == sparse_size[0]
rusty1s's avatar
rusty1s committed
101

rusty1s's avatar
rusty1s committed
102
103
104
105
106
        if colcount is not None:
            assert colcount.dtype == torch.long
            assert colcount.device == index.device
            assert colcount.dim() == 1 and colcount.numel() == sparse_size[1]

rusty1s's avatar
rusty1s committed
107
        if colptr is not None:
rusty1s's avatar
rusty1s committed
108
109
            assert colptr.dtype == torch.long
            assert colptr.device == index.device
rusty1s's avatar
rusty1s committed
110
            assert colptr.dim() == 1 and colptr.numel() - 1 == sparse_size[1]
rusty1s's avatar
rusty1s committed
111

rusty1s's avatar
rusty1s committed
112
113
114
115
116
        if csr2csc is not None:
            assert csr2csc.dtype == torch.long
            assert csr2csc.device == index.device
            assert csr2csc.dim() == 1
            assert csr2csc.numel() == index.size(1)
rusty1s's avatar
rusty1s committed
117

rusty1s's avatar
rusty1s committed
118
119
120
121
122
        if csc2csr is not None:
            assert csc2csr.dtype == torch.long
            assert csc2csr.device == index.device
            assert csc2csr.dim() == 1
            assert csc2csr.numel() == index.size(1)
rusty1s's avatar
rusty1s committed
123

rusty1s's avatar
rusty1s committed
124
125
126
        if not is_sorted:
            idx = sparse_size[1] * index[0] + index[1]
            # Only sort if necessary...
rusty1s's avatar
rusty1s committed
127
            if (idx < torch.cat([idx.new_zeros(1), idx[:-1]], dim=0)).any():
rusty1s's avatar
rusty1s committed
128
129
130
                perm = idx.argsort()
                index = index[:, perm]
                value = None if value is None else value[perm]
rusty1s's avatar
rusty1s committed
131
132
                csr2csc = None
                csc2csr = None
rusty1s's avatar
rusty1s committed
133

rusty1s's avatar
rusty1s committed
134
135
136
        self._index = index
        self._value = value
        self._sparse_size = sparse_size
rusty1s's avatar
rusty1s committed
137
        self._rowcount = rowcount
rusty1s's avatar
rusty1s committed
138
        self._rowptr = rowptr
rusty1s's avatar
rusty1s committed
139
        self._colcount = colcount
rusty1s's avatar
rusty1s committed
140
        self._colptr = colptr
rusty1s's avatar
rusty1s committed
141
142
        self._csr2csc = csr2csc
        self._csc2csr = csc2csr
rusty1s's avatar
rusty1s committed
143
144

    @property
rusty1s's avatar
rusty1s committed
145
146
    def index(self):
        return self._index
rusty1s's avatar
rusty1s committed
147
148

    @property
rusty1s's avatar
rusty1s committed
149
150
    def row(self):
        return self._index[0]
rusty1s's avatar
rusty1s committed
151
152

    @property
rusty1s's avatar
rusty1s committed
153
154
    def col(self):
        return self._index[1]
rusty1s's avatar
rusty1s committed
155

rusty1s's avatar
rusty1s committed
156
157
    def has_value(self):
        return self._value is not None
rusty1s's avatar
rusty1s committed
158
159

    @property
rusty1s's avatar
rusty1s committed
160
161
162
163
164
165
    def value(self):
        return self._value

    def set_value_(self, value, layout=None):
        assert value.device == self._index.device
        assert value.size(0) == self._index.size(1)
rusty1s's avatar
rusty1s committed
166
167
        if value is not None and get_layout(layout) == 'csc':
            value = value[self.csc2csr]
rusty1s's avatar
rusty1s committed
168
169
        self._value = value
        return self
rusty1s's avatar
rusty1s committed
170
171
172
173

    def set_value(self, value, layout=None):
        assert value.device == self._index.device
        assert value.size(0) == self._index.size(1)
rusty1s's avatar
rusty1s committed
174
175
        if value is not None and get_layout(layout) == 'csc':
            value = value[self.csc2csr]
rusty1s's avatar
rusty1s committed
176
177
178
179
180
181
182
183
184
185
186
187
        return self.__class__(
            self._index,
            value,
            self._sparse_size,
            self._rowcount,
            self._rowptr,
            self._colcount,
            self._colptr,
            self._csr2csc,
            self._csc2csr,
            is_sorted=True,
        )
rusty1s's avatar
rusty1s committed
188
189

    def sparse_size(self, dim=None):
rusty1s's avatar
rusty1s committed
190
        return self._sparse_size if dim is None else self._sparse_size[dim]
rusty1s's avatar
rusty1s committed
191
192
193

    def sparse_resize_(self, *sizes):
        assert len(sizes) == 2
rusty1s's avatar
rusty1s committed
194
        self._sparse_size = sizes
rusty1s's avatar
rusty1s committed
195
        return self
rusty1s's avatar
rusty1s committed
196

rusty1s's avatar
rusty1s committed
197
198
    @cached_property
    def rowcount(self):
rusty1s's avatar
rusty1s committed
199
        # TODO
rusty1s's avatar
rusty1s committed
200
201
202
        one = torch.ones_like(self.row)
        return segment_add(one, self.row, dim=0, dim_size=self._sparse_size[0])

rusty1s's avatar
rusty1s committed
203
204
    @cached_property
    def rowptr(self):
rusty1s's avatar
rusty1s committed
205
        # TODO
rusty1s's avatar
rusty1s committed
206
        rowcount = self.rowcount
rusty1s's avatar
rusty1s committed
207
208
209
        rowptr = rowcount.new_zeros(rowcount.numel() + 1)
        torch.cumsum(rowcount, dim=0, out=rowptr[1:])
        return rowptr
rusty1s's avatar
rusty1s committed
210
211
212

    @cached_property
    def colcount(self):
rusty1s's avatar
rusty1s committed
213
        # TODO
rusty1s's avatar
rusty1s committed
214
215
        one = torch.ones_like(self.col)
        return scatter_add(one, self.col, dim=0, dim_size=self._sparse_size[1])
rusty1s's avatar
rusty1s committed
216
217
218

    @cached_property
    def colptr(self):
rusty1s's avatar
rusty1s committed
219
        # TODO
rusty1s's avatar
rusty1s committed
220
        colcount = self.colcount
rusty1s's avatar
rusty1s committed
221
222
223
        colptr = colcount.new_zeros(colcount.numel() + 1)
        torch.cumsum(colcount, dim=0, out=colptr[1:])
        return colptr
rusty1s's avatar
rusty1s committed
224
225

    @cached_property
rusty1s's avatar
rusty1s committed
226
    def csr2csc(self):
rusty1s's avatar
rusty1s committed
227
228
229
230
        idx = self._sparse_size[0] * self.col + self.row
        return idx.argsort()

    @cached_property
rusty1s's avatar
rusty1s committed
231
232
    def csc2csr(self):
        return self.csr2csc.argsort()
rusty1s's avatar
rusty1s committed
233

rusty1s's avatar
rusty1s committed
234
    def is_coalesced(self):
rusty1s's avatar
rusty1s committed
235
        idx = self.sparse_size(1) * self.row + self.col
rusty1s's avatar
rusty1s committed
236
237
        mask = idx > torch.cat([idx.new_full((1, ), -1), idx[:-1]], dim=0)
        return mask.all().item()
rusty1s's avatar
rusty1s committed
238

rusty1s's avatar
rusty1s committed
239
240
241
242
243
244
245
246
247
248
249
250
    def coalesce(self, reduce='add'):
        idx = self.sparse_size(1) * self.row + self.col
        mask = idx > torch.cat([idx.new_full((1, ), -1), idx[:-1]], dim=0)

        if mask.all():  # Already coalesced
            return self

        index = self.index[:, mask]

        value = self.value
        if self.has_value():
            idx = mask.cumsum(0) - 1
rusty1s's avatar
rusty1s committed
251
252
            dim_size = idx[-1].item() + 1
            value = segment_csr(idx, value, dim_size=dim_size, reduce=reduce)
rusty1s's avatar
rusty1s committed
253
254
255
            value = value[0] if isinstance(value, tuple) else value

        return self.__class__(index, value, self.sparse_size(), is_sorted=True)
rusty1s's avatar
rusty1s committed
256

rusty1s's avatar
rusty1s committed
257
258
259
260
261
262
    def cached_keys(self):
        return [
            key for key in self.cache_keys
            if getattr(self, f'_{key}', None) is not None
        ]

rusty1s's avatar
rusty1s committed
263
    def fill_cache_(self, *args):
rusty1s's avatar
rusty1s committed
264
265
        for arg in args or self.cache_keys:
            getattr(self, arg)
rusty1s's avatar
rusty1s committed
266
        return self
rusty1s's avatar
rusty1s committed
267

rusty1s's avatar
rusty1s committed
268
269
270
271
    def clear_cache_(self, *args):
        for arg in args or self.cache_keys:
            setattr(self, f'_{arg}', None)
        return self
rusty1s's avatar
rusty1s committed
272

rusty1s's avatar
rusty1s committed
273
274
275
    def __copy__(self):
        return self.apply(lambda x: x)

rusty1s's avatar
test  
rusty1s committed
276
277
278
279
280
    def clone(self):
        return self.apply(lambda x: x.clone())

    def __deepcopy__(self, memo):
        new_storage = self.clone()
rusty1s's avatar
rusty1s committed
281
        memo[id(self)] = new_storage
rusty1s's avatar
test  
rusty1s committed
282
283
        return new_storage

rusty1s's avatar
rusty1s committed
284
285
    def apply_value_(self, func):
        self._value = optional(func, self._value)
rusty1s's avatar
rusty1s committed
286
        return self
rusty1s's avatar
rusty1s committed
287

rusty1s's avatar
rusty1s committed
288
289
290
291
292
    def apply_value(self, func):
        return self.__class__(
            self._index,
            optional(func, self._value),
            self._sparse_size,
rusty1s's avatar
rusty1s committed
293
            self._rowcount,
rusty1s's avatar
rusty1s committed
294
            self._rowptr,
rusty1s's avatar
rusty1s committed
295
            self._colcount,
rusty1s's avatar
rusty1s committed
296
            self._colptr,
rusty1s's avatar
rusty1s committed
297
298
            self._csr2csc,
            self._csc2csr,
rusty1s's avatar
rusty1s committed
299
300
301
302
303
304
            is_sorted=True,
        )

    def apply_(self, func):
        self._index = func(self._index)
        self._value = optional(func, self._value)
rusty1s's avatar
rusty1s committed
305
        for key in self.cached_keys():
rusty1s's avatar
rusty1s committed
306
            setattr(self, f'_{key}', func(getattr(self, f'_{key}')))
rusty1s's avatar
rusty1s committed
307
        return self
rusty1s's avatar
rusty1s committed
308
309
310
311
312
313

    def apply(self, func):
        return self.__class__(
            func(self._index),
            optional(func, self._value),
            self._sparse_size,
rusty1s's avatar
rusty1s committed
314
            optional(func, self._rowcount),
rusty1s's avatar
rusty1s committed
315
            optional(func, self._rowptr),
rusty1s's avatar
rusty1s committed
316
            optional(func, self._colcount),
rusty1s's avatar
rusty1s committed
317
            optional(func, self._colptr),
rusty1s's avatar
rusty1s committed
318
319
            optional(func, self._csr2csc),
            optional(func, self._csc2csr),
rusty1s's avatar
rusty1s committed
320
321
322
            is_sorted=True,
        )

rusty1s's avatar
rusty1s committed
323
324
325
326
    def map(self, func):
        data = [func(self.index)]
        if self.has_value():
            data += [func(self.value)]
rusty1s's avatar
rusty1s committed
327
        data += [func(getattr(self, f'_{key}')) for key in self.cached_keys()]
rusty1s's avatar
rusty1s committed
328
        return data