test/experimental/solve.py

   1 #!/usr/bin/env python
   2 # -*- coding: ISO-8859-1 -*-
   3 #
   4 #
   5 # Copyright (C) 2004 André Wobst <wobsta@users.sourceforge.net>
   6 #
   7 # This file is part of PyX (http://pyx.sourceforge.net/).
   8 #
   9 # PyX is free software; you can redistribute it and/or modify
  10 # it under the terms of the GNU General Public License as published by
  11 # the Free Software Foundation; either version 2 of the License, or
  12 # (at your option) any later version.
  13 #
  14 # PyX is distributed in the hope that it will be useful,
  15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  17 # GNU General Public License for more details.
  18 #
  19 # You should have received a copy of the GNU General Public License
  20 # along with PyX; if not, write to the Free Software
  21 # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  22
  23
  24 import Numeric, LinearAlgebra
  25
  26
  27 def sum(list):
  28     # we can assume len(list) != 0 here (and do not start from the scalar 0)
  29     sum = list[0]
  30     for item in list[1:]:
  31         sum += item
  32     return sum
  33
  34 def product(list):
  35     # we can assume len(list) != 0 here (and do not start from the scalar 1)
  36     product = list[0]
  37     for item in list[1:]:
  38         product *= item
  39     return product
  40
  41
  42 class scalar:
  43     # represents a scalar variable or constant
  44
  45     def __init__(self, value=None, name="unnamed_scalar"):
  46         self._scalar = None
  47         if value is not None:
  48             self.set(value)
  49         self.name = name
  50
  51     def scalar(self):
  52         return self
  53
  54     def addend(self):
  55         return addend([self])
  56
  57     def polynom(self):
  58         return self.addend().polynom()
  59
  60     def __neg__(self):
  61         return -self.addend()
  62
  63     def __add__(self, other):
  64         return self.polynom() + other
  65
  66     __radd__ = __add__
  67
  68     def __sub__(self, other):
  69         return self.polynom() - other
  70
  71     def __rsub__(self, other):
  72         return -self.polynom() + other
  73
  74     def __mul__(self, other):
  75         return self.addend()*other
  76
  77     __rmul__ = __mul__
  78
  79     def __div__(self, other):
  80         return self.addend()/other
  81
  82     def is_set(self):
  83         return self._scalar is not None
  84
  85     def set(self, value):
  86         if self.is_set():
  87             raise RuntimeError("scalar already defined")
  88         try:
  89             self._scalar = float(value)
  90         except:
  91             raise RuntimeError("float expected")
  92
  93     def get(self):
  94         if not self.is_set():
  95             raise RuntimeError("scalar not yet defined")
  96         return self._scalar
  97
  98     def __float__(self):
  99         return self.get()
 100
 101     def __str__(self):
 102         if self.is_set():
 103             return "%s{=%s}" % (self.name, self._scalar)
 104         else:
 105             return self.name
 106
 107
 108 class addend:
 109     # represents a addend, i.e. list of scalars to be multiplied by each other
 110
 111     def __init__(self, scalars):
 112         self._scalars = [scalar.scalar() for scalar in scalars]
 113         if not len(self._scalars):
 114             raise RuntimeError("empty scalars not allowed")
 115
 116     def addend(self):
 117         return self
 118
 119     def polynom(self):
 120         return polynom([self])
 121
 122     def __neg__(self):
 123         return addend([scalar(-1)] + self._scalars)
 124
 125     def __add__(self, other):
 126         return self.polynom() + other
 127
 128     __radd__ = __add__
 129
 130     def __sub__(self, other):
 131         return self.polynom() - other
 132
 133     def __rsub__(self, other):
 134         return -self.polynom() + other
 135
 136     def __mul__(self, other):
 137         try:
 138             other = other.addend()
 139         except (TypeError, AttributeError):
 140             try:
 141                 other = scalar(other)
 142             except RuntimeError:
 143                 return other * self
 144             else:
 145                 return addend(self._scalars + [other])
 146         else:
 147             return addend(self._scalars + other._scalars)
 148
 149     __rmul__ = __mul__
 150
 151     def __div__(self, other):
 152         return addend([scalar(1/other)] + self._scalars)
 153
 154     def __float__(self):
 155         return product([float(scalar) for scalar in self._scalars])
 156
 157     def is_linear(self):
 158         return len([scalar for scalar in self._scalars if not scalar.is_set()]) < 2
 159
 160     def prefactor(self):
 161         assert self.is_linear()
 162         setscalars = [scalar for scalar in self._scalars if scalar.is_set()]
 163         if len(setscalars):
 164             return float(addend(setscalars))
 165         else:
 166             return 1
 167
 168     def variable(self):
 169         assert self.is_linear()
 170         unsetscalars = [scalar for scalar in self._scalars if not scalar.is_set()]
 171         if len(unsetscalars):
 172             assert len(unsetscalars) == 1
 173             return unsetscalars[0]
 174         else:
 175             return None
 176
 177     def __str__(self):
 178         return " * ".join([str(scalar) for scalar in self._scalars])
 179
 180
 181 class polynom:
 182     # represents a polynom, i.e. a list of addends to be summed up
 183
 184     def __init__(self, polynom):
 185         self._addends = [addend.addend() for addend in polynom]
 186         if not len(self._addends):
 187             raise RuntimeError("empty polynom not allowed")
 188
 189     def polynom(self):
 190         return self
 191
 192     def __neg__(self):
 193         return polynom([-addend for addend in self._addends])
 194
 195     def __add__(self, other):
 196         try:
 197             other = other.polynom()
 198         except (TypeError, AttributeError):
 199             other = scalar(other).polynom()
 200         return polynom(self._addends + other._addends)
 201
 202     __radd__ = __add__
 203
 204     def __sub__(self, other):
 205         return -other + self
 206
 207     def __rsub__(self, other):
 208         return -self + other
 209
 210     def __mul__(self, other):
 211         return sum([addend*other for addend in self._addends])
 212
 213     __rmul__ = __mul__
 214
 215     def __div__(self, other):
 216         return polynom([addend/other for addend in self._addends])
 217
 218     def __float__(self):
 219         return sum([float(addend) for addend in self._addends])
 220
 221     def is_linear(self):
 222         is_linear = 1
 223         for addend in self._addends:
 224             is_linear = is_linear and addend.is_linear()
 225         return is_linear
 226
 227     def __str__(self):
 228         return "  +  ".join([str(addend) for addend in self._addends])
 229
 230     def solve(self, solver):
 231         solver.addequation(self)
 232
 233
 234 class vector:
 235     # represents a vector, i.e. a list of scalars (or polynoms)
 236
 237     def __init__(self, dimension_or_values, name="unnamed_vector"):
 238         try:
 239             name + ""
 240         except:
 241             raise RuntimeError("a vectors name should be a string (you probably wanted to write vector([x, y]) instead of vector(x, y))")
 242         try:
 243             for value in dimension_or_values:
 244                 pass
 245         except:
 246             # dimension
 247             self._items = [scalar(name="%s[%i]" % (name, i))
 248                            for i in range(dimension_or_values)]
 249         else:
 250             # values
 251             self._items = []
 252             for value in dimension_or_values:
 253                 try:
 254                     value.polynom()
 255                 except (TypeError, AttributeError):
 256                     self._items.append(scalar(value=value, name="%s[%i]" % (name, len(self._items))))
 257                 else:
 258                     self._items.append(value)
 259         if not len(self._items):
 260             raise RuntimeError("empty vector not allowed")
 261         self.name = name
 262
 263     def __len__(self):
 264         return len(self._items)
 265
 266     def __getitem__(self, i):
 267         return self._items[i]
 268
 269     def __getattr__(self, attr):
 270         if attr == "x":
 271             return self[0]
 272         if attr == "y":
 273             return self[1]
 274         if attr == "z":
 275             return self[2]
 276         else:
 277             raise AttributeError(attr)
 278
 279     def vector(self):
 280         return self
 281
 282     def __neg__(self):
 283         return vector([-item for item in self._items])
 284
 285     def __add__(self, other):
 286         other = other.vector()
 287         if len(self) != len(other):
 288             raise RuntimeError("vector length mismatch in add")
 289         return vector([selfitem + otheritem for selfitem, otheritem in zip(self._items, other._items)])
 290
 291     __radd__ = __add__
 292
 293     def __sub__(self, other):
 294         return -other + self
 295
 296     def __rsub__(self, other):
 297         return -self + other
 298
 299     def __mul__(self, other):
 300         try:
 301             other = other.vector()
 302         except (TypeError, AttributeError):
 303             return vector([item*other for item in self._items])
 304         else:
 305             # scalar product
 306             if len(self) != len(other):
 307                 raise RuntimeError("vector length mismatch in scalar product")
 308             return sum([selfitem*otheritem for selfitem, otheritem in zip(self._items, other._items)])
 309
 310     def __rmul__(self, other):
 311         return vector([other*item for item in self._items])
 312         # We do not allow for vector * <matrix-like-object> here (i.e. a
 313         # simulating the behaviour of a dual vector). In principle we could do
 314         # that, but for transformations it would lead to (a*X)*c != a*(X*c).
 315         # Hence we disallow vector*X for X being something else that a scalar.
 316
 317     def __div__(self, other):
 318         return vector([item/other for item in self._items])
 319
 320     def __str__(self):
 321         return "%s{=(%s)}" % (self.name, ", ".join([str(item) for item in self._items]))
 322
 323     def solve(self, solver):
 324         for item in self._items:
 325             solver.addequation(item)
 326
 327
 328 class zerovector(vector):
 329
 330     def __init__(self, dimension, name="0"):
 331         vector.__init__(self, [0 for i in range(dimension)], name)
 332
 333
 334 class matrix:
 335     # represents a matrix, i.e. a 2d list of scalars (or polynoms)
 336
 337     def __init__(self, dimensions_or_values, name="unnamed_matrix"):
 338         try:
 339             name + ""
 340         except:
 341             raise RuntimeError("a matrix name should be a string (you probably wanted to write matrix([x, y]) instead of matrix(x, y))")
 342         try:
 343             for row in dimensions_or_values:
 344                 for col in row:
 345                     pass
 346         except:
 347             # dimension
 348             self._numberofrows, self._numberofcols = [int(x) for x in dimensions_or_values]
 349             self._rows = [[scalar(name="%s[%i, %i]" % (name, row, col))
 350                            for col in range(self._numberofcols)]
 351                           for row in range(self._numberofrows)]
 352         else:
 353             # values
 354             self._rows = []
 355             self._numberofcols = None
 356             for row in dimensions_or_values:
 357                 _cols = []
 358                 for col in row:
 359                     try:
 360                         col.polynom()
 361                     except (TypeError, AttributeError):
 362                         _cols.append(scalar(value=col, name="%s[%i, %i]" % (name, len(self._rows), len(_cols))))
 363                     else:
 364                         _cols.append(col)
 365                 self._rows.append(_cols)
 366                 if self._numberofcols is None:
 367                     self._numberofcols = len(_cols)
 368                 elif self._numberofcols != len(_cols):
 369                     raise RuntimeError("column length mismatch")
 370             self._numberofrows = len(self._rows)
 371         if not self._numberofrows or not self._numberofcols:
 372             raise RuntimeError("empty matrix not allowed")
 373         self.name = name
 374
 375     # instead of __len__ two methods to fetch the matrix dimensions
 376     def getnumberofrows(self):
 377         return self._numberofrows
 378
 379     def getnumberofcols(self):
 380         return self._numberofcols
 381
 382     def __getitem__(self, (row, col)):
 383         return self._rows[row][col]
 384
 385     def matrix(self):
 386         return self
 387
 388     def __neg__(self):
 389         return matrix([[-col for col in row] for row in self._rows])
 390
 391     def __add__(self, other):
 392         other = other.matrix()
 393         if self._numberofrows != other._numberofrows or self._numberofcols != other._numberofcols:
 394             raise RuntimeError("matrix geometry mismatch in add")
 395         return matrix([[selfcol + othercol
 396                         for selfcol, othercol in zip(selfrow, otherrow)]
 397                        for selfrow, otherrow in zip(self._rows, other._rows)])
 398
 399     __radd__ = __add__
 400
 401     def __sub__(self, other):
 402         return -other + self
 403
 404     def __rsub__(self, other):
 405         return -self + other
 406
 407     def __mul__(self, other):
 408         try:
 409             other = other.matrix()
 410         except (TypeError, AttributeError):
 411             try:
 412                 other = other.vector()
 413             except (TypeError, AttributeError):
 414                 return matrix([[col*other for col in row] for row in self._rows])
 415             else:
 416                 if self._numberofcols != len(other):
 417                     raise RuntimeError("size mismatch in matrix vector product")
 418                 return vector([sum([col*otheritem
 419                                     for col, otheritem in zip(row, other)])
 420                                for row in self._rows])
 421         else:
 422             if self._numberofcols != other._numberofrows:
 423                 raise RuntimeError("size mismatch in matrix product")
 424             return matrix([[sum([self._rows[row][i]*other._rows[i][col]
 425                                  for i in range(self._numberofcols)])
 426                             for col in range(other._numberofcols)]
 427                            for row in range(self._numberofrows)])
 428
 429     def __rmul__(self, other):
 430         try:
 431             other = other.vector()
 432         except (TypeError, AttributeError):
 433             return matrix([[other*col for col in row] for row in self._rows])
 434         else:
 435             if self._numberofrows != len(other):
 436                 raise RuntimeError("size mismatch in matrix vector product")
 437             return vector([sum([other[i]*self._rows[i][col]
 438                                 for i in range(self._numberofrows)])
 439                            for col in range(self._numberofcols)])
 440
 441     def __div__(self, other):
 442         return matrix([[col/other for col in row] for row in self._rows])
 443
 444     def __str__(self):
 445         return "%s{=(%s)}" % (self.name, ", ".join(["(" + ", ".join([str(col) for col in row]) + ")" for row in self._rows]))
 446
 447     def solve(self, solver):
 448         for row in self._rows:
 449             for col in row:
 450                 solver.addequation(col)
 451
 452
 453 class identitymatrix(matrix):
 454
 455     def __init__(self, dimension, name="I"):
 456         def eq(row, col):
 457             if row == col:
 458                 return 1
 459             else:
 460                 return 0
 461         matrix.__init__(self, [[eq(row, col) for col in range(dimension)] for row in range(dimension)], name)
 462
 463
 464 class trafo:
 465     # represents a transformation, i.e. matrix and a constant vector
 466
 467     def __init__(self, dimensions_or_values, name="unnamed_trafo"):
 468         try:
 469             name + ""
 470         except:
 471             raise RuntimeError("a trafo name should be a string (you probably wanted to write trafo([x, y]) instead of trafo(x, y))")
 472         if len(dimensions_or_values) != 2:
 473             raise RuntimeError("first parameter of a trafo must contain two elements: either two dimensions or a matrix and a vector")
 474         try:
 475             numberofrows, numberofcols = [int(x) for x in dimensions_or_values]
 476         except:
 477             self._matrix = dimensions_or_values[0].matrix()
 478             self._vector = dimensions_or_values[1].vector()
 479             if self._matrix.getnumberofrows() != len(self._vector):
 480                 raise RuntimeError("size mismatch between matrix and vector")
 481         else:
 482             self._matrix = matrix((numberofrows, numberofcols), name=name + "_matrix")
 483             self._vector = vector(numberofrows, name=name + "_vector")
 484         self.name = name
 485
 486     def trafo(self):
 487         return self
 488
 489     def getmatrix(self):
 490         return self._matrix
 491
 492     def getvector(self):
 493         return self._vector
 494
 495     def __neg__(self):
 496         return trafo((-self._matrix, -self._vector))
 497
 498     def __add__(self, other):
 499         other = other.trafo()
 500         return trafo((self._matrix + other._matrix, self._vector + other._vector))
 501
 502     __radd__ = __add__
 503
 504     def __sub__(self, other):
 505         return -other + self
 506
 507     def __rsub__(self, other):
 508         return -self + other
 509
 510     def __mul__(self, other):
 511         try:
 512             other = other.trafo()
 513         except (TypeError, AttributeError):
 514             try:
 515                 other = other.vector()
 516             except (TypeError, AttributeError):
 517                 return trafo((self._matrix * other, self._vector * other))
 518             else:
 519                 return self._matrix * other + self._vector
 520         else:
 521             return trafo((self._matrix * other._matrix, self._vector + self._matrix * other._vector))
 522
 523     def __rmul__(self, other):
 524         return trafo((other * self._matrix, other * self._vector))
 525
 526     def __div__(self, other):
 527         return trafo((self._matrix/other, self._vector/other))
 528
 529     def __str__(self):
 530         return "%s{=(matrix: %s, vector: %s)}" % (self.name, self._matrix, self._vector)
 531
 532     def solve(self, solver):
 533         self._matrix.solve(solver)
 534         self._vector.solve(solver)
 535
 536
 537 class identitytrafo(trafo):
 538
 539     def __init__(self, dimension, name="I"):
 540         trafo.__init__(self, (identitymatrix(dimension, name=name),
 541                               zerovector(dimension, name=name)), name)
 542
 543
 544 class Solver:
 545     # linear equation solver
 546
 547     def __init__(self):
 548         self.eqs = [] # scalar equations not yet solved (a equation is a polynom to be zero)
 549
 550     def eq(self, lhs, rhs=None):
 551         if rhs is None:
 552             eq = lhs
 553         else:
 554             eq = lhs - rhs
 555         eq.solve(self)
 556
 557     def addequation(self, equation):
 558         # the equation is just a polynom which should be zero
 559         self.eqs.append(equation.polynom())
 560
 561         # try to solve some combinations of linear equations
 562         while 1:
 563             for eqs in self.combine(self.eqs):
 564                 if self.solve(eqs):
 565                     break # restart for loop
 566             else:
 567                 break # quit while loop
 568
 569     def combine(self, eqs):
 570         # create combinations of linear equations
 571         if not len(eqs):
 572             yield []
 573         else:
 574             for x in self.combine(eqs[1:]):
 575                 yield x
 576             if eqs[0].is_linear():
 577                 for x in self.combine(eqs[1:]):
 578                     yield [eqs[0]] + x
 579
 580     def solve(self, eqs):
 581         # try to solve a set of linear equations
 582         l = len(eqs)
 583         if l:
 584             vars = []
 585             for eq in eqs:
 586                 for addend in eq._addends:
 587                     var = addend.variable()
 588                     if var is not None and var not in vars:
 589                         vars.append(var)
 590             if len(vars) == l:
 591                 a = Numeric.zeros((l, l), Numeric.Float)
 592                 b = Numeric.zeros((l, ), Numeric.Float)
 593                 for i, eq in enumerate(eqs):
 594                     for addend in eq._addends:
 595                         var = addend.variable()
 596                         if var is not None:
 597                             a[i, vars.index(var)] += addend.prefactor()
 598                         else:
 599                             b[i] -= addend.prefactor()
 600                 for i, value in enumerate(LinearAlgebra.solve_linear_equations(a, b)):
 601                     vars[i].set(value)
 602                 for eq in eqs:
 603                     i, = [i for i, selfeq in enumerate(self.eqs) if selfeq == eq]
 604                     del self.eqs[i]
 605                 return 1
 606             elif len(vars) < l:
 607                 raise RuntimeError("equations are overdetermined")
 608         return 0
 609
 610 solver = Solver()
 611
 612