Skip to content

Commit f304bd9

Browse files
chopan050chopan050
authored
Fix return type of Polygram.get_vertex_groups() and rename variables in .round_corners() (#4063)
1 parent 384350c commit f304bd9

File tree

2 files changed

+126
-34
lines changed

2 files changed

+126
-34
lines changed

‎manim/mobject/geometry/polygram.py

+36-33
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@
3838
from typing_extensions import Self
3939

4040
from manim.typing import (
41-
ManimFloat,
4241
Point3D,
4342
Point3D_Array,
4443
Point3DLike,
@@ -122,39 +121,45 @@ def get_vertices(self) -> Point3D_Array:
122121
"""
123122
return self.get_start_anchors()
124123

125-
def get_vertex_groups(self) -> npt.NDArray[ManimFloat]:
124+
def get_vertex_groups(self) -> list[Point3D_Array]:
126125
"""Gets the vertex groups of the :class:`Polygram`.
127126
128127
Returns
129128
-------
130-
:class:`numpy.ndarray`
131-
The vertex groups of the :class:`Polygram`.
129+
list[Point3D_Array]
130+
The list of vertex groups of the :class:`Polygram`.
132131
133132
Examples
134133
--------
135134
::
136135
137-
>>> poly = Polygram([ORIGIN, RIGHT, UP], [LEFT, LEFT + UP, 2 * LEFT])
138-
>>> poly.get_vertex_groups()
139-
array([[[ 0., 0., 0.],
140-
[ 1., 0., 0.],
141-
[ 0., 1., 0.]],
142-
<BLANKLINE>
143-
[[-1., 0., 0.],
144-
[-1., 1., 0.],
145-
[-2., 0., 0.]]])
136+
>>> poly = Polygram([ORIGIN, RIGHT, UP, LEFT + UP], [LEFT, LEFT + UP, 2 * LEFT])
137+
>>> groups = poly.get_vertex_groups()
138+
>>> len(groups)
139+
2
140+
>>> groups[0]
141+
array([[ 0., 0., 0.],
142+
[ 1., 0., 0.],
143+
[ 0., 1., 0.],
144+
[-1., 1., 0.]])
145+
>>> groups[1]
146+
array([[-1., 0., 0.],
147+
[-1., 1., 0.],
148+
[-2., 0., 0.]])
146149
"""
147150
vertex_groups = []
148151

152+
# TODO: If any of the original vertex groups contained the starting vertex N
153+
# times, then .get_vertex_groups() splits it into N vertex groups.
149154
group = []
150155
for start, end in zip(self.get_start_anchors(), self.get_end_anchors()):
151156
group.append(start)
152157

153158
if self.consider_points_equals(end, group[0]):
154-
vertex_groups.append(group)
159+
vertex_groups.append(np.array(group))
155160
group = []
156161

157-
return np.array(vertex_groups)
162+
return vertex_groups
158163

159164
def round_corners(
160165
self,
@@ -223,18 +228,18 @@ def construct(self):
223228

224229
new_points: list[Point3D] = []
225230

226-
for vertices in self.get_vertex_groups():
231+
for vertex_group in self.get_vertex_groups():
227232
arcs = []
228233

229234
# Repeat the radius list as necessary in order to provide a radius
230235
# for each vertex.
231236
if isinstance(radius, (int, float)):
232-
radius_list = [radius] * len(vertices)
237+
radius_list = [radius] * len(vertex_group)
233238
else:
234-
radius_list = radius * ceil(len(vertices) / len(radius))
239+
radius_list = radius * ceil(len(vertex_group) / len(radius))
235240

236-
for currentRadius, (v1, v2, v3) in zip(
237-
radius_list, adjacent_n_tuples(vertices, 3)
241+
for current_radius, (v1, v2, v3) in zip(
242+
radius_list, adjacent_n_tuples(vertex_group, 3)
238243
):
239244
vect1 = v2 - v1
240245
vect2 = v3 - v2
@@ -243,10 +248,10 @@ def construct(self):
243248

244249
angle = angle_between_vectors(vect1, vect2)
245250
# Negative radius gives concave curves
246-
angle *= np.sign(currentRadius)
251+
angle *= np.sign(current_radius)
247252

248253
# Distance between vertex and start of the arc
249-
cut_off_length = currentRadius * np.tan(angle / 2)
254+
cut_off_length = current_radius * np.tan(angle / 2)
250255

251256
# Determines counterclockwise vs. clockwise
252257
sign = np.sign(np.cross(vect1, vect2)[2])
@@ -261,17 +266,17 @@ def construct(self):
261266

262267
if evenly_distribute_anchors:
263268
# Determine the average length of each curve
264-
nonZeroLengthArcs = [arc for arc in arcs if len(arc.points) > 4]
265-
if len(nonZeroLengthArcs):
266-
totalArcLength = sum(
267-
[arc.get_arc_length() for arc in nonZeroLengthArcs]
269+
nonzero_length_arcs = [arc for arc in arcs if len(arc.points) > 4]
270+
if len(nonzero_length_arcs) > 0:
271+
total_arc_length = sum(
272+
[arc.get_arc_length() for arc in nonzero_length_arcs]
268273
)
269-
totalCurveCount = (
270-
sum([len(arc.points) for arc in nonZeroLengthArcs]) / 4
274+
num_curves = (
275+
sum([len(arc.points) for arc in nonzero_length_arcs]) / 4
271276
)
272-
averageLengthPerCurve = totalArcLength / totalCurveCount
277+
average_arc_length = total_arc_length / num_curves
273278
else:
274-
averageLengthPerCurve = 1
279+
average_arc_length = 1.0
275280

276281
# To ensure that we loop through starting with last
277282
arcs = [arcs[-1], *arcs[:-1]]
@@ -284,9 +289,7 @@ def construct(self):
284289

285290
# Make sure anchors are evenly distributed, if necessary
286291
if evenly_distribute_anchors:
287-
line.insert_n_curves(
288-
ceil(line.get_length() / averageLengthPerCurve)
289-
)
292+
line.insert_n_curves(ceil(line.get_length() / average_arc_length))
290293

291294
new_points.extend(line.points)
292295

‎tests/module/mobject/geometry/test_unit_geometry.py

+90-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,14 @@
44

55
import numpy as np
66

7-
from manim import BackgroundRectangle, Circle, Sector, Square, SurroundingRectangle
7+
from manim import (
8+
BackgroundRectangle,
9+
Circle,
10+
Polygram,
11+
Sector,
12+
Square,
13+
SurroundingRectangle,
14+
)
815

916
logger = logging.getLogger(__name__)
1017

@@ -15,6 +22,88 @@ def test_get_arc_center():
1522
)
1623

1724

25+
def test_Polygram_get_vertex_groups():
26+
# Test that, once a Polygram polygram is created with some vertex groups,
27+
# polygram.get_vertex_groups() (usually) returns the same vertex groups.
28+
vertex_groups_arr = [
29+
# 2 vertex groups for polygram 1
30+
[
31+
# Group 1: Triangle
32+
np.array(
33+
[
34+
[2, 1, 0],
35+
[0, 2, 0],
36+
[-2, 1, 0],
37+
]
38+
),
39+
# Group 2: Square
40+
np.array(
41+
[
42+
[1, 0, 0],
43+
[0, 1, 0],
44+
[-1, 0, 0],
45+
[0, -1, 0],
46+
]
47+
),
48+
],
49+
# 3 vertex groups for polygram 1
50+
[
51+
# Group 1: Quadrilateral
52+
np.array(
53+
[
54+
[2, 0, 0],
55+
[0, -1, 0],
56+
[0, 0, -2],
57+
[0, 1, 0],
58+
]
59+
),
60+
# Group 2: Triangle
61+
np.array(
62+
[
63+
[3, 1, 0],
64+
[0, 0, 2],
65+
[2, 0, 0],
66+
]
67+
),
68+
# Group 3: Pentagon
69+
np.array(
70+
[
71+
[1, -1, 0],
72+
[1, 1, 0],
73+
[0, 2, 0],
74+
[-1, 1, 0],
75+
[-1, -1, 0],
76+
]
77+
),
78+
],
79+
]
80+
81+
for vertex_groups in vertex_groups_arr:
82+
polygram = Polygram(*vertex_groups)
83+
poly_vertex_groups = polygram.get_vertex_groups()
84+
for poly_group, group in zip(poly_vertex_groups, vertex_groups):
85+
np.testing.assert_array_equal(poly_group, group)
86+
87+
# If polygram is a Polygram of a vertex group containing the start vertex N times,
88+
# then polygram.get_vertex_groups() splits it into N vertex groups.
89+
splittable_vertex_group = np.array(
90+
[
91+
[0, 1, 0],
92+
[1, -2, 0],
93+
[1, 2, 0],
94+
[0, 1, 0], # same vertex as start
95+
[-1, 2, 0],
96+
[-1, -2, 0],
97+
[0, 1, 0], # same vertex as start
98+
[0.5, 2, 0],
99+
[-0.5, 2, 0],
100+
]
101+
)
102+
103+
polygram = Polygram(splittable_vertex_group)
104+
assert len(polygram.get_vertex_groups()) == 3
105+
106+
18107
def test_SurroundingRectangle():
19108
circle = Circle()
20109
square = Square()

0 commit comments

Comments
 (0)