diff --git a/manim/camera/camera.py b/manim/camera/camera.py index e2137fc858..cab8c5ac0f 100644 --- a/manim/camera/camera.py +++ b/manim/camera/camera.py @@ -14,11 +14,19 @@ import cairo import numpy as np +import numpy.typing as npt from PIL import Image from scipy.spatial.distance import pdist from typing_extensions import Self -from manim.typing import MatrixMN, PixelArray, Point3D, Point3D_Array +from manim.typing import ( + FloatRGBA_Array, + FloatRGBALike_Array, + ManimInt, + PixelArray, + Point3D, + Point3D_Array, +) from .. import config, logger from ..constants import * @@ -211,8 +219,8 @@ def type_or_raise( type[Mobject], Callable[[list[Mobject], PixelArray], Any] ] = { VMobject: self.display_multiple_vectorized_mobjects, # type: ignore[dict-item] - PMobject: self.display_multiple_point_cloud_mobjects, - AbstractImageMobject: self.display_multiple_image_mobjects, + PMobject: self.display_multiple_point_cloud_mobjects, # type: ignore[dict-item] + AbstractImageMobject: self.display_multiple_image_mobjects, # type: ignore[dict-item] Mobject: lambda batch, pa: batch, # Do nothing } # We have to check each type in turn because we are dealing with @@ -723,7 +731,7 @@ def set_cairo_context_path(self, ctx: cairo.Context, vmobject: VMobject) -> Self return self def set_cairo_context_color( - self, ctx: cairo.Context, rgbas: MatrixMN, vmobject: VMobject + self, ctx: cairo.Context, rgbas: FloatRGBALike_Array, vmobject: VMobject ) -> Self: """Sets the color of the cairo context @@ -818,7 +826,7 @@ def apply_stroke( def get_stroke_rgbas( self, vmobject: VMobject, background: bool = False - ) -> PixelArray: + ) -> FloatRGBA_Array: """Gets the RGBA array for the stroke of the passed VMobject. @@ -837,7 +845,7 @@ def get_stroke_rgbas( """ return vmobject.get_stroke_rgbas(background) - def get_fill_rgbas(self, vmobject: VMobject) -> PixelArray: + def get_fill_rgbas(self, vmobject: VMobject) -> FloatRGBA_Array: """Returns the RGBA array of the fill of the passed VMobject Parameters @@ -898,7 +906,7 @@ def display_multiple_background_colored_vmobjects( # As a result, the other methods do not have as detailed docstrings as would be preferred. def display_multiple_point_cloud_mobjects( - self, pmobjects: list, pixel_array: PixelArray + self, pmobjects: Iterable[PMobject], pixel_array: PixelArray ) -> None: """Displays multiple PMobjects by modifying the passed pixel array. @@ -921,8 +929,8 @@ def display_multiple_point_cloud_mobjects( def display_point_cloud( self, pmobject: PMobject, - points: list, - rgbas: np.ndarray, + points: Point3D_Array, + rgbas: FloatRGBA_Array, thickness: float, pixel_array: PixelArray, ) -> None: @@ -972,7 +980,9 @@ def display_point_cloud( pixel_array[:, :] = new_pa.reshape((ph, pw, rgba_len)) def display_multiple_image_mobjects( - self, image_mobjects: list, pixel_array: np.ndarray + self, + image_mobjects: Iterable[AbstractImageMobject], + pixel_array: PixelArray, ) -> None: """Displays multiple image mobjects by modifying the passed pixel_array. @@ -1121,8 +1131,8 @@ def transform_points_pre_display( def points_to_pixel_coords( self, mobject: Mobject, - points: np.ndarray, - ) -> np.ndarray: # TODO: Write more detailed docstrings for this method. + points: Point3D_Array, + ) -> npt.NDArray[ManimInt]: # TODO: Write more detailed docstrings for this method. points = self.transform_points_pre_display(mobject, points) shifted_points = points - self.frame_center diff --git a/manim/camera/three_d_camera.py b/manim/camera/three_d_camera.py index 3d44b3e910..e20512ab9e 100644 --- a/manim/camera/three_d_camera.py +++ b/manim/camera/three_d_camera.py @@ -20,6 +20,7 @@ from manim.mobject.types.vectorized_mobject import VMobject from manim.mobject.value_tracker import ValueTracker from manim.typing import ( + FloatRGBA_Array, MatrixMN, Point3D, Point3D_Array, @@ -109,7 +110,9 @@ def get_value_trackers(self) -> list[ValueTracker]: self.zoom_tracker, ] - def modified_rgbas(self, vmobject: VMobject, rgbas: MatrixMN) -> MatrixMN: + def modified_rgbas( + self, vmobject: VMobject, rgbas: FloatRGBA_Array + ) -> FloatRGBA_Array: if not self.should_apply_shading: return rgbas if vmobject.shade_in_3d and (vmobject.get_num_points() > 0): @@ -137,12 +140,12 @@ def get_stroke_rgbas( self, vmobject: VMobject, background: bool = False, - ) -> MatrixMN: # NOTE : DocStrings From parent + ) -> FloatRGBA_Array: # NOTE : DocStrings From parent return self.modified_rgbas(vmobject, vmobject.get_stroke_rgbas(background)) def get_fill_rgbas( self, vmobject: VMobject - ) -> MatrixMN: # NOTE : DocStrings From parent + ) -> FloatRGBA_Array: # NOTE : DocStrings From parent return self.modified_rgbas(vmobject, vmobject.get_fill_rgbas()) def get_mobjects_to_display( diff --git a/manim/mobject/opengl/opengl_mobject.py b/manim/mobject/opengl/opengl_mobject.py index c8312eddc9..5fd7def78f 100644 --- a/manim/mobject/opengl/opengl_mobject.py +++ b/manim/mobject/opengl/opengl_mobject.py @@ -53,6 +53,8 @@ from manim.renderer.shader_wrapper import ShaderWrapper from manim.typing import ( + FloatRGB_Array, + FloatRGBA_Array, ManimFloat, MappingFunction, MatrixMN, @@ -328,9 +330,9 @@ def init_data(self) -> None: """Initializes the ``points``, ``bounding_box`` and ``rgbas`` attributes and groups them into self.data. Subclasses can inherit and overwrite this method to extend `self.data`. """ - self.points = np.zeros((0, 3)) - self.bounding_box = np.zeros((3, 3)) - self.rgbas = np.zeros((1, 4)) + self.points: Point3D_Array = np.zeros((0, 3)) + self.bounding_box: Point3D_Array = np.zeros((3, 3)) + self.rgbas: FloatRGBA_Array = np.zeros((1, 4)) def init_colors(self) -> object: """Initializes the colors. @@ -2082,7 +2084,7 @@ def set_rgba_array( recurse: bool = True, ) -> Self: if color is not None: - rgbs = np.array([color_to_rgb(c) for c in listify(color)]) + rgbs: FloatRGB_Array = np.array([color_to_rgb(c) for c in listify(color)]) if opacity is not None: opacities = listify(opacity) @@ -2105,14 +2107,16 @@ def set_rgba_array( # Color and opacity if color is not None and opacity is not None: - rgbas = np.array([[*rgb, o] for rgb, o in zip(*make_even(rgbs, opacities))]) + rgbas: FloatRGBA_Array = np.array( + [[*rgb, o] for rgb, o in zip(*make_even(rgbs, opacities))] + ) for mob in self.get_family(recurse): mob.data[name] = rgbas.copy() return self def set_rgba_array_direct( self, - rgbas: npt.NDArray[RGBA_Array_Float], + rgbas: FloatRGBA_Array, name: str = "rgbas", recurse: bool = True, ) -> Self: @@ -2794,6 +2798,7 @@ def set_color_by_xyz_func( # of the shader code for char in "xyz": glsl_snippet = glsl_snippet.replace(char, "point." + char) + # TODO: get_colormap_list does not exist rgb_list = get_colormap_list(colormap) self.set_color_by_code( f"color.rgb = float_to_color({glsl_snippet}, {float(min_value)}, {float(max_value)}, {get_colormap_code(rgb_list)});", diff --git a/manim/mobject/opengl/opengl_point_cloud_mobject.py b/manim/mobject/opengl/opengl_point_cloud_mobject.py index 72e196fb9b..ecdfc39469 100644 --- a/manim/mobject/opengl/opengl_point_cloud_mobject.py +++ b/manim/mobject/opengl/opengl_point_cloud_mobject.py @@ -2,9 +2,10 @@ __all__ = ["OpenGLPMobject", "OpenGLPGroup", "OpenGLPMPoint"] +from typing import TYPE_CHECKING + import moderngl import numpy as np -from typing_extensions import Self from manim.constants import * from manim.mobject.opengl.opengl_mobject import OpenGLMobject @@ -20,6 +21,16 @@ from manim.utils.config_ops import _Uniforms from manim.utils.iterables import resize_with_interpolation +if TYPE_CHECKING: + from typing_extensions import Self + + from manim.typing import ( + FloatRGBA_Array, + FloatRGBALike_Array, + Point3D_Array, + Point3DLike_Array, + ) + __all__ = ["OpenGLPMobject", "OpenGLPGroup", "OpenGLPMPoint"] @@ -48,14 +59,20 @@ def __init__( ) def reset_points(self) -> Self: - self.rgbas = np.zeros((1, 4)) - self.points = np.zeros((0, 3)) + self.rgbas: FloatRGBA_Array = np.zeros((1, 4)) + self.points: Point3D_Array = np.zeros((0, 3)) return self def get_array_attrs(self): return ["points", "rgbas"] - def add_points(self, points, rgbas=None, color=None, opacity=None): + def add_points( + self, + points: Point3DLike_Array, + rgbas: FloatRGBALike_Array | None = None, + color: ParsableManimColor | None = None, + opacity: float | None = None, + ) -> Self: """Add points. Points must be a Nx3 numpy array. diff --git a/manim/mobject/opengl/opengl_vectorized_mobject.py b/manim/mobject/opengl/opengl_vectorized_mobject.py index d8ced06d65..3f8dccbff8 100644 --- a/manim/mobject/opengl/opengl_vectorized_mobject.py +++ b/manim/mobject/opengl/opengl_vectorized_mobject.py @@ -84,6 +84,9 @@ class OpenGLVMobject(OpenGLMobject): stroke_shader_folder = "quadratic_bezier_stroke" fill_shader_folder = "quadratic_bezier_fill" + # TODO: although these are called "rgba" in singular, they are used as + # FloatRGBA_Arrays and should be called instead "rgbas" in plural for consistency. + # The same should probably apply for "stroke_width" and "unit_normal". fill_rgba = _Data() stroke_rgba = _Data() stroke_width = _Data() diff --git a/manim/mobject/types/point_cloud_mobject.py b/manim/mobject/types/point_cloud_mobject.py index f820f49bfc..dcef684dfe 100644 --- a/manim/mobject/types/point_cloud_mobject.py +++ b/manim/mobject/types/point_cloud_mobject.py @@ -33,7 +33,14 @@ import numpy.typing as npt from typing_extensions import Self - from manim.typing import ManimFloat, Point3DLike + from manim.typing import ( + FloatRGBA_Array, + FloatRGBALike_Array, + ManimFloat, + Point3D_Array, + Point3DLike, + Point3DLike_Array, + ) class PMobject(Mobject, metaclass=ConvertToOpenGL): @@ -70,8 +77,8 @@ def __init__(self, stroke_width: int = DEFAULT_STROKE_WIDTH, **kwargs: Any) -> N super().__init__(**kwargs) def reset_points(self) -> Self: - self.rgbas = np.zeros((0, 4)) - self.points = np.zeros((0, 3)) + self.rgbas: FloatRGBA_Array = np.zeros((0, 4)) + self.points: Point3D_Array = np.zeros((0, 3)) return self def get_array_attrs(self) -> list[str]: @@ -79,10 +86,10 @@ def get_array_attrs(self) -> list[str]: def add_points( self, - points: npt.NDArray, - rgbas: npt.NDArray | None = None, + points: Point3DLike_Array, + rgbas: FloatRGBALike_Array | None = None, color: ParsableManimColor | None = None, - alpha: float = 1, + alpha: float = 1.0, ) -> Self: """Add points. diff --git a/manim/mobject/types/vectorized_mobject.py b/manim/mobject/types/vectorized_mobject.py index f937d8ab58..cc2ecc2565 100644 --- a/manim/mobject/types/vectorized_mobject.py +++ b/manim/mobject/types/vectorized_mobject.py @@ -54,6 +54,8 @@ CubicBezierPath, CubicBezierPointsLike, CubicSpline, + FloatRGBA, + FloatRGBA_Array, ManimFloat, MappingFunction, Point2DLike, @@ -61,10 +63,8 @@ Point3D_Array, Point3DLike, Point3DLike_Array, - RGBA_Array_Float, Vector3D, Vector3DLike, - Zeros, ) # TODO @@ -215,8 +215,10 @@ def init_colors(self, propagate_colors: bool = True) -> Self: return self def generate_rgbas_array( - self, color: ManimColor | list[ManimColor], opacity: float | Iterable[float] - ) -> RGBA_Array_Float: + self, + color: ParsableManimColor | Iterable[ManimColor] | None, + opacity: float | Iterable[float], + ) -> FloatRGBA: """ First arg can be either a color, or a tuple/list of colors. Likewise, opacity can either be a float, or a tuple of floats. @@ -230,7 +232,7 @@ def generate_rgbas_array( opacities: list[float] = [ o if (o is not None) else 0.0 for o in tuplify(opacity) ] - rgbas: npt.NDArray[RGBA_Array_Float] = np.array( + rgbas: FloatRGBA_Array = np.array( [c.to_rgba_with_alpha(o) for c, o in zip(*make_even(colors, opacities))], ) @@ -245,7 +247,7 @@ def generate_rgbas_array( def update_rgbas_array( self, array_name: str, - color: ManimColor | None = None, + color: ParsableManimColor | Iterable[ManimColor] | None = None, opacity: float | None = None, ) -> Self: rgbas = self.generate_rgbas_array(color, opacity) @@ -313,7 +315,7 @@ def construct(self): for submobject in self.submobjects: submobject.set_fill(color, opacity, family) self.update_rgbas_array("fill_rgbas", color, opacity) - self.fill_rgbas: RGBA_Array_Float + self.fill_rgbas: FloatRGBA_Array if opacity is not None: self.fill_opacity = opacity return self @@ -539,7 +541,7 @@ def fade(self, darkness: float = 0.5, family: bool = True) -> Self: super().fade(darkness, family) return self - def get_fill_rgbas(self) -> RGBA_Array_Float | Zeros: + def get_fill_rgbas(self) -> FloatRGBA_Array: try: return self.fill_rgbas except AttributeError: @@ -572,13 +574,13 @@ def get_fill_colors(self) -> list[ManimColor | None]: def get_fill_opacities(self) -> npt.NDArray[ManimFloat]: return self.get_fill_rgbas()[:, 3] - def get_stroke_rgbas(self, background: bool = False) -> RGBA_Array_float | Zeros: + def get_stroke_rgbas(self, background: bool = False) -> FloatRGBA_Array: try: if background: - self.background_stroke_rgbas: RGBA_Array_Float + self.background_stroke_rgbas: FloatRGBA_Array rgbas = self.background_stroke_rgbas else: - self.stroke_rgbas: RGBA_Array_Float + self.stroke_rgbas: FloatRGBA_Array rgbas = self.stroke_rgbas return rgbas except AttributeError: diff --git a/manim/mobject/vector_field.py b/manim/mobject/vector_field.py index 44b3be6c0d..ff153c7a57 100644 --- a/manim/mobject/vector_field.py +++ b/manim/mobject/vector_field.py @@ -12,6 +12,7 @@ import random from collections.abc import Callable, Iterable, Sequence from math import ceil, floor +from typing import TYPE_CHECKING import numpy as np from PIL import Image @@ -42,6 +43,15 @@ from ..utils.rate_functions import ease_out_sine, linear from ..utils.simple_functions import sigmoid +if TYPE_CHECKING: + from manim.typing import ( + FloatRGB, + FloatRGB_Array, + FloatRGBA_Array, + Point3D, + Vector3D, + ) + DEFAULT_SCALAR_FIELD_COLORS: list = [BLUE_E, GREEN, YELLOW, RED] @@ -73,9 +83,9 @@ class VectorField(VGroup): def __init__( self, - func: Callable[[np.ndarray], np.ndarray], + func: Callable[[Point3D], Vector3D], color: ParsableManimColor | None = None, - color_scheme: Callable[[np.ndarray], float] | None = None, + color_scheme: Callable[[Vector3D], float] | None = None, min_color_scheme_value: float = 0, max_color_scheme_value: float = 2, colors: Sequence[ParsableManimColor] = DEFAULT_SCALAR_FIELD_COLORS, @@ -87,13 +97,13 @@ def __init__( self.single_color = False if color_scheme is None: - def color_scheme(p): - return np.linalg.norm(p) + def color_scheme(vec: Vector3D) -> float: + return np.linalg.norm(vec) self.color_scheme = color_scheme # TODO maybe other default for direction? - self.rgbs = np.array(list(map(color_to_rgb, colors))) + self.rgbs: FloatRGB_Array = np.array(list(map(color_to_rgb, colors))) - def pos_to_rgb(pos: np.ndarray) -> tuple[float, float, float, float]: + def pos_to_rgb(pos: Point3D) -> FloatRGB: vec = self.func(pos) color_value = np.clip( self.color_scheme(vec), @@ -106,8 +116,8 @@ def pos_to_rgb(pos: np.ndarray) -> tuple[float, float, float, float]: color_value, ) alpha *= len(self.rgbs) - 1 - c1 = self.rgbs[int(alpha)] - c2 = self.rgbs[min(int(alpha + 1), len(self.rgbs) - 1)] + c1: FloatRGB = self.rgbs[int(alpha)] + c2: FloatRGB = self.rgbs[min(int(alpha + 1), len(self.rgbs) - 1)] alpha %= 1 return interpolate(c1, c2, alpha) @@ -417,7 +427,7 @@ def get_vectorized_rgba_gradient_function( start: float, end: float, colors: Iterable[ParsableManimColor], - ): + ) -> Callable[[Sequence[float], float], FloatRGBA_Array]: """ Generates a gradient of rgbas as a numpy array @@ -434,9 +444,9 @@ def get_vectorized_rgba_gradient_function( ------- function to generate the gradients as numpy arrays representing rgba values """ - rgbs = np.array([color_to_rgb(c) for c in colors]) + rgbs: FloatRGB_Array = np.array([color_to_rgb(c) for c in colors]) - def func(values, opacity=1): + def func(values: Sequence[float], opacity: float = 1.0) -> FloatRGBA_Array: alphas = inverse_interpolate(start, end, np.array(values)) alphas = np.clip(alphas, 0, 1) scaled_alphas = alphas * (len(rgbs) - 1) @@ -444,12 +454,14 @@ def func(values, opacity=1): next_indices = np.clip(indices + 1, 0, len(rgbs) - 1) inter_alphas = scaled_alphas % 1 inter_alphas = inter_alphas.repeat(3).reshape((len(indices), 3)) - result = interpolate(rgbs[indices], rgbs[next_indices], inter_alphas) - result = np.concatenate( - (result, np.full([len(result), 1], opacity)), + new_rgbs: FloatRGB_Array = interpolate( + rgbs[indices], rgbs[next_indices], inter_alphas + ) + new_rgbas: FloatRGBA_Array = np.concatenate( + (new_rgbs, np.full([len(new_rgbs), 1], opacity)), axis=1, ) - return result + return new_rgbas return func diff --git a/manim/renderer/shader_wrapper.py b/manim/renderer/shader_wrapper.py index 8a2b0d1fbe..84975b4f59 100644 --- a/manim/renderer/shader_wrapper.py +++ b/manim/renderer/shader_wrapper.py @@ -4,10 +4,14 @@ import logging import re from pathlib import Path +from typing import TYPE_CHECKING import moderngl import numpy as np +if TYPE_CHECKING: + from manim.typing import FloatRGBLike_Array + # Mobjects that should be rendered with # the same shader will be organized and # clumped together based on keeping track @@ -193,6 +197,6 @@ def get_shader_code_from_file(filename: Path) -> str | None: return result -def get_colormap_code(rgb_list): +def get_colormap_code(rgb_list: FloatRGBLike_Array) -> str: data = ",".join("vec3({}, {}, {})".format(*rgb) for rgb in rgb_list) return f"vec3[{len(rgb_list)}]({data})" diff --git a/manim/typing.py b/manim/typing.py index 5682ee4eed..3242781c78 100644 --- a/manim/typing.py +++ b/manim/typing.py @@ -32,20 +32,24 @@ "ManimFloat", "ManimInt", "ManimColorDType", - "RGB_Array_Float", - "RGB_Tuple_Float", - "RGB_Array_Int", - "RGB_Tuple_Int", - "RGBA_Array_Float", - "RGBA_Tuple_Float", - "RGBA_Array_Int", - "RGBA_Tuple_Int", - "HSV_Array_Float", - "HSV_Tuple_Float", - "HSL_Array_Float", - "HSL_Tuple_Float", - "HSVA_Array_Float", - "HSVA_Tuple_Float", + "FloatRGB", + "FloatRGBLike", + "FloatRGB_Array", + "FloatRGBLike_Array", + "IntRGB", + "IntRGBLike", + "FloatRGBA", + "FloatRGBALike", + "FloatRGBA_Array", + "FloatRGBALike_Array", + "IntRGBA", + "IntRGBALike", + "FloatHSV", + "FloatHSVLike", + "FloatHSL", + "FloatHSLLike", + "FloatHSVA", + "FloatHSVALike", "ManimColorInternal", "PointDType", "Point2D", @@ -143,7 +147,7 @@ double-precision float between 0 and 1. """ -RGB_Array_Float: TypeAlias = npt.NDArray[ManimColorDType] +FloatRGB: TypeAlias = npt.NDArray[ManimColorDType] """``shape: (3,)`` A :class:`numpy.ndarray` of 3 floats between 0 and 1, representing a @@ -153,17 +157,32 @@ Blue in the represented color. """ -RGB_Tuple_Float: TypeAlias = tuple[float, float, float] +FloatRGBLike: TypeAlias = Union[FloatRGB, tuple[float, float, float]] """``shape: (3,)`` -A tuple of 3 floats between 0 and 1, representing a color in RGB +An array of 3 floats between 0 and 1, representing a color in RGB format. -Its components describe, in order, the intensity of Red, Green, and -Blue in the represented color. +This represents anything which can be converted to a :class:`FloatRGB` NumPy +array. """ -RGB_Array_Int: TypeAlias = npt.NDArray[ManimInt] +FloatRGB_Array: TypeAlias = npt.NDArray[ManimColorDType] +"""``shape: (M, 3)`` + +A :class:`numpy.ndarray` of many rows of 3 floats representing RGB colors. +""" + +FloatRGBLike_Array: TypeAlias = Union[FloatRGB_Array, Sequence[FloatRGBLike]] +"""``shape: (M, 3)`` + +An array of many rows of 3 floats representing RGB colors. + +This represents anything which can be converted to a :class:`FloatRGB_Array` NumPy +array. +""" + +IntRGB: TypeAlias = npt.NDArray[ManimInt] """``shape: (3,)`` A :class:`numpy.ndarray` of 3 integers between 0 and 255, @@ -173,17 +192,17 @@ Blue in the represented color. """ -RGB_Tuple_Int: TypeAlias = tuple[int, int, int] +IntRGBLike: TypeAlias = Union[IntRGB, tuple[int, int, int]] """``shape: (3,)`` -A tuple of 3 integers between 0 and 255, representing a color in RGB +An array of 3 integers between 0 and 255, representing a color in RGB format. -Its components describe, in order, the intensity of Red, Green, and -Blue in the represented color. +This represents anything which can be converted to an :class:`IntRGB` NumPy +array. """ -RGBA_Array_Float: TypeAlias = npt.NDArray[ManimColorDType] +FloatRGBA: TypeAlias = npt.NDArray[ManimColorDType] """``shape: (4,)`` A :class:`numpy.ndarray` of 4 floats between 0 and 1, representing a @@ -193,17 +212,32 @@ and Alpha (opacity) in the represented color. """ -RGBA_Tuple_Float: TypeAlias = tuple[float, float, float, float] +FloatRGBALike: TypeAlias = Union[FloatRGBA, tuple[float, float, float, float]] """``shape: (4,)`` -A tuple of 4 floats between 0 and 1, representing a color in RGBA +An array of 4 floats between 0 and 1, representing a color in RGBA format. -Its components describe, in order, the intensity of Red, Green, Blue -and Alpha (opacity) in the represented color. +This represents anything which can be converted to a :class:`FloatRGBA` NumPy +array. """ -RGBA_Array_Int: TypeAlias = npt.NDArray[ManimInt] +FloatRGBA_Array: TypeAlias = npt.NDArray[ManimColorDType] +"""``shape: (M, 4)`` + +A :class:`numpy.ndarray` of many rows of 4 floats representing RGBA colors. +""" + +FloatRGBALike_Array: TypeAlias = Union[FloatRGBA_Array, Sequence[FloatRGBALike]] +"""``shape: (M, 4)`` + +An array of many rows of 4 floats representing RGBA colors. + +This represents anything which can be converted to a :class:`FloatRGBA_Array` NumPy +array. +""" + +IntRGBA: TypeAlias = npt.NDArray[ManimInt] """``shape: (4,)`` A :class:`numpy.ndarray` of 4 integers between 0 and 255, @@ -213,17 +247,17 @@ and Alpha (opacity) in the represented color. """ -RGBA_Tuple_Int: TypeAlias = tuple[int, int, int, int] +IntRGBALike: TypeAlias = Union[IntRGBA, tuple[int, int, int, int]] """``shape: (4,)`` -A tuple of 4 integers between 0 and 255, representing a color in RGBA +An array of 4 integers between 0 and 255, representing a color in RGBA format. -Its components describe, in order, the intensity of Red, Green, Blue -and Alpha (opacity) in the represented color. +This represents anything which can be converted to an :class:`IntRGBA` NumPy +array. """ -HSV_Array_Float: TypeAlias = RGB_Array_Float +FloatHSV: TypeAlias = FloatRGB """``shape: (3,)`` A :class:`numpy.ndarray` of 3 floats between 0 and 1, representing a @@ -233,17 +267,17 @@ Brightness) in the represented color. """ -HSV_Tuple_Float: TypeAlias = RGB_Tuple_Float +FloatHSVLike: TypeAlias = FloatRGBLike """``shape: (3,)`` -A tuple of 3 floats between 0 and 1, representing a color in HSV (or +An array of 3 floats between 0 and 1, representing a color in HSV (or HSB) format. -Its components describe, in order, the Hue, Saturation and Value (or -Brightness) in the represented color. +This represents anything which can be converted to a :class:`FloatHSV` NumPy +array. """ -HSVA_Array_Float: TypeAlias = RGBA_Array_Float +FloatHSVA: TypeAlias = FloatRGBA """``shape: (4,)`` A :class:`numpy.ndarray` of 4 floats between 0 and 1, representing a @@ -253,17 +287,17 @@ Brightness) in the represented color. """ -HSVA_Tuple_Float: TypeAlias = RGBA_Tuple_Float +FloatHSVALike: TypeAlias = FloatRGBALike """``shape: (4,)`` -A tuple of 4 floats between 0 and 1, representing a color in HSVA (or +An array of 4 floats between 0 and 1, representing a color in HSVA (or HSBA) format. -Its components describe, in order, the Hue, Saturation and Value (or -Brightness) in the represented color. +This represents anything which can be converted to a :class:`FloatHSVA` NumPy +array. """ -HSL_Array_Float: TypeAlias = RGB_Array_Float +FloatHSL: TypeAlias = FloatRGB """``shape: (3,)`` A :class:`numpy.ndarray` of 3 floats between 0 and 1, representing a @@ -273,17 +307,16 @@ in the represented color. """ -HSL_Tuple_Float: TypeAlias = RGB_Tuple_Float +FloatHSLLike: TypeAlias = FloatRGBLike """``shape: (3,)`` -A :class:`numpy.ndarray` of 3 floats between 0 and 1, representing a -color in HSL format. +An array of 3 floats between 0 and 1, representing a color in HSL format. -Its components describe, in order, the Hue, Saturation and Lightness -in the represented color. +This represents anything which can be converted to a :class:`FloatHSL` NumPy +array. """ -ManimColorInternal: TypeAlias = RGBA_Array_Float +ManimColorInternal: TypeAlias = FloatRGBA """``shape: (4,)`` Internal color representation used by :class:`~.ManimColor`, diff --git a/manim/utils/color/core.py b/manim/utils/color/core.py index af25992e59..7bc7b08033 100644 --- a/manim/utils/color/core.py +++ b/manim/utils/color/core.py @@ -77,24 +77,24 @@ from typing_extensions import Self, TypeAlias, TypeIs, override from manim.typing import ( - HSL_Array_Float, - HSL_Tuple_Float, - HSV_Array_Float, - HSV_Tuple_Float, - HSVA_Array_Float, - HSVA_Tuple_Float, + FloatHSL, + FloatHSLLike, + FloatHSV, + FloatHSVA, + FloatHSVALike, + FloatHSVLike, + FloatRGB, + FloatRGBA, + FloatRGBALike, + FloatRGBLike, + IntRGB, + IntRGBA, + IntRGBALike, + IntRGBLike, ManimColorDType, ManimColorInternal, ManimFloat, Point3D, - RGB_Array_Float, - RGB_Array_Int, - RGB_Tuple_Float, - RGB_Tuple_Int, - RGBA_Array_Float, - RGBA_Array_Int, - RGBA_Tuple_Float, - RGBA_Tuple_Int, Vector3D, ) @@ -190,9 +190,9 @@ def __init__( length = len(value) if all(isinstance(x, float) for x in value): if length == 3: - self._internal_value = ManimColor._internal_from_rgb(value, alpha) # type: ignore[arg-type] + self._internal_value = ManimColor._internal_from_rgb(value, alpha) elif length == 4: - self._internal_value = ManimColor._internal_from_rgba(value) # type: ignore[arg-type] + self._internal_value = ManimColor._internal_from_rgba(value) else: raise ValueError( f"ManimColor only accepts lists/tuples/arrays of length 3 or 4, not {length}" @@ -200,11 +200,10 @@ def __init__( else: if length == 3: self._internal_value = ManimColor._internal_from_int_rgb( - value, # type: ignore[arg-type] - alpha, + value, alpha ) elif length == 4: - self._internal_value = ManimColor._internal_from_int_rgba(value) # type: ignore[arg-type] + self._internal_value = ManimColor._internal_from_int_rgba(value) else: raise ValueError( f"ManimColor only accepts lists/tuples/arrays of length 3 or 4, not {length}" @@ -336,7 +335,7 @@ def _internal_from_hex_string(hex_: str, alpha: float) -> ManimColorInternal: @staticmethod def _internal_from_int_rgb( - rgb: RGB_Tuple_Int, alpha: float = 1.0 + rgb: IntRGBLike, alpha: float = 1.0 ) -> ManimColorInternal: """Internal function for converting an RGB tuple of integers into the internal representation of a :class:`ManimColor`. @@ -361,9 +360,7 @@ def _internal_from_int_rgb( return value @staticmethod - def _internal_from_rgb( - rgb: RGB_Tuple_Float, alpha: float = 1.0 - ) -> ManimColorInternal: + def _internal_from_rgb(rgb: FloatRGBLike, alpha: float = 1.0) -> ManimColorInternal: """Internal function for converting a rgb tuple of floats into the internal representation of a :class:`ManimColor`. @@ -387,7 +384,7 @@ def _internal_from_rgb( return value @staticmethod - def _internal_from_int_rgba(rgba: RGBA_Tuple_Int) -> ManimColorInternal: + def _internal_from_int_rgba(rgba: IntRGBALike) -> ManimColorInternal: """Internal function for converting an RGBA tuple of integers into the internal representation of a :class:`ManimColor`. @@ -406,7 +403,7 @@ def _internal_from_int_rgba(rgba: RGBA_Tuple_Int) -> ManimColorInternal: return np.asarray(rgba, dtype=ManimColorDType) / 255 @staticmethod - def _internal_from_rgba(rgba: RGBA_Tuple_Float) -> ManimColorInternal: + def _internal_from_rgba(rgba: FloatRGBALike) -> ManimColorInternal: """Internal function for converting an RGBA tuple of floats into the internal representation of a :class:`ManimColor`. @@ -470,7 +467,7 @@ def to_integer(self) -> int: tmp = (self._internal_value[:3] * 255).astype(dtype=np.byte).tobytes() return int.from_bytes(tmp, "big") - def to_rgb(self) -> RGB_Array_Float: + def to_rgb(self) -> FloatRGB: """Convert the current :class:`ManimColor` into an RGB array of floats. Returns @@ -480,7 +477,7 @@ def to_rgb(self) -> RGB_Array_Float: """ return self._internal_value[:3] - def to_int_rgb(self) -> RGB_Array_Int: + def to_int_rgb(self) -> IntRGB: """Convert the current :class:`ManimColor` into an RGB array of integers. Returns @@ -490,7 +487,7 @@ def to_int_rgb(self) -> RGB_Array_Int: """ return (self._internal_value[:3] * 255).astype(int) - def to_rgba(self) -> RGBA_Array_Float: + def to_rgba(self) -> FloatRGBA: """Convert the current :class:`ManimColor` into an RGBA array of floats. Returns @@ -500,7 +497,7 @@ def to_rgba(self) -> RGBA_Array_Float: """ return self._internal_value - def to_int_rgba(self) -> RGBA_Array_Int: + def to_int_rgba(self) -> IntRGBA: """Convert the current ManimColor into an RGBA array of integers. @@ -511,7 +508,7 @@ def to_int_rgba(self) -> RGBA_Array_Int: """ return (self._internal_value * 255).astype(int) - def to_rgba_with_alpha(self, alpha: float) -> RGBA_Array_Float: + def to_rgba_with_alpha(self, alpha: float) -> FloatRGBA: """Convert the current :class:`ManimColor` into an RGBA array of floats. This is similar to :meth:`to_rgba`, but you can change the alpha value. @@ -527,7 +524,7 @@ def to_rgba_with_alpha(self, alpha: float) -> RGBA_Array_Float: """ return np.fromiter((*self._internal_value[:3], alpha), dtype=ManimColorDType) - def to_int_rgba_with_alpha(self, alpha: float) -> RGBA_Array_Int: + def to_int_rgba_with_alpha(self, alpha: float) -> IntRGBA: """Convert the current :class:`ManimColor` into an RGBA array of integers. This is similar to :meth:`to_int_rgba`, but you can change the alpha value. @@ -570,7 +567,7 @@ def to_hex(self, with_alpha: bool = False) -> str: tmp += f"{int(self._internal_value[3] * 255):02X}" return tmp - def to_hsv(self) -> HSV_Array_Float: + def to_hsv(self) -> FloatHSV: """Convert the :class:`ManimColor` to an HSV array. .. note:: @@ -587,7 +584,7 @@ def to_hsv(self) -> HSV_Array_Float: """ return np.array(colorsys.rgb_to_hsv(*self.to_rgb())) - def to_hsl(self) -> HSL_Array_Float: + def to_hsl(self) -> FloatHSL: """Convert the :class:`ManimColor` to an HSL array. .. note:: @@ -797,7 +794,7 @@ def _from_internal(cls, value: ManimColorInternal) -> Self: @classmethod def from_rgb( cls, - rgb: RGB_Array_Float | RGB_Tuple_Float | RGB_Array_Int | RGB_Tuple_Int, + rgb: FloatRGBLike | IntRGBLike, alpha: float = 1.0, ) -> Self: """Create a ManimColor from an RGB array. Automagically decides which type it @@ -824,9 +821,7 @@ def from_rgb( return cls._from_internal(ManimColor(rgb, alpha)._internal_value) @classmethod - def from_rgba( - cls, rgba: RGBA_Array_Float | RGBA_Tuple_Float | RGBA_Array_Int | RGBA_Tuple_Int - ) -> Self: + def from_rgba(cls, rgba: FloatRGBALike | IntRGBALike) -> Self: """Create a ManimColor from an RGBA Array. Automagically decides which type it is: ``int`` or ``float``. @@ -868,9 +863,7 @@ def from_hex(cls, hex_str: str, alpha: float = 1.0) -> Self: return cls._from_internal(ManimColor(hex_str, alpha)._internal_value) @classmethod - def from_hsv( - cls, hsv: HSV_Array_Float | HSV_Tuple_Float, alpha: float = 1.0 - ) -> Self: + def from_hsv(cls, hsv: FloatHSVLike, alpha: float = 1.0) -> Self: """Create a :class:`ManimColor` from an HSV array. Parameters @@ -890,9 +883,7 @@ def from_hsv( return cls._from_internal(ManimColor(rgb, alpha)._internal_value) @classmethod - def from_hsl( - cls, hsl: HSL_Array_Float | HSL_Tuple_Float, alpha: float = 1.0 - ) -> Self: + def from_hsl(cls, hsl: FloatHSLLike, alpha: float = 1.0) -> Self: """Create a :class:`ManimColor` from an HSL array. Parameters @@ -1102,11 +1093,11 @@ class HSV(ManimColor): def __init__( self, - hsv: HSV_Array_Float | HSV_Tuple_Float | HSVA_Array_Float | HSVA_Tuple_Float, + hsv: FloatHSVLike | FloatHSVALike, alpha: float = 1.0, ) -> None: super().__init__(None) - self.__hsv: HSVA_Array_Float + self.__hsv: FloatHSVA if len(hsv) == 3: self.__hsv = np.asarray((*hsv, alpha)) elif len(hsv) == 4: @@ -1224,14 +1215,10 @@ def _internal_value(self, value: ManimColorInternal) -> None: ManimColor, int, str, - RGB_Tuple_Int, - RGB_Tuple_Float, - RGBA_Tuple_Int, - RGBA_Tuple_Float, - RGB_Array_Int, - RGB_Array_Float, - RGBA_Array_Int, - RGBA_Array_Float, + IntRGBLike, + FloatRGBLike, + IntRGBALike, + FloatRGBALike, ] """`ParsableManimColor` represents all the types which can be parsed to a :class:`ManimColor` in Manim. @@ -1241,7 +1228,7 @@ def _internal_value(self, value: ManimColorInternal) -> None: ManimColorT = TypeVar("ManimColorT", bound=ManimColor) -def color_to_rgb(color: ParsableManimColor) -> RGB_Array_Float: +def color_to_rgb(color: ParsableManimColor) -> FloatRGB: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.to_rgb`. @@ -1258,7 +1245,7 @@ def color_to_rgb(color: ParsableManimColor) -> RGB_Array_Float: return ManimColor(color).to_rgb() -def color_to_rgba(color: ParsableManimColor, alpha: float = 1.0) -> RGBA_Array_Float: +def color_to_rgba(color: ParsableManimColor, alpha: float = 1.0) -> FloatRGBA: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.to_rgba_with_alpha`. @@ -1278,7 +1265,7 @@ def color_to_rgba(color: ParsableManimColor, alpha: float = 1.0) -> RGBA_Array_F return ManimColor(color).to_rgba_with_alpha(alpha) -def color_to_int_rgb(color: ParsableManimColor) -> RGB_Array_Int: +def color_to_int_rgb(color: ParsableManimColor) -> IntRGB: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.to_int_rgb`. @@ -1295,7 +1282,7 @@ def color_to_int_rgb(color: ParsableManimColor) -> RGB_Array_Int: return ManimColor(color).to_int_rgb() -def color_to_int_rgba(color: ParsableManimColor, alpha: float = 1.0) -> RGBA_Array_Int: +def color_to_int_rgba(color: ParsableManimColor, alpha: float = 1.0) -> IntRGBA: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.to_int_rgba_with_alpha`. @@ -1315,9 +1302,7 @@ def color_to_int_rgba(color: ParsableManimColor, alpha: float = 1.0) -> RGBA_Arr return ManimColor(color).to_int_rgba_with_alpha(alpha) -def rgb_to_color( - rgb: RGB_Array_Float | RGB_Tuple_Float | RGB_Array_Int | RGB_Tuple_Int, -) -> ManimColor: +def rgb_to_color(rgb: FloatRGBLike | IntRGBLike) -> ManimColor: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.from_rgb`. @@ -1334,9 +1319,7 @@ def rgb_to_color( return ManimColor.from_rgb(rgb) -def rgba_to_color( - rgba: RGBA_Array_Float | RGBA_Tuple_Float | RGBA_Array_Int | RGBA_Tuple_Int, -) -> ManimColor: +def rgba_to_color(rgba: FloatRGBALike | IntRGBALike) -> ManimColor: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.from_rgba`. @@ -1353,9 +1336,7 @@ def rgba_to_color( return ManimColor.from_rgba(rgba) -def rgb_to_hex( - rgb: RGB_Array_Float | RGB_Tuple_Float | RGB_Array_Int | RGB_Tuple_Int, -) -> str: +def rgb_to_hex(rgb: FloatRGBLike | IntRGBLike) -> str: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.from_rgb` and :meth:`ManimColor.to_hex`. @@ -1372,7 +1353,7 @@ def rgb_to_hex( return ManimColor.from_rgb(rgb).to_hex() -def hex_to_rgb(hex_code: str) -> RGB_Array_Float: +def hex_to_rgb(hex_code: str) -> FloatRGB: """Helper function for use in functional style programming. Refer to :meth:`ManimColor.to_rgb`. @@ -1622,11 +1603,11 @@ def _random_color(cls) -> ManimColor: def get_shaded_rgb( - rgb: RGB_Array_Float, + rgb: FloatRGB, point: Point3D, unit_normal_vect: Vector3D, light_source: Point3D, -) -> RGB_Array_Float: +) -> FloatRGB: """Add light or shadow to the ``rgb`` color of some surface which is located at a given ``point`` in space and facing in the direction of ``unit_normal_vect``, depending on whether the surface is facing a ``light_source`` or away from it. @@ -1652,7 +1633,7 @@ def get_shaded_rgb( light = 0.5 * np.dot(unit_normal_vect, to_sun) ** 3 if light < 0: light *= 0.5 - shaded_rgb: RGB_Array_Float = rgb + light + shaded_rgb: FloatRGB = rgb + light return shaded_rgb diff --git a/manim/utils/images.py b/manim/utils/images.py index dd5c57858e..f3f49268f2 100644 --- a/manim/utils/images.py +++ b/manim/utils/images.py @@ -15,7 +15,7 @@ import numpy as np from PIL import Image -from manim.typing import RGBPixelArray +from manim.typing import RGBAPixelArray, RGBPixelArray from .. import config from ..utils.file_ops import seek_full_path_from_defaults @@ -55,7 +55,7 @@ def invert_image(image: np.array) -> Image: return Image.fromarray(arr) -def change_to_rgba_array(image: RGBPixelArray, dtype: str = "uint8") -> RGBPixelArray: +def change_to_rgba_array(image: RGBPixelArray, dtype: str = "uint8") -> RGBAPixelArray: """Converts an RGB array into RGBA with the alpha value opacity maxed.""" pa = image if len(pa.shape) == 2: diff --git a/tests/opengl/test_color_opengl.py b/tests/opengl/test_color_opengl.py index 3aeb2d6021..48339f9fbf 100644 --- a/tests/opengl/test_color_opengl.py +++ b/tests/opengl/test_color_opengl.py @@ -37,58 +37,58 @@ def test_background_color(using_opengl_renderer): def test_set_color(using_opengl_renderer): m = OpenGLMobject() assert m.color.to_hex() == "#FFFFFF" - np.all(m.rgbas == np.array((0.0, 0.0, 0.0, 1.0))) + np.all(m.rgbas == np.array([[0.0, 0.0, 0.0, 1.0]])) m.set_color(BLACK) assert m.color.to_hex() == "#000000" - np.all(m.rgbas == np.array((1.0, 1.0, 1.0, 1.0))) + np.all(m.rgbas == np.array([[1.0, 1.0, 1.0, 1.0]])) m.set_color(PURE_GREEN, opacity=0.5) assert m.color.to_hex() == "#00FF00" - np.all(m.rgbas == np.array((0.0, 1.0, 0.0, 0.5))) + np.all(m.rgbas == np.array([[0.0, 1.0, 0.0, 0.5]])) m = OpenGLVMobject() assert m.color.to_hex() == "#FFFFFF" - np.all(m.fill_rgba == np.array((0.0, 0.0, 0.0, 1.0))) - np.all(m.stroke_rgba == np.array((0.0, 0.0, 0.0, 1.0))) + np.all(m.fill_rgba == np.array([[0.0, 0.0, 0.0, 1.0]])) + np.all(m.stroke_rgba == np.array([[0.0, 0.0, 0.0, 1.0]])) m.set_color(BLACK) assert m.color.to_hex() == "#000000" - np.all(m.fill_rgba == np.array((1.0, 1.0, 1.0, 1.0))) - np.all(m.stroke_rgba == np.array((1.0, 1.0, 1.0, 1.0))) + np.all(m.fill_rgba == np.array([[1.0, 1.0, 1.0, 1.0]])) + np.all(m.stroke_rgba == np.array([[1.0, 1.0, 1.0, 1.0]])) m.set_color(PURE_GREEN, opacity=0.5) assert m.color.to_hex() == "#00FF00" - np.all(m.fill_rgba == np.array((0.0, 1.0, 0.0, 0.5))) - np.all(m.stroke_rgba == np.array((0.0, 1.0, 0.0, 0.5))) + np.all(m.fill_rgba == np.array([[0.0, 1.0, 0.0, 0.5]])) + np.all(m.stroke_rgba == np.array([[0.0, 1.0, 0.0, 0.5]])) def test_set_fill_color(using_opengl_renderer): m = OpenGLVMobject() assert m.fill_color.to_hex() == "#FFFFFF" - np.all(m.fill_rgba == np.array((0.0, 1.0, 0.0, 0.5))) + np.all(m.fill_rgba == np.array([[0.0, 1.0, 0.0, 0.5]])) m.set_fill(BLACK) assert m.fill_color.to_hex() == "#000000" - np.all(m.fill_rgba == np.array((1.0, 1.0, 1.0, 1.0))) + np.all(m.fill_rgba == np.array([[1.0, 1.0, 1.0, 1.0]])) m.set_fill(PURE_GREEN, opacity=0.5) assert m.fill_color.to_hex() == "#00FF00" - np.all(m.fill_rgba == np.array((0.0, 1.0, 0.0, 0.5))) + np.all(m.fill_rgba == np.array([[0.0, 1.0, 0.0, 0.5]])) def test_set_stroke_color(using_opengl_renderer): m = OpenGLVMobject() assert m.stroke_color.to_hex() == "#FFFFFF" - np.all(m.stroke_rgba == np.array((0.0, 1.0, 0.0, 0.5))) + np.all(m.stroke_rgba == np.array([[0.0, 1.0, 0.0, 0.5]])) m.set_stroke(BLACK) assert m.stroke_color.to_hex() == "#000000" - np.all(m.stroke_rgba == np.array((1.0, 1.0, 1.0, 1.0))) + np.all(m.stroke_rgba == np.array([[1.0, 1.0, 1.0, 1.0]])) m.set_stroke(PURE_GREEN, opacity=0.5) assert m.stroke_color.to_hex() == "#00FF00" - np.all(m.stroke_rgba == np.array((0.0, 1.0, 0.0, 0.5))) + np.all(m.stroke_rgba == np.array([[0.0, 1.0, 0.0, 0.5]])) def test_set_fill(using_opengl_renderer):