shapes.py 33 KB

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