shapes.py 81 KB

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