diff --git a/manim/mobject/graphing/functions.py b/manim/mobject/graphing/functions.py index 83c48b1092..24d21daaf6 100644 --- a/manim/mobject/graphing/functions.py +++ b/manim/mobject/graphing/functions.py @@ -104,7 +104,7 @@ def construct(self): def __init__( self, - function: Callable[[float], Point3DLike], + function: Callable[[float | np.ndarray], Point3DLike], t_range: tuple[float, float] | tuple[float, float, float] = (0, 1), scaling: _ScaleBase = LinearBase(), dt: float = 1e-8, @@ -113,7 +113,7 @@ def __init__( use_vectorized: bool = False, **kwargs, ): - def internal_parametric_function(t: float) -> Point3D: + def internal_parametric_function(t: float | np.ndarray) -> Point3D: """Wrap ``function``'s output inside a NumPy array.""" return np.asarray(function(t)) @@ -143,23 +143,23 @@ def generate_points(self) -> Self: lambda t: self.t_min <= t <= self.t_max, self.discontinuities, ) - discontinuities = np.array(list(discontinuities)) + discontinuities_array = np.array(list(discontinuities)) boundary_times = np.array( [ self.t_min, self.t_max, - *(discontinuities - self.dt), - *(discontinuities + self.dt), + *(discontinuities_array + self.dt), + *(discontinuities_array - self.dt), ], ) boundary_times.sort() else: - boundary_times = [self.t_min, self.t_max] + boundary_times = np.array([self.t_min, self.t_max]) for t1, t2 in zip(boundary_times[0::2], boundary_times[1::2]): t_range = np.array( [ - *self.scaling.function(np.arange(t1, t2, self.t_step)), + *(self.scaling.function(t) for t in np.arange(t1, t2, self.t_step)), self.scaling.function(t2), ], ) diff --git a/manim/mobject/graphing/scale.py b/manim/mobject/graphing/scale.py index 78ffa2308b..972fceeb17 100644 --- a/manim/mobject/graphing/scale.py +++ b/manim/mobject/graphing/scale.py @@ -26,7 +26,7 @@ class _ScaleBase: def __init__(self, custom_labels: bool = False): self.custom_labels = custom_labels - def function(self, value: float) -> float: + def function(self, value: float | np.ndarray) -> float: """The function that will be used to scale the values. Parameters