Source code for irfpy.util.ndsparse

''' Multidimensional sparse array.

Interface could be very similar to :class:`scipy:scipy.sparse.dok_matrix`.
However, this class is developed "on demand", so that the implementation
is not fully completed.

.. codeauthor:: Yoshifumi Futaana
'''
import itertools
import copy
import numpy as np

from irfpy.util.exception import PyanaError

[docs]class ndsparse_container: ''' N-dimensional sparse array. This class provide a container part. ''' def __init__(self, arg1): ''' Construct multi-dimensional sparse array. 1. From a dense array: >>> mat = ndsparse([[0, 0, 0, 0, 2], [0, 3, 0, 0, 1], [0, 1, -1, 0, 0]]) will create (3, 5) shaped array. >>> print(mat.shape) (3, 5) 2. From a sparse array >>> mat2 = ndsparse(mat) >>> print(mat2.shape) (3, 5) 3. By specifying the shape, if ``arg1`` is a tuple. >>> mat3 = ndsparse((5, 3, 2)) >>> print(mat3.shape) (5, 3, 2) ''' self.data = {} if isinstance(arg1, tuple): self.shape = copy.copy(arg1) self.ndim = len(self.shape) elif isinstance(arg1, ndsparse): self.data = copy.copy(arg1.sparsearray) self.shape = copy.copy(arg1.shape) self.ndim = arg1.ndim else: arr = np.array(arg1) self.shape = arr.shape self.ndim = arr.ndim idx = np.array(np.where(arr != 0)) # (array(...), array(...), ...) -> (ndim, ndat) array ndat = idx.shape[1] for idat in range(ndat): iidx = tuple(idx[:, idat]) self.data[iidx] = arr[iidx] if self.ndim <= 1: raise PyanaError('Use scipy.sparse module for 1-dimension purpose...')
[docs] def todense(self): ''' Convert to dense matrix. >>> mat = ndsparse_container([[0, 0, 0, 0, 2], [0, 3, 0, 0, 1], [0, 1, -1, 0, 0]]) >>> print(mat.todense()) [[ 0. 0. 0. 0. 2.] [ 0. 3. 0. 0. 1.] [ 0. 1. -1. 0. 0.]] ''' densearray = np.zeros(self.shape) for idx in self.data: densearray[idx] = self.data[idx] return densearray
def __getitem__(self, key): ''' Get item by index. >>> mat = ndsparse_container([[0, 0, 0, 0, 2], [0, 3, 0, 0, 1], [0, 1, -1, 0, 0]]) >>> print(mat.shape) (3, 5) >>> print(mat[0, 0]) 0 >>> print(mat[0, 4]) 2 >>> print(mat[2, 2]) -1 >>> try: ... print(mat[3, 2]) ... print('Error not raised') ... except IndexError as e: ... print('Error raised') Error raised >>> print(mat[1:3, 2:-1]) [[ 0 0 1] [-1 0 0]] .. note:: Due to itertools.islice limitation, negative step is not supported... ''' if len(key) != self.ndim: raise PyanaError('Dimension mismatch: Given={}/Mustbe={}'.format( len(key), self.ndim)) key_slice = [] for idim, kk in enumerate(key): if isinstance(kk, int): if kk < 0: k = self.shape[idim] - kk else: k = kk key_slice.append([i for i in itertools.islice(range(self.shape[idim]), k, k + 1)]) elif isinstance(kk, slice): if kk.start < 0: k0 = self.shape[idim] - kk.start else: k0 = kk.start if kk.stop < 0: k1 = self.shape[idim] - kk.stop else: k1 = kk.stop key_slice.append([i for i in itertools.islice(range(self.shape[idim]), k0, k1, kk.step)]) else: raise PyanaError('Whats up? Type of key={}'.format(type(key))) shape = np.array(self.shape, dtype=np.int) val = [] for idx in itertools.product(*key_slice): if idx in self.data: val.append(self.data[idx]) else: npkey = np.array(idx, dtype=np.int) if (npkey < 0).any() or (npkey >= shape).any(): raise IndexError('No index, {}'.format(key)) val.append(0) newshape = [len(k) for k in key_slice] if (np.array(newshape) == 0).any(): raise IndexError('No index, {}'.format(key)) if (np.array(newshape) == 1).all(): return np.array(val).flatten()[0] else: return np.array(val).reshape(newshape) # if key in self.data: # return self.data[key] # else: # shape = np.array(self.shape, dtype=np.int) # npkey = np.array(key, dtype=np.int) # if (npkey < 0).any() or (npkey >= shape).any(): # raise IndexError('No index, {}'.format(key)) # else: # return 0 def __setitem__(self, key, val): ''' Set item by index. >>> mat = ndsparse_container([[0, 0, 0, 0, 2], [0, 3, 0, 0, 1], [0, 1, -1, 0, 0]]) >>> print(mat.shape) (3, 5) >>> print(mat[1, 4]) 1 >>> mat[1, 4] = 3.5 >>> print(mat[1, 4]) 3.5 >>> mat[1, 4] = 0 >>> print(mat[1, 4]) 0 .. todo:: SLICE support. ''' if len(key) != self.ndim: raise PyanaError('Dimension mismatch: Given={}/Mustbe={}'.format( len(key), self.ndim)) shape = np.array(self.shape, dtype=np.int) npkey = np.array(key, dtype=np.int) if (npkey < 0).any() or (npkey >= shape).any(): raise IndexError('No index, {}'.format(key)) if val == 0: if key in self.data: self.data.pop(key) else: self.data[key] = val
[docs]class ndsparse: ''' N-dimensional sparse array. ''' def __init__(self, arg1, **argv): ''' Create a sparse array container. See :class:`ndsparse_container`. ''' self.sparsearray = ndsparse_container(arg1, **argv) self.shape = self.sparsearray.shape self.ndim = self.sparsearray.ndim def __getitem__(self, key): ''' Return the stored value. ''' return self.sparsearray[key] def __setitem__(self, key, value): self.sparsearray[key] = value
[docs] def todense(self): return self.sparsearray.todense()
[docs] def sum(self, axis=None): raise NotImplementedError('')
import unittest import doctest
[docs]def doctests(): return unittest.TestSuite(( doctest.DocTestSuite(), ))
if __name__ == '__main__': unittest.main(defaultTest='doctests')