diff --git a/manim/mobject/three_d/three_dimensions.py b/manim/mobject/three_d/three_dimensions.py index ccbc4dea91..53d8e98c83 100644 --- a/manim/mobject/three_d/three_dimensions.py +++ b/manim/mobject/three_d/three_dimensions.py @@ -17,7 +17,7 @@ ] from collections.abc import Callable, Iterable, Sequence -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import numpy as np from typing_extensions import Self @@ -40,15 +40,22 @@ ParsableManimColor, interpolate_color, ) -from manim.utils.iterables import tuplify from manim.utils.space_ops import normalize, perpendicular_bisector, z_to_vector if TYPE_CHECKING: - from manim.typing import Point3D, Point3DLike, Vector3DLike + from manim.mobject.graphing.coordinate_systems import ThreeDAxes + from manim.typing import Point3D, Point3DLike, Vector3D, Vector3DLike class ThreeDVMobject(VMobject, metaclass=ConvertToOpenGL): - def __init__(self, shade_in_3d: bool = True, **kwargs): + u_index: int + v_index: int + u1: float + u2: float + v1: float + v2: float + + def __init__(self, shade_in_3d: bool = True, **kwargs: Any): super().__init__(shade_in_3d=shade_in_3d, **kwargs) @@ -107,13 +114,16 @@ def construct(self): def __init__( self, func: Callable[[float, float], np.ndarray], - u_range: Sequence[float] = [0, 1], - v_range: Sequence[float] = [0, 1], - resolution: Sequence[int] = 32, + u_range: tuple[float, float] = (0, 1), + v_range: tuple[float, float] = (0, 1), + resolution: int | Sequence[int] = 32, surface_piece_config: dict = {}, fill_color: ParsableManimColor = BLUE_D, fill_opacity: float = 1.0, - checkerboard_colors: Sequence[ParsableManimColor] | bool = [BLUE_D, BLUE_E], + checkerboard_colors: Iterable[ParsableManimColor] | Literal[False] = [ + BLUE_D, + BLUE_E, + ], stroke_color: ParsableManimColor = LIGHT_GREY, stroke_width: float = 0.5, should_make_jagged: bool = False, @@ -131,12 +141,11 @@ def __init__( ) self.resolution = resolution self.surface_piece_config = surface_piece_config - if checkerboard_colors: - self.checkerboard_colors: list[ManimColor] = [ - ManimColor(x) for x in checkerboard_colors - ] - else: + self.checkerboard_colors: list[ManimColor] | Literal[False] + if checkerboard_colors is False: self.checkerboard_colors = checkerboard_colors + else: + self.checkerboard_colors = [ManimColor(i) for i in checkerboard_colors] self.should_make_jagged = should_make_jagged self.pre_function_handle_to_anchor_scale_factor = ( pre_function_handle_to_anchor_scale_factor @@ -151,11 +160,10 @@ def func(self, u: float, v: float) -> np.ndarray: return self._func(u, v) def _get_u_values_and_v_values(self) -> tuple[np.ndarray, np.ndarray]: - res = tuplify(self.resolution) - if len(res) == 1: - u_res = v_res = res[0] + if isinstance(self.resolution, int): + u_res = v_res = self.resolution else: - u_res, v_res = res + u_res, v_res = self.resolution u_values = np.linspace(*self.u_range, u_res + 1) v_values = np.linspace(*self.v_range, v_res + 1) @@ -197,7 +205,7 @@ def _setup_in_uv_space(self) -> None: self.set_fill_by_checkerboard(*self.checkerboard_colors) def set_fill_by_checkerboard( - self, *colors: Iterable[ParsableManimColor], opacity: float | None = None + self, *colors: ParsableManimColor, opacity: float | None = None ) -> Self: """Sets the fill_color of each face of :class:`Surface` in an alternating pattern. @@ -223,10 +231,12 @@ def set_fill_by_checkerboard( def set_fill_by_value( self, - axes: Mobject, - colorscale: list[ParsableManimColor] | ParsableManimColor | None = None, + axes: ThreeDAxes, + colorscale: Iterable[ParsableManimColor] + | Iterable[tuple[ParsableManimColor, float]] + | None = None, axis: int = 2, - **kwargs, + **kwargs: Any, ) -> Self: """Sets the color of each mobject of a parametric surface to a color relative to its axis-value. @@ -286,19 +296,23 @@ def param_surface(u, v): "the surface fill color has not been changed" ) return self + colorscale_list = list(colorscale) ranges = [axes.x_range, axes.y_range, axes.z_range] - - if type(colorscale[0]) is tuple: + assert isinstance(colorscale_list, list) + new_colors: list[ManimColor] + if type(colorscale_list[0]) is tuple and len(colorscale_list[0]) == 2: new_colors, pivots = [ - [i for i, j in colorscale], - [j for i, j in colorscale], + [ManimColor(i) for i, j in colorscale_list], + [j for i, j in colorscale_list], ] else: - new_colors = colorscale + new_colors = [ManimColor(i) for i in colorscale_list] + current_range = ranges[axis] - pivot_min = ranges[axis][0] - pivot_max = ranges[axis][1] + assert current_range is not None + pivot_min = current_range[0] + pivot_max = current_range[1] pivot_frequency = (pivot_max - pivot_min) / (len(new_colors) - 1) pivots = np.arange( start=pivot_min, @@ -325,6 +339,7 @@ def param_surface(u, v): color_index, ) if config.renderer == RendererType.OPENGL: + assert isinstance(mob, OpenGLMobject) mob.set_color(mob_color, recurse=False) elif config.renderer == RendererType.CAIRO: mob.set_color(mob_color, family=False) @@ -383,10 +398,10 @@ def __init__( self, center: Point3DLike = ORIGIN, radius: float = 1, - resolution: Sequence[int] | None = None, - u_range: Sequence[float] = (0, TAU), - v_range: Sequence[float] = (0, PI), - **kwargs, + resolution: int | Sequence[int] | None = None, + u_range: tuple[float, float] = (0, TAU), + v_range: tuple[float, float] = (0, PI), + **kwargs: Any, ) -> None: if config.renderer == RendererType.OPENGL: res_value = (101, 51) @@ -409,12 +424,12 @@ def __init__( self.shift(center) - def func(self, u: float, v: float) -> np.ndarray: + def func(self, u: float, v: float) -> Point3D: """The z values defining the :class:`Sphere` being plotted. Returns ------- - :class:`numpy.array` + :class:`Point3D` The z values defining the :class:`Sphere`. """ return self.radius * np.array( @@ -456,11 +471,11 @@ def construct(self): def __init__( self, - point: list | np.ndarray = ORIGIN, + point: Point3D = ORIGIN, radius: float = DEFAULT_DOT_RADIUS, color: ParsableManimColor = WHITE, - resolution: tuple[int, int] = (8, 8), - **kwargs, + resolution: int | tuple[int, int] | None = (8, 8), + **kwargs: Any, ) -> None: super().__init__(center=point, radius=radius, resolution=resolution, **kwargs) self.set_color(color) @@ -502,7 +517,7 @@ def __init__( fill_opacity: float = 0.75, fill_color: ParsableManimColor = BLUE, stroke_width: float = 0, - **kwargs, + **kwargs: Any, ) -> None: self.side_length = side_length super().__init__( @@ -554,7 +569,9 @@ def construct(self): """ def __init__( - self, dimensions: tuple[float, float, float] | np.ndarray = [3, 2, 1], **kwargs + self, + dimensions: Vector3DLike = [3, 2, 1], + **kwargs: Any, ) -> None: self.dimensions = dimensions super().__init__(**kwargs) @@ -608,20 +625,20 @@ def __init__( self, base_radius: float = 1, height: float = 1, - direction: np.ndarray = Z_AXIS, + direction: Vector3DLike = Z_AXIS, show_base: bool = False, - v_range: Sequence[float] = [0, TAU], + v_range: tuple[float, float] = (0, TAU), u_min: float = 0, - checkerboard_colors: bool = False, + checkerboard_colors: Iterable[ParsableManimColor] | Literal[False] = False, **kwargs: Any, ) -> None: - self.direction = direction + self.direction = np.array(direction) self.theta = PI - np.arctan(base_radius / height) super().__init__( self.func, v_range=v_range, - u_range=[u_min, np.sqrt(base_radius**2 + height**2)], + u_range=(u_min, np.sqrt(base_radius**2 + height**2)), checkerboard_colors=checkerboard_colors, **kwargs, ) @@ -642,7 +659,7 @@ def __init__( self._rotate_to_direction() - def func(self, u: float, v: float) -> np.ndarray: + def func(self, u: float, v: float) -> Point3D: """Converts from spherical coordinates to cartesian. Parameters @@ -667,10 +684,10 @@ def func(self, u: float, v: float) -> np.ndarray: ], ) - def get_start(self) -> np.ndarray: + def get_start(self) -> Point3D: return self.start_point.get_center() - def get_end(self) -> np.ndarray: + def get_end(self) -> Point3D: return self.end_point.get_center() def _rotate_to_direction(self) -> None: @@ -703,7 +720,7 @@ def _rotate_to_direction(self) -> None: self._current_theta = theta self._current_phi = phi - def set_direction(self, direction: np.ndarray) -> None: + def set_direction(self, direction: Vector3DLike) -> None: """Changes the direction of the apex of the :class:`Cone`. Parameters @@ -711,10 +728,10 @@ def set_direction(self, direction: np.ndarray) -> None: direction The direction of the apex. """ - self.direction = direction + self.direction = np.array(direction) self._rotate_to_direction() - def get_direction(self) -> np.ndarray: + def get_direction(self) -> Vector3D: """Returns the current direction of the apex of the :class:`Cone`. Returns @@ -724,7 +741,7 @@ def get_direction(self) -> np.ndarray: """ return self.direction - def _set_start_and_end_attributes(self, direction): + def _set_start_and_end_attributes(self, direction: Vector3D) -> None: normalized_direction = direction * np.linalg.norm(direction) start = self.base_circle.get_center() @@ -770,18 +787,18 @@ def __init__( self, radius: float = 1, height: float = 2, - direction: np.ndarray = Z_AXIS, - v_range: Sequence[float] = [0, TAU], + direction: Vector3DLike = Z_AXIS, + v_range: tuple[float, float] = (0, TAU), show_ends: bool = True, - resolution: Sequence[int] = (24, 24), - **kwargs, + resolution: int | tuple[int, int] = (24, 24), + **kwargs: Any, ) -> None: self._height = height self.radius = radius super().__init__( self.func, resolution=resolution, - u_range=[-self._height / 2, self._height / 2], + u_range=(-self._height / 2, self._height / 2), v_range=v_range, **kwargs, ) @@ -813,7 +830,9 @@ def func(self, u: float, v: float) -> np.ndarray: def add_bases(self) -> None: """Adds the end caps of the cylinder.""" + opacity: float if config.renderer == RendererType.OPENGL: + assert isinstance(self, OpenGLMobject) color = self.color opacity = self.opacity elif config.renderer == RendererType.CAIRO: @@ -868,7 +887,7 @@ def _rotate_to_direction(self) -> None: self._current_theta = theta self._current_phi = phi - def set_direction(self, direction: np.ndarray) -> None: + def set_direction(self, direction: Vector3DLike) -> None: """Sets the direction of the central axis of the :class:`Cylinder`. Parameters @@ -927,15 +946,17 @@ def construct(self): def __init__( self, - start: np.ndarray = LEFT, - end: np.ndarray = RIGHT, + start: Point3DLike = LEFT, + end: Point3DLike = RIGHT, thickness: float = 0.02, color: ParsableManimColor | None = None, - resolution: int | Sequence[int] = 24, - **kwargs, + resolution: int | tuple[int, int] = 24, + **kwargs: Any, ): self.thickness = thickness - self.resolution = (2, resolution) if isinstance(resolution, int) else resolution + self.resolution: tuple[int, int] = ( + (2, resolution) if isinstance(resolution, int) else resolution + ) start = np.array(start, dtype=np.float64) end = np.array(end, dtype=np.float64) @@ -945,7 +966,7 @@ def __init__( self.set_color(color) def set_start_and_end_attrs( - self, start: np.ndarray, end: np.ndarray, **kwargs + self, start: Point3DLike, end: Point3DLike, **kwargs: Any ) -> None: """Sets the start and end points of the line. @@ -963,7 +984,7 @@ def set_start_and_end_attrs( rough_end = self.pointify(end) self.vect = rough_end - rough_start self.length = np.linalg.norm(self.vect) - self.direction = normalize(self.vect) + self.direction: Vector3D = normalize(self.vect) # Now that we know the direction between them, # we can the appropriate boundary point from # start and end, if they're mobjects @@ -1005,7 +1026,7 @@ def pointify( return mob.get_boundary_point(direction) return np.array(mob_or_point) - def get_start(self) -> np.ndarray: + def get_start(self) -> Point3D: """Returns the starting point of the :class:`Line3D`. Returns @@ -1015,7 +1036,7 @@ def get_start(self) -> np.ndarray: """ return self.start - def get_end(self) -> np.ndarray: + def get_end(self) -> Point3D: """Returns the ending point of the :class:`Line3D`. Returns @@ -1031,7 +1052,7 @@ def parallel_to( line: Line3D, point: Point3DLike = ORIGIN, length: float = 5, - **kwargs, + **kwargs: Any, ) -> Line3D: """Returns a line parallel to another line going through a given point. @@ -1077,9 +1098,9 @@ def construct(self): def perpendicular_to( cls, line: Line3D, - point: Vector3DLike = ORIGIN, + point: Point3DLike = ORIGIN, length: float = 5, - **kwargs, + **kwargs: Any, ) -> Line3D: """Returns a line perpendicular to another line going through a given point. @@ -1167,14 +1188,14 @@ def construct(self): def __init__( self, - start: np.ndarray = LEFT, - end: np.ndarray = RIGHT, + start: Point3DLike = LEFT, + end: Point3DLike = RIGHT, thickness: float = 0.02, height: float = 0.3, base_radius: float = 0.08, color: ParsableManimColor = WHITE, - resolution: int | Sequence[int] = 24, - **kwargs, + resolution: int | tuple[int, int] = 24, + **kwargs: Any, ) -> None: super().__init__( start=start, @@ -1241,10 +1262,10 @@ def __init__( self, major_radius: float = 3, minor_radius: float = 1, - u_range: Sequence[float] = (0, TAU), - v_range: Sequence[float] = (0, TAU), - resolution: tuple[int, int] | None = None, - **kwargs, + u_range: tuple[float, float] = (0, TAU), + v_range: tuple[float, float] = (0, TAU), + resolution: int | tuple[int, int] | None = None, + **kwargs: Any, ) -> None: if config.renderer == RendererType.OPENGL: res_value = (101, 101) @@ -1263,7 +1284,7 @@ def __init__( **kwargs, ) - def func(self, u: float, v: float) -> np.ndarray: + def func(self, u: float, v: float) -> Point3D: """The z values defining the :class:`Torus` being plotted. Returns diff --git a/mypy.ini b/mypy.ini index 4a96a6d7f8..91ce5479c8 100644 --- a/mypy.ini +++ b/mypy.ini @@ -109,9 +109,6 @@ ignore_errors = True [mypy-manim.mobject.table] ignore_errors = True -[mypy-manim.mobject.three_d.three_dimensions] -ignore_errors = True - [mypy-manim.mobject.types.image_mobject] ignore_errors = True