geometry.py 33 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078
  1. from __future__ import annotations
  2. import math
  3. import numbers
  4. import numpy as np
  5. from manimlib.constants import DL, DOWN, DR, LEFT, ORIGIN, OUT, RIGHT, UL, UP, UR
  6. from manimlib.constants import GREY_A, RED, WHITE, BLACK
  7. from manimlib.constants import MED_SMALL_BUFF, SMALL_BUFF
  8. from manimlib.constants import DEGREES, PI, TAU
  9. from manimlib.mobject.mobject import Mobject
  10. from manimlib.mobject.types.vectorized_mobject import DashedVMobject
  11. from manimlib.mobject.types.vectorized_mobject import VGroup
  12. from manimlib.mobject.types.vectorized_mobject import VMobject
  13. from manimlib.utils.bezier import bezier
  14. from manimlib.utils.bezier import quadratic_bezier_points_for_arc
  15. from manimlib.utils.bezier import partial_quadratic_bezier_points
  16. from manimlib.utils.iterables import adjacent_n_tuples
  17. from manimlib.utils.iterables import adjacent_pairs
  18. from manimlib.utils.simple_functions import clip
  19. from manimlib.utils.simple_functions import fdiv
  20. from manimlib.utils.space_ops import angle_between_vectors
  21. from manimlib.utils.space_ops import angle_of_vector
  22. from manimlib.utils.space_ops import cross2d
  23. from manimlib.utils.space_ops import compass_directions
  24. from manimlib.utils.space_ops import find_intersection
  25. from manimlib.utils.space_ops import get_norm
  26. from manimlib.utils.space_ops import normalize
  27. from manimlib.utils.space_ops import rotate_vector
  28. from manimlib.utils.space_ops import rotation_matrix_transpose
  29. from manimlib.utils.space_ops import rotation_between_vectors
  30. from manimlib.utils.space_ops import rotation_about_z
  31. from typing import TYPE_CHECKING
  32. if TYPE_CHECKING:
  33. from typing import Iterable, Optional
  34. from manimlib.typing import ManimColor, Vect3, Vect3Array, Self
  35. DEFAULT_DOT_RADIUS = 0.08
  36. DEFAULT_SMALL_DOT_RADIUS = 0.04
  37. DEFAULT_DASH_LENGTH = 0.05
  38. DEFAULT_ARROW_TIP_LENGTH = 0.35
  39. DEFAULT_ARROW_TIP_WIDTH = 0.35
  40. # Deprecate?
  41. class TipableVMobject(VMobject):
  42. """
  43. Meant for shared functionality between Arc and Line.
  44. Functionality can be classified broadly into these groups:
  45. * Adding, Creating, Modifying tips
  46. - add_tip calls create_tip, before pushing the new tip
  47. into the TipableVMobject's list of submobjects
  48. - stylistic and positional configuration
  49. * Checking for tips
  50. - Boolean checks for whether the TipableVMobject has a tip
  51. and a starting tip
  52. * Getters
  53. - Straightforward accessors, returning information pertaining
  54. to the TipableVMobject instance's tip(s), its length etc
  55. """
  56. tip_config: dict = dict(
  57. fill_opacity=1.0,
  58. stroke_width=0.0,
  59. tip_style=0.0, # triangle=0, inner_smooth=1, dot=2
  60. )
  61. # Adding, Creating, Modifying tips
  62. def add_tip(self, at_start: bool = False, **kwargs) -> Self:
  63. """
  64. Adds a tip to the TipableVMobject instance, recognising
  65. that the endpoints might need to be switched if it's
  66. a 'starting tip' or not.
  67. """
  68. tip = self.create_tip(at_start, **kwargs)
  69. self.reset_endpoints_based_on_tip(tip, at_start)
  70. self.asign_tip_attr(tip, at_start)
  71. tip.set_color(self.get_stroke_color())
  72. self.add(tip)
  73. return self
  74. def create_tip(self, at_start: bool = False, **kwargs) -> ArrowTip:
  75. """
  76. Stylises the tip, positions it spacially, and returns
  77. the newly instantiated tip to the caller.
  78. """
  79. tip = self.get_unpositioned_tip(**kwargs)
  80. self.position_tip(tip, at_start)
  81. return tip
  82. def get_unpositioned_tip(self, **kwargs) -> ArrowTip:
  83. """
  84. Returns a tip that has been stylistically configured,
  85. but has not yet been given a position in space.
  86. """
  87. config = dict()
  88. config.update(self.tip_config)
  89. config.update(kwargs)
  90. return ArrowTip(**config)
  91. def position_tip(self, tip: ArrowTip, at_start: bool = False) -> ArrowTip:
  92. # Last two control points, defining both
  93. # the end, and the tangency direction
  94. if at_start:
  95. anchor = self.get_start()
  96. handle = self.get_first_handle()
  97. else:
  98. handle = self.get_last_handle()
  99. anchor = self.get_end()
  100. tip.rotate(angle_of_vector(handle - anchor) - PI - tip.get_angle())
  101. tip.shift(anchor - tip.get_tip_point())
  102. return tip
  103. def reset_endpoints_based_on_tip(self, tip: ArrowTip, at_start: bool) -> Self:
  104. if self.get_length() == 0:
  105. # Zero length, put_start_and_end_on wouldn't
  106. # work
  107. return self
  108. if at_start:
  109. start = tip.get_base()
  110. end = self.get_end()
  111. else:
  112. start = self.get_start()
  113. end = tip.get_base()
  114. self.put_start_and_end_on(start, end)
  115. return self
  116. def asign_tip_attr(self, tip: ArrowTip, at_start: bool) -> Self:
  117. if at_start:
  118. self.start_tip = tip
  119. else:
  120. self.tip = tip
  121. return self
  122. # Checking for tips
  123. def has_tip(self) -> bool:
  124. return hasattr(self, "tip") and self.tip in self
  125. def has_start_tip(self) -> bool:
  126. return hasattr(self, "start_tip") and self.start_tip in self
  127. # Getters
  128. def pop_tips(self) -> VGroup:
  129. start, end = self.get_start_and_end()
  130. result = VGroup()
  131. if self.has_tip():
  132. result.add(self.tip)
  133. self.remove(self.tip)
  134. if self.has_start_tip():
  135. result.add(self.start_tip)
  136. self.remove(self.start_tip)
  137. self.put_start_and_end_on(start, end)
  138. return result
  139. def get_tips(self) -> VGroup:
  140. """
  141. Returns a VGroup (collection of VMobjects) containing
  142. the TipableVMObject instance's tips.
  143. """
  144. result = VGroup()
  145. if hasattr(self, "tip"):
  146. result.add(self.tip)
  147. if hasattr(self, "start_tip"):
  148. result.add(self.start_tip)
  149. return result
  150. def get_tip(self) -> ArrowTip:
  151. """Returns the TipableVMobject instance's (first) tip,
  152. otherwise throws an exception."""
  153. tips = self.get_tips()
  154. if len(tips) == 0:
  155. raise Exception("tip not found")
  156. else:
  157. return tips[0]
  158. def get_default_tip_length(self) -> float:
  159. return self.tip_length
  160. def get_first_handle(self) -> Vect3:
  161. return self.get_points()[1]
  162. def get_last_handle(self) -> Vect3:
  163. return self.get_points()[-2]
  164. def get_end(self) -> Vect3:
  165. if self.has_tip():
  166. return self.tip.get_start()
  167. else:
  168. return VMobject.get_end(self)
  169. def get_start(self) -> Vect3:
  170. if self.has_start_tip():
  171. return self.start_tip.get_start()
  172. else:
  173. return VMobject.get_start(self)
  174. def get_length(self) -> float:
  175. start, end = self.get_start_and_end()
  176. return get_norm(start - end)
  177. class Arc(TipableVMobject):
  178. def __init__(
  179. self,
  180. start_angle: float = 0,
  181. angle: float = TAU / 4,
  182. radius: float = 1.0,
  183. n_components: int = 8,
  184. arc_center: Vect3 = ORIGIN,
  185. **kwargs
  186. ):
  187. super().__init__(**kwargs)
  188. self.set_points(quadratic_bezier_points_for_arc(angle, n_components))
  189. self.rotate(start_angle, about_point=ORIGIN)
  190. self.scale(radius, about_point=ORIGIN)
  191. self.shift(arc_center)
  192. def get_arc_center(self) -> Vect3:
  193. """
  194. Looks at the normals to the first two
  195. anchors, and finds their intersection points
  196. """
  197. # First two anchors and handles
  198. a1, h, a2 = self.get_points()[:3]
  199. # Tangent vectors
  200. t1 = h - a1
  201. t2 = h - a2
  202. # Normals
  203. n1 = rotate_vector(t1, TAU / 4)
  204. n2 = rotate_vector(t2, TAU / 4)
  205. return find_intersection(a1, n1, a2, n2)
  206. def get_start_angle(self) -> float:
  207. angle = angle_of_vector(self.get_start() - self.get_arc_center())
  208. return angle % TAU
  209. def get_stop_angle(self) -> float:
  210. angle = angle_of_vector(self.get_end() - self.get_arc_center())
  211. return angle % TAU
  212. def move_arc_center_to(self, point: Vect3) -> Self:
  213. self.shift(point - self.get_arc_center())
  214. return self
  215. class ArcBetweenPoints(Arc):
  216. def __init__(
  217. self,
  218. start: Vect3,
  219. end: Vect3,
  220. angle: float = TAU / 4,
  221. **kwargs
  222. ):
  223. super().__init__(angle=angle, **kwargs)
  224. if angle == 0:
  225. self.set_points_as_corners([LEFT, RIGHT])
  226. self.put_start_and_end_on(start, end)
  227. class CurvedArrow(ArcBetweenPoints):
  228. def __init__(
  229. self,
  230. start_point: Vect3,
  231. end_point: Vect3,
  232. **kwargs
  233. ):
  234. super().__init__(start_point, end_point, **kwargs)
  235. self.add_tip()
  236. class CurvedDoubleArrow(CurvedArrow):
  237. def __init__(
  238. self,
  239. start_point: Vect3,
  240. end_point: Vect3,
  241. **kwargs
  242. ):
  243. super().__init__(start_point, end_point, **kwargs)
  244. self.add_tip(at_start=True)
  245. class Circle(Arc):
  246. def __init__(
  247. self,
  248. start_angle: float = 0,
  249. stroke_color: ManimColor = RED,
  250. **kwargs
  251. ):
  252. super().__init__(
  253. start_angle, TAU,
  254. stroke_color=stroke_color,
  255. **kwargs
  256. )
  257. def surround(
  258. self,
  259. mobject: Mobject,
  260. dim_to_match: int = 0,
  261. stretch: bool = False,
  262. buff: float = MED_SMALL_BUFF
  263. ) -> Self:
  264. self.replace(mobject, dim_to_match, stretch)
  265. self.stretch((self.get_width() + 2 * buff) / self.get_width(), 0)
  266. self.stretch((self.get_height() + 2 * buff) / self.get_height(), 1)
  267. return self
  268. def point_at_angle(self, angle: float) -> Vect3:
  269. start_angle = self.get_start_angle()
  270. return self.point_from_proportion(
  271. ((angle - start_angle) % TAU) / TAU
  272. )
  273. def get_radius(self) -> float:
  274. return get_norm(self.get_start() - self.get_center())
  275. class Dot(Circle):
  276. def __init__(
  277. self,
  278. point: Vect3 = ORIGIN,
  279. radius: float = DEFAULT_DOT_RADIUS,
  280. stroke_color: ManimColor = BLACK,
  281. stroke_width: float = 0.0,
  282. fill_opacity: float = 1.0,
  283. fill_color: ManimColor = WHITE,
  284. **kwargs
  285. ):
  286. super().__init__(
  287. arc_center=point,
  288. radius=radius,
  289. stroke_color=stroke_color,
  290. stroke_width=stroke_width,
  291. fill_opacity=fill_opacity,
  292. fill_color=fill_color,
  293. **kwargs
  294. )
  295. class SmallDot(Dot):
  296. def __init__(
  297. self,
  298. point: Vect3 = ORIGIN,
  299. radius: float = DEFAULT_SMALL_DOT_RADIUS,
  300. **kwargs
  301. ):
  302. super().__init__(point, radius=radius, **kwargs)
  303. class Ellipse(Circle):
  304. def __init__(
  305. self,
  306. width: float = 2.0,
  307. height: float = 1.0,
  308. **kwargs
  309. ):
  310. super().__init__(**kwargs)
  311. self.set_width(width, stretch=True)
  312. self.set_height(height, stretch=True)
  313. class AnnularSector(VMobject):
  314. def __init__(
  315. self,
  316. angle: float = TAU / 4,
  317. start_angle: float = 0.0,
  318. inner_radius: float = 1.0,
  319. outer_radius: float = 2.0,
  320. arc_center: Vect3 = ORIGIN,
  321. fill_color: ManimColor = GREY_A,
  322. fill_opacity: float = 1.0,
  323. stroke_width: float = 0.0,
  324. **kwargs,
  325. ):
  326. super().__init__(
  327. fill_color=fill_color,
  328. fill_opacity=fill_opacity,
  329. stroke_width=stroke_width,
  330. **kwargs,
  331. )
  332. # Initialize points
  333. inner_arc, outer_arc = [
  334. Arc(
  335. start_angle=start_angle,
  336. angle=angle,
  337. radius=radius,
  338. arc_center=arc_center,
  339. )
  340. for radius in (inner_radius, outer_radius)
  341. ]
  342. self.set_points(inner_arc.get_points()[::-1]) # Reverse
  343. self.add_line_to(outer_arc.get_points()[0])
  344. self.add_subpath(outer_arc.get_points())
  345. self.add_line_to(inner_arc.get_points()[-1])
  346. class Sector(AnnularSector):
  347. def __init__(
  348. self,
  349. angle: float = TAU / 4,
  350. radius: float = 1.0,
  351. **kwargs
  352. ):
  353. super().__init__(
  354. angle,
  355. inner_radius=0,
  356. outer_radius=radius,
  357. **kwargs
  358. )
  359. class Annulus(VMobject):
  360. def __init__(
  361. self,
  362. inner_radius: float = 1.0,
  363. outer_radius: float = 2.0,
  364. fill_opacity: float = 1.0,
  365. stroke_width: float = 0.0,
  366. fill_color: ManimColor = GREY_A,
  367. center: Vect3 = ORIGIN,
  368. **kwargs,
  369. ):
  370. super().__init__(
  371. fill_color=fill_color,
  372. fill_opacity=fill_opacity,
  373. stroke_width=stroke_width,
  374. **kwargs,
  375. )
  376. self.radius = outer_radius
  377. outer_path = outer_radius * quadratic_bezier_points_for_arc(TAU)
  378. inner_path = inner_radius * quadratic_bezier_points_for_arc(-TAU)
  379. self.add_subpath(outer_path)
  380. self.add_subpath(inner_path)
  381. self.shift(center)
  382. class Line(TipableVMobject):
  383. def __init__(
  384. self,
  385. start: Vect3 | Mobject = LEFT,
  386. end: Vect3 | Mobject = RIGHT,
  387. buff: float = 0.0,
  388. path_arc: float = 0.0,
  389. **kwargs
  390. ):
  391. super().__init__(**kwargs)
  392. self.path_arc = path_arc
  393. self.buff = buff
  394. self.set_start_and_end_attrs(start, end)
  395. self.set_points_by_ends(self.start, self.end, buff, path_arc)
  396. def set_points_by_ends(
  397. self,
  398. start: Vect3,
  399. end: Vect3,
  400. buff: float = 0,
  401. path_arc: float = 0
  402. ) -> Self:
  403. self.clear_points()
  404. self.start_new_path(start)
  405. self.add_arc_to(end, path_arc)
  406. # Apply buffer
  407. if buff > 0:
  408. length = self.get_arc_length()
  409. alpha = min(buff / length, 0.5)
  410. self.pointwise_become_partial(self, alpha, 1 - alpha)
  411. return self
  412. def set_path_arc(self, new_value: float) -> Self:
  413. self.path_arc = new_value
  414. self.init_points()
  415. return self
  416. def set_start_and_end_attrs(self, start: Vect3 | Mobject, end: Vect3 | Mobject):
  417. # If either start or end are Mobjects, this
  418. # gives their centers
  419. rough_start = self.pointify(start)
  420. rough_end = self.pointify(end)
  421. vect = normalize(rough_end - rough_start)
  422. # Now that we know the direction between them,
  423. # we can find the appropriate boundary point from
  424. # start and end, if they're mobjects
  425. self.start = self.pointify(start, vect)
  426. self.end = self.pointify(end, -vect)
  427. def pointify(
  428. self,
  429. mob_or_point: Mobject | Vect3,
  430. direction: Vect3 | None = None
  431. ) -> Vect3:
  432. """
  433. Take an argument passed into Line (or subclass) and turn
  434. it into a 3d point.
  435. """
  436. if isinstance(mob_or_point, Mobject):
  437. mob = mob_or_point
  438. if direction is None:
  439. return mob.get_center()
  440. else:
  441. return mob.get_continuous_bounding_box_point(direction)
  442. else:
  443. point = mob_or_point
  444. result = np.zeros(self.dim)
  445. result[:len(point)] = point
  446. return result
  447. def put_start_and_end_on(self, start: Vect3, end: Vect3) -> Self:
  448. curr_start, curr_end = self.get_start_and_end()
  449. if np.isclose(curr_start, curr_end).all():
  450. # Handle null lines more gracefully
  451. self.set_points_by_ends(start, end, buff=0, path_arc=self.path_arc)
  452. return self
  453. return super().put_start_and_end_on(start, end)
  454. def get_vector(self) -> Vect3:
  455. return self.get_end() - self.get_start()
  456. def get_unit_vector(self) -> Vect3:
  457. return normalize(self.get_vector())
  458. def get_angle(self) -> float:
  459. return angle_of_vector(self.get_vector())
  460. def get_projection(self, point: Vect3) -> Vect3:
  461. """
  462. Return projection of a point onto the line
  463. """
  464. unit_vect = self.get_unit_vector()
  465. start = self.get_start()
  466. return start + np.dot(point - start, unit_vect) * unit_vect
  467. def get_slope(self) -> float:
  468. return np.tan(self.get_angle())
  469. def set_angle(self, angle: float, about_point: Optional[Vect3] = None) -> Self:
  470. if about_point is None:
  471. about_point = self.get_start()
  472. self.rotate(
  473. angle - self.get_angle(),
  474. about_point=about_point,
  475. )
  476. return self
  477. def set_length(self, length: float, **kwargs):
  478. self.scale(length / self.get_length(), **kwargs)
  479. return self
  480. def get_arc_length(self) -> float:
  481. arc_len = get_norm(self.get_vector())
  482. if self.path_arc > 0:
  483. arc_len *= self.path_arc / (2 * math.sin(self.path_arc / 2))
  484. return arc_len
  485. class DashedLine(Line):
  486. def __init__(
  487. self,
  488. start: Vect3 = LEFT,
  489. end: Vect3 = RIGHT,
  490. dash_length: float = DEFAULT_DASH_LENGTH,
  491. positive_space_ratio: float = 0.5,
  492. **kwargs
  493. ):
  494. super().__init__(start, end, **kwargs)
  495. num_dashes = self.calculate_num_dashes(dash_length, positive_space_ratio)
  496. dashes = DashedVMobject(
  497. self,
  498. num_dashes=num_dashes,
  499. positive_space_ratio=positive_space_ratio
  500. )
  501. self.clear_points()
  502. self.add(*dashes)
  503. def calculate_num_dashes(self, dash_length: float, positive_space_ratio: float) -> int:
  504. try:
  505. full_length = dash_length / positive_space_ratio
  506. return int(np.ceil(self.get_length() / full_length))
  507. except ZeroDivisionError:
  508. return 1
  509. def get_start(self) -> Vect3:
  510. if len(self.submobjects) > 0:
  511. return self.submobjects[0].get_start()
  512. else:
  513. return Line.get_start(self)
  514. def get_end(self) -> Vect3:
  515. if len(self.submobjects) > 0:
  516. return self.submobjects[-1].get_end()
  517. else:
  518. return Line.get_end(self)
  519. def get_first_handle(self) -> Vect3:
  520. return self.submobjects[0].get_points()[1]
  521. def get_last_handle(self) -> Vect3:
  522. return self.submobjects[-1].get_points()[-2]
  523. class TangentLine(Line):
  524. def __init__(
  525. self,
  526. vmob: VMobject,
  527. alpha: float,
  528. length: float = 2,
  529. d_alpha: float = 1e-6,
  530. **kwargs
  531. ):
  532. a1 = clip(alpha - d_alpha, 0, 1)
  533. a2 = clip(alpha + d_alpha, 0, 1)
  534. super().__init__(vmob.pfp(a1), vmob.pfp(a2), **kwargs)
  535. self.scale(length / self.get_length())
  536. class Elbow(VMobject):
  537. def __init__(
  538. self,
  539. width: float = 0.2,
  540. angle: float = 0,
  541. **kwargs
  542. ):
  543. super().__init__(**kwargs)
  544. self.set_points_as_corners([UP, UR, RIGHT])
  545. self.set_width(width, about_point=ORIGIN)
  546. self.rotate(angle, about_point=ORIGIN)
  547. class StrokeArrow(Line):
  548. def __init__(
  549. self,
  550. start: Vect3 | Mobject,
  551. end: Vect3 | Mobject,
  552. stroke_color: ManimColor = GREY_A,
  553. stroke_width: float = 5,
  554. buff: float = 0.25,
  555. tip_width_ratio: float = 5,
  556. tip_len_to_width: float = 0.0075,
  557. max_tip_length_to_length_ratio: float = 0.3,
  558. max_width_to_length_ratio: float = 8.0,
  559. **kwargs,
  560. ):
  561. self.tip_width_ratio = tip_width_ratio
  562. self.tip_len_to_width = tip_len_to_width
  563. self.max_tip_length_to_length_ratio = max_tip_length_to_length_ratio
  564. self.max_width_to_length_ratio = max_width_to_length_ratio
  565. self.n_tip_points = 3
  566. self.original_stroke_width = stroke_width
  567. super().__init__(
  568. start, end,
  569. stroke_color=stroke_color,
  570. stroke_width=stroke_width,
  571. buff=buff,
  572. **kwargs
  573. )
  574. def set_points_by_ends(
  575. self,
  576. start: Vect3,
  577. end: Vect3,
  578. buff: float = 0,
  579. path_arc: float = 0
  580. ) -> Self:
  581. super().set_points_by_ends(start, end, buff, path_arc)
  582. self.insert_tip_anchor()
  583. self.create_tip_with_stroke_width()
  584. return self
  585. def insert_tip_anchor(self) -> Self:
  586. prev_end = self.get_end()
  587. arc_len = self.get_arc_length()
  588. tip_len = self.get_stroke_width() * self.tip_width_ratio * self.tip_len_to_width
  589. if tip_len >= self.max_tip_length_to_length_ratio * arc_len or arc_len == 0:
  590. alpha = self.max_tip_length_to_length_ratio
  591. else:
  592. alpha = tip_len / arc_len
  593. if self.path_arc > 0 and self.buff > 0:
  594. self.insert_n_curves(10) # Is this needed?
  595. self.pointwise_become_partial(self, 0.0, 1.0 - alpha)
  596. self.add_line_to(self.get_end())
  597. self.add_line_to(prev_end)
  598. self.n_tip_points = 3
  599. return self
  600. @Mobject.affects_data
  601. def create_tip_with_stroke_width(self) -> Self:
  602. if self.get_num_points() < 3:
  603. return self
  604. stroke_width = min(
  605. self.original_stroke_width,
  606. self.max_width_to_length_ratio * self.get_length(),
  607. )
  608. tip_width = self.tip_width_ratio * stroke_width
  609. ntp = self.n_tip_points
  610. self.data['stroke_width'][:-ntp] = self.data['stroke_width'][0]
  611. self.data['stroke_width'][-ntp:, 0] = tip_width * np.linspace(1, 0, ntp)
  612. return self
  613. def reset_tip(self) -> Self:
  614. self.set_points_by_ends(
  615. self.get_start(), self.get_end(),
  616. path_arc=self.path_arc
  617. )
  618. return self
  619. def set_stroke(
  620. self,
  621. color: ManimColor | Iterable[ManimColor] | None = None,
  622. width: float | Iterable[float] | None = None,
  623. *args, **kwargs
  624. ) -> Self:
  625. super().set_stroke(color=color, width=width, *args, **kwargs)
  626. self.original_stroke_width = self.get_stroke_width()
  627. if self.has_points():
  628. self.reset_tip()
  629. return self
  630. def _handle_scale_side_effects(self, scale_factor: float) -> Self:
  631. if scale_factor != 1.0:
  632. self.reset_tip()
  633. return self
  634. class Arrow(Line):
  635. tickness_multiplier = 0.015
  636. def __init__(
  637. self,
  638. start: Vect3 | Mobject = LEFT,
  639. end: Vect3 | Mobject = LEFT,
  640. buff: float = MED_SMALL_BUFF,
  641. path_arc: float = 0,
  642. fill_color: ManimColor = GREY_A,
  643. fill_opacity: float = 1.0,
  644. stroke_width: float = 0.0,
  645. thickness: float = 3.0,
  646. tip_width_ratio: float = 5,
  647. tip_angle: float = PI / 3,
  648. max_tip_length_to_length_ratio: float = 0.5,
  649. max_width_to_length_ratio: float = 0.1,
  650. **kwargs,
  651. ):
  652. self.thickness = thickness
  653. self.tip_width_ratio = tip_width_ratio
  654. self.tip_angle = tip_angle
  655. self.max_tip_length_to_length_ratio = max_tip_length_to_length_ratio
  656. self.max_width_to_length_ratio = max_width_to_length_ratio
  657. super().__init__(
  658. start, end,
  659. fill_color=fill_color,
  660. fill_opacity=fill_opacity,
  661. stroke_width=stroke_width,
  662. buff=buff,
  663. path_arc=path_arc,
  664. **kwargs
  665. )
  666. def get_key_dimensions(self, length):
  667. width = self.thickness * self.tickness_multiplier
  668. w_ratio = fdiv(self.max_width_to_length_ratio, fdiv(width, length))
  669. if w_ratio < 1:
  670. width *= w_ratio
  671. tip_width = self.tip_width_ratio * width
  672. tip_length = tip_width / (2 * np.tan(self.tip_angle / 2))
  673. t_ratio = fdiv(self.max_tip_length_to_length_ratio, fdiv(tip_length, length))
  674. if t_ratio < 1:
  675. tip_length *= t_ratio
  676. tip_width *= t_ratio
  677. return width, tip_width, tip_length
  678. def set_points_by_ends(
  679. self,
  680. start: Vect3,
  681. end: Vect3,
  682. buff: float = 0,
  683. path_arc: float = 0
  684. ) -> Self:
  685. vect = end - start
  686. length = max(get_norm(vect), 1e-8) # More systematic min?
  687. unit_vect = normalize(vect)
  688. # Find the right tip length and thickness
  689. width, tip_width, tip_length = self.get_key_dimensions(length - buff)
  690. # Adjust start and end based on buff
  691. if path_arc == 0:
  692. start = start + buff * unit_vect
  693. end = end - buff * unit_vect
  694. else:
  695. R = length / 2 / math.sin(path_arc / 2)
  696. midpoint = 0.5 * (start + end)
  697. center = midpoint + rotate_vector(0.5 * vect, PI / 2) / math.tan(path_arc / 2)
  698. sign = 1
  699. start = center + rotate_vector(start - center, buff / R)
  700. end = center + rotate_vector(end - center, -buff / R)
  701. path_arc -= (2 * buff + tip_length) / R
  702. vect = end - start
  703. length = get_norm(vect)
  704. # Find points for the stem, imagining an arrow pointed to the left
  705. if path_arc == 0:
  706. points1 = (length - tip_length) * np.array([RIGHT, 0.5 * RIGHT, ORIGIN])
  707. points1 += width * UP / 2
  708. points2 = points1[::-1] + width * DOWN
  709. else:
  710. # Find arc points
  711. points1 = quadratic_bezier_points_for_arc(path_arc)
  712. points2 = np.array(points1[::-1])
  713. points1 *= (R + width / 2)
  714. points2 *= (R - width / 2)
  715. rot_T = rotation_matrix_transpose(PI / 2 - path_arc, OUT)
  716. for points in points1, points2:
  717. points[:] = np.dot(points, rot_T)
  718. points += R * DOWN
  719. self.set_points(points1)
  720. # Tip
  721. self.add_line_to(tip_width * UP / 2)
  722. self.add_line_to(tip_length * LEFT)
  723. self.tip_index = len(self.get_points()) - 1
  724. self.add_line_to(tip_width * DOWN / 2)
  725. self.add_line_to(points2[0])
  726. # Close it out
  727. self.add_subpath(points2)
  728. self.add_line_to(points1[0])
  729. # Reposition to match proper start and end
  730. self.rotate(angle_of_vector(vect) - self.get_angle())
  731. self.rotate(
  732. PI / 2 - np.arccos(normalize(vect)[2]),
  733. axis=rotate_vector(self.get_unit_vector(), -PI / 2),
  734. )
  735. self.shift(start - self.get_start())
  736. return self
  737. def reset_points_around_ends(self) -> Self:
  738. self.set_points_by_ends(
  739. self.get_start().copy(),
  740. self.get_end().copy(),
  741. path_arc=self.path_arc
  742. )
  743. return self
  744. def get_start(self) -> Vect3:
  745. points = self.get_points()
  746. return 0.5 * (points[0] + points[-3])
  747. def get_end(self) -> Vect3:
  748. return self.get_points()[self.tip_index]
  749. def get_start_and_end(self):
  750. return (self.get_start(), self.get_end())
  751. def put_start_and_end_on(self, start: Vect3, end: Vect3) -> Self:
  752. self.set_points_by_ends(start, end, buff=0, path_arc=self.path_arc)
  753. return self
  754. def scale(self, *args, **kwargs) -> Self:
  755. super().scale(*args, **kwargs)
  756. self.reset_points_around_ends()
  757. return self
  758. def set_thickness(self, thickness: float) -> Self:
  759. self.thickness = thickness
  760. self.reset_points_around_ends()
  761. return self
  762. def set_path_arc(self, path_arc: float) -> Self:
  763. self.path_arc = path_arc
  764. self.reset_points_around_ends()
  765. return self
  766. def set_perpendicular_to_camera(self, camera_frame):
  767. to_cam = camera_frame.get_implied_camera_location() - self.get_center()
  768. normal = self.get_unit_normal()
  769. axis = normalize(self.get_vector())
  770. # Project to be perpendicular to axis
  771. trg_normal = to_cam - np.dot(to_cam, axis) * axis
  772. mat = rotation_between_vectors(normal, trg_normal)
  773. self.apply_matrix(mat, about_point=self.get_start())
  774. return self
  775. class Vector(Arrow):
  776. def __init__(
  777. self,
  778. direction: Vect3 = RIGHT,
  779. buff: float = 0.0,
  780. **kwargs
  781. ):
  782. if len(direction) == 2:
  783. direction = np.hstack([direction, 0])
  784. super().__init__(ORIGIN, direction, buff=buff, **kwargs)
  785. class CubicBezier(VMobject):
  786. def __init__(
  787. self,
  788. a0: Vect3,
  789. h0: Vect3,
  790. h1: Vect3,
  791. a1: Vect3,
  792. **kwargs
  793. ):
  794. super().__init__(**kwargs)
  795. self.add_cubic_bezier_curve(a0, h0, h1, a1)
  796. class Polygon(VMobject):
  797. def __init__(
  798. self,
  799. *vertices: Vect3,
  800. **kwargs
  801. ):
  802. super().__init__(**kwargs)
  803. self.set_points_as_corners([*vertices, vertices[0]])
  804. def get_vertices(self) -> Vect3Array:
  805. return self.get_start_anchors()
  806. def round_corners(self, radius: Optional[float] = None) -> Self:
  807. if radius is None:
  808. verts = self.get_vertices()
  809. min_edge_length = min(
  810. get_norm(v1 - v2)
  811. for v1, v2 in zip(verts, verts[1:])
  812. if not np.isclose(v1, v2).all()
  813. )
  814. radius = 0.25 * min_edge_length
  815. vertices = self.get_vertices()
  816. arcs = []
  817. for v1, v2, v3 in adjacent_n_tuples(vertices, 3):
  818. vect1 = normalize(v2 - v1)
  819. vect2 = normalize(v3 - v2)
  820. angle = angle_between_vectors(vect1, vect2)
  821. # Distance between vertex and start of the arc
  822. cut_off_length = radius * np.tan(angle / 2)
  823. # Negative radius gives concave curves
  824. sign = float(np.sign(radius * cross2d(vect1, vect2)))
  825. arc = ArcBetweenPoints(
  826. v2 - vect1 * cut_off_length,
  827. v2 + vect2 * cut_off_length,
  828. angle=sign * angle,
  829. n_components=2,
  830. )
  831. arcs.append(arc)
  832. self.clear_points()
  833. # To ensure that we loop through starting with last
  834. arcs = [arcs[-1], *arcs[:-1]]
  835. for arc1, arc2 in adjacent_pairs(arcs):
  836. self.add_subpath(arc1.get_points())
  837. self.add_line_to(arc2.get_start())
  838. return self
  839. class Polyline(VMobject):
  840. def __init__(
  841. self,
  842. *vertices: Vect3,
  843. **kwargs
  844. ):
  845. super().__init__(**kwargs)
  846. self.set_points_as_corners(vertices)
  847. class RegularPolygon(Polygon):
  848. def __init__(
  849. self,
  850. n: int = 6,
  851. radius: float = 1.0,
  852. start_angle: float | None = None,
  853. **kwargs
  854. ):
  855. # Defaults to 0 for odd, 90 for even
  856. if start_angle is None:
  857. start_angle = (n % 2) * 90 * DEGREES
  858. start_vect = rotate_vector(radius * RIGHT, start_angle)
  859. vertices = compass_directions(n, start_vect)
  860. super().__init__(*vertices, **kwargs)
  861. class Triangle(RegularPolygon):
  862. def __init__(self, **kwargs):
  863. super().__init__(n=3, **kwargs)
  864. class ArrowTip(Triangle):
  865. def __init__(
  866. self,
  867. angle: float = 0,
  868. width: float = DEFAULT_ARROW_TIP_WIDTH,
  869. length: float = DEFAULT_ARROW_TIP_LENGTH,
  870. fill_opacity: float = 1.0,
  871. fill_color: ManimColor = WHITE,
  872. stroke_width: float = 0.0,
  873. tip_style: int = 0, # triangle=0, inner_smooth=1, dot=2
  874. **kwargs
  875. ):
  876. super().__init__(
  877. start_angle=0,
  878. fill_opacity=fill_opacity,
  879. fill_color=fill_color,
  880. stroke_width=stroke_width,
  881. **kwargs
  882. )
  883. self.set_height(width)
  884. self.set_width(length, stretch=True)
  885. if tip_style == 1:
  886. self.set_height(length * 0.9, stretch=True)
  887. self.data["point"][4] += np.array([0.6 * length, 0, 0])
  888. elif tip_style == 2:
  889. h = length / 2
  890. self.set_points(Dot().set_width(h).get_points())
  891. self.rotate(angle)
  892. def get_base(self) -> Vect3:
  893. return self.point_from_proportion(0.5)
  894. def get_tip_point(self) -> Vect3:
  895. return self.get_points()[0]
  896. def get_vector(self) -> Vect3:
  897. return self.get_tip_point() - self.get_base()
  898. def get_angle(self) -> float:
  899. return angle_of_vector(self.get_vector())
  900. def get_length(self) -> float:
  901. return get_norm(self.get_vector())
  902. class Rectangle(Polygon):
  903. def __init__(
  904. self,
  905. width: float = 4.0,
  906. height: float = 2.0,
  907. **kwargs
  908. ):
  909. super().__init__(UR, UL, DL, DR, **kwargs)
  910. self.set_width(width, stretch=True)
  911. self.set_height(height, stretch=True)
  912. def surround(self, mobject, buff=SMALL_BUFF) -> Self:
  913. target_shape = np.array(mobject.get_shape()) + 2 * buff
  914. self.set_shape(*target_shape)
  915. self.move_to(mobject)
  916. return self
  917. class Square(Rectangle):
  918. def __init__(self, side_length: float = 2.0, **kwargs):
  919. super().__init__(side_length, side_length, **kwargs)
  920. class RoundedRectangle(Rectangle):
  921. def __init__(
  922. self,
  923. width: float = 4.0,
  924. height: float = 2.0,
  925. corner_radius: float = 0.5,
  926. **kwargs
  927. ):
  928. super().__init__(width, height, **kwargs)
  929. self.round_corners(corner_radius)