shapes.py 34 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037
  1. from numpy import linspace, sin, cos, pi, array, asarray, ndarray, sqrt, abs
  2. import pprint, copy, glob, os
  3. from MatplotlibDraw import MatplotlibDraw
  4. drawing_tool = MatplotlibDraw()
  5. def point(x, y):
  6. return array((x, y), dtype=float)
  7. def animate(fig, time_points, user_action, moviefiles=False,
  8. pause_per_frame=0.5):
  9. if moviefiles:
  10. # Clean up old frame files
  11. framefilestem = 'tmp_frame_'
  12. framefiles = glob.glob('%s*.png' % framefilestem)
  13. for framefile in framefiles:
  14. os.remove(framefile)
  15. for n, t in enumerate(time_points):
  16. drawing_tool.erase()
  17. user_action(t, fig)
  18. #could demand returning fig, but in-place modifications
  19. #are done anyway
  20. #fig = user_action(t, fig)
  21. #if fig is None:
  22. # raise TypeError(
  23. # 'animate: user_action returns None, not fig\n'
  24. # '(a Shape object with the whole figure)')
  25. fig.draw()
  26. drawing_tool.display()
  27. if moviefiles:
  28. drawing_tool.savefig('%s%04d.png' % (framefilestem, n))
  29. if moviefiles:
  30. return '%s*.png' % framefilestem
  31. class Shape:
  32. """
  33. Superclass for drawing different geometric shapes.
  34. Subclasses define shapes, but drawing, rotation, translation,
  35. etc. are done in generic functions in this superclass.
  36. """
  37. def __init__(self):
  38. """
  39. Until new version of shapes.py is ready:
  40. Never to be called from subclasses.
  41. """
  42. raise NotImplementedError(
  43. 'class %s must implement __init__,\nwhich defines '
  44. 'self.shapes as a list of Shape objects\n'
  45. '(and preferably self._repr string).\n'
  46. 'Do not call Shape.__init__!' % \
  47. self.__class__.__name__)
  48. def __iter__(self):
  49. # We iterate over self.shapes many places, and will
  50. # get here if self.shapes is just a Shape object and
  51. # not the assumed list.
  52. print 'Warning: class %s does not define self.shapes\n'\
  53. 'as a *list* of Shape objects'
  54. return [self] # Make the iteration work
  55. def copy(self):
  56. return copy.deepcopy(self)
  57. def __getitem__(self, name):
  58. """
  59. Allow indexing like::
  60. obj1['name1']['name2']
  61. all the way down to ``Curve`` or ``Point`` (``Text``)
  62. objects.
  63. """
  64. if hasattr(self, 'shapes'):
  65. if name in self.shapes:
  66. return self.shapes[name]
  67. else:
  68. for shape in self.shapes:
  69. if isinstance(self.shapes[shape], (Curve,Point)):
  70. # Indexing of Curve/Point/Text is not possible
  71. raise TypeError(
  72. 'Index "%s" is illegal' % name)
  73. return self.shapes[shape][name]
  74. else:
  75. raise Exception('This is a bug')
  76. def for_all_shapes(self, func, *args, **kwargs):
  77. if not hasattr(self, 'shapes'):
  78. # When self.shapes is lacking, we either come to
  79. # a special implementation of func or we come here
  80. # because Shape.func is just inherited. This is
  81. # an error if the class is not Curve or Point
  82. if isinstance(self, (Curve, Point)):
  83. return # ok: no shapes and no func
  84. else:
  85. raise AttributeError('class %s has no shapes attribute!' %
  86. self.__class__.__name__)
  87. is_dict = True if isinstance(self.shapes, dict) else False
  88. for shape in self.shapes:
  89. if is_dict:
  90. shape = self.shapes[shape]
  91. if not isinstance(shape, Shape):
  92. if isinstance(shape, (dict,list,tuple)):
  93. raise TypeError(
  94. 'class %s has a shapes attribute containing '
  95. 'dict/list/tuple objects (nested shapes),\n'
  96. 'which is not allowed - all object must be '
  97. 'derived from Shape and the shapes dict/list\n'
  98. 'cannot be nested.' % self.__class__.__name__)
  99. else:
  100. raise TypeError(
  101. 'class %s has a shapes attribute where not all '
  102. 'values are Shape objects:\n%s' %
  103. (self.__class__.__name__, pprint.pformat(self.shapes)))
  104. getattr(shape, func)(*args, **kwargs)
  105. def draw(self):
  106. self.for_all_shapes('draw')
  107. def rotate(self, angle, center=point(0,0)):
  108. self.for_all_shapes('rotate', angle, center)
  109. def translate(self, vec):
  110. self.for_all_shapes('translate', vec)
  111. def scale(self, factor):
  112. self.for_all_shapes('scale', factor)
  113. def set_linestyle(self, style):
  114. self.for_all_shapes('set_linestyle', style)
  115. def set_linewidth(self, width):
  116. self.for_all_shapes('set_linewidth', width)
  117. def set_linecolor(self, color):
  118. self.for_all_shapes('set_linecolor', color)
  119. def set_arrow(self, style):
  120. self.for_all_shapes('set_arrow', style)
  121. def set_filled_curves(self, color='', pattern=''):
  122. self.for_all_shapes('set_filled_curves', color, pattern)
  123. def show_hierarchy(self, indent=0, format='std'):
  124. """Recursive pretty print of hierarchy of objects."""
  125. if not isinstance(self.shapes, dict):
  126. print 'cannot print hierarchy when %s.shapes is not a dict' % \
  127. self.__class__.__name__
  128. s = ''
  129. if format == 'dict':
  130. s += '{'
  131. for shape in self.shapes:
  132. if format == 'dict':
  133. shape_str = repr(shape) + ':'
  134. elif format == 'plain':
  135. shape_str = shape
  136. else:
  137. shape_str = shape + ':'
  138. if format == 'dict' or format == 'plain':
  139. class_str = ''
  140. else:
  141. class_str = ' (%s)' % \
  142. self.shapes[shape].__class__.__name__
  143. s += '\n%s%s%s %s' % (
  144. ' '*indent,
  145. shape_str,
  146. class_str,
  147. self.shapes[shape].show_hierarchy(indent+4, format))
  148. if format == 'dict':
  149. s += '}'
  150. return s
  151. def __str__(self):
  152. """Display hierarchy with minimum information (just object names)."""
  153. return self.show_hierarchy(format='plain')
  154. def __repr__(self):
  155. """Display hierarchy as a dictionary."""
  156. return self.show_hierarchy(format='dict')
  157. #return pprint.pformat(self.shapes)
  158. class Curve(Shape):
  159. """General curve as a sequence of (x,y) coordintes."""
  160. def __init__(self, x, y):
  161. """
  162. `x`, `y`: arrays holding the coordinates of the curve.
  163. """
  164. self.x, self.y = x, y
  165. # Turn to numpy arrays
  166. self.x = asarray(self.x, dtype=float)
  167. self.y = asarray(self.y, dtype=float)
  168. #self.shapes must not be defined in this class
  169. #as self.shapes holds children objects:
  170. #Curve has no children (end leaf of self.shapes tree)
  171. self.linestyle = None
  172. self.linewidth = None
  173. self.linecolor = None
  174. self.fillcolor = None
  175. self.fillpattern = None
  176. self.arrow = None
  177. def inside_plot_area(self, verbose=True):
  178. """Check that all coordinates are within drawing_tool's area."""
  179. xmin, xmax = self.x.min(), self.x.max()
  180. ymin, ymax = self.y.min(), self.y.max()
  181. t = drawing_tool
  182. inside = True
  183. if xmin < t.xmin:
  184. inside = False
  185. if verbose:
  186. print 'x_min=%g < plot area x_min=%g' % (xmin, t.xmin)
  187. if xmax > t.xmax:
  188. inside = False
  189. if verbose:
  190. print 'x_max=%g > plot area x_max=%g' % (xmax, t.xmax)
  191. if ymin < t.ymin:
  192. inside = False
  193. if verbose:
  194. print 'y_min=%g < plot area y_min=%g' % (ymin, t.ymin)
  195. if xmax > t.xmax:
  196. inside = False
  197. if verbose:
  198. print 'y_max=%g > plot area y_max=%g' % (ymax, t.ymax)
  199. return inside
  200. def draw(self):
  201. """
  202. Send the curve to the plotting engine. That is, convert
  203. coordinate information in self.x and self.y, together
  204. with optional settings of linestyles, etc., to
  205. plotting commands for the chosen engine.
  206. """
  207. self.inside_plot_area()
  208. drawing_tool.define_curve(
  209. self.x, self.y,
  210. self.linestyle, self.linewidth, self.linecolor,
  211. self.arrow, self.fillcolor, self.fillpattern)
  212. def rotate(self, angle, center=point(0,0)):
  213. """
  214. Rotate all coordinates: `angle` is measured in degrees and
  215. (`x`,`y`) is the "origin" of the rotation.
  216. """
  217. angle = angle*pi/180
  218. x, y = center
  219. c = cos(angle); s = sin(angle)
  220. xnew = x + (self.x - x)*c - (self.y - y)*s
  221. ynew = y + (self.x - x)*s + (self.y - y)*c
  222. self.x = xnew
  223. self.y = ynew
  224. def scale(self, factor):
  225. """Scale all coordinates by `factor`: ``x = factor*x``, etc."""
  226. self.x = factor*self.x
  227. self.y = factor*self.y
  228. def translate(self, vec):
  229. """Translate all coordinates by a vector `vec`."""
  230. self.x += vec[0]
  231. self.y += vec[1]
  232. def set_linecolor(self, color):
  233. self.linecolor = color
  234. def set_linewidth(self, width):
  235. self.linewidth = width
  236. def set_linestyle(self, style):
  237. self.linestyle = style
  238. def set_arrow(self, style=None):
  239. styles = ('->', '<-', '<->')
  240. if not style in styles:
  241. raise ValueError('style=%s must be in %s' % (style, styles))
  242. self.arrow = style
  243. def set_name(self, name):
  244. self.name = name
  245. def set_filled_curves(self, color='', pattern=''):
  246. self.fillcolor = color
  247. self.fillpattern = pattern
  248. def show_hierarchy(self, indent=0, format='std'):
  249. if format == 'dict':
  250. return '"%s"' % str(self)
  251. elif format == 'plain':
  252. return ''
  253. else:
  254. return str(self)
  255. def __str__(self):
  256. """Compact pretty print of a Curve object."""
  257. s = '%d coords' % self.x.size
  258. if not self.inside_plot_area(verbose=False):
  259. s += ', some coordinates are outside plotting area!\n'
  260. props = ('linecolor', 'linewidth', 'linestyle', 'arrow',
  261. 'fillcolor', 'fillpattern')
  262. for prop in props:
  263. value = getattr(self, prop)
  264. if value is not None:
  265. s += ' %s=%s' % (prop, repr(value))
  266. return s
  267. def __repr__(self):
  268. return str(self)
  269. class Point(Shape):
  270. """A point (x,y) which can be rotated, translated, and scaled."""
  271. def __init__(self, x, y):
  272. self.x, self.y = x, y
  273. #self.shapes is not needed in this class
  274. def __add__(self, other):
  275. if isinstance(other, (list,tuple)):
  276. other = Point(other)
  277. return Point(self.x+other.x, self.y+other.y)
  278. # class Point is an abstract class - only subclasses are useful
  279. # and must implement draw
  280. def draw(self):
  281. raise NotImplementedError(
  282. 'class %s must implement the draw method' %
  283. self.__class__.__name__)
  284. def rotate(self, angle, center=point(0,0)):
  285. """Rotate point an `angle` (in degrees) around (`x`,`y`)."""
  286. angle = angle*pi/180
  287. x, y = center
  288. c = cos(angle); s = sin(angle)
  289. xnew = x + (self.x - x)*c - (self.y - y)*s
  290. ynew = y + (self.x - x)*s + (self.y - y)*c
  291. self.x = xnew
  292. self.y = ynew
  293. def scale(self, factor):
  294. """Scale point coordinates by `factor`: ``x = factor*x``, etc."""
  295. self.x = factor*self.x
  296. self.y = factor*self.y
  297. def translate(self, vec):
  298. """Translate point by a vector `vec`."""
  299. self.x += vec[0]
  300. self.y += vec[1]
  301. def show_hierarchy(self, indent=0, format='std'):
  302. s = '%s at (%g,%g)' % (self.__class__.__name__, self.x, self.y)
  303. if format == 'dict':
  304. return '"%s"' % s
  305. elif format == 'plain':
  306. return ''
  307. else:
  308. return s
  309. # no need to store input data as they are invalid after rotations etc.
  310. class Rectangle(Shape):
  311. def __init__(self, lower_left_corner, width, height):
  312. ll = lower_left_corner # short form
  313. x = [ll[0], ll[0] + width,
  314. ll[0] + width, ll[0], ll[0]]
  315. y = [ll[1], ll[1], ll[1] + height,
  316. ll[1] + height, ll[1]]
  317. self.shapes = {'rectangle': Curve(x,y)}
  318. class Triangle(Shape):
  319. """Triangle defined by its three vertices p1, p2, and p3."""
  320. def __init__(self, p1, p2, p3):
  321. x = [p1[0], p2[0], p3[0], p1[0]]
  322. y = [p1[1], p2[1], p3[1], p1[1]]
  323. self.shapes = {'triangle': Curve(x,y)}
  324. class Line(Shape):
  325. def __init__(self, start, stop):
  326. x = [start[0], stop[0]]
  327. y = [start[1], stop[1]]
  328. self.shapes = {'line': Curve(x, y)}
  329. self.compute_formulas()
  330. def compute_formulas(self):
  331. x, y = self.shapes['line'].x, self.shapes['line'].y
  332. # Define equations for line:
  333. # y = a*x + b, x = c*y + d
  334. try:
  335. self.a = (y[1] - y[0])/(x[1] - x[0])
  336. self.b = y[0] - self.a*x[0]
  337. except ZeroDivisionError:
  338. # Vertical line, y is not a function of x
  339. self.a = None
  340. self.b = None
  341. try:
  342. if self.a is None:
  343. self.c = 0
  344. else:
  345. self.c = 1/float(self.a)
  346. if self.b is None:
  347. self.d = x[1]
  348. except ZeroDivisionError:
  349. # Horizontal line, x is not a function of y
  350. self.c = None
  351. self.d = None
  352. def compute_formulas(self):
  353. x, y = self.shapes['line'].x, self.shapes['line'].y
  354. tol = 1E-14
  355. # Define equations for line:
  356. # y = a*x + b, x = c*y + d
  357. if abs(x[1] - x[0]) > tol:
  358. self.a = (y[1] - y[0])/(x[1] - x[0])
  359. self.b = y[0] - self.a*x[0]
  360. else:
  361. # Vertical line, y is not a function of x
  362. self.a = None
  363. self.b = None
  364. if self.a is None:
  365. self.c = 0
  366. elif abs(self.a) > tol:
  367. self.c = 1/float(self.a)
  368. self.d = x[1]
  369. else: # self.a is 0
  370. # Horizontal line, x is not a function of y
  371. self.c = None
  372. self.d = None
  373. def __call__(self, x=None, y=None):
  374. """Given x, return y on the line, or given y, return x."""
  375. self.compute_formulas()
  376. if x is not None and self.a is not None:
  377. return self.a*x + self.b
  378. elif y is not None and self.c is not None:
  379. return self.c*y + self.d
  380. else:
  381. raise ValueError(
  382. 'Line.__call__(x=%s, y=%s) not meaningful' % \
  383. (x, y))
  384. # First implementation of class Circle
  385. class Circle(Shape):
  386. def __init__(self, center, radius, resolution=180):
  387. self.center, self.radius = center, radius
  388. self.resolution = resolution
  389. t = linspace(0, 2*pi, resolution+1)
  390. x0 = center[0]; y0 = center[1]
  391. R = radius
  392. x = x0 + R*cos(t)
  393. y = y0 + R*sin(t)
  394. self.shapes = {'circle': Curve(x, y)}
  395. def __call__(self, theta):
  396. """Return (x, y) point corresponding to theta."""
  397. return self.center[0] + self.radius*cos(theta), \
  398. self.center[1] + self.radius*sin(theta)
  399. class Arc(Shape):
  400. def __init__(self, center, radius,
  401. start_degrees, opening_degrees,
  402. resolution=180):
  403. self.center = center
  404. self.radius = radius
  405. self.start_degrees = start_degrees*pi/180 # radians
  406. self.opening_degrees = opening_degrees*pi/180
  407. self.resolution = resolution
  408. t = linspace(self.start_degrees,
  409. self.start_degrees + self.opening_degrees,
  410. resolution+1)
  411. x0 = center[0]; y0 = center[1]
  412. R = radius
  413. x = x0 + R*cos(t)
  414. y = y0 + R*sin(t)
  415. self.shapes = {'arc': Curve(x, y)}
  416. def __call__(self, theta):
  417. """Return (x,y) point at start_degrees + theta."""
  418. theta = theta*pi/180
  419. t = self.start_degrees + theta
  420. x0 = self.center[0]
  421. y0 = self.center[1]
  422. R = self.radius
  423. x = x0 + R*cos(t)
  424. y = y0 + R*sin(t)
  425. return (x, y)
  426. # Alternative for small arcs: Parabola
  427. class Parabola(Shape):
  428. def __init__(self, start, mid, stop, resolution=21):
  429. self.p1, self.p2, self.p3 = start, mid, stop
  430. # y as function of x? (no point on line x=const?)
  431. tol = 1E-14
  432. if abs(self.p1[0] - self.p2[0]) > 1E-14 and \
  433. abs(self.p2[0] - self.p3[0]) > 1E-14 and \
  434. abs(self.p3[0] - self.p1[0]) > 1E-14:
  435. self.y_of_x = True
  436. else:
  437. self.y_of_x = False
  438. # x as function of y? (no point on line y=const?)
  439. tol = 1E-14
  440. if abs(self.p1[1] - self.p2[1]) > 1E-14 and \
  441. abs(self.p2[1] - self.p3[1]) > 1E-14 and \
  442. abs(self.p3[1] - self.p1[1]) > 1E-14:
  443. self.x_of_y = True
  444. else:
  445. self.x_of_y = False
  446. if self.y_of_x:
  447. x = linspace(start[0], end[0], resolution)
  448. y = self(x=x)
  449. elif self.x_of_y:
  450. y = linspace(start[1], end[1], resolution)
  451. x = self(y=y)
  452. else:
  453. raise ValueError(
  454. 'Parabola: two or more points lie on x=const '
  455. 'or y=const - not allowed')
  456. self.shapes = {'parabola': Curve(x, y)}
  457. def __call__(self, x=None, y=None):
  458. if x is not None and self.y_of_x:
  459. return self._L2x(self.p1, self.p2)*self.p3[1] + \
  460. self._L2x(self.p2, self.p3)*self.p1[1] + \
  461. self._L2x(self.p3, self.p1)*self.p2[1]
  462. elif y is not None and self.x_of_y:
  463. return self._L2y(self.p1, self.p2)*self.p3[0] + \
  464. self._L2y(self.p2, self.p3)*self.p1[0] + \
  465. self._L2y(self.p3, self.p1)*self.p2[0]
  466. else:
  467. raise ValueError(
  468. 'Parabola.__call__(x=%s, y=%s) not meaningful' % \
  469. (x, y))
  470. def _L2x(self, x, pi, pj, pk):
  471. return (x - pi[0])*(x - pj[0])/((pk[0] - pi[0])*(pk[0] - pj[0]))
  472. def _L2y(self, y, pi, pj, pk):
  473. return (y - pi[1])*(y - pj[1])/((pk[1] - pi[1])*(pk[1] - pj[1]))
  474. class Circle(Arc):
  475. def __init__(self, center, radius, resolution=180):
  476. Arc.__init__(self, center, radius, 0, 360, resolution)
  477. # class Wall: horizontal Line with many small Lines 45 degrees
  478. class XWall(Shape):
  479. def __init__(start, length, dx, below=True):
  480. n = int(round(length/float(dx))) # no of intervals
  481. x = linspace(start[0], start[0] + length, n+1)
  482. y = start[1]
  483. dy = dx
  484. if below:
  485. taps = [Line((xi,y-dy), (xi+dx, y)) for xi in x[:-1]]
  486. else:
  487. taps = [Line((xi,y), (xi+dx, y+dy)) for xi in x[:-1]]
  488. self.shapes = [Line(start, (start[0]+length, start[1]))] + taps
  489. class Wall(Shape):
  490. def __init__(self, start, length, thickness, rotation_angle=0):
  491. p1 = asarray(start)
  492. p2 = p1 + asarray([length, 0])
  493. p3 = p2 + asarray([0, thickness])
  494. p4 = p1 + asarray([0, thickness])
  495. p5 = p1
  496. x = [p[0] for p in p1, p2, p3, p4, p5]
  497. y = [p[1] for p in p1, p2, p3, p4, p5]
  498. wall = Curve(x, y)
  499. wall.set_filled_curves('white', '/')
  500. wall.rotate(rotation_angle, start)
  501. self.shapes = {'wall': wall}
  502. """
  503. def draw(self):
  504. x = self.shapes['wall'].x
  505. y = self.shapes['wall'].y
  506. drawing_tool.ax.fill(x, y, 'w',
  507. edgecolor=drawing_tool.linecolor,
  508. hatch='/')
  509. """
  510. class CurveWall(Shape):
  511. def __init__(self, x, y, thickness):
  512. x1 = asarray(x, float)
  513. y1 = asarray(y, float)
  514. x2 = x1
  515. y2 = y1 + thickness
  516. from numpy import concatenate
  517. # x1/y1 + reversed x2/y2
  518. x = concatenate((x1, x2[-1::-1]))
  519. y = concatenate((y1, y2[-1::-1]))
  520. wall = Curve(x, y)
  521. wall.set_filled_curves('white', '/')
  522. self.shapes = {'wall': wall}
  523. """
  524. def draw(self):
  525. x = self.shapes['wall'].x
  526. y = self.shapes['wall'].y
  527. drawing_tool.ax.fill(x, y, 'w',
  528. edgecolor=drawing_tool.linecolor,
  529. hatch='/')
  530. """
  531. class Text(Point):
  532. def __init__(self, text, position, alignment='center', fontsize=18):
  533. self.text = text
  534. self.alignment, self.fontsize = alignment, fontsize
  535. is_sequence(position, length=2, can_be_None=True)
  536. Point.__init__(self, position[0], position[1])
  537. #no need for self.shapes here
  538. def draw(self):
  539. drawing_tool.text(self.text, (self.x, self.y),
  540. self.alignment, self.fontsize)
  541. def __str__(self):
  542. return 'text "%s" at (%g,%g)' % (self.text, self.x, self.y)
  543. def __repr__(self):
  544. return str(self)
  545. class Text_wArrow(Text):
  546. def __init__(self, text, position, arrow_tip,
  547. alignment='center', fontsize=18):
  548. is_sequence(arrow_tip, length=2, can_be_None=True)
  549. self.arrow_tip = arrow_tip
  550. Text.__init__(self, text, position, alignment, fontsize)
  551. def draw(self):
  552. drawing_tool.text(self.text, self.position,
  553. self.alignment, self.fontsize,
  554. self.arrow_tip)
  555. def __str__(self):
  556. return 'annotation "%s" at (%g,%g) with arrow to (%g,%g)' % \
  557. (self.text, self.x, self.y,
  558. self.arrow_tip[0], self.arrow_tip[1])
  559. def __repr__(self):
  560. return str(self)
  561. class Axis(Shape):
  562. def __init__(self, bottom_point, length, label, below=True,
  563. rotation_angle=0, label_spacing=1./25):
  564. """
  565. Draw axis from bottom_point with `length` to the right
  566. (x axis). Place label below (True) or above (False) axis.
  567. Then return `rotation_angle` (in degrees).
  568. To make a standard x axis, call with ``below=True`` and
  569. ``rotation_angle=0``. To make a standard y axis, call with
  570. ``below=False`` and ``rotation_angle=90``.
  571. A tilted axis can also be drawn.
  572. The `label_spacing` denotes the space between the symbol
  573. and the arrow tip as a fraction of the length of the plot
  574. in x direction.
  575. """
  576. # Arrow is vertical arrow, make it horizontal
  577. arrow = Arrow(bottom_point, length, rotation_angle=-90)
  578. arrow.rotate(rotation_angle, bottom_point)
  579. spacing = drawing_tool.xrange*label_spacing
  580. if below:
  581. spacing = - spacing
  582. label_pos = [bottom_point[0] + length, bottom_point[1] + spacing]
  583. symbol = Text(label, position=label_pos)
  584. symbol.rotate(rotation_angle, bottom_point)
  585. self.shapes = {'arrow': arrow, 'symbol': symbol}
  586. class Gravity(Axis):
  587. """Downward-pointing gravity arrow with the symbol g."""
  588. def __init__(self, start, length):
  589. Axis.__init__(self, start, length, '$g$', below=False,
  590. rotation_angle=-90, label_spacing=1./30)
  591. def test_Axis():
  592. set_coordinate_system(xmin=0, xmax=15, ymin=0, ymax=15, axis=True)
  593. x_axis = Axis((7.5,2), 5, 'x', rotation_angle=0)
  594. y_axis = Axis((7.5,2), 5, 'y', below=False, rotation_angle=90)
  595. system = Compose({'x axis': x_axis, 'y axis': y_axis})
  596. system.draw()
  597. drawing_tool.display()
  598. set_linestyle('dashed')
  599. system.shapes['x axis'].rotate(40, (7.5, 2))
  600. system.shapes['y axis'].rotate(40, (7.5, 2))
  601. system.draw()
  602. drawing_tool.display()
  603. print repr(system)
  604. class DistanceSymbol(Shape):
  605. """
  606. Arrow with symbol at the midpoint,
  607. for identifying a distance with a symbol.
  608. """
  609. def __init__(self, start, end, symbol, fontsize=14):
  610. start = asarray(start, float)
  611. end = asarray(end, float)
  612. mid = 0.5*(start + end) # midpoint of start-end line
  613. tangent = end - start
  614. normal = asarray([-tangent[1], tangent[0]])/\
  615. sqrt(tangent[0]**2 + tangent[1]**2)
  616. symbol_pos = mid + normal*drawing_tool.xrange/60.
  617. self.shapes = {'arrow': Arrow1(start, end, style='<->'),
  618. 'symbol': Text(symbol, symbol_pos, fontsize=fontsize)}
  619. class ArcSymbol(Shape):
  620. def __init__(self, symbol, center, radius,
  621. start_degrees, opening_degrees,
  622. resolution=180, fontsize=14):
  623. arc = Arc(center, radius, start_degrees, opening_degrees,
  624. resolution)
  625. mid = asarray(arc(opening_degrees/2.))
  626. normal = mid - asarray(center, float)
  627. normal = normal/sqrt(normal[0]**2 + normal[1]**2)
  628. symbol_pos = mid + normal*drawing_tool.xrange/60.
  629. self.shapes = {'arc': arc,
  630. 'symbol': Text(symbol, symbol_pos, fontsize=fontsize)}
  631. class Compose(Shape):
  632. def __init__(self, shapes):
  633. """shapes: list or dict of Shape objects."""
  634. self.shapes = shapes
  635. # can make help methods: Line.midpoint, Line.normal(pt, dir='left') -> (x,y)
  636. # list annotations in each class? contains extra annotations for explaining
  637. # important parameters to the constructor, e.g., Line.annotations holds
  638. # start and end as Text objects. Shape.demo calls shape.draw and
  639. # for annotation in self.demo: annotation.draw() YES!
  640. # Can make overall demo of classes by making objects and calling demo
  641. # Could include demo fig in each constructor
  642. class Arrow1(Shape):
  643. """Draw an arrow as Line with arrow."""
  644. def __init__(self, start, end, style='->'):
  645. self.start, self.end, self.style = start, end, style
  646. self.shapes = {'arrow': Line(start, end, arrow=style)}
  647. class Arrow3(Shape):
  648. """Draw a vertical line and arrow head. Then rotate `rotation_angle`."""
  649. def __init__(self, bottom_point, length, rotation_angle=0):
  650. self.bottom = bottom_point
  651. self.length = length
  652. self.angle = rotation_angle
  653. top = (self.bottom[0], self.bottom[1] + self.length)
  654. main = Line(self.bottom, top)
  655. #head_length = self.length/8.0
  656. head_length = drawing_tool.xrange/50.
  657. head_degrees = 30*pi/180
  658. head_left_pt = (top[0] - head_length*sin(head_degrees),
  659. top[1] - head_length*cos(head_degrees))
  660. head_right_pt = (top[0] + head_length*sin(head_degrees),
  661. top[1] - head_length*cos(head_degrees))
  662. head_left = Line(head_left_pt, top)
  663. head_right = Line(head_right_pt, top)
  664. head_left.set_linestyle('solid')
  665. head_right.set_linestyle('solid')
  666. self.shapes = {'line': main, 'head left': head_left,
  667. 'head right': head_right}
  668. # rotate goes through self.shapes so this must be initialized first
  669. self.rotate(rotation_angle, bottom_point)
  670. Arrow = Arrow3 # backward compatibility
  671. class Wheel(Shape):
  672. def __init__(self, center, radius, inner_radius=None, nlines=10):
  673. self.center = center
  674. self.radius = radius
  675. if inner_radius is None:
  676. self.inner_radius = radius/5.0
  677. else:
  678. self.inner_radius = inner_radius
  679. self.nlines = nlines
  680. outer = Circle(self.center, self.radius)
  681. inner = Circle(self.center, self.inner_radius)
  682. lines = []
  683. # Draw nlines+1 since the first and last coincide
  684. # (then nlines lines will be visible)
  685. t = linspace(0, 2*pi, self.nlines+1)
  686. Ri = self.inner_radius; Ro = self.radius
  687. x0 = self.center[0]; y0 = self.center[1]
  688. xinner = x0 + Ri*cos(t)
  689. yinner = y0 + Ri*sin(t)
  690. xouter = x0 + Ro*cos(t)
  691. youter = y0 + Ro*sin(t)
  692. lines = [Line((xi,yi),(xo,yo)) for xi, yi, xo, yo in \
  693. zip(xinner, yinner, xouter, youter)]
  694. self.shapes = [outer, inner] + lines
  695. class Wave(Shape):
  696. def __init__(self, xstart, xstop,
  697. wavelength, amplitude, mean_level):
  698. self.xstart = xstart
  699. self.xstop = xstop
  700. self.wavelength = wavelength
  701. self.amplitude = amplitude
  702. self.mean_level = mean_level
  703. npoints = (self.xstop - self.xstart)/(self.wavelength/61.0)
  704. x = linspace(self.xstart, self.xstop, npoints)
  705. k = 2*pi/self.wavelength # frequency
  706. y = self.mean_level + self.amplitude*sin(k*x)
  707. self.shapes = {'waves': Curve(x,y)}
  708. # make a version of Spring using Point class
  709. class Spring(Shape):
  710. def __init__(self, bottom_point, length, tagwidth, ntags=4):
  711. """
  712. Specify a vertical spring, starting at bottom_point and
  713. having a specified lengths. In the middle third of the
  714. spring there are ntags tags.
  715. """
  716. self.B = bottom_point
  717. self.n = ntags - 1 # n counts tag intervals
  718. # n must be odd:
  719. if self.n % 2 == 0:
  720. self.n = self.n+1
  721. self.L = length
  722. self.w = tagwidth
  723. B, L, n, w = self.B, self.L, self.n, self.w # short forms
  724. t = L/(3.0*n) # must be better worked out
  725. P0 = (B[0], B[1]+L/3.0)
  726. P1 = (B[0], B[1]+L/3.0+t/2.0)
  727. P2 = (B[0], B[1]+L*2/3.0)
  728. P3 = (B[0], B[1]+L)
  729. line1 = Line(B, P1)
  730. lines = [line1]
  731. #line2 = Line(P2, P3)
  732. T1 = P1
  733. T2 = (T1[0] + w, T1[1] + t/2.0)
  734. lines.append(Line(T1,T2))
  735. T1 = (T2[0], T2[1])
  736. for i in range(n):
  737. T2 = (T1[0] + (-1)**(i+1)*2*w, T1[1] + t/2.0)
  738. lines.append(Line(T1, T2))
  739. T1 = (T2[0], T2[1])
  740. T2 = (T1[0] + w, T1[1] + t/2.0)
  741. lines.append(Line(T1,T2))
  742. #print P2, T2
  743. lines.append(Line(T2, P3))
  744. self.shapes = lines
  745. # COMPOSITE types:
  746. # MassSpringForce: Line(horizontal), Spring, Rectangle, Arrow/Line(w/arrow)
  747. # must be easy to find the tip of the arrow
  748. # Maybe extra dict: self.name['mass'] = Rectangle object - YES!
  749. def _test1():
  750. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  751. l1 = Line((0,0), (1,1))
  752. l1.draw()
  753. input(': ')
  754. c1 = Circle((5,2), 1)
  755. c2 = Circle((6,2), 1)
  756. w1 = Wheel((7,2), 1)
  757. c1.draw()
  758. c2.draw()
  759. w1.draw()
  760. hardcopy()
  761. display() # show the plot
  762. def _test2():
  763. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  764. l1 = Line((0,0), (1,1))
  765. l1.draw()
  766. input(': ')
  767. c1 = Circle((5,2), 1)
  768. c2 = Circle((6,2), 1)
  769. w1 = Wheel((7,2), 1)
  770. filled_curves(True)
  771. set_linecolor('blue')
  772. c1.draw()
  773. set_linecolor('aqua')
  774. c2.draw()
  775. filled_curves(False)
  776. set_linecolor('red')
  777. w1.draw()
  778. hardcopy()
  779. display() # show the plot
  780. def _test3():
  781. """Test example from the book."""
  782. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  783. l1 = Line(start=(0,0), stop=(1,1)) # define line
  784. l1.draw() # make plot data
  785. r1 = Rectangle(lower_left_corner=(0,1), width=3, height=5)
  786. r1.draw()
  787. Circle(center=(5,7), radius=1).draw()
  788. Wheel(center=(6,2), radius=2, inner_radius=0.5, nlines=7).draw()
  789. hardcopy()
  790. display()
  791. def _test4():
  792. """Second example from the book."""
  793. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  794. r1 = Rectangle(lower_left_corner=(0,1), width=3, height=5)
  795. c1 = Circle(center=(5,7), radius=1)
  796. w1 = Wheel(center=(6,2), radius=2, inner_radius=0.5, nlines=7)
  797. c2 = Circle(center=(7,7), radius=1)
  798. filled_curves(True)
  799. c1.draw()
  800. set_linecolor('blue')
  801. r1.draw()
  802. set_linecolor('aqua')
  803. c2.draw()
  804. # Add thick aqua line around rectangle:
  805. filled_curves(False)
  806. set_linewidth(4)
  807. r1.draw()
  808. set_linecolor('red')
  809. # Draw wheel with thick lines:
  810. w1.draw()
  811. hardcopy('tmp_colors')
  812. display()
  813. def _test5():
  814. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  815. c = 6. # center point of box
  816. w = 2. # size of box
  817. L = 3
  818. r1 = Rectangle((c-w/2, c-w/2), w, w)
  819. l1 = Line((c,c-w/2), (c,c-w/2-L))
  820. linecolor('blue')
  821. filled_curves(True)
  822. r1.draw()
  823. linecolor('aqua')
  824. filled_curves(False)
  825. l1.draw()
  826. hardcopy()
  827. display() # show the plot
  828. def rolling_wheel(total_rotation_angle):
  829. """Animation of a rotating wheel."""
  830. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  831. import time
  832. center = (6,2)
  833. radius = 2.0
  834. angle = 2.0
  835. pngfiles = []
  836. w1 = Wheel(center=center, radius=radius, inner_radius=0.5, nlines=7)
  837. for i in range(int(total_rotation_angle/angle)):
  838. w1.draw()
  839. print 'XXXXXXXXXXXXXXXXXXXXXX BIG PROBLEM WITH ANIMATE!!!'
  840. display()
  841. filename = 'tmp_%03d' % i
  842. pngfiles.append(filename + '.png')
  843. hardcopy(filename)
  844. time.sleep(0.3) # pause
  845. L = radius*angle*pi/180 # translation = arc length
  846. w1.rotate(angle, center)
  847. w1.translate((-L, 0))
  848. center = (center[0] - L, center[1])
  849. erase()
  850. cmd = 'convert -delay 50 -loop 1000 %s tmp_movie.gif' \
  851. % (' '.join(pngfiles))
  852. print 'converting PNG files to animated GIF:\n', cmd
  853. import commands
  854. failure, output = commands.getstatusoutput(cmd)
  855. if failure: print 'Could not run', cmd
  856. def is_sequence(seq, length=None,
  857. can_be_None=False, error_message=True):
  858. if can_be_None:
  859. legal_types = (list,tuple,ndarray,None)
  860. else:
  861. legal_types = (list,tuple,ndarray)
  862. if isinstance(seq, legal_types):
  863. if length is not None:
  864. if length == len(seq):
  865. return True
  866. elif error_message:
  867. raise TypeError('%s is %s; must be %s of length %d' %
  868. (str(point), type(point),
  869. ', '.join([str(t) for t in legal_types]),
  870. len(seq)))
  871. else:
  872. return False
  873. else:
  874. return True
  875. elif error_message:
  876. raise TypeError('%s is %s; must be %s' %
  877. str(point), type(point),
  878. ', '.join([str(t) for t in legal_types]))
  879. else:
  880. return False
  881. if __name__ == '__main__':
  882. #rolling_wheel(40)
  883. #_test1()
  884. #_test3()
  885. funcs = [
  886. #test_Axis,
  887. test_inclined_plane,
  888. ]
  889. for func in funcs:
  890. func()
  891. raw_input('Type Return: ')