shapes.py 29 KB

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