Use MappingFunction

ManimCommunity · Apr 30, 2024 · 9bc32e2 · 9bc32e2
1 parent 8331c71
commit 9bc32e2
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/manim/mobject/vector_field.py b/manim/mobject/vector_field.py
@@ -46,7 +46,7 @@
     import numpy.typing as npt
     from typing_extensions import Self
 
-    from manim.typing import ManimFloat, Point3D, Vector3D
+    from manim.typing import ManimFloat, Point3D, Vector3D, MappingFunction
 
 DEFAULT_SCALAR_FIELD_COLORS: list = [BLUE_E, GREEN, YELLOW, RED]
 
@@ -79,7 +79,7 @@ class VectorField(VGroup[VMobjectT]):
 
     def __init__(
         self,
-        func: Callable[[Point3D], npt.NDArray[ManimFloat]],
+        func: MappingFunction,
         color: ParsableManimColor | None = None,
         color_scheme: Callable[[Vector3D], float] | None = None,
         min_color_scheme_value: float = 0,
@@ -126,9 +126,9 @@ def pos_to_rgb(pos: np.ndarray) -> tuple[float, float, float, float]:
 
     @staticmethod
     def shift_func(
-        func: Callable[[npt.NDArray[ManimFloat]], npt.NDArray],
+        func: MappingFunction,
         shift_vector: Vector3D,
-    ) -> Callable[[np.ndarray], np.ndarray]:
+    ) -> MappingFunction:
         """Shift a vector field function.
 
         Parameters
@@ -148,9 +148,9 @@ def shift_func(
 
     @staticmethod
     def scale_func(
-        func: Callable[[npt.NDArray[ManimFloat]], npt.NDArray],
+        func: MappingFunction,
         scalar: float,
-    ) -> Callable[[npt.NDArray[ManimFloat]], npt.NDArray]:
+    ) -> MappingFunction:
         """Scale a vector field function.
 
         Parameters
@@ -546,7 +546,7 @@ def construct(self):
 
     def __init__(
         self,
-        func: Callable[[npt.NDArray[ManimFloat]], npt.NDArray],
+        func: MappingFunction,
         color: ParsableManimColor | None = None,
         color_scheme: Callable[[npt.NDArray], float] | None = None,
         min_color_scheme_value: float = 0,
@@ -720,7 +720,7 @@ def construct(self):
 
     def __init__(
         self,
-        func: Callable[[npt.NDArray[ManimFloat]], npt.NDArray],
+        func: MappingFunction,
         color: ParsableManimColor | None = None,
         color_scheme: Callable[[npt.NDArray], float] | None = None,
         min_color_scheme_value: float = 0,