sympycore/matrices/matrix_operations.py

   1
   2 from ..utils import MATRIX, MATRIX_DICT
   3 from .algebra import Matrix, MatrixBase, MatrixDict
   4
   5 from ..core import init_module
   6 init_module.import_lowlevel_operations()
   7
   8 def MATRIX_DICT_iadd(self, other):
   9     """ Inplace matrix add.
  10     """
  11     t = type(other)
  12     if t is list or t is tuple:
  13         other = Matrix(other)
  14         t = type(other)
  15     if t is MatrixDict:
  16         if self.is_writable:
  17             ret = self
  18         else:
  19             ret = self.copy()
  20         head1, data1 = ret.pair
  21         head2, data2 = other.pair
  22         assert head1.shape==head2.shape,`head1, head2`
  23         if head1.is_transpose:
  24             if head2.is_transpose:
  25                 iadd_MATRIX_MATRIX_TT(data1, data2)
  26             else:
  27                 iadd_MATRIX_MATRIX_TA(data1, data2)
  28         elif head2.is_transpose:
  29             iadd_MATRIX_MATRIX_AT(data1, data2)
  30         else:
  31             iadd_MATRIX_MATRIX_AA(data1, data2)
  32         if head2.is_array:
  33             return ret.A
  34         return ret.M
  35     elif isinstance(other, MatrixBase):
  36         raise NotImplementedError(`type(other)`)
  37     else:
  38         if other:
  39             if self.is_writable:
  40                 ret = self
  41             else:
  42                 ret = self.copy()
  43             head, data = ret.pair
  44             rows, cols = head.shape
  45             if head.is_transpose:
  46                 iadd_MATRIX_T_SCALAR(rows, cols, data, other)
  47             else:
  48                 iadd_MATRIX_SCALAR(rows, cols, data, other)
  49             return ret
  50         else:
  51             return self
  52
  53 def MATRIX_DICT_imul(self, other):
  54     """ Inplace matrix multiplication.
  55     """
  56     t = type(other)
  57     if t is list or t is tuple:
  58         other = Matrix(other)
  59         t = type(other)
  60     if t is MatrixDict:
  61         head1, data1 = self.pair
  62         head2, data2 = other.pair
  63         if head1.is_array or head2.is_array:
  64             assert head1.shape==head2.shape,`head1, head2`
  65             if self.is_writable:
  66                 ret = self
  67             else:
  68                 ret = self.copy()
  69             data1 = ret.data
  70             if head1.is_transpose:
  71                 if head2.is_transpose:
  72                     imul_MATRIX_MATRIX_ATT(data1, data2)
  73                 else:
  74                     imul_MATRIX_MATRIX_TA(data1, data2)
  75             elif head2.is_transpose:
  76                 imul_MATRIX_MATRIX_AT(data1, data2)
  77             else:
  78                 imul_MATRIX_MATRIX_AA(data1, data2)
  79             if head2.is_array:
  80                 return ret.A
  81             return ret.M
  82         else:
  83             assert head1.cols==head2.rows,`head1, head2`
  84             args = data1, data2, head1.rows, head2.cols, head1.cols
  85             if head1.is_transpose:
  86                 if head2.is_transpose:
  87                     ret = mul_MATRIX_MATRIX_MTT(*args)
  88                 else:
  89                     ret = mul_MATRIX_MATRIX_TM(*args)
  90             elif head2.is_transpose:
  91                 ret = mul_MATRIX_MATRIX_MT(*args)
  92             else:
  93                 ret = mul_MATRIX_MATRIX_MM(*args)
  94             return ret
  95     elif isinstance(other, MatrixBase):
  96         raise NotImplementedError(`type(other)`)
  97     else:
  98         if self.is_writable:
  99             ret = self
 100         else:
 101             ret = self.copy()
 102         if other:
 103             head, data = ret.pair
 104             for key in data:
 105                 data[key] *= other
 106         else:
 107             head, data = ret.pair
 108             data.clear()
 109         return ret
 110
 111 def iadd_MATRIX_SCALAR(rows, cols, data, value):
 112     col_indices = range(cols)
 113     for i in xrange(rows):
 114         for j in col_indices:
 115             key = (i,j)
 116             b = data.get(key)
 117             if b is None:
 118                 data[key] = value
 119             else:
 120                 b += value
 121                 if b:
 122                     data[key] = b
 123                 else:
 124                     del data[key]
 125
 126 def iadd_MATRIX_T_SCALAR(rows, cols, data, value):
 127     col_indices = range(cols)
 128     for i in xrange(rows):
 129         for j in col_indices:
 130             key = (j,i)
 131             b = data.get(key)
 132             if b is None:
 133                 data[key] = value
 134             else:
 135                 b += value
 136                 if b:
 137                     data[key] = b
 138                 else:
 139                     del data[key]
 140
 141 def iadd_MATRIX_MATRIX_AA(data1, data2):
 142     for key,x in data2.items():
 143         b = data1.get(key)
 144         if b is None:
 145             data1[key] = x
 146         else:
 147             b += x
 148             if b:
 149                 data1[key] = b
 150             else:
 151                 del data1[key]
 152
 153 def iadd_MATRIX_MATRIX_AT(data1, data2):
 154     for (j,i),x in data2.items():
 155         key = i,j
 156         b = data1.get(key)
 157         if b is None:
 158             data1[key] = x
 159         else:
 160             b += x
 161             if b:
 162                 data1[key] = b
 163             else:
 164                 del data1[key]
 165
 166 def iadd_MATRIX_MATRIX_TA(data1, data2):
 167     for (i,j),x in data2.items():
 168         key = j,i
 169         b = data1.get(key)
 170         if b is None:
 171             data1[key] = x
 172         else:
 173             b += x
 174             if b:
 175                 data1[key] = b
 176             else:
 177                 del data1[key]
 178
 179 iadd_MATRIX_MATRIX_TT = iadd_MATRIX_MATRIX_AA
 180
 181 def imul_MATRIX_MATRIX_AA(data1, data2):
 182     for key in data1:
 183         b = data2.get(key)
 184         if b is None:
 185             del data1[key]
 186         else:
 187             data1[key] *= b
 188
 189 def imul_MATRIX_MATRIX_AT(data1, data2):
 190     for key in data1:
 191         i,j = key
 192         b = data2.get((j,i))
 193         if b is None:
 194             del data1[key]
 195         else:
 196             data1[key] *= b
 197
 198 imul_MATRIX_MATRIX_TA = imul_MATRIX_MATRIX_AT
 199 imul_MATRIX_MATRIX_ATT = imul_MATRIX_MATRIX_AA
 200
 201 def mul_MATRIX_MATRIX_AA(data1, data2, rows, cols):
 202     indices = range(n)
 203     d = {}
 204     data1_get = data1.get
 205     data2_get = data2.get
 206     for i in xrange(rows):
 207         for j in xrange(cols):
 208             key = i,j
 209             a_ij = data1_get(key)
 210             if a_ij is None:
 211                 continue
 212             b_ij = data2_get(key)
 213             if b_ij is None:
 214                 continue
 215             d[key] = a_ij * b_ij
 216     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 217
 218 def mul_MATRIX_MATRIX_AT(data1, data2, rows, cols):
 219     indices = range(n)
 220     d = {}
 221     data1_get = data1.get
 222     data2_get = data2.get
 223     for i in xrange(rows):
 224         for j in xrange(cols):
 225             key = i,j
 226             a_ij = data1_get(key)
 227             if a_ij is None:
 228                 continue
 229             b_ij = data2_get((j,i))
 230             if b_ij is None:
 231                 continue
 232             d[key] = a_ij * b_ij
 233     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 234
 235 def mul_MATRIX_MATRIX_TA(data1, data2, rows, cols):
 236     indices = range(n)
 237     d = {}
 238     data1_get = data1.get
 239     data2_get = data2.get
 240     for i in xrange(rows):
 241         for j in xrange(cols):
 242             key = i,j
 243             a_ij = data1_get((j,i))
 244             if a_ij is None:
 245                 continue
 246             b_ij = data2_get(key)
 247             if b_ij is None:
 248                 continue
 249             d[key] = a_ij * b_ij
 250     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 251
 252 def mul_MATRIX_MATRIX_ATT(data1, data2, rows, cols):
 253     indices = range(n)
 254     d = {}
 255     data1_get = data1.get
 256     data2_get = data2.get
 257     for i in xrange(rows):
 258         for j in xrange(cols):
 259             ikey = j,i
 260             a_ij = data1_get(ikey)
 261             if a_ij is None:
 262                 continue
 263             b_ij = data2_get(ikey)
 264             if b_ij is None:
 265                 continue
 266             d[i,j] = a_ij * b_ij
 267     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 268
 269 def mul_MATRIX_MATRIX_MM(data1, data2, rows, cols, n):
 270     d = {}
 271     left_indices = {}
 272     right_indices = {}
 273     for i,j in data1.keys():
 274         s = left_indices.get(i)
 275         if s is None:
 276             s = left_indices[i] = set()
 277         s.add(j)
 278     for j,k in data2.keys():
 279         s = right_indices.get(k)
 280         if s is None:
 281             s = right_indices[k] = set()
 282         s.add(j)
 283
 284     for i,ji in left_indices.items ():
 285         for k,jk in right_indices.items ():
 286             for j in ji.intersection(jk):
 287                 dict_add_item(None, d, (i,k), data1[(i,j)] * data2[(j,k)])
 288     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 289
 290 def mul_MATRIX_MATRIX_TM(data1, data2, rows, cols, n):
 291     d = {}
 292     left_indices = {}
 293     right_indices = {}
 294     for j,i in data1.keys():
 295         s = left_indices.get(i)
 296         if s is None:
 297             s = left_indices[i] = set()
 298         s.add(j)
 299     for j,k in data2.keys():
 300         s = right_indices.get(k)
 301         if s is None:
 302             s = right_indices[k] = set()
 303         s.add(j)
 304
 305     for i,ji in left_indices.items ():
 306         for k,jk in right_indices.items ():
 307             for j in ji.intersection(jk):
 308                 dict_add_item(None, d, (i,k), data1[(j,i)] * data2[(j,k)])
 309     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 310
 311 def mul_MATRIX_MATRIX_MT(data1, data2, rows, cols, n):
 312     d = {}
 313     left_indices = {}
 314     right_indices = {}
 315     for i,j in data1.keys():
 316         s = left_indices.get(i)
 317         if s is None:
 318             s = left_indices[i] = set()
 319         s.add(j)
 320     for k,j in data2.keys():
 321         s = right_indices.get(k)
 322         if s is None:
 323             s = right_indices[k] = set()
 324         s.add(j)
 325
 326     for i,ji in left_indices.items ():
 327         for k,jk in right_indices.items ():
 328             for j in ji.intersection(jk):
 329                 dict_add_item(None, d, (i,k), data1[(i,j)] * data2[(k,j)])
 330     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 331
 332 def mul_MATRIX_MATRIX_MTT(data1, data2, rows, cols, n):
 333     d = {}
 334     left_indices = {}
 335     right_indices = {}
 336     for j,i in data1.keys():
 337         s = left_indices.get(i)
 338         if s is None:
 339             s = left_indices[i] = set()
 340         s.add(j)
 341     for k,j in data2.keys():
 342         s = right_indices.get(k)
 343         if s is None:
 344             s = right_indices[k] = set()
 345         s.add(j)
 346
 347     for i,ji in left_indices.items ():
 348         for k,jk in right_indices.items ():
 349             for j in ji.intersection(jk):
 350                 dict_add_item(None, d, (i,k), data1[(j,i)] * data2[(k,j)])
 351     return MatrixDict(MATRIX(rows, cols, MATRIX_DICT), d)
 352