Source code for manim.mobject.types.point_cloud_mobject

"""Mobjects representing point clouds."""

from __future__ import annotations

__all__ = ["PMobject", "Mobject1D", "Mobject2D", "PGroup", "PointCloudDot", "Point"]

from collections.abc import Callable
from typing import TYPE_CHECKING, Any

import numpy as np

from manim.mobject.opengl.opengl_compatibility import ConvertToOpenGL
from manim.mobject.opengl.opengl_point_cloud_mobject import OpenGLPMobject

from ...constants import *
from ...mobject.mobject import Mobject
from ...utils.bezier import interpolate
from ...utils.color import (
    BLACK,
    PURE_YELLOW,
    WHITE,
    ManimColor,
    ParsableManimColor,
    color_gradient,
    color_to_rgba,
    rgba_to_color,
)
from ...utils.iterables import stretch_array_to_length

__all__ = ["PMobject", "Mobject1D", "Mobject2D", "PGroup", "PointCloudDot", "Point"]

if TYPE_CHECKING:
    from typing import Self

    import numpy.typing as npt

    from manim.typing import (
        FloatRGBA_Array,
        FloatRGBALike_Array,
        ManimFloat,
        Point3D_Array,
        Point3DLike,
        Point3DLike_Array,
    )


[docs] class PMobject(Mobject, metaclass=ConvertToOpenGL): """A disc made of a cloud of Dots Examples -------- .. manim:: PMobjectExample :save_last_frame: class PMobjectExample(Scene): def construct(self): pG = PGroup() # This is just a collection of PMobject's # As the scale factor increases, the number of points # removed increases. for sf in range(1, 9 + 1): p = PointCloudDot(density=20, radius=1).thin_out(sf) # PointCloudDot is a type of PMobject # and can therefore be added to a PGroup pG.add(p) # This organizes all the shapes in a grid. pG.arrange_in_grid() self.add(pG) """ def __init__(self, stroke_width: int = DEFAULT_STROKE_WIDTH, **kwargs: Any) -> None: self.stroke_width = stroke_width super().__init__(**kwargs)
[docs] def reset_points(self) -> Self: 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]: return super().get_array_attrs() + ["rgbas"]
[docs] def add_points( self, points: Point3DLike_Array, rgbas: FloatRGBALike_Array | None = None, color: ParsableManimColor | None = None, alpha: float = 1.0, ) -> Self: """Add points. Points must be a Nx3 numpy array. Rgbas must be a Nx4 numpy array if it is not None. """ if not isinstance(points, np.ndarray): points = np.array(points) num_new_points = len(points) self.points = np.append(self.points, points, axis=0) if rgbas is None: color = ManimColor(color) if color else self.color rgbas = np.repeat([color_to_rgba(color, alpha)], num_new_points, axis=0) elif len(rgbas) != len(points): raise ValueError("points and rgbas must have same length") self.rgbas = np.append(self.rgbas, rgbas, axis=0) return self
[docs] def set_color( self, color: ParsableManimColor = PURE_YELLOW, family: bool = True ) -> Self: rgba = color_to_rgba(color) mobs = self.family_members_with_points() if family else [self] for mob in mobs: mob.rgbas[:, :] = rgba self.color = ManimColor.parse(color) return self
def get_stroke_width(self) -> int: return self.stroke_width def set_stroke_width(self, width: int, family: bool = True) -> Self: mobs = self.family_members_with_points() if family else [self] for mob in mobs: mob.stroke_width = width return self
[docs] def set_color_by_gradient(self, *colors: ParsableManimColor) -> Self: self.rgbas = np.array( list(map(color_to_rgba, color_gradient(colors, len(self.points)))), ) return self
def set_colors_by_radial_gradient( self, center: Point3DLike | None = None, radius: float = 1, inner_color: ParsableManimColor = WHITE, outer_color: ParsableManimColor = BLACK, ) -> Self: start_rgba, end_rgba = list(map(color_to_rgba, [inner_color, outer_color])) if center is None: center = self.get_center() for mob in self.family_members_with_points(): distances = np.abs(self.points - center) alphas = np.linalg.norm(distances, axis=1) / radius mob.rgbas = np.array( np.array( [interpolate(start_rgba, end_rgba, alpha) for alpha in alphas], ), ) return self def match_colors(self, mobject: Mobject) -> Self: Mobject.align_data(self, mobject) self.rgbas = np.array(mobject.rgbas) return self def filter_out(self, condition: npt.NDArray) -> Self: for mob in self.family_members_with_points(): to_eliminate = ~np.apply_along_axis(condition, 1, mob.points) mob.points = mob.points[to_eliminate] mob.rgbas = mob.rgbas[to_eliminate] return self
[docs] def thin_out(self, factor: int = 5) -> Self: """Removes all but every nth point for n = factor""" for mob in self.family_members_with_points(): num_points = self.get_num_points() mob.apply_over_attr_arrays( lambda arr, n=num_points: arr[np.arange(0, n, factor)], # type: ignore[misc] ) return self
[docs] def sort_points( self, function: Callable[[npt.NDArray[ManimFloat]], float] = lambda p: p[0] ) -> Self: """Function is any map from R^3 to R""" for mob in self.family_members_with_points(): indices = np.argsort(np.apply_along_axis(function, 1, mob.points)) mob.apply_over_attr_arrays(lambda arr, idx=indices: arr[idx]) # type: ignore[misc] return self
def fade_to( self, color: ParsableManimColor, alpha: float, family: bool = True ) -> Self: self.rgbas = interpolate(self.rgbas, color_to_rgba(color), alpha) for mob in self.submobjects: mob.fade_to(color, alpha, family) return self def get_all_rgbas(self) -> npt.NDArray: return self.get_merged_array("rgbas") def ingest_submobjects(self) -> Self: attrs = self.get_array_attrs() arrays = list(map(self.get_merged_array, attrs)) for attr, array in zip(attrs, arrays, strict=True): setattr(self, attr, array) self.submobjects = [] return self
[docs] def get_color(self) -> ManimColor: return rgba_to_color(self.rgbas[0, :])
def point_from_proportion(self, alpha: float) -> Any: index = alpha * (self.get_num_points() - 1) return self.points[np.floor(index)]
[docs] @staticmethod def get_mobject_type_class() -> type[PMobject]: return PMobject
# Alignment def align_points_with_larger(self, larger_mobject: Mobject) -> None: assert isinstance(larger_mobject, PMobject) self.apply_over_attr_arrays( lambda a: stretch_array_to_length(a, larger_mobject.get_num_points()), )
[docs] def get_point_mobject(self, center: Point3DLike | None = None) -> Point: if center is None: center = self.get_center() return Point(center)
def interpolate_color( self, mobject1: Mobject, mobject2: Mobject, alpha: float ) -> Self: self.rgbas = interpolate(mobject1.rgbas, mobject2.rgbas, alpha) self.set_stroke_width( interpolate( mobject1.get_stroke_width(), mobject2.get_stroke_width(), alpha, ), ) return self def pointwise_become_partial(self, mobject: Mobject, a: float, b: float) -> None: lower_index, upper_index = (int(x * mobject.get_num_points()) for x in (a, b)) for attr in self.get_array_attrs(): full_array = getattr(mobject, attr) partial_array = full_array[lower_index:upper_index] setattr(self, attr, partial_array)
# TODO, Make the two implementations below non-redundant
[docs] class Mobject1D(PMobject, metaclass=ConvertToOpenGL): def __init__(self, density: int = DEFAULT_POINT_DENSITY_1D, **kwargs: Any) -> None: self.density = density self.epsilon = 1.0 / self.density super().__init__(**kwargs) def add_line( self, start: npt.NDArray, end: npt.NDArray, color: ParsableManimColor | None = None, ) -> None: start, end = list(map(np.array, [start, end])) length = np.linalg.norm(end - start) if length == 0: points = np.array([start]) else: epsilon = self.epsilon / length points = np.array( [interpolate(start, end, t) for t in np.arange(0, 1, epsilon)] ) self.add_points(points, color=color)
[docs] class Mobject2D(PMobject, metaclass=ConvertToOpenGL): def __init__(self, density: int = DEFAULT_POINT_DENSITY_2D, **kwargs: Any) -> None: self.density = density self.epsilon = 1.0 / self.density super().__init__(**kwargs)
[docs] class PGroup(PMobject): """A group for several point mobjects. Examples -------- .. manim:: PgroupExample :save_last_frame: class PgroupExample(Scene): def construct(self): p1 = PointCloudDot(radius=1, density=20, color=BLUE) p1.move_to(4.5 * LEFT) p2 = PointCloudDot() p3 = PointCloudDot(radius=1.5, stroke_width=2.5, color=PINK) p3.move_to(4.5 * RIGHT) pList = PGroup(p1, p2, p3) self.add(pList) """ def __init__(self, *pmobs: Any, **kwargs: Any) -> None: if not all(isinstance(m, (PMobject, OpenGLPMobject)) for m in pmobs): raise ValueError( "All submobjects must be of type PMobject or OpenGLPMObject" " if using the opengl renderer", ) super().__init__(**kwargs) self.add(*pmobs) def fade_to( self, color: ParsableManimColor, alpha: float, family: bool = True ) -> Self: if family: for mob in self.submobjects: mob.fade_to(color, alpha, family) return self
[docs] class PointCloudDot(Mobject1D): """A disc made of a cloud of dots. Examples -------- .. manim:: PointCloudDotExample :save_last_frame: class PointCloudDotExample(Scene): def construct(self): cloud_1 = PointCloudDot(color=RED) cloud_2 = PointCloudDot(stroke_width=4, radius=1) cloud_3 = PointCloudDot(density=15) group = Group(cloud_1, cloud_2, cloud_3).arrange() self.add(group) .. manim:: PointCloudDotExample2 class PointCloudDotExample2(Scene): def construct(self): plane = ComplexPlane() cloud = PointCloudDot(color=RED) self.add( plane, cloud ) self.wait() self.play( cloud.animate.apply_complex_function(lambda z: np.exp(z)) ) """ def __init__( self, center: Point3DLike = ORIGIN, radius: float = 2.0, stroke_width: int = 2, density: int = DEFAULT_POINT_DENSITY_1D, color: ManimColor = PURE_YELLOW, **kwargs: Any, ) -> None: self.radius = radius self.epsilon = 1.0 / density super().__init__( stroke_width=stroke_width, density=density, color=color, **kwargs ) self.shift(center) def init_points(self) -> None: self.reset_points() self.generate_points()
[docs] def generate_points(self) -> None: self.add_points( np.array( [ r * (np.cos(theta) * RIGHT + np.sin(theta) * UP) for r in np.arange(self.epsilon, self.radius, self.epsilon) # Num is equal to int(stop - start)/ (step + 1) reformulated. for theta in np.linspace( 0, 2 * np.pi, num=int(2 * np.pi * (r + self.epsilon) / self.epsilon), ) ] ), )
[docs] class Point(PMobject): """A mobject representing a point. Examples -------- .. manim:: ExamplePoint :save_last_frame: class ExamplePoint(Scene): def construct(self): colorList = [RED, GREEN, BLUE, YELLOW] for i in range(200): point = Point(location=[0.63 * np.random.randint(-4, 4), 0.37 * np.random.randint(-4, 4), 0], color=np.random.choice(colorList)) self.add(point) for i in range(200): point = Point(location=[0.37 * np.random.randint(-4, 4), 0.63 * np.random.randint(-4, 4), 0], color=np.random.choice(colorList)) self.add(point) self.add(point) """ def __init__( self, location: Point3DLike = ORIGIN, color: ManimColor = BLACK, **kwargs: Any ) -> None: self.location = location super().__init__(color=color, **kwargs) def init_points(self) -> None: self.reset_points() self.generate_points() self.set_points([self.location])
[docs] def generate_points(self) -> None: self.add_points(np.array([self.location]))