shapes.py 69 KB

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