linalg.py (13057B)
1 #Embedded file name: C:\BuildAgent\work\20d73f1b24eed06e\eve\release\V24.01\carbon\common\script\util\linalg.py 2 from future.utils import raise_ 3 import math 4 import types 5 import sys 6 TINY = 0.0 7 8 class VectorOP(object): 9 10 def __Coarce(self, o): 11 if hasattr(o, '__len__') or hasattr(o, '__iter__'): 12 return o 13 return [ o for i in xrange(len(self)) ] 14 15 def Set(self, other): 16 self[:] = other 17 18 def __iadd__(self, other): 19 other = self.__Coarce(other) 20 for i in xrange(len(self)): 21 self[i] += other[i] 22 23 return self 24 25 def __isub__(self, other): 26 other = self.__Coarce(other) 27 for i in xrange(len(self)): 28 self[i] -= other[i] 29 30 return self 31 32 def __imul__(self, other): 33 other = self.__Coarce(other) 34 for i in xrange(len(self)): 35 self[i] *= other[i] 36 37 return self 38 39 def __idiv__(self, other): 40 other = self.__Coarce(other) 41 for i in xrange(len(self)): 42 self[i] /= other[i] 43 44 return self 45 46 def Negate(self): 47 for i in xrange(len(self)): 48 self[i] = -self[i] 49 50 def Dot(self, other): 51 r = self[0] * other[0] 52 for i in range(1, len(self)): 53 r += self[i] * other[i] 54 55 return r 56 57 def Len2(self): 58 return self.Dot(self) 59 60 def Len(self): 61 return math.sqrt(self.Len2()) 62 63 def __abs__(self): 64 return self.Len() 65 66 def Normalize(self): 67 self /= self.Len() 68 69 70 def Mixin(current, cls): 71 current.update(cls.__dict__) 72 73 74 class Vector(list): 75 Mixin(locals(), VectorOP) 76 77 def __init__(self, a = None, b = 0): 78 try: 79 list.__init__(self, [ b for i in range(a) ]) 80 except TypeError: 81 list.__init__(self, a) 82 sys.exc_clear() 83 84 def Ones(length, item = 1.0): 85 return Vector(length, item) 86 87 Ones = staticmethod(Ones) 88 89 def Copy(self): 90 return Vector(self) 91 92 def LMatrix(self): 93 return Matrix((1, len(self)), self) 94 95 def CMatrix(self): 96 return Matrix((len(self), 1), self) 97 98 def __MatchSize(self, other): 99 if len(self) != len(other): 100 raise TypeError('operands have mismatching sizes') 101 102 def __Vectorize(self, obj): 103 if not isinstance(obj, Vector): 104 return Vector(len(self), obj) 105 return obj 106 107 def __add__(self, op): 108 r = Vector(self) 109 r += op 110 return r 111 112 def __radd__(self, op): 113 return self + op 114 115 def __sub__(self, op): 116 r = Vector(self) 117 r -= op 118 return r 119 120 def __rsub__(self, op): 121 return self.__Vectorize(op) - self 122 123 def __mul__(self, op): 124 r = Vector(self) 125 r *= op 126 return r 127 128 def __rmul__(self, op): 129 return self.__Vectorize(op) * self 130 131 def __div__(self, op): 132 r = Vector(self) 133 r /= op 134 return r 135 136 def __rdiv__(self, other): 137 return self.__Vectorize(other) / self 138 139 def __neg__(self): 140 return Vector([ -each for each in self ]) 141 142 def Cross(self, other): 143 return Vector([self[1] * other[2] - other[1] * self[2], -(self[0] * other[2] - other[0] * self[2]), self[0] * other[1] - other[0] * self[1]]) 144 145 def Normal(self): 146 return self / abs(self) 147 148 149 class Matrix(Vector): 150 __slots__ = 'dim' 151 152 def __init__(self, a, b = None): 153 if isinstance(a, Matrix) and b is None: 154 super(Matrix, self).__init__(a) 155 self.dim = a.dim 156 return 157 self.dim = a 158 l = a[0] * a[1] 159 if b is None: 160 super(Matrix, self).__init__(l) 161 else: 162 if len(b) != l: 163 raise ValueError('operand has incorrect length') 164 super(Matrix, self).__init__(b) 165 166 def __getitem__(self, idx): 167 if isinstance(idx, types.TupleType): 168 return super(Matrix, self).__getitem__(idx[0] * self.dim[1] + idx[1]) 169 return super(Matrix, self).__getitem__(idx) 170 171 def __setitem__(self, idx, val): 172 if isinstance(idx, types.TupleType): 173 super(Matrix, self).__setitem__(idx[0] * self.dim[1] + idx[1], val) 174 else: 175 super(Matrix, self).__setitem__(idx, val) 176 177 def Ones(dim, item = 1.0): 178 return Matrix(dim, [item] * dim[0] * dim[1]) 179 180 Ones = staticmethod(Ones) 181 182 def I(n, item = 1.0): 183 tmp = [item] + ([0.0] * n + [item]) * (n - 1) 184 return Matrix((n, n), tmp) 185 186 I = staticmethod(I) 187 188 def FromLines(lines): 189 l = len(lines) 190 c = len(lines[0]) 191 v = [] 192 for line in lines: 193 v.extend(line) 194 195 return Matrix((l, c), v) 196 197 FromLines = staticmethod(FromLines) 198 199 def FromCols(cols): 200 c = len(cols) 201 l = len(cols[0]) 202 v = [] 203 for i in range(l): 204 v.extend([ cols[j][i] for j in range(c) ]) 205 206 return Matrix((l, c), v) 207 208 FromCols = staticmethod(FromCols) 209 210 def __repr__(self): 211 return '{Matrix %s %s}' % (self.dim, super(Matrix, self).__repr__()) 212 213 def __str__(self): 214 return '\n'.join([ str(self.Line(l)) for l in range(self.dim[0]) ]) 215 216 def Copy(self): 217 return Matrix(self) 218 219 def Set(self, other): 220 self[:] = other 221 self.dim = other.dim 222 223 def Dim(self): 224 return self.dim 225 226 def __Matrixize(self, other): 227 if isinstance(other, Matrix): 228 return other 229 if isinstance(other, Vector): 230 return other.LMatrix() 231 return self.Ones(self.dim, other) 232 233 def __add__(self, other): 234 r = self.Copy() 235 r += other 236 return r 237 238 def __radd__(self, other): 239 return self.__Matrixize(other) + self 240 241 def __sub__(self, other): 242 r = self.Copy() 243 r -= other 244 return r 245 246 def __rsub__(self, other): 247 return self.__Matrixize(other) - self 248 249 def __neg__(self): 250 r = self.Copy() 251 r.Negate() 252 return r 253 254 def __mul__(self, other): 255 if not isinstance(other, Matrix): 256 if not isinstance(other, Vector): 257 return Matrix(self.dim, super(Matrix, self).__mul__(other)) 258 other = other.CMatrix() 259 d = self.dim[1] 260 if d != other.dim[0]: 261 raise_(ValueError, 'invalid operand (dim%s * dim%s)' % (self.dim, other.dim)) 262 nl, nc = self.dim[0], other.dim[1] 263 dim = (self.dim[0], other.dim[1]) 264 r = Matrix((nl, nc)) 265 nc = range(nc) 266 d = range(1, d) 267 for l in xrange(nl): 268 for c in nc: 269 v = self[l, 0] * other[0, c] 270 for i in d: 271 v += self[l, i] * other[i, c] 272 273 r[l, c] = v 274 275 return r 276 for l in xrange(nl): 277 line = Line(self, l) 278 for c in xrange(nc): 279 r[l, c] = line.Dot(Column(other, c)) 280 281 def __imul__(self, other): 282 self.Set(self * other) 283 return self 284 285 def __rmul__(self, other): 286 return self.__Matrixize(other) * self 287 288 def __div__(self, other): 289 if isinstance(other, Matrix): 290 if other.dim[0] < 3: 291 return self * other.Inv() 292 return other.Trans().Solve(self.Trans()).Trans() 293 super(Matrix, self).__div__(other) 294 295 def __idiv__(self, other): 296 self.Set(self / other) 297 return self 298 299 def __rdiv__(self, other): 300 return self.__Matrixize(other) / self 301 302 def Line(self, l): 303 return Vector(self[l * self.dim[1]:(l + 1) * self.dim[1]]) 304 305 def Lines(self): 306 return [ self.Line(l) for l in range(self.dim[0]) ] 307 308 def Col(self, c): 309 v = [ self[i] for i in range(c, len(self), self.dim[1]) ] 310 return Vector(v) 311 312 def Cols(self): 313 return [ self.Col(c) for c in range(self.dim[1]) ] 314 315 def Sub(self, start, size): 316 end = (start[0] + size[0], start[1] + size[1]) 317 v = [] 318 for l in range(start[0], end[0]): 319 v.extend(self.Line(l)[start[1]:end[1]]) 320 321 return Matrix(size, v) 322 323 def JoinH(self, other): 324 r = [] 325 for l in range(self.dim[0]): 326 a = self.Line(l) 327 a.extend(other.Line(l)) 328 r.extend(a) 329 330 return Matrix((self.dim[0], self.dim[1] + other.dim[1]), r) 331 332 def JoinV(self, other): 333 a = self.Copy() 334 a.extend(other) 335 a.dim[0] += other.dim[0] 336 return a 337 338 def Det(self): 339 dim = self.dim 340 if dim[0] == 1: 341 return self[0, 0] 342 f = 1 343 s = 0 344 for c in range(self.dim[1]): 345 a = self.Sub((1, 0), (dim[0] - 1, c)) 346 b = self.Sub((1, c + 1), (dim[0] - 1, dim[1] - c - 1)) 347 subm = a.JoinH(b) 348 s += f * self[0, c] * subm.Det() 349 f *= -1 350 351 return s 352 353 def Det2(self): 354 a = self.Copy() 355 idx, even = a.LUDecompose() 356 d = 1.0 357 for i in range(a.dim[0]): 358 d *= a[i, i] 359 360 if not even: 361 d *= -1.0 362 return d 363 364 def Transpose(self): 365 l, c = self.dim[1], self.dim[0] 366 if l < 2 or c < 2: 367 return Matrix((l, c), self) 368 r = Matrix((l, c)) 369 for i in range(l * c): 370 r[i % l * c + i / l] = self[i] 371 372 return r 373 374 Trans = Transpose 375 376 def Solve(self, other): 377 dc = self.Copy() 378 index, even = dc.LUDecompose() 379 if isinstance(other, Matrix): 380 cols = other.Cols() 381 for col in cols: 382 dc.LUBacksubstitute(index, col) 383 384 return self.FromCols(cols) 385 else: 386 v = Vector(other) 387 dc.LUBacksubstitute(index, v) 388 return v 389 390 def Inverse(self): 391 if self.dim == (1, 1): 392 det = self[0, 0] 393 if not det: 394 raise_(ValueError, 'singular matrix (%s)' % self) 395 return Matrix((1, 1), [1.0 / self[0, 0]]) 396 if self.dim == (2, 2): 397 det = self[0, 0] * self[1, 1] - self[1, 0] * self[0, 1] 398 if not det: 399 raise_(ValueError, 'singular matrix (%s)' % self) 400 return Matrix((2, 2), [self[1, 1] / det, 401 -self[0, 1] / det, 402 -self[1, 0] / det, 403 self[0, 0] / det]) 404 return self.Solve(self.I(self.dim[0])) 405 406 Inv = Inverse 407 408 def Invert(self): 409 self.Set(self.Inverse()) 410 411 def LUDecompose(self): 412 n = self.dim[0] 413 if n != self.dim[1]: 414 raise ValueError('Matrix must be square') 415 vv = [] 416 even = True 417 for i in range(n): 418 big = 0.0 419 for j in range(n): 420 temp = abs(self[i, j]) 421 if temp > big: 422 big = temp 423 424 if not big: 425 raise_(ValueError, 'singular matrix (%s)' % self) 426 vv.append(1.0 / big) 427 428 indx = [] 429 for j in range(n): 430 for i in range(j): 431 sum = self[i, j] 432 for k in range(i): 433 sum -= self[i, k] * self[k, j] 434 435 self[i, j] = sum 436 437 big = 0.0 438 for i in range(j, n): 439 sum = self[i, j] 440 for k in range(j): 441 sum -= self[i, k] * self[k, j] 442 443 self[i, j] = sum 444 dum = vv[i] * abs(sum) 445 if dum >= big: 446 big = dum 447 imax = i 448 449 if j != imax: 450 for k in range(n): 451 self[imax, k], self[j, k] = self[j, k], self[imax, k] 452 453 even = not even 454 vv[imax] = vv[j] 455 indx.append(imax) 456 if not self[j, j]: 457 raise ValueError('singular matrix') 458 self[j, j] = TINY 459 if j != n: 460 dum = 1.0 / self[j, j] 461 for i in range(j + 1, n): 462 self[i, j] *= dum 463 464 return (indx, even) 465 466 def LUBacksubstitute(self, indx, b): 467 ii = None 468 n = self.dim[0] 469 for i in range(n): 470 ip = indx[i] 471 sum, b[ip] = b[ip], b[i] 472 if ii is not None: 473 for j in range(ii, i): 474 sum -= self[i, j] * b[j] 475 476 elif sum: 477 ii = i 478 b[i] = sum 479 480 for i in range(n - 1, -1, -1): 481 sum = b[i] 482 for j in range(i + 1, n): 483 sum -= self[i, j] * b[j] 484 485 b[i] = float(sum) / self[i, i] 486 487 488 class Column(VectorOP): 489 __slots__ = ['m', 'c'] 490 491 def __init__(self, m, c): 492 self.m, self.c = m, c 493 494 def __getitem__(self, i): 495 return self.m[i, self.c] 496 497 def __setitem__(self, i, v): 498 self.m[i, self.c] = v 499 500 def __len__(self): 501 return self.m.dim[0] 502 503 504 class Line(VectorOP): 505 __slots__ = ['m', 'l'] 506 507 def __init__(self, m, l): 508 self.m, self.l = m, l 509 510 def __getitem__(self, i): 511 return self.m[self.l, i] 512 513 def __setitem__(self, i, v): 514 self.m[self.l, i] = v 515 516 def __len__(self): 517 return self.m.dim[1]