shapes.py 55 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604
  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, check_inside=False):
  6. if isinstance(x, (float,int)) and isinstance(y, (float,int)):
  7. pass
  8. else:
  9. raise TypeError('x=%s,y=%s must be float,float, not %s,%s' %
  10. (x, y, type(x), type(y)))
  11. if check_inside:
  12. ok, msg = drawing_tool.inside((x,y), exception=True)
  13. if not ok:
  14. print msg
  15. return array((x, y), dtype=float)
  16. def arr2D(x, check_inside=False):
  17. if isinstance(x, (tuple,list,ndarray)):
  18. if len(x) == 2:
  19. pass
  20. else:
  21. raise ValueError('x=%s has length %d, not 2' % (x, len(x)))
  22. else:
  23. raise TypeError('x=%s must be list/tuple/ndarray, not %s' %
  24. (x, type(x)))
  25. if check_inside:
  26. ok, msg = drawing_tool.inside(x, exception=True)
  27. if not ok:
  28. print msg
  29. return asarray(x, dtype=float)
  30. def _is_sequence(seq, length=None,
  31. can_be_None=False, error_message=True):
  32. if can_be_None:
  33. legal_types = (list,tuple,ndarray,None)
  34. else:
  35. legal_types = (list,tuple,ndarray)
  36. if isinstance(seq, legal_types):
  37. if length is not None:
  38. if length == len(seq):
  39. return True
  40. elif error_message:
  41. raise TypeError('%s is %s; must be %s of length %d' %
  42. (str(seq), type(seq),
  43. ', '.join([str(t) for t in legal_types]),
  44. len(seq)))
  45. else:
  46. return False
  47. else:
  48. return True
  49. elif error_message:
  50. raise TypeError('%s is %s; must be %s' %
  51. (str(seq), type(seq),
  52. ','.join([str(t)[5:-1] for t in legal_types])))
  53. else:
  54. return False
  55. def is_sequence(*sequences, **kwargs):
  56. length = kwargs.get('length', 2)
  57. can_be_None = kwargs.get('can_be_None', False)
  58. error_message = kwargs.get('error_message', True)
  59. check_inside = kwargs.get('check_inside', False)
  60. for x in sequences:
  61. _is_sequence(x, length=length, can_be_None=can_be_None,
  62. error_message=error_message)
  63. if check_inside:
  64. ok, msg = drawing_tool.inside(x, exception=True)
  65. if not ok:
  66. print msg
  67. def animate(fig, time_points, user_action, moviefiles=False,
  68. pause_per_frame=0.5):
  69. if moviefiles:
  70. # Clean up old frame files
  71. framefilestem = 'tmp_frame_'
  72. framefiles = glob.glob('%s*.png' % framefilestem)
  73. for framefile in framefiles:
  74. os.remove(framefile)
  75. for n, t in enumerate(time_points):
  76. drawing_tool.erase()
  77. user_action(t, fig)
  78. #could demand returning fig, but in-place modifications
  79. #are done anyway
  80. #fig = user_action(t, fig)
  81. #if fig is None:
  82. # raise TypeError(
  83. # 'animate: user_action returns None, not fig\n'
  84. # '(a Shape object with the whole figure)')
  85. fig.draw()
  86. drawing_tool.display()
  87. if moviefiles:
  88. drawing_tool.savefig('%s%04d.png' % (framefilestem, n))
  89. if moviefiles:
  90. return '%s*.png' % framefilestem
  91. class Shape:
  92. """
  93. Superclass for drawing different geometric shapes.
  94. Subclasses define shapes, but drawing, rotation, translation,
  95. etc. are done in generic functions in this superclass.
  96. """
  97. def __init__(self):
  98. """
  99. Until new version of shapes.py is ready:
  100. Never to be called from subclasses.
  101. """
  102. raise NotImplementedError(
  103. 'class %s must implement __init__,\nwhich defines '
  104. 'self.shapes as a list of Shape objects\n'
  105. '(and preferably self._repr string).\n'
  106. 'Do not call Shape.__init__!' % \
  107. self.__class__.__name__)
  108. def __iter__(self):
  109. # We iterate over self.shapes many places, and will
  110. # get here if self.shapes is just a Shape object and
  111. # not the assumed list.
  112. print 'Warning: class %s does not define self.shapes\n'\
  113. 'as a *list* of Shape objects'
  114. return [self] # Make the iteration work
  115. def copy(self):
  116. return copy.deepcopy(self)
  117. def __getitem__(self, name):
  118. """
  119. Allow indexing like::
  120. obj1['name1']['name2']
  121. all the way down to ``Curve`` or ``Point`` (``Text``)
  122. objects.
  123. """
  124. if hasattr(self, 'shapes'):
  125. if name in self.shapes:
  126. return self.shapes[name]
  127. else:
  128. for shape in self.shapes:
  129. if isinstance(self.shapes[shape], (Curve,Point)):
  130. # Indexing of Curve/Point/Text is not possible
  131. raise TypeError(
  132. 'Index "%s" (%s) is illegal' %
  133. (name, self.__class__.__name__))
  134. return self.shapes[shape][name]
  135. else:
  136. raise Exception('This is a bug')
  137. def _for_all_shapes(self, func, *args, **kwargs):
  138. if not hasattr(self, 'shapes'):
  139. # When self.shapes is lacking, we either come to
  140. # a special implementation of func or we come here
  141. # because Shape.func is just inherited. This is
  142. # an error if the class is not Curve or Point
  143. if isinstance(self, (Curve, Point)):
  144. return # ok: no shapes and no func
  145. else:
  146. raise AttributeError('class %s has no shapes attribute!' %
  147. self.__class__.__name__)
  148. is_dict = True if isinstance(self.shapes, dict) else False
  149. for shape in self.shapes:
  150. if is_dict:
  151. shape = self.shapes[shape]
  152. if not isinstance(shape, Shape):
  153. if isinstance(shape, dict):
  154. raise TypeError(
  155. 'class %s has a shapes attribute that is just\n'
  156. 'a plain dictionary,\n%s\n'
  157. 'Did you mean to embed this dict in a Compose\n'
  158. 'object?' % (self.__class__.__name__,
  159. str(shape)))
  160. elif isinstance(shape, (list,tuple)):
  161. raise TypeError(
  162. 'class %s has a shapes attribute containing\n'
  163. 'a %s object %s,\n'
  164. 'Did you mean to embed this list in a Compose\n'
  165. 'object?' % (self.__class__.__name__,
  166. type(shape), str(shape)))
  167. else:
  168. raise TypeError(
  169. 'class %s has a shapes attribute %s which is not'
  170. 'a Shape objects\n%s' %
  171. (self.__class__.__name__, type(shape),
  172. pprint.pformat(self.shapes)))
  173. getattr(shape, func)(*args, **kwargs)
  174. def draw(self):
  175. self._for_all_shapes('draw')
  176. return self
  177. def draw_dimensions(self):
  178. if hasattr(self, 'dimensions'):
  179. for shape in self.dimensions:
  180. self.dimensions[shape].draw()
  181. return self
  182. else:
  183. #raise AttributeError('no self.dimensions dict for defining dimensions of class %s' % self.__classname__.__name__)
  184. return self
  185. def rotate(self, angle, center):
  186. is_sequence(center, length=2)
  187. self._for_all_shapes('rotate', angle, center)
  188. return self
  189. def translate(self, vec):
  190. is_sequence(vec, length=2)
  191. self._for_all_shapes('translate', vec)
  192. return self
  193. def scale(self, factor):
  194. self._for_all_shapes('scale', factor)
  195. return self
  196. def set_linestyle(self, style):
  197. styles = ('solid', 'dashed', 'dashdot', 'dotted')
  198. if style not in styles:
  199. raise ValueError('%s: style=%s must be in %s' %
  200. (self.__class__.__name__ + '.set_linestyle:',
  201. style, str(styles)))
  202. self._for_all_shapes('set_linestyle', style)
  203. return self
  204. def set_linewidth(self, width):
  205. if not isinstance(width, int) and width >= 0:
  206. raise ValueError('%s: width=%s must be positive integer' %
  207. (self.__class__.__name__ + '.set_linewidth:',
  208. width))
  209. self._for_all_shapes('set_linewidth', width)
  210. return self
  211. def set_linecolor(self, color):
  212. if color in drawing_tool.line_colors:
  213. color = drawing_tool.line_colors[color]
  214. elif color in drawing_tool.line_colors.values():
  215. pass # color is ok
  216. else:
  217. raise ValueError('%s: invalid color "%s", must be in %s' %
  218. (self.__class__.__name__ + '.set_linecolor:',
  219. color, list(drawing_tool.line_colors.keys())))
  220. self._for_all_shapes('set_linecolor', color)
  221. return self
  222. def set_arrow(self, style):
  223. styles = ('->', '<-', '<->')
  224. if not style in styles:
  225. raise ValueError('%s: style=%s must be in %s' %
  226. (self.__class__.__name__ + '.set_arrow:',
  227. style, styles))
  228. self._for_all_shapes('set_arrow', style)
  229. return self
  230. def set_filled_curves(self, color='', pattern=''):
  231. if color in drawing_tool.line_colors:
  232. color = drawing_tool.line_colors[color]
  233. elif color in drawing_tool.line_colors.values():
  234. pass # color is ok
  235. else:
  236. raise ValueError('%s: invalid color "%s", must be in %s' %
  237. (self.__class__.__name__ + '.set_filled_curves:',
  238. color, list(drawing_tool.line_colors.keys())))
  239. self._for_all_shapes('set_filled_curves', color, pattern)
  240. return self
  241. def show_hierarchy(self, indent=0, format='std'):
  242. """Recursive pretty print of hierarchy of objects."""
  243. if not isinstance(self.shapes, dict):
  244. print 'cannot print hierarchy when %s.shapes is not a dict' % \
  245. self.__class__.__name__
  246. s = ''
  247. if format == 'dict':
  248. s += '{'
  249. for shape in self.shapes:
  250. if format == 'dict':
  251. shape_str = repr(shape) + ':'
  252. elif format == 'plain':
  253. shape_str = shape
  254. else:
  255. shape_str = shape + ':'
  256. if format == 'dict' or format == 'plain':
  257. class_str = ''
  258. else:
  259. class_str = ' (%s)' % \
  260. self.shapes[shape].__class__.__name__
  261. s += '\n%s%s%s %s' % (
  262. ' '*indent,
  263. shape_str,
  264. class_str,
  265. self.shapes[shape].show_hierarchy(indent+4, format))
  266. if format == 'dict':
  267. s += '}'
  268. return s
  269. def __str__(self):
  270. """Display hierarchy with minimum information (just object names)."""
  271. return self.show_hierarchy(format='plain')
  272. def __repr__(self):
  273. """Display hierarchy as a dictionary."""
  274. return self.show_hierarchy(format='dict')
  275. #return pprint.pformat(self.shapes)
  276. class Curve(Shape):
  277. """General curve as a sequence of (x,y) coordintes."""
  278. def __init__(self, x, y):
  279. """
  280. `x`, `y`: arrays holding the coordinates of the curve.
  281. """
  282. self.x = asarray(x, dtype=float)
  283. self.y = asarray(y, dtype=float)
  284. #self.shapes must not be defined in this class
  285. #as self.shapes holds children objects:
  286. #Curve has no children (end leaf of self.shapes tree)
  287. self.linestyle = None
  288. self.linewidth = None
  289. self.linecolor = None
  290. self.fillcolor = None
  291. self.fillpattern = None
  292. self.arrow = None
  293. def inside_plot_area(self, verbose=True):
  294. """Check that all coordinates are within drawing_tool's area."""
  295. xmin, xmax = self.x.min(), self.x.max()
  296. ymin, ymax = self.y.min(), self.y.max()
  297. t = drawing_tool
  298. inside = True
  299. if xmin < t.xmin:
  300. inside = False
  301. if verbose:
  302. print 'x_min=%g < plot area x_min=%g' % (xmin, t.xmin)
  303. if xmax > t.xmax:
  304. inside = False
  305. if verbose:
  306. print 'x_max=%g > plot area x_max=%g' % (xmax, t.xmax)
  307. if ymin < t.ymin:
  308. inside = False
  309. if verbose:
  310. print 'y_min=%g < plot area y_min=%g' % (ymin, t.ymin)
  311. if xmax > t.xmax:
  312. inside = False
  313. if verbose:
  314. print 'y_max=%g > plot area y_max=%g' % (ymax, t.ymax)
  315. return inside
  316. def draw(self):
  317. """
  318. Send the curve to the plotting engine. That is, convert
  319. coordinate information in self.x and self.y, together
  320. with optional settings of linestyles, etc., to
  321. plotting commands for the chosen engine.
  322. """
  323. self.inside_plot_area()
  324. drawing_tool.plot_curve(
  325. self.x, self.y,
  326. self.linestyle, self.linewidth, self.linecolor,
  327. self.arrow, self.fillcolor, self.fillpattern)
  328. def rotate(self, angle, center):
  329. """
  330. Rotate all coordinates: `angle` is measured in degrees and
  331. (`x`,`y`) is the "origin" of the rotation.
  332. """
  333. angle = angle*pi/180
  334. x, y = center
  335. c = cos(angle); s = sin(angle)
  336. xnew = x + (self.x - x)*c - (self.y - y)*s
  337. ynew = y + (self.x - x)*s + (self.y - y)*c
  338. self.x = xnew
  339. self.y = ynew
  340. return self
  341. def scale(self, factor):
  342. """Scale all coordinates by `factor`: ``x = factor*x``, etc."""
  343. self.x = factor*self.x
  344. self.y = factor*self.y
  345. return self
  346. def translate(self, vec):
  347. """Translate all coordinates by a vector `vec`."""
  348. self.x += vec[0]
  349. self.y += vec[1]
  350. return self
  351. def set_linecolor(self, color):
  352. self.linecolor = color
  353. return self
  354. def set_linewidth(self, width):
  355. self.linewidth = width
  356. return self
  357. def set_linestyle(self, style):
  358. self.linestyle = style
  359. return self
  360. def set_arrow(self, style=None):
  361. self.arrow = style
  362. return self
  363. def set_name(self, name):
  364. self.name = name
  365. return self
  366. def set_filled_curves(self, color='', pattern=''):
  367. self.fillcolor = color
  368. self.fillpattern = pattern
  369. return self
  370. def show_hierarchy(self, indent=0, format='std'):
  371. if format == 'dict':
  372. return '"%s"' % str(self)
  373. elif format == 'plain':
  374. return ''
  375. else:
  376. return str(self)
  377. def __str__(self):
  378. """Compact pretty print of a Curve object."""
  379. s = '%d coords' % self.x.size
  380. if not self.inside_plot_area(verbose=False):
  381. s += ', some coordinates are outside plotting area!\n'
  382. props = ('linecolor', 'linewidth', 'linestyle', 'arrow',
  383. 'fillcolor', 'fillpattern')
  384. for prop in props:
  385. value = getattr(self, prop)
  386. if value is not None:
  387. s += ' %s=%s' % (prop, repr(value))
  388. return s
  389. def __repr__(self):
  390. return str(self)
  391. class Point(Shape):
  392. """A point (x,y) which can be rotated, translated, and scaled."""
  393. def __init__(self, x, y):
  394. self.x, self.y = x, y
  395. #self.shapes is not needed in this class
  396. def __add__(self, other):
  397. if isinstance(other, (list,tuple)):
  398. other = Point(other)
  399. return Point(self.x+other.x, self.y+other.y)
  400. # class Point is an abstract class - only subclasses are useful
  401. # and must implement draw
  402. def draw(self):
  403. raise NotImplementedError(
  404. 'class %s must implement the draw method' %
  405. self.__class__.__name__)
  406. def rotate(self, angle, center):
  407. """Rotate point an `angle` (in degrees) around (`x`,`y`)."""
  408. angle = angle*pi/180
  409. x, y = center
  410. c = cos(angle); s = sin(angle)
  411. xnew = x + (self.x - x)*c - (self.y - y)*s
  412. ynew = y + (self.x - x)*s + (self.y - y)*c
  413. self.x = xnew
  414. self.y = ynew
  415. return self
  416. def scale(self, factor):
  417. """Scale point coordinates by `factor`: ``x = factor*x``, etc."""
  418. self.x = factor*self.x
  419. self.y = factor*self.y
  420. return self
  421. def translate(self, vec):
  422. """Translate point by a vector `vec`."""
  423. self.x += vec[0]
  424. self.y += vec[1]
  425. return self
  426. def show_hierarchy(self, indent=0, format='std'):
  427. s = '%s at (%g,%g)' % (self.__class__.__name__, self.x, self.y)
  428. if format == 'dict':
  429. return '"%s"' % s
  430. elif format == 'plain':
  431. return ''
  432. else:
  433. return s
  434. # no need to store input data as they are invalid after rotations etc.
  435. class Rectangle(Shape):
  436. """
  437. Rectangle specified by the point `lower_left_corner`, `width`,
  438. and `height`.
  439. Recorded geometric features:
  440. ==================== =============================================
  441. Attribute Description
  442. ==================== =============================================
  443. lower_left Lower left corner point.
  444. upper_left Upper left corner point.
  445. lower_right Lower right corner point.
  446. upper_right Upper right corner point.
  447. lower_mid Middle point on lower side.
  448. upper_mid Middle point on upper side.
  449. ==================== =============================================
  450. """
  451. def __init__(self, lower_left_corner, width, height):
  452. is_sequence(lower_left_corner)
  453. p = arr2D(lower_left_corner) # short form
  454. x = [p[0], p[0] + width,
  455. p[0] + width, p[0], p[0]]
  456. y = [p[1], p[1], p[1] + height,
  457. p[1] + height, p[1]]
  458. self.shapes = {'rectangle': Curve(x,y)}
  459. # Geometric features
  460. self.lower_left = lower_left_corner
  461. self.lower_right = lower_left_corner + point(width,0)
  462. self.upper_left = lower_left_corner + point(0,height)
  463. self.upper_right = lower_left_corner + point(width,height)
  464. self.lower_mid = 0.5*(self.lower_left + self.lower_right)
  465. self.upper_mid = 0.5*(self.upper_left + self.upper_right)
  466. class Triangle(Shape):
  467. """
  468. Triangle defined by its three vertices p1, p2, and p3.
  469. Recorded geometric features:
  470. ==================== =============================================
  471. Attribute Description
  472. ==================== =============================================
  473. p1, p2, p3 Corners as given to the constructor.
  474. ==================== =============================================
  475. """
  476. def __init__(self, p1, p2, p3):
  477. is_sequence(p1, p2, p3)
  478. x = [p1[0], p2[0], p3[0], p1[0]]
  479. y = [p1[1], p2[1], p3[1], p1[1]]
  480. self.shapes = {'triangle': Curve(x,y)}
  481. # Geometric features
  482. self.p1 = arr2D(p1)
  483. self.p2 = arr2D(p2)
  484. self.p3 = arr2D(p3)
  485. class Line(Shape):
  486. def __init__(self, start, end):
  487. is_sequence(start, end)
  488. x = [start[0], end[0]]
  489. y = [start[1], end[1]]
  490. self.shapes = {'line': Curve(x, y)}
  491. def compute_formulas(self):
  492. x, y = self.shapes['line'].x, self.shapes['line'].y
  493. # Define equations for line:
  494. # y = a*x + b, x = c*y + d
  495. try:
  496. self.a = (y[1] - y[0])/(x[1] - x[0])
  497. self.b = y[0] - self.a*x[0]
  498. except ZeroDivisionError:
  499. # Vertical line, y is not a function of x
  500. self.a = None
  501. self.b = None
  502. try:
  503. if self.a is None:
  504. self.c = 0
  505. else:
  506. self.c = 1/float(self.a)
  507. if self.b is None:
  508. self.d = x[1]
  509. except ZeroDivisionError:
  510. # Horizontal line, x is not a function of y
  511. self.c = None
  512. self.d = None
  513. def compute_formulas(self):
  514. x, y = self.shapes['line'].x, self.shapes['line'].y
  515. tol = 1E-14
  516. # Define equations for line:
  517. # y = a*x + b, x = c*y + d
  518. if abs(x[1] - x[0]) > tol:
  519. self.a = (y[1] - y[0])/(x[1] - x[0])
  520. self.b = y[0] - self.a*x[0]
  521. else:
  522. # Vertical line, y is not a function of x
  523. self.a = None
  524. self.b = None
  525. if self.a is None:
  526. self.c = 0
  527. elif abs(self.a) > tol:
  528. self.c = 1/float(self.a)
  529. self.d = x[1]
  530. else: # self.a is 0
  531. # Horizontal line, x is not a function of y
  532. self.c = None
  533. self.d = None
  534. def __call__(self, x=None, y=None):
  535. """Given x, return y on the line, or given y, return x."""
  536. self.compute_formulas()
  537. if x is not None and self.a is not None:
  538. return self.a*x + self.b
  539. elif y is not None and self.c is not None:
  540. return self.c*y + self.d
  541. else:
  542. raise ValueError(
  543. 'Line.__call__(x=%s, y=%s) not meaningful' % \
  544. (x, y))
  545. # First implementation of class Circle
  546. class Circle(Shape):
  547. def __init__(self, center, radius, resolution=180):
  548. self.center, self.radius = center, radius
  549. self.resolution = resolution
  550. t = linspace(0, 2*pi, resolution+1)
  551. x0 = center[0]; y0 = center[1]
  552. R = radius
  553. x = x0 + R*cos(t)
  554. y = y0 + R*sin(t)
  555. self.shapes = {'circle': Curve(x, y)}
  556. def __call__(self, theta):
  557. """
  558. Return (x, y) point corresponding to angle theta.
  559. Not valid after a translation, rotation, or scaling.
  560. """
  561. return self.center[0] + self.radius*cos(theta), \
  562. self.center[1] + self.radius*sin(theta)
  563. class Arc(Shape):
  564. def __init__(self, center, radius,
  565. start_angle, arc_angle,
  566. resolution=180):
  567. is_sequence(center)
  568. self.center = center
  569. self.radius = radius
  570. self.start_angle = start_angle*pi/180 # radians
  571. self.arc_angle = arc_angle*pi/180
  572. self.resolution = resolution
  573. t = linspace(self.start_angle,
  574. self.start_angle + self.arc_angle,
  575. resolution+1)
  576. x0 = center[0]; y0 = center[1]
  577. R = radius
  578. x = x0 + R*cos(t)
  579. y = y0 + R*sin(t)
  580. self.shapes = {'arc': Curve(x, y)}
  581. def __call__(self, theta):
  582. """
  583. Return (x,y) point at start_angle + theta.
  584. Not valid after translation, rotation, or scaling.
  585. """
  586. theta = theta*pi/180
  587. t = self.start_angle + theta
  588. x0 = self.center[0]
  589. y0 = self.center[1]
  590. R = self.radius
  591. x = x0 + R*cos(t)
  592. y = y0 + R*sin(t)
  593. return (x, y)
  594. # Alternative for small arcs: Parabola
  595. class Parabola(Shape):
  596. def __init__(self, start, mid, stop, resolution=21):
  597. self.p1, self.p2, self.p3 = start, mid, stop
  598. # y as function of x? (no point on line x=const?)
  599. tol = 1E-14
  600. if abs(self.p1[0] - self.p2[0]) > 1E-14 and \
  601. abs(self.p2[0] - self.p3[0]) > 1E-14 and \
  602. abs(self.p3[0] - self.p1[0]) > 1E-14:
  603. self.y_of_x = True
  604. else:
  605. self.y_of_x = False
  606. # x as function of y? (no point on line y=const?)
  607. tol = 1E-14
  608. if abs(self.p1[1] - self.p2[1]) > 1E-14 and \
  609. abs(self.p2[1] - self.p3[1]) > 1E-14 and \
  610. abs(self.p3[1] - self.p1[1]) > 1E-14:
  611. self.x_of_y = True
  612. else:
  613. self.x_of_y = False
  614. if self.y_of_x:
  615. x = linspace(start[0], end[0], resolution)
  616. y = self(x=x)
  617. elif self.x_of_y:
  618. y = linspace(start[1], end[1], resolution)
  619. x = self(y=y)
  620. else:
  621. raise ValueError(
  622. 'Parabola: two or more points lie on x=const '
  623. 'or y=const - not allowed')
  624. self.shapes = {'parabola': Curve(x, y)}
  625. def __call__(self, x=None, y=None):
  626. if x is not None and self.y_of_x:
  627. return self._L2x(self.p1, self.p2)*self.p3[1] + \
  628. self._L2x(self.p2, self.p3)*self.p1[1] + \
  629. self._L2x(self.p3, self.p1)*self.p2[1]
  630. elif y is not None and self.x_of_y:
  631. return self._L2y(self.p1, self.p2)*self.p3[0] + \
  632. self._L2y(self.p2, self.p3)*self.p1[0] + \
  633. self._L2y(self.p3, self.p1)*self.p2[0]
  634. else:
  635. raise ValueError(
  636. 'Parabola.__call__(x=%s, y=%s) not meaningful' % \
  637. (x, y))
  638. def _L2x(self, x, pi, pj, pk):
  639. return (x - pi[0])*(x - pj[0])/((pk[0] - pi[0])*(pk[0] - pj[0]))
  640. def _L2y(self, y, pi, pj, pk):
  641. return (y - pi[1])*(y - pj[1])/((pk[1] - pi[1])*(pk[1] - pj[1]))
  642. class Circle(Arc):
  643. def __init__(self, center, radius, resolution=180):
  644. Arc.__init__(self, center, radius, 0, 360, resolution)
  645. class Wall(Shape):
  646. def __init__(self, x, y, thickness, pattern='/'):
  647. is_sequence(x, y, length=len(x))
  648. if isinstance(x[0], (tuple,list,ndarray)):
  649. # x is list of curves
  650. x1 = concatenate(x)
  651. else:
  652. x1 = asarray(x, float)
  653. if isinstance(y[0], (tuple,list,ndarray)):
  654. # x is list of curves
  655. y = concatenate(y)
  656. else:
  657. y1 = asarray(y, float)
  658. # Displaced curve (according to thickness)
  659. x2 = x1
  660. y2 = y1 + thickness
  661. # Combine x1,y1 with x2,y2 reversed
  662. from numpy import concatenate
  663. x = concatenate((x1, x2[-1::-1]))
  664. y = concatenate((y1, y2[-1::-1]))
  665. wall = Curve(x, y)
  666. wall.set_filled_curves(color='white', pattern=pattern)
  667. x = [x1[-1]] + x2[-1::-1].tolist() + [x1[0]]
  668. y = [y1[-1]] + y2[-1::-1].tolist() + [y1[0]]
  669. white_eraser = Curve(x, y)
  670. white_eraser.set_linecolor('white')
  671. from collections import OrderedDict
  672. self.shapes = OrderedDict()
  673. self.shapes['wall'] = wall
  674. self.shapes['eraser'] = white_eraser
  675. class VelocityProfile(Shape):
  676. def __init__(self, start, height, profile, num_arrows, scaling=1):
  677. # vx, vy = profile(y)
  678. shapes = {}
  679. # Draw left line
  680. shapes['start line'] = Line(start, (start[0], start[1]+height))
  681. # Draw velocity arrows
  682. dy = float(height)/(num_arrows-1)
  683. x = start[0]
  684. y = start[1]
  685. r = profile(y) # Test on return type
  686. if not isinstance(r, (list,tuple,ndarray)) and len(r) != 2:
  687. raise TypeError('VelocityProfile constructor: profile(y) function must return velocity vector (vx,vy), not %s' % type(r))
  688. for i in range(num_arrows):
  689. y = i*dy
  690. vx, vy = profile(y)
  691. if abs(vx) < 1E-8:
  692. continue
  693. vx *= scaling
  694. vy *= scaling
  695. arr = Arrow1((x,y), (x+vx, y+vy), '->')
  696. shapes['arrow%d' % i] = arr
  697. # Draw smooth profile
  698. xs = []
  699. ys = []
  700. n = 100
  701. dy = float(height)/n
  702. for i in range(n+2):
  703. y = i*dy
  704. vx, vy = profile(y)
  705. vx *= scaling
  706. vy *= scaling
  707. xs.append(x+vx)
  708. ys.append(y+vy)
  709. shapes['smooth curve'] = Curve(xs, ys)
  710. self.shapes = shapes
  711. class Arrow1(Shape):
  712. """Draw an arrow as Line with arrow."""
  713. def __init__(self, start, end, style='->'):
  714. arrow = Line(start, end)
  715. arrow.set_arrow(style)
  716. self.shapes = {'arrow': arrow}
  717. class Arrow3(Shape):
  718. """
  719. Build a vertical line and arrow head from Line objects.
  720. Then rotate `rotation_angle`.
  721. """
  722. def __init__(self, start, length, rotation_angle=0):
  723. self.bottom = start
  724. self.length = length
  725. self.angle = rotation_angle
  726. top = (self.bottom[0], self.bottom[1] + self.length)
  727. main = Line(self.bottom, top)
  728. #head_length = self.length/8.0
  729. head_length = drawing_tool.xrange/50.
  730. head_degrees = 30*pi/180
  731. head_left_pt = (top[0] - head_length*sin(head_degrees),
  732. top[1] - head_length*cos(head_degrees))
  733. head_right_pt = (top[0] + head_length*sin(head_degrees),
  734. top[1] - head_length*cos(head_degrees))
  735. head_left = Line(head_left_pt, top)
  736. head_right = Line(head_right_pt, top)
  737. head_left.set_linestyle('solid')
  738. head_right.set_linestyle('solid')
  739. self.shapes = {'line': main, 'head left': head_left,
  740. 'head right': head_right}
  741. # rotate goes through self.shapes so self.shapes
  742. # must be initialized first
  743. self.rotate(rotation_angle, start)
  744. class Text(Point):
  745. """
  746. Place `text` at the (x,y) point `position`, with the given
  747. fontsize. The text is centered around `position` if `alignment` is
  748. 'center'; if 'left', the text starts at `position`, and if
  749. 'right', the right and of the text is located at `position`.
  750. """
  751. def __init__(self, text, position, alignment='center', fontsize=14):
  752. is_sequence(position)
  753. is_sequence(position, length=2, can_be_None=True)
  754. self.text = text
  755. self.position = position
  756. self.alignment = alignment
  757. self.fontsize = fontsize
  758. Point.__init__(self, position[0], position[1])
  759. #no need for self.shapes here
  760. def draw(self):
  761. drawing_tool.text(self.text, (self.x, self.y),
  762. self.alignment, self.fontsize)
  763. def __str__(self):
  764. return 'text "%s" at (%g,%g)' % (self.text, self.x, self.y)
  765. def __repr__(self):
  766. return str(self)
  767. class Text_wArrow(Text):
  768. """
  769. As class Text, but an arrow is drawn from the mid part of the text
  770. to some point `arrow_tip`.
  771. """
  772. def __init__(self, text, position, arrow_tip,
  773. alignment='center', fontsize=14):
  774. is_sequence(arrow_tip, length=2, can_be_None=True)
  775. is_sequence(position)
  776. self.arrow_tip = arrow_tip
  777. Text.__init__(self, text, position, alignment, fontsize)
  778. def draw(self):
  779. drawing_tool.text(self.text, self.position,
  780. self.alignment, self.fontsize,
  781. self.arrow_tip)
  782. def __str__(self):
  783. return 'annotation "%s" at (%g,%g) with arrow to (%g,%g)' % \
  784. (self.text, self.x, self.y,
  785. self.arrow_tip[0], self.arrow_tip[1])
  786. def __repr__(self):
  787. return str(self)
  788. class Axis(Shape):
  789. def __init__(self, start, length, label, below=True,
  790. rotation_angle=0, fontsize=14,
  791. label_spacing=1./30):
  792. """
  793. Draw axis from start with `length` to the right
  794. (x axis). Place label below (True) or above (False) axis.
  795. Then return `rotation_angle` (in degrees).
  796. To make a standard x axis, call with ``below=True`` and
  797. ``rotation_angle=0``. To make a standard y axis, call with
  798. ``below=False`` and ``rotation_angle=90``.
  799. A tilted axis can also be drawn.
  800. The `label_spacing` denotes the space between the symbol
  801. and the arrow tip as a fraction of the length of the plot
  802. in x direction.
  803. """
  804. # Arrow is vertical arrow, make it horizontal
  805. arrow = Arrow3(start, length, rotation_angle=-90)
  806. arrow.rotate(rotation_angle, start)
  807. spacing = drawing_tool.xrange*label_spacing
  808. if below:
  809. spacing = - spacing
  810. label_pos = [start[0] + length, start[1] + spacing]
  811. symbol = Text(label, position=label_pos, fontsize=fontsize)
  812. symbol.rotate(rotation_angle, start)
  813. self.shapes = {'arrow': arrow, 'symbol': symbol}
  814. class Gravity(Axis):
  815. """Downward-pointing gravity arrow with the symbol g."""
  816. def __init__(self, start, length):
  817. Axis.__init__(self, start, length, '$g$', below=False,
  818. rotation_angle=-90, label_spacing=1./30)
  819. class Force(Arrow1):
  820. """
  821. Indication of a force by an arrow and a symbol.
  822. Draw an arrow, starting at `start` and with the tip at `end`.
  823. The symbol is placed at the `start` point, in a distance
  824. `symbol_spacing` times the width of the total plotting area.
  825. """
  826. def __init__(self, start, end, symbol, symbol_spacing=1./60,
  827. fontsize=14):
  828. Arrow1.__init__(self, start, end, style='->')
  829. spacing = drawing_tool.xrange*symbol_spacing
  830. start, end = arr2D(start), arr2D(end)
  831. spacing_dir = start - end
  832. spacing_dir /= sqrt(spacing_dir[0]**2 + spacing_dir[1]**2)
  833. symbol_pos = start + spacing*spacing_dir
  834. self.shapes['symbol'] = Text(symbol, symbol_pos, fontsize=fontsize)
  835. class Distance_wText(Shape):
  836. """
  837. Arrow <-> with text (usually a symbol) at the midpoint, used for
  838. identifying a some distance in a figure. The text is placed
  839. slightly to the right of vertical-like arrows, with text displaced
  840. `text_spacing` times to total distance in x direction of the plot
  841. area. The text is by default aligned 'left' in this case. For
  842. horizontal-like arrows, the text is placed the same distance
  843. above, but aligned 'center' by default (when `alignment` is None).
  844. """
  845. def __init__(self, start, end, text, fontsize=14, text_spacing=1/60.,
  846. alignment=None, text_pos='mid'):
  847. start = arr2D(start)
  848. end = arr2D(end)
  849. # Decide first if we have a vertical or horizontal arrow
  850. vertical = abs(end[0]-start[0]) < 2*abs(end[1]-start[1])
  851. if vertical:
  852. # Assume end above start
  853. if end[1] < start[1]:
  854. start, end = end, start
  855. if alignment is None:
  856. alignment = 'left'
  857. else: # horizontal arrow
  858. # Assume start to the right of end
  859. if start[0] < end[0]:
  860. start, end = end, start
  861. if alignment is None:
  862. alignment = 'center'
  863. tangent = end - start
  864. # Tangeng goes always to the left and upward
  865. normal = arr2D([tangent[1], -tangent[0]])/\
  866. sqrt(tangent[0]**2 + tangent[1]**2)
  867. mid = 0.5*(start + end) # midpoint of start-end line
  868. if text_pos == 'mid':
  869. text_pos = mid + normal*drawing_tool.xrange*text_spacing
  870. text = Text(text, text_pos, fontsize=fontsize,
  871. alignment=alignment)
  872. else:
  873. is_sequence(text_pos, length=2)
  874. text = Text_wArrow(text, text_pos, mid, alignment='left',
  875. fontsize=fontsize)
  876. arrow = Arrow1(start, end, style='<->')
  877. arrow.set_linecolor('black')
  878. arrow.set_linewidth(1)
  879. self.shapes = {'arrow': arrow, 'text': text}
  880. class ArcSymbol(Shape):
  881. def __init__(self, symbol, center, radius,
  882. start_angle, arc_angle, fontsize=14,
  883. resolution=180, symbol_spacing=1/60.):
  884. arc = Arc(center, radius, start_angle, arc_angle,
  885. resolution)
  886. mid = arr2D(arc(arc_angle/2.))
  887. normal = mid - arr2D(center)
  888. normal = normal/sqrt(normal[0]**2 + normal[1]**2)
  889. symbol_pos = mid + normal*drawing_tool.xrange*symbol_spacing
  890. self.shapes = {'arc': arc,
  891. 'symbol': Text(symbol, symbol_pos, fontsize=fontsize)}
  892. class Compose(Shape):
  893. def __init__(self, shapes):
  894. """shapes: list or dict of Shape objects."""
  895. self.shapes = shapes
  896. # can make help methods: Line.midpoint, Line.normal(pt, dir='left') -> (x,y)
  897. # list annotations in each class? contains extra annotations for explaining
  898. # important parameters to the constructor, e.g., Line.annotations holds
  899. # start and end as Text objects. Shape.demo calls shape.draw and
  900. # for annotation in self.demo: annotation.draw() YES!
  901. # Can make overall demo of classes by making objects and calling demo
  902. # Could include demo fig in each constructor
  903. class SimplySupportedBeam(Shape):
  904. def __init__(self, pos, size):
  905. pos = arr2D(pos)
  906. P0 = (pos[0] - size/2., pos[1]-size)
  907. P1 = (pos[0] + size/2., pos[1]-size)
  908. triangle = Triangle(P0, P1, pos)
  909. gap = size/5.
  910. h = size/4. # height of rectangle
  911. P2 = (P0[0], P0[1]-gap-h)
  912. rectangle = Rectangle(P2, size, h).set_filled_curves(pattern='/')
  913. self.shapes = {'triangle': triangle, 'rectangle': rectangle}
  914. self.dimensions = {'pos': Text('pos', pos),
  915. 'size': Distance_wText((P2[0], P2[1]-size),
  916. (P2[0]+size, P2[1]-size),
  917. 'size')}
  918. # Geometric features
  919. self.mid_support = point(P2[0] + size/2., P2[1]) # lower center
  920. self.top = pos
  921. class ConstantBeamLoad(Shape):
  922. """
  923. Downward-pointing arrows indicating a vertical load.
  924. The arrows are of equal length and filling a rectangle
  925. specified as in the :class:`Rectangle` class.
  926. Recorded geometric features:
  927. ==================== =============================================
  928. Attribute Description
  929. ==================== =============================================
  930. mid_point Middle point at the top of the row of
  931. arrows (often used for positioning a text).
  932. ==================== =============================================
  933. """
  934. def __init__(self, lower_left_corner, width, height, num_arrows=10):
  935. box = Rectangle(lower_left_corner, width, height)
  936. self.shapes = {'box': box}
  937. dx = float(width)/(num_arrows-1)
  938. y_top = lower_left_corner[1] + height
  939. y_tip = lower_left_corner[1]
  940. for i in range(num_arrows):
  941. x = lower_left_corner[0] + i*dx
  942. self.shapes['arrow%d' % i] = Arrow1((x, y_top), (x, y_tip))
  943. # Geometric features
  944. self.mid_top = arr2D(lower_left_corner) + point(width/2., height)
  945. class Moment(ArcSymbol):
  946. def __init__(self, symbol, center, radius,
  947. left=True, counter_clockwise=True,
  948. fontsize=14, symbol_spacing=1/60.):
  949. style = '->' if counter_clockwise else '<-'
  950. start_angle = 90 if left else -90
  951. ArcSymbol.__init__(self, symbol, center, radius,
  952. start_angle=start_angle,
  953. arc_angle=180, fontsize=fontsize,
  954. symbol_spacing=symbol_spacing,
  955. resolution=180)
  956. self.shapes['arc'].set_arrow(style)
  957. class Wheel(Shape):
  958. def __init__(self, center, radius, inner_radius=None, nlines=10):
  959. if inner_radius is None:
  960. inner_radius = radius/5.0
  961. outer = Circle(center, radius)
  962. inner = Circle(center, inner_radius)
  963. lines = []
  964. # Draw nlines+1 since the first and last coincide
  965. # (then nlines lines will be visible)
  966. t = linspace(0, 2*pi, self.nlines+1)
  967. Ri = inner_radius; Ro = radius
  968. x0 = center[0]; y0 = center[1]
  969. xinner = x0 + Ri*cos(t)
  970. yinner = y0 + Ri*sin(t)
  971. xouter = x0 + Ro*cos(t)
  972. youter = y0 + Ro*sin(t)
  973. lines = [Line((xi,yi),(xo,yo)) for xi, yi, xo, yo in \
  974. zip(xinner, yinner, xouter, youter)]
  975. self.shapes = {'inner': inner, 'outer': outer,
  976. 'spokes': Compose(
  977. {'spoke%d' % i: lines[i]
  978. for i in range(len(lines))})}
  979. class SineWave(Shape):
  980. def __init__(self, xstart, xstop,
  981. wavelength, amplitude, mean_level):
  982. self.xstart = xstart
  983. self.xstop = xstop
  984. self.wavelength = wavelength
  985. self.amplitude = amplitude
  986. self.mean_level = mean_level
  987. npoints = (self.xstop - self.xstart)/(self.wavelength/61.0)
  988. x = linspace(self.xstart, self.xstop, npoints)
  989. k = 2*pi/self.wavelength # frequency
  990. y = self.mean_level + self.amplitude*sin(k*x)
  991. self.shapes = {'waves': Curve(x,y)}
  992. class Spring1(Shape):
  993. """
  994. Specify a vertical spring, starting at `start`, with
  995. given vertical `length`. In the middle of the
  996. spring there are `num_teeth` saw teeth.
  997. Recorded geometric features:
  998. ==================== =============================================
  999. Attribute Description
  1000. ==================== =============================================
  1001. start Start point of spring.
  1002. end End point of spring.
  1003. ==================== =============================================
  1004. """
  1005. spring_fraction = 1./2 # fraction of total length occupied by spring
  1006. def __init__(self, start, length, tooth_width, num_teeth=8):
  1007. B = start
  1008. n = num_teeth - 1 # n counts teeth intervals
  1009. # n must be odd:
  1010. if n % 2 == 0:
  1011. n = n+1
  1012. L = length
  1013. w = tooth_width
  1014. # [0, x, L-x, L], f = (L-2*x)/L
  1015. # x = L*(1-f)/2.
  1016. shapes = {}
  1017. f = Spring1.spring_fraction
  1018. t = f*L/n # distance between teeth
  1019. s = L*(1-f)/2. # start of spring
  1020. P0 = (B[0], B[1]+s)
  1021. shapes['line start'] = Line(B, P0)
  1022. T1 = P0
  1023. T2 = (T1[0] + w, T1[1] + t/2.0)
  1024. k = 1
  1025. shapes['line%d' % k] = Line(T1,T2)
  1026. T1 = T2[:] # copy
  1027. for i in range(2*n-3):
  1028. T2 = (T1[0] + (-1)**(i+1)*2*w, T1[1] + t/2.0)
  1029. k += 1
  1030. shapes['line%d' % k] = Line(T1, T2)
  1031. T1 = (T2[0], T2[1])
  1032. T2 = (T1[0] + w, T1[1] + t/2.0)
  1033. k += 1
  1034. shapes['line%d' % k] = Line(T1,T2)
  1035. P2 = (B[0], B[1]+L)
  1036. shapes['line end'] = Line(T2, P2)
  1037. self.shapes = shapes
  1038. # Dimensions
  1039. start = Text_wArrow('start', (B[0]-1.5*w,B[1]-1.5*w), B)
  1040. width = Distance_wText((B[0]-w, B[1]-3.5*w), (B[0]+w, B[1]-3.5*w),
  1041. 'tooth_width')
  1042. length = Distance_wText((B[0]+3*w, B[1]), (B[0]+3*w, B[1]+L),
  1043. 'length')
  1044. num_teeth = Text_wArrow('num_teeth',
  1045. (B[0]+2*w,P2[1]+w),
  1046. (B[0]+1.2*w, B[1]+L/2.))
  1047. dims = {'start': start, 'width': width, 'length': length,
  1048. 'num_teeth': num_teeth}
  1049. self.dimensions = dims
  1050. # Geometric features
  1051. self.start = B
  1052. self.end = point(B[0], B[1]+L)
  1053. class Spring2(Shape):
  1054. """
  1055. Specify a vertical spring, starting at `start` and,
  1056. with vertical `length`. In the middle of the
  1057. spring there are `num_windings` circular windings to illustrate
  1058. the spring.
  1059. Recorded geometric features:
  1060. ==================== =============================================
  1061. Attribute Description
  1062. ==================== =============================================
  1063. start Start point of spring.
  1064. end End point of spring.
  1065. ==================== =============================================
  1066. """
  1067. spring_fraction = 1./2 # fraction of total length occupied by spring
  1068. def __init__(self, start, length, width, num_windings=11):
  1069. B = start
  1070. n = num_windings - 1 # n counts teeth intervals
  1071. if n <= 6:
  1072. n = 7
  1073. # n must be odd:
  1074. if n % 2 == 0:
  1075. n = n+1
  1076. L = length
  1077. w = width
  1078. # [0, x, L-x, L], f = (L-2*x)/L
  1079. # x = L*(1-f)/2.
  1080. shapes = {}
  1081. f = Spring2.spring_fraction
  1082. t = f*L/n # must be better worked out
  1083. s = L*(1-f)/2. # start of spring
  1084. P0 = (B[0], B[1]+s)
  1085. shapes['line start'] = Line(B, P0)
  1086. q = linspace(0, n, n*180 + 1)
  1087. x = P0[0] + w*sin(2*pi*q)
  1088. y = P0[1] + q*t
  1089. shapes['sprial'] = Curve(x, y)
  1090. P1 = (B[0], L-s)
  1091. P2 = (B[0], B[1]+L)
  1092. shapes['line end'] = Line(P1,P2)
  1093. self.shapes = shapes
  1094. # Dimensions
  1095. start = Text_wArrow('start', (B[0]-1.5*w,B[1]-1.5*w), B)
  1096. width = Distance_wText((B[0]-w, B[1]-3.5*w), (B[0]+w, B[1]-3.5*w),
  1097. 'width')
  1098. length = Distance_wText((B[0]+3*w, B[1]), (B[0]+3*w, B[1]+L),
  1099. 'length')
  1100. num_windings = Text_wArrow('num_windings',
  1101. (B[0]+2*w,P2[1]+w),
  1102. (B[0]+1.2*w, B[1]+L/2.))
  1103. spring_length = Distance_wText((B[0]-2*w, P0[1]), (B[0]-2*w, P1[1]),
  1104. 'Spring2.spring_fraction*length',
  1105. text_pos=(B[0]-6*w, P2[1]+2.5*w))
  1106. dims = {'start': start, 'width': width, 'length': length,
  1107. 'num_windings': num_windings, 'spring_length': spring_length}
  1108. self.dimensions = dims
  1109. # Geometric features
  1110. self.start = B
  1111. self.end = point(B[0], B[1]+L)
  1112. class Dashpot(Shape):
  1113. """
  1114. Specify a vertical dashpot of height `total_length` and
  1115. `start` as bottom/starting point. The rectangular dashpot part
  1116. has width `width` and height `dashpot_length`. If the latter
  1117. is not given (None), it becomes
  1118. ``Dashpot.dashpot_fraction*total_length`` (default
  1119. ``total_length/2```). The piston position inside the
  1120. rectangular dashpot, can be specified as `piston_pos`, (the
  1121. default value None places it at 1/3 from the bottom of the
  1122. dashpot).
  1123. Recorded geometric features:
  1124. ==================== =============================================
  1125. Attribute Description
  1126. ==================== =============================================
  1127. start Start point of dashpot.
  1128. end End point of dashpot.
  1129. ==================== =============================================
  1130. """
  1131. dashpot_fraction = 1./2
  1132. piston_gap_fraction = 1./6
  1133. piston_thickness_fraction = 1./8
  1134. def __init__(self, start, total_length,
  1135. width, dashpot_length=None, piston_pos=None):
  1136. B = start
  1137. L = total_length
  1138. w = width
  1139. # [0, x, L-x, L], f = (L-2*x)/L
  1140. # x = L*(1-f)/2.
  1141. shapes = {}
  1142. # dashpot is P0-P1 in y and width 2*w
  1143. if dashpot_length is None:
  1144. f = Dashpot.dashpot_fraction
  1145. s = L*(1-f)/2. # start of dashpot
  1146. P1 = (B[0], B[1]+L-s)
  1147. dashpot_length = f*L
  1148. else:
  1149. f = 1./2
  1150. s = f*dashpot_length # start of dashpot
  1151. P1 = (B[0], B[1]+s+dashpot_length)
  1152. P0 = (B[0], B[1]+s)
  1153. P2 = (B[0], B[1]+L)
  1154. shapes['line start'] = Line(B, P0)
  1155. shapes['pot'] = Curve([P1[0]-w, P0[0]-w, P0[0]+w, P1[0]+w],
  1156. [P1[1], P0[1], P0[1], P1[1]])
  1157. piston_thickness = f*L*Dashpot.piston_thickness_fraction
  1158. if piston_pos is None:
  1159. piston_pos = P0[1] + 1/3.*dashpot_length
  1160. if piston_pos < P0[1]:
  1161. piston_pos = P0[1]
  1162. if piston_pos > P1[1]-piston_thickness:
  1163. piston_pos = P1[1]-piston_thickness
  1164. gap = w*Dashpot.piston_gap_fraction
  1165. shapes['piston'] = Compose(
  1166. {'line': Line(P2, (B[0], piston_pos + piston_thickness)),
  1167. 'rectangle': Rectangle((B[0] - w+gap, piston_pos),
  1168. 2*w-2*gap, piston_thickness),
  1169. })
  1170. shapes['piston']['rectangle'].set_filled_curves(pattern='X')
  1171. self.shapes = shapes
  1172. # Dimensions
  1173. start = Text_wArrow('start', (B[0]-1.5*w,B[1]-1.5*w), B)
  1174. width = Distance_wText((B[0]-w, B[1]-3.5*w), (B[0]+w, B[1]-3.5*w),
  1175. 'width')
  1176. dplength = Distance_wText((B[0]+2*w, P0[1]), (B[0]+2*w, P1[1]),
  1177. 'dashpot_length', text_pos=(B[0]+w,B[1]-w))
  1178. tlength = Distance_wText((B[0]+4*w, B[1]), (B[0]+4*w, B[1]+L),
  1179. 'total_length',
  1180. text_pos=(B[0]+4.5*w, B[1]+L-2*w))
  1181. line = Line((B[0]+w, piston_pos), (B[0]+7*w, piston_pos)).set_linestyle('dotted').set_linecolor('black').set_linewidth(1)
  1182. pp = Text('piston_pos', (B[0]+7*w, piston_pos), alignment='left')
  1183. dims = {'start': start, 'width': width, 'dashpot_length': dplength,
  1184. 'total_length': tlength,
  1185. 'piston_pos': Compose({'line': line, 'text': pp})}
  1186. self.dimensions = dims
  1187. # Geometric features
  1188. self.start = B
  1189. self.end = point(B[0], B[1]+L)
  1190. # COMPOSITE types:
  1191. # MassSpringForce: Line(horizontal), Spring, Rectangle, Arrow/Line(w/arrow)
  1192. # must be easy to find the tip of the arrow
  1193. # Maybe extra dict: self.name['mass'] = Rectangle object - YES!
  1194. def test_Axis():
  1195. set_coordinate_system(xmin=0, xmax=15, ymin=0, ymax=15, axis=True)
  1196. x_axis = Axis((7.5,2), 5, 'x', rotation_angle=0)
  1197. y_axis = Axis((7.5,2), 5, 'y', below=False, rotation_angle=90)
  1198. system = Compose({'x axis': x_axis, 'y axis': y_axis})
  1199. system.draw()
  1200. drawing_tool.display()
  1201. set_linestyle('dashed')
  1202. #system.shapes['x axis'].rotate(40, (7.5, 2))
  1203. #system.shapes['y axis'].rotate(40, (7.5, 2))
  1204. system.rotate(40, (7.5,2))
  1205. system.draw()
  1206. drawing_tool.display('Axis')
  1207. drawing_tool.savefig('tmp_Axis.png')
  1208. print repr(system)
  1209. def test_Distance_wText():
  1210. drawing_tool.set_coordinate_system(xmin=0, xmax=10,
  1211. ymin=0, ymax=6,
  1212. axis=True,
  1213. instruction_file='tmp_mpl.py')
  1214. #drawing_tool.arrow_head_width = 0.1
  1215. fontsize=14
  1216. t = r'$ 2\pi R^2 $'
  1217. dims2 = Compose({
  1218. 'a0': Distance_wText((4,5), (8, 5), t, fontsize),
  1219. 'a6': Distance_wText((4,5), (4, 4), t, fontsize),
  1220. 'a1': Distance_wText((0,2), (2, 4.5), t, fontsize),
  1221. 'a2': Distance_wText((0,2), (2, 0), t, fontsize),
  1222. 'a3': Distance_wText((2,4.5), (0, 5.5), t, fontsize),
  1223. 'a4': Distance_wText((8,4), (10, 3), t, fontsize,
  1224. text_spacing=-1./60),
  1225. 'a5': Distance_wText((8,2), (10, 1), t, fontsize,
  1226. text_spacing=-1./40, alignment='right'),
  1227. 'c1': Text_wArrow('text_spacing=-1./60',
  1228. (4, 3.5), (9, 3.2),
  1229. fontsize=10, alignment='left'),
  1230. 'c2': Text_wArrow('text_spacing=-1./40, alignment="right"',
  1231. (4, 0.5), (9, 1.2),
  1232. fontsize=10, alignment='left'),
  1233. })
  1234. dims2.draw()
  1235. drawing_tool.display('Distance_wText and text positioning')
  1236. drawing_tool.savefig('tmp_Distance_wText.png')
  1237. def test_Springs():
  1238. L = 5
  1239. W = 2
  1240. drawing_tool.set_coordinate_system(xmin=0, xmax=7*W,
  1241. ymin=-2, ymax=L+2,
  1242. axis=True)
  1243. drawing_tool.set_linecolor('blue')
  1244. drawing_tool.set_grid(True)
  1245. xpos = W
  1246. s1 = Spring1((W,0), L, W/4.)
  1247. s1.draw()
  1248. s1.draw_dimensions()
  1249. xpos += 3*W
  1250. s2 = Spring2((xpos,0), L, W/4.)
  1251. s2.draw()
  1252. s2.draw_dimensions()
  1253. drawing_tool.display('Spring1 (left) and Spring2 (right)')
  1254. drawing_tool.savefig('tmp_springs.png')
  1255. def test_Dashpot():
  1256. L = 5
  1257. W = 2
  1258. drawing_tool.set_coordinate_system(xmin=xpos, xmax=xpos+5*W,
  1259. ymin=-2, ymax=L+2,
  1260. axis=True)
  1261. drawing_tool.set_linecolor('blue')
  1262. drawing_tool.set_grid(True)
  1263. # Default (simple) dashpot
  1264. xpos = 2
  1265. d1 = Dashpot(start=(xpos,0), total_length=L, width=W/4.)
  1266. text1 = Text('Dashpot (default)', (xpos, 1.1*L))
  1267. d1.draw()
  1268. text1.draw()
  1269. # Dashpot for animation with fixed dashpot_length and
  1270. # prescribed piston_pos
  1271. xpos += 1.5*W
  1272. d2 = Dashpot(start=(xpos,0), total_length=L+1.5, width=W/4.,
  1273. dashpot_length=2.5, piston_pos=L/2.)
  1274. d2.draw()
  1275. d2.draw_dimensions()
  1276. drawing_tool.display('Dashpot')
  1277. drawing_tool.savefig('tmp_dashpot.png')
  1278. def _test1():
  1279. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  1280. l1 = Line((0,0), (1,1))
  1281. l1.draw()
  1282. input(': ')
  1283. c1 = Circle((5,2), 1)
  1284. c2 = Circle((6,2), 1)
  1285. w1 = Wheel((7,2), 1)
  1286. c1.draw()
  1287. c2.draw()
  1288. w1.draw()
  1289. hardcopy()
  1290. display() # show the plot
  1291. def _test2():
  1292. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  1293. l1 = Line((0,0), (1,1))
  1294. l1.draw()
  1295. input(': ')
  1296. c1 = Circle((5,2), 1)
  1297. c2 = Circle((6,2), 1)
  1298. w1 = Wheel((7,2), 1)
  1299. filled_curves(True)
  1300. set_linecolor('blue')
  1301. c1.draw()
  1302. set_linecolor('aqua')
  1303. c2.draw()
  1304. filled_curves(False)
  1305. set_linecolor('red')
  1306. w1.draw()
  1307. hardcopy()
  1308. display() # show the plot
  1309. def _test3():
  1310. """Test example from the book."""
  1311. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  1312. l1 = Line(start=(0,0), stop=(1,1)) # define line
  1313. l1.draw() # make plot data
  1314. r1 = Rectangle(lower_left_corner=(0,1), width=3, height=5)
  1315. r1.draw()
  1316. Circle(center=(5,7), radius=1).draw()
  1317. Wheel(center=(6,2), radius=2, inner_radius=0.5, nlines=7).draw()
  1318. hardcopy()
  1319. display()
  1320. def _test4():
  1321. """Second example from the book."""
  1322. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  1323. r1 = Rectangle(lower_left_corner=(0,1), width=3, height=5)
  1324. c1 = Circle(center=(5,7), radius=1)
  1325. w1 = Wheel(center=(6,2), radius=2, inner_radius=0.5, nlines=7)
  1326. c2 = Circle(center=(7,7), radius=1)
  1327. filled_curves(True)
  1328. c1.draw()
  1329. set_linecolor('blue')
  1330. r1.draw()
  1331. set_linecolor('aqua')
  1332. c2.draw()
  1333. # Add thick aqua line around rectangle:
  1334. filled_curves(False)
  1335. set_linewidth(4)
  1336. r1.draw()
  1337. set_linecolor('red')
  1338. # Draw wheel with thick lines:
  1339. w1.draw()
  1340. hardcopy('tmp_colors')
  1341. display()
  1342. def _test5():
  1343. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  1344. c = 6. # center point of box
  1345. w = 2. # size of box
  1346. L = 3
  1347. r1 = Rectangle((c-w/2, c-w/2), w, w)
  1348. l1 = Line((c,c-w/2), (c,c-w/2-L))
  1349. linecolor('blue')
  1350. filled_curves(True)
  1351. r1.draw()
  1352. linecolor('aqua')
  1353. filled_curves(False)
  1354. l1.draw()
  1355. hardcopy()
  1356. display() # show the plot
  1357. def rolling_wheel(total_rotation_angle):
  1358. """Animation of a rotating wheel."""
  1359. set_coordinate_system(xmin=0, xmax=10, ymin=0, ymax=10)
  1360. import time
  1361. center = (6,2)
  1362. radius = 2.0
  1363. angle = 2.0
  1364. pngfiles = []
  1365. w1 = Wheel(center=center, radius=radius, inner_radius=0.5, nlines=7)
  1366. for i in range(int(total_rotation_angle/angle)):
  1367. w1.draw()
  1368. print 'XXXXXXXXXXXXXXXXXXXXXX BIG PROBLEM WITH ANIMATE!!!'
  1369. display()
  1370. filename = 'tmp_%03d' % i
  1371. pngfiles.append(filename + '.png')
  1372. hardcopy(filename)
  1373. time.sleep(0.3) # pause
  1374. L = radius*angle*pi/180 # translation = arc length
  1375. w1.rotate(angle, center)
  1376. w1.translate((-L, 0))
  1377. center = (center[0] - L, center[1])
  1378. erase()
  1379. cmd = 'convert -delay 50 -loop 1000 %s tmp_movie.gif' \
  1380. % (' '.join(pngfiles))
  1381. print 'converting PNG files to animated GIF:\n', cmd
  1382. import commands
  1383. failure, output = commands.getstatusoutput(cmd)
  1384. if failure: print 'Could not run', cmd
  1385. if __name__ == '__main__':
  1386. #rolling_wheel(40)
  1387. #_test1()
  1388. #_test3()
  1389. funcs = [
  1390. #test_Axis,
  1391. test_inclined_plane,
  1392. ]
  1393. for func in funcs:
  1394. func()
  1395. raw_input('Type Return: ')