pyx/mathtree.py

   1 #!/usr/bin/env python
   2 #
   3 #
   4 # Copyright (C) 2002 Jörg Lehmann <joergl@users.sourceforge.net>
   5 # Copyright (C) 2002 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 string, re, math
  25
  26
  27 class ParseStr:
  28
  29     def __init__(self, StrToParse, Pos = 0):
  30         self.StrToParse = StrToParse
  31         self.Pos = Pos
  32
  33     def __repr__(self):
  34         return "ParseStr('" + self.StrToParse + "', " + str(self.Pos) + ")"
  35
  36     def __str__(self, Indent = ""):
  37         WhiteSpaces = ""
  38         for i in range(self.Pos):
  39             WhiteSpaces = WhiteSpaces + " "
  40         return Indent + self.StrToParse + "\n" + Indent + WhiteSpaces + "^"
  41
  42     def NextNonWhiteSpace(self, i = None):
  43         if i == None:
  44             i = self.Pos
  45         while self.StrToParse[i] in string.whitespace:
  46             i = i + 1
  47         return i
  48
  49     def MatchStr(self, Str):
  50         try:
  51             i = self.NextNonWhiteSpace()
  52             if self.StrToParse[i: i + len(Str)] == Str:
  53                 self.Pos = i + len(Str)
  54                 return Str
  55         except IndexError:
  56             pass
  57
  58     def MatchStrParenthesis(self, Str):
  59         try:
  60             i = self.NextNonWhiteSpace()
  61             if self.StrToParse[i: i + len(Str)] == Str:
  62                 i = i + len(Str)
  63                 i = self.NextNonWhiteSpace(i)
  64                 if self.StrToParse[i: i + 1] == "(":
  65                     self.Pos = i + 1
  66                     return Str
  67         except IndexError:
  68             pass
  69
  70     def MatchPattern(self, Pat):
  71         try:
  72             i = self.NextNonWhiteSpace()
  73             Match = Pat.match(self.StrToParse[i:])
  74             if Match:
  75                 self.Pos = i + Match.end()
  76                 return Match.group()
  77         except IndexError:
  78             pass
  79
  80     def AllDone(self):
  81         try:
  82             self.NextNonWhiteSpace()
  83         except IndexError:
  84             return 1
  85         return 0
  86
  87
  88 class ArgCountError(Exception): pass
  89 class DerivativeError(Exception): pass
  90
  91 class MathTree:
  92
  93     def __init__(self, ArgCount, *Args):
  94         self.ArgCount = ArgCount
  95         self.Args = []
  96         for arg in Args:
  97             self.AddArg(arg)
  98
  99     def __repr__(self, depth = 0):
 100         indent = ""
 101         SingleIndent = "    "
 102         for i in range(depth):
 103             indent = indent + SingleIndent
 104         result = indent + self.__class__.__name__ + "(\n"
 105         for SubTree in self.Args:
 106             if isinstance(SubTree, MathTree):
 107                 result = result + SubTree.__repr__(depth + 1)
 108             else:
 109                 result = result + indent + SingleIndent + repr(SubTree)
 110
 111             if SubTree != self.Args[-1]:
 112                 result = result + ",\n"
 113             else:
 114                 result = result + ")"
 115         return result
 116
 117     def AddArg(self, Arg, Last=0, NotLast=0):
 118         if len(self.Args) == self.ArgCount:
 119             raise ArgCountError
 120         self.Args.append(Arg)
 121         if NotLast and len(self.Args) == self.ArgCount:
 122             raise ArgCountError
 123         if Last and len(self.Args) != self.ArgCount:
 124             raise ArgCountError
 125
 126     def DependOn(self, arg):
 127         for Arg in self.Args:
 128             if Arg.DependOn(arg):
 129                 return 1
 130         return 0
 131
 132     def Derivative(self, arg):
 133         if not self.DependOn(arg):
 134             return MathTreeValConst(0.0)
 135         return self.CalcDerivative(arg)
 136
 137     def VarList(self):
 138         list = [ ]
 139         for Arg in self.Args:
 140             newlist = Arg.VarList()
 141             for x in newlist:
 142                 if x not in list:
 143                     list.append(x)
 144         return list
 145
 146
 147 class MathTreeVal(MathTree):
 148
 149     def __init__(self, *args):
 150         MathTree.__init__(self, 1, *args)
 151
 152     def __str__(self):
 153         return str(self.Args[0])
 154
 155
 156 ConstPattern = re.compile(r"-?\d*((\d\.?)|(\.?\d))\d*(E[+-]?\d+)?",
 157                           re.IGNORECASE)
 158
 159 class MathTreeValConst(MathTreeVal):
 160
 161     def InitByParser(self, arg):
 162         Match = arg.MatchPattern(ConstPattern)
 163         if Match:
 164             self.AddArg(string.atof(Match))
 165             return 1
 166
 167     def CalcDerivative(self, arg):
 168         raise DerivativeError("expression doesn't depend on \"%s\"" % arg)
 169
 170     def Derivative(self, arg):
 171         return MathTreeValConst(0.0)
 172
 173     def DependOn(self, arg):
 174         return 0
 175
 176     def VarList(self):
 177         return [ ]
 178
 179     def Calc(HIDDEN_self, **args):
 180         return HIDDEN_self.Args[0]
 181
 182
 183 VarPattern = re.compile(r"[a-z_][a-z0-9_]*", re.IGNORECASE)
 184 MathConst = {"pi": math.pi, "e": math.e}
 185
 186 class MathTreeValVar(MathTreeVal):
 187
 188     def InitByParser(self, arg):
 189         Match = arg.MatchPattern(VarPattern)
 190         if Match:
 191             self.AddArg(Match)
 192             return 1
 193
 194     def CalcDerivative(self, arg):
 195         if arg != self.Args[0]:
 196             raise DerivativeError("expression doesn't depend on \"%s\"" % arg)
 197         return MathTreeValConst(1.0)
 198
 199     def DependOn(self, arg):
 200         if arg == self.Args[0]:
 201             return 1
 202         return 0
 203
 204     def VarList(self):
 205         if self.Args[0] in MathConst.keys():
 206             return []
 207         return [self.Args[0]]
 208
 209     def Calc(HIDDEN_self, **args):
 210         if HIDDEN_self.Args[0] in args.keys():
 211             return float(args[HIDDEN_self.Args[0]])
 212         return MathConst[HIDDEN_self.Args[0]]
 213
 214
 215 class MathTreeFunc(MathTree):
 216
 217     def __init__(self, name, ArgCount, *args):
 218         self.name = name
 219         MathTree.__init__(self, ArgCount, *args)
 220
 221     def InitByParser(self, arg):
 222         return arg.MatchStrParenthesis(self.name)
 223
 224     def __str__(self):
 225         args = ""
 226         for SubTree in self.Args:
 227             args = args + str(SubTree)
 228             if SubTree != self.Args[-1]:
 229                 args = args + ","
 230         return self.name + "(" + args + ")"
 231
 232
 233 class MathTreeFunc1(MathTreeFunc):
 234
 235     def __init__(self, name, *args):
 236         MathTreeFunc.__init__(self, name, 1, *args)
 237
 238
 239 class MathTreeFunc1Neg(MathTreeFunc1):
 240
 241     def __init__(self, *args):
 242         MathTreeFunc1.__init__(self, "neg", *args)
 243
 244     def CalcDerivative(self, arg):
 245         return MathTreeFunc1Neg(self.Args[0].CalcDerivative(arg))
 246
 247     def Calc(HIDDEN_self, **args):
 248         return -HIDDEN_self.Args[0].Calc(**args)
 249
 250
 251 class MathTreeFunc1Abs(MathTreeFunc1):
 252
 253     def __init__(self, *args):
 254         MathTreeFunc1.__init__(self, "abs", *args)
 255
 256     def CalcDerivative(self, arg):
 257         return MathTreeOpMul(
 258                    MathTreeFunc1Sgn(self.Args[0]),
 259                    self.Args[0].CalcDerivative(arg))
 260
 261     def Calc(HIDDEN_self, **args):
 262         return abs(HIDDEN_self.Args[0].Calc(**args))
 263
 264
 265 class MathTreeFunc1Sgn(MathTreeFunc1):
 266
 267     def __init__(self, *args):
 268         MathTreeFunc1.__init__(self, "sgn", *args)
 269
 270     def CalcDerivative(self, arg):
 271         return MathTreeValConst(0.0)
 272
 273     def Calc(HIDDEN_self, **args):
 274         if HIDDEN_self.Args[0].Calc(**args) < 0:
 275             return -1.0
 276         return 1.0
 277
 278
 279 class MathTreeFunc1Sqrt(MathTreeFunc1):
 280
 281     def __init__(self, *args):
 282         MathTreeFunc1.__init__(self, "sqrt", *args)
 283
 284     def CalcDerivative(self, arg):
 285         return MathTreeOpMul(
 286                    MathTreeOpDiv(
 287                        MathTreeValConst(0.5),
 288                        self),
 289                    self.Args[0].CalcDerivative(arg))
 290
 291     def Calc(HIDDEN_self, **args):
 292         return math.sqrt(HIDDEN_self.Args[0].Calc(**args))
 293
 294
 295 class MathTreeFunc1Exp(MathTreeFunc1):
 296
 297     def __init__(self, *args):
 298         MathTreeFunc1.__init__(self, "exp", *args)
 299
 300     def CalcDerivative(self, arg):
 301         return MathTreeOpMul(self, self.Args[0].CalcDerivative(arg))
 302
 303     def Calc(HIDDEN_self, **args):
 304         return math.exp(HIDDEN_self.Args[0].Calc(**args))
 305
 306
 307 class MathTreeFunc1Log(MathTreeFunc1):
 308
 309     def __init__(self, *args):
 310         MathTreeFunc1.__init__(self, "log", *args)
 311
 312     def CalcDerivative(self, arg):
 313         return MathTreeOpDiv(self.Args[0].CalcDerivative(arg), self.Args[0])
 314
 315     def Calc(HIDDEN_self, **args):
 316         return math.log(HIDDEN_self.Args[0].Calc(**args))
 317
 318
 319 class MathTreeFunc1Sin(MathTreeFunc1):
 320
 321     def __init__(self, *args):
 322         MathTreeFunc1.__init__(self, "sin", *args)
 323
 324     def CalcDerivative(self, arg):
 325         return MathTreeOpMul(
 326                    MathTreeFunc1Cos(self.Args[0]),
 327                    self.Args[0].CalcDerivative(arg))
 328
 329     def Calc(HIDDEN_self, **args):
 330         return math.sin(HIDDEN_self.Args[0].Calc(**args))
 331
 332
 333 class MathTreeFunc1Cos(MathTreeFunc1):
 334
 335     def __init__(self, *args):
 336         MathTreeFunc1.__init__(self, "cos", *args)
 337
 338     def CalcDerivative(self, arg):
 339         return MathTreeOpMul(
 340                    MathTreeFunc1Neg(MathTreeFunc1Sin(self.Args[0])),
 341                    self.Args[0].CalcDerivative(arg))
 342
 343     def Calc(HIDDEN_self, **args):
 344         return math.cos(HIDDEN_self.Args[0].Calc(**args))
 345
 346
 347 class MathTreeFunc1Tan(MathTreeFunc1):
 348
 349     def __init__(self, *args):
 350         MathTreeFunc1.__init__(self, "tan", *args)
 351
 352     def CalcDerivative(self, arg):
 353         return MathTreeOpDiv(
 354                    self.Args[0].CalcDerivative(arg),
 355                    MathTreeOpPow(
 356                        MathTreeFunc1Cos(self.Args[0]),
 357                        MathTreeValConst(2.0)))
 358
 359     def Calc(HIDDEN_self, **args):
 360         return math.tan(HIDDEN_self.Args[0].Calc(**args))
 361
 362
 363 class MathTreeFunc1ASin(MathTreeFunc1):
 364
 365     def __init__(self, *args):
 366         MathTreeFunc1.__init__(self, "asin", *args)
 367
 368     def CalcDerivative(self, arg):
 369         return MathTreeOpDiv(
 370                    self.Args[0].CalcDerivative(arg),
 371                    MathTreeFunc1Sqrt(
 372                        MathTreeOpSub(
 373                            MathTreeValConst(1.0),
 374                            MathTreeOpPow(
 375                                self.Args[0],
 376                                MathTreeValConst(2.0)))))
 377
 378     def Calc(HIDDEN_self, **args):
 379         return math.asin(HIDDEN_self.Args[0].Calc(**args))
 380
 381
 382 class MathTreeFunc1ACos(MathTreeFunc1):
 383
 384     def __init__(self, *args):
 385         MathTreeFunc1.__init__(self, "acos", *args)
 386
 387     def CalcDerivate(self, arg):
 388         return MathTreeOpDiv(
 389                    MathTreeFunc1Neg(self.Args[0].CalcDerivative(arg)),
 390                    MathTreeFunc1Sqrt(
 391                        MathTreeOpSub(
 392                            MathTreeValConst(1.0),
 393                            MathTreeOpPow(
 394                                self.Args[0],
 395                                MathTreeValConst(2.0)))))
 396
 397     def Calc(HIDDEN_self, **args):
 398         return math.acos(HIDDEN_self.Args[0].Calc(**args))
 399
 400
 401 class MathTreeFunc1ATan(MathTreeFunc1):
 402
 403     def __init__(self, *args):
 404         MathTreeFunc1.__init__(self, "atan", *args)
 405
 406     def CalcDerivate(self, arg):
 407         return MathTreeOpDiv(
 408                    self.Args[0].CalcDerivative(arg),
 409                    MathTreeOpAdd(
 410                        MathTreeValConst(1.0),
 411                        MathTreeOpPow(
 412                            self.Args[0],
 413                            MathTreeValConst(2.0))))
 414
 415     def Calc(HIDDEN_self, **args):
 416         return math.atan(HIDDEN_self.Args[0].Calc(**args))
 417
 418
 419 class MathTreeFunc1SinD(MathTreeFunc1):
 420
 421     def __init__(self, *args):
 422         MathTreeFunc1.__init__(self, "sind", *args)
 423
 424     def CalcDerivative(self, arg):
 425         return MathTreeOpMul(
 426                    MathTreeFunc1CosD(self.Args[0]),
 427                    MathTreeOpMul(
 428                        MathTreeValConst(math.pi/180.0),
 429                        self.Args[0].CalcDerivative(arg)))
 430
 431     def Calc(HIDDEN_self, **args):
 432         return math.sin(math.pi/180.0*HIDDEN_self.Args[0].Calc(**args))
 433
 434
 435 class MathTreeFunc1CosD(MathTreeFunc1):
 436
 437     def __init__(self, *args):
 438         MathTreeFunc1.__init__(self, "cosd", *args)
 439
 440     def CalcDerivative(self, arg):
 441         return MathTreeOpMul(
 442                    MathTreeFunc1Neg(MathTreeFunc1Sin(self.Args[0])),
 443                    MathTreeOpMul(
 444                        MathTreeValConst(math.pi/180.0),
 445                        self.Args[0].CalcDerivative(arg)))
 446
 447     def Calc(HIDDEN_self, **args):
 448         return math.cos(math.pi/180.0*HIDDEN_self.Args[0].Calc(**args))
 449
 450
 451 class MathTreeFunc1TanD(MathTreeFunc1):
 452
 453     def __init__(self, *args):
 454         MathTreeFunc1.__init__(self, "tand", *args)
 455
 456     def CalcDerivative(self, arg):
 457         return MathTreeOpDiv(
 458                    MathTreeOpMul(
 459                        MathTreeValConst(math.pi/180.0),
 460                        self.Args[0].CalcDerivative(arg)),
 461                    MathTreeOpPow(
 462                        MathTreeFunc1Cos(self.Args[0]),
 463                        MathTreeValConst(2.0)))
 464
 465     def Calc(HIDDEN_self, **args):
 466         return math.tan(math.pi/180.0*HIDDEN_self.Args[0].Calc(**args))
 467
 468
 469 class MathTreeFunc1ASinD(MathTreeFunc1):
 470
 471     def __init__(self, *args):
 472         MathTreeFunc1.__init__(self, "asind", *args)
 473
 474     def CalcDerivative(self, arg):
 475         return MathTreeOpDiv(
 476                    MathTreeOpMul(
 477                        MathTreeValConst(180.0/math.pi),
 478                        self.Args[0].CalcDerivative(arg)),
 479                    MathTreeFunc1Sqrt(
 480                        MathTreeOpSub(
 481                            MathTreeValConst(1.0),
 482                            MathTreeOpPow(
 483                                self.Args[0],
 484                                MathTreeValConst(2.0)))))
 485
 486     def Calc(HIDDEN_self, **args):
 487         return 180.0/math.pi*math.asin(HIDDEN_self.Args[0].Calc(**args))
 488
 489
 490 class MathTreeFunc1ACosD(MathTreeFunc1):
 491
 492     def __init__(self, *args):
 493         MathTreeFunc1.__init__(self, "acosd", *args)
 494
 495     def CalcDerivate(self, arg):
 496         return MathTreeOpDiv(
 497                    MathTreeFunc1Neg(
 498                        MathTreeOpMul(
 499                            MathTreeValConst(180.0/math.pi),
 500                            self.Args[0].CalcDerivative(arg))),
 501                    MathTreeFunc1Sqrt(
 502                        MathTreeOpSub(
 503                            MathTreeValConst(1.0),
 504                            MathTreeOpPow(
 505                                self.Args[0],
 506                                MathTreeValConst(2.0)))))
 507
 508     def Calc(HIDDEN_self, **args):
 509         return 180.0/math.pi*math.acos(HIDDEN_self.Args[0].Calc(**args))
 510
 511
 512 class MathTreeFunc1ATanD(MathTreeFunc1):
 513
 514     def __init__(self, *args):
 515         MathTreeFunc1.__init__(self, "atand", *args)
 516
 517     def CalcDerivate(self, arg):
 518         return MathTreeOpDiv(
 519                    MathTreeOpMul(
 520                        MathTreeValConst(180.0/math.pi),
 521                        self.Args[0].CalcDerivative(arg)),
 522                    MathTreeOpAdd(
 523                        MathTreeValConst(1.0),
 524                        MathTreeOpPow(
 525                            self.Args[0],
 526                            MathTreeValConst(2.0))))
 527
 528     def Calc(HIDDEN_self, **args):
 529         return 180.0/math.pi*math.atan(HIDDEN_self.Args[0].Calc(**args))
 530
 531
 532 class MathTreeFunc2(MathTreeFunc):
 533
 534     def __init__(self, name, *args):
 535         MathTreeFunc.__init__(self, name, 2, *args)
 536
 537
 538 class MathTreeFunc2Norm(MathTreeFunc2):
 539
 540     def __init__(self, *args):
 541         MathTreeFunc2.__init__(self, "norm", *args)
 542
 543     def CalcDerivative(self, arg):
 544         if self.Args[0].DependOn(arg):
 545             if self.Args[1].DependOn(arg):
 546                 return MathTreeOpDiv(
 547                            MathTreeOpAdd(
 548                                MathTreeOpMul(
 549                                    self.Args[0],
 550                                    self.Args[0].CalcDerivative(arg)),
 551                                MathTreeOpMul(
 552                                    self.Args[1],
 553                                    self.Args[1].CalcDerivative(arg))),
 554                            self)
 555             else:
 556                 return MathTreeOpDiv(
 557                            MathTreeOpMul(
 558                                self.Args[0],
 559                                self.Args[0].CalcDerivative(arg)),
 560                            self)
 561         else:
 562             if self.Args[1].DependOn(arg):
 563                 return MathTreeOpDiv(
 564                            MathTreeOpMul(
 565                                self.Args[1],
 566                                self.Args[1].CalcDerivative(arg)),
 567                            self)
 568
 569     def Calc(HIDDEN_self, **args):
 570         return math.sqrt(HIDDEN_self.Args[0].Calc(**args) ** 2 +
 571                          HIDDEN_self.Args[1].Calc(**args) ** 2)
 572
 573
 574 FuncExternPattern = re.compile(r"([a-z_][a-z0-9_]*)\s*\(", re.IGNORECASE)
 575
 576 class MathTreeFuncExtern(MathTreeFunc):
 577
 578     def __init__(self, *args):
 579         MathTreeFunc.__init__(self, None, -1, *args)
 580
 581     def InitByParser(self, arg):
 582         Match = arg.MatchPattern(FuncExternPattern)
 583         if Match:
 584             self.name = Match[:-1].strip()
 585             return self.name
 586
 587     def Calc(HIDDEN_self, **args):
 588         return args[HIDDEN_self.name](*[arg.Calc(**args) for arg in HIDDEN_self.Args])
 589
 590
 591 class MathTreeOp(MathTree):
 592
 593     def __init__(self, level, symbol, *args):
 594         self.ParenthesisBarrier = 0
 595         self.level = level
 596         self.symbol = symbol
 597         MathTree.__init__(self, 2, *args)
 598
 599     def __str__(self):
 600         result = ""
 601         if isinstance(self.Args[0], MathTreeOp) and\
 602             self.level > self.Args[0].level:
 603             result = result + "(" + str(self.Args[0]) + ")"
 604         else:
 605             result = result + str(self.Args[0])
 606         result = result + self.symbol
 607         if isinstance(self.Args[1], MathTreeOp) and\
 608             self.level >= self.Args[1].level:
 609             result = result + "(" + str(self.Args[1]) + ")"
 610         else:
 611             result = result + str(self.Args[1])
 612         return result
 613
 614     def InitByParser(self, arg):
 615         return arg.MatchStr(self.symbol)
 616
 617
 618 class MathTreeOpAdd(MathTreeOp):
 619
 620     def __init__(self, *args):
 621         MathTreeOp.__init__(self, 1, "+", *args)
 622
 623     def CalcDerivative(self, arg):
 624         if self.Args[0].DependOn(arg):
 625             if self.Args[1].DependOn(arg):
 626                 return MathTreeOpAdd(
 627                            self.Args[0].CalcDerivative(arg),
 628                            self.Args[1].CalcDerivative(arg))
 629             else:
 630                 return self.Args[0].CalcDerivative(arg)
 631         else:
 632             if self.Args[1].DependOn(arg):
 633                 return self.Args[1].CalcDerivative(arg)
 634
 635     def Calc(HIDDEN_self, **args):
 636         return HIDDEN_self.Args[0].Calc(**args) + HIDDEN_self.Args[1].Calc(**args)
 637
 638
 639 class MathTreeOpSub(MathTreeOp):
 640
 641     def __init__(self, *args):
 642         MathTreeOp.__init__(self, 1, "-", *args)
 643
 644     def CalcDerivative(self, arg):
 645         if self.Args[0].DependOn(arg):
 646             if self.Args[1].DependOn(arg):
 647                 return MathTreeOpSub(
 648                            self.Args[0].CalcDerivative(arg),
 649                            self.Args[1].CalcDerivative(arg))
 650             else:
 651                 return self.Args[0].CalcDerivative(arg)
 652         else:
 653             if self.Args[1].DependOn(arg):
 654                 return MathTreeFunc1Neg(self.Args[1].CalcDerivative(arg))
 655
 656     def Calc(HIDDEN_self, **args):
 657         return HIDDEN_self.Args[0].Calc(**args) - HIDDEN_self.Args[1].Calc(**args)
 658
 659
 660 class MathTreeOpMul(MathTreeOp):
 661
 662     def __init__(self, *args):
 663         MathTreeOp.__init__(self, 2, "*", *args)
 664
 665     def CalcDerivative(self, arg):
 666         if self.Args[0].DependOn(arg):
 667             if self.Args[1].DependOn(arg):
 668                 return MathTreeOpAdd(
 669                            MathTreeOpMul(
 670                                self.Args[0],
 671                                self.Args[1].CalcDerivative(arg)),
 672                            MathTreeOpMul(
 673                                self.Args[0].CalcDerivative(arg),
 674                                self.Args[1]))
 675             else:
 676                 return MathTreeOpMul(
 677                            self.Args[0].CalcDerivative(arg),
 678                            self.Args[1])
 679         else:
 680             if self.Args[1].DependOn(arg):
 681                 return MathTreeOpMul(
 682                            self.Args[0],
 683                            self.Args[1].CalcDerivative(arg))
 684
 685     def Calc(HIDDEN_self, **args):
 686         return HIDDEN_self.Args[0].Calc(**args) * HIDDEN_self.Args[1].Calc(**args)
 687
 688
 689 class MathTreeOpDiv(MathTreeOp):
 690
 691     def __init__(self, *args):
 692         MathTreeOp.__init__(self, 2, "/", *args)
 693
 694     def CalcDerivative(self, arg):
 695         if self.Args[0].DependOn(arg):
 696             if self.Args[1].DependOn(arg):
 697                 return MathTreeOpMul(
 698                            MathTreeOpSub(
 699                                MathTreeOpMul(
 700                                    self.Args[0].CalcDerivative(arg),
 701                                    self.Args[1]),
 702                                MathTreeOpMul(
 703                                    self.Args[0],
 704                                    self.Args[1].CalcDerivative(arg))),
 705                            MathTreeOpPow(
 706                                self.Args[1],
 707                                MathTreeValConst(-2.0)))
 708             else:
 709                 return MathTreeOpDiv(
 710                            self.Args[0].CalcDerivative(arg),
 711                            self.Args[1])
 712         else:
 713             if self.Args[1].DependOn(arg):
 714                 return MathTreeOpMul(
 715                            MathTreeOpMul(
 716                                MathTreeFunc1Neg(self.Args[0]),
 717                                MathTreeOpPow(
 718                                    self.Args[1],
 719                                    MathTreeValConst(-2.0))),
 720                            self.Args[1].CalcDerivative(arg))
 721
 722     def Calc(HIDDEN_self, **args):
 723         return HIDDEN_self.Args[0].Calc(**args) / HIDDEN_self.Args[1].Calc(**args)
 724
 725
 726 class MathTreeOpPow(MathTreeOp):
 727
 728     def __init__(self, *args):
 729         MathTreeOp.__init__(self, 3, "^", *args)
 730
 731     def InitByParser(self, arg):
 732         pos = arg.MatchStr("^")
 733         if pos:
 734             return 1
 735         else:
 736             return arg.MatchStr("**")
 737
 738     def CalcDerivative(self, arg):
 739         if self.Args[0].DependOn(arg):
 740             if self.Args[1].DependOn(arg):
 741                 return MathTreeOpMul(
 742                            MathTreeOpPow(
 743                                self.Args[0],
 744                                self.Args[1]),
 745                            MathTreeOpAdd(
 746                                MathTreeOpMul(
 747                                    MathTreeFunc1Log(self.Args[0]),
 748                                    self.Args[1].CalcDerivative(arg)),
 749                                MathTreeOpMul(
 750                                    self.Args[1],
 751                                    MathTreeOpDiv(
 752                                        self.Args[0].CalcDerivative(arg),
 753                                        self.Args[0]))))
 754             else:
 755                 return MathTreeOpMul(
 756                            self.Args[1],
 757                            MathTreeOpMul(
 758                                MathTreeOpPow(
 759                                    self.Args[0],
 760                                    MathTreeOpSub(
 761                                        self.Args[1],
 762                                        MathTreeValConst(1.0))),
 763                                self.Args[0].CalcDerivative(arg)))
 764         else:
 765             if self.Args[1].DependOn(arg):
 766                 return MathTreeOpMul(
 767                            MathTreeOpMul(
 768                                MathTreeOpPow(
 769                                    self.Args[0],
 770                                    self.Args[1]),
 771                                MathTreeFunc1Log(self.Args[0])),
 772                            self.Args[1].CalcDerivative(arg))
 773
 774     def Calc(HIDDEN_self, **args):
 775         return HIDDEN_self.Args[0].Calc(**args) ** HIDDEN_self.Args[1].Calc(**args)
 776
 777
 778
 779 class UndefinedMathTreeParseError(Exception):
 780
 781     def __init__(self, ParseStr, MathTree):
 782         self.ParseStr = ParseStr
 783         self.MathTree = MathTree
 784
 785     def __str__(self):
 786         return "\n" + str(self.ParseStr)
 787
 788
 789 class RightParenthesisExpectedMathTreeParseError(UndefinedMathTreeParseError): pass
 790 class RightParenthesisFoundMathTreeParseError(UndefinedMathTreeParseError): pass
 791 class CommaExpectedMathTreeParseError(UndefinedMathTreeParseError): pass
 792 class CommaFoundMathTreeParseError(UndefinedMathTreeParseError): pass
 793 class OperatorExpectedMathTreeParseError(UndefinedMathTreeParseError): pass
 794 class OperandExpectedMathTreeParseError(UndefinedMathTreeParseError): pass
 795
 796
 797 DefaultMathTreeOps = (MathTreeOpPow, MathTreeOpDiv, MathTreeOpMul, MathTreeOpSub, MathTreeOpAdd)
 798 DefaultMathTreeFuncs = (MathTreeFunc1Neg, MathTreeFunc1Abs, MathTreeFunc1Sgn, MathTreeFunc1Sqrt,
 799                         MathTreeFunc1Exp, MathTreeFunc1Log,
 800                         MathTreeFunc1Sin, MathTreeFunc1Cos, MathTreeFunc1Tan,
 801                         MathTreeFunc1ASin, MathTreeFunc1ACos, MathTreeFunc1ATan,
 802                         MathTreeFunc1SinD, MathTreeFunc1CosD, MathTreeFunc1TanD,
 803                         MathTreeFunc1ASinD, MathTreeFunc1ACosD, MathTreeFunc1ATanD,
 804                         MathTreeFunc2Norm)
 805
 806 DefaultMathTreeVals = (MathTreeValConst, MathTreeValVar)
 807
 808 class parser:
 809
 810     def __init__(self, MathTreeOps=DefaultMathTreeOps,
 811                        MathTreeFuncs=DefaultMathTreeFuncs,
 812                        MathTreeVals=DefaultMathTreeVals):
 813         self.MathTreeOps = MathTreeOps
 814         self.MathTreeFuncs = MathTreeFuncs
 815         self.MathTreeVals = MathTreeVals
 816
 817     def parse(self, str):
 818         return self.ParseMathTree(ParseStr(str))
 819
 820     def ParseMathTree(self, arg):
 821         Tree = None
 822         Match = arg.MatchPattern(re.compile(r"\s*-(?![0-9\.])"))
 823         if Match:
 824             Tree = MathTreeFunc1Neg()
 825         while 1:
 826             i = arg.MatchStr("(")
 827             if i:
 828                 try:
 829                     self.ParseMathTree(arg)
 830                     raise RightParenthesisExpectedMathTreeParseError(arg, Tree)
 831                 except RightParenthesisFoundMathTreeParseError, e:
 832                     if isinstance(e.MathTree, MathTreeOp):
 833                         e.MathTree.ParenthesisBarrier = 1
 834                     if Tree is not None:
 835                         SubTree = Tree # XXX: four lines code dublication
 836                         while isinstance(SubTree, MathTreeOp) and len(SubTree.Args) == 2:
 837                             SubTree = SubTree.Args[1]
 838                         SubTree.AddArg(e.MathTree)
 839                     else:
 840                         Tree = e.MathTree
 841             else:
 842                 for FuncClass in self.MathTreeFuncs:
 843                     Func = FuncClass()
 844                     if Func.InitByParser(arg):
 845                         if Tree is not None:
 846                             SubTree = Tree # XXX: four lines code dublication
 847                             while isinstance(SubTree, MathTreeOp) and len(SubTree.Args) == 2:
 848                                 SubTree = SubTree.Args[1]
 849                             SubTree.AddArg(Func)
 850                         else:
 851                             Tree = Func
 852                         while 1:
 853                             try:
 854                                 self.ParseMathTree(arg)
 855                                 raise RightParenthesisExpectedMathTreeParseError(arg, Tree)
 856                             except CommaFoundMathTreeParseError, e:
 857                                 try:
 858                                     Func.AddArg(e.MathTree, NotLast=1)
 859                                 except ArgCountError:
 860                                     raise RightParenthesisExpectedMathTreeParseError(arg, Tree)
 861                                 continue
 862                             except RightParenthesisFoundMathTreeParseError, e:
 863                                 try:
 864                                     Func.AddArg(e.MathTree, Last=1)
 865                                 except ArgCountError:
 866                                     raise CommaExpectedMathTreeParseError(arg, Tree)
 867                                 break
 868                         break
 869                 else:
 870                     FuncExtern = MathTreeFuncExtern()
 871                     if FuncExtern.InitByParser(arg):
 872                         if Tree is not None:
 873                             SubTree = Tree # XXX: four lines code dublication
 874                             while isinstance(SubTree, MathTreeOp) and len(SubTree.Args) == 2:
 875                                 SubTree = SubTree.Args[1]
 876                             SubTree.AddArg(FuncExtern)
 877                         else:
 878                             Tree = FuncExtern
 879                         while 1:
 880                             try:
 881                                 self.ParseMathTree(arg)
 882                                 raise RightParenthesisExpectedMathTreeParseError(arg, Tree)
 883                             except CommaFoundMathTreeParseError, e:
 884                                 FuncExtern.AddArg(e.MathTree)
 885                                 continue
 886                             except RightParenthesisFoundMathTreeParseError, e:
 887                                 FuncExtern.AddArg(e.MathTree)
 888                                 break
 889                     else:
 890                         for ValClass in self.MathTreeVals:
 891                             Val = ValClass()
 892                             if Val.InitByParser(arg):
 893                                 if Tree is not None:
 894                                     SubTree = Tree # XXX: four lines code dublication
 895                                     while isinstance(SubTree, MathTreeOp) and len(SubTree.Args) == 2:
 896                                         SubTree = SubTree.Args[1]
 897                                     SubTree.AddArg(Val)
 898                                 else:
 899                                     Tree = Val
 900                                 break
 901                         else:
 902                             raise OperandExpectedMathTreeParseError(arg, Tree)
 903             if arg.AllDone():
 904                 return Tree
 905             i = arg.MatchStr(")")
 906             if i:
 907                 raise RightParenthesisFoundMathTreeParseError(arg, Tree)
 908             i = arg.MatchStr(",")
 909             if i:
 910                 raise CommaFoundMathTreeParseError(arg, Tree)
 911             for OpClass in self.MathTreeOps:
 912                 Op = OpClass()
 913                 if Op.InitByParser(arg):
 914                     SubTree = Tree
 915                     SubTreeRoot = None
 916                     while isinstance(SubTree, MathTreeOp) and\
 917                           Op.level > SubTree.level and\
 918                           not SubTree.ParenthesisBarrier:
 919                         SubTreeRoot = SubTree
 920                         SubTree = SubTree.Args[1]
 921                     if SubTreeRoot:
 922                         Op.AddArg(SubTree)
 923                         SubTreeRoot.Args[1] = Op
 924                     else:
 925                         Op.AddArg(Tree)
 926                         Tree = Op
 927                     break
 928             else:
 929                 raise OperatorExpectedMathTreeParseError(arg, Tree)