branches/2d_boundary/Tests/svgfig/curve.py

   1 import math, cmath, copy, re, sys, new
   2 import defaults, svg, trans, pathdata, glyphs, _curve
   3
   4 ############################### generic curve with marks (tick marks, arrows, etc)
   5
   6 class Curve(svg.SVG):
   7   attrib = {"stroke": "black", "fill": "none"}
   8   smooth = False
   9   marks = []
  10   random_sampling = True
  11   random_seed = 12345
  12   recursion_limit = 15
  13   linearity_limit = 0.05
  14   discontinuity_limit = 5.
  15   text_offsetx = 0.
  16   text_offsety = -2.5
  17   text_attrib = {}
  18
  19   _varlist = ["attrib", "smooth", "marks", "random_sampling", "random_seed", "recursion_limit", "linearity_limit", "discontinuity_limit", "text_offsetx", "text_offsety", "text_attrib"]
  20
  21   def __init__(self, expr, low, high, **kwds):
  22     self.__dict__["tag"] = None
  23     self.__dict__["children"] = []
  24     self.__dict__["_svg"] = []
  25     self.__dict__["f"] = svg.cannonical_parametric(expr)
  26     self.__dict__["low"] = low
  27     self.__dict__["high"] = high
  28
  29     for var in self._varlist:
  30       if not callable(eval("self.%s" % var)) and var[:1] != "__" and var[-1:] != "__":
  31         if var in kwds:
  32           self.__dict__[var] = kwds[var]
  33           del kwds[var]
  34         else:
  35           self.__dict__[var] = copy.deepcopy(eval("self.%s" % var))
  36
  37     # needed to set arrow color
  38     if "stroke" in kwds:
  39       self.attrib["stroke"] = kwds["stroke"]
  40       del kwds["stroke"]
  41
  42     if "farrow" in kwds:
  43       if kwds["farrow"] is True:
  44         self.add(self.high, glyphs.farrowhead)
  45       else:
  46         self.add(self.high, kwds["farrow"])
  47       self.marks[-1][1]["fill"] = self["stroke"]
  48       del kwds["farrow"]
  49
  50     if "barrow" in kwds:
  51       if kwds["barrow"] is True:
  52         self.add(self.low, glyphs.barrowhead)
  53       else:
  54         self.add(self.low, kwds["barrow"])
  55       self.marks[-1][1]["fill"] = self["stroke"]
  56       del kwds["barrow"]
  57
  58     self.clean_arrows()
  59
  60     # now the rest of the attributes (other than stroke)
  61     self.__dict__["attrib"].update(kwds)
  62     self.__dict__["trans"] = []
  63
  64   def __call__(self, t, transformed=True):
  65     x, y = self.f(t)
  66     if transformed:
  67       for trans in self.trans:
  68         x, y = trans(x, y)
  69     return x, y
  70
  71   def angle(self, t, transformed=True):
  72     x, y = self(t, transformed)
  73
  74     tprime = t + trans.epsilon * abs(self.high - self.low)
  75     xprime, yprime = self(tprime, transformed)
  76
  77     delx, dely = xprime - x, yprime - y
  78     return math.atan2(dely, delx)
  79
  80   ### pickleability and access issues
  81   def __getattr__(self, name): return self.__dict__[name]
  82   def __setattr__(self, name, value): self.__dict__[name] = value
  83
  84   def __getstate__(self):
  85     mostdict = copy.copy(self.__dict__)
  86     del mostdict["f"]
  87     del mostdict["trans"]
  88     transcode = map(lambda f: (f.func_code, f.func_name), self.trans)
  89     fcode = self.f.func_code, self.f.func_name
  90     return (sys.version_info, defaults.version_info, mostdict, transcode, fcode)
  91
  92   def __setstate__(self, state):
  93     self.__dict__ = state[2]
  94     self.__dict__["trans"] = []
  95
  96     for code, name in state[3]:
  97       context = globals()
  98       if "z" in code.co_names:
  99         context.update(cmath.__dict__)
 100       else:
 101         context.update(math.__dict__)
 102       f = new.function(code, context)
 103       f.func_name = name
 104       self.__dict__["trans"].append(f)
 105
 106     context = globals()
 107     if "z" in state[4][0].co_names:
 108       context.update(cmath.__dict__)
 109     else:
 110       context.update(math.__dict__)
 111     self.__dict__["f"] = new.function(state[4][0], context)
 112     self.__dict__["f"].func_name = state[4][1]
 113
 114   def __deepcopy__(self, memo={}):
 115     mostdict = copy.copy(self.__dict__)
 116     del mostdict["trans"]
 117     del mostdict["f"]
 118     if "repr" in mostdict: del mostdict["repr"]
 119     output = new.instance(self.__class__)
 120     output.__dict__ = copy.deepcopy(mostdict, memo)
 121     output.__dict__["trans"] = copy.copy(self.trans)
 122     output.__dict__["f"] = self.f
 123
 124     memo[id(self)] = output
 125     return output
 126
 127   ### presentation
 128   def __repr__(self):
 129     marks = ""
 130     if len(self.marks) == 1: marks = " (1 mark)"
 131     elif len(self.marks) > 1: marks = " (%d marks)" % len(self.marks)
 132
 133     trans = ""
 134     if len(self.trans) > 0: trans = " (%d trans)" % len(self.trans)
 135
 136     attrib = ""
 137     for var in "stroke", "fill":
 138       if var in self.attrib:
 139         attrib += " %s=%s" % (var, repr(self.attrib[var]))
 140
 141     return "<%s %s from %g to %g%s%s%s>" % (self.__class__.__name__, self.f, self.low, self.high, marks, trans, attrib)
 142
 143   ### transformation is like Delay
 144   def transform(self, t): self.trans.append(svg.cannonical_transformation(t))
 145
 146   def bbox(self): return pathdata.bbox(self.d())
 147
 148   ### construct the SVG path
 149   def svg(self):
 150     obj = new.instance(svg.SVG)
 151     obj.__dict__["tag"] = "path"
 152     obj.__dict__["attrib"] = self.attrib
 153     obj.__dict__["children"] = self.children
 154     obj.__dict__["_svg"] = obj
 155     obj.attrib["d"] = self.d()
 156
 157     if self.marks == []:
 158       output = obj
 159     else:
 160       output = svg.SVG("g", obj)
 161
 162       lowX, lowY = self(self.low)
 163       highX, highY = self(self.high)
 164
 165       for item in self.marks:
 166         if isinstance(item, (int, long, float)):
 167           t, mark = item, glyphs.tick
 168         else:
 169           t, mark = item # marks should be (pos, mark) pairs or just pos
 170
 171         X, Y = self(t)
 172         if self.low <= t <= self.high or \
 173            math.sqrt((X - lowX)**2 + (Y - lowY)**2) < trans.epsilon or \
 174            math.sqrt((X - highX)**2 + (Y - highY)**2) < trans.epsilon:
 175
 176           angle = self.angle(t)
 177
 178           if isinstance(mark, basestring):
 179             mark = self._render_text(X, Y, angle, mark)
 180
 181           else:
 182             mark = trans.transform(lambda x, y: (X + math.cos(angle)*x - math.sin(angle)*y,
 183                                                  Y + math.sin(angle)*x + math.cos(angle)*y), mark)
 184           output.append(mark)
 185
 186     self._svg = output
 187
 188   def d(self):
 189     data = _curve.curve(self.f, self.trans, self.low, self.high,
 190                         self.random_sampling, self.random_seed, self.recursion_limit, self.linearity_limit, self.discontinuity_limit)
 191
 192     segments = []
 193     last_d = None
 194     for d in data:
 195       if d is not None:
 196         if last_d is None: segments.append([])
 197         segments[-1].append(d)
 198       last_d = d
 199
 200     output = []
 201     if self.smooth:
 202       for seg in segments:
 203         output.extend(pathdata.smooth(*seg))
 204     else:
 205       for seg in segments:
 206         output.extend(pathdata.poly(*seg))
 207
 208     return output
 209
 210   def _render_text(self, X, Y, angle, text):
 211     text_attrib = {"transform": "translate(%g, %g) rotate(%g)" %
 212                    (X + self.text_offsetx*math.cos(angle) - self.text_offsety*math.sin(angle),
 213                     Y + self.text_offsetx*math.sin(angle) + self.text_offsety*math.cos(angle), 180./math.pi*angle),
 214                    "text-anchor": "middle"}
 215     text_attrib.update(self.text_attrib)
 216     return svg.SVG("text", 0., 0., **text_attrib)(text)
 217
 218   ### lots of functions for adding/removing marks
 219   def _matches(self, matching, mark):
 220     if matching is None: return True
 221
 222     try:
 223       if isinstance(mark, matching): return True
 224     except TypeError: pass
 225
 226     if isinstance(mark, svg.SVG) and isinstance(matching, basestring):
 227       if re.search(matching, mark.tag): return True
 228       if "repr" in mark.__dict__ and re.search(matching, mark.repr): return True
 229
 230     if isinstance(matching, basestring) and isinstance(mark, basestring):
 231       if re.search(matching, mark): return True
 232
 233     return matching == mark
 234
 235   def wipe(self, low=None, high=None, matching=None):
 236     if low is None: low = self.low
 237     if high is None: high = self.high
 238
 239     newmarks = []
 240     for item in self.marks:
 241       if isinstance(item, (int, long, float)):
 242         if self._matches(matching, item):
 243           if not low <= item <= high: newmarks.append(item)
 244         else:
 245           newmarks.append(item)
 246
 247       else:
 248         pos, mark = item # marks should be (pos, mark) pairs or just pos
 249         if self._matches(matching, mark):
 250           if not low <= pos <= high: newmarks.append(item)
 251         else:
 252           newmarks.append(item)
 253
 254     self.marks = newmarks
 255
 256   def keep(self, low=None, high=None, matching=None):
 257     if low is None: low = self.low
 258     if high is None: high = self.high
 259
 260     newmarks = []
 261     for item in self.marks:
 262       if isinstance(item, (int, long, float)):
 263         if self._matches(matching, item):
 264           if low <= item <= high: newmarks.append(item)
 265
 266       else:
 267         pos, mark = item # marks should be (pos, mark) pairs or just pos
 268         if self._matches(matching, mark):
 269           if low <= pos <= high: newmarks.append(item)
 270
 271     self.marks = newmarks
 272
 273   def drop(self, t, tolerance=None, matching=None):
 274     if tolerance is None: tolerance = trans.epsilon * abs(self.high - self.low)
 275     self.wipe(t - tolerance, t + tolerance, matching=matching)
 276
 277   def add(self, t, mark, angle=0., dx=0., dy=0.):
 278     if not isinstance(mark, basestring):
 279       mark = trans.transform(lambda x, y: (dx + math.cos(angle)*x - math.sin(angle)*y,
 280                                            dy + math.sin(angle)*x + math.cos(angle)*y), mark)
 281     self.marks.append((t, mark))
 282
 283   def tick(self, t, mark=None):
 284     if mark is None:
 285       self.add(t, glyphs.tick)
 286       self.marks[-1][1]["stroke"] = self["stroke"]
 287     elif isinstance(mark, basestring):
 288       self.add(t, glyphs.tick)
 289       self.marks[-1][1]["stroke"] = self["stroke"]
 290       self.add(t, mark)
 291     else:
 292       self.add(t, mark)
 293       self.marks[-1][1]["stroke"] = self["stroke"]
 294
 295   def minitick(self, t): self.tick(t, glyphs.minitick)
 296
 297   def _markorder(self, a, b):
 298     if isinstance(a, (int, long, float)):
 299       posa, marka = a, None
 300     else:
 301       posa, marka = a # marks should be (pos, mark) pairs or just pos
 302     if isinstance(b, (int, long, float)):
 303       posb, markb = b, None
 304     else:
 305       posb, markb = b # marks should be (pos, mark) pairs or just pos
 306
 307     if marka is None and markb is not None: return 1
 308     if marka is not None and markb is None: return -1
 309     return cmp(posa, posb)
 310
 311   def sort(self, order=None):
 312     if order is None: order = lambda a, b: self._markorder(a, b)
 313     self.marks.sort(order)
 314
 315   def closest(self, t, tolerance=None, matching=None):
 316     if tolerance is None: tolerance = trans.epsilon * abs(self.high - self.low)
 317
 318     candidates = []
 319     for item in self.marks:
 320       if isinstance(item, (int, long, float)):
 321         if self._matches(matching, item) and abs(t - item) < tolerance:
 322           candidates.append(item)
 323       else:
 324         pos, mark = item
 325         if self._matches(matching, mark) and abs(t - pos) < tolerance:
 326           candidates.append(item)
 327
 328     def closecmp(a, b):
 329       if isinstance(a, (int, long, float)):
 330         posa, marka = a, None
 331       else:
 332         posa, marka = a # marks should be (pos, mark) pairs or just pos
 333       if isinstance(b, (int, long, float)):
 334         posb, markb = b, None
 335       else:
 336         posb, markb = b # marks should be (pos, mark) pairs or just pos
 337       return cmp(abs(posa - t), abs(posb - t))
 338
 339     candidates.sort(closecmp)
 340     return candidates
 341
 342   def clean_arrows(self):
 343     end1, end2 = None, None
 344     for item in self.marks:
 345       if not isinstance(item, (int, long, float)):
 346         pos, mark = item # marks should be (pos, mark) pairs or just pos
 347         if mark not in (glyphs.tick, glyphs.minitick, glyphs.frtick, glyphs.frminitick) and \
 348            not isinstance(mark, basestring) and not (isinstance(mark, svg.SVG) and self.tag == "text"):
 349           if pos == self.low: end1 = item
 350           if pos == self.high: end2 = item
 351
 352     newmarks = []
 353     for item in self.marks:
 354       if isinstance(item, (int, long, float)):
 355         if (item != self.low and item != self.high) or \
 356            (item == self.low and end1 is None) or \
 357            (item == self.high and end2 is None):
 358           newmarks.append(item)
 359
 360       else:
 361         pos, mark = item
 362         if (mark not in (glyphs.tick, glyphs.minitick, glyphs.frtick, glyphs.frminitick)) or \
 363            (pos != self.low and pos != self.high) or \
 364            (pos == self.low and end1 is None) or \
 365            (pos == self.high and end2 is None):
 366           newmarks.append(item)
 367
 368     self.marks = newmarks
 369
 370   ### act like a list
 371   def append(self, other): self.marks.append(other)
 372   def prepend(self, other): self.marks[0:0] = [other]
 373   def insert(self, i, other): self.marks.insert(i, other)
 374   def remove(self, other): self.marks.remove(other)
 375   def __len__(self): return len(self.marks)
 376
 377   def extend(self, other):
 378     if isinstance(other, SVG):
 379       self.marks.extend(other.children)
 380     elif isinstance(other, basestring):
 381       self.marks.append(other)
 382     else:
 383       self.marks.extend(other)
 384
 385   def __add__(self, other):
 386     output = copy.deepcopy(self)
 387     output += other
 388     return output
 389
 390   def __iadd__(self, other):
 391     self.marks.append(other)
 392     return self
 393
 394   def __mul__(self, other):
 395     output = copy.deepcopy(self)
 396     output *= other
 397     return output
 398
 399   def __rmul__(self, other):
 400     return self * other
 401
 402   def __imul__(self, other):
 403     self.marks *= other
 404     return self
 405
 406   def count(self, *args, **kwds): return self.marks.count(*args, **kwds)
 407   def index(self, *args, **kwds): return self.marks.index(*args, **kwds)
 408   def pop(self, *args, **kwds): return self.marks.pop(*args, **kwds)
 409   def reverse(self, *args, **kwds): return self.marks.reverse(*args, **kwds)
 410
 411 ############################### plot axes
 412
 413 class XAxis(Curve):
 414   text_offsety = 2.5 + 3.
 415   xlogbase = None
 416
 417   _varlist = Curve._varlist + ["xlogbase"]
 418
 419   def _reassign_marks(self):
 420     if self.xlogbase is not None:
 421       return logticks(self.low, self.high)
 422     else:
 423       return ticks(self.low, self.high)
 424
 425   def _reassign_f(self):
 426     output = eval("lambda t: (t, %s)" % repr(self.y))
 427     output.func_name = "x-value"
 428     return output
 429
 430   def __init__(self, low, high, y, **kwds):
 431     self.__dict__["low"] = low
 432     self.__dict__["high"] = high
 433     self.y = y
 434
 435     if "marks" not in kwds:
 436       kwds["marks"] = self._reassign_marks()
 437
 438     Curve.__init__(self, self._reassign_f(), low, high, **kwds)
 439
 440   def _render_text(self, X, Y, angle, text):
 441     text_attrib = {"transform": "translate(%g, %g) rotate(%g)" %
 442                    (X + self.text_offsetx*math.cos(angle) - self.text_offsety*math.sin(angle),
 443                     Y + self.text_offsetx*math.sin(angle) + self.text_offsety*math.cos(angle), 180./math.pi*angle),
 444                    "text-anchor": "middle"}
 445     text_attrib.update(self.text_attrib)
 446     return svg.SVG("text", 0., 0., **text_attrib)(text)
 447
 448   def __setattr__(self, name, value):
 449     self.__dict__[name] = value
 450     if name == "xlogbase":
 451       self.marks = self._reassign_marks()
 452     if name in ("xlogbase", "low", "high"):
 453       self.f = self._reassign_f()
 454
 455 class YAxis(Curve):
 456   text_offsetx = -2.5
 457   text_offsety = 1.5 # when dominant-baseline is implemented everywhere, this hack will no longer be necessary
 458   ylogbase = None
 459
 460   _varlist = Curve._varlist + ["ylogbase"]
 461
 462   def _reassign_marks(self):
 463     if self.ylogbase is not None:
 464       return logticks(self.low, self.high)
 465     else:
 466       return ticks(self.low, self.high)
 467
 468   def _reassign_f(self):
 469     output = eval("lambda t: (t, %s)" % repr(self.y))
 470     output.func_name = "y-value"
 471     return output
 472
 473   def __init__(self, low, high, y, **kwds):
 474     self.__dict__["low"] = low
 475     self.__dict__["high"] = high
 476     self.y = y
 477
 478     if "marks" not in kwds:
 479       kwds["marks"] = self._reassign_marks()
 480
 481     Curve.__init__(self, self._reassign_f(), low, high, **kwds)
 482
 483   def _render_text(self, X, Y, angle, text):
 484     angle += math.pi/2.
 485     text_attrib = {"transform": "translate(%g, %g) rotate(%g)" %
 486                    (X + self.text_offsetx*math.cos(angle) - self.text_offsety*math.sin(angle),
 487                     Y + self.text_offsetx*math.sin(angle) + self.text_offsety*math.cos(angle), 180./math.pi*angle),
 488                    "text-anchor": "end"}
 489     text_attrib.update(self.text_attrib)
 490     return svg.SVG("text", 0., 0., **text_attrib)(text)
 491
 492   def __setattr__(self, name, value):
 493     self.__dict__[name] = value
 494     if name == "ylogbase":
 495       self.marks = self._reassign_marks()
 496     if name in ("ylogbase", "low", "high"):
 497       self.f = self._reassign_f()
 498
 499 ############################### functions for making ticks
 500
 501 def format_number(x, format="%g", scale=1.):
 502   eps = trans.epsilon * abs(scale)
 503   if abs(x) < eps: return "0"
 504   return format % x
 505
 506 def unicode_number(x, scale=1.):
 507   """Converts numbers to a Unicode string, taking advantage of special
 508 Unicode characters to make nice minus signs and scientific notation."""
 509   output = format_number(x, u"%g", scale)
 510
 511   if output[0] == u"-":
 512     output = u"\u2013" + output[1:]
 513
 514   index = output.find(u"e")
 515   if index != -1:
 516     uniout = unicode(output[:index]) + u"\u00d710"
 517     saw_nonzero = False
 518     for n in output[index+1:]:
 519       if n == u"+": pass # uniout += u"\u207a"
 520       elif n == u"-": uniout += u"\u207b"
 521       elif n == u"0":
 522         if saw_nonzero: uniout += u"\u2070"
 523       elif n == u"1":
 524         saw_nonzero = True
 525         uniout += u"\u00b9"
 526       elif n == u"2":
 527         saw_nonzero = True
 528         uniout += u"\u00b2"
 529       elif n == u"3":
 530         saw_nonzero = True
 531         uniout += u"\u00b3"
 532       elif u"4" <= n <= u"9":
 533         saw_nonzero = True
 534         if saw_nonzero: uniout += eval("u\"\\u%x\"" % (0x2070 + ord(n) - ord(u"0")))
 535       else: uniout += n
 536
 537     if uniout[:2] == u"1\u00d7": uniout = uniout[2:]
 538     return uniout
 539
 540   return output
 541
 542 def ticks(low, high, maximum=None, exactly=None, format=unicode_number):
 543   if exactly is not None:
 544     output = []
 545     t = low
 546     for i in xrange(exactly):
 547       output.append((t, glyphs.tick))
 548       output.append((t, format(t, scale=abs(high - low))))
 549       t += (high - low)/(exactly - 1.)
 550     return output
 551
 552   if maximum is None: maximum = 10
 553
 554   counter = 0
 555   granularity = 10**math.ceil(math.log10(max(abs(low), abs(high))))
 556   lowN = math.ceil(1.*low / granularity)
 557   highN = math.floor(1.*high / granularity)
 558
 559   def subdivide(counter, granularity, low, high, lowN, highN):
 560     countermod3 = counter % 3
 561     if countermod3 == 0: granularity *= 0.5
 562     elif countermod3 == 1: granularity *= 0.4
 563     elif countermod3 == 2: granularity *= 0.5
 564     counter += 1
 565     lowN = math.ceil(1.*low / granularity)
 566     highN = math.floor(1.*high / granularity)
 567     return counter, granularity, low, high, lowN, highN
 568
 569   while lowN > highN:
 570     counter, granularity, low, high, lowN, highN = \
 571              subdivide(counter, granularity, low, high, lowN, highN)
 572
 573   last_granularity = granularity
 574   last_trial = None
 575
 576   while True:
 577     trial = []
 578     for n in range(int(lowN), int(highN)+1):
 579       t = n * granularity
 580       trial.append(t)
 581
 582     if len(trial) > maximum:
 583       if last_trial is None:
 584         v1, v2 = low, high
 585         return [(v1, format(v1, scale=abs(high - low))), (v2, format(v2, scale=abs(high - low)))]
 586
 587       else:
 588         if counter % 3 == 2:
 589           counter, granularity, low, high, lowN, highN = \
 590                    subdivide(counter, granularity, low, high, lowN, highN)
 591         trial = []
 592         for n in range(int(lowN), int(highN)+1):
 593           t = n * granularity
 594           trial.append(t)
 595
 596         output = []
 597         for t in last_trial:
 598           output.append((t, glyphs.tick))
 599           output.append((t, format(t, scale=abs(high - low))))
 600         for t in trial:
 601           if t not in last_trial:
 602             output.append((t, glyphs.minitick))
 603         return output
 604
 605     last_granularity = granularity
 606     last_trial = trial
 607
 608     counter, granularity, low, high, lowN, highN = \
 609              subdivide(counter, granularity, low, high, lowN, highN)
 610
 611 def logticks(low, high, base=10., maximum=None, format=unicode_number):
 612   if maximum is None: maximum = 10
 613
 614   lowN = math.floor(math.log(low, base))
 615   highN = math.ceil(math.log(high, base))
 616
 617   trial = []
 618   for n in range(int(lowN), int(highN)+1):
 619     trial.append(base**n)
 620
 621   output = []
 622
 623   # don't need every decade if the ticks cover too many
 624   for i in range(1, len(trial)):
 625     subtrial = trial[::i]
 626     if len(subtrial) <= maximum:
 627       for t in trial:
 628         output.append((t, glyphs.tick))
 629         if t in subtrial:
 630           output.append((t, format(t)))
 631       break
 632
 633   if len(trial) <= 2:
 634     output2 = ticks(low, high, maximum=maximum, format=format)
 635
 636     lowest = min(output2)
 637     for t, mark in output:
 638       if t < lowest: output2.append((t, mark))
 639
 640     return output2
 641
 642   for n in range(int(lowN), int(highN)+1):
 643     t = base**n
 644     for m in range(2, int(math.ceil(base))):
 645       output.append((m * t, glyphs.minitick))
 646
 647   return output
 648