diff --git a/.gitignore b/.gitignore index abec5da495..b294075283 100644 --- a/.gitignore +++ b/.gitignore @@ -132,5 +132,9 @@ dist/ /media_dir.txt # ^TODO: Remove the need for this with a proper config file +#uv lock file +uv.lock +*.lock + # Ignore the built dependencies third_party/* diff --git a/manim/__init__.py b/manim/__init__.py index 9fd8e65f1d..66fc11a30c 100644 --- a/manim/__init__.py +++ b/manim/__init__.py @@ -75,6 +75,7 @@ from .mobject.text.tex_mobject import * from .mobject.text.text_mobject import * from .mobject.text.typst_mobject import * +from .mobject.three_d.light_source import * from .mobject.three_d.polyhedra import * from .mobject.three_d.three_d_utils import * from .mobject.three_d.three_dimensions import * diff --git a/manim/_config/utils.py b/manim/_config/utils.py index 8d1de7a9d4..297cb2143d 100644 --- a/manim/_config/utils.py +++ b/manim/_config/utils.py @@ -857,7 +857,10 @@ def digest_args(self, args: argparse.Namespace) -> Self: if args.tex_template: self.tex_template = TexTemplate.from_file(args.tex_template) - if self.renderer == RendererType.OPENGL and args.write_to_movie is None: + if ( + self.renderer in (RendererType.OPENGL, RendererType.WEBGPU) + and args.write_to_movie is None + ): # --write_to_movie was not passed on the command line, so don't generate video. self["write_to_movie"] = False diff --git a/manim/cli/render/commands.py b/manim/cli/render/commands.py index 5c36b41a94..1f81fb7052 100644 --- a/manim/cli/render/commands.py +++ b/manim/cli/render/commands.py @@ -120,6 +120,18 @@ def render(**kwargs: Any) -> ClickArgs | dict[str, Any]: except Exception: error_console.print_exception() sys.exit(1) + elif config.renderer == RendererType.WEBGPU: + from manim.renderer.webgpu.webgpu_renderer import WebGPURenderer + + renderer = WebGPURenderer() + for SceneClass in scene_classes_from_file(file): + try: + with tempconfig({}): + scene = SceneClass(renderer) + scene.render() + except Exception: + error_console.print_exception() + sys.exit(1) else: for SceneClass in scene_classes_from_file(file): try: diff --git a/manim/constants.py b/manim/constants.py index ccf99a0293..9774e98acf 100644 --- a/manim/constants.py +++ b/manim/constants.py @@ -273,6 +273,7 @@ class RendererType(Enum): CAIRO = "cairo" #: A renderer based on the cairo backend. OPENGL = "opengl" #: An OpenGL-based renderer. + WEBGPU = "webgpu" #: A WebGPU-based renderer (wgpu-py). class LineJointType(Enum): diff --git a/manim/mobject/opengl/opengl_compatibility.py b/manim/mobject/opengl/opengl_compatibility.py index 761cd32918..c607c4cb64 100644 --- a/manim/mobject/opengl/opengl_compatibility.py +++ b/manim/mobject/opengl/opengl_compatibility.py @@ -4,10 +4,6 @@ from typing import Any from manim import config -from manim.mobject.opengl.opengl_mobject import OpenGLMobject -from manim.mobject.opengl.opengl_point_cloud_mobject import OpenGLPMobject -from manim.mobject.opengl.opengl_three_dimensions import OpenGLSurface -from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject from ...constants import RendererType @@ -26,6 +22,11 @@ def __new__( mcls, name: str, bases: tuple[type, ...], namespace: dict[str, Any] ) -> type: if config.renderer == RendererType.OPENGL: + from manim.mobject.opengl.opengl_mobject import OpenGLMobject + from manim.mobject.opengl.opengl_point_cloud_mobject import OpenGLPMobject + from manim.mobject.opengl.opengl_three_dimensions import OpenGLSurface + from manim.mobject.opengl.opengl_vectorized_mobject import OpenGLVMobject + # Must check class names to prevent # cyclic importing. base_names_to_opengl: dict[str, type] = { diff --git a/manim/mobject/svg/svg_mobject.py b/manim/mobject/svg/svg_mobject.py index 75cfd621dd..4c0ab1bda8 100644 --- a/manim/mobject/svg/svg_mobject.py +++ b/manim/mobject/svg/svg_mobject.py @@ -382,6 +382,9 @@ def apply_style_to_mobject(mob: VMobject, shape: se.GraphicObject) -> VMobject: fill_color=shape.fill.hexrgb, fill_opacity=shape.fill.opacity, ) + # Extract fill-rule for WebGPU renderer (0=nonzero, 1=evenodd). + fill_rule_str = shape.values.get("fill-rule", "nonzero") + mob.fill_rule = 1 if fill_rule_str == "evenodd" else 0 return mob def path_to_mobject(self, path: se.Path) -> VMobjectFromSVGPath: diff --git a/manim/mobject/three_d/__init__.py b/manim/mobject/three_d/__init__.py index 98d295a24e..c3236a3a0e 100644 --- a/manim/mobject/three_d/__init__.py +++ b/manim/mobject/three_d/__init__.py @@ -6,6 +6,7 @@ .. autosummary:: :toctree: ../reference + ~light_source ~polyhedra ~three_d_utils ~three_dimensions diff --git a/manim/mobject/three_d/dot_cloud.py b/manim/mobject/three_d/dot_cloud.py new file mode 100644 index 0000000000..093b2afcbe --- /dev/null +++ b/manim/mobject/three_d/dot_cloud.py @@ -0,0 +1,193 @@ +"""Point-cloud mobject for WebGPU TrueDot rendering. + +``PointDot`` is a single dot rendered as a 3-D lit sphere. +``DotCloud3D`` is an N-point cloud rendered as N lit spheres. + +These classes are ``Mobject``-based (not ``OpenGLMobject``-based) so they +work transparently with both Cairo scenes (skipped silently) and WebGPU +scenes (routed to the TrueDot pipeline). +""" + +from __future__ import annotations + +__all__ = ["DotCloud3D", "PointDot"] + +from typing import Any + +import numpy as np + +from manim.constants import ORIGIN +from manim.mobject.mobject import Mobject +from manim.typing import Point3DLike +from manim.utils.color import WHITE, ParsableManimColor, color_to_rgba + + +class DotCloud3D(Mobject): + """A cloud of points, each rendered as a lit sphere by the WebGPU renderer. + + In Cairo / OpenGL renderers, ``DotCloud3D`` objects are silently ignored + (they produce no geometry for those pipelines). + + Parameters + ---------- + points + Array of world-space positions, shape (N, 3). + color + Base colour of all dots (can be overridden per-point via ``set_rgbas``). + radius + World-space radius of each sphere in scene units. + gloss + Specular shininess (Cairo-style): 0 = matte, 1 = very shiny. + shadow + Diffuse darkening strength: 0 = no shadow, 1 = full Lambert shading. + """ + + def __init__( + self, + points: np.ndarray | list | None = None, + color: ParsableManimColor = WHITE, + radius: float = 0.05, + gloss: float = 0.3, + shadow: float = 0.3, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + pts = ( + np.zeros((0, 3), dtype=np.float32) + if points is None + else np.asarray(points, dtype=np.float32) + ) + if pts.ndim == 1: + pts = pts.reshape(1, 3) + self._cloud_points: np.ndarray = pts.astype(np.float32) + self._rgbas: np.ndarray = np.tile( + np.asarray(color_to_rgba(color), dtype=np.float32), (max(len(pts), 1), 1) + ) + self.dot_radius: float = float(radius) + self.gloss: float = float(gloss) + self.shadow: float = float(shadow) + # Set Mobject.points to the cloud positions so bounding-box helpers work. + if len(pts) > 0: + self.set_points(pts) + + # ------------------------------------------------------------------ + # Cloud-specific API + # ------------------------------------------------------------------ + + def get_cloud_points(self) -> np.ndarray: + """Return the (N, 3) float32 array of dot centres.""" + return self._cloud_points + + def set_cloud_points(self, points: np.ndarray) -> DotCloud3D: + pts = np.asarray(points, dtype=np.float32) + if pts.ndim == 1: + pts = pts.reshape(1, 3) + self._cloud_points = pts + if len(pts) > 0: + self.set_points(pts) + return self + + def get_rgbas(self) -> np.ndarray: + """Return the (N, 4) float32 RGBA array for all dots.""" + return self._rgbas + + def set_rgbas(self, rgbas: np.ndarray) -> DotCloud3D: + self._rgbas = np.asarray(rgbas, dtype=np.float32) + return self + + def set_color(self, color: ParsableManimColor, family: bool = True) -> DotCloud3D: # type: ignore[override] + rgba = np.asarray(color_to_rgba(color), dtype=np.float32) + self._rgbas = np.tile(rgba, (max(len(self._cloud_points), 1), 1)) + if family: + for sub in self.submobjects: + if isinstance(sub, DotCloud3D): + sub.set_color(color, family=False) + return self + + def set_opacity(self, opacity: float, family: bool = True) -> DotCloud3D: # type: ignore[override] + self._rgbas[:, 3] = float(opacity) + if family: + for sub in self.submobjects: + if isinstance(sub, DotCloud3D): + sub.set_opacity(opacity, family=False) + return self + + # ------------------------------------------------------------------ + # Animation support — required Mobject overrides + # ------------------------------------------------------------------ + + def align_points_with_larger(self, larger_mobject: Mobject) -> None: + """Tile _cloud_points and _rgbas to match the size of *larger_mobject*.""" + if not isinstance(larger_mobject, DotCloud3D): + return + n_target = len(larger_mobject._cloud_points) + n_self = len(self._cloud_points) + if n_self == 0 or n_self >= n_target: + return + reps = -(-n_target // n_self) # ceiling division + self._cloud_points = np.tile(self._cloud_points, (reps, 1))[:n_target] + self._rgbas = np.tile(self._rgbas, (reps, 1))[:n_target] + self.set_points(self._cloud_points) + + def interpolate_color( + self, mobject1: Mobject, mobject2: Mobject, alpha: float + ) -> None: + """Linearly interpolate _rgbas between *mobject1* and *mobject2*.""" + if not isinstance(mobject1, DotCloud3D) or not isinstance(mobject2, DotCloud3D): + return + self._rgbas = ((1 - alpha) * mobject1._rgbas + alpha * mobject2._rgbas).astype( + np.float32 + ) + + def interpolate( + self, + mobject1: Mobject, + mobject2: Mobject, + alpha: float, + path_func: Any = None, + ) -> DotCloud3D: + """Interpolate position and colour; keep _cloud_points in sync with points.""" + from manim.utils.bezier import interpolate as lerp + + if path_func is None: + path_func = lerp + super().interpolate(mobject1, mobject2, alpha, path_func) + # Mobject.interpolate writes into self.points; mirror that into _cloud_points. + self._cloud_points = np.asarray(self.points, dtype=np.float32) + return self + + +class PointDot(DotCloud3D): + """A single dot at *center* rendered as a lit sphere by the WebGPU renderer. + + Parameters + ---------- + center + World-space position of the dot. + color + Base colour. + radius + World-space radius of the sphere in scene units. + gloss + Specular shininess: 0 = matte, 1 = very shiny. + shadow + Diffuse darkening: 0 = flat, 1 = full Lambert shading. + """ + + def __init__( + self, + center: Point3DLike = ORIGIN, + color: ParsableManimColor = WHITE, + radius: float = 0.05, + gloss: float = 0.3, + shadow: float = 0.3, + **kwargs: Any, + ) -> None: + super().__init__( + points=np.asarray(center, dtype=np.float32).reshape(1, 3), + color=color, + radius=radius, + gloss=gloss, + shadow=shadow, + **kwargs, + ) diff --git a/manim/mobject/three_d/light_source.py b/manim/mobject/three_d/light_source.py new file mode 100644 index 0000000000..8e67f15361 --- /dev/null +++ b/manim/mobject/three_d/light_source.py @@ -0,0 +1,301 @@ +"""Light source mobjects for WebGPU 3-D rendering. + +.. warning:: + + **WebGPU renderer only.** All classes in this module are silently ignored + by the Cairo and OpenGL renderers. They have no visual effect outside of + scenes rendered with ``--renderer=webgpu``. + +Classes +------- +LightSource + Abstract base for all light types. Extends :class:`~.Mobject` so it + participates in scene management (``add``, ``remove``, ``play``). + +AmbientLight + Uniform omnidirectional light that brightens every surface equally. + Only **one** ambient light may exist per scene; ``ThreeDScene`` adds a + default one automatically. + +DirectionalLight + Parallel light from a fixed direction (like sunlight). Intensity is + constant regardless of position. + +PointLight + Omnidirectional light that radiates from a point in world space. Falls + off with the inverse-square of distance. + +SpotLight + Cone-shaped light from a point in a direction. Same attenuation as + ``PointLight`` but only illuminates within *cone_angle* of the direction. + Soft penumbra can be controlled via the *penumbra* parameter. +""" + +from __future__ import annotations + +__all__ = ["AmbientLight", "DirectionalLight", "LightSource", "PointLight", "SpotLight"] + +from typing import Any + +import numpy as np + +from manim.mobject.mobject import Mobject +from manim.typing import Point3DLike, Vector3D +from manim.utils.color import WHITE, ParsableManimColor, color_to_rgb + +# ── Light kind constants (must match WGSL shader) ───────────────────────────── +_KIND_AMBIENT = 0 +_KIND_DIRECTIONAL = 1 +_KIND_POINT = 2 +_KIND_SPOT = 3 + + +class LightSource(Mobject): + """Base class for all WebGPU light sources. + + Parameters + ---------- + color + Light colour. + intensity + Brightness scalar. Typical range is [0, 1] for ambient/directional; + higher values (e.g. 300) are suitable for point/spot lights with + distance attenuation. + **kwargs + Forwarded to :class:`~.Mobject`. + + .. note:: + + **WebGPU renderer only** — ignored by Cairo and OpenGL renderers. + """ + + # Subclasses must set this before calling super().__init__. + _kind: int = -1 + + def __init__( + self, + color: ParsableManimColor = WHITE, + intensity: float = 1.0, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + self.light_color: np.ndarray = np.asarray(color_to_rgb(color), dtype=np.float32) + self.intensity: float = float(intensity) + self.should_render: bool = False + + # ------------------------------------------------------------------ + # Packing helpers (used by the renderer) + # ------------------------------------------------------------------ + + def pack(self) -> bytes: + """Return the 64-byte binary representation of this light. + + Matches the WGSL ``Light`` struct layout: + + .. code-block:: text + + offset 0 position vec3 12 B + offset 12 kind u32 4 B + offset 16 direction vec3 12 B + offset 28 intensity f32 4 B + offset 32 color vec3 12 B + offset 44 cone_angle f32 4 B + offset 48 penumbra f32 4 B + offset 52 _pad0-2 f32×3 12 B (alignment padding) + """ + buf = np.zeros(16, dtype=np.float32) # 16 × 4 B = 64 B + buf[0:3] = self._get_position() + buf[3] = np.float32(self._kind).view(np.float32) + buf[4:7] = self._get_direction() + buf[7] = self.intensity + buf[8:11] = self.light_color + buf[11] = self._get_cone_angle() + buf[12] = self._get_penumbra() + # buf[13], buf[14], buf[15] remain zero (padding) + + # Reinterpret index 3 as u32 so we get exact integer bit pattern. + raw = buf.tobytes() + kind_bytes = np.uint32(self._kind).tobytes() + return raw[:12] + kind_bytes + raw[16:] + + # Subclass hooks — override as needed. + def _get_position(self) -> np.ndarray: + return np.zeros(3, dtype=np.float32) + + def _get_direction(self) -> np.ndarray: + return np.zeros(3, dtype=np.float32) + + def _get_cone_angle(self) -> float: + return 0.0 + + def _get_penumbra(self) -> float: + return 0.0 + + +class AmbientLight(LightSource): + """Uniform ambient light — illuminates every surface equally from all sides. + + Only **one** ambient light is allowed per scene. ``ThreeDScene`` adds one + by default (white, intensity 0.5). Replacing it or adjusting its intensity + gives global brightness control. + + Parameters + ---------- + color + Light colour. Default: white. + intensity + Ambient brightness. Default: ``0.5``. + + .. note:: + + **WebGPU renderer only** — ignored by Cairo and OpenGL renderers. + """ + + _kind = _KIND_AMBIENT + + def __init__( + self, + color: ParsableManimColor = WHITE, + intensity: float = 0.5, + **kwargs: Any, + ) -> None: + super().__init__(color=color, intensity=intensity, **kwargs) + + +class DirectionalLight(LightSource): + """Parallel directional light (like sunlight) — constant intensity everywhere. + + Parameters + ---------- + direction + World-space vector the light travels *toward* (points from light toward + the scene). Does not need to be normalised. Default: ``[0, -1, -1]`` + (down-forward). + color + Light colour. Default: white. + intensity + Brightness scalar. Default: ``0.8``. + + .. note:: + + **WebGPU renderer only** — ignored by Cairo and OpenGL renderers. + """ + + _kind = _KIND_DIRECTIONAL + + def __init__( + self, + direction: Vector3D = np.array([0.0, -1.0, -1.0]), + color: ParsableManimColor = WHITE, + intensity: float = 1.0, + **kwargs: Any, + ) -> None: + super().__init__(color=color, intensity=intensity, **kwargs) + d = np.asarray(direction, dtype=np.float32) + norm = np.linalg.norm(d) + self._direction: np.ndarray = ( + (d / norm) if norm > 1e-8 else np.array([0.0, 0.0, -1.0], dtype=np.float32) + ) + + def _get_direction(self) -> np.ndarray: + return self._direction + + +class PointLight(LightSource): + """Omnidirectional point light — radiates from a fixed world-space position. + + Intensity falls off with the inverse-square of the distance to the surface + (``attenuation = intensity / dot(light_dir, light_dir)``). + + Parameters + ---------- + position + World-space centre of the light. Default: ``[10, 10, -10]``. + color + Light colour. Default: white. + intensity + Source brightness (before distance attenuation). Default: ``300``. + + .. note:: + + **WebGPU renderer only** — ignored by Cairo and OpenGL renderers. + """ + + _kind = _KIND_POINT + + def __init__( + self, + position: Point3DLike = np.array([10.0, 10.0, -10.0]), + color: ParsableManimColor = WHITE, + intensity: float = 300.0, + **kwargs: Any, + ) -> None: + super().__init__(color=color, intensity=intensity, **kwargs) + self._position: np.ndarray = np.asarray(position, dtype=np.float32) + + def _get_position(self) -> np.ndarray: + return self._position + + +class SpotLight(LightSource): + """Cone-shaped point light. + + Like ``PointLight`` but only illuminates within *cone_angle* degrees of the + *direction* vector. A soft penumbra region of width *penumbra* degrees + linearly fades the outer rim. + + Parameters + ---------- + position + World-space origin of the spot. Default: ``[10, 10, -10]``. + direction + World-space vector the cone points toward. Does not need to be + normalised. Default: ``[0, -1, -1]``. + cone_angle + Half-angle of the inner (full-brightness) cone, in **degrees**. + Default: ``30``. + penumbra + Width of the soft penumbra region in **degrees**. Default: ``5``. + color + Light colour. Default: white. + intensity + Source brightness (before distance attenuation). Default: ``300``. + + .. note:: + + **WebGPU renderer only** — ignored by Cairo and OpenGL renderers. + """ + + _kind = _KIND_SPOT + + def __init__( + self, + position: Point3DLike = np.array([10.0, 10.0, -10.0]), + direction: Vector3D = np.array([0.0, -1.0, -1.0]), + cone_angle: float = 30.0, + penumbra: float = 5.0, + color: ParsableManimColor = WHITE, + intensity: float = 300.0, + **kwargs: Any, + ) -> None: + super().__init__(color=color, intensity=intensity, **kwargs) + self._position: np.ndarray = np.asarray(position, dtype=np.float32) + d = np.asarray(direction, dtype=np.float32) + norm = np.linalg.norm(d) + self._direction: np.ndarray = ( + (d / norm) if norm > 1e-8 else np.array([0.0, 0.0, -1.0], dtype=np.float32) + ) + self._cone_angle: float = float(cone_angle) + self._penumbra: float = float(penumbra) + + def _get_position(self) -> np.ndarray: + return self._position + + def _get_direction(self) -> np.ndarray: + return self._direction + + def _get_cone_angle(self) -> float: + return self._cone_angle + + def _get_penumbra(self) -> float: + return self._penumbra diff --git a/manim/mobject/three_d/three_dimensions.py b/manim/mobject/three_d/three_dimensions.py index 4508cc0da6..6479a5e4d8 100644 --- a/manim/mobject/three_d/three_dimensions.py +++ b/manim/mobject/three_d/three_dimensions.py @@ -88,6 +88,27 @@ class Surface(VGroup, metaclass=ConvertToOpenGL): should_make_jagged Changes the anchor mode of the Bézier curves from smooth to jagged. Defaults to ``False``. + diffuse_strength + Strength of the diffuse (Lambertian) lighting component, in [0, 1]. + Defaults to 0.8. + specular_strength + Strength of the specular (Phong) highlight, in [0, ∞]. + Defaults to 0.9. + specular_exponent + Phong shininess exponent — higher values produce a tighter, sharper + specular highlight; lower values produce a broad, soft one. + Defaults to 16.0. + + .. warning:: + + The ``diffuse_strength``, ``specular_strength``, and + ``specular_exponent`` parameters — and all material setter methods + (:meth:`set_diffuse_strength`, :meth:`set_specular_strength`, + :meth:`set_specular_exponent`, :meth:`set_material`, + :meth:`set_diffuse_by_func`, :meth:`set_specular_by_func`, + :meth:`set_specular_exponent_by_func`, :meth:`set_material_by_func`) + — are **WebGPU renderer only**. They are silently ignored by the + Cairo and OpenGL renderers. Examples -------- @@ -127,10 +148,21 @@ def __init__( stroke_width: float = 0.5, should_make_jagged: bool = False, pre_function_handle_to_anchor_scale_factor: float = 0.00001, + diffuse_strength: float = 0.8, + specular_strength: float = 0.9, + specular_exponent: float = 16.0, **kwargs: Any, ) -> None: + # If `color` is explicitly passed, use it as fill_color and disable + # checkerboard so the explicit color isn't silently overridden. + if "color" in kwargs: + fill_color = kwargs.pop("color") + checkerboard_colors = False self.u_range = u_range self.v_range = v_range + self.diffuse_strength = float(diffuse_strength) + self.specular_strength = float(specular_strength) + self.specular_exponent = float(specular_exponent) super().__init__( fill_color=fill_color, fill_opacity=fill_opacity, @@ -159,6 +191,139 @@ def __init__( def func(self, u: float, v: float) -> np.ndarray: return self._func(u, v) + # ------------------------------------------------------------------ + # Material setters (WebGPU renderer only) + # ------------------------------------------------------------------ + + def set_diffuse_strength(self, value: float) -> Surface: + """Set the Lambertian diffuse strength in [0, 1]. + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + self.diffuse_strength = float(value) + return self + + def set_specular_strength(self, value: float) -> Surface: + """Set the Phong specular highlight strength. + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + self.specular_strength = float(value) + return self + + def set_specular_exponent(self, value: float) -> Surface: + """Set the Phong shininess exponent. + + Higher values give a tighter highlight; lower values give a broad, + soft one. Typical range: 4 (very soft) to 128 (mirror-like). + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + self.specular_exponent = float(value) + return self + + def set_material( + self, + diffuse_strength: float | None = None, + specular_strength: float | None = None, + specular_exponent: float | None = None, + ) -> Surface: + """Set material parameters uniformly across the whole surface. + + Any parameter left as ``None`` is unchanged. Per-patch overrides + set via :meth:`set_diffuse_by_func` etc. take precedence at render + time. + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + if diffuse_strength is not None: + self.diffuse_strength = float(diffuse_strength) + if specular_strength is not None: + self.specular_strength = float(specular_strength) + if specular_exponent is not None: + self.specular_exponent = float(specular_exponent) + return self + + # ------------------------------------------------------------------ + # Per-patch material — function-based assignment + # ------------------------------------------------------------------ + + def set_diffuse_by_func(self, func: Callable[[float, float], float]) -> Surface: + """Assign a per-patch diffuse strength using a ``(u, v)`` function. + + *func* is called with the centre ``(u, v)`` coordinates of each + patch and must return a float in [0, 1]. + + Example — gradient from matte at the bottom to reflective at top:: + + surface.set_diffuse_by_func(lambda u, v: v / v_max) + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + for face in self.list_of_faces: + face.diffuse_strength = float(func(face.u_center, face.v_center)) + return self + + def set_specular_by_func(self, func: Callable[[float, float], float]) -> Surface: + """Assign a per-patch specular strength using a ``(u, v)`` function. + + *func* is called with the centre ``(u, v)`` coordinates of each + patch and must return a float. + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + for face in self.list_of_faces: + face.specular_strength = float(func(face.u_center, face.v_center)) + return self + + def set_specular_exponent_by_func( + self, func: Callable[[float, float], float] + ) -> Surface: + """Assign a per-patch specular exponent (shininess) using a ``(u, v)`` + function. + + *func* is called with the centre ``(u, v)`` coordinates of each + patch and must return a float (typical range: 4 to 128). + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + for face in self.list_of_faces: + face.specular_exponent = float(func(face.u_center, face.v_center)) + return self + + def set_material_by_func(self, func: Callable[[float, float], dict]) -> Surface: + """Assign per-patch material parameters using a ``(u, v)`` function. + + *func* is called with the centre ``(u, v)`` of each patch and must + return a :class:`dict` with any subset of the keys + ``"diffuse_strength"``, ``"specular_strength"``, + ``"specular_exponent"``. Missing keys leave the corresponding + attribute unchanged on that patch. + + Example — shinier at the equator, matte at the poles:: + + def mat(u, v): + t = abs(np.sin(u)) # 0 at poles, 1 at equator + return { + "specular_exponent": 8 + 120 * t, + "specular_strength": 0.2 + 0.8 * t, + } + + + sphere.set_material_by_func(mat) + + .. warning:: **WebGPU renderer only** — ignored by Cairo and OpenGL. + """ + for face in self.list_of_faces: + params = func(face.u_center, face.v_center) + if "diffuse_strength" in params: + face.diffuse_strength = float(params["diffuse_strength"]) + if "specular_strength" in params: + face.specular_strength = float(params["specular_strength"]) + if "specular_exponent" in params: + face.specular_exponent = float(params["specular_exponent"]) + return self + def _get_u_values_and_v_values(self) -> tuple[np.ndarray, np.ndarray]: if isinstance(self.resolution, int): u_res = v_res = self.resolution @@ -195,6 +360,8 @@ def _setup_in_uv_space(self) -> None: face.u2 = u2 face.v1 = v1 face.v2 = v2 + face.u_center = float(u1 + u2) * 0.5 + face.v_center = float(v1 + v2) * 0.5 self.list_of_faces.append(face) faces.set_fill(color=self.fill_color, opacity=self.fill_opacity) faces.set_stroke( @@ -343,7 +510,10 @@ def param_surface(u, v): if config.renderer == RendererType.OPENGL: assert isinstance(mob, OpenGLMobject) mob.set_color(mob_color, recurse=False) - elif config.renderer == RendererType.CAIRO: + elif config.renderer in { + RendererType.CAIRO, + RendererType.WEBGPU, + }: mob.set_color(mob_color, family=False) break @@ -452,7 +622,7 @@ def __init__( ) -> None: if config.renderer == RendererType.OPENGL: res_value = (101, 51) - elif config.renderer == RendererType.CAIRO: + elif config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: res_value = (24, 12) else: raise Exception("Unknown renderer") @@ -886,7 +1056,7 @@ def add_bases(self) -> Self: assert isinstance(self, OpenGLMobject) color = self.color opacity = self.opacity - elif config.renderer == RendererType.CAIRO: + elif config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: color = self.fill_color opacity = self.fill_opacity @@ -1323,7 +1493,7 @@ def __init__( ) -> None: if config.renderer == RendererType.OPENGL: res_value = (101, 101) - elif config.renderer == RendererType.CAIRO: + elif config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: res_value = (24, 24) resolution = resolution if resolution is not None else res_value diff --git a/manim/renderer/base_renderer.py b/manim/renderer/base_renderer.py new file mode 100644 index 0000000000..7db1072e0f --- /dev/null +++ b/manim/renderer/base_renderer.py @@ -0,0 +1,126 @@ +"""Structural interface shared by all Manim renderers. + +Using ``typing.Protocol`` (rather than an ABC) means the existing +``CairoRenderer`` and ``OpenGLRenderer`` satisfy the interface automatically +without any inheritance changes. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable + +import numpy as np + +if TYPE_CHECKING: + from PIL import Image + + from manim.mobject.mobject import Mobject + from manim.mobject.value_tracker import ValueTracker + from manim.scene.scene import Scene + + +# --------------------------------------------------------------------------- +# Camera protocols +# --------------------------------------------------------------------------- + + +@runtime_checkable +class ThreeDCameraProtocol(Protocol): + """Interface that every 3D camera must satisfy. + + ``ThreeDScene`` calls these attributes and methods on + ``renderer.camera``; declaring them here makes the contract explicit + for all three renderers (Cairo ``ThreeDCamera``, ``OpenGLCamera``, + and ``WebGPUCamera``). + """ + + # ── angle trackers (Cairo uses ValueTrackers; OpenGL/WebGPU store + # the angles directly and expose them as plain attributes) ───────────── + theta_tracker: ValueTracker + phi_tracker: ValueTracker + gamma_tracker: ValueTracker + focal_distance_tracker: ValueTracker + zoom_tracker: ValueTracker + + # Mobject used as the camera frame-centre (moved to pan the scene). + _frame_center: Mobject + + # ── orientation setters ────────────────────────────────────────────────── + def set_phi(self, phi: float) -> None: ... + def set_theta(self, theta: float) -> None: ... + def set_gamma(self, gamma: float) -> None: ... + def set_zoom(self, zoom: float) -> None: ... + def set_focal_distance(self, focal_distance: float) -> None: ... + + # ── incremental rotation (OpenGL / WebGPU ambient rotation) ───────────── + def increment_theta(self, dtheta: float) -> None: ... + def increment_phi(self, dphi: float) -> None: ... + def increment_gamma(self, dgamma: float) -> None: ... + + # ── updater support (camera is a Mobject in OpenGL / WebGPU) ──────────── + def add_updater(self, func: Any, **kwargs: Any) -> None: ... + def clear_updaters(self) -> None: ... + + # ── moving-mobject tracking ────────────────────────────────────────────── + def get_value_trackers(self) -> list[ValueTracker]: ... + + # ── fixed-orientation / fixed-in-frame helpers (Cairo) ────────────────── + def add_fixed_orientation_mobjects( + self, *mobjects: Mobject, **kwargs: Any + ) -> None: ... + def remove_fixed_orientation_mobjects(self, *mobjects: Mobject) -> None: ... + def add_fixed_in_frame_mobjects(self, *mobjects: Mobject) -> None: ... + def remove_fixed_in_frame_mobjects(self, *mobjects: Mobject) -> None: ... + + +# --------------------------------------------------------------------------- +# Renderer protocol +# --------------------------------------------------------------------------- + + +@runtime_checkable +class RendererProtocol(Protocol): + """Protocol that every Manim renderer must satisfy. + + ``scene.py`` accesses these attributes and methods on the renderer object. + Declaring them here makes the contract explicit and enables static type + checking without breaking existing renderer classes. + """ + + # ── core attributes ────────────────────────────────────────────────────── + camera: Any # ThreeDCameraProtocol for 3D renderers; Any for 2D + skip_animations: bool + num_plays: int + time: float + file_writer: Any + window: Any # WebGPUWindow | pyglet window | None + animation_start_time: float + static_image: Any + + # Camera configuration dict (pixel_width, pixel_height, …). + # Accessed by SpecialThreeDScene to choose low/high quality config. + camera_config: dict + + # Set of currently-held key codes; polled by Scene.interact() and + # the WebGPU window event handler. + pressed_keys: set + + # ── lifecycle ──────────────────────────────────────────────────────────── + def init_scene(self, scene: Scene) -> None: ... + + def play(self, scene: Scene, *args: Any, **kwargs: Any) -> None: ... + + def render( + self, scene: Scene, frame_offset: float, moving_mobjects: list + ) -> None: ... + + def update_frame(self, scene: Scene) -> None: ... + + def scene_finished(self, scene: Scene) -> None: ... + + def clear_screen(self) -> None: ... + + # ── frame access ───────────────────────────────────────────────────────── + def get_image(self) -> Image.Image: ... + + def get_frame(self) -> np.ndarray: ... diff --git a/manim/renderer/webgpu/__init__.py b/manim/renderer/webgpu/__init__.py new file mode 100644 index 0000000000..e431e91996 --- /dev/null +++ b/manim/renderer/webgpu/__init__.py @@ -0,0 +1,5 @@ +"""WebGPU rendering backend for Manim (wgpu-py).""" + +from .webgpu_renderer import WebGPURenderer + +__all__ = ["WebGPURenderer"] diff --git a/manim/renderer/webgpu/shaders/cubic_to_quads.wgsl b/manim/renderer/webgpu/shaders/cubic_to_quads.wgsl new file mode 100644 index 0000000000..465a7d3f8a --- /dev/null +++ b/manim/renderer/webgpu/shaders/cubic_to_quads.wgsl @@ -0,0 +1,88 @@ +// GPU compute shader: cubic Bezier → quadratic approximations. +// +// Each thread converts one cubic Bezier (4 × 3-D control points) into four +// quadratic Beziers (3 points each) using two levels of de Casteljau +// subdivision at t = 0.5 followed by midpoint degree-reduction. +// +// This runs in the same command encoder as the render pass (before it), +// so WebGPU's implicit pass ordering gives a barrier — the render shader +// safely reads the output quads without an explicit synchronisation step. +// +// Buffer layout +// ------------- +// binding 0 in_cubics read-only-storage array +// 12 floats per cubic: b0.xyz, b1.xyz, b2.xyz, b3.xyz (tightly packed) +// +// binding 1 out_quads storage (read_write) array +// 36 floats per input cubic (4 quads × 9 floats): +// quad k: p0.xyz, pmid.xyz, p1.xyz (start, control, end) +// Order: [sub-cubic 0, sub-cubic 1, sub-cubic 2, sub-cubic 3] +// +// binding 2 params uniform +// offset 0: n_cubics u32 +// (padded to 16 bytes) +// +// Dispatch: ceil(n_cubics / 64) workgroups × 1 × 1, workgroup_size = 64. + +struct Params { n_cubics : u32, _pad0: u32, _pad1: u32, _pad2: u32 }; + +@group(0) @binding(0) var in_cubics : array; +@group(0) @binding(1) var out_quads : array; +@group(0) @binding(2) var params : Params; + +// Write one quadratic (p0, pmid, p2) as 9 consecutive floats at base. +fn write_quad(base: u32, p0: vec3, pmid: vec3, p2: vec3) { + out_quads[base ] = p0.x; out_quads[base + 1u] = p0.y; out_quads[base + 2u] = p0.z; + out_quads[base + 3u] = pmid.x; out_quads[base + 4u] = pmid.y; out_quads[base + 5u] = pmid.z; + out_quads[base + 6u] = p2.x; out_quads[base + 7u] = p2.y; out_quads[base + 8u] = p2.z; +} + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if idx >= params.n_cubics { return; } + + // Read 4 control points of this cubic. + let bi = idx * 12u; + let b0 = vec3(in_cubics[bi ], in_cubics[bi + 1u], in_cubics[bi + 2u]); + let b1 = vec3(in_cubics[bi + 3u], in_cubics[bi + 4u], in_cubics[bi + 5u]); + let b2 = vec3(in_cubics[bi + 6u], in_cubics[bi + 7u], in_cubics[bi + 8u]); + let b3 = vec3(in_cubics[bi + 9u], in_cubics[bi + 10u], in_cubics[bi + 11u]); + + // ── Level 1: split [b0, b1, b2, b3] at t = 0.5 ────────────────────── + // Left half: [b0, m01, m012, m0123] + // Right half: [m0123, m123, m23, b3] + let m01 = (b0 + b1) * 0.5; + let m12 = (b1 + b2) * 0.5; + let m23 = (b2 + b3) * 0.5; + let m012 = (m01 + m12) * 0.5; + let m123 = (m12 + m23) * 0.5; + let m0123 = (m012 + m123) * 0.5; + + // ── Level 2a: split left half [b0, m01, m012, m0123] at t = 0.5 ────── + let lm01 = (b0 + m01) * 0.5; + let lm12 = (m01 + m012) * 0.5; + let lm23 = (m012 + m0123) * 0.5; + let lm012 = (lm01 + lm12) * 0.5; + let lm123 = (lm12 + lm23) * 0.5; + let lm0123 = (lm012 + lm123) * 0.5; + // Quad 0: [b0, lm01, lm012, lm0123] + // Quad 1: [lm0123, lm123, lm23, m0123] + + // ── Level 2b: split right half [m0123, m123, m23, b3] at t = 0.5 ───── + let rm01 = (m0123 + m123) * 0.5; + let rm12 = (m123 + m23) * 0.5; + let rm23 = (m23 + b3) * 0.5; + let rm012 = (rm01 + rm12) * 0.5; + let rm123 = (rm12 + rm23) * 0.5; + let rm0123 = (rm012 + rm123) * 0.5; + // Quad 2: [m0123, rm01, rm012, rm0123] + // Quad 3: [rm0123, rm123, rm23, b3] + + // Write 4 quadratics. Degree reduction: mid-handle = (h0 + h1) * 0.5. + let bo = idx * 36u; + write_quad(bo , b0, (lm01 + lm012) * 0.5, lm0123); + write_quad(bo + 9u, lm0123, (lm123 + lm23 ) * 0.5, m0123 ); + write_quad(bo + 18u, m0123, (rm01 + rm012) * 0.5, rm0123); + write_quad(bo + 27u, rm0123, (rm123 + rm23 ) * 0.5, b3 ); +} diff --git a/manim/renderer/webgpu/shaders/image.wgsl b/manim/renderer/webgpu/shaders/image.wgsl new file mode 100644 index 0000000000..e6c859041d --- /dev/null +++ b/manim/renderer/webgpu/shaders/image.wgsl @@ -0,0 +1,68 @@ +// Image quad shader. +// +// Renders a textured quad from four world-space corner vertices. +// Used by WebGPURenderer to draw ImageMobject instances. +// +// Uniform layout (group 0, binding 0) — same 656-byte block as the surface +// shaders; only projection and view are used here: +// offset 0 — projection mat4x4 64 B +// offset 64 — view mat4x4 64 B +// (remaining bytes are lighting fields, unused by this shader) +// +// Texture / sampler (group 1): +// binding 0 — texture_2d (rgba8unorm uploaded as f32 [0,1] per channel) +// binding 1 — sampler (linear, clamp-to-edge) +// +// UV origin is top-left (matches WebGPU texture layout and Manim pixel_array +// row-major order: row 0 = top). No y-flip is needed. +// +// Tint uniform (group 2, binding 0) — 16-byte block: +// offset 0 — rgb vec3 12 B per-channel colour multiplier +// offset 12 — _pad f32 4 B alignment padding (unused) +// +// Default white (1, 1, 1) is identity — texture is returned unchanged. +// Set via mob.color; populated by the renderer from mob.color.to_rgb(). +// +// Vertex attributes (stride 20 bytes): +// location 0 — in_pos float32x3 offset 0 +// location 1 — in_uv float32x2 offset 12 + +struct Uniforms { + projection : mat4x4, + view : mat4x4, +}; +@group(0) @binding(0) var u : Uniforms; + +@group(1) @binding(0) var img_texture : texture_2d; +@group(1) @binding(1) var img_sampler : sampler; + +struct TintUniforms { + rgb : vec3, + _pad : f32, +}; +@group(2) @binding(0) var tint : TintUniforms; + +struct VertexInput { + @location(0) in_pos : vec3, + @location(1) in_uv : vec2, +}; + +struct VertexOutput { + @builtin(position) position : vec4, + @location(0) uv : vec2, +}; + +@vertex +fn vs_main(in: VertexInput) -> VertexOutput { + var out: VertexOutput; + out.position = u.projection * u.view * vec4(in.in_pos, 1.0); + out.uv = in.in_uv; + return out; +} + +@fragment +fn fs_main(in: VertexOutput) -> @location(0) vec4 { + let sample = textureSample(img_texture, img_sampler, in.uv); + // Multiply RGB by the tint; alpha is taken from the texture as-is. + return vec4(sample.rgb * tint.rgb, sample.a); +} diff --git a/manim/renderer/webgpu/shaders/oit_compose.wgsl b/manim/renderer/webgpu/shaders/oit_compose.wgsl new file mode 100644 index 0000000000..409651126b --- /dev/null +++ b/manim/renderer/webgpu/shaders/oit_compose.wgsl @@ -0,0 +1,53 @@ +// WebGPU OIT composition shader. +// +// Reads the two OIT accumulation textures produced by surface_oit.wgsl and +// composites the transparent geometry result onto the existing opaque framebuffer. +// +// No vertex buffer is needed — three vertices are generated from the built-in +// vertex index, covering the full screen with a single oversized triangle. +// +// The output is alpha-blended onto the main render texture using standard +// {src-alpha, one-minus-src-alpha} blending, so the pipeline must be +// configured with that blend mode. + +@group(0) @binding(0) var oit_accum : texture_2d; // rgba16float +@group(0) @binding(1) var oit_reveal : texture_2d; // rgba16float + +struct VertexOutput { + @builtin(position) clip_position : vec4, +}; + +// Full-screen triangle: 3 vertices cover [-1,1]x[-1,1] without a VBO. +@vertex +fn vs_main(@builtin(vertex_index) vi: u32) -> VertexOutput { + var pos = array, 3>( + vec2(-1.0, -1.0), + vec2( 3.0, -1.0), + vec2(-1.0, 3.0), + ); + var out: VertexOutput; + out.clip_position = vec4(pos[vi], 0.0, 1.0); + return out; +} + +@fragment +fn fs_main(in: VertexOutput) -> @location(0) vec4 { + let coord = vec2(in.clip_position.xy); + let accum = textureLoad(oit_accum, coord, 0); + let reveal = textureLoad(oit_reveal, coord, 0).r; + + // Nothing accumulated at this pixel — don't touch the framebuffer. + if (accum.a < 1e-5) { discard; } + + // Weighted average colour. + let avg_color = accum.rgb / accum.a; + + // Overall opacity: 1 − ∏(1 − αᵢ). + // `reveal` started at 1 and each fragment multiplied it by (1 − α), + // so reveal == ∏(1 − αᵢ) and 1 − reveal is the accumulated opacity. + let alpha = clamp(1.0 - reveal, 0.0, 1.0); + + // Output with premultiplied alpha so the {src-alpha, one-minus-src-alpha} + // pipeline blend correctly composites over the opaque framebuffer. + return vec4(avg_color, alpha); +} diff --git a/manim/renderer/webgpu/shaders/readback_compact.wgsl b/manim/renderer/webgpu/shaders/readback_compact.wgsl new file mode 100644 index 0000000000..a9682c3020 --- /dev/null +++ b/manim/renderer/webgpu/shaders/readback_compact.wgsl @@ -0,0 +1,48 @@ +// GPU compact-readback compute shader. +// +// Reads every pixel from the bgra8unorm render texture and writes tightly-packed +// RGBA bytes into a storage buffer — one u32 per pixel, little-endian: +// byte 0 = R, byte 1 = G, byte 2 = B, byte 3 = A +// +// Two CPU operations are eliminated in a single pass: +// +// 1. Row-padding strip +// copy_texture_to_buffer requires bytes_per_row to be a multiple of 256. +// The CPU previously looped over every row to remove the padding bytes. +// Here each thread writes directly to the tight index y*width + x, +// so the output buffer is already compact — no post-processing needed. +// +// 2. B↔R channel swap +// The render texture is bgra8unorm (GPU memory layout: B G R A). +// textureLoad() always returns components as (r, g, b, a) regardless of +// the physical layout, so the output is already in RGBA byte order. +// The CPU numpy channel-swap is no longer required. +// +// Workgroup size 16×16 = 256 threads. Each thread handles one pixel. +// Caller dispatches ceil(width/16) × ceil(height/16) workgroups; the +// out-of-bounds guard below is a no-op for tiles that fit exactly. + +@group(0) @binding(0) var src_tex : texture_2d; +@group(0) @binding(1) var dst : array; + +@compute @workgroup_size(16, 16) +fn main(@builtin(global_invocation_id) gid : vec3) { + let dims = textureDimensions(src_tex); + + // Discard threads outside the image boundary (last tile edge). + if (gid.x >= dims.x || gid.y >= dims.y) { return; } + + // textureLoad returns (r, g, b, a) as normalized f32 regardless of bgra + // memory layout — no manual component swap needed. + let c = textureLoad(src_tex, vec2(i32(gid.x), i32(gid.y)), 0); + + let r = u32(clamp(c.r * 255.0 + 0.5, 0.0, 255.0)); + let g = u32(clamp(c.g * 255.0 + 0.5, 0.0, 255.0)); + let b = u32(clamp(c.b * 255.0 + 0.5, 0.0, 255.0)); + let a = u32(clamp(c.a * 255.0 + 0.5, 0.0, 255.0)); + + // Pack as little-endian u32: byte0=R, byte1=G, byte2=B, byte3=A. + // NumPy / PIL both read this as RGBA when the buffer is reinterpreted as + // uint8 in row-major order. + dst[gid.y * dims.x + gid.x] = r | (g << 8u) | (b << 16u) | (a << 24u); +} diff --git a/manim/renderer/webgpu/shaders/surface_combined.wgsl b/manim/renderer/webgpu/shaders/surface_combined.wgsl new file mode 100644 index 0000000000..d90776e26a --- /dev/null +++ b/manim/renderer/webgpu/shaders/surface_combined.wgsl @@ -0,0 +1,223 @@ +// Combined opaque surface fill + barycentric wireframe shader. +// +// One draw call per surface face renders both Phong-lit fill and mesh-grid +// lines without a separate stroke pass. +// +// Technique +// --------- +// Each triangle (centroid, anchor_i, anchor_{i+1}) in the centroid fan +// carries barycentric coordinates: +// centroid → bary = (1, 0, 0) bary.x = 0 on the outer edge +// anchor_i → bary = (0, 1, 0) +// anchor_{i+1} → bary = (0, 0, 1) +// +// bary.x is 0 on the "outer" edge (anchor_i ↔ anchor_{i+1}), which is the +// visible mesh-grid edge. The inner spoke edges (centroid ↔ anchor_*) have +// bary.y = 0 or bary.z = 0 but are NOT rendered as wireframe. +// +// fwidth(bary.x) gives the screen-space derivative of bary.x, so +// edge_dist_px = bary.x / fwidth(bary.x) +// is approximately the distance from the outer edge in screen pixels. +// A smooth SDF step at stroke_half_px produces anti-aliased grid lines. +// +// Compositing: wireframe stroke "over" Phong fill (Porter-Duff). +// +// Vertex layout (must match _SURFACE_COMBINED_DTYPE, stride 84 bytes): +// location 0 — in_vert float32x3 offset 0 +// location 1 — in_normal float32x3 offset 12 +// location 2 — in_fill_color float32x4 offset 24 +// location 3 — in_stroke_color float32x4 offset 40 +// location 4 — in_bary float32x3 offset 56 +// location 5 — stroke_half_px float32 offset 68 +// location 6 — diffuse_strength float32 offset 72 +// location 7 — specular_strength float32 offset 76 +// location 8 — specular_exponent float32 offset 80 + +// ── Lighting uniform layout (656 bytes total) ───────────────────────────── +// +// offset 0 — projection mat4x4 64 B +// offset 64 — view mat4x4 64 B +// offset 128 — num_lights u32 4 B +// offset 132 — _pad u32 × 3 12 B (align array to 16 B) +// offset 144 — lights Light × 8 512 B +// +// Light struct (64 bytes): +// offset 0 position vec3 12 B — point / spot world position +// offset 12 kind u32 4 B — 0=ambient,1=directional,2=point,3=spot +// offset 16 direction vec3 12 B — directional / spot direction +// offset 28 intensity f32 4 B +// offset 32 color vec3 12 B +// offset 44 cone_angle f32 4 B — spot inner half-angle (degrees) +// offset 48 penumbra f32 4 B — spot penumbra width (degrees) +// offset 52 _pad0-2 f32 × 3 12 B + +const MAX_LIGHTS : u32 = 8u; + +struct Light { + position : vec3, + kind : u32, + direction : vec3, + intensity : f32, + color : vec3, + cone_angle : f32, + penumbra : f32, + _pad0 : f32, + _pad1 : f32, + _pad2 : f32, +}; + +struct Uniforms { + projection : mat4x4, + view : mat4x4, + num_lights : u32, + _pad0 : u32, + _pad1 : u32, + _pad2 : u32, + lights : array, +}; +@group(0) @binding(0) var u : Uniforms; + +struct VertexInput { + @location(0) in_vert : vec3, + @location(1) in_normal : vec3, + @location(2) in_fill_color : vec4, + @location(3) in_stroke_color : vec4, + @location(4) in_bary : vec3, + @location(5) stroke_half_px : f32, + @location(6) diffuse_strength : f32, + @location(7) specular_strength : f32, + @location(8) specular_exponent : f32, +}; + +struct VertexOutput { + @builtin(position) clip_position : vec4, + @location(0) v_fill_color : vec4, + @location(1) v_stroke_color : vec4, + @location(2) v_view_normal : vec3, + @location(3) v_view_pos : vec3, + @location(4) v_bary : vec3, + @location(5) @interpolate(flat) v_stroke_half : f32, + @location(6) @interpolate(flat) v_diffuse : f32, + @location(7) @interpolate(flat) v_specular : f32, + @location(8) @interpolate(flat) v_spec_exp : f32, +}; + +@vertex +fn vs_main(in: VertexInput) -> VertexOutput { + var out: VertexOutput; + let view_pos = u.view * vec4(in.in_vert, 1.0); + out.clip_position = u.projection * view_pos; + out.v_view_pos = view_pos.xyz; + let view3 = mat3x3(u.view[0].xyz, u.view[1].xyz, u.view[2].xyz); + out.v_view_normal = view3 * in.in_normal; + out.v_fill_color = in.in_fill_color; + out.v_stroke_color = in.in_stroke_color; + out.v_bary = in.in_bary; + out.v_stroke_half = in.stroke_half_px; + out.v_diffuse = in.diffuse_strength; + out.v_specular = in.specular_strength; + out.v_spec_exp = in.specular_exponent; + return out; +} + +// ── Lighting helpers ────────────────────────────────────────────────────────── + +fn compute_lighting( + view_pos : vec3, + view_normal : vec3, + base_rgb : vec3, + diff_str : f32, + spec_str : f32, + spec_exp : f32, +) -> vec3 { + let specular_exp = spec_exp; + let view_dir = normalize(-view_pos); + let view3 = mat3x3(u.view[0].xyz, u.view[1].xyz, u.view[2].xyz); + + var acc_rgb = vec3(0.0); + + for (var i = 0u; i < u.num_lights; i++) { + let L = u.lights[i]; + + switch L.kind { + // ── Ambient ─────────────────────────────────────────────────── + case 0u: { + acc_rgb += base_rgb * L.color * L.intensity; + } + // ── Directional ─────────────────────────────────────────────── + case 1u: { + // direction is the direction the light *travels toward* (world space). + // Transform to view space and negate to get the "to-light" direction. + let light_dir = normalize(-(view3 * L.direction)); + let half_vec = normalize(light_dir + view_dir); + let diff = clamp(dot(view_normal, light_dir), 0.0, 1.0); + let spec = pow(max(dot(view_normal, half_vec), 0.0), specular_exp); + acc_rgb += base_rgb * L.color * (diff_str * diff * L.intensity); + acc_rgb += L.color * (spec_str * spec * L.intensity); + } + // ── Point ───────────────────────────────────────────────────── + case 2u: { + let light_view_pos = (u.view * vec4(L.position, 1.0)).xyz; + let light_dir_v = light_view_pos - view_pos; + let light_dir = normalize(light_dir_v); + let half_vec = normalize(light_dir + view_dir); + let diff = clamp(dot(view_normal, light_dir), 0.0, 1.0); + let spec = pow(max(dot(view_normal, half_vec), 0.0), specular_exp); + let attenuation = L.intensity / dot(light_dir_v, light_dir_v); + acc_rgb += base_rgb * L.color * (diff_str * diff * attenuation); + acc_rgb += L.color * (spec_str * spec * attenuation); + } + // ── Spot ────────────────────────────────────────────────────── + case 3u: { + let light_view_pos = (u.view * vec4(L.position, 1.0)).xyz; + let light_dir_v = light_view_pos - view_pos; + let light_dir = normalize(light_dir_v); + let half_vec = normalize(light_dir + view_dir); + let diff = clamp(dot(view_normal, light_dir), 0.0, 1.0); + let spec = pow(max(dot(view_normal, half_vec), 0.0), specular_exp); + let attenuation = L.intensity / dot(light_dir_v, light_dir_v); + + // Cone falloff: compare angle between -light_dir and spot direction. + let spot_dir = normalize(view3 * L.direction); + let cos_theta = dot(-light_dir, spot_dir); + let cos_inner = cos(radians(L.cone_angle)); + let cos_outer = cos(radians(L.cone_angle + L.penumbra)); + let spot_factor = clamp((cos_theta - cos_outer) / (cos_inner - cos_outer + 1e-6), 0.0, 1.0); + + acc_rgb += base_rgb * L.color * (diff_str * diff * attenuation * spot_factor); + acc_rgb += L.color * (spec_str * spec * attenuation * spot_factor); + } + default: {} + } + } + + return clamp(acc_rgb, vec3(0.0), vec3(1.0)); +} + +@fragment +fn fs_main(in: VertexOutput, @builtin(front_facing) front_facing: bool) -> @location(0) vec4 { + let diffuse_strength = in.v_diffuse; + let specular_strength = in.v_specular; + + let raw_normal = select(-in.v_view_normal, in.v_view_normal, front_facing); + let norm = normalize(raw_normal); + + let lit_rgb = compute_lighting( + in.v_view_pos, norm, + in.v_fill_color.rgb, + diffuse_strength, specular_strength, in.v_spec_exp, + ); + let fill_a = in.v_fill_color.a; + + // ── Barycentric wireframe ────────────────────────────────────────────── + let edge_dist_px = in.v_bary.x / max(fwidth(in.v_bary.x), 1e-6); + let stroke_cov = clamp(in.v_stroke_half + 0.5 - edge_dist_px, 0.0, 1.0); + let stroke_a = in.v_stroke_color.a * stroke_cov; + + // ── Porter-Duff "over": stroke on top of fill ───────────────────────── + let total_a = stroke_a + fill_a * (1.0 - stroke_a); + if total_a <= 0.001 { discard; } + + let out_rgb = (stroke_a * in.v_stroke_color.rgb + fill_a * (1.0 - stroke_a) * lit_rgb) / total_a; + return vec4(out_rgb, total_a); +} diff --git a/manim/renderer/webgpu/shaders/surface_oit.wgsl b/manim/renderer/webgpu/shaders/surface_oit.wgsl new file mode 100644 index 0000000000..b2c8470725 --- /dev/null +++ b/manim/renderer/webgpu/shaders/surface_oit.wgsl @@ -0,0 +1,220 @@ +// WebGPU OIT accumulation shader — Weighted Blended Order-Independent Transparency. +// +// McGuire & Bavoil 2013. Renders transparent surface fragments into two +// accumulation targets instead of the main framebuffer: +// +// location 0 accum rgba16float weighted colour + alpha sum +// location 1 reveal rgba16float per-channel transmittance product +// +// The main pipeline blend modes for these targets are set to: +// accum: {src: one, dst: one} — additive accumulation +// reveal: {src: zero, dst: one-minus-src-alpha} — transmittance multiplication +// +// A subsequent full-screen composition pass reads both textures and composites +// the result onto the opaque framebuffer. +// +// Vertex layout matches surface_combined.wgsl (stride 84 bytes): +// location 0 — in_vert float32x3 offset 0 +// location 1 — in_normal float32x3 offset 12 +// location 2 — in_fill_color float32x4 offset 24 +// location 3 — in_stroke_color float32x4 offset 40 +// location 4 — in_bary float32x3 offset 56 +// location 5 — stroke_half_px float32 offset 68 +// location 6 — diffuse_strength float32 offset 72 +// location 7 — specular_strength float32 offset 76 +// location 8 — specular_exponent float32 offset 80 + +// ── Lighting uniform layout (656 bytes total) ───────────────────────────── +// +// offset 0 — projection mat4x4 64 B +// offset 64 — view mat4x4 64 B +// offset 128 — num_lights u32 4 B +// offset 132 — _pad u32 × 3 12 B +// offset 144 — lights Light × 8 512 B +// +// Light struct (64 bytes): +// offset 0 position vec3 12 B +// offset 12 kind u32 4 B — 0=ambient,1=directional,2=point,3=spot +// offset 16 direction vec3 12 B +// offset 28 intensity f32 4 B +// offset 32 color vec3 12 B +// offset 44 cone_angle f32 4 B +// offset 48 penumbra f32 4 B +// offset 52 _pad0-2 f32 × 3 12 B + +const MAX_LIGHTS : u32 = 8u; + +struct Light { + position : vec3, + kind : u32, + direction : vec3, + intensity : f32, + color : vec3, + cone_angle : f32, + penumbra : f32, + _pad0 : f32, + _pad1 : f32, + _pad2 : f32, +}; + +struct Uniforms { + projection : mat4x4, + view : mat4x4, + num_lights : u32, + _pad0 : u32, + _pad1 : u32, + _pad2 : u32, + lights : array, +}; +@group(0) @binding(0) var u : Uniforms; + +struct VertexInput { + @location(0) in_vert : vec3, + @location(1) in_normal : vec3, + @location(2) in_fill_color : vec4, + @location(3) in_stroke_color : vec4, + @location(4) in_bary : vec3, + @location(5) stroke_half_px : f32, + @location(6) diffuse_strength : f32, + @location(7) specular_strength : f32, + @location(8) specular_exponent : f32, +}; + +struct VertexOutput { + @builtin(position) clip_position : vec4, + @location(0) v_fill_color : vec4, + @location(1) v_stroke_color : vec4, + @location(2) v_view_normal : vec3, + @location(3) v_view_pos : vec3, + @location(4) v_bary : vec3, + @location(5) @interpolate(flat) v_stroke_half : f32, + @location(6) @interpolate(flat) v_diffuse : f32, + @location(7) @interpolate(flat) v_specular : f32, + @location(8) @interpolate(flat) v_spec_exp : f32, +}; + +@vertex +fn vs_main(in: VertexInput) -> VertexOutput { + var out: VertexOutput; + let view_pos = u.view * vec4(in.in_vert, 1.0); + out.clip_position = u.projection * view_pos; + out.v_view_pos = view_pos.xyz; + let view3 = mat3x3(u.view[0].xyz, u.view[1].xyz, u.view[2].xyz); + out.v_view_normal = view3 * in.in_normal; + out.v_fill_color = in.in_fill_color; + out.v_stroke_color = in.in_stroke_color; + out.v_bary = in.in_bary; + out.v_stroke_half = in.stroke_half_px; + out.v_diffuse = in.diffuse_strength; + out.v_specular = in.specular_strength; + out.v_spec_exp = in.specular_exponent; + return out; +} + +// ── Lighting helpers ────────────────────────────────────────────────────────── + +fn compute_lighting( + view_pos : vec3, + view_normal : vec3, + base_rgb : vec3, + diff_str : f32, + spec_str : f32, + spec_exp : f32, +) -> vec3 { + let specular_exp = spec_exp; + let view_dir = normalize(-view_pos); + let view3 = mat3x3(u.view[0].xyz, u.view[1].xyz, u.view[2].xyz); + + var acc_rgb = vec3(0.0); + + for (var i = 0u; i < u.num_lights; i++) { + let L = u.lights[i]; + + switch L.kind { + case 0u: { + acc_rgb += base_rgb * L.color * L.intensity; + } + case 1u: { + let light_dir = normalize(-(view3 * L.direction)); + let half_vec = normalize(light_dir + view_dir); + let diff = clamp(dot(view_normal, light_dir), 0.0, 1.0); + let spec = pow(max(dot(view_normal, half_vec), 0.0), specular_exp); + acc_rgb += base_rgb * L.color * (diff_str * diff * L.intensity); + acc_rgb += L.color * (spec_str * spec * L.intensity); + } + case 2u: { + let light_view_pos = (u.view * vec4(L.position, 1.0)).xyz; + let light_dir_v = light_view_pos - view_pos; + let light_dir = normalize(light_dir_v); + let half_vec = normalize(light_dir + view_dir); + let diff = clamp(dot(view_normal, light_dir), 0.0, 1.0); + let spec = pow(max(dot(view_normal, half_vec), 0.0), specular_exp); + let attenuation = L.intensity / dot(light_dir_v, light_dir_v); + acc_rgb += base_rgb * L.color * (diff_str * diff * attenuation); + acc_rgb += L.color * (spec_str * spec * attenuation); + } + case 3u: { + let light_view_pos = (u.view * vec4(L.position, 1.0)).xyz; + let light_dir_v = light_view_pos - view_pos; + let light_dir = normalize(light_dir_v); + let half_vec = normalize(light_dir + view_dir); + let diff = clamp(dot(view_normal, light_dir), 0.0, 1.0); + let spec = pow(max(dot(view_normal, half_vec), 0.0), specular_exp); + let attenuation = L.intensity / dot(light_dir_v, light_dir_v); + let spot_dir = normalize(view3 * L.direction); + let cos_theta = dot(-light_dir, spot_dir); + let cos_inner = cos(radians(L.cone_angle)); + let cos_outer = cos(radians(L.cone_angle + L.penumbra)); + let spot_factor = clamp((cos_theta - cos_outer) / (cos_inner - cos_outer + 1e-6), 0.0, 1.0); + acc_rgb += base_rgb * L.color * (diff_str * diff * attenuation * spot_factor); + acc_rgb += L.color * (spec_str * spec * attenuation * spot_factor); + } + default: {} + } + } + + return clamp(acc_rgb, vec3(0.0), vec3(1.0)); +} + +struct FragOutput { + @location(0) accum : vec4, // weighted colour sum → rgba16float + @location(1) reveal : vec4, // transmittance product → rgba16float +}; + +@fragment +fn fs_main(in: VertexOutput, @builtin(front_facing) front_facing: bool) -> FragOutput { + let diffuse_strength = in.v_diffuse; + let specular_strength = in.v_specular; + + let raw_normal = select(-in.v_view_normal, in.v_view_normal, front_facing); + let norm = normalize(raw_normal); + + let lit_rgb = compute_lighting( + in.v_view_pos, norm, + in.v_fill_color.rgb, + diffuse_strength, specular_strength, in.v_spec_exp, + ); + let fill_a = in.v_fill_color.a; + + // ── Barycentric wireframe ────────────────────────────────────────────── + let edge_dist_px = in.v_bary.x / max(fwidth(in.v_bary.x), 1e-6); + let stroke_cov = clamp(in.v_stroke_half + 0.5 - edge_dist_px, 0.0, 1.0); + let stroke_a = in.v_stroke_color.a * stroke_cov; + + // ── Porter-Duff "over": stroke on top of fill ───────────────────────── + let total_a = stroke_a + fill_a * (1.0 - stroke_a); + if total_a <= 0.001 { discard; } + let out_rgb = (stroke_a * in.v_stroke_color.rgb + fill_a * (1.0 - stroke_a) * lit_rgb) / total_a; + + // ── Weighted Blended OIT ─────────────────────────────────────────────── + let z = in.v_view_pos.z; + let w = clamp( + pow(total_a, 3.0) / (1e-5 + pow(abs(z) / 5.0, 4.0)), + 1e-2, 3e3 + ); + + var out: FragOutput; + out.accum = vec4(out_rgb * total_a * w, total_a * w); + out.reveal = vec4(total_a, total_a, total_a, total_a); + return out; +} diff --git a/manim/renderer/webgpu/shaders/true_dot.wgsl b/manim/renderer/webgpu/shaders/true_dot.wgsl new file mode 100644 index 0000000000..edf609f6bc --- /dev/null +++ b/manim/renderer/webgpu/shaders/true_dot.wgsl @@ -0,0 +1,164 @@ +// WebGPU TrueDot shader — screen-aligned sphere dot rendering. +// +// Each dot is expanded CPU-side into 2 triangles (6 vertices) forming a +// screen-aligned quad. UV coords span (-1,-1) → (1,1) across the quad. +// The fragment shader treats the quad as a sphere projected onto the screen: +// • Pixels outside the unit disc are discarded (anti-aliased edge). +// • The sphere normal is reconstructed from the UV position. +// • Multi-light Phong shading is applied (same light array as surface shaders). +// gloss / shadow parameters blend the result toward the Cairo-style look. +// +// The same camera Uniforms struct and bind group layout as the surface +// shaders are reused (binding 0, group 0). +// +// Vertex layout (stride 48 bytes) — must match _TRUE_DOT_DTYPE: +// location 0 — center float32x3 offset 0 (12 B) +// location 1 — color float32x4 offset 12 (16 B) +// location 2 — uv float32x2 offset 28 ( 8 B) +// location 3 — radius float32 offset 36 ( 4 B) +// location 4 — gloss float32 offset 40 ( 4 B) +// location 5 — shadow float32 offset 44 ( 4 B) + +// ── Shared uniform (same layout as surface shaders, 656 bytes) ──────────────── + +const MAX_LIGHTS : u32 = 8u; + +struct Light { + position : vec3, + kind : u32, + direction : vec3, + intensity : f32, + color : vec3, + cone_angle : f32, + penumbra : f32, + _pad0 : f32, + _pad1 : f32, + _pad2 : f32, +}; + +struct Uniforms { + projection : mat4x4, + view : mat4x4, + num_lights : u32, + _pad0 : u32, + _pad1 : u32, + _pad2 : u32, + lights : array, +}; +@group(0) @binding(0) var u : Uniforms; + +struct VertexInput { + @location(0) center : vec3, + @location(1) color : vec4, + @location(2) uv : vec2, + @location(3) radius : f32, + @location(4) gloss : f32, + @location(5) shadow : f32, +}; + +struct VertexOutput { + @builtin(position) clip_position : vec4, + @location(0) v_color : vec4, + @location(1) v_uv : vec2, + @location(2) @interpolate(flat) v_gloss : f32, + @location(3) @interpolate(flat) v_shadow : f32, + @location(4) v_center_view : vec3, +}; + +@vertex +fn vs_main(in: VertexInput) -> VertexOutput { + // Project dot center into view space. + let cv = u.view * vec4(in.center, 1.0); + + // Expand quad in view space: move corner by radius × UV along x/y. + let expanded = cv + vec4(in.uv.x * in.radius, in.uv.y * in.radius, 0.0, 0.0); + + var out: VertexOutput; + out.clip_position = u.projection * expanded; + out.v_color = in.color; + out.v_uv = in.uv; + out.v_gloss = in.gloss; + out.v_shadow = in.shadow; + out.v_center_view = cv.xyz; + return out; +} + +@fragment +fn fs_main(in: VertexOutput) -> @location(0) vec4 { + let d = length(in.v_uv); + + // Anti-aliased disc edge. + let fw = fwidth(d); + let alpha_mult = 1.0 - smoothstep(1.0 - fw, 1.0 + fw, d); + if alpha_mult <= 0.001 { discard; } + + // Reconstruct sphere surface normal in view space. + let z2 = max(0.0, 1.0 - d * d); + let sphere_normal = normalize(vec3(in.v_uv.x, in.v_uv.y, sqrt(z2))); + let view_dir = normalize(-in.v_center_view); + let view3 = mat3x3(u.view[0].xyz, u.view[1].xyz, u.view[2].xyz); + + // ── Multi-light accumulation ────────────────────────────────────────── + var acc_rgb = vec3(0.0); + + for (var i = 0u; i < u.num_lights; i++) { + let L = u.lights[i]; + + switch L.kind { + // Ambient + case 0u: { + acc_rgb += in.v_color.rgb * L.color * L.intensity; + } + // Directional + case 1u: { + let to_light = normalize(-(view3 * L.direction)); + let dot_ln = clamp(dot(sphere_normal, to_light), 0.0, 1.0); + // Cairo-style shadow: darken by Lambertian term + let darkening = mix(1.0, dot_ln, in.v_shadow); + // Cairo-style specular gloss + let reflect_l = reflect(-to_light, sphere_normal); + let dot_rv = clamp(dot(reflect_l, view_dir), 0.0, 1.0); + let shine = in.v_gloss * exp(-3.0 * pow(1.0 - dot_rv, 2.0)); + let lit = darkening * mix(in.v_color.rgb, vec3(1.0), shine); + acc_rgb += lit * L.color * L.intensity; + } + // Point + case 2u: { + let light_vpos = (u.view * vec4(L.position, 1.0)).xyz; + let light_dir_v = light_vpos - in.v_center_view; + let to_light = normalize(light_dir_v); + let attenuation = L.intensity / dot(light_dir_v, light_dir_v); + let dot_ln = clamp(dot(sphere_normal, to_light), 0.0, 1.0); + let darkening = mix(1.0, dot_ln, in.v_shadow); + let reflect_l = reflect(-to_light, sphere_normal); + let dot_rv = clamp(dot(reflect_l, view_dir), 0.0, 1.0); + let shine = in.v_gloss * exp(-3.0 * pow(1.0 - dot_rv, 2.0)); + let lit = darkening * mix(in.v_color.rgb, vec3(1.0), shine); + acc_rgb += lit * L.color * attenuation; + } + // Spot + case 3u: { + let light_vpos = (u.view * vec4(L.position, 1.0)).xyz; + let light_dir_v = light_vpos - in.v_center_view; + let to_light = normalize(light_dir_v); + let attenuation = L.intensity / dot(light_dir_v, light_dir_v); + let spot_dir = normalize(view3 * L.direction); + let cos_theta = dot(-to_light, spot_dir); + let cos_inner = cos(radians(L.cone_angle)); + let cos_outer = cos(radians(L.cone_angle + L.penumbra)); + let spot_factor = clamp((cos_theta - cos_outer) / (cos_inner - cos_outer + 1e-6), 0.0, 1.0); + let dot_ln = clamp(dot(sphere_normal, to_light), 0.0, 1.0); + let darkening = mix(1.0, dot_ln, in.v_shadow); + let reflect_l = reflect(-to_light, sphere_normal); + let dot_rv = clamp(dot(reflect_l, view_dir), 0.0, 1.0); + let shine = in.v_gloss * exp(-3.0 * pow(1.0 - dot_rv, 2.0)); + let lit = darkening * mix(in.v_color.rgb, vec3(1.0), shine); + acc_rgb += lit * L.color * attenuation * spot_factor; + } + default: {} + } + } + + let out_rgb = clamp(acc_rgb, vec3(0.0), vec3(1.0)); + return vec4(out_rgb, in.v_color.a * alpha_mult); +} diff --git a/manim/renderer/webgpu/shaders/vmobject_fill_stroke.wgsl b/manim/renderer/webgpu/shaders/vmobject_fill_stroke.wgsl new file mode 100644 index 0000000000..10bd9f5e7b --- /dev/null +++ b/manim/renderer/webgpu/shaders/vmobject_fill_stroke.wgsl @@ -0,0 +1,309 @@ +// Combined VMobject fill + stroke shader. +// +// One bounding quad per object; one fragment loop simultaneously accumulates: +// 1. Fill coverage — Slug winding-number algorithm in NDC space +// (Lengyel 2017, patent-dedicated public domain, code MIT) +// 2. Stroke coverage — SDF minimum distance to each curve in pixel space +// +// Result is Porter-Duff "over" compositing: stroke painted on top of fill. +// +// The quadratic Bezier control points are written by cubic_to_quads.wgsl +// into a single shared storage buffer. Each object's fill and stroke +// curves occupy separate contiguous regions of that buffer referenced by +// (fill_curve_start, n_fill_curves) and (stroke_curve_start, n_stroke_curves). +// +// Objects with no fill: pass fill_color.a = 0 or n_fill_curves = 0. +// Objects with no stroke: pass stroke_half_ndc = 0 or n_stroke_curves = 0. +// +// Uniform layout (group 0, binding 0) — 656-byte block shared with surface shaders +// (only the first two fields are used here): +// offset 0 — projection mat4x4 (64 B) +// offset 64 — view mat4x4 (64 B) +// offset 128 — ... (lighting data, unused by this shader) +// +// Storage buffer (group 0, binding 1) — array, 9 floats per quadratic: +// [p0.x p0.y p0.z pmid.x pmid.y pmid.z p2.x p2.y p2.z] +// +// Vertex attributes (must match _FILL_STROKE_DTYPE, stride 68 bytes): +// location 0 — in_pos float32x3 offset 0 +// location 1 — in_fill_color float32x4 offset 12 +// location 2 — in_stroke_color float32x4 offset 28 +// location 3 — stroke_half_ndc float32 offset 44 +// location 4 — fill_curve_start uint32 offset 48 +// location 5 — n_fill_curves uint32 offset 52 +// location 6 — stroke_curve_start uint32 offset 56 +// location 7 — n_stroke_curves uint32 offset 60 +// location 8 — fill_rule uint32 offset 64 (0=nonzero, 1=evenodd) + +struct Uniforms { + projection : mat4x4, + view : mat4x4, +}; +@group(0) @binding(0) var u : Uniforms; +@group(0) @binding(1) var quads : array; + +struct VertexInput { + @location(0) in_pos : vec3, + @location(1) in_fill_color : vec4, + @location(2) in_stroke_color : vec4, + @location(3) stroke_half_ndc : f32, + @location(4) fill_curve_start : u32, + @location(5) n_fill_curves : u32, + @location(6) stroke_curve_start : u32, + @location(7) n_stroke_curves : u32, + @location(8) fill_rule : u32, +}; + +struct VertexOutput { + @builtin(position) clip_pos : vec4, + @location(0) ndc_xy : vec2, + @location(1) v_fill_color : vec4, + @location(2) v_stroke_color : vec4, + @location(3) @interpolate(flat) v_stroke_half_ndc : f32, + @location(4) @interpolate(flat) fill_curve_start : u32, + @location(5) @interpolate(flat) n_fill_curves : u32, + @location(6) @interpolate(flat) stroke_curve_start: u32, + @location(7) @interpolate(flat) n_stroke_curves : u32, + @location(8) @interpolate(flat) fill_rule : u32, +}; + +@vertex +fn vs_main(in: VertexInput) -> VertexOutput { + var out: VertexOutput; + let view_pos = u.view * vec4(in.in_pos, 1.0); + let clip = u.projection * view_pos; + out.clip_pos = clip; + // Perspective divide: ortho gives w=1 (no change); perspective maps + // the vertex to the correct 2-D screen-proportional position. + out.ndc_xy = clip.xy / clip.w; + out.v_fill_color = in.in_fill_color; + out.v_stroke_color = in.in_stroke_color; + out.v_stroke_half_ndc = in.stroke_half_ndc; + out.fill_curve_start = in.fill_curve_start; + out.n_fill_curves = in.n_fill_curves; + out.stroke_curve_start = in.stroke_curve_start; + out.n_stroke_curves = in.n_stroke_curves; + out.fill_rule = in.fill_rule; + return out; +} + +// --------------------------------------------------------------------------- +// Slug helpers — winding-number fill (NDC space) +// Adapted from Lengyel 2017 (HLSL → WGSL). +// --------------------------------------------------------------------------- + +fn calc_root_code(y1: f32, y2: f32, y3: f32) -> u32 { + let i1 = (bitcast(y1) >> 31u) & 1u; + let i2 = (bitcast(y2) >> 30u) & 2u; + let i3 = (bitcast(y3) >> 29u) & 4u; + return (0x2E74u >> (i3 | i2 | i1)) & 0x0101u; +} + +fn solve_horiz(p1: vec2, p2: vec2, p3: vec2) -> vec2 { + let ay = p1.y - 2.0*p2.y + p3.y; + let by = p1.y - p2.y; + let ax = p1.x - 2.0*p2.x + p3.x; + let bx = p1.x - p2.x; + var t1: f32; var t2: f32; + if abs(ay) < (1.0 / 65536.0) { + let denom = select(1.0, by, abs(by) > 1e-10); + t1 = p1.y * 0.5 / denom; t2 = t1; + } else { + let ra = 1.0 / ay; + let d = sqrt(max(by*by - ay*p1.y, 0.0)); + t1 = (by - d) * ra; t2 = (by + d) * ra; + } + return vec2((ax*t1 - bx*2.0)*t1 + p1.x, (ax*t2 - bx*2.0)*t2 + p1.x); +} + +fn solve_vert(p1: vec2, p2: vec2, p3: vec2) -> vec2 { + let ax = p1.x - 2.0*p2.x + p3.x; + let bx = p1.x - p2.x; + let ay = p1.y - 2.0*p2.y + p3.y; + let by = p1.y - p2.y; + var t1: f32; var t2: f32; + if abs(ax) < (1.0 / 65536.0) { + let denom = select(1.0, bx, abs(bx) > 1e-10); + t1 = p1.x * 0.5 / denom; t2 = t1; + } else { + let ra = 1.0 / ax; + let d = sqrt(max(bx*bx - ax*p1.x, 0.0)); + t1 = (bx - d) * ra; t2 = (bx + d) * ra; + } + return vec2((ay*t1 - by*2.0)*t1 + p1.y, (ay*t2 - by*2.0)*t2 + p1.y); +} + +fn calc_coverage(xcov: f32, ycov: f32, xwgt: f32, ywgt: f32) -> f32 { + let blended = abs(xcov*xwgt + ycov*ywgt) / max(xwgt + ywgt, 1.0/65536.0); + return clamp(max(blended, min(abs(xcov), abs(ycov))), 0.0, 1.0); +} + +// Even-odd fill coverage: triangle wave — inside when winding count is odd. +// The raw winding accumulator (xcov or ycov) is a signed integer at stable +// interiors. A triangle wave with period 2 maps even integers → 0 (outside) +// and odd integers → 1 (inside), with half-pixel AA transitions. +fn calc_coverage_evenodd(xcov: f32, ycov: f32, xwgt: f32, ywgt: f32) -> f32 { + let w = (xcov*xwgt + ycov*ywgt) / max(xwgt + ywgt, 1.0/65536.0); + // Triangle wave: period 2, peak at odd integers. + let tri = 1.0 - abs(2.0 * fract(abs(w) * 0.5) - 1.0); + // Fallback: take the max of both axis coverages independently. + let tx = 1.0 - abs(2.0 * fract(abs(xcov) * 0.5) - 1.0); + let ty = 1.0 - abs(2.0 * fract(abs(ycov) * 0.5) - 1.0); + return clamp(max(tri, min(tx, ty)), 0.0, 1.0); +} + +// --------------------------------------------------------------------------- +// Stroke SDF helper — min distance from origin to a 2-D quadratic Bezier. +// +// B(t) = a·t² + b·t + c, a = p1−2·p2+p3, b = 2(p2−p1), c = p1. +// Fragment is at the origin, so B(t)−origin = B(t). +// +// Minimise |B(t)|² by Newton on f(t) = B(t)·B'(t). +// f'(t) = |B'(t)|² + 2a·B(t). +// --------------------------------------------------------------------------- + +fn min_dist_to_quad_px(p1: vec2, p2: vec2, p3: vec2) -> f32 { + let a = p1 - 2.0*p2 + p3; + let b = 2.0*(p2 - p1); + let c = p1; + + // Coarse: sample t = 0, 0.25, 0.5, 0.75, 1.0. + var best_t = 0.0; + var best_d2 = dot(c, c); + for (var i = 1u; i <= 4u; i++) { + let t = f32(i) * 0.25; + let bt = a*t*t + b*t + c; + let d2 = dot(bt, bt); + if d2 < best_d2 { best_d2 = d2; best_t = t; } + } + + // Newton refinement (4 iterations). + for (var iter = 0u; iter < 4u; iter++) { + let t = clamp(best_t, 0.0, 1.0); + let Bt = a*t*t + b*t + c; + let Bpt = 2.0*a*t + b; + let f = dot(Bt, Bpt); + let fp = dot(Bpt, Bpt) + dot(2.0*a, Bt); + if abs(fp) < 1e-10 { break; } + best_t = clamp(t - f/fp, 0.0, 1.0); + } + + let t_f = clamp(best_t, 0.0, 1.0); + let Bf = a*t_f*t_f + b*t_f + c; + return sqrt(max(dot(Bf, Bf), 0.0)); +} + +// --------------------------------------------------------------------------- +// Fragment shader +// --------------------------------------------------------------------------- + +@fragment +fn fs_main(in: VertexOutput) -> @location(0) vec4 { + // NDC units per screen pixel (non-zero denominator guard). + let ndc_per_pixel = fwidth(in.ndc_xy); + let pixels_per_ndc = 1.0 / max(ndc_per_pixel, vec2(1e-9)); + + let pv = u.projection * u.view; + + // ── Fill: Slug winding-number accumulation in NDC space ─────────────── + var xcov = 0.0; var xwgt = 0.0; + var ycov = 0.0; var ywgt = 0.0; + + for (var i = 0u; i < in.n_fill_curves; i++) { + let f = (in.fill_curve_start + i) * 9u; + let p1w = vec3(quads[f ], quads[f + 1u], quads[f + 2u]); + let p2w = vec3(quads[f + 3u], quads[f + 4u], quads[f + 5u]); + let p3w = vec3(quads[f + 6u], quads[f + 7u], quads[f + 8u]); + + // Transform world → NDC, shift so the current fragment is origin. + let c1 = pv * vec4(p1w, 1.0); + let c2 = pv * vec4(p2w, 1.0); + let c3 = pv * vec4(p3w, 1.0); + let p1 = c1.xy/c1.w - in.ndc_xy; + let p2 = c2.xy/c2.w - in.ndc_xy; + let p3 = c3.xy/c3.w - in.ndc_xy; + + // Horizontal ray (x-coverage accumulation). + let hcode = calc_root_code(p1.y, p2.y, p3.y); + if hcode != 0u { + let r = solve_horiz(p1, p2, p3) * pixels_per_ndc.x; + if (hcode & 1u) != 0u { + xcov += clamp(r.x + 0.5, 0.0, 1.0); + xwgt = max(xwgt, clamp(1.0 - abs(r.x)*2.0, 0.0, 1.0)); + } + if hcode > 1u { + xcov -= clamp(r.y + 0.5, 0.0, 1.0); + xwgt = max(xwgt, clamp(1.0 - abs(r.y)*2.0, 0.0, 1.0)); + } + } + + // Vertical ray (y-coverage accumulation). + let vcode = calc_root_code(p1.x, p2.x, p3.x); + if vcode != 0u { + let r = solve_vert(p1, p2, p3) * pixels_per_ndc.y; + if (vcode & 1u) != 0u { + ycov -= clamp(r.x + 0.5, 0.0, 1.0); + ywgt = max(ywgt, clamp(1.0 - abs(r.x)*2.0, 0.0, 1.0)); + } + if vcode > 1u { + ycov += clamp(r.y + 0.5, 0.0, 1.0); + ywgt = max(ywgt, clamp(1.0 - abs(r.y)*2.0, 0.0, 1.0)); + } + } + } + + var fill_cov = 0.0; + if in.n_fill_curves > 0u { + if in.fill_rule == 1u { + fill_cov = calc_coverage_evenodd(xcov, ycov, xwgt, ywgt); + } else { + fill_cov = calc_coverage(xcov, ycov, xwgt, ywgt); + } + } + + // ── Stroke: SDF minimum distance in physical pixel space ────────────── + // stroke_half_ndc is in NDC units; pixels_per_ndc.x converts to pixels. + // (stroke_half_ndc was calibrated using pm[0,0], the NDC x-scale.) + let stroke_half_px = in.v_stroke_half_ndc * pixels_per_ndc.x; + var min_dist_px = 1e9; + + for (var i = 0u; i < in.n_stroke_curves; i++) { + let f = (in.stroke_curve_start + i) * 9u; + let p1w = vec3(quads[f ], quads[f + 1u], quads[f + 2u]); + let p2w = vec3(quads[f + 3u], quads[f + 4u], quads[f + 5u]); + let p3w = vec3(quads[f + 6u], quads[f + 7u], quads[f + 8u]); + + let c1 = pv * vec4(p1w, 1.0); + let c2 = pv * vec4(p2w, 1.0); + let c3 = pv * vec4(p3w, 1.0); + // NDC-relative coordinates (fragment at origin), then scaled to pixels. + let n1 = c1.xy/c1.w - in.ndc_xy; + let n2 = c2.xy/c2.w - in.ndc_xy; + let n3 = c3.xy/c3.w - in.ndc_xy; + let p1_px = n1 * pixels_per_ndc; + let p2_px = n2 * pixels_per_ndc; + let p3_px = n3 * pixels_per_ndc; + + let d = min_dist_to_quad_px(p1_px, p2_px, p3_px); + min_dist_px = min(min_dist_px, d); + } + + // Smooth SDF: 1 within stroke, 0 outside, ½-pixel anti-aliased transition. + var stroke_cov = 0.0; + if in.n_stroke_curves > 0u && stroke_half_px > 0.0 { + stroke_cov = clamp(stroke_half_px + 0.5 - min_dist_px, 0.0, 1.0); + } + + // ── Porter-Duff "over": stroke on top of fill ───────────────────────── + let fill_a = in.v_fill_color.a * fill_cov; + let stroke_a = in.v_stroke_color.a * stroke_cov; + let total_a = stroke_a + fill_a * (1.0 - stroke_a); + + if total_a <= 0.001 { discard; } + + let fill_rgb = in.v_fill_color.rgb; + let stroke_rgb = in.v_stroke_color.rgb; + let out_rgb = (stroke_a * stroke_rgb + fill_a * (1.0 - stroke_a) * fill_rgb) / total_a; + + return vec4(out_rgb, total_a); +} diff --git a/manim/renderer/webgpu/webgpu_interactive.py b/manim/renderer/webgpu/webgpu_interactive.py new file mode 100644 index 0000000000..f5a387fe26 --- /dev/null +++ b/manim/renderer/webgpu/webgpu_interactive.py @@ -0,0 +1,260 @@ +"""Interactive IPython embed loop for the WebGPU renderer. + +Called by :meth:`~manim.scene.scene.Scene.interactive_embed` when the WebGPU +renderer is active. Mirrors the OpenGL :meth:`~manim.scene.scene.Scene.interact` +loop but drives the ``rendercanvas`` event loop instead of moderngl-window. + +Thread model +------------ +* **Main thread** — window event loop + scene method execution. Drains + ``scene.queue`` and calls scene methods (``play``, ``add``, …) so that all + GPU work stays on the thread that owns the WebGPU device. +* **IPython thread** (daemon) — blocking readline / prompt_toolkit. Scene + method calls are not executed here; they are posted to ``scene.queue`` and + picked up by the main thread. + +After every IPython cell ``post_run_cell`` schedules a re-render sentinel on +``scene.queue`` so the window immediately reflects any changes (``add``, +``remove``, property mutations …) that did not go through ``play`` / ``wait``. + +File watching +------------- +A ``watchdog.Observer`` watches the scene's source file. When the file is +saved on disk the observer posts ``SceneInteractRerun("file")`` to +``scene.queue``. The main loop then tears down the IPython session and raises +:class:`~manim.utils.exceptions.RerunSceneException`, which propagates through +``construct()`` back to the render command's rerun loop. + +``rerun()`` in the IPython shell triggers the same mechanism manually. +""" + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from manim.scene.scene import Scene + + from .webgpu_renderer import WebGPURenderer + +# Frames per second at which the main loop polls OS events when idle (no +# scene method is running). Higher values make mouse / keyboard more +# responsive; lower values save CPU. +_POLL_HZ: int = 60 + + +def interactive_embed( + scene: Scene, + renderer: WebGPURenderer, + local_namespace: dict[str, Any], +) -> bool: + """Run an interactive IPython session alongside the WebGPU preview window. + + Parameters + ---------- + scene: + The running scene instance. + renderer: + The active :class:`~.WebGPURenderer`. + local_namespace: + The caller's (``construct()``'s) local variables — captured in + :meth:`~manim.scene.scene.Scene.interactive_embed` before this + function is called. Scene shortcuts (``play``, ``wait``, ``add``, + ``remove``) and the full ``manim`` namespace are injected here so + the user can type commands without a ``self.`` prefix. + + Returns + ------- + bool + ``True`` if the session ended because ``rerun()`` was called or the + source file changed on disk — the caller should raise + :class:`~manim.utils.exceptions.RerunSceneException` in that case. + ``False`` for a normal exit (shell closed / window closed). + """ + import threading + + import manim + from manim import config, logger + from manim.data_structures import MethodWithArgs + from manim.scene.scene import SceneInteractContinue, SceneInteractRerun + + window = renderer.window + + # ── IPython imports ────────────────────────────────────────────────── + try: + from sqlite3 import connect + + from IPython.core.getipython import get_ipython + from IPython.terminal.embed import InteractiveShellEmbed + from traitlets.config import Config as IPConfig + except ImportError: + logger.error( + "IPython is required for the interactive WebGPU embed.\n" + "Install it with: pip install ipython", + ) + return False + + # ── File watcher ───────────────────────────────────────────────────── + try: + from watchdog.events import FileSystemEventHandler + from watchdog.observers import Observer + + class _FileHandler(FileSystemEventHandler): + def on_modified(self, event: Any) -> None: + scene.queue.put(SceneInteractRerun("file")) + + file_observer = Observer() + file_observer.schedule(_FileHandler(), config["input_file"], recursive=True) + file_observer.start() + except Exception: + logger.debug("watchdog not available — file-watching disabled.") + file_observer = None + + # ── Build shell ────────────────────────────────────────────────────── + ipcfg = IPConfig() + ipcfg.TerminalInteractiveShell.confirm_exit = False + + existing = get_ipython() + if existing is None: + shell = InteractiveShellEmbed.instance(config=ipcfg) + else: + shell = InteractiveShellEmbed(config=ipcfg) + + # Make the SQLite history database thread-safe so IPython history works + # correctly from the daemon keyboard thread. + hist = get_ipython().history_manager + hist.db = connect(hist.hist_file, check_same_thread=False) + + # ── Populate namespace ─────────────────────────────────────────────── + # Pre-import the full manim namespace so users don't need import + # statements inside the session. + for name in dir(manim): + local_namespace[name] = getattr(manim, name) + + # Proxy scene methods: posting to scene.queue keeps all GPU work on the + # main thread (matching the OpenGL embedded_method pattern). + def _make_proxy(method_name: str): + method = getattr(scene, method_name) + + def _proxy(*args: Any, **kwargs: Any) -> None: + scene.queue.put(MethodWithArgs(method, args, kwargs)) + + _proxy.__name__ = method_name + return _proxy + + for _name in ("play", "wait", "add", "remove"): + local_namespace[_name] = _make_proxy(_name) + + # rerun(): tear down this session and re-run the scene from scratch, + # picking up any edits saved to the source file. + def _rerun(*args: Any, **kwargs: Any) -> None: + scene.queue.put(SceneInteractRerun("keyboard")) + shell.exiter() + + local_namespace["rerun"] = _rerun + + # ── After every cell: schedule a re-render on the main thread ──────── + # _post_cell runs in the IPython thread — GPU calls must not happen here. + # Posting a sentinel to scene.queue ensures update_frame() is called on + # the main thread, which owns the WebGPU device. + _RENDER = object() # sentinel: "please re-render" + + def _post_cell(*_a: Any, **_kw: Any) -> None: + scene.queue.put(_RENDER) + + shell.events.register("post_run_cell", _post_cell) + + # ── IPython thread ─────────────────────────────────────────────────── + def _keyboard_thread() -> None: + shell(local_ns=local_namespace) + # Signal the main loop that the user closed the shell (not a rerun). + scene.queue.put(SceneInteractContinue("keyboard")) + + keyboard_thread = threading.Thread(target=_keyboard_thread) + # Run as a daemon so the thread is killed if the main thread exits + # (e.g. the window is closed before the shell prompt is answered). + if not shell.pt_app: + keyboard_thread.daemon = True + keyboard_thread.start() + + # ── Helpers ────────────────────────────────────────────────────────── + def _stop_file_observer() -> None: + if file_observer is not None: + file_observer.unschedule_all() + file_observer.stop() + file_observer.join() + + def _exit_keyboard_thread() -> None: + if shell.pt_app: + try: + shell.pt_app.app.exit(exception=EOFError) + except Exception: + pass + keyboard_thread.join() + while not scene.queue.empty(): + scene.queue.get() + + # ── Main thread: event loop ────────────────────────────────────────── + scene.quit_interaction = False + keyboard_thread_needs_join = shell.pt_app is not None + sleep_s = 1.0 / _POLL_HZ + rerun_requested = False + + while not (window.is_closing or scene.quit_interaction): + if not scene.queue.empty(): + action = scene.queue.get_nowait() + + if isinstance(action, SceneInteractRerun): + rerun_requested = True + _stop_file_observer() + if action.sender == "keyboard": + # rerun() was called from the shell — thread already + # exiting via shell.exiter(); just join it. + keyboard_thread.join() + else: + # File changed — kill the prompt and join. + _exit_keyboard_thread() + # Drain stale queue items and signal rerun to the caller. + while not scene.queue.empty(): + scene.queue.get() + break + + elif isinstance(action, SceneInteractContinue): + # IPython shell exited normally (user typed exit / Ctrl-D). + keyboard_thread.join() + while not scene.queue.empty(): + scene.queue.get() + keyboard_thread_needs_join = False + break + + elif isinstance(action, MethodWithArgs): + # Execute the proxied scene method on the main thread. + action.method(*action.args, **action.kwargs) + # Re-render so the result is visible immediately. + # (play/wait already render internally; add/remove do not.) + if renderer._device is not None: + renderer.update_frame(scene) + window._canvas.force_draw() + + elif action is _RENDER: + # Triggered by _post_cell — re-render after any IPython cell + # that mutated the scene without going through a proxy method. + if renderer._device is not None: + renderer.update_frame(scene) + window._canvas.force_draw() + else: + # Idle — process OS events so mouse/keyboard controls stay live. + window._canvas._process_events() + time.sleep(sleep_s) + + # ── Teardown ───────────────────────────────────────────────────────── + _stop_file_observer() + + if keyboard_thread_needs_join: + _exit_keyboard_thread() + + if window.is_closing: + window.destroy() + + return rerun_requested diff --git a/manim/renderer/webgpu/webgpu_renderer.py b/manim/renderer/webgpu/webgpu_renderer.py new file mode 100644 index 0000000000..7141512bdc --- /dev/null +++ b/manim/renderer/webgpu/webgpu_renderer.py @@ -0,0 +1,3283 @@ +"""WebGPU renderer for Manim — Phase 1. + +Phase 1 scope +------------- +* Headless rendering (no preview window). +* VMobject fill only. +* ``config.save_last_frame = True`` → saves a PNG. +* ``config.write_to_movie = True`` → writes video frames. + +Design +------ +Reads geometry directly from Cairo ``VMobject``. No dependency on OpenGL +classes (``OpenGLCamera``, ``OpenGLVMobject``, ``moderngl``). + +Camera +------ +A simple orthographic projection matrix maps Manim's frame coordinate system +(centre at origin, width = config.frame_width, height = config.frame_height) +to WebGPU NDC (x, y ∈ [-1, 1], z ∈ [0, 1]). +""" + +from __future__ import annotations + +import collections +import time +import weakref +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +from PIL import Image + +from manim import config, logger +from manim.constants import OUT, PI, RIGHT +from manim.mobject.mobject import Mobject +from manim.mobject.three_d.light_source import LightSource +from manim.mobject.three_d.three_dimensions import Surface +from manim.mobject.types.image_mobject import ( + AbstractImageMobject, + ImageMobjectFromCamera, +) +from manim.mobject.types.vectorized_mobject import VMobject +from manim.scene.scene_file_writer import SceneFileWriter +from manim.utils.color import color_to_rgba +from manim.utils.exceptions import EndSceneEarlyException +from manim.utils.hashing import get_hash_from_play_call +from manim.utils.iterables import list_update +from manim.utils.simple_functions import clip +from manim.utils.space_ops import ( + quaternion_from_angle_axis, + quaternion_mult, + rotation_matrix_transpose_from_quaternion, +) + +from .webgpu_vmobject_rendering import ( + FILL_STROKE_VERTEX_LAYOUT, + SURFACE_COMBINED_VERTEX_LAYOUT, + TRUE_DOT_VERTEX_LAYOUT, + DotCloud3D, + build_true_dot_vbo, + collect_frame_data, + draw_frame_data, + draw_frame_data_subcam, +) + +if TYPE_CHECKING: + import wgpu as wgpu_t + + from manim.scene.scene import Scene + + from .webgpu_renderer_window import WebGPUWindow + +try: + import wgpu +except ImportError as exc: + msg = ( + "wgpu-py is required for the WebGPU renderer. " + "Install it with: pip install wgpu" + ) + raise ImportError(msg) from exc + + +# --------------------------------------------------------------------------- +# Camera — feature-parity with OpenGLCamera for 2-D + 3-D scenes. +# --------------------------------------------------------------------------- + + +class WebGPUCamera(Mobject): + """Camera for the WebGPU renderer. + + Inherits from ``Mobject`` so it can carry updaters and be added to the + scene, matching the pattern used by ``OpenGLCamera(OpenGLMobject)``. + + Matches the attribute / method surface of ``OpenGLCamera`` so that + scene code that inspects ``renderer.camera`` works without changes. + + Projection + ---------- + * 2-D scenes: orthographic, z mapped to the WebGPU [0, 1] NDC + range. + * 3-D scenes: perspective projection driven by ``focal_distance`` + and the Euler-angle view matrix. + + Interactive controls (preview window only) + ------------------------------------------ + When a preview window is open, :class:`WebGPUWindow` handles mouse + events and mutates the camera in real time: + + * **Left-drag** — orbit: horizontal motion changes ``theta``; vertical + motion changes ``phi`` (clamped to ``[minimum_polar_angle, + maximum_polar_angle]``). + * **Right-drag / middle-drag** — pan: translates the view laterally in + camera space. The pan offset is owned by the window and injected into + ``view_matrix`` via ``_cam_pan_x`` / ``_cam_pan_y`` just before each + render; it is not part of the scripted camera model. + * **Scroll wheel** — zoom: adjusts ``focal_distance`` for perspective + cameras or scales ``frame_shape`` for orthographic cameras. + * **Key** ``r`` — reset: restores ``euler_angles`` and ``focal_distance`` + to their defaults and clears the window pan offset. + + Parameters + ---------- + frame_shape + (width, height) of the rendered frame. Defaults to + ``(config.frame_width, config.frame_height)``. + frame_center + World-space origin of the camera frame. Defaults to the origin. + euler_angles + (theta, phi, gamma) camera orientation angles in radians. + Defaults to (0, 0, 0) — looking straight down the −Z axis. + focal_distance + Perspective focal distance expressed as a multiple of ``frame_height``. + Only used when ``orthographic=False``. + orthographic + Use orthographic (True) or perspective (False) projection. + Default is True (matching Manim's default 2-D look). + minimum_polar_angle / maximum_polar_angle + Clamp range for the phi Euler angle during interactive orbit. + """ + + near: float = -100.0 + far: float = 100.0 + use_z_index: bool = True + + def __init__( + self, + frame_shape: tuple[float, float] | None = None, + frame_center: np.ndarray | None = None, + euler_angles: np.ndarray | None = None, + focal_distance: float = 20.0, + orthographic: bool = False, + minimum_polar_angle: float = -PI / 2, + maximum_polar_angle: float = PI / 2, + ) -> None: + super().__init__() + self.use_z_index = True + self.frame_rate: int = config.get("frame_rate", 60) + self.orthographic = orthographic + self.minimum_polar_angle = minimum_polar_angle + self.maximum_polar_angle = maximum_polar_angle + self.focal_distance = focal_distance + + self.frame_shape: tuple[float, float] = ( + frame_shape + if frame_shape is not None + else (float(config["frame_width"]), float(config["frame_height"])) + ) + self.frame_center: np.ndarray = ( + np.asarray(frame_center, dtype=float) + if frame_center is not None + else np.array([0.0, 0.0, focal_distance], dtype=float) + ) + # Default theta matches Cairo's default (-90°) so that the initial + # rotation formula (theta + 90°) gives identity for 2-D scenes. + self.euler_angles: np.ndarray = np.asarray( + euler_angles if euler_angles is not None else [-PI / 2, 0.0, 0.0], + dtype=float, + ) + self.reset_rotation_matrix() + + # Fixed-mobject registries — populated by ThreeDScene helpers. + # fixed_in_frame: objects rendered with identity rotation + ortho + # projection as a 2-D overlay on top of the 3-D scene (e.g. title text). + # fixed_orientation: objects rendered with identity rotation + current + # projection so they don't tilt as the camera orbits (e.g. 3-D labels). + self.fixed_in_frame_mobjects: set[Mobject] = set() + self.fixed_orientation_mobjects: set[Mobject] = set() + + # ThreeDScene.get_moving_mobjects() checks _frame_center and + # get_value_trackers() to detect camera-driven animation. + # These are defined on ThreeDCamera (Cairo) but not on Mobject, + # so we provide equivalent stubs here. + self._frame_center: Mobject = Mobject() + + # ImageMobjectFromCamera registration — for ZoomedScene support. + # Each frame the renderer renders the sub-camera view into a dedicated + # GPU texture which is then composited as a regular image quad. + self.image_mobjects_from_cameras: list = [] + + def get_value_trackers(self) -> list: + """Required by ThreeDScene.get_moving_mobjects. + + Returning ``[self]`` ensures that when the camera has updaters (e.g. + ambient rotation), ThreeDScene.get_moving_mobjects() detects the camera + in ``moving_mobjects`` and returns all scene mobjects — preventing the + static-frame optimisation from freezing the 3-D scene under camera + motion. + """ + return [self] + + def get_mobjects_indicating_movement(self) -> list: + """Return mobjects whose movement implies the whole scene is moving. + + Called by :class:`~.MovingCameraScene` to detect whether the camera + frame (or any registered sub-camera frame) is animated, which forces + all scene mobjects to be treated as moving so the static-frame + optimisation is skipped. + + Mirrors ``MultiCamera.get_mobjects_indicating_movement`` so that + :class:`~.ZoomedScene` works with the WebGPU renderer. + """ + return [imfc.camera.frame for imfc in self.image_mobjects_from_cameras] + + # ------------------------------------------------------------------ + # Frame geometry helpers (mirrors OpenGLCamera) + # ------------------------------------------------------------------ + + def get_width(self) -> float: + """Width of the camera frame in scene units.""" + return self.frame_shape[0] + + def get_height(self) -> float: + """Height of the camera frame in scene units.""" + return self.frame_shape[1] + + def get_shape(self) -> tuple[float, float]: + """(width, height) of the camera frame in scene units.""" + return self.frame_shape + + def get_center(self) -> np.ndarray: + """World-space centre of the camera frame.""" + return self.frame_center.copy() + + def get_focal_distance(self) -> float: + """Perspective focal distance in scene units.""" + return self.focal_distance * self.get_height() + + # ------------------------------------------------------------------ + # Camera reset + # ------------------------------------------------------------------ + + def to_default_state(self) -> WebGPUCamera: + """Reset frame size, position, and orientation to config defaults.""" + self.frame_shape = ( + float(config["frame_width"]), + float(config["frame_height"]), + ) + self.frame_center = np.array([0.0, 0.0, self.focal_distance], dtype=float) + self.euler_angles = np.array([-PI / 2, 0.0, 0.0]) + self.reset_rotation_matrix() + return self + + # ------------------------------------------------------------------ + # Rotation — matches OpenGLCamera.set/increment_* interface + # ------------------------------------------------------------------ + + def reset_rotation_matrix(self) -> None: + """Refresh the camera's inverse rotation matrix based on its Euler angles. + + The formula replicates Cairo's ThreeDCamera so that the same (theta, phi, + gamma) values produce the same view in both renderers. + + Cairo's generate_rotation_matrix builds: + R = R_z(gamma) @ R_x(-phi) @ R_z(-theta - 90°) (np.dot loop order) + and applies it to world column vectors in project_points. + + The WebGPU view matrix stores ``inverse_rotation_matrix`` and applies it + directly. ``rotation_matrix_transpose_from_quaternion(q)`` returns R_q^T + where R_q is the rotation for quaternion q. To get R_q^T == R_cairo we + need: + R_q = R_cairo^T = R_z(theta + 90°) @ R_x(phi) @ R_z(-gamma) + i.e. the quaternion that rotates: first by -gamma around Z, then by phi + around X, then by (theta + 90°) around Z: + q = q(theta + PI/2, OUT) * q(phi, RIGHT) * q(-gamma, OUT) + """ + theta, phi, gamma = self.euler_angles + quat = quaternion_mult( + quaternion_from_angle_axis(theta + PI / 2, OUT, axis_normalized=True), + quaternion_from_angle_axis(phi, RIGHT, axis_normalized=True), + quaternion_from_angle_axis(-gamma, OUT, axis_normalized=True), + ) + self.inverse_rotation_matrix: np.ndarray = np.array( + rotation_matrix_transpose_from_quaternion(np.asarray(quat, dtype=float)), + dtype=float, + ) + + def set_euler_angles( + self, + theta: float | None = None, + phi: float | None = None, + gamma: float | None = None, + ) -> WebGPUCamera: + if theta is not None: + self.euler_angles[0] = theta + if phi is not None: + self.euler_angles[1] = phi + if gamma is not None: + self.euler_angles[2] = gamma + self.reset_rotation_matrix() + return self + + def set_theta(self, theta: float) -> WebGPUCamera: + return self.set_euler_angles(theta=theta) + + def set_phi(self, phi: float) -> WebGPUCamera: + return self.set_euler_angles(phi=phi) + + def set_gamma(self, gamma: float) -> WebGPUCamera: + return self.set_euler_angles(gamma=gamma) + + _PERSPECTIVE_FAR: float = 200.0 + + def set_focal_distance(self, focal_distance: float) -> WebGPUCamera: + """Set the perspective focal distance. + + Matches Cairo's ``ThreeDCamera.focal_distance`` convention: larger + values push the camera further from the scene (same FOV, objects + appear further away); smaller values pull it closer. + + The near plane is derived as ``focal_distance / 6`` so that the + frustum height at depth ``focal_distance`` exactly equals the frame + height — matching Cairo's perspective formula + ``factor = focal_distance / (focal_distance - z_cam)``. + + ``focal_distance`` must be positive and less than ``_PERSPECTIVE_FAR``. + Values outside that range are clamped. + """ + max_fd = self._PERSPECTIVE_FAR * (1.0 - 1e-4) + clamped = float(np.clip(focal_distance, 1e-4, max_fd)) + if clamped != focal_distance: + logger.warning( + "WebGPUCamera.set_focal_distance: value %.4g clamped to %.4g " + "(must be in (0, far=%.4g))", + focal_distance, + clamped, + self._PERSPECTIVE_FAR, + ) + self.focal_distance = clamped + # Keep the virtual camera position (frame_center z) in sync so that + # the perspective projection exactly matches Cairo's formula at all depths. + self.frame_center[2] = clamped + return self + + def increment_theta(self, dtheta: float) -> WebGPUCamera: + self.euler_angles[0] += dtheta + self.reset_rotation_matrix() + return self + + def increment_phi(self, dphi: float) -> WebGPUCamera: + self.euler_angles[1] = clip( + self.euler_angles[1] + dphi, + self.minimum_polar_angle, + self.maximum_polar_angle, + ) + self.reset_rotation_matrix() + return self + + def increment_gamma(self, dgamma: float) -> WebGPUCamera: + self.euler_angles[2] += dgamma + self.reset_rotation_matrix() + return self + + # ------------------------------------------------------------------ + # View matrix (world → camera space) + # ------------------------------------------------------------------ + + @property + def view_matrix(self) -> np.ndarray: + """4×4 float32 view matrix: rotates and translates world space into + camera space. + + Uses T(-c) @ R_inv, which rotates the world around the origin (matches + OpenGLCamera behavior where the camera orbits the focal point). + + ``_cam_pan_x`` / ``_cam_pan_y`` are injected by :class:`WebGPUWindow` + just before each render call. They are not initialised in ``__init__`` + so that the camera model stays free of window/interaction state. + ``getattr`` defaults to 0 when no window is attached. + """ + R = np.asarray(self.inverse_rotation_matrix, dtype=np.float32) # 3×3 + c = self.frame_center.astype(np.float32) + view = np.eye(4, dtype=np.float32) + view[:3, :3] = R + # Camera-space translation: orbit distance along -Z plus lateral pan. + # Positive pan_x shifts the scene right (camera moves left), matching + # the "drag scene to the right" expectation for right-drag pan. + pan_x = float(getattr(self, "_cam_pan_x", 0.0)) + pan_y = float(getattr(self, "_cam_pan_y", 0.0)) + view[:3, 3] = [-pan_x, -pan_y, -c[2]] + return view + + @property + def fixed_view_matrix(self) -> np.ndarray: + """View matrix with camera rotation stripped — z-translation only. + + Used for fixed-orientation and fixed-in-frame mobjects so they don't + tilt or spin when the camera orbits. The z-translation is preserved so + depth ordering within the fixed layer is consistent with the main scene. + """ + view = np.eye(4, dtype=np.float32) + view[2, 3] = -float(self.frame_center[2]) + return view + + @property + def ortho_projection_matrix(self) -> np.ndarray: + """Forced orthographic projection matrix, regardless of self.orthographic. + + Fixed-in-frame overlays always use orthographic so that screen-space + coordinates map directly to Manim scene units (matching 2-D scenes). + """ + # Orthographic: map frame to NDC with z ∈ [0, 1]. + # Note: Manim's +Z is out of the screen (towards the viewer). + # So +Z should map to 0 (near) and -Z should map to 1 (far). + # Z_clip = -1/(far-near) * Z + far/(far-near) + fw, fh = self.frame_shape + near, far = self.near, self.far + return np.array( + [ + [2.0 / fw, 0.0, 0.0, 0.0], + [0.0, 2.0 / fh, 0.0, 0.0], + [0.0, 0.0, -1.0 / (far - near), far / (far - near)], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + + # ------------------------------------------------------------------ + # Projection matrix (used by the shader uniform upload) + # ------------------------------------------------------------------ + + @property + def projection_matrix(self) -> np.ndarray: + """4×4 float32 projection matrix in WebGPU NDC convention (z ∈ [0, 1]). + + Perspective when ``self.orthographic`` is False (default). + + Design: the near plane is derived as ``focal_distance / 6`` so that the + visible height at camera depth ``focal_distance`` (where world_z = 0 + maps to) exactly equals the frame height. Combined with + ``frame_center_z = focal_distance``, this exactly replicates Cairo's + perspective formula ``factor = focal_distance / (focal_distance - z_cam)`` + (for all depths, not just at world_z = 0). + """ + fw, fh = self.frame_shape + + if self.orthographic: + return self.ortho_projection_matrix + else: + # n = fd/6, w = fw/6, h = fh/6 → 2n/w = 2*fd/fw, 2n/h = 2*fd/fh + # → NDC_y = (2*fd/fh) * y / (-z_view) + # = (2*fd/fh) * y / (fd - z_cairo) [with z_view = z_cairo - fd] + # Cairo: NDC_y = fd/(fd-z_cairo) * y / (fh/2) = 2*fd/fh * y / (fd-z_cairo) ✓ + f = self._PERSPECTIVE_FAR + fd = self.focal_distance + n = float(np.clip(fd / 6.0, 1e-6, f * (1.0 - 1e-4))) + w, h = fw / 6.0, fh / 6.0 + return np.array( + [ + [2.0 * n / w, 0.0, 0.0, 0.0], + [0.0, 2.0 * n / h, 0.0, 0.0], + [0.0, 0.0, f / (n - f), n * f / (n - f)], + [0.0, 0.0, -1.0, 0.0], + ], + dtype=np.float32, + ) + + # ------------------------------------------------------------------ + # Fixed-mobject registry (used by ThreeDScene) + # ------------------------------------------------------------------ + + def add_fixed_in_frame_mobjects(self, *mobjects: Mobject) -> None: + """Register mobjects to be rendered as 2-D screen-space overlays. + + These objects are drawn after the 3-D scene with a fresh depth buffer, + identity camera rotation, and an orthographic projection so they always + appear on top at their 2-D screen-space coordinates. + """ + self.fixed_in_frame_mobjects.update(mobjects) + + def remove_fixed_in_frame_mobjects(self, *mobjects: Mobject) -> None: + """Unregister mobjects previously added with add_fixed_in_frame_mobjects.""" + self.fixed_in_frame_mobjects.difference_update(mobjects) + + def add_fixed_orientation_mobjects(self, *mobjects: Mobject) -> None: + """Register mobjects whose orientation is frozen relative to the camera. + + These objects still move in 3-D space (their world coordinates are used + normally) but the camera rotation is not applied — they remain upright as + the camera orbits. Useful for 3-D labels that should always face forward. + """ + self.fixed_orientation_mobjects.update(mobjects) + + def remove_fixed_orientation_mobjects(self, *mobjects: Mobject) -> None: + """Unregister mobjects previously added with add_fixed_orientation_mobjects.""" + self.fixed_orientation_mobjects.difference_update(mobjects) + + def add_image_mobject_from_camera(self, image_mob_from_camera: Any) -> None: + """Register an ImageMobjectFromCamera for sub-camera rendering. + + Called by :class:`~.ZoomedScene` when zooming is activated. Each + registered mob is rendered from its associated ``MovingCamera``'s + perspective into a dedicated GPU texture every frame. + + **WebGPU renderer only** — this method is a no-op for Cairo / OpenGL. + """ + if image_mob_from_camera not in self.image_mobjects_from_cameras: + self.image_mobjects_from_cameras.append(image_mob_from_camera) + + def remove_image_mobject_from_camera(self, image_mob_from_camera: Any) -> None: + """Unregister an ImageMobjectFromCamera previously added via + ``add_image_mobject_from_camera``. + """ + if image_mob_from_camera in self.image_mobjects_from_cameras: + self.image_mobjects_from_cameras.remove(image_mob_from_camera) + + +# --------------------------------------------------------------------------- +# Main renderer class +# --------------------------------------------------------------------------- + + +class WebGPURenderer: + """WebGPU renderer for Manim — headless and interactive preview. + + Supports PNG / video output in headless mode and a live preview window + (enabled by ``config.preview = True`` or the ``-p`` CLI flag). + + Preview window controls + ----------------------- + When the preview window is open, the camera can be navigated interactively + with the mouse. All interactions take effect immediately without restarting + the scene. + + ======================== ================================================ + Input Action + ======================== ================================================ + Left-drag **Orbit** — horizontal drag rotates theta (yaw); + vertical drag tilts phi (pitch), clamped to ±90°. + Right-drag / Middle-drag **Pan** — translates the view laterally in camera + space, proportional to the current frame size. + Scroll wheel **Zoom** — perspective: adjusts ``focal_distance`` + exponentially (~12 % per notch); orthographic: + scales ``frame_shape`` by the same factor. + Key ``r`` **Reset** — restores the camera's default orbit, + zoom, and clears any accumulated pan offset. + Key ``q`` **Quit** — closes the preview window. + ======================== ================================================ + + Customising window interaction + ------------------------------ + Subclass :class:`~.WebGPUWindow` and pass it via ``window_class`` to + override any interaction hook without modifying the renderer itself: + + ======================== ================================================ + Hook to override Triggered by + ======================== ================================================ + ``on_mouse_drag`` Pointer move while a button is held. + ``on_scroll`` Wheel / scroll event. + ``on_key_press`` Key pressed down. + ``on_key_release`` Key released. + ``on_mouse_left_click`` Left button clicked (no drag). + ``on_mouse_right_click`` Right button clicked (no drag). + ======================== ================================================ + + Building-block helpers ``orbit(dx, dy)``, ``pan(dx, dy)``, and + ``zoom(scroll_dy)`` are also overridable for finer control. See + :class:`~.WebGPUWindow` for a full example. + + MSAA + ---- + Pass ``msaa_samples=4`` to enable 4× multisample anti-aliasing. This + smooths geometric edges on surfaces, images, and dot-clouds (VMobjects + already use SDF/coverage-based AA so the gain for pure-2D scenes is + modest). MSAA bypasses the static-frame optimisation, re-rendering every + frame in full. + """ + + def __init__( + self, + file_writer_class: type[SceneFileWriter] = SceneFileWriter, + skip_animations: bool = False, + msaa_samples: int = 1, + window_class: type | None = None, + ) -> None: + """Create a WebGPU renderer. + + Parameters + ---------- + file_writer_class: + Class used to write frames to disk. + skip_animations: + When True the renderer skips all animations (used for caching). + msaa_samples: + Multisample Anti-Aliasing sample count. Must be 1 (disabled) or 4. + MSAA smooths geometric edges of surfaces, images, and dot-clouds. + VMobjects already use SDF/coverage-based AA, so the visual gain for + pure-2D scenes is modest; the benefit is most visible for 3-D + scenes with surface meshes and ``DotCloud3D``. + + When MSAA is enabled the static-frame optimisation is bypassed + (every frame is fully re-rendered), which increases per-frame GPU + work. This is acceptable because MSAA already implies a quality- + over-speed trade-off. + window_class: + Class used to create the preview window. Must be + :class:`~.WebGPUWindow` or a subclass of it. Defaults to + :class:`~.WebGPUWindow`. Pass a subclass to customise mouse/ + keyboard interaction by overriding :meth:`~.WebGPUWindow.on_mouse_drag`, + :meth:`~.WebGPUWindow.on_scroll`, :meth:`~.WebGPUWindow.on_key_press`, + or :meth:`~.WebGPUWindow.on_key_release`. + """ + if msaa_samples not in (1, 4): + msg = f"msaa_samples must be 1 or 4, got {msaa_samples}" + raise ValueError(msg) + self._msaa_samples = msaa_samples + self._window_class = window_class + self._file_writer_class = file_writer_class + self._original_skipping_status = skip_animations + self.skip_animations = skip_animations + + self.animation_start_time: float = 0.0 + self.animation_elapsed_time: float = 0.0 + self.time: float = 0.0 + self.num_plays: int = 0 + self.animations_hashes: list[str | None] = [] + + self.camera: WebGPUCamera = WebGPUCamera() + self.window: WebGPUWindow | None = None + self.pressed_keys: set[int] = set() + self._static_image: Any = None + self.file_writer: SceneFileWriter | None = None # set by init_scene() + + # Static-frame compositing (WP1). + # save_static_frame_data() renders static mobjects once into + # _static_texture; update_frame() blits it as the background and only + # re-draws the moving subset each animation frame. + self._static_texture: wgpu_t.GPUTexture | None = None + self._static_texture_view: wgpu_t.GPUTextureView | None = None + self._has_static_frame: bool = False + # IDs (id()) of the top-level mobjects that belong to the static layer. + # Used in render() to partition scene.mobjects into static vs dynamic. + self._static_mob_ids: set[int] = set() + + # SpecialThreeDScene reads renderer.camera_config["pixel_width"] to decide + # whether to apply low-quality overrides. Mirrors the pattern used by + # OpenGLRenderer so that SpecialThreeDScene works unchanged with WebGPU. + self.camera_config: dict = { + "pixel_width": config.pixel_width, + "pixel_height": config.pixel_height, + } + + self.background_color = config["background_color"] + + # Filled by init_scene(): + self._device: wgpu_t.GPUDevice | None = None + self._render_texture: wgpu_t.GPUTexture | None = None + self._render_texture_view: wgpu_t.GPUTextureView | None = None + self._depth_texture: wgpu_t.GPUTexture | None = None + self._depth_texture_view: wgpu_t.GPUTextureView | None = None + self._proj_bgl: wgpu_t.GPUBindGroupLayout | None = None + + # MSAA textures — created only when msaa_samples > 1. + # _msaa_texture: bgra8unorm, sample_count=N, RENDER_ATTACHMENT only. + # Used as the draw target in Pass 1; resolves to + # _render_texture at the end of each pass. + # _msaa_depth_texture: depth24plus, sample_count=N, RENDER_ATTACHMENT only. + # Must match the sample count of all MSAA pipelines. + self._msaa_texture: wgpu_t.GPUTexture | None = None + self._msaa_texture_view: wgpu_t.GPUTextureView | None = None + self._msaa_depth_texture: wgpu_t.GPUTexture | None = None + self._msaa_depth_texture_view: wgpu_t.GPUTextureView | None = None + + # Combined fill+stroke pipelines (vmobject_fill_stroke.wgsl). + # _fill_stroke_bgl is reused for both compute output and render input + # (camera uniform + read-only quads storage). + self._fill_stroke_bgl: wgpu_t.GPUBindGroupLayout | None = None + # Main pipelines — multisample count matches self._msaa_samples. + # Used in Pass 1 (the MSAA main-render pass). + self._fill_stroke_pipeline: wgpu_t.GPURenderPipeline | None = ( + None # 2-D, no depth write + ) + self._fill_stroke_3d_pipeline: wgpu_t.GPURenderPipeline | None = ( + None # 3-D, depth write + ) + # Overlay pipelines — always count=1. + # Used in Pass 4 (fixed-in-frame overlay, renders to _render_texture_view + # at sample_count=1 after the MSAA resolve has already completed). + self._fill_stroke_pipeline_1x: wgpu_t.GPURenderPipeline | None = None + self._fill_stroke_3d_pipeline_1x: wgpu_t.GPURenderPipeline | None = None + + # Compute pipeline: cubic_to_quads.wgsl. + self._compute_bgl: wgpu_t.GPUBindGroupLayout | None = None + self._cubic_to_quads_pipeline: wgpu_t.GPUComputePipeline | None = None + + # Surface pipelines: opaque (depth write, combined fill+wireframe) and OIT. + self._surface_pipeline: wgpu_t.GPURenderPipeline | None = None # opaque + self._surface_oit_pipeline: wgpu_t.GPURenderPipeline | None = None + # OIT accumulation textures (rgba16float each). + self._oit_accum_texture: wgpu_t.GPUTexture | None = None + self._oit_accum_view: wgpu_t.GPUTextureView | None = None + self._oit_reveal_texture: wgpu_t.GPUTexture | None = None + self._oit_reveal_view: wgpu_t.GPUTextureView | None = None + # OIT composition pipeline + bind group. + self._oit_compose_pipeline: wgpu_t.GPURenderPipeline | None = None + self._oit_compose_bgl: wgpu_t.GPUBindGroupLayout | None = None + self._oit_compose_bind_group: wgpu_t.GPUBindGroup | None = None + + # TrueDot pipeline (true_dot.wgsl) — renders DotCloud3D/PointDot as + # screen-aligned lit sphere quads (CPU-expanded, 6 verts per dot). + self._true_dot_pipeline: wgpu_t.GPURenderPipeline | None = None + + # Sub-camera pipelines — identical to the main pipelines but target + # rgba8unorm instead of bgra8unorm so the rendered texture can be + # sampled by the image pipeline without B↔R channel confusion. + # Used for ImageMobjectFromCamera (ZoomedScene support). + self._sub_cam_fill_stroke_pipeline: wgpu_t.GPURenderPipeline | None = None + self._sub_cam_fill_stroke_3d_pipeline: wgpu_t.GPURenderPipeline | None = None + self._sub_cam_surface_pipeline: wgpu_t.GPURenderPipeline | None = None + # Per-mob GPU resource cache: mob id → dict with render texture, + # depth texture, uniform buffer, camera bind group, tex bind group. + self._sub_cam_resources: dict[int, dict] = {} + + # Image pipeline (image.wgsl) — renders ImageMobject pixel arrays as + # textured quads before the VMobject pass (painter's algorithm). + self._image_pipeline: wgpu_t.GPURenderPipeline | None = None + self._image_tex_bgl: wgpu_t.GPUBindGroupLayout | None = None + self._image_tint_bgl: wgpu_t.GPUBindGroupLayout | None = None + # Cache: ImageMobject → (fingerprint, GPUTexture, GPUBindGroup). + # Keyed weakly so destroyed mobs release their GPU textures. + self._image_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + # Cache: ImageMobject → (points_fingerprint, GPUBuffer). + # VBO is reused across frames as long as mob.points[:4] hasn't changed. + self._image_vbo_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + # Cache: ImageMobject → (tint_fingerprint, GPUBuffer, GPUBindGroup). + # Tint bind group is rebuilt only when mob.color changes. + self._image_tint_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + # Compact readback compute pipeline (GPU row-depadding + B↔R fix). + self._readback_compute_pipeline: wgpu_t.GPUComputePipeline | None = None + self._readback_compute_bgl: wgpu_t.GPUBindGroupLayout | None = None + self._readback_compute_bind_group: wgpu_t.GPUBindGroup | None = None + # Storage buffer the compute shader writes into (STORAGE | COPY_SRC). + # One shared buffer — the compute shader always writes here first. + self._readback_storage_buf: wgpu_t.GPUBuffer | None = None + + # ── Staging-buffer pool ────────────────────────────────────────── + # Instead of a single MAP_READ buffer we keep a ring of N buffers. + # update_frame() writes the readback into the next free slot and + # starts map_async immediately. _get_mapped_frame_array() dequeues + # the oldest in-flight slot; by the time N frames have been submitted + # the GPU has had N frame-times to finish and sync_wait() returns + # instantly with no pipeline stall. + _READBACK_POOL = 3 # triple-buffering + self._READBACK_POOL: int = _READBACK_POOL + self._readback_pool: list[Any] = [] # GPUBuffer × _READBACK_POOL + # FIFO of slot indices submitted but not yet read. + self._readback_queue: collections.deque = collections.deque() + # Ring write pointer: next pool slot for update_frame to fill. + self._readback_write_slot: int = 0 + + # Cache the last successfully read pixel array. Cleared by + # update_frame() so that get_image() / get_frame() hit the GPU path + # only once per rendered frame; repeated calls within the same frame + # are served from this cache without any GPU interaction. + self._readback_cache: np.ndarray | None = None + + # Per-frame state (set during update_frame, cleared after submit). + self.current_render_pass: wgpu_t.GPURenderPassEncoder | None = None + self.camera_bind_group: wgpu_t.GPUBindGroup | None = None + self._camera_uniform_buf: wgpu_t.GPUBuffer | None = None + # Fixed-mobject bind groups (rebuilt each frame with stripped-rotation view). + # fixed_camera_bind_group: identity rotation + current projection (fixed-orientation) + # fixed_frame_bind_group: identity rotation + orthographic projection (fixed-in-frame) + self.fixed_camera_bind_group: wgpu_t.GPUBindGroup | None = None + self._fixed_orient_uniform_buf: wgpu_t.GPUBuffer | None = None + self.fixed_frame_bind_group: wgpu_t.GPUBindGroup | None = None + self._fixed_frame_uniform_buf: wgpu_t.GPUBuffer | None = None + self.frame_vbos: list[wgpu_t.GPUBuffer] = [] + + # _FrameData cache: keyed by cache slot name ("normal", "orient", "frame"). + # Each entry is (fingerprint_bytes, _FrameData). On a fingerprint hit we + # return the cached _FrameData, skipping all tessellation and buffer uploads. + self._fd_cache: dict[str, tuple[bytes, Any]] = {} + + # ------------------------------------------------------------------ + # static_image property — scene.py sets this to None at end of play() + # ------------------------------------------------------------------ + + @property + def static_image(self) -> Any: + return self._static_image + + @static_image.setter + def static_image(self, value: Any) -> None: + self._static_image = value + if value is None: + self._has_static_frame = False + self._static_mob_ids = set() + + # ------------------------------------------------------------------ + # Initialisation + # ------------------------------------------------------------------ + + def init_scene(self, scene: Scene) -> None: + """Create the wgpu device, offscreen texture, and file writer.""" + self.scene = scene + + if self._device is not None: + # Rerun path — reuse the existing GPU device, textures, and window. + # Only reset the per-scene state so the renderer is ready for a + # fresh construct() call without tearing down and recreating GPU + # resources (which would lose the live preview window). + self.partial_movie_files = [] + self.file_writer = self._file_writer_class(self, scene.__class__.__name__) + self.background_color = config["background_color"] + return + + self.partial_movie_files: list[str | None] = [] + self.file_writer: SceneFileWriter = self._file_writer_class( + self, + scene.__class__.__name__, + ) + + self.background_color = config["background_color"] + + adapter = wgpu.gpu.request_adapter_sync(power_preference="high-performance") + self._device = adapter.request_device_sync( + required_features=[], + required_limits={}, + ) + logger.debug("WebGPU adapter: %s", adapter.info) + + width = config.pixel_width + height = config.pixel_height + # bgra8unorm matches the window surface format on all major platforms + # (Metal/Vulkan/DX12), enabling copy_texture_to_texture without a + # blit shader. Readback in _get_mapped_frame_array() swaps B↔R to + # produce the RGBA output expected by PIL / numpy callers. + self._render_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.bgra8unorm, + usage=( + wgpu.TextureUsage.RENDER_ATTACHMENT + | wgpu.TextureUsage.COPY_SRC + | wgpu.TextureUsage.COPY_DST # receives blit from _static_texture + | wgpu.TextureUsage.TEXTURE_BINDING # read by compact-readback compute shader + ), + ) + self._render_texture_view = self._render_texture.create_view() + + # Static-frame texture: stores the pre-rendered static layer. + # Populated once per animation by save_static_frame_data(); blitted + # back into _render_texture each frame by update_frame(blit_static=True). + self._static_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.bgra8unorm, + usage=( + wgpu.TextureUsage.COPY_DST # written by copy from _render_texture + | wgpu.TextureUsage.COPY_SRC # read back into _render_texture each frame + ), + ) + self._static_texture_view = self._static_texture.create_view() + + self._depth_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.depth24plus, + usage=wgpu.TextureUsage.RENDER_ATTACHMENT, + ) + self._depth_texture_view = self._depth_texture.create_view() + + # MSAA textures — only when msaa_samples > 1. + if self._msaa_samples > 1: + self._msaa_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.bgra8unorm, + usage=wgpu.TextureUsage.RENDER_ATTACHMENT, + sample_count=self._msaa_samples, + ) + self._msaa_texture_view = self._msaa_texture.create_view() + self._msaa_depth_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.depth24plus, + usage=wgpu.TextureUsage.RENDER_ATTACHMENT, + sample_count=self._msaa_samples, + ) + self._msaa_depth_texture_view = self._msaa_depth_texture.create_view() + + self._proj_bgl = self._create_camera_bgl() + + # Surface mesh lines sit exactly on the surface triangles. A negative + # depth bias pulls each fragment slightly toward the camera so the mesh + # always wins the depth test without visually offsetting the lines. + # depth_bias=-10000 gives ~6e-4 constant offset in [0,1] depth space + # (depth24plus unit ≈ 6e-8), which is large enough to reliably beat + # floating-point depth jitter on flat/low-slope surface regions where + # depth_bias_slope_scale alone contributes nearly zero. + self._surface_pipeline = self._create_surface_pipeline( + self._proj_bgl, + cull_mode="none", + depth_write=True, + msaa_samples=self._msaa_samples, + ) + + # Combined fill+stroke pipeline (replaces separate slug + stroke pipelines). + self._fill_stroke_bgl, self._fill_stroke_pipeline = ( + self._create_fill_stroke_pipeline( + depth_test=False, msaa_samples=self._msaa_samples + ) + ) + _, self._fill_stroke_3d_pipeline = self._create_fill_stroke_pipeline( + depth_test=True, msaa_samples=self._msaa_samples + ) + + # Overlay pipelines — always count=1. Used in Pass 4 (fixed-in-frame) + # which renders directly into _render_texture_view after the MSAA resolve. + if self._msaa_samples > 1: + _, self._fill_stroke_pipeline_1x = self._create_fill_stroke_pipeline( + depth_test=False, msaa_samples=1 + ) + _, self._fill_stroke_3d_pipeline_1x = self._create_fill_stroke_pipeline( + depth_test=True, msaa_samples=1 + ) + else: + # When MSAA is off the overlay pipelines are the same objects. + self._fill_stroke_pipeline_1x = self._fill_stroke_pipeline + self._fill_stroke_3d_pipeline_1x = self._fill_stroke_3d_pipeline + + # GPU compute: cubic → quadratic conversion. + self._compute_bgl, self._cubic_to_quads_pipeline = ( + self._create_cubic_to_quads_pipeline() + ) + + self._create_oit_resources(width, height) + self._create_readback_pipeline(width, height) + self._image_tex_bgl, self._image_tint_bgl, self._image_pipeline = ( + self._create_image_pipeline(msaa_samples=self._msaa_samples) + ) + self._true_dot_pipeline = self._create_true_dot_pipeline( + self._proj_bgl, + msaa_samples=self._msaa_samples, + ) + + # Sub-camera pipelines (rgba8unorm target) for ZoomedScene support. + # Sub-camera targets are always count=1 (they render into their own + # rgba8unorm textures, not into the MSAA main buffer). + _, self._sub_cam_fill_stroke_pipeline = self._create_fill_stroke_pipeline( + depth_test=False, target_format="rgba8unorm", msaa_samples=1 + ) + _, self._sub_cam_fill_stroke_3d_pipeline = self._create_fill_stroke_pipeline( + depth_test=True, target_format="rgba8unorm", msaa_samples=1 + ) + self._sub_cam_surface_pipeline = self._create_surface_pipeline( + self._proj_bgl, + cull_mode="none", + depth_write=True, + target_format="rgba8unorm", + msaa_samples=1, + ) + + # Persistent camera uniform buffers — created once, updated each frame via + # write_buffer. Using COPY_DST so queue.write_buffer can write into them. + # These stable GPU objects let cached _FrameData bind groups remain valid + # across frames: the bind group references the same buffer; write_buffer + # updates its contents so the shader always sees the current camera. + # Uniform buffer size: proj(64) + view(64) + num_lights+pad(16) + Light×8(512) = 656 B + _UBO_SIZE = 656 + self._camera_uniform_buf = self._device.create_buffer( + size=_UBO_SIZE, + usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST, + ) + self._fixed_orient_uniform_buf = self._device.create_buffer( + size=_UBO_SIZE, + usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST, + ) + self._fixed_frame_uniform_buf = self._device.create_buffer( + size=_UBO_SIZE, + usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST, + ) + + # Persistent camera bind groups — constant layout + constant buffer objects, + # so they never need to be recreated. + def _make_persistent_bg(buf: wgpu_t.GPUBuffer) -> wgpu_t.GPUBindGroup: + return self._device.create_bind_group( + layout=self._proj_bgl, + entries=[ + { + "binding": 0, + "resource": {"buffer": buf, "offset": 0, "size": _UBO_SIZE}, + } + ], + ) + + self.camera_bind_group = _make_persistent_bg(self._camera_uniform_buf) + self.fixed_camera_bind_group = _make_persistent_bg( + self._fixed_orient_uniform_buf + ) + self.fixed_frame_bind_group = _make_persistent_bg(self._fixed_frame_uniform_buf) + + if self.should_create_window(): + from .webgpu_renderer_window import WebGPUWindow + + wclass = self._window_class or WebGPUWindow + self.window = wclass(self) + + # ------------------------------------------------------------------ + # Pipeline creation + # ------------------------------------------------------------------ + + def _create_camera_bgl(self) -> wgpu_t.GPUBindGroupLayout: + """Create the bind group layout shared by stroke, surface, and Slug pipelines. + + Layout: binding 0 — one uniform buffer (656 bytes total): + offset 0 — projection mat4x4 64 B + offset 64 — view mat4x4 64 B + offset 128 — num_lights u32 4 B + offset 132 — _pad u32 × 3 12 B + offset 144 — lights Light × 8 512 B (each Light = 64 B) + """ + assert self._device is not None + return self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.VERTEX | wgpu.ShaderStage.FRAGMENT, + "buffer": {"type": "uniform"}, + } + ] + ) + + def _create_fill_stroke_pipeline( + self, + depth_test: bool = False, + target_format: str = "bgra8unorm", + msaa_samples: int = 1, + ) -> tuple[wgpu_t.GPUBindGroupLayout, wgpu_t.GPURenderPipeline]: + """Create the combined fill+stroke pipeline (vmobject_fill_stroke.wgsl). + + The bind group layout mirrors the slug fill layout: + binding 0 — camera uniform (656 bytes) + binding 1 — quads storage buffer (read-only, output of compute shader) + + depth_test=False — 2-D objects: depth-read-only (painter's algorithm). + depth_test=True — 3-D objects: depth-write + depth-test. + """ + assert self._device is not None + shader_path = Path(__file__).parent / "shaders" / "vmobject_fill_stroke.wgsl" + shader_module = self._device.create_shader_module( + code=shader_path.read_text(encoding="utf-8") + ) + + bgl = self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.VERTEX | wgpu.ShaderStage.FRAGMENT, + "buffer": {"type": "uniform"}, + }, + { + "binding": 1, + "visibility": wgpu.ShaderStage.FRAGMENT, + "buffer": { + "type": "read-only-storage", + "has_dynamic_offset": False, + }, + }, + ] + ) + + _blend = { + "color": { + "src_factor": "src-alpha", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + "alpha": { + "src_factor": "one", + "dst_factor": "one", + "operation": "add", + }, + } + + pipeline = self._device.create_render_pipeline( + layout=self._device.create_pipeline_layout(bind_group_layouts=[bgl]), + vertex={ + "module": shader_module, + "entry_point": "vs_main", + "buffers": [FILL_STROKE_VERTEX_LAYOUT], + }, + fragment={ + "module": shader_module, + "entry_point": "fs_main", + "targets": [ + { + "format": getattr(wgpu.TextureFormat, target_format), + "blend": _blend, + } + ], + }, + primitive={"topology": "triangle-list", "cull_mode": "none"}, + depth_stencil={ + "format": wgpu.TextureFormat.depth24plus, + "depth_write_enabled": depth_test, + "depth_compare": "less", + "stencil_front": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_back": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_read_mask": 0, + "stencil_write_mask": 0, + }, + multisample={ + "count": msaa_samples, + "mask": 0xFFFF_FFFF, + "alpha_to_coverage_enabled": False, + }, + ) + return bgl, pipeline + + def _create_cubic_to_quads_pipeline( + self, + ) -> tuple[wgpu_t.GPUBindGroupLayout, wgpu_t.GPUComputePipeline]: + """Create the compute pipeline that converts cubics → quadratics. + + Bind group layout: + binding 0 — input cubics (read-only-storage, 12 floats/cubic) + binding 1 — output quads (storage read_write, 36 floats/cubic) + binding 2 — params uniform (n_cubics u32, padded to 16 bytes) + + Dispatch: ceil(n_cubics / 64) × 1 × 1 workgroups. + """ + assert self._device is not None + shader_path = Path(__file__).parent / "shaders" / "cubic_to_quads.wgsl" + shader_module = self._device.create_shader_module( + code=shader_path.read_text(encoding="utf-8") + ) + + bgl = self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.COMPUTE, + "buffer": {"type": "read-only-storage"}, + }, + { + "binding": 1, + "visibility": wgpu.ShaderStage.COMPUTE, + "buffer": {"type": "storage"}, + }, + { + "binding": 2, + "visibility": wgpu.ShaderStage.COMPUTE, + "buffer": {"type": "uniform"}, + }, + ] + ) + + pipeline = self._device.create_compute_pipeline( + layout=self._device.create_pipeline_layout(bind_group_layouts=[bgl]), + compute={"module": shader_module, "entry_point": "main"}, + ) + return bgl, pipeline + + def _create_surface_pipeline( + self, + proj_bgl: wgpu_t.GPUBindGroupLayout, + cull_mode: str = "none", + depth_write: bool = True, + target_format: str = "bgra8unorm", + msaa_samples: int = 1, + ) -> wgpu_t.GPURenderPipeline: + """Create a surface (mesh) pipeline. + + cull_mode — WebGPU cull mode passed directly to the pipeline. + Use "back" for opaque surfaces (back faces are never + visible and culling them halves fragment work). + Use "none" for OIT transparent surfaces (both faces + must contribute so the interior is visible through + the front face). + depth_write — True for opaque surfaces so they occlude later geometry. + False for OIT surfaces so transparent layers do not block + each other (they still depth-test against opaque geometry). + """ + assert self._device is not None + shader_path = Path(__file__).parent / "shaders" / "surface_combined.wgsl" + shader_module = self._device.create_shader_module( + code=shader_path.read_text(encoding="utf-8") + ) + _blend = { + "color": { + "src_factor": "src-alpha", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + "alpha": { + "src_factor": "one", + "dst_factor": "one", + "operation": "add", + }, + } + return self._device.create_render_pipeline( + layout=self._device.create_pipeline_layout(bind_group_layouts=[proj_bgl]), + vertex={ + "module": shader_module, + "entry_point": "vs_main", + "buffers": [SURFACE_COMBINED_VERTEX_LAYOUT], + }, + fragment={ + "module": shader_module, + "entry_point": "fs_main", + "targets": [ + { + "format": getattr(wgpu.TextureFormat, target_format), + "blend": _blend, + } + ], + }, + primitive={"topology": "triangle-list", "cull_mode": cull_mode}, + depth_stencil={ + "format": wgpu.TextureFormat.depth24plus, + "depth_write_enabled": depth_write, + "depth_compare": "less", + "stencil_front": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_back": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_read_mask": 0, + "stencil_write_mask": 0, + }, + multisample={ + "count": msaa_samples, + "mask": 0xFFFF_FFFF, + "alpha_to_coverage_enabled": False, + }, + ) + + def _create_true_dot_pipeline( + self, + proj_bgl: wgpu_t.GPUBindGroupLayout, + msaa_samples: int = 1, + ) -> wgpu_t.GPURenderPipeline: + """Create the TrueDot pipeline (true_dot.wgsl). + + Reuses the camera bind group layout (``proj_bgl``) at group 0. + Depth write is enabled so dots occlude each other and other geometry. + Alpha blending is on so the anti-aliased disc edge fades smoothly. + """ + assert self._device is not None + shader_path = Path(__file__).parent / "shaders" / "true_dot.wgsl" + shader = self._device.create_shader_module( + code=shader_path.read_text(encoding="utf-8") + ) + _blend = { + "color": { + "src_factor": "src-alpha", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + "alpha": { + "src_factor": "one", + "dst_factor": "one", + "operation": "add", + }, + } + return self._device.create_render_pipeline( + layout=self._device.create_pipeline_layout(bind_group_layouts=[proj_bgl]), + vertex={ + "module": shader, + "entry_point": "vs_main", + "buffers": [TRUE_DOT_VERTEX_LAYOUT], + }, + fragment={ + "module": shader, + "entry_point": "fs_main", + "targets": [{"format": wgpu.TextureFormat.bgra8unorm, "blend": _blend}], + }, + primitive={"topology": "triangle-list", "cull_mode": "none"}, + depth_stencil={ + "format": wgpu.TextureFormat.depth24plus, + "depth_write_enabled": True, + "depth_compare": "less", + "stencil_front": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_back": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_read_mask": 0, + "stencil_write_mask": 0, + }, + multisample={ + "count": msaa_samples, + "mask": 0xFFFF_FFFF, + "alpha_to_coverage_enabled": False, + }, + ) + + def _create_image_pipeline( + self, + msaa_samples: int = 1, + ) -> tuple[ + wgpu_t.GPUBindGroupLayout, wgpu_t.GPUBindGroupLayout, wgpu_t.GPURenderPipeline + ]: + """Create the render pipeline for ImageMobject textured quads. + + Layout + ------ + group 0 — camera uniform (reuses ``_proj_bgl``, same as VMobject shaders) + group 1 — texture_2d at binding 0, sampler at binding 1 + group 2 — tint uniform: vec3 rgb colour multiplier (16-byte block) + + Vertex buffer (stride 20 B): + location 0 — in_pos float32x3 (12 B) + location 1 — in_uv float32x2 ( 8 B) + """ + assert self._device is not None + assert self._proj_bgl is not None + + shader_path = Path(__file__).parent / "shaders" / "image.wgsl" + shader = self._device.create_shader_module( + code=shader_path.read_text(encoding="utf-8") + ) + + # Group 1: texture + sampler + tex_bgl = self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.FRAGMENT, + "texture": { + "sample_type": "float", + "view_dimension": "2d", + "multisampled": False, + }, + }, + { + "binding": 1, + "visibility": wgpu.ShaderStage.FRAGMENT, + "sampler": {"type": "filtering"}, + }, + ] + ) + + # Group 2: tint colour uniform (16 bytes: rgb vec3 + 4-byte pad) + tint_bgl = self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.FRAGMENT, + "buffer": {"type": "uniform", "min_binding_size": 16}, + }, + ] + ) + + layout = self._device.create_pipeline_layout( + bind_group_layouts=[self._proj_bgl, tex_bgl, tint_bgl] + ) + + blend = { + "color": { + "src_factor": "src-alpha", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + "alpha": { + "src_factor": "one", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + } + + pipeline = self._device.create_render_pipeline( + layout=layout, + vertex={ + "module": shader, + "entry_point": "vs_main", + "buffers": [ + { + "array_stride": 20, # 3+2 floats × 4 B + "step_mode": "vertex", + "attributes": [ + {"format": "float32x3", "offset": 0, "shader_location": 0}, + {"format": "float32x2", "offset": 12, "shader_location": 1}, + ], + } + ], + }, + fragment={ + "module": shader, + "entry_point": "fs_main", + "targets": [ + { + "format": wgpu.TextureFormat.bgra8unorm, + "blend": blend, + } + ], + }, + primitive={"topology": "triangle-list", "cull_mode": "none"}, + depth_stencil={ + # Must match the render pass's depth format. + # Images use painter's algorithm (draw order), not depth test. + "format": wgpu.TextureFormat.depth24plus, + "depth_write_enabled": False, + "depth_compare": "always", + "stencil_front": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_back": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_read_mask": 0, + "stencil_write_mask": 0, + }, + multisample={ + "count": msaa_samples, + "mask": 0xFFFF_FFFF, + "alpha_to_coverage_enabled": False, + }, + ) + + return tex_bgl, tint_bgl, pipeline + + def _image_fingerprint(self, pixel_array: np.ndarray) -> int: + """Cheap dirty-check fingerprint for a pixel array. + + Samples the first 64, middle 64, and last 64 bytes of the flattened + array plus the shape tuple — fast enough for typical image sizes and + reliably detects in-place changes such as set_opacity(). + """ + flat = pixel_array.ravel() + n = len(flat) + if n <= 192: + return hash((pixel_array.shape, flat.tobytes())) + mid = n // 2 + sample = np.concatenate([flat[:64], flat[mid : mid + 64], flat[-64:]]) + return hash((pixel_array.shape, sample.tobytes())) + + def _get_image_gpu_resources( + self, mob: Any + ) -> tuple[wgpu_t.GPUTexture, wgpu_t.GPUBindGroup] | None: + """Return (texture, bind_group) for *mob*, re-uploading if pixel_array changed. + + Returns None if the mob has no valid pixel array. + """ + assert self._device is not None + assert self._image_tex_bgl is not None + + pixel_array: np.ndarray | None = getattr(mob, "pixel_array", None) + if pixel_array is None or pixel_array.ndim != 3 or pixel_array.shape[2] < 4: + return None + + fp = self._image_fingerprint(pixel_array) + cached = self._image_cache.get(mob) + if cached is not None and cached[0] == fp: + return cached[1], cached[2] + + # (Re-)upload texture. + h, w = pixel_array.shape[:2] + # Ensure RGBA uint8. + if pixel_array.dtype != np.uint8: + pixel_array = pixel_array.astype(np.uint8) + + # bytes_per_row must be a multiple of 256. + bytes_per_row = w * 4 + aligned_bpr = (bytes_per_row + 255) & ~255 + if aligned_bpr == bytes_per_row: + data = pixel_array.tobytes() + else: + rows = [ + pixel_array[r].ravel().tobytes() + + b"\x00" * (aligned_bpr - bytes_per_row) + for r in range(h) + ] + data = b"".join(rows) + + tex = self._device.create_texture( + size=(w, h, 1), + format=wgpu.TextureFormat.rgba8unorm, + usage=wgpu.TextureUsage.TEXTURE_BINDING | wgpu.TextureUsage.COPY_DST, + ) + self._device.queue.write_texture( + {"texture": tex, "mip_level": 0, "origin": (0, 0, 0)}, + data, + {"bytes_per_row": aligned_bpr, "rows_per_image": h}, + (w, h, 1), + ) + + sampler = self._device.create_sampler( + min_filter="linear", + mag_filter="linear", + address_mode_u="clamp-to-edge", + address_mode_v="clamp-to-edge", + ) + + bg = self._device.create_bind_group( + layout=self._image_tex_bgl, + entries=[ + {"binding": 0, "resource": tex.create_view()}, + {"binding": 1, "resource": sampler}, + ], + ) + + self._image_cache[mob] = (fp, tex, bg) + return tex, bg + + def _get_image_vbo(self, mob: Any) -> wgpu_t.GPUBuffer | None: + """Return a cached 6-vertex (20 B/vertex) VBO for *mob*'s bounding quad. + + The VBO is rebuilt only when mob.points[:4] changes; otherwise the + same GPUBuffer is reused across frames without any CPU/GPU allocation. + + Corner layout from AbstractImageMobject.reset_points(): + points[0] = UP + LEFT → UV (0, 0) + points[1] = UP + RIGHT → UV (1, 0) + points[2] = DOWN + LEFT → UV (0, 1) + points[3] = DOWN + RIGHT→ UV (1, 1) + + Two CCW triangles: [0,1,2] and [1,3,2]. + + ``scale_to_resolution`` semantics are enforced at the Mobject level: + ``AbstractImageMobject.reset_points()`` converts pixel dimensions to + world-space units using ``scale_to_resolution`` during ``__init__``. + The renderer reads the resulting world-space corner positions directly + from ``mob.points`` — no additional scaling is applied here. + """ + assert self._device is not None + + pts = getattr(mob, "points", None) + if pts is None or len(pts) < 4: + return None + + # Fingerprint the 4 corner points exactly (48 bytes — cheap). + corners = pts[:4].astype(np.float32) + fp = hash(corners.tobytes()) + + cached = self._image_vbo_cache.get(mob) + if cached is not None and cached[0] == fp: + return cached[1] + + uvs = np.array( + [[0.0, 0.0], [1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=np.float32 + ) + idx = [0, 1, 2, 1, 3, 2] + data = np.empty((6, 5), dtype=np.float32) + data[:, :3] = corners[idx] + data[:, 3:] = uvs[idx] + + buf = self._device.create_buffer_with_data( + data=data.tobytes(), + usage=wgpu.BufferUsage.VERTEX, + ) + # Store persistently — NOT in frame_vbos; WeakKeyDictionary releases + # the buffer when the mob is garbage-collected. + self._image_vbo_cache[mob] = (fp, buf) + return buf + + def _get_image_tint_bind_group(self, mob: Any) -> wgpu_t.GPUBindGroup | None: + """Return a cached GPUBindGroup for *mob*'s tint colour uniform. + + The bind group is rebuilt only when ``mob.color`` changes. The + default WHITE tint ``(1, 1, 1)`` is identity — texture is unchanged. + """ + assert self._device is not None + assert self._image_tint_bgl is not None + + # Read mob.color → (r, g, b) in [0, 1]. Fall back to white. + try: + rgb = np.asarray(mob.color.to_rgb(), dtype=np.float32) + except Exception: + rgb = np.ones(3, dtype=np.float32) + + fp = hash(rgb.tobytes()) + cached = self._image_tint_cache.get(mob) + if cached is not None and cached[0] == fp: + return cached[2] + + # 16-byte block: rgb (12 B) + 4-byte pad. + data = np.array([rgb[0], rgb[1], rgb[2], 0.0], dtype=np.float32) + buf = self._device.create_buffer_with_data( + data=data.tobytes(), + usage=wgpu.BufferUsage.UNIFORM, + ) + bg = self._device.create_bind_group( + layout=self._image_tint_bgl, + entries=[ + {"binding": 0, "resource": {"buffer": buf, "offset": 0, "size": 16}} + ], + ) + self._image_tint_cache[mob] = (fp, buf, bg) + return bg + + def _create_oit_resources(self, width: int, height: int) -> None: + """Create OIT accumulation textures, pipelines, and bind groups.""" + assert self._device is not None + assert self._proj_bgl is not None + + # ── Accumulation textures ────────────────────────────────────────── + oit_usage = ( + wgpu.TextureUsage.RENDER_ATTACHMENT | wgpu.TextureUsage.TEXTURE_BINDING + ) + self._oit_accum_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.rgba16float, + usage=oit_usage, + ) + self._oit_accum_view = self._oit_accum_texture.create_view() + + self._oit_reveal_texture = self._device.create_texture( + size=(width, height, 1), + format=wgpu.TextureFormat.rgba16float, + usage=oit_usage, + ) + self._oit_reveal_view = self._oit_reveal_texture.create_view() + + # ── OIT accumulation pipeline ────────────────────────────────────── + oit_shader_path = Path(__file__).parent / "shaders" / "surface_oit.wgsl" + oit_shader = self._device.create_shader_module( + code=oit_shader_path.read_text(encoding="utf-8") + ) + _accum_blend = { + "color": {"src_factor": "one", "dst_factor": "one", "operation": "add"}, + "alpha": {"src_factor": "one", "dst_factor": "one", "operation": "add"}, + } + _reveal_blend = { + "color": { + "src_factor": "zero", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + "alpha": {"src_factor": "zero", "dst_factor": "one", "operation": "add"}, + } + self._surface_oit_pipeline = self._device.create_render_pipeline( + layout=self._device.create_pipeline_layout( + bind_group_layouts=[self._proj_bgl] + ), + vertex={ + "module": oit_shader, + "entry_point": "vs_main", + "buffers": [SURFACE_COMBINED_VERTEX_LAYOUT], + }, + fragment={ + "module": oit_shader, + "entry_point": "fs_main", + "targets": [ + {"format": wgpu.TextureFormat.rgba16float, "blend": _accum_blend}, + {"format": wgpu.TextureFormat.rgba16float, "blend": _reveal_blend}, + ], + }, + primitive={"topology": "triangle-list", "cull_mode": "none"}, + depth_stencil={ + "format": wgpu.TextureFormat.depth24plus, + "depth_write_enabled": False, + "depth_compare": "less", + "stencil_front": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_back": { + "compare": "always", + "fail_op": "keep", + "depth_fail_op": "keep", + "pass_op": "keep", + }, + "stencil_read_mask": 0, + "stencil_write_mask": 0, + }, + multisample={ + "count": 1, + "mask": 0xFFFF_FFFF, + "alpha_to_coverage_enabled": False, + }, + ) + + # ── OIT composition pipeline ─────────────────────────────────────── + compose_path = Path(__file__).parent / "shaders" / "oit_compose.wgsl" + compose_shader = self._device.create_shader_module( + code=compose_path.read_text(encoding="utf-8") + ) + self._oit_compose_bgl = self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.FRAGMENT, + "texture": { + "sample_type": "unfilterable-float", + "view_dimension": "2d", + "multisampled": False, + }, + }, + { + "binding": 1, + "visibility": wgpu.ShaderStage.FRAGMENT, + "texture": { + "sample_type": "unfilterable-float", + "view_dimension": "2d", + "multisampled": False, + }, + }, + ] + ) + _compose_blend = { + "color": { + "src_factor": "src-alpha", + "dst_factor": "one-minus-src-alpha", + "operation": "add", + }, + "alpha": {"src_factor": "one", "dst_factor": "one", "operation": "add"}, + } + self._oit_compose_pipeline = self._device.create_render_pipeline( + layout=self._device.create_pipeline_layout( + bind_group_layouts=[self._oit_compose_bgl] + ), + vertex={"module": compose_shader, "entry_point": "vs_main", "buffers": []}, + fragment={ + "module": compose_shader, + "entry_point": "fs_main", + "targets": [ + {"format": wgpu.TextureFormat.bgra8unorm, "blend": _compose_blend} + ], + }, + primitive={"topology": "triangle-list", "cull_mode": "none"}, + multisample={ + "count": 1, + "mask": 0xFFFF_FFFF, + "alpha_to_coverage_enabled": False, + }, + ) + self._oit_compose_bind_group = self._device.create_bind_group( + layout=self._oit_compose_bgl, + entries=[ + {"binding": 0, "resource": self._oit_accum_view}, + {"binding": 1, "resource": self._oit_reveal_view}, + ], + ) + + # ------------------------------------------------------------------ + # Camera bind group (rebuilt each frame when projection changes) + # ------------------------------------------------------------------ + + def _collect_lights(self) -> list: + """Return all LightSource instances in the current scene. + + Lights are always added directly to the scene via ``self.add(light)``, + so scanning the top-level ``scene.mobjects`` list is sufficient and + avoids a deep O(N_submobjects) traversal through thousands of Surface + patches every frame. + """ + if self.scene is None: + return [] + return [m for m in self.scene.mobjects if isinstance(m, LightSource)] + + _MAX_LIGHTS = 8 + + def _pack_camera_uniforms_bytes( + self, + proj: np.ndarray, + view: np.ndarray, + ) -> bytes: + """Return a 656-byte camera+lighting uniform payload from explicit proj/view. + + Layout (matches Uniforms struct in surface_combined.wgsl / surface_oit.wgsl): + offset 0 — projection mat4x4 64 B + offset 64 — view mat4x4 64 B + offset 128 — num_lights u32 4 B + offset 132 — _pad u32 × 3 12 B + offset 144 — lights Light × 8 512 B (each Light = 64 B) + """ + proj_bytes = proj.T.flatten().astype(np.float32).tobytes() + view_bytes = view.T.flatten().astype(np.float32).tobytes() + + lights = self._collect_lights() + n = min(len(lights), self._MAX_LIGHTS) + + # num_lights (u32) + 3× padding u32 + header = np.array([n, 0, 0, 0], dtype=np.uint32).tobytes() + + # Pack up to MAX_LIGHTS light structs; pad the rest with zeros. + light_data = b"" + for i in range(self._MAX_LIGHTS): + if i < n: + light_data += lights[i].pack() + else: + light_data += b"\x00" * 64 + + return proj_bytes + view_bytes + header + light_data + + # ------------------------------------------------------------------ + # Sub-camera rendering (ZoomedScene / ImageMobjectFromCamera) + # ------------------------------------------------------------------ + + def _sub_camera_proj_view( + self, sub_cam_frame: Any + ) -> tuple[np.ndarray, np.ndarray]: + """Return (proj, view) matrices for a MovingCamera's frame viewport. + + The sub-camera is always orthographic. Its viewport is defined by the + ``frame`` mobject's current center and size. + """ + fw = float(sub_cam_frame.get_width()) + fh = float(sub_cam_frame.get_height()) + cen = sub_cam_frame.get_center() + cx, cy = float(cen[0]), float(cen[1]) + near, far = -100.0, 100.0 + + # Orthographic projection using the sub-camera's frame dimensions + # (centered at origin — the view matrix handles the translation). + proj = np.array( + [ + [2.0 / fw, 0.0, 0.0, 0.0], + [0.0, 2.0 / fh, 0.0, 0.0], + [0.0, 0.0, -1.0 / (far - near), far / (far - near)], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + + # View: translate the world so frame center lands at the origin. + view = np.eye(4, dtype=np.float32) + view[0, 3] = -cx + view[1, 3] = -cy + view[2, 3] = -float(self.camera.focal_distance) + return proj, view + + def _get_sub_cam_resources(self, mob: Any) -> dict: + """Return (and lazily create) per-mob GPU resources for sub-camera rendering. + + Returns a dict with keys: + render_tex, render_view — rgba8unorm render target + (RENDER_ATTACHMENT | TEXTURE_BINDING | COPY_SRC) + depth_tex, depth_view — depth24plus depth buffer + uniform_buf — 656-byte camera uniform buffer (COPY_DST | UNIFORM) + cam_bg — camera-only bind group (proj_bgl, binding 0 = uniform_buf) + tex_bg — image display bind group (image_tex_bgl, binding 0 = render_view) + staging_buf — row-aligned COPY_DST | MAP_READ buffer for CPU readback + staging_aligned_bpr — aligned bytes-per-row used by the staging buffer + """ + assert self._device is not None + assert self._proj_bgl is not None + assert self._image_tex_bgl is not None + + mob_id = id(mob) + if mob_id in self._sub_cam_resources: + return self._sub_cam_resources[mob_id] + + w, h = config.pixel_width, config.pixel_height + _UBO_SIZE = 656 + + # bytes_per_row must be a multiple of 256 for copy_texture_to_buffer. + aligned_bpr = ((w * 4) + 255) & ~255 + + render_tex = self._device.create_texture( + size=(w, h, 1), + format=wgpu.TextureFormat.rgba8unorm, + usage=( + wgpu.TextureUsage.RENDER_ATTACHMENT + | wgpu.TextureUsage.TEXTURE_BINDING + | wgpu.TextureUsage.COPY_SRC + ), + ) + render_view = render_tex.create_view() + + depth_tex = self._device.create_texture( + size=(w, h, 1), + format=wgpu.TextureFormat.depth24plus, + usage=wgpu.TextureUsage.RENDER_ATTACHMENT, + ) + depth_view = depth_tex.create_view() + + uniform_buf = self._device.create_buffer( + size=_UBO_SIZE, + usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST, + ) + cam_bg = self._device.create_bind_group( + layout=self._proj_bgl, + entries=[ + { + "binding": 0, + "resource": {"buffer": uniform_buf, "offset": 0, "size": _UBO_SIZE}, + } + ], + ) + + sampler = self._device.create_sampler( + min_filter="linear", + mag_filter="linear", + address_mode_u="clamp-to-edge", + address_mode_v="clamp-to-edge", + ) + tex_bg = self._device.create_bind_group( + layout=self._image_tex_bgl, + entries=[ + {"binding": 0, "resource": render_view}, + {"binding": 1, "resource": sampler}, + ], + ) + + # Staging buffer for CPU readback of the rendered sub-camera texture. + # rgba8unorm — no B↔R swap needed (unlike the main bgra8unorm target). + staging_buf = self._device.create_buffer( + size=aligned_bpr * h, + usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ, + ) + + resources = { + "render_tex": render_tex, + "render_view": render_view, + "depth_tex": depth_tex, + "depth_view": depth_view, + "uniform_buf": uniform_buf, + "cam_bg": cam_bg, + "tex_bg": tex_bg, + "staging_buf": staging_buf, + "staging_aligned_bpr": aligned_bpr, + } + self._sub_cam_resources[mob_id] = resources + return resources + + def _render_sub_camera_pass( + self, + mob: Any, + encoder: Any, + normal_fds: list, + ) -> None: + """Render a sub-camera view for *mob* (an ImageMobjectFromCamera). + + The sub-camera's view of the main scene is drawn into the mob's + persistent rgba8unorm render texture using sub-camera pipelines. + The result is available as ``resources["tex_bg"]`` in the same frame. + + Parameters + ---------- + mob + An ``ImageMobjectFromCamera`` instance with a ``camera`` attribute + that is a ``MovingCamera``. + encoder + Active ``GPUCommandEncoder`` (compute pass must already be ended). + normal_fds + List of ``_FrameData`` objects from the main frame's normal-camera + render queue. The same GPU geometry buffers are reused here with + a different camera uniform. + """ + assert self._device is not None + assert self._fill_stroke_bgl is not None + assert self._sub_cam_fill_stroke_pipeline is not None + assert self._sub_cam_fill_stroke_3d_pipeline is not None + assert self._sub_cam_surface_pipeline is not None + + sub_cam_frame = mob.camera.frame + proj, view = self._sub_camera_proj_view(sub_cam_frame) + ubo_bytes = self._pack_camera_uniforms_bytes(proj, view) + + res = self._get_sub_cam_resources(mob) + self._device.queue.write_buffer(res["uniform_buf"], 0, ubo_bytes) + + bg = self._background_color + + sub_pass = encoder.begin_render_pass( + color_attachments=[ + { + "view": res["render_view"], + "load_op": "clear", + "store_op": "store", + "clear_value": tuple(float(c) for c in bg), + } + ], + depth_stencil_attachment={ + "view": res["depth_view"], + "depth_clear_value": 1.0, + "depth_load_op": "clear", + "depth_store_op": "store", + }, + ) + + for fd in normal_fds: + if fd is None: + continue + # Build a sub-camera bind group that combines the sub-camera + # uniform buffer (binding 0) with the same quads output buffer + # (binding 1) used in the main render. This reuses the already- + # computed quadratic Bezier data without re-running the compute shader. + if fd.quads_out_buf is not None: + sub_fill_render_bg = self._device.create_bind_group( + layout=self._fill_stroke_bgl, + entries=[ + { + "binding": 0, + "resource": { + "buffer": res["uniform_buf"], + "offset": 0, + "size": res["uniform_buf"].size, + }, + }, + { + "binding": 1, + "resource": { + "buffer": fd.quads_out_buf, + "offset": 0, + "size": fd.quads_out_buf.size, + }, + }, + ], + ) + else: + sub_fill_render_bg = None + + draw_frame_data_subcam( + sub_pass, + fd, + sub_fill_render_bg=sub_fill_render_bg, + sub_cam_bg=res["cam_bg"], + fill_2d_pipeline=self._sub_cam_fill_stroke_pipeline, + fill_3d_pipeline=self._sub_cam_fill_stroke_3d_pipeline, + surf_pipeline=self._sub_cam_surface_pipeline, + ) + + sub_pass.end() + + # Keep the old name as a shim so any external callers don't break. + def _pack_camera_uniforms( + self, proj: np.ndarray, view: np.ndarray + ) -> wgpu_t.GPUBuffer: + """Create a throw-away 656-byte uniform buffer (legacy path, rarely used).""" + assert self._device is not None + buf = self._device.create_buffer_with_data( + data=self._pack_camera_uniforms_bytes(proj, view), + usage=wgpu.BufferUsage.UNIFORM, + ) + self.frame_vbos.append(buf) + return buf + + def _build_camera_bind_group(self) -> wgpu_t.GPUBindGroup: + """Update all three persistent camera uniform buffers for the current frame. + + The three uniform buffers and their bind groups are created once in + init_scene. Each frame we write fresh matrix data into the buffers via + queue.write_buffer so the shaders see the updated camera. + + normal (camera_bind_group) + Full camera rotation + current projection. Used for all regular + mobjects. + + fixed_camera_bind_group + Rotation-stripped view + current projection. Used for + fixed-orientation mobjects: they don't tilt with the camera but + are still depth-sorted with the rest of the scene. + + fixed_frame_bind_group + Rotation-stripped view + forced orthographic projection. Used for + fixed-in-frame mobjects: 2-D overlays rendered after the 3-D scene + with a fresh depth buffer so they always appear on top. + """ + assert self._device is not None + assert self._camera_uniform_buf is not None + assert self._fixed_orient_uniform_buf is not None + assert self._fixed_frame_uniform_buf is not None + + fixed_view = self.camera.fixed_view_matrix + + self._device.queue.write_buffer( + self._camera_uniform_buf, + 0, + self._pack_camera_uniforms_bytes( + self.camera.projection_matrix, self.camera.view_matrix + ), + ) + self._device.queue.write_buffer( + self._fixed_orient_uniform_buf, + 0, + self._pack_camera_uniforms_bytes(self.camera.projection_matrix, fixed_view), + ) + self._device.queue.write_buffer( + self._fixed_frame_uniform_buf, + 0, + self._pack_camera_uniforms_bytes( + self.camera.ortho_projection_matrix, fixed_view + ), + ) + + # Return the persistent normal bind group (unchanged object). + return self.camera_bind_group + + # ------------------------------------------------------------------ + # Pipeline / device accessors (used by webgpu_vmobject_rendering) + # ------------------------------------------------------------------ + + @property + def device(self) -> wgpu_t.GPUDevice: + assert self._device is not None, "init_scene() has not been called" + return self._device + + @property + def fill_stroke_pipeline(self) -> wgpu_t.GPURenderPipeline: + """Combined fill+stroke pipeline — 2-D (no depth write).""" + assert self._fill_stroke_pipeline is not None, ( + "init_scene() has not been called" + ) + return self._fill_stroke_pipeline + + @property + def fill_stroke_3d_pipeline(self) -> wgpu_t.GPURenderPipeline: + """Combined fill+stroke pipeline — 3-D (depth write + test).""" + assert self._fill_stroke_3d_pipeline is not None, ( + "init_scene() has not been called" + ) + return self._fill_stroke_3d_pipeline + + @property + def fill_stroke_pipeline_1x(self) -> wgpu_t.GPURenderPipeline: + """Combined fill+stroke pipeline — 2-D, always count=1. + + Used for the fixed-in-frame overlay pass (Pass 4) which renders directly + into ``_render_texture_view`` after the MSAA resolve has completed. + """ + assert self._fill_stroke_pipeline_1x is not None, ( + "init_scene() has not been called" + ) + return self._fill_stroke_pipeline_1x + + @property + def fill_stroke_3d_pipeline_1x(self) -> wgpu_t.GPURenderPipeline: + """Combined fill+stroke pipeline — 3-D, always count=1. + + Used for the fixed-in-frame overlay pass (Pass 4) which renders directly + into ``_render_texture_view`` after the MSAA resolve has completed. + """ + assert self._fill_stroke_3d_pipeline_1x is not None, ( + "init_scene() has not been called" + ) + return self._fill_stroke_3d_pipeline_1x + + @property + def surface_pipeline(self) -> wgpu_t.GPURenderPipeline: + """Opaque surface pipeline (depth_write=True).""" + assert self._surface_pipeline is not None, "init_scene() has not been called" + return self._surface_pipeline + + @property + def surface_oit_pipeline(self) -> wgpu_t.GPURenderPipeline: + assert self._surface_oit_pipeline is not None, ( + "init_scene() has not been called" + ) + return self._surface_oit_pipeline + + # ------------------------------------------------------------------ + # Frame rendering + # ------------------------------------------------------------------ + + def update_frame( + self, + scene: Scene, + mob_list: list | None = None, + blit_static: bool = False, + _readback: bool = True, + ) -> None: + """Render one frame into the offscreen texture. + + Parameters + ---------- + mob_list: + When provided, render only these top-level mobjects instead of + all of ``scene.mobjects``. Used by ``save_static_frame_data`` + (static subset) and by ``render`` (moving subset). + blit_static: + When True, blit ``_static_texture`` → ``_render_texture`` before + the main render pass so the static background is preserved. The + main pass then uses ``load_op="load"`` to composite moving mobs + on top. If False (the default), the frame is cleared to the + background colour first. + + Pass structure + -------------- + 0. **Texture blit** (when *blit_static*) — copies the pre-rendered + static layer into the render texture before any render passes. + 1. **Compute pass** — cubic_to_quads.wgsl converts raw cubic Bezier + control points to quadratic approximations for all three mobject + groups (normal, fixed-orientation, fixed-in-frame). This runs + before any render pass in the same command encoder, so WebGPU's + implicit pass ordering provides the barrier. + 2. **Main pass** — clears (or loads) the frame; draws normal and + fixed-orientation mobjects (shared depth buffer; fixed-orient uses + a rotation-stripped camera bind group). + 3. **OIT accumulation pass** — transparent surfaces use Weighted + Blended OIT into two rgba16float textures. + 4. **OIT composition pass** — full-screen triangle composites the OIT + result onto the main texture. + 5. **Fixed-in-frame overlay pass** — 2-D overlays rendered with a + fresh depth buffer so they always appear on top. + """ + assert self._device is not None + assert self._render_texture_view is not None + assert self._depth_texture_view is not None + assert self._cubic_to_quads_pipeline is not None + + bg = self._background_color + + # Build all three per-frame camera uniform buffers + bind groups. + self.camera_bind_group = self._build_camera_bind_group() + self.frame_vbos = [] + + # ── Partition and z-sort mobjects ──────────────────────────────── + cam = self.camera + fixed_in_frame = cam.fixed_in_frame_mobjects + fixed_orient = cam.fixed_orientation_mobjects + fixed_view = self.camera.fixed_view_matrix + + assert self._camera_uniform_buf is not None + assert self._fixed_orient_uniform_buf is not None + assert self._fixed_frame_uniform_buf is not None + + if mob_list is not None: + # Caller (save_static_frame_data, render) already sorted the list. + source = mob_list + else: + # Full-frame path: merge mobjects + foreground_mobjects and apply + # z_index ordering (Bug 1). foreground_mobjects are already present + # in scene.mobjects (add_foreground_mobjects calls add()), so + # list_update just removes the duplicates from the left side. + # We then sort with a two-key tuple so that: + # key[0] = 0 for normal mobs, 1 for foreground mobs + # key[1] = z_index + # This ensures foreground mobs always draw last (on top) even when + # they share z_index=0 with regular mobs (Bug 3). + all_mobs = list_update( + list(scene.mobjects), list(scene.foreground_mobjects) + ) + if self.camera.use_z_index: + foreground_ids = {id(m) for m in scene.foreground_mobjects} + source = sorted( + all_mobs, + key=lambda m: (1 if id(m) in foreground_ids else 0, m.z_index), + ) + else: + source = all_mobs + + # Build a z-ordered render queue by walking `source` in order. + # + # Rules: + # • fixed_in_frame mobs → skipped here, collected separately below + # • fixed_orient mobs → VMobject batch with stripped-rotation camera + # • normal VMobjects → VMobject batch with full camera + # • ImageMobjects → image draw item, flushing any pending + # VMobject runs first so z-order is respected + # • containers (Group…) → recursed + # + # The resulting queue is a list of items: + # ('vmobs', _FrameData, camera_bind_group) + # ('image', ImageMobject) + # + # Within the main render pass these are drawn in queue order, giving + # correct painter's-algorithm depth for any interleaving of images and + # VMobjects in scene.mobjects. + + render_queue: list[tuple] = [] + _run_normal: list = [] + _run_orient: list = [] + _seen: set[int] = set() + + def _flush_runs() -> None: + if _run_normal: + fd = collect_frame_data( + self, + list(_run_normal), + self._camera_uniform_buf, + cache_slot="normal", + ) + if fd is not None: + render_queue.append(("vmobs", fd, self.camera_bind_group)) + _run_normal.clear() + if _run_orient: + fd = collect_frame_data( + self, + list(_run_orient), + self._fixed_orient_uniform_buf, + view_matrix_override=fixed_view, + center_view_matrix=self.camera.view_matrix, + cache_slot="orient", + ) + if fd is not None: + render_queue.append(("vmobs", fd, self.fixed_camera_bind_group)) + _run_orient.clear() + + def _walk(mob: Any) -> None: + if id(mob) in _seen: + return + _seen.add(id(mob)) + if isinstance(mob, AbstractImageMobject): + _flush_runs() + render_queue.append(("image", mob)) + # Recurse into submobjects (e.g. the display_frame SurroundingRectangle + # added by ImageMobjectFromCamera.add_display_frame()) so they are + # rendered as VMobject overlays on top of the image quad. + for sub in mob.submobjects: + _walk(sub) + elif isinstance(mob, DotCloud3D): + # WebGPU dot cloud — rendered as screen-aligned sphere quads. + _flush_runs() + render_queue.append(("truedot", mob)) + elif isinstance(mob, Surface): + # Parametric Surface: add individually so collect_frame_data + # uses the Surface lighting/geometry cache path. + if mob in fixed_in_frame: + pass # handled in overlay pass below + elif mob in fixed_orient: + _run_orient.append(mob) + else: + _run_normal.append(mob) + elif isinstance(mob, VMobject): + # Container check: if any *direct* child is a Surface, Image, or + # DotCloud, recurse rather than treating this mob as a monolithic + # VMobject. This ensures e.g. VGroup(Sphere(), Sphere()) routes + # each Sphere through the surface rendering path. + if mob.submobjects and any( + isinstance(s, (Surface, AbstractImageMobject, DotCloud3D)) + for s in mob.submobjects + ): + for sub in mob.submobjects: + _walk(sub) + elif mob in fixed_in_frame: + pass # handled in overlay pass below + elif mob in fixed_orient: + _run_orient.append(mob) + else: + _run_normal.append(mob) + else: + for sub in mob.submobjects: + _walk(sub) + + for mob in source: + _walk(mob) + _flush_runs() + + # Pre-fetch image GPU resources (texture upload, VBO) before the + # command encoder starts. Replace ('image', mob) queue items with + # ('image', vbo, tex_bg) so the render loop has no CPU work left. + # Similarly expand TrueDot mobs into vertex arrays and GPU buffers. + # + # ImageMobjectFromCamera mobs are handled specially: instead of reading + # their pixel_array (which is never populated by the WebGPU renderer), + # we use the pre-rendered sub-camera texture produced in Pass 0.5. + resolved_queue: list[tuple] = [] + for item in render_queue: + if item[0] == "image": + mob = item[1] + vbo = self._get_image_vbo(mob) + tint_bg = self._get_image_tint_bind_group(mob) + if isinstance(mob, ImageMobjectFromCamera): + res = self._sub_cam_resources.get(id(mob)) + tex_bg = res["tex_bg"] if res is not None else None + if vbo is not None and tex_bg is not None and tint_bg is not None: + resolved_queue.append(("image", vbo, tex_bg, tint_bg)) + else: + resources = self._get_image_gpu_resources(mob) + if ( + vbo is not None + and resources is not None + and tint_bg is not None + ): + resolved_queue.append(("image", vbo, resources[1], tint_bg)) + elif item[0] == "truedot": + mob = item[1] + arr = build_true_dot_vbo(mob) + if arr is not None and len(arr) > 0: + buf = self._device.create_buffer_with_data( + data=arr.tobytes(), + usage=wgpu.BufferUsage.VERTEX, + ) + self.frame_vbos.append(buf) + resolved_queue.append(("truedot", buf, len(arr))) + else: + resolved_queue.append(item) + + # Fixed-in-frame: always last, separate overlay pass. + fixed_frame_mobs = [ + m + for m in _seen + if False # placeholder — rebuilt below from source flatten + ] + + # Re-flatten source to get all VMobjects (including those inside containers) + # and filter to the fixed_in_frame set. + def _flatten_vmobjects(src: list) -> list: + out: list = [] + seen2: set[int] = set() + + def _f(m: Any) -> None: + if id(m) in seen2: + return + seen2.add(id(m)) + if isinstance(m, VMobject): + out.append(m) + else: + for s in m.submobjects: + _f(s) + + for m in src: + _f(m) + return out + + fixed_frame_mobs = [ + m for m in _flatten_vmobjects(source) if m in fixed_in_frame + ] + fixed_frame_fd = collect_frame_data( + self, + fixed_frame_mobs, + self._fixed_frame_uniform_buf, + view_matrix_override=fixed_view, + proj_matrix_override=self.camera.ortho_projection_matrix, + cache_slot="frame", + ) + + # OIT surfaces come from all normal VMobject batches in the queue. + all_normal_fds = [ + item[1] + for item in resolved_queue + if item[0] == "vmobs" and item[2] is self.camera_bind_group + ] + + encoder = self._device.create_command_encoder() + + # ── Pre-pass: blit static background ───────────────────────────── + # When compositing moving mobs on top of the pre-rendered static layer, + # copy the static texture into the render texture before any render + # passes. The subsequent main pass uses load_op="load" so the static + # pixels are preserved under the newly drawn moving mobs. + # + # With MSAA the static optimisation is bypassed: the MSAA texture is an + # intermediate buffer that always resolves into _render_texture at the + # end of Pass 1, overwriting whatever was there. Every frame is fully + # re-rendered instead. (MSAA already implies a quality-over-speed + # trade-off, so the extra work per frame is acceptable.) + _do_static_blit = ( + blit_static + and self._has_static_frame + and self._static_texture is not None + and self._msaa_samples == 1 + ) + if _do_static_blit: + encoder.copy_texture_to_texture( + {"texture": self._static_texture, "mip_level": 0, "origin": (0, 0, 0)}, + {"texture": self._render_texture, "mip_level": 0, "origin": (0, 0, 0)}, + (config.pixel_width, config.pixel_height, 1), + ) + + # ── Pass 0: compute — cubic → quadratic conversion ──────────────── + # Runs before any render pass; WebGPU guarantees the output buffer is + # ready by the time the fragment shader reads it in Pass 1. + cp = encoder.begin_compute_pass() + cp.set_pipeline(self._cubic_to_quads_pipeline) + all_fds = [item[1] for item in resolved_queue if item[0] == "vmobs"] + ( + [fixed_frame_fd] if fixed_frame_fd is not None else [] + ) + for fd in all_fds: + if fd.n_cubics_total > 0 and fd.compute_bg is not None: + cp.set_bind_group(0, fd.compute_bg, [], 0, 0) + cp.dispatch_workgroups((fd.n_cubics_total + 63) // 64, 1, 1) + cp.end() + + # ── Pass 0.5: sub-camera render passes ─────────────────────────── + # Render scene geometry into each ImageMobjectFromCamera's private + # rgba8unorm texture so it is ready to be sampled as an image in + # Pass 1. These passes share the already-computed quads buffers + # from the compute pass above; no re-dispatch is needed. + if self.camera.image_mobjects_from_cameras: + normal_fds = [item[1] for item in resolved_queue if item[0] == "vmobs"] + for sub_mob in self.camera.image_mobjects_from_cameras: + self._get_sub_cam_resources(sub_mob) # ensure resources created + self._render_sub_camera_pass(sub_mob, encoder, normal_fds) + + # Encode texture → staging-buffer copies so that after submit the + # rendered sub-camera pixels are available for CPU readback. + # This allows ImageMobjectFromCamera.get_pixel_array() to return + # current frame data instead of the stale initial pixel_array. + w, h = config.pixel_width, config.pixel_height + for sub_mob in self.camera.image_mobjects_from_cameras: + res = self._sub_cam_resources.get(id(sub_mob)) + if res is None: + continue + aligned_bpr = res["staging_aligned_bpr"] + encoder.copy_texture_to_buffer( + {"texture": res["render_tex"], "mip_level": 0, "origin": (0, 0, 0)}, + { + "buffer": res["staging_buf"], + "offset": 0, + "bytes_per_row": aligned_bpr, + "rows_per_image": h, + }, + (w, h, 1), + ) + + # ── Pass 1: main render ─────────────────────────────────────────── + # Draw the z-ordered render queue (VMobject batches and images + # interleaved in scene.mobjects order) so painter's-algorithm depth + # is respected for any combination of images and geometry. + # + # MSAA path (msaa_samples > 1): + # color view — _msaa_texture_view (intermediate MSAA buffer) + # resolve_target — _render_texture_view (receives the resolved pixels) + # store_op — "discard" for the MSAA buffer (transient, never read back) + # depth view — _msaa_depth_texture_view (sample_count must match pipeline) + # depth store — "discard" (MSAA depth is not sampled later; OIT uses + # _depth_texture_view at count=1 separately) + # + # Non-MSAA path (msaa_samples == 1): + # color view — _render_texture_view directly + # depth view — _depth_texture_view (reused by OIT pass below) + if self._msaa_samples > 1: + assert self._msaa_texture_view is not None + assert self._msaa_depth_texture_view is not None + _color_attachment = { + "view": self._msaa_texture_view, + "resolve_target": self._render_texture_view, + "load_op": "clear", + "store_op": "discard", + "clear_value": tuple(float(c) for c in bg), + } + _depth_attachment = { + "view": self._msaa_depth_texture_view, + "depth_clear_value": 1.0, + "depth_load_op": "clear", + "depth_store_op": "discard", + } + else: + color_load_op = "load" if _do_static_blit else "clear" + _color_attachment = { + "view": self._render_texture_view, + "load_op": color_load_op, + "store_op": "store", + "clear_value": tuple(float(c) for c in bg), + } + _depth_attachment = { + "view": self._depth_texture_view, + "depth_clear_value": 1.0, + "depth_load_op": "clear", + "depth_store_op": "store", + } + + main_pass = encoder.begin_render_pass( + color_attachments=[_color_attachment], + depth_stencil_attachment=_depth_attachment, + ) + self.current_render_pass = main_pass + + current_pipeline = None + for item in resolved_queue: + if item[0] == "image": + _, vbo, tex_bg, tint_bg = item + if current_pipeline != "image": + main_pass.set_pipeline(self._image_pipeline) + main_pass.set_bind_group(0, self.camera_bind_group, [], 0, 0) + current_pipeline = "image" + main_pass.set_bind_group(1, tex_bg, [], 0, 0) + main_pass.set_bind_group(2, tint_bg, [], 0, 0) + main_pass.set_vertex_buffer(0, vbo) + main_pass.draw(6) + elif item[0] == "truedot": + _, buf, n_verts = item + if current_pipeline != "truedot": + main_pass.set_pipeline(self._true_dot_pipeline) + main_pass.set_bind_group(0, self.camera_bind_group, [], 0, 0) + current_pipeline = "truedot" + main_pass.set_vertex_buffer(0, buf) + main_pass.draw(n_verts) + elif item[0] == "vmobs": + _, fd, cam_bg = item + current_pipeline = "vmobs" # draw_frame_data sets its own pipelines + draw_frame_data(self, fd, cam_bg) + + main_pass.end() + + # ── Pass 2: OIT accumulation ────────────────────────────────────── + # Collect OIT surfaces from all normal-camera VMobject batches. + oit_fds = [fd for fd in all_normal_fds if fd.oit_indices] + if oit_fds: + # Use the first fd with OIT surfaces as representative; actual + # OIT draw loops over all of them below. + oit_fd = oit_fds[0] + else: + oit_fd = None + if oit_fds: + # When MSAA is enabled the opaque depth lives in _msaa_depth_texture + # (sample_count=N). _depth_texture (count=1) was not written in + # Pass 1, so OIT transparent surfaces cannot depth-test against + # opaque geometry — use "clear" to avoid reading stale data. + # In practice this means transparent surfaces are not clipped by + # opaque geometry when MSAA is on, which is acceptable for the + # typical use case (transparent surfaces in 3-D scenes). + oit_depth_load_op = "clear" if self._msaa_samples > 1 else "load" + oit_pass = encoder.begin_render_pass( + color_attachments=[ + { + "view": self._oit_accum_view, + "load_op": "clear", + "store_op": "store", + "clear_value": (0.0, 0.0, 0.0, 0.0), + }, + { + "view": self._oit_reveal_view, + "load_op": "clear", + "store_op": "store", + "clear_value": (1.0, 1.0, 1.0, 1.0), + }, + ], + depth_stencil_attachment={ + "view": self._depth_texture_view, + "depth_load_op": oit_depth_load_op, + "depth_store_op": "discard", + }, + ) + oit_pass.set_pipeline(self.surface_oit_pipeline) + oit_pass.set_bind_group(0, self.camera_bind_group, [], 0, 0) + for oit_fd in oit_fds: + for idx in oit_fd.oit_indices: + arr = oit_fd.surface_parts[idx] + oit_pass.set_vertex_buffer( + 0, + oit_fd.surface_buf, + oit_fd.surface_byte_offsets[idx], + arr.nbytes, + ) + oit_pass.draw(len(arr), 1, 0, 0) + oit_pass.end() + + # ── Pass 3: OIT composition ────────────────────────────────── + compose_pass = encoder.begin_render_pass( + color_attachments=[ + { + "view": self._render_texture_view, + "load_op": "load", + "store_op": "store", + } + ], + ) + compose_pass.set_pipeline(self._oit_compose_pipeline) + compose_pass.set_bind_group(0, self._oit_compose_bind_group, [], 0, 0) + compose_pass.draw(3, 1, 0, 0) + compose_pass.end() + + # ── Pass 4: fixed-in-frame overlay ─────────────────────────────── + # Rendered after OIT so overlays always appear on top of the 3-D scene. + # Fresh depth buffer: overlays only depth-test against each other. + # + # This pass always renders directly into _render_texture_view at + # sample_count=1, regardless of the MSAA setting. The MSAA resolve + # (end of Pass 1) has already completed, so _render_texture contains + # the full scene. Fixed-in-frame mobjects are 2-D overlays whose SDF + # anti-aliasing is already excellent; MSAA adds nothing here. + # + # When MSAA is enabled the fill+stroke pipelines have multisample + # count=N and cannot be used in a count=1 pass. We temporarily swap + # in the _1x pipeline variants so that draw_frame_data issues + # compatible GPU commands. + if fixed_frame_fd is not None: + fixed_pass = encoder.begin_render_pass( + color_attachments=[ + { + "view": self._render_texture_view, + "load_op": "load", + "store_op": "store", + } + ], + depth_stencil_attachment={ + "view": self._depth_texture_view, + "depth_clear_value": 1.0, + "depth_load_op": "clear", + "depth_store_op": "discard", + }, + ) + self.current_render_pass = fixed_pass + if self._msaa_samples > 1: + # Temporarily expose count=1 pipelines so draw_frame_data + # records compatible draw calls for this count=1 pass. + _saved_2d = self._fill_stroke_pipeline + _saved_3d = self._fill_stroke_3d_pipeline + self._fill_stroke_pipeline = self._fill_stroke_pipeline_1x + self._fill_stroke_3d_pipeline = self._fill_stroke_3d_pipeline_1x + draw_frame_data(self, fixed_frame_fd, self.fixed_frame_bind_group) + if self._msaa_samples > 1: + self._fill_stroke_pipeline = _saved_2d + self._fill_stroke_3d_pipeline = _saved_3d + fixed_pass.end() + + # ── Invalidate readback cache ──────────────────────────────────── + # A new frame is being rendered; cached pixel data from the previous + # frame is no longer valid. + self._readback_cache = None + + # ── Pool-based pre-staged readback ─────────────────────────────── + # Append the readback compute dispatch + buffer copy into the current + # pool slot, then call map_async immediately after submit. The GPU + # processes the readback asynchronously while the CPU continues. + # When get_image() / get_frame() dequeues the slot, it calls + # sync_wait() — which returns instantly when the pool has been filled + # (the oldest slot has been in-flight for >= _READBACK_POOL frames). + # + # Guard: only pre-stage if the pool has a free slot. A slot is free + # once it has been read (unmap called) by _get_mapped_frame_array. + # If all slots are still in-flight we skip pre-staging for this frame; + # _get_mapped_frame_array falls back to a fresh submit+map_sync. + # + # _readback=False skips readback staging entirely (used by + # save_static_frame_data, which renders for texture capture only and + # must not pollute the pool with non-movie frames). + if ( + _readback + and self._readback_compute_pipeline is not None + and self._readback_compute_bind_group is not None + and self._readback_storage_buf is not None + and self._readback_pool + and len(self._readback_queue) < self._READBACK_POOL + ): + width = config.pixel_width + height = config.pixel_height + packed_size = width * height * 4 + slot = self._readback_write_slot + + cp = encoder.begin_compute_pass() + cp.set_pipeline(self._readback_compute_pipeline) + cp.set_bind_group(0, self._readback_compute_bind_group) + cp.dispatch_workgroups((width + 15) // 16, (height + 15) // 16) + cp.end() + + encoder.copy_buffer_to_buffer( + self._readback_storage_buf, + 0, + self._readback_pool[slot], + 0, + packed_size, + ) + + self._device.queue.submit([encoder.finish()]) + + self._readback_queue.append(slot) + self._readback_write_slot = (slot + 1) % self._READBACK_POOL + else: + self._device.queue.submit([encoder.finish()]) + + # ── Sub-camera CPU readback ─────────────────────────────────────── + # Populate mob.camera.pixel_array from the staged sub-camera texture + # data so that ImageMobjectFromCamera.get_pixel_array() returns the + # current frame rather than the stale initial array. + # rgba8unorm needs no B↔R channel swap (unlike the main bgra8unorm target). + if self.camera.image_mobjects_from_cameras: + w, h = config.pixel_width, config.pixel_height + for sub_mob in self.camera.image_mobjects_from_cameras: + res = self._sub_cam_resources.get(id(sub_mob)) + if res is None: + continue + staging_buf = res["staging_buf"] + aligned_bpr = res["staging_aligned_bpr"] + staging_buf.map_sync(wgpu.MapMode.READ) + raw = bytes(staging_buf.read_mapped()) + staging_buf.unmap() + if aligned_bpr == w * 4: + arr = np.frombuffer(raw, dtype=np.uint8).reshape(h, w, 4).copy() + else: + # Strip row padding before reshaping. + rows = [ + raw[r * aligned_bpr : r * aligned_bpr + w * 4] for r in range(h) + ] + arr = ( + np.frombuffer(b"".join(rows), dtype=np.uint8) + .reshape(h, w, 4) + .copy() + ) + # Write into the Cairo sub-camera's pixel_array so that + # ImageMobjectFromCamera.get_pixel_array() returns current data. + try: + sub_mob.camera.pixel_array = arr + except Exception: + pass + + self.current_render_pass = None + # camera_bind_group is now persistent (created once in init_scene) — + # do NOT null it here. + self.frame_vbos = [] + + self.animation_elapsed_time = time.time() - self.animation_start_time + + # ------------------------------------------------------------------ + # Frame readback + # ------------------------------------------------------------------ + + def _create_readback_pipeline(self, width: int, height: int) -> None: + """Create the compact-readback compute pipeline and its persistent buffers. + + The compute shader (readback_compact.wgsl) reads every pixel from the + bgra8unorm render texture and writes tightly-packed RGBA u32 values into + a storage buffer — eliminating two CPU operations that previously ran on + every frame: + + * Row-padding strip — copy_texture_to_buffer requires bytes_per_row + to be a multiple of 256; the shader writes directly to tight index + ``y * width + x``, so the output is already compact. + + * B↔R channel swap — textureLoad() returns components as (r, g, b, a) + regardless of the bgra physical layout, so the output is already in + RGBA byte order. + """ + assert self._device is not None + assert self._render_texture_view is not None + + shader_path = Path(__file__).parent / "shaders" / "readback_compact.wgsl" + shader = self._device.create_shader_module( + code=shader_path.read_text(encoding="utf-8") + ) + + self._readback_compute_bgl = self._device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.COMPUTE, + "texture": { + "sample_type": "float", + "view_dimension": "2d", + "multisampled": False, + }, + }, + { + "binding": 1, + "visibility": wgpu.ShaderStage.COMPUTE, + "buffer": {"type": "storage"}, + }, + ] + ) + + self._readback_compute_pipeline = self._device.create_compute_pipeline( + layout=self._device.create_pipeline_layout( + bind_group_layouts=[self._readback_compute_bgl] + ), + compute={"module": shader, "entry_point": "main"}, + ) + + packed_size = width * height * 4 + self._readback_storage_buf = self._device.create_buffer( + size=packed_size, + usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC, + ) + # Allocate all pool slots up front — no per-frame allocation ever. + self._readback_pool = [ + self._device.create_buffer( + size=packed_size, + usage=wgpu.BufferUsage.COPY_DST | wgpu.BufferUsage.MAP_READ, + ) + for _ in range(self._READBACK_POOL) + ] + self._readback_compute_bind_group = self._device.create_bind_group( + layout=self._readback_compute_bgl, + entries=[ + {"binding": 0, "resource": self._render_texture_view}, + { + "binding": 1, + "resource": { + "buffer": self._readback_storage_buf, + "offset": 0, + "size": packed_size, + }, + }, + ], + ) + + def _get_mapped_frame_array(self) -> np.ndarray: + """Return the current frame as a flat uint8 numpy array (H*W*4 elements). + + Cache path (fastest): update_frame() has not been called since the last + readback, so the frame is unchanged — return the cached array directly + with zero GPU interaction. + + Pool path: update_frame() enqueued a slot index into _readback_queue. + We dequeue the oldest slot and call map_sync(). When the pool has + been filled (_READBACK_POOL frames submitted before the first read), + the GPU work is already complete; map_sync() only pays the driver's + buffer-mapping overhead (< 1 ms) rather than a full GPU pipeline + stall. + + Fallback path: nothing is in the queue (get_image called without a + preceding update_frame, or the pool was full at submit time). We + submit fresh readback work to the next pool slot and block with + map_sync — same behaviour as before, but still using pool buffers so + no new GPU allocation occurs. + """ + assert self._device is not None + assert self._readback_compute_pipeline is not None + assert self._readback_compute_bind_group is not None + assert self._readback_storage_buf is not None + assert self._readback_pool + + # ── Cache hit ──────────────────────────────────────────────────── + if self._readback_cache is not None: + return self._readback_cache + + width = config.pixel_width + height = config.pixel_height + packed_size = width * height * 4 + + if self._readback_queue: + # ── Pool path: dequeue oldest in-flight slot ───────────────── + # The slot was submitted >= _READBACK_POOL frames ago, so the GPU + # has had enough time to finish. map_sync initiates the mapping + # and waits; since the GPU work is already done, only the driver's + # buffer-mapping overhead remains (typically < 1 ms). + slot = self._readback_queue.popleft() + buf = self._readback_pool[slot] + buf.map_sync(wgpu.MapMode.READ) + else: + # ── Fallback path: submit now, block ───────────────────────── + slot = self._readback_write_slot + buf = self._readback_pool[slot] + + encoder = self._device.create_command_encoder() + + cp = encoder.begin_compute_pass() + cp.set_pipeline(self._readback_compute_pipeline) + cp.set_bind_group(0, self._readback_compute_bind_group) + cp.dispatch_workgroups((width + 15) // 16, (height + 15) // 16) + cp.end() + + encoder.copy_buffer_to_buffer( + self._readback_storage_buf, + 0, + buf, + 0, + packed_size, + ) + + self._device.queue.submit([encoder.finish()]) + buf.map_sync(wgpu.MapMode.READ) + self._readback_write_slot = (slot + 1) % self._READBACK_POOL + + # numpy copy is ~3× faster than bytes() for large buffers (SIMD path). + arr = np.frombuffer(buf.read_mapped(), dtype=np.uint8).copy() + buf.unmap() + + self._readback_cache = arr + return arr + + def get_image(self) -> Image.Image: + """Return the current frame as a PIL Image (RGBA).""" + arr = self._get_mapped_frame_array() + return Image.fromarray( + arr.reshape(config.pixel_height, config.pixel_width, 4), "RGBA" + ) + + def get_frame(self) -> np.ndarray: + """Return the current frame as a (height, width, 4) uint8 NumPy array.""" + return self._get_mapped_frame_array().reshape( + config.pixel_height, config.pixel_width, 4 + ) + + # ------------------------------------------------------------------ + # Window helpers + # ------------------------------------------------------------------ + + def should_create_window(self) -> bool: + """Mirror of ``OpenGLRenderer.should_create_window``. + + A preview window is opened when ``--preview`` is active and the + renderer is not writing a movie or saving a still frame. + """ + if config["force_window"]: + logger.warning( + "'--force_window' is enabled; this is intended for debugging " + "and may impact performance when combined with file output.", + ) + return True + return ( + config["preview"] + and not config["save_last_frame"] + and not config["format"] + and not config["write_to_movie"] + and not config["dry_run"] + ) + + def pixel_coords_to_space_coords( + self, + px: float, + py: float, + relative: bool = False, + top_left: bool = False, + ) -> np.ndarray: + """Convert pixel coordinates to Manim scene-space coordinates. + + Parameters + ---------- + px, py: + Pixel position. For ``relative=False``, these are absolute + pixel coordinates within the render texture. + relative: + When True, treat *px*/*py* as a delta and return the + corresponding scene-space delta (normalised to ``[-1, 1]`` + then scaled). + top_left: + When True (the default for ``rendercanvas``), the origin is + at the top-left corner; y increases downward. + """ + pixel_width = config.pixel_width + pixel_height = config.pixel_height + frame_height = config.frame_height + frame_center = self.camera.get_center() + + if relative: + return 2.0 * np.array([px / pixel_width, py / pixel_height, 0.0]) + + scale = frame_height / pixel_height + y_direction = -1 if top_left else 1 + return frame_center + scale * np.array( + [(px - pixel_width / 2), y_direction * (py - pixel_height / 2), 0.0] + ) + + # ------------------------------------------------------------------ + # Scene rendering + # ------------------------------------------------------------------ + + def render(self, scene: Scene, frame_offset: float, moving_mobjects: list) -> None: + if self._has_static_frame: + # Use the family-level moving list produced by begin_animations() + # directly. That list is already z_index-sorted by + # extract_mobject_family_members and is at the correct granularity + # (same as what Cairo passes to its camera). Filtering + # scene.mobjects by static IDs was wrong because it operated at + # top-level container granularity while _static_mob_ids stores + # family-member IDs (Bug 2 fix). + self.update_frame(scene, mob_list=list(moving_mobjects), blit_static=True) + else: + self.update_frame(scene) + if self.skip_animations: + return + self.file_writer.write_frame(self) + if self.window is not None: + if self.window.is_closing: + self.window = None + return + self.window.present() + while self.animation_elapsed_time < frame_offset: + if self.window.is_closing: + self.window = None + break + if self._has_static_frame: + self.update_frame( + scene, mob_list=list(moving_mobjects), blit_static=True + ) + else: + self.update_frame(scene) + self.window.present() + + def play(self, scene: Scene, *animations: Any, **kwargs: Any) -> None: + self.animation_start_time = time.time() + self.skip_animations = self._original_skipping_status + self.update_skipping_status() + + # Compile first so we can compute a real hash (same order as CairoRenderer). + scene.compile_animation_data(*animations, **kwargs) + + if self.skip_animations: + hash_current_animation = None + self.time += scene.duration + elif config["disable_caching"]: + hash_current_animation = f"uncached_{self.num_plays:05}" + else: + assert scene.animations is not None + hash_current_animation = get_hash_from_play_call( + scene, self.camera, scene.animations, scene.mobjects + ) + if self.file_writer.is_already_cached(hash_current_animation): + logger.info( + "Animation %d: using cached data (hash: %s)", + self.num_plays, + hash_current_animation, + ) + self.skip_animations = True + self.time += scene.duration + + self.animations_hashes.append(hash_current_animation) + self.file_writer.add_partial_movie_file(hash_current_animation) + + self.file_writer.begin_animation(not self.skip_animations) + scene.begin_animations() + + # Pre-render static mobjects once, matching Cairo's optimisation. + # scene.static_mobjects is populated by begin_animations() above. + self.save_static_frame_data(scene, scene.static_mobjects) + + if scene.is_current_animation_frozen_frame(): + self.update_frame(scene) + if not self.skip_animations: + self.file_writer.write_frame( + self, num_frames=int(config.frame_rate * scene.duration) + ) + if self.window is not None: + if self.window.is_closing: + self.window = None + else: + self.window.present() + while time.time() - self.animation_start_time < scene.duration: + if self.window.is_closing: + self.window = None + break + self.window.present() + self.animation_elapsed_time = scene.duration + else: + scene.play_internal() + + self.file_writer.end_animation(not self.skip_animations) + self.time += scene.duration + self.num_plays += 1 + + def scene_finished(self, scene: Scene) -> None: + if self.num_plays > 0: + self.file_writer.finish() + elif self.num_plays == 0 and config.write_to_movie: + config.save_last_frame = True + config.write_to_movie = False + + if self._should_save_last_frame(): + config.save_last_frame = True + self.update_frame(scene) + self.file_writer.save_image(self.get_image()) + + # Explicitly close the preview window so that GlfwRenderCanvas._rc_close() + # destroys the GLFW window handle while GLFW is still alive. Without this, + # Python shutdown calls GlfwCanvasGroup.__del__ → glfw.terminate() first, + # then the GlfwRenderCanvas.__del__ fires and tries glfw.destroy_window() + # on an already-terminated GLFW library, producing a GLFWError warning. + if self.window is not None: + try: + self.window.destroy() + except Exception: + pass + self.window = None + + def save_static_frame_data(self, scene: Scene, static_mobjects: Any) -> None: + """Render *static_mobjects* once and cache the result in ``_static_texture``. + + Called by ``play()`` after ``begin_animations()``, before the per-frame + loop starts. Subsequent calls to ``render()`` blit this cached texture + as the background and only re-draw the moving subset, matching Cairo's + static-image compositing optimisation. + + When *static_mobjects* is empty (all mobs are moving, or no mobs at all), + the static frame is cleared so that ``render()`` falls back to a full + redraw each frame. + """ + assert self._device is not None + assert self._static_texture is not None + + static_list = list(static_mobjects) if static_mobjects else [] + + if not static_list: + self._has_static_frame = False + self._static_mob_ids = set() + return + + # If the camera itself is animated (e.g. begin_ambient_camera_rotation), + # any pre-rendered static texture is stale on the very next frame because + # the projection changes. Disable the optimisation so every frame is + # fully re-rendered with the correct camera orientation. + if self.camera.has_time_based_updater(): + self._has_static_frame = False + self._static_mob_ids = set() + return + + self._static_mob_ids = set(id(m) for m in static_list) + + # Render the static mob list into _render_texture (full clear + draw). + # _readback=False: this render is for texture capture only — it must not + # enqueue a readback pool slot, because save_static_frame_data is not + # writing a movie frame and a stale slot in the pool would cause the next + # write_frame() call to read wrong pixel data. + self.update_frame( + scene, mob_list=static_list, blit_static=False, _readback=False + ) + + # Copy _render_texture → _static_texture for later per-frame blits. + encoder = self._device.create_command_encoder() + encoder.copy_texture_to_texture( + {"texture": self._render_texture, "mip_level": 0, "origin": (0, 0, 0)}, + {"texture": self._static_texture, "mip_level": 0, "origin": (0, 0, 0)}, + (config.pixel_width, config.pixel_height, 1), + ) + self._device.queue.submit([encoder.finish()]) + + self._has_static_frame = True + + def clear_screen(self) -> None: + if self.window is not None: + self.window.present() + + # ------------------------------------------------------------------ + # Skipping helpers + # ------------------------------------------------------------------ + + def update_skipping_status(self) -> None: + """Check and update ``skip_animations`` for the current animation. + + Mirrors ``CairoRenderer.update_skipping_status`` and + ``OpenGLRenderer.update_skipping_status`` so the WebGPU renderer + honours the same configuration knobs: + + * ``file_writer.sections[-1].skip_animations`` — section-level skip + (e.g. the section was marked skip via ``scene.next_section``). + * ``config.save_last_frame`` — only the final frame matters; all + intermediate animation frames can be skipped. + * ``config.from_animation_number`` — skip animations before the given + index (useful for scrubbing to a specific animation). + * ``config.upto_animation_number`` — stop rendering after the given + index and raise ``EndSceneEarlyException``. + """ + # there is always at least one section → no out-of-bounds here + if self.file_writer.sections[-1].skip_animations: + self.skip_animations = True + if config["save_last_frame"]: + self.skip_animations = True + if ( + config.from_animation_number > 0 + and self.num_plays < config.from_animation_number + ): + self.skip_animations = True + if ( + config.upto_animation_number >= 0 + and self.num_plays > config.upto_animation_number + ): + self.skip_animations = True + raise EndSceneEarlyException() + + def _should_save_last_frame(self) -> bool: + if config["save_last_frame"]: + return True + if self.scene.interactive_mode: + return False + return self.num_plays == 0 + + # ------------------------------------------------------------------ + # Background colour + # ------------------------------------------------------------------ + + @property + def background_color(self): + return self._background_color + + @background_color.setter + def background_color(self, value) -> None: + self._background_color = color_to_rgba(value, 1.0) + + def get_pixel_shape(self) -> tuple[int, int]: + return (config.pixel_width, config.pixel_height) diff --git a/manim/renderer/webgpu/webgpu_renderer_window.py b/manim/renderer/webgpu/webgpu_renderer_window.py new file mode 100644 index 0000000000..ba6ee7abc7 --- /dev/null +++ b/manim/renderer/webgpu/webgpu_renderer_window.py @@ -0,0 +1,906 @@ +"""Preview window for the WebGPU renderer. + +Uses ``rendercanvas`` to open a native OS window. The offscreen render +texture (``bgra8unorm``) is copied directly to the window surface via +``copy_texture_to_texture`` — no format conversion needed. + +Interactive camera controls +--------------------------- + +======================== ================================================ +Input Action +======================== ================================================ +Left-drag **Orbit** — horizontal drag rotates theta (yaw); + vertical drag tilts phi (pitch), clamped to ±90°. +Right-drag / Middle-drag **Pan** — translates the view laterally in camera + space; one screen-width drag = one frame width. +Scroll wheel **Zoom** — perspective: ``focal_distance`` scales + exponentially (~12 % per notch); orthographic: + ``frame_shape`` scales by the same factor. +Key ``r`` **Reset** — restores default orbit, zoom, and + clears the accumulated pan offset. +Key ``q`` **Quit** — closes the preview window. +======================== ================================================ + +Pan state (``_pan_x``, ``_pan_y``) is stored on :class:`WebGPUWindow` and +injected into the camera just before each render; it is not part of the +scripted camera model. + +Event mapping +------------- +rendercanvas delivers Web-standard key strings (``"ArrowLeft"``, ``"q"`` +…). A mapping table converts these to pyglet-compatible integer codes so +that ``scene.on_key_press`` callbacks work unchanged. +""" + +from __future__ import annotations + +import math +import re +from typing import TYPE_CHECKING + +import numpy as np + +from manim import __version__, config, logger + +if TYPE_CHECKING: + from .webgpu_renderer import WebGPURenderer + +try: + import wgpu + from rendercanvas.auto import RenderCanvas +except ImportError as exc: + raise ImportError( + "wgpu-py (with rendercanvas) is required for the preview window. " + "Install it with: pip install wgpu" + ) from exc + + +# --------------------------------------------------------------------------- +# Key / modifier mapping +# --------------------------------------------------------------------------- + +# rendercanvas key strings → pyglet-compatible integer codes. +# Printable single chars use ord() directly (see _key_to_int below). +_SPECIAL_KEY_MAP: dict[str, int] = { + "ArrowLeft": 65361, + "ArrowRight": 65363, + "ArrowUp": 65362, + "ArrowDown": 65364, + "Escape": 65307, + "Enter": 65293, + "Backspace": 65288, + "Tab": 65289, + "Delete": 65535, + "Home": 65360, + "End": 65367, + "PageUp": 65365, + "PageDown": 65366, + "Insert": 65379, + "F1": 65470, + "F2": 65471, + "F3": 65472, + "F4": 65473, + "F5": 65474, + "F6": 65475, + "F7": 65476, + "F8": 65477, + "F9": 65478, + "F10": 65479, + "F11": 65480, + "F12": 65481, + "Shift": 65505, # SHIFT_VALUE in manim/constants.py + "Control": 65507, + "Alt": 65513, + "Meta": 65511, + "CapsLock": 65509, + "NumLock": 65407, + "ScrollLock": 65300, +} + +# rendercanvas modifier strings → pyglet modifier bitmask bits +_MODIFIER_BITS: dict[str, int] = { + "Shift": 1, + "Control": 4, + "Alt": 8, + "Meta": 16, +} + + +def _key_to_int(key: str) -> int: + """Convert a rendercanvas key string to a pyglet-compatible integer.""" + if key in _SPECIAL_KEY_MAP: + return _SPECIAL_KEY_MAP[key] + if len(key) == 1: + return ord(key) + # Unknown multi-char key — use a stable hash to avoid collisions with + # printable chars. + return (hash(key) & 0x7FFF_FFFF) | 0x8000_0000 + + +def _modifiers_to_int(modifiers: tuple | list) -> int: + """Convert a rendercanvas modifiers collection to a pyglet modifier int.""" + result = 0 + for m in modifiers: + result |= _MODIFIER_BITS.get(m, 0) + return result + + +# --------------------------------------------------------------------------- +# Window configuration helpers +# --------------------------------------------------------------------------- + + +def _compute_window_size() -> tuple[int, int]: + """Return the initial canvas size in logical pixels. + + If ``config.window_size`` is ``"default"`` the canvas is made exactly + ``(pixel_width, pixel_height)`` so the offscreen render texture and the + window surface are always the same size — keeping + ``copy_texture_to_texture`` safe with no extra bookkeeping. + + If the user specified an explicit size (e.g. ``--window_size 960,540``) + that size is used for the *display* window instead. + """ + win_size = config.window_size + if win_size != "default": + return int(win_size[0]), int(win_size[1]) + return config.pixel_width, config.pixel_height + + +def _resolve_window_position( + pos: str, + monitor: object, + win_w: int, + win_h: int, +) -> tuple[int, int]: + """Convert a ``config.window_position`` string to absolute ``(x, y)``. + + Accepts direction strings (``"UL"``, ``"UR"``, ``"DL"``, ``"DR"``, + ``"ORIGIN"``, ``"LEFT"``, ``"RIGHT"``, ``"UP"``, ``"DOWN"``) or a pixel + coordinate pair in ``"x,y"`` / ``"x;y"`` format. + + Parameters + ---------- + pos: + The raw ``config.window_position`` string. + monitor: + A ``screeninfo.Monitor`` (or any object with ``.x``, ``.y``, + ``.width``, ``.height`` attributes). + win_w, win_h: + Current canvas logical width / height in pixels. + """ + mx: int = monitor.x # type: ignore[attr-defined] + my: int = monitor.y # type: ignore[attr-defined] + mw: int = monitor.width # type: ignore[attr-defined] + mh: int = monitor.height # type: ignore[attr-defined] + + # Numeric "x,y" or "x;y" coordinate pair + m = re.match(r"^(\d+)\s*[,;]\s*(\d+)$", pos.strip()) + if m: + return int(m.group(1)), int(m.group(2)) + + pos_u = pos.strip().upper() + right = mx + mw - win_w + bottom = my + mh - win_h + h_center = mx + (mw - win_w) // 2 + v_center = my + (mh - win_h) // 2 + + return { + "UL": (mx, my), + "UR": (right, my), + "DL": (mx, bottom), + "DR": (right, bottom), + "ORIGIN": (h_center, v_center), + "LEFT": (mx, v_center), + "RIGHT": (right, v_center), + "UP": (h_center, my), + "DOWN": (h_center, bottom), + }.get(pos_u, (h_center, v_center)) + + +def _apply_window_config(canvas) -> None: + """Apply all window-related manim config options to a live canvas. + + Handles ``window_size``, ``window_position``, ``window_monitor``, and + ``fullscreen``. Errors are caught and logged as debug messages so that a + missing ``screeninfo`` installation or an unsupported backend never + prevents the window from opening. + """ + # ── Window display size ────────────────────────────────────────────── + win_size = config.window_size + if win_size != "default": + try: + canvas.set_logical_size(float(win_size[0]), float(win_size[1])) + except Exception as exc: + logger.debug("WebGPU: could not set window size: %s", exc) + + # ── Monitor list ───────────────────────────────────────────────────── + try: + import screeninfo + + monitors = screeninfo.get_monitors() + except Exception: + monitors = [] + + mon_idx = int(config.window_monitor) if config.window_monitor is not None else 0 + monitor = None + if monitors: + monitor = monitors[mon_idx] if mon_idx < len(monitors) else monitors[0] + + # ── Backend-specific placement ─────────────────────────────────────── + # Try GLFW backend first (canvas has a raw ``_window`` handle). + glfw_window = getattr(canvas, "_window", None) + if glfw_window is not None: + _apply_glfw_placement(canvas, glfw_window, monitor, mon_idx) + return + + # Try Qt backend (canvas IS a QWidget — move() / showFullScreen() work). + if hasattr(canvas, "showFullScreen") and hasattr(canvas, "move"): + _apply_qt_placement(canvas, monitor) + + +def _apply_glfw_placement(canvas, glfw_window, monitor, mon_idx: int) -> None: + """Apply position / fullscreen via the raw GLFW window handle.""" + try: + import glfw # pyGLFW — installed as a rendercanvas dependency + except ImportError: + logger.debug("WebGPU: glfw not importable; skipping window placement.") + return + + if config.fullscreen: + glfw_monitors = glfw.get_monitors() + if not glfw_monitors: + return + glfw_mon = ( + glfw_monitors[mon_idx] if mon_idx < len(glfw_monitors) else glfw_monitors[0] + ) + mode = glfw.get_video_mode(glfw_mon) + glfw.set_window_monitor( + glfw_window, + glfw_mon, + 0, + 0, + mode.size.width, + mode.size.height, + mode.refresh_rate, + ) + return + + if monitor is None: + return + + try: + lw, lh = canvas.get_logical_size() + except Exception: + lw, lh = config.pixel_width, config.pixel_height + + x, y = _resolve_window_position( + str(config.window_position), monitor, int(lw), int(lh) + ) + glfw.set_window_pos(glfw_window, x, y) + + +def _apply_qt_placement(canvas, monitor) -> None: + """Apply position / fullscreen on a Qt-backed canvas (QWidget).""" + if config.fullscreen: + canvas.showFullScreen() + return + + if monitor is None: + return + + try: + lw, lh = canvas.get_logical_size() + except Exception: + lw, lh = config.pixel_width, config.pixel_height + + x, y = _resolve_window_position( + str(config.window_position), monitor, int(lw), int(lh) + ) + canvas.move(x, y) + + +# --------------------------------------------------------------------------- +# Interactive camera sensitivity constants +# --------------------------------------------------------------------------- + +# Wheel: zoom sensitivity. +# rendercanvas delivers wheel deltas in CSS pixels (≈100–120 per notch on +# most platforms). A factor of 0.001 gives ≈10 % zoom per notch +# (exp(120 * 0.001) ≈ 1.13). +_ZOOM_SCROLL_FACTOR: float = 0.001 + +# Minimum/maximum focal distance (perspective only). +_ZOOM_MIN_FD: float = 0.5 +_ZOOM_MAX_FD: float = 100.0 + +# Maximum pointer movement (in window pixels) between press and release that +# is still classified as a click rather than a drag. +_CLICK_THRESHOLD_PX: float = 4.0 + + +# --------------------------------------------------------------------------- +# Window class +# --------------------------------------------------------------------------- + + +class WebGPUWindow: + """Preview window wrapping a ``rendercanvas.RenderCanvas``. + + Each call to :meth:`present` polls OS events and blits the offscreen + render texture to the window surface via ``copy_texture_to_texture``. + + Interface expected by ``scene.py`` and ``WebGPURenderer``: + + * ``is_closing`` — True once the OS window has been closed. + * ``destroy()`` — tear down the underlying canvas. + + Interactive camera controls + --------------------------- + Mouse events are translated into camera mutations and an immediate + re-render so the view updates in real time. Pan state (``_pan_x``, + ``_pan_y``) is owned by this class and pushed to the camera just before + each render via :meth:`_sync_pan_to_camera`; orbit and zoom are applied + directly to :class:`~.WebGPUCamera` fields. + + ======================== ================================================ + Input Action + ======================== ================================================ + Left-drag **Orbit** — horizontal drag: ``increment_theta``; + vertical drag: ``increment_phi`` (clamped). + Right-drag / Middle-drag **Pan** — accumulates ``_pan_x`` / ``_pan_y`` + proportional to ``frame_shape / pixel_size``. + Scroll wheel **Zoom** — perspective: ``focal_distance *= exp(dy + * 0.001)``; orthographic: ``frame_shape *= same``. + Key ``r`` **Reset** — clears ``_pan_x``, ``_pan_y`` and + calls ``camera.to_default_state()``. + Key ``q`` **Quit** — closes the preview window. + ======================== ================================================ + + Customising interaction + ----------------------- + Subclass :class:`WebGPUWindow` and override any of the four interaction + hooks to change behaviour without touching internal event dispatch: + + ======================== ================================================ + Method When called + ======================== ================================================ + :meth:`on_mouse_drag` Pointer moves while a button is held. + :meth:`on_scroll` Wheel (scroll) event. + :meth:`on_key_press` Key pressed down. + :meth:`on_key_release` Key released. + :meth:`on_mouse_left_click` Left button pressed and released in place. + :meth:`on_mouse_right_click` Right button pressed and released in place. + ======================== ================================================ + + Each hook can call the building-block helpers :meth:`orbit`, :meth:`pan`, + and :meth:`zoom` and finish with :meth:`_render_from_window` to trigger a + re-render. Pass the subclass to the renderer via + ``WebGPURenderer(window_class=MyWindow)``. + + Example — swap orbit and pan, double zoom speed:: + + class MyWindow(WebGPUWindow): + def on_mouse_drag(self, x, y, dx, dy, button): + if button == 3: # right-drag → orbit + self.orbit(dx, dy) + elif button == 1: # left-drag → pan + self.pan(dx, dy) + else: + return + self._render_from_window() + + def on_scroll(self, x, y, dy): + self.zoom(dy * 2) # 2× sensitivity + self._render_from_window() + + + renderer = WebGPURenderer(window_class=MyWindow) + """ + + def __init__(self, renderer: WebGPURenderer) -> None: + self._renderer = renderer + + win_w, win_h = _compute_window_size() + self._canvas = RenderCanvas( + size=(win_w, win_h), + title=f"Manim Community {__version__}", + # "manual" means we call force_draw() ourselves; the scheduler + # never schedules draws on its own. + update_mode="manual", + ) + + # Configure the wgpu context. + # bgra8unorm matches the render texture format → copy_texture_to_texture + # works without any format conversion. + # COPY_DST is needed on the surface texture so we can copy into it. + self._context = self._canvas.get_wgpu_context() + self._context.configure( + device=renderer._device, + format=wgpu.TextureFormat.bgra8unorm, + usage=wgpu.TextureUsage.RENDER_ATTACHMENT, + ) + + # Register the draw callback (executed inside the rendercanvas lifecycle + # on every force_draw() call). + self._canvas.request_draw(self._draw_frame) + + # ── Drag / click state ──────────────────────────────────────────── + # _drag_button: 1 = left (orbit), 2 = middle (pan), 3 = right (pan) + # None when no button is held. + self._drag_button: int | None = None + self._last_px: float = 0.0 + self._last_py: float = 0.0 + # _press_x/y: pointer position at the moment the button was pressed, + # used to distinguish a click (≤ _CLICK_THRESHOLD_PX movement) from a drag. + self._press_x: float = 0.0 + self._press_y: float = 0.0 + + # ── Pan state ───────────────────────────────────────────────────── + # Camera-space lateral offset in scene units, accumulated from + # right/middle-drag events. Stored here (not on the camera) because + # pan is interactive view navigation, not part of the scripted camera + # model. Synced to camera._cam_pan_x/y just before every render. + self._pan_x: float = 0.0 + self._pan_y: float = 0.0 + + # Blit pipeline — scales the offscreen render texture to whatever size + # the window surface happens to be (supports config.window_size). + self._blit_pipeline, self._blit_bgl, self._blit_sampler = ( + self._create_blit_pipeline() + ) + + # Register event handlers. + self._canvas.add_event_handler(self._on_key_down, "key_down") + self._canvas.add_event_handler(self._on_key_up, "key_up") + self._canvas.add_event_handler(self._on_pointer_move, "pointer_move") + self._canvas.add_event_handler(self._on_pointer_down, "pointer_down") + self._canvas.add_event_handler(self._on_pointer_up, "pointer_up") + self._canvas.add_event_handler(self._on_wheel, "wheel") + + # Apply window configuration: size, position, monitor, fullscreen. + _apply_window_config(self._canvas) + + # ------------------------------------------------------------------ + # Public interface (consumed by WebGPURenderer and scene.py) + # ------------------------------------------------------------------ + + @property + def is_closing(self) -> bool: + """True once the OS window has been closed.""" + return self._canvas.get_closed() + + def destroy(self) -> None: + """Close the OS window.""" + self._canvas.close() + + def present(self) -> None: + """Poll OS events and blit the current render texture to the window. + + Call this after every :meth:`~WebGPURenderer.update_frame` that + should be visible in the preview window. + """ + # Process pending OS events (keyboard, mouse, resize, close …). + self._canvas._process_events() + # Trigger _draw_frame → copy_texture_to_texture → present to screen. + self._canvas.force_draw() + + # ------------------------------------------------------------------ + # Blit pipeline — scales render texture → window surface + # ------------------------------------------------------------------ + + _BLIT_SHADER = """ + struct VertOut { + @builtin(position) pos : vec4, + @location(0) uv : vec2, + }; + + // Full-screen quad from 4 vertices (triangle-strip). + // NDC y=+1 is the top of the screen; UV y=0 is the top of the texture. + @vertex + fn vs_main(@builtin(vertex_index) vi: u32) -> VertOut { + var pos = array, 4>( + vec2(-1.0, 1.0), // top-left + vec2( 1.0, 1.0), // top-right + vec2(-1.0, -1.0), // bottom-left + vec2( 1.0, -1.0), // bottom-right + ); + var uv = array, 4>( + vec2(0.0, 0.0), // top-left + vec2(1.0, 0.0), // top-right + vec2(0.0, 1.0), // bottom-left + vec2(1.0, 1.0), // bottom-right + ); + var out: VertOut; + out.pos = vec4(pos[vi], 0.0, 1.0); + out.uv = uv[vi]; + return out; + } + + @group(0) @binding(0) var tex : texture_2d; + @group(0) @binding(1) var samp : sampler; + + @fragment + fn fs_main(in: VertOut) -> @location(0) vec4 { + return textureSample(tex, samp, in.uv); + } + """ + + def _create_blit_pipeline(self): + """Build the pipeline used to blit the render texture to the window. + + Returns ``(pipeline, bind_group_layout, sampler)``. All three are + reused every frame; only the per-frame bind group (which wraps the + current render texture view) is created anew in ``_draw_frame``. + """ + device = self._renderer._device + + shader = device.create_shader_module(code=self._BLIT_SHADER) + + bgl = device.create_bind_group_layout( + entries=[ + { + "binding": 0, + "visibility": wgpu.ShaderStage.FRAGMENT, + "texture": { + "sample_type": "float", + "view_dimension": "2d", + "multisampled": False, + }, + }, + { + "binding": 1, + "visibility": wgpu.ShaderStage.FRAGMENT, + "sampler": {"type": "filtering"}, + }, + ] + ) + + pipeline = device.create_render_pipeline( + layout=device.create_pipeline_layout(bind_group_layouts=[bgl]), + vertex={"module": shader, "entry_point": "vs_main"}, + fragment={ + "module": shader, + "entry_point": "fs_main", + "targets": [{"format": wgpu.TextureFormat.bgra8unorm}], + }, + primitive={ + "topology": wgpu.PrimitiveTopology.triangle_strip, + "strip_index_format": wgpu.IndexFormat.uint32, + }, + ) + + # Linear filtering gives a smooth downscale when the window is + # smaller than the render texture; nearest would produce aliasing. + sampler = device.create_sampler( + mag_filter="linear", + min_filter="linear", + ) + + return pipeline, bgl, sampler + + # ------------------------------------------------------------------ + # Draw callback (runs inside the rendercanvas present lifecycle) + # ------------------------------------------------------------------ + + def _draw_frame(self) -> None: + """Blit the offscreen render texture to the window surface. + + A full-screen-quad render pass scales the render texture to whatever + size the window surface currently is, so the image always fills the + window correctly regardless of ``config.window_size``. + """ + renderer = self._renderer + if renderer._render_texture is None or renderer._device is None: + return + + device = renderer._device + surface_tex = self._context.get_current_texture() + surface_view = surface_tex.create_view() + + render_tex_view = renderer._render_texture.create_view() + bind_group = device.create_bind_group( + layout=self._blit_bgl, + entries=[ + {"binding": 0, "resource": render_tex_view}, + {"binding": 1, "resource": self._blit_sampler}, + ], + ) + + encoder = device.create_command_encoder() + rp = encoder.begin_render_pass( + color_attachments=[ + { + "view": surface_view, + "load_op": "clear", + "store_op": "store", + "clear_value": (0.0, 0.0, 0.0, 1.0), + } + ] + ) + rp.set_pipeline(self._blit_pipeline) + rp.set_bind_group(0, bind_group) + rp.draw(4) # 4 vertices → one triangle-strip quad + rp.end() + + device.queue.submit([encoder.finish()]) + + # ------------------------------------------------------------------ + # Interactive camera helpers + # ------------------------------------------------------------------ + + def _sync_pan_to_camera(self) -> None: + """Write the window's pan offset into the camera before rendering. + + ``WebGPUCamera.view_matrix`` reads ``_cam_pan_x`` / ``_cam_pan_y`` + via ``getattr`` so they don't need to be initialised in ``__init__``. + This call makes the camera pick up the current window pan without + the camera model needing to know about interactive navigation. + """ + cam = self._renderer.camera + cam._cam_pan_x = self._pan_x + cam._cam_pan_y = self._pan_y + + def _render_from_window(self) -> None: + """Sync pan, re-render the current scene, and present to the window. + + Called after every camera mutation so the preview updates immediately — + even when no animation is running and the main loop is not calling + ``update_frame`` in a tight loop. + """ + renderer = self._renderer + if renderer._device is None: + return + scene = getattr(renderer, "scene", None) + if scene is None: + return + self._sync_pan_to_camera() + renderer.update_frame(scene) + self._canvas.force_draw() + + # ------------------------------------------------------------------ + # Overridable interaction building blocks + # ------------------------------------------------------------------ + + def orbit(self, dx: float, dy: float) -> None: + """Rotate the camera by *dx* horizontal and *dy* vertical pixel deltas. + + Override to change orbit behaviour or sensitivity. + + Horizontal drag rotates around the vertical axis (theta). + Vertical drag tilts up/down (phi, clamped to ±90°). + + A full horizontal swipe (pixel_width pixels) = one full revolution. + A full vertical swipe (pixel_height pixels) = 180° tilt. + + Sign convention (rendercanvas y-axis points downward): + * Drag right (dx > 0) → scene rotates to the right → theta increases. + * Drag down (dy > 0) → scene tilts down → phi decreases. + """ + cam = self._renderer.camera + pw = max(config.pixel_width, 1) + ph = max(config.pixel_height, 1) + dtheta = dx * (2.0 * math.pi / pw) + dphi = -dy * (math.pi / ph) + cam.increment_theta(dtheta) + cam.increment_phi(dphi) + + def pan(self, dx: float, dy: float) -> None: + """Translate the camera laterally by *dx* / *dy* pixel deltas. + + Override to change pan behaviour or sensitivity. + + Updates the window-owned pan state. The new values are pushed to the + camera by ``_render_from_window`` → ``_sync_pan_to_camera``. + + One full horizontal swipe (pixel_width pixels) shifts the scene by + exactly one ``frame_width`` scene unit. + + Signs (rendercanvas y-axis downward; Manim y-axis upward): + * Drag right (dx > 0) → scene moves right → _pan_x increases. + * Drag down (dy > 0) → scene moves down → _pan_y decreases. + """ + fw, fh = self._renderer.camera.frame_shape + pw = max(config.pixel_width, 1) + ph = max(config.pixel_height, 1) + self._pan_x += dx * (fw / pw) + self._pan_y -= dy * (fh / ph) + + def zoom(self, scroll_dy: float) -> None: + """Zoom in/out by *scroll_dy* CSS-pixel scroll units. + + Override to change zoom behaviour, sensitivity, or limits. + + Positive *scroll_dy* (scroll down) zooms out; negative zooms in. + + Perspective camera: adjusts ``focal_distance``. + Orthographic camera: scales ``frame_shape`` proportionally. + + The zoom is exponential so that successive zoom steps are perceptually + uniform: each notch (≈ 120 CSS pixels) changes the scale by ~12 %. + """ + cam = self._renderer.camera + factor = math.exp(scroll_dy * _ZOOM_SCROLL_FACTOR) + + if cam.orthographic: + fw, fh = cam.frame_shape + new_fw = max(fw * factor, 0.01) + new_fh = max(fh * factor, 0.01) + cam.frame_shape = (new_fw, new_fh) + else: + new_fd = float( + np.clip( + cam.focal_distance * factor, + _ZOOM_MIN_FD, + _ZOOM_MAX_FD, + ) + ) + cam.set_focal_distance(new_fd) + + # ------------------------------------------------------------------ + # Overridable interaction hooks + # ------------------------------------------------------------------ + + def on_mouse_drag( + self, x: float, y: float, dx: float, dy: float, button: int + ) -> None: + """Called on every pointer-move event while a mouse button is held. + + Override to customise drag behaviour. The default implementation + maps button 1 (left) to :meth:`orbit` and buttons 2/3 + (middle/right) to :meth:`pan`. + + Parameters + ---------- + x, y: + Current pointer position in window pixels (y-axis downward). + dx, dy: + Delta from the previous pointer position in window pixels. + button: + Web-standard button code: 1 = left, 2 = middle, 3 = right. + """ + if button == 1: + self.orbit(dx, dy) + elif button in (2, 3): + self.pan(dx, dy) + else: + return + self._render_from_window() + + def on_scroll(self, x: float, y: float, dy: float) -> None: + """Called on every wheel (scroll) event. + + Override to customise scroll behaviour. The default implementation + calls :meth:`zoom`. + + Parameters + ---------- + x, y: + Pointer position at the time of the scroll, in window pixels. + dy: + Vertical scroll delta in CSS pixels (positive = scroll down = + zoom out). + """ + self.zoom(dy) + self._render_from_window() + + def on_key_press(self, key: str, modifiers: int) -> None: + """Called when a key is pressed. + + Override to add or replace key bindings. Call ``super().on_key_press(key, + modifiers)`` to keep the default ``r`` → reset and ``q`` → quit bindings. + + Parameters + ---------- + key: + Web-standard key string (e.g. ``"r"``, ``"ArrowLeft"``, + ``"Escape"``). + modifiers: + Pyglet-compatible modifier bitmask (Shift=1, Control=4, Alt=8). + """ + if key == "r": + self._pan_x = 0.0 + self._pan_y = 0.0 + self._renderer.camera.to_default_state() + self._render_from_window() + elif key == "q": + self._canvas.close() + + def on_key_release(self, key: str, modifiers: int) -> None: + """Called when a key is released. + + Override to react to key-release events. The default implementation + does nothing (key tracking is handled internally). + + Parameters + ---------- + key: + Web-standard key string. + modifiers: + Pyglet-compatible modifier bitmask. + """ + + def on_mouse_left_click(self, x: float, y: float) -> None: + """Called when the left mouse button is clicked (pressed and released + without dragging more than ``_CLICK_THRESHOLD_PX`` pixels). + + Override to add left-click behaviour. The default implementation + does nothing. + + Parameters + ---------- + x, y: + Pointer position at release, in window pixels (y-axis downward). + """ + + def on_mouse_right_click(self, x: float, y: float) -> None: + """Called when the right mouse button is clicked (pressed and released + without dragging more than ``_CLICK_THRESHOLD_PX`` pixels). + + Override to add right-click behaviour. The default implementation + does nothing. + + Parameters + ---------- + x, y: + Pointer position at release, in window pixels (y-axis downward). + """ + + # ------------------------------------------------------------------ + # Raw event handlers — dispatch to the overridable hooks above. + # Subclasses should override the hooks, not these methods. + # ------------------------------------------------------------------ + + def _on_key_down(self, event: dict) -> None: + key = event.get("key", "") + modifiers = _modifiers_to_int(event.get("modifiers", [])) + self._renderer.pressed_keys.add(_key_to_int(key)) + self.on_key_press(key, modifiers) + + def _on_key_up(self, event: dict) -> None: + key = event.get("key", "") + modifiers = _modifiers_to_int(event.get("modifiers", [])) + self._renderer.pressed_keys.discard(_key_to_int(key)) + self.on_key_release(key, modifiers) + + def _on_pointer_down(self, event: dict) -> None: + # button: 1 = left, 2 = middle, 3 = right (Web standard) + self._drag_button = event.get("button", 1) + x = float(event.get("x", 0)) + y = float(event.get("y", 0)) + self._last_px = x + self._last_py = y + self._press_x = x + self._press_y = y + + def _on_pointer_up(self, event: dict) -> None: + button = self._drag_button + x = float(event.get("x", 0)) + y = float(event.get("y", 0)) + self._drag_button = None + # Fire a click hook only when the pointer barely moved (not a drag). + if math.hypot(x - self._press_x, y - self._press_y) < _CLICK_THRESHOLD_PX: + if button == 1: + self.on_mouse_left_click(x, y) + elif button == 3: + self.on_mouse_right_click(x, y) + + def _on_pointer_move(self, event: dict) -> None: + if self._drag_button is None: + return + px = float(event.get("x", 0)) + py = float(event.get("y", 0)) + dx = px - self._last_px + dy = py - self._last_py + self._last_px = px + self._last_py = py + if dx == 0.0 and dy == 0.0: + return + self.on_mouse_drag(px, py, dx, dy, self._drag_button) + + def _on_wheel(self, event: dict) -> None: + dy = float(event.get("dy", 0)) + if dy == 0.0: + return + self.on_scroll(float(event.get("x", 0)), float(event.get("y", 0)), dy) diff --git a/manim/renderer/webgpu/webgpu_vmobject_rendering.py b/manim/renderer/webgpu/webgpu_vmobject_rendering.py new file mode 100644 index 0000000000..6186279ca3 --- /dev/null +++ b/manim/renderer/webgpu/webgpu_vmobject_rendering.py @@ -0,0 +1,1947 @@ +"""WebGPU draw calls for VMobject fill + stroke + surface rendering. + +Fill + Stroke (combined pipeline) +---------------------------------- +``cubic_to_quads.wgsl`` — GPU compute shader converts raw cubic Bezier +control points to quadratic approximations (4 per cubic, two-level de +Casteljau subdivision). + +``vmobject_fill_stroke.wgsl`` — combined render shader: one bounding quad +per object, one fragment loop accumulates both Slug winding-number fill +coverage (in NDC space) and SDF stroke distance (in pixel space). Porter- +Duff "over" compositing produces the final colour. + +Closing segments +~~~~~~~~~~~~~~~~ +Every open subpath gets a linear closing cubic (degree-elevated line from +the last anchor back to the first) appended to the fill cubic list. This +makes the winding-number integral correct for partial paths (e.g. during +``Create`` animations). The closing cubic is NOT added to the stroke cubic +list — strokes should follow the visible part of the curve only. + +Surfaces +-------- +Parametric surfaces use a combined triangle-mesh pipeline +(``surface_combined.wgsl`` / ``surface_oit.wgsl``) with Phong lighting and +barycentric wireframe in a single draw call. The centroid vertex of each +triangle fan carries bary=(1,0,0); the outer edge (anchor_i ↔ anchor_{i+1}) +has bary.x=0 — this is the visible mesh-grid edge. Transparent surfaces go +through the OIT accumulation + composition passes. + +Batching +-------- +``collect_frame_data`` tessellates *all* scene mobjects on the CPU, uploads +one cubics buffer (fill then stroke, all objects) and one vertex buffer (one +bounding quad per object), then returns a ``_FrameData`` ready for the GPU. +``draw_frame_data`` records draw calls into the active render pass. +""" + +from __future__ import annotations + +import struct +import weakref +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import numpy as np + +from manim.mobject.three_d.dot_cloud import DotCloud3D +from manim.mobject.three_d.three_dimensions import Surface +from manim.mobject.types.vectorized_mobject import VMobject + +if TYPE_CHECKING: + import wgpu as wgpu_t + + from manim.renderer.webgpu.webgpu_renderer import WebGPURenderer + + +# --------------------------------------------------------------------------- +# Combined surface vertex layout — must match surface_combined.wgsl / +# surface_oit.wgsl locations. +# +# location 0 — in_vert float32x3 offset 0 (12 B) +# location 1 — in_normal float32x3 offset 12 (12 B) +# location 2 — in_fill_color float32x4 offset 24 (16 B) +# location 3 — in_stroke_color float32x4 offset 40 (16 B) +# location 4 — in_bary float32x3 offset 56 (12 B) +# location 5 — stroke_half_px float32 offset 68 ( 4 B) +# location 6 — diffuse_strength float32 offset 72 ( 4 B) +# location 7 — specular_strength float32 offset 76 ( 4 B) +# location 8 — specular_exponent float32 offset 80 ( 4 B) +# stride: 84 bytes +# --------------------------------------------------------------------------- + +_SURFACE_COMBINED_DTYPE = np.dtype( + [ + ("in_vert", np.float32, (3,)), + ("in_normal", np.float32, (3,)), + ("in_fill_color", np.float32, (4,)), + ("in_stroke_color", np.float32, (4,)), + ("in_bary", np.float32, (3,)), + ("stroke_half_px", np.float32), + ("diffuse_strength", np.float32), + ("specular_strength", np.float32), + ("specular_exponent", np.float32), + ] +) +_SURFACE_COMBINED_STRIDE: int = _SURFACE_COMBINED_DTYPE.itemsize # 84 bytes + +_SURFACE_COMBINED_OFFSETS: dict[str, int] = { + name: _SURFACE_COMBINED_DTYPE.fields[name][1] # type: ignore[index] + for name in _SURFACE_COMBINED_DTYPE.names +} + +SURFACE_COMBINED_VERTEX_LAYOUT: dict = { + "array_stride": _SURFACE_COMBINED_STRIDE, + "step_mode": "vertex", + "attributes": [ + { + "format": "float32x3", + "offset": _SURFACE_COMBINED_OFFSETS["in_vert"], + "shader_location": 0, + }, + { + "format": "float32x3", + "offset": _SURFACE_COMBINED_OFFSETS["in_normal"], + "shader_location": 1, + }, + { + "format": "float32x4", + "offset": _SURFACE_COMBINED_OFFSETS["in_fill_color"], + "shader_location": 2, + }, + { + "format": "float32x4", + "offset": _SURFACE_COMBINED_OFFSETS["in_stroke_color"], + "shader_location": 3, + }, + { + "format": "float32x3", + "offset": _SURFACE_COMBINED_OFFSETS["in_bary"], + "shader_location": 4, + }, + { + "format": "float32", + "offset": _SURFACE_COMBINED_OFFSETS["stroke_half_px"], + "shader_location": 5, + }, + { + "format": "float32", + "offset": _SURFACE_COMBINED_OFFSETS["diffuse_strength"], + "shader_location": 6, + }, + { + "format": "float32", + "offset": _SURFACE_COMBINED_OFFSETS["specular_strength"], + "shader_location": 7, + }, + { + "format": "float32", + "offset": _SURFACE_COMBINED_OFFSETS["specular_exponent"], + "shader_location": 8, + }, + ], +} + + +# --------------------------------------------------------------------------- +# Combined fill+stroke vertex layout — must match vmobject_fill_stroke.wgsl. +# +# location 0 — in_pos float32x3 offset 0 (12 B) +# location 1 — in_fill_color float32x4 offset 12 (16 B) +# location 2 — in_stroke_color float32x4 offset 28 (16 B) +# location 3 — stroke_half_ndc float32 offset 44 ( 4 B) +# location 4 — fill_curve_start uint32 offset 48 ( 4 B) +# location 5 — n_fill_curves uint32 offset 52 ( 4 B) +# location 6 — stroke_curve_start uint32 offset 56 ( 4 B) +# location 7 — n_stroke_curves uint32 offset 60 ( 4 B) +# location 8 — fill_rule uint32 offset 64 ( 4 B) 0=nonzero, 1=evenodd +# stride: 68 bytes +# --------------------------------------------------------------------------- + +_FILL_STROKE_DTYPE = np.dtype( + [ + ("in_pos", np.float32, (3,)), + ("in_fill_color", np.float32, (4,)), + ("in_stroke_color", np.float32, (4,)), + ("stroke_half_ndc", np.float32), + ("fill_curve_start", np.uint32), + ("n_fill_curves", np.uint32), + ("stroke_curve_start", np.uint32), + ("n_stroke_curves", np.uint32), + ("fill_rule", np.uint32), + ] +) +_FILL_STROKE_STRIDE: int = _FILL_STROKE_DTYPE.itemsize # 64 bytes + +_FILL_STROKE_OFFSETS: dict[str, int] = { + name: _FILL_STROKE_DTYPE.fields[name][1] # type: ignore[index] + for name in _FILL_STROKE_DTYPE.names +} + +FILL_STROKE_VERTEX_LAYOUT: dict = { + "array_stride": _FILL_STROKE_STRIDE, + "step_mode": "vertex", + "attributes": [ + { + "format": "float32x3", + "offset": _FILL_STROKE_OFFSETS["in_pos"], + "shader_location": 0, + }, + { + "format": "float32x4", + "offset": _FILL_STROKE_OFFSETS["in_fill_color"], + "shader_location": 1, + }, + { + "format": "float32x4", + "offset": _FILL_STROKE_OFFSETS["in_stroke_color"], + "shader_location": 2, + }, + { + "format": "float32", + "offset": _FILL_STROKE_OFFSETS["stroke_half_ndc"], + "shader_location": 3, + }, + { + "format": "uint32", + "offset": _FILL_STROKE_OFFSETS["fill_curve_start"], + "shader_location": 4, + }, + { + "format": "uint32", + "offset": _FILL_STROKE_OFFSETS["n_fill_curves"], + "shader_location": 5, + }, + { + "format": "uint32", + "offset": _FILL_STROKE_OFFSETS["stroke_curve_start"], + "shader_location": 6, + }, + { + "format": "uint32", + "offset": _FILL_STROKE_OFFSETS["n_stroke_curves"], + "shader_location": 7, + }, + { + "format": "uint32", + "offset": _FILL_STROKE_OFFSETS["fill_rule"], + "shader_location": 8, + }, + ], +} + + +# --------------------------------------------------------------------------- +# TrueDot vertex layout — must match true_dot.wgsl locations. +# +# location 0 — center float32x3 offset 0 (12 B) +# location 1 — color float32x4 offset 12 (16 B) +# location 2 — uv float32x2 offset 28 ( 8 B) +# location 3 — radius float32 offset 36 ( 4 B) +# location 4 — gloss float32 offset 40 ( 4 B) +# location 5 — shadow float32 offset 44 ( 4 B) +# stride: 48 bytes +# --------------------------------------------------------------------------- + +_TRUE_DOT_DTYPE = np.dtype( + [ + ("center", np.float32, (3,)), + ("color", np.float32, (4,)), + ("uv", np.float32, (2,)), + ("radius", np.float32), + ("gloss", np.float32), + ("shadow", np.float32), + ] +) +_TRUE_DOT_STRIDE: int = _TRUE_DOT_DTYPE.itemsize # 48 bytes + +_TRUE_DOT_OFFSETS: dict[str, int] = { + name: _TRUE_DOT_DTYPE.fields[name][1] # type: ignore[index] + for name in _TRUE_DOT_DTYPE.names +} + +TRUE_DOT_VERTEX_LAYOUT: dict = { + "array_stride": _TRUE_DOT_STRIDE, + "step_mode": "vertex", + "attributes": [ + { + "format": "float32x3", + "offset": _TRUE_DOT_OFFSETS["center"], + "shader_location": 0, + }, + { + "format": "float32x4", + "offset": _TRUE_DOT_OFFSETS["color"], + "shader_location": 1, + }, + { + "format": "float32x2", + "offset": _TRUE_DOT_OFFSETS["uv"], + "shader_location": 2, + }, + { + "format": "float32", + "offset": _TRUE_DOT_OFFSETS["radius"], + "shader_location": 3, + }, + { + "format": "float32", + "offset": _TRUE_DOT_OFFSETS["gloss"], + "shader_location": 4, + }, + { + "format": "float32", + "offset": _TRUE_DOT_OFFSETS["shadow"], + "shader_location": 5, + }, + ], +} + +# Corner UV offsets for the two triangles that form a screen-aligned quad: +# triangle 0: (BL, BR, TL) → corners 0,1,2 +# triangle 1: (BR, TR, TL) → corners 1,3,2 +# x_sign: -1 +1 -1 +1 y_sign: -1 -1 +1 +1 +_QUAD_UVS = np.array( + [ + [-1.0, -1.0], # BL (0) + [1.0, -1.0], # BR (1) + [-1.0, 1.0], # TL (2) + [1.0, -1.0], # BR (1) ← repeated for 2nd triangle + [1.0, 1.0], # TR (3) + [-1.0, 1.0], # TL (2) ← repeated + ], + dtype=np.float32, +) # shape (6, 2) + + +def build_true_dot_vbo( + mob: DotCloud3D, +) -> np.ndarray | None: + """Expand a ``DotCloud3D`` into a flat vertex array for TrueDot rendering. + + Each point becomes 6 vertices (2 triangles) forming a screen-aligned quad. + UV coords span (−1,−1) → (1,1); the radius is in world-space scene units. + + Returns ``None`` if the mob has no renderable points. + """ + pts = mob.get_cloud_points() + rgbas = mob.get_rgbas() + radius = mob.dot_radius + gloss = mob.gloss + shadow = mob.shadow + + pts = np.asarray(pts, dtype=np.float32) # (N, 3) + N = len(pts) + if N == 0: + return None + + # Broadcast rgbas to (N, 4). + if rgbas is None or len(rgbas) == 0: + rgba = np.ones((N, 4), dtype=np.float32) + else: + rgbas = np.asarray(rgbas, dtype=np.float32) + if len(rgbas) == 1: + rgba = np.repeat(rgbas[:1], N, axis=0) + elif len(rgbas) < N: + # Resize with interpolation (matches OpenGL behaviour). + indices = np.round(np.linspace(0, len(rgbas) - 1, N)).astype(int) + rgba = rgbas[indices] + else: + rgba = rgbas[:N] + + # Expand N points → N×6 vertices. + pts_rep = np.repeat(pts, 6, axis=0) # (N*6, 3) + rgba_rep = np.repeat(rgba, 6, axis=0) # (N*6, 4) + uvs = np.tile(_QUAD_UVS, (N, 1)) # (N*6, 2) + + arr = np.zeros(N * 6, dtype=_TRUE_DOT_DTYPE) + arr["center"] = pts_rep + arr["color"] = rgba_rep + arr["uv"] = uvs + arr["radius"] = radius + arr["gloss"] = gloss + arr["shadow"] = shadow + return arr + + +# --------------------------------------------------------------------------- +# Per-frame data container +# --------------------------------------------------------------------------- + + +@dataclass +class _FrameData: + """All GPU-ready data for one group of mobjects (one camera bind group). + + Produced by ``collect_frame_data``; consumed by ``draw_frame_data`` and + the caller's OIT / fixed-frame passes. + """ + + # VMobject fill+stroke via combined pipeline + fs_parts: list[np.ndarray] # _FILL_STROKE_DTYPE arrays, one per draw call + fs_buf: wgpu_t.GPUBuffer | None # concatenated vertex buffer + fs_byte_offsets: list[int] # byte offset of each part in fs_buf + + # GPU compute: cubic → quadratic conversion + cubics_buf: wgpu_t.GPUBuffer | None # input (12 floats/cubic), all objects + quads_out_buf: wgpu_t.GPUBuffer | None # output (36 floats/cubic = 4 quads × 9) + n_cubics_total: int + compute_bg: wgpu_t.GPUBindGroup | None # compute pass bind group + render_bg: wgpu_t.GPUBindGroup | None # fragment bind group (camera + quads) + + # Parametric surfaces (combined fill + barycentric wireframe pipeline) + surface_parts: list[np.ndarray] + surface_buf: wgpu_t.GPUBuffer | None + surface_byte_offsets: list[int] + + # Ordered draw commands: + # "fill_stroke_2d" — 2-D VMobject (no depth write) + # "fill_stroke_3d" — shade_in_3d VMobject (depth write + test) + # "surface_opaque" — opaque parametric surface + # "surface_oit" — transparent parametric surface (OIT pass, caller handles) + draw_plan: list[tuple[str, int]] + + # Indices into surface_parts that need OIT (handled by the caller). + oit_indices: list[int] + + +# --------------------------------------------------------------------------- +# Geometry caches +# --------------------------------------------------------------------------- + +# fill_stroke_cache: vmobject → (points_hash, (fill_cubics, stroke_cubics)) +# fill_cubics : (N, 4, 3) float32 — includes closing segments for winding +# stroke_cubics: (M, 4, 3) float32 — no closing segments (visible curve only) +# Geometry only; colors/widths are fetched fresh every frame. +_fill_stroke_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + +# surface_mob_cache: Surface mob → +# (geom_hash, color_hash, big_template, seg_starts, seg_ends, +# sw_arr, sa_arr, has_sw, draw_cmds) +# +# geom_hash : bytes — hash of patch geometry + material ONLY (no colors). +# Unchanged when FadeIn/set_fill changes opacity. +# color_hash : bytes — hash of fill_rgba, stroke_rgba, stroke_width per patch. +# Changes on FadeIn/set_fill without requiring retessellation. +# big_template: single concatenated _SURFACE_COMBINED_DTYPE array with +# already-smoothed normals; colors reflect the last stored frame; +# stroke_half_px == 0.0 (recomputed per-frame on each hit). +# seg_starts/ends: numpy intp arrays of part boundaries in big_template. +# sw_arr/sa_arr/has_sw: stroke metadata, updated in-place on color misses. +# draw_cmds : list[str] — "surface_opaque" or "surface_oit" per part; +# updated in-place on color misses (opacity class can change). +# +# Cache hit states: +# geom HIT + color HIT → just recompute stroke_half_px (camera rotation path) +# geom HIT + color MISS → patch colors in big_template + recompute stroke_half_px +# (FadeIn / set_fill without geometry change; O(N_parts)) +# geom MISS → full retessellation + normal smoothing +_surface_mob_cache: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + +# surface_color_memo: Surface mob → +# (fill_cols, stroke_cols, sw_arr, sa_arr) +# fill_cols : float32 (N_active, 4) — fill RGBA per active part +# stroke_cols : float32 (N_active, 4) — stroke RGBA per active part +# sw_arr : float32 (N_active,) — stroke width per active part +# sa_arr : float32 (N_active,) — stroke alpha per active part +# +# Populated by _surface_hash_pair whenever the color hash is recomputed +# (color-only miss or full miss path). The color-only update path in +# collect_frame_data reads from here instead of re-iterating all submobjects, +# saving ~2 full O(N_submobs) passes per Surface per color-changing frame. +_surface_color_memo: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + +def _points_hash(vmobject: VMobject) -> int: + pts = vmobject.points + if pts.size == 0: + return 0 + # Cast to float32 before hashing so that sub-epsilon float64 noise + # (e.g. from FadeIn / Transform's straight_path interpolation, which + # produces ~1e-17 differences when start == end) does not create + # spurious cache misses. Genuine geometry changes are at least + # float32-epsilon (~1e-7) in magnitude and are still detected. + return hash(pts.astype(np.float32).tobytes()) + + +# _surface_geom_hash_memo: Surface mob → (fast_geom_id, fast_color_id, geom_hash, color_hash) +# fast_geom_id — XOR of id(s.points) for all submobs. +# fast_color_id — XOR of id(fill_rgbas) ^ id(stroke_rgbas) for all submobs. +# Separately tracking the two fast IDs lets us skip recomputing the geometry hash +# on a color-only change without missing a genuine geometry update. +_surface_geom_hash_memo: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + +def _surface_hash_pair( + mob: Surface, + submobs: list | None = None, +) -> tuple[bytes, bytes]: + """Return ``(geom_hash, color_hash)`` for *mob*. + + ``geom_hash`` covers patch point positions + material params (diffuse, + specular, specular_exp). It does NOT include fill/stroke colors, so a + FadeIn that only changes opacity does not invalidate it. + + ``color_hash`` is ``fast_color_id`` packed as 8 bytes, where + ``fast_color_id`` is the XOR of ``id(fill_rgbas)`` / ``id(stroke_rgbas)`` + for all submobject patches. It changes whenever Manim replaces any color + array (which happens on every FadeIn / set_fill frame). + + Two-level memoisation avoids full per-submob rehashing on every frame: + the fast IDs (XOR of array ``id()``s) detect changes in O(N_submobs) + without touching array data; the slow geometry hash runs only on a + geometry miss (first call or actual point change). + + *submobs* — optional pre-computed ``mob.family_members_with_points()``. + Pass this from ``collect_frame_data`` to avoid a redundant tree walk. + """ + if submobs is None: + submobs = mob.family_members_with_points() + + fast_geom_id = 0 + fast_color_id = 0 + for s in submobs: + fast_geom_id ^= id(s.points) + fast_color_id ^= id(getattr(s, "fill_rgbas", None)) + fast_color_id ^= id(getattr(s, "stroke_rgbas", None)) + + memo = _surface_geom_hash_memo.get(mob) + if memo is not None: + cached_fgi, cached_fci, cached_gh, cached_ch = memo + if cached_fgi == fast_geom_id and cached_fci == fast_color_id: + # Both geometry and colors unchanged. + return cached_gh, cached_ch + if cached_fgi == fast_geom_id: + # Geometry unchanged, colors changed — use fast_color_id as the + # color hash (no submob method calls needed here). The actual + # color arrays are read lazily by collect_frame_data when it + # applies the per-vertex update, so we skip the second O(N_submobs) + # iteration entirely. + new_ch = struct.pack(" bytes: + """Compute a compact fingerprint of the mobject set + camera state. + + Captures: view/projection matrices, per-submobject geometry hash, fill + color, stroke color, and stroke width. Two frames with identical + fingerprints are guaranteed to produce pixel-identical renders. + + *center_view_matrix* — when provided (fixed-orientation path), it is + included in the fingerprint so that camera rotation invalidates the cache + even though *view_matrix* (the stripped fixed_view) is constant. + + Cost: O(n) over all leaf submobjects, but only does scalar reads and bytes + operations — no numpy matrix math or buffer allocations. Much cheaper + than a full tessellation pass. + """ + parts: list[bytes] = [view_matrix.tobytes(), proj_matrix.tobytes()] + if center_view_matrix is not None: + parts.append(center_view_matrix.tobytes()) + for mob in mobjects: + for submob in mob.family_members_with_points(): + phash = _points_hash(submob) + fill_rgba = submob.get_fill_rgbas() + stroke_rgba = submob.get_stroke_rgbas() + sw = float(submob.get_stroke_width()) if stroke_rgba.shape[0] > 0 else 0.0 + parts.append(struct.pack(" _FrameData | None: + """Tessellate *mobjects*, upload to GPU, return a ``_FrameData``. + + Does NOT record any GPU commands — only uploads buffers and creates bind + groups. The caller must run the compute pass (via ``_FrameData.compute_bg``) + before the render pass. + + *camera_uniform_buf* is the 656-byte uniform buffer for this camera group. + It is stored in the render bind group so the fragment shader can project + world-space curve data into the correct NDC space. + + *view_matrix_override* / *proj_matrix_override* replace the camera's normal + view and projection matrices for CPU-side bounding-quad computation. Use + these for fixed-in-frame and fixed-orientation mobjects so that the + world-space quad vertices are consistent with the bind group the GPU will + use to rasterise them. + + *center_view_matrix* — when provided, enables fixed-orientation rendering. + Each submobject's bezier control points are pre-translated by + ``(R_full - I) @ center_w`` so the object appears at the 3D-projected + position of its world center while its local orientation stays upright + (no camera rotation applied to the local shape). The GPU shader still + uses *view_matrix_override* (typically the rotation-stripped fixed_view), + and the pre-translation makes the combined effect equivalent to Cairo's + ``transform_points_pre_display`` for fixed-orientation objects. + + *cache_slot* — when not None, enables ``_FrameData`` caching for this + call. On a fingerprint hit the cached ``_FrameData`` is returned + immediately, skipping all tessellation and GPU buffer uploads. The + ``camera_uniform_buf`` must be a *persistent* buffer (same Python object + across frames) so that cached ``render_bg`` bind groups remain valid. + """ + import wgpu + + view_matrix: np.ndarray = ( + view_matrix_override + if view_matrix_override is not None + else renderer.camera.view_matrix + ) + proj_matrix: np.ndarray = ( + proj_matrix_override + if proj_matrix_override is not None + else renderer.camera.projection_matrix + ) + + # ── _FrameData cache check ──────────────────────────────────────────── + # When all mobs are Surface objects the _surface_mob_cache handles + # geometry caching. Skip _fd_fingerprint (it always misses when the + # camera rotates) and disable _FrameData caching for this call — the + # surface GPU buffer must be regenerated each frame to update + # stroke_half_px. Non-surface or mixed calls still use _fd_fingerprint. + _all_surface_call: bool = bool(mobjects) and all( + isinstance(m, Surface) for m in mobjects if isinstance(m, VMobject) + ) + fp: bytes = b"" # populated below on non-surface-only paths + if cache_slot is not None and mobjects and not _all_surface_call: + fp = _fd_fingerprint(mobjects, view_matrix, proj_matrix, center_view_matrix) + cached = renderer._fd_cache.get(cache_slot) + if cached is not None and cached[0] == fp: + # Scene + camera unchanged — return the cached GPU data directly. + # The compute pass will re-dispatch into the same quads_out_buf + # (safe: identical input → identical output; the render pass reads + # it after the compute pass completes within the same encoder). + return cached[1] + # Cache miss — tessellate below, then store result before returning. + + # Per-draw-call data collected across all mobjects. + fs_parts: list[np.ndarray] = [] + # Cubics: fill first (all objects), then stroke (all objects). + all_fill_cubics: list[np.ndarray] = [] # (Ni, 4, 3) per draw call + all_stroke_cubics: list[np.ndarray] = [] # (Mi, 4, 3) per draw call + n_fill_cubics_per: list[int] = [] # Ni per draw call + n_stroke_cubics_per: list[int] = [] # Mi per draw call + + surface_parts: list[np.ndarray] = [] + draw_plan: list[tuple[str, int]] = [] + + # Tracks newly-tessellated Surface mobs (cache misses) so we can store + # their parts in _surface_mob_cache after _smooth_surface_normals runs. + # Each entry: (mob, geom_hash, parts_start_idx, parts_end_idx) + _new_surface_mobs: list[tuple] = [] + + # Guard against double-processing the same submobject. This can happen + # when mob_list is a flat family list (e.g. moving_mobjects from + # begin_animations) that contains both a VGroup and its children: without + # the guard, family_members_with_points() on the VGroup would process the + # children, and then those children would be processed again individually. + _seen_submobs: set[int] = set() + + use_z_index: bool = renderer.camera.use_z_index + + for mob in mobjects: + if not isinstance(mob, VMobject): + continue + + # ── Parametric Surface ──────────────────────────────────────────── + if isinstance(mob, Surface): + surface_submobs = mob.family_members_with_points() + if use_z_index: + surface_submobs = sorted(surface_submobs, key=lambda m: m.z_index) + # Read material params from the parent Surface as defaults; each + # submobject patch may override them individually by carrying its + # own diffuse_strength / specular_strength / specular_exponent + # instance attribute (set via set_*_by_func or direct assignment). + surf_diffuse = float(getattr(mob, "diffuse_strength", 0.8)) + surf_specular = float(getattr(mob, "specular_strength", 0.9)) + surf_spec_exp = float(getattr(mob, "specular_exponent", 16.0)) + + # ── Surface geometry cache ──────────────────────────────────── + # The geometry (verts, normals, bary, colors, material) does NOT + # depend on view_matrix / proj_matrix — only stroke_half_px does. + # Cache the fully-tessellated + smoothed arrays per mob so that + # a camera-only change (ambient rotation, etc.) skips the expensive + # per-patch Python loop and just updates stroke_half_px. + geom_hash, color_hash = _surface_hash_pair(mob, submobs=surface_submobs) + cached_entry = _surface_mob_cache.get(mob) + + if cached_entry is not None and cached_entry[0] == geom_hash: + # ── Geometry HIT ────────────────────────────────────────── + # cached_entry = (geom_hash, color_hash, big_template, + # seg_starts, seg_ends, + # sw_arr, sa_arr, has_sw, draw_cmds) + ( + _, + cached_color_hash, + big_template, + seg_starts, + seg_ends, + sw_arr, + sa_arr, + has_sw, + draw_cmds, + ) = cached_entry + + if cached_color_hash != color_hash: + # ── Color-only miss (FadeIn / set_fill / set_stroke) ── + # Read current colors from submobjects directly using fast + # attribute access (avoids method call overhead). Build + # per-part color arrays then write them in a single + # vectorized np.repeat call instead of N_parts slice writes. + fill_list: list[np.ndarray] = [] + stroke_list: list[np.ndarray] = [] + sw_list: list[float] = [] + sa_list: list[float] = [] + for submob in surface_submobs: + if id(submob) in _seen_submobs: + continue + f_rgba = getattr(submob, "fill_rgbas", None) + s_rgba = getattr(submob, "stroke_rgbas", None) + if f_rgba is None or f_rgba.shape[0] == 0: + continue + if float(f_rgba[0, 3]) <= 0.0: + continue + has_stroke = s_rgba is not None and s_rgba.shape[0] > 0 + fill_list.append(f_rgba[0].astype(np.float32)) + stroke_list.append( + s_rgba[0].astype(np.float32) + if has_stroke + else np.zeros(4, dtype=np.float32) + ) + sw_list.append( + float(submob.stroke_width) if has_stroke else 0.0 + ) + sa_list.append(float(s_rgba[0, 3]) if has_stroke else 0.0) + + if len(fill_list) == len(draw_cmds): + fill_cols = np.array(fill_list, dtype=np.float32) + stroke_cols = np.array(stroke_list, dtype=np.float32) + sw_arr = np.array(sw_list, dtype=np.float32) + sa_arr = np.array(sa_list, dtype=np.float32) + has_sw = (sw_arr > 0.0) & (sa_arr > 0.001) + draw_cmds = [ + "surface_opaque" + if float(fill_cols[i, 3]) >= 0.99 + else "surface_oit" + for i in range(len(draw_cmds)) + ] + # Vectorized color write: expand per-part colors to + # per-vertex with repeat counts, then assign in one op. + rep_counts = (seg_ends - seg_starts).astype(np.intp) + big_template["in_fill_color"] = np.repeat( + fill_cols, rep_counts, axis=0 + ) + big_template["in_stroke_color"] = np.repeat( + stroke_cols, rep_counts, axis=0 + ) + _surface_mob_cache[mob] = ( + geom_hash, + color_hash, + big_template, + seg_starts, + seg_ends, + sw_arr, + sa_arr, + has_sw, + draw_cmds, + ) + else: + # Part count changed — treat as full miss. + cached_entry = None + + if cached_entry is not None and cached_entry[0] == geom_hash: + # ── Full HIT: copy template and recompute stroke_half_px ── + big_copy = big_template.copy() + vm = view_matrix.astype(np.float32) + pm = proj_matrix.astype(np.float32) + R, t = vm[:3, :3], vm[:3, 3] + from manim import config as _cfg + + px_half = _cfg.pixel_width * 0.5 + pm_00 = abs(float(pm[0, 0])) + pm_32 = float(pm[3, 2]) + pm_33 = float(pm[3, 3]) + + # Vectorized stroke_half_px: one matrix multiply + reduceat + # instead of per-part Python slice+mean inside a loop. + z_vals = ((R @ big_copy["in_vert"].T).T + t)[:, 2] + part_sizes = (seg_ends - seg_starts).astype(np.float32) + + # Per-part average view-space z via reduceat sum / count. + z_sums = np.add.reduceat(z_vals, seg_starts) + avg_z = z_sums / part_sizes # (n_parts,) + clip_w = pm_32 * avg_z + pm_33 # (n_parts,) + clip_w = np.where(np.abs(clip_w) < 1e-8, 1.0, clip_w) + + shp = np.where( + has_sw, + 0.004 * sw_arr * pm_00 / np.abs(clip_w) * px_half, + 0.0, + ).astype(np.float32) # (n_parts,) + + # Write per-part stroke_half_px into the copy using index ranges. + for i in range(len(draw_cmds)): + big_copy["stroke_half_px"][seg_starts[i] : seg_ends[i]] = shp[i] + + # Slice views for draw_plan / surface_parts. + for i, cmd in enumerate(draw_cmds): + draw_plan.append((cmd, len(surface_parts))) + surface_parts.append(big_copy[seg_starts[i] : seg_ends[i]]) + + # Mark submobs as seen so they aren't re-processed as VMobjects. + for submob in surface_submobs: + _seen_submobs.add(id(submob)) + continue + + # ── Full MISS: tessellation + smoothing (original path) ─────── + # Collect (stroke_width, stroke_color_alpha) per part so we can + # recompute stroke_half_px on future cache hits. + new_parts_start = len(surface_parts) + stroke_per_part_new: list[tuple[float, float]] = [] + + for submob in surface_submobs: + if id(submob) in _seen_submobs: + continue + _seen_submobs.add(id(submob)) + stroke_rgba_sub = submob.get_stroke_rgbas() + sw_sub = ( + float(submob.get_stroke_width()) + if stroke_rgba_sub.shape[0] > 0 + else 0.0 + ) + s_alpha_sub = ( + float(stroke_rgba_sub[0, 3]) + if stroke_rgba_sub.shape[0] > 0 + else 0.0 + ) + data = _collect_surface_geometry( + submob, + view_matrix, + proj_matrix, + diffuse_strength=float( + getattr(submob, "diffuse_strength", surf_diffuse) + ), + specular_strength=float( + getattr(submob, "specular_strength", surf_specular) + ), + specular_exponent=float( + getattr(submob, "specular_exponent", surf_spec_exp) + ), + ) + if data is not None: + cls = _surface_opacity_class(data) + cmd = "surface_opaque" if cls == "opaque" else "surface_oit" + draw_plan.append((cmd, len(surface_parts))) + surface_parts.append(data) + stroke_per_part_new.append((sw_sub, s_alpha_sub)) + + # Record this mob so we can cache its smoothed parts later. + _new_surface_mobs.append( + ( + mob, + geom_hash, + color_hash, + new_parts_start, + len(surface_parts), + stroke_per_part_new, + ) + ) + continue + + # ── Regular VMobject (2-D or shade_in_3d) ──────────────────────── + vmob_submobs = mob.family_members_with_points() + if use_z_index: + vmob_submobs = sorted(vmob_submobs, key=lambda m: m.z_index) + for submob in vmob_submobs: + if id(submob) in _seen_submobs: + continue + _seen_submobs.add(id(submob)) + phash = _points_hash(submob) + cached = _fill_stroke_cache.get(submob) + if cached is None or cached[0] != phash: + result = _collect_cubics(submob) + if result is not None: + _fill_stroke_cache[submob] = (phash, result) + else: + _fill_stroke_cache.pop(submob, None) + continue + cached = _fill_stroke_cache.get(submob) + if cached is None: + continue + fill_cubics, stroke_cubics = cached[1] + + # Fixed-orientation pre-transform: translate control points so the + # submob appears at its full 3D-projected center position while + # preserving local orientation (no rotation of the local shape). + # + # Cairo's equivalent: transform_points_pre_display() computes + # new_center = project_point(center) (full camera rotation) + # points = points + (new_center - center) + # i.e. translate all points by the difference between the + # camera-space center and the world-space center. + # + # In WebGPU, with t_full == t_fixed == [0,0,-11], the offset + # simplifies to: + # offset = R_full @ center_w - center_w = (R_full - I) @ center_w + # After adding this offset, the GPU shader applies fixed_view + # (identity rotation + z-translation), giving: + # view_pos = (point_w + offset) + [0,0,-11] + # = (point_w - center_w) + (R_full @ center_w + [0,0,-11]) + # which is the local shape centred at the full-projection center. ✓ + if center_view_matrix is not None: + R_full = center_view_matrix[:3, :3].astype(np.float32) + c_w = submob.get_center().astype(np.float32) + offset = R_full @ c_w - c_w # shape (3,) + fill_cubics = fill_cubics + offset # broadcast (N,4,3)+(3,) + stroke_cubics = stroke_cubics + offset + + # Fetch current colors every frame (they change during animations). + fill_rgba = submob.get_fill_rgbas() + stroke_rgba = submob.get_stroke_rgbas() + fill_color = ( + fill_rgba[0].astype(np.float32) + if fill_rgba.shape[0] > 0 + else np.zeros(4, dtype=np.float32) + ) + stroke_color = ( + stroke_rgba[0].astype(np.float32) + if stroke_rgba.shape[0] > 0 + else np.zeros(4, dtype=np.float32) + ) + stroke_width = ( + float(submob.get_stroke_width()) if stroke_rgba.shape[0] > 0 else 0.0 + ) + + # Skip entirely invisible objects (both fill and stroke transparent). + if fill_color[3] < 0.001 and ( + stroke_color[3] < 0.001 or stroke_width < 0.001 + ): + continue + + # 0 = nonzero (default), 1 = evenodd (set by SVG parser) + fill_rule = int(getattr(submob, "fill_rule", 0)) + + # Gradient fill: pass all colour stops and gradient axis endpoints. + gradient_start = gradient_end = None + if fill_rgba.shape[0] > 1: + try: + gradient_start, gradient_end = ( + submob.get_gradient_start_and_end_points() + ) + except Exception: + pass # fall back to solid fill_color[0] + + # Gradient stroke: same pattern as fill gradient. + stroke_gradient_start = stroke_gradient_end = None + if stroke_rgba.shape[0] > 1: + try: + stroke_gradient_start, stroke_gradient_end = ( + submob.get_gradient_start_and_end_points() + ) + except Exception: + pass # fall back to solid stroke_color[0] + + # Build bounding quad with placeholder curve indices. + quad_verts = _build_fill_stroke_quad( + fill_cubics=fill_cubics, + stroke_cubics=stroke_cubics, + fill_color=fill_color, + stroke_color=stroke_color, + stroke_width=stroke_width, + fill_curve_start=0, # assigned below after all objects are collected + stroke_curve_start=0, # assigned below + view_matrix=view_matrix, + proj_matrix=proj_matrix, + fill_rule=fill_rule, + fill_rgbas=fill_rgba if fill_rgba.shape[0] > 1 else None, + gradient_start=gradient_start, + gradient_end=gradient_end, + stroke_rgbas=stroke_rgba if stroke_rgba.shape[0] > 1 else None, + stroke_gradient_start=stroke_gradient_start, + stroke_gradient_end=stroke_gradient_end, + ) + if len(quad_verts) == 0: + continue + + is_3d = getattr(submob, "shade_in_3d", False) + draw_plan.append( + ("fill_stroke_3d" if is_3d else "fill_stroke_2d", len(fs_parts)) + ) + fs_parts.append(quad_verts) + all_fill_cubics.append(fill_cubics) + all_stroke_cubics.append(stroke_cubics) + n_fill_cubics_per.append(len(fill_cubics)) + n_stroke_cubics_per.append(len(stroke_cubics)) + + if not draw_plan: + return None + + device: wgpu_t.GPUDevice = renderer.device + + # ── Assign global curve start indices ──────────────────────────────── + # Cubics buffer layout: [fill_cubics_obj0, fill_cubics_obj1, ..., + # stroke_cubics_obj0, stroke_cubics_obj1, ...] + # Quads output layout: [fill_quads_obj0, fill_quads_obj1, ..., + # stroke_quads_obj0, stroke_quads_obj1, ...] + total_fill_cubics = sum(n_fill_cubics_per) + total_stroke_cubics = sum(n_stroke_cubics_per) + n_cubics_total = total_fill_cubics + total_stroke_cubics + + fill_global = 0 # running fill cubic index + stroke_global = total_fill_cubics # stroke cubics follow all fill cubics + + for i, part in enumerate(fs_parts): + part["fill_curve_start"] = fill_global * 4 + part["n_fill_curves"] = n_fill_cubics_per[i] * 4 + part["stroke_curve_start"] = stroke_global * 4 + part["n_stroke_curves"] = n_stroke_cubics_per[i] * 4 + fill_global += n_fill_cubics_per[i] + stroke_global += n_stroke_cubics_per[i] + + # ── Upload vertex data ─────────────────────────────────────────────── + fs_buf, fs_byte_offsets = None, [] + if fs_parts: + fs_buf, fs_byte_offsets = _batch_upload(device, fs_parts) + renderer.frame_vbos.append(fs_buf) + + # ── Upload cubics and create compute/render bind groups ────────────── + cubics_buf = quads_out_buf = compute_bg = render_bg = None + + if n_cubics_total > 0: + # Build flat float32 array: [all fill cubics..., all stroke cubics...] + fill_arrays = [c for c in all_fill_cubics if len(c) > 0] + stroke_arrays = [c for c in all_stroke_cubics if len(c) > 0] + all_arrays = fill_arrays + stroke_arrays + all_cubics = np.concatenate(all_arrays, axis=0) # (N, 4, 3) + cubics_flat = all_cubics.astype(np.float32).ravel() # N*12 floats + + cubics_buf = device.create_buffer_with_data( + data=cubics_flat.tobytes(), + usage=wgpu.BufferUsage.STORAGE, + ) + renderer.frame_vbos.append(cubics_buf) + + quads_size = n_cubics_total * 36 * 4 # 4 quads × 9 floats × 4 bytes + quads_out_buf = device.create_buffer( + size=max(quads_size, 16), # WebGPU minimum binding size + usage=wgpu.BufferUsage.STORAGE, + ) + renderer.frame_vbos.append(quads_out_buf) + + # Params uniform (n_cubics, padded to 16 bytes for WebGPU alignment). + params_bytes = struct.pack("<4I", n_cubics_total, 0, 0, 0) + params_buf = device.create_buffer_with_data( + data=params_bytes, + usage=wgpu.BufferUsage.UNIFORM, + ) + renderer.frame_vbos.append(params_buf) + + compute_bg = device.create_bind_group( + layout=renderer._compute_bgl, + entries=[ + { + "binding": 0, + "resource": { + "buffer": cubics_buf, + "offset": 0, + "size": cubics_buf.size, + }, + }, + { + "binding": 1, + "resource": { + "buffer": quads_out_buf, + "offset": 0, + "size": quads_out_buf.size, + }, + }, + { + "binding": 2, + "resource": {"buffer": params_buf, "offset": 0, "size": 16}, + }, + ], + ) + + render_bg = device.create_bind_group( + layout=renderer._fill_stroke_bgl, + entries=[ + { + "binding": 0, + "resource": { + "buffer": camera_uniform_buf, + "offset": 0, + "size": camera_uniform_buf.size, + }, + }, + { + "binding": 1, + "resource": { + "buffer": quads_out_buf, + "offset": 0, + "size": quads_out_buf.size, + }, + }, + ], + ) + + # ── Upload surface data ────────────────────────────────────────────── + # Apply normal smoothing only to parts from cache-miss mobs; cached + # parts already carry correctly smoothed normals. + surface_buf, surface_byte_offsets = None, [] + if surface_parts: + if _new_surface_mobs: + # Smooth normals for newly-tessellated slices. + # _new_surface_mobs entries: + # (mob, geom_hash, color_hash, start, end, stroke_per_part_new) + new_slices: list[np.ndarray] = [] + for mob, geom_hash, color_hash, start, end, _ in _new_surface_mobs: + new_slices.extend(surface_parts[start:end]) + _smooth_surface_normals(new_slices) + + # Cache each newly-tessellated mob's smoothed parts. + for ( + mob, + geom_hash, + color_hash, + start, + end, + stroke_per_part_new, + ) in _new_surface_mobs: + parts_for_mob = surface_parts[start:end] + if not parts_for_mob: + continue + # Collect draw commands for this mob's global part indices. + # Must filter on cmd type to exclude VMobject ("fill_stroke_*") + # entries: draw_plan is shared between VMobject and Surface paths, + # and both index from 0 (fs_parts vs surface_parts respectively), + # so index-range-only filtering incorrectly includes VMobject entries + # whose fs_parts index happens to fall inside [start, end). + draw_cmds_for_mob = [ + cmd + for cmd, idx in draw_plan + if start <= idx < end and cmd in ("surface_opaque", "surface_oit") + ] + # Cache as a SINGLE concatenated array so a hit can copy the + # whole mob's geometry in one numpy operation. stroke_half_px + # is zeroed in the template; it is recomputed on every hit. + # Colors in big_template reflect the current frame's colors so + # that a color-only miss can patch them in O(N_parts). + big_template = np.concatenate(parts_for_mob, axis=0) + big_template["stroke_half_px"] = 0.0 + # Part boundary offsets as numpy arrays (avoid Python list ops on hit). + sizes = np.array([len(p) for p in parts_for_mob], dtype=np.intp) + starts = np.concatenate([[0], np.cumsum(sizes[:-1])]).astype(np.intp) + ends = starts + sizes + # Precompute stroke metadata as numpy arrays for vectorised hit path. + sw_arr_c = np.array( + [sw for sw, _ in stroke_per_part_new], dtype=np.float32 + ) + sa_arr_c = np.array( + [alpha for _, alpha in stroke_per_part_new], dtype=np.float32 + ) + has_sw_c = (sw_arr_c > 0.0) & (sa_arr_c > 0.001) + _surface_mob_cache[mob] = ( + geom_hash, + color_hash, + big_template, + starts, + ends, + sw_arr_c, + sa_arr_c, + has_sw_c, + draw_cmds_for_mob, + ) + + surface_buf, surface_byte_offsets = _batch_upload(device, surface_parts) + renderer.frame_vbos.append(surface_buf) + + oit_indices = [idx for cmd, idx in draw_plan if cmd == "surface_oit"] + + result = _FrameData( + fs_parts=fs_parts, + fs_buf=fs_buf, + fs_byte_offsets=fs_byte_offsets, + cubics_buf=cubics_buf, + quads_out_buf=quads_out_buf, + n_cubics_total=n_cubics_total, + compute_bg=compute_bg, + render_bg=render_bg, + surface_parts=surface_parts, + surface_buf=surface_buf, + surface_byte_offsets=surface_byte_offsets, + draw_plan=draw_plan, + oit_indices=oit_indices, + ) + + # Store in cache so the NEXT frame can skip tessellation on a fingerprint hit. + # frame_vbos are NOT added for cached buffers — the cache itself is the owner. + # Remove the just-uploaded buffers from frame_vbos so they aren't released at + # end-of-frame (the cache needs them to survive across frames). + # _all_surface_call paths are excluded: the surface GPU buffer changes every + # frame (stroke_half_px update), so the _FrameData cache cannot help there. + if cache_slot is not None and not _all_surface_call: + cached_bufs = { + id(result.fs_buf), + id(result.cubics_buf), + id(result.quads_out_buf), + id(result.surface_buf), + } - {id(None)} + renderer.frame_vbos = [ + b for b in renderer.frame_vbos if id(b) not in cached_bufs + ] + renderer._fd_cache[cache_slot] = (fp, result) + + return result + + +def draw_frame_data( + renderer: WebGPURenderer, + fd: _FrameData, + cam_bg: wgpu_t.GPUBindGroup, +) -> None: + """Record draw commands for *fd* into ``renderer.current_render_pass``. + + Draw order + ---------- + 1. 2-D fill+stroke objects — interleaved in ``draw_plan`` order (painter's + algorithm; no depth write so objects paint over each other correctly). + 2. 3-D fill+stroke objects — depth write + test (shade_in_3d). + 3. Opaque parametric surfaces — depth write (includes barycentric wireframe). + + OIT surfaces are NOT drawn here; the caller reads ``fd.oit_indices`` and + handles them in a separate accumulation pass. + """ + rp = renderer.current_render_pass + + # ── 1. 2-D fill+stroke: painter's algorithm ─────────────────────────── + # All 2-D quads live in fs_buf in their original draw_plan order. + # Instead of N separate set_vertex_buffer+draw calls we issue one draw + # per *contiguous run* of "fill_stroke_2d" entries, dramatically reducing + # the number of wgpu API calls (typically from 800 to 1 for a pure-2D scene). + if fd.fs_buf is not None and fd.render_bg is not None: + rp.set_pipeline(renderer.fill_stroke_pipeline) + rp.set_bind_group(0, fd.render_bg, [], 0, 0) + rp.set_vertex_buffer(0, fd.fs_buf) + + # Walk draw_plan and batch consecutive fill_stroke_2d entries. + run_first_vertex: int = -1 + run_vertex_count: int = 0 + + for cmd, idx in fd.draw_plan: + if cmd != "fill_stroke_2d": + # Flush the current 2D run (if any) before breaking the batch. + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_vertex_count = 0 + run_first_vertex = -1 + continue + + arr = fd.fs_parts[idx] + byte_offset = fd.fs_byte_offsets[idx] + first_vert = byte_offset // _FILL_STROKE_STRIDE + + if run_first_vertex < 0: + # Start a new run. + run_first_vertex = first_vert + run_vertex_count = len(arr) + elif first_vert == run_first_vertex + run_vertex_count: + # Extend the current contiguous run. + run_vertex_count += len(arr) + else: + # Gap in the buffer (shouldn't happen for pure-2D scenes but + # guard anyway). Flush old run, start new one. + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_first_vertex = first_vert + run_vertex_count = len(arr) + + # Flush final run. + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + + # ── 2. 3-D fill+stroke: depth-tested and depth-written ──────────────── + # Same batching strategy for shade_in_3d VMobjects. + if fd.fs_buf is not None and fd.render_bg is not None: + rp.set_pipeline(renderer.fill_stroke_3d_pipeline) + rp.set_bind_group(0, fd.render_bg, [], 0, 0) + rp.set_vertex_buffer(0, fd.fs_buf) + + run_first_vertex = -1 + run_vertex_count = 0 + + for cmd, idx in fd.draw_plan: + if cmd != "fill_stroke_3d": + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_vertex_count = 0 + run_first_vertex = -1 + continue + + arr = fd.fs_parts[idx] + byte_offset = fd.fs_byte_offsets[idx] + first_vert = byte_offset // _FILL_STROKE_STRIDE + + if run_first_vertex < 0: + run_first_vertex = first_vert + run_vertex_count = len(arr) + elif first_vert == run_first_vertex + run_vertex_count: + run_vertex_count += len(arr) + else: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_first_vertex = first_vert + run_vertex_count = len(arr) + + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + + # ── 3. Opaque parametric surfaces ───────────────────────────────────── + if fd.surface_buf is not None: + rp.set_pipeline(renderer.surface_pipeline) + rp.set_bind_group(0, cam_bg, [], 0, 0) + rp.set_vertex_buffer(0, fd.surface_buf) + + run_first_vertex = -1 + run_vertex_count = 0 + + for cmd, idx in fd.draw_plan: + if cmd != "surface_opaque": + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_vertex_count = 0 + run_first_vertex = -1 + continue + + arr = fd.surface_parts[idx] + byte_offset = fd.surface_byte_offsets[idx] + first_vert = byte_offset // _SURFACE_COMBINED_STRIDE + + if run_first_vertex < 0: + run_first_vertex = first_vert + run_vertex_count = len(arr) + elif first_vert == run_first_vertex + run_vertex_count: + run_vertex_count += len(arr) + else: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_first_vertex = first_vert + run_vertex_count = len(arr) + + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + + +def draw_frame_data_subcam( + rp: Any, + fd: _FrameData, + sub_fill_render_bg: Any, + sub_cam_bg: Any, + fill_2d_pipeline: Any, + fill_3d_pipeline: Any, + surf_pipeline: Any, +) -> None: + """Draw *fd* into render pass *rp* using sub-camera pipelines and bind groups. + + Mirrors ``draw_frame_data`` but accepts explicit pipeline objects and bind + groups instead of reading them from the renderer. Used by + ``_render_sub_camera_pass`` to reuse cached geometry (quads buffer, vertex + buffer, surface buffer) with a different camera uniform. + + Parameters + ---------- + rp + Active render pass encoder targeting the sub-camera texture. + fd + Cached geometry from the main frame's ``collect_frame_data`` call. + sub_fill_render_bg + Bind group with sub-camera uniform (binding 0) + quads storage (binding 1). + Replaces ``fd.render_bg`` for fill-stroke draw calls. + sub_cam_bg + Camera-only bind group with sub-camera uniform (binding 0). + Used for surface draw calls. + fill_2d_pipeline / fill_3d_pipeline / surf_pipeline + Sub-camera render pipelines targeting the sub-camera texture format. + """ + # ── 2-D fill+stroke ─────────────────────────────────────────────────── + if fd.fs_buf is not None and sub_fill_render_bg is not None: + rp.set_pipeline(fill_2d_pipeline) + rp.set_bind_group(0, sub_fill_render_bg, [], 0, 0) + rp.set_vertex_buffer(0, fd.fs_buf) + + run_first_vertex: int = -1 + run_vertex_count: int = 0 + for cmd, idx in fd.draw_plan: + if cmd != "fill_stroke_2d": + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_vertex_count = 0 + run_first_vertex = -1 + continue + arr = fd.fs_parts[idx] + byte_offset = fd.fs_byte_offsets[idx] + first_vert = byte_offset // _FILL_STROKE_STRIDE + if run_first_vertex < 0: + run_first_vertex = first_vert + run_vertex_count = len(arr) + elif first_vert == run_first_vertex + run_vertex_count: + run_vertex_count += len(arr) + else: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_first_vertex = first_vert + run_vertex_count = len(arr) + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + + # ── 3-D fill+stroke ─────────────────────────────────────────────────── + if fd.fs_buf is not None and sub_fill_render_bg is not None: + rp.set_pipeline(fill_3d_pipeline) + rp.set_bind_group(0, sub_fill_render_bg, [], 0, 0) + rp.set_vertex_buffer(0, fd.fs_buf) + + run_first_vertex = -1 + run_vertex_count = 0 + for cmd, idx in fd.draw_plan: + if cmd != "fill_stroke_3d": + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_vertex_count = 0 + run_first_vertex = -1 + continue + arr = fd.fs_parts[idx] + byte_offset = fd.fs_byte_offsets[idx] + first_vert = byte_offset // _FILL_STROKE_STRIDE + if run_first_vertex < 0: + run_first_vertex = first_vert + run_vertex_count = len(arr) + elif first_vert == run_first_vertex + run_vertex_count: + run_vertex_count += len(arr) + else: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_first_vertex = first_vert + run_vertex_count = len(arr) + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + + # ── Opaque surfaces ─────────────────────────────────────────────────── + if fd.surface_buf is not None and surf_pipeline is not None: + rp.set_pipeline(surf_pipeline) + rp.set_bind_group(0, sub_cam_bg, [], 0, 0) + rp.set_vertex_buffer(0, fd.surface_buf) + + run_first_vertex = -1 + run_vertex_count = 0 + for cmd, idx in fd.draw_plan: + if cmd != "surface_opaque": + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_vertex_count = 0 + run_first_vertex = -1 + continue + arr = fd.surface_parts[idx] + byte_offset = fd.surface_byte_offsets[idx] + first_vert = byte_offset // _SURFACE_COMBINED_STRIDE + if run_first_vertex < 0: + run_first_vertex = first_vert + run_vertex_count = len(arr) + elif first_vert == run_first_vertex + run_vertex_count: + run_vertex_count += len(arr) + else: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + run_first_vertex = first_vert + run_vertex_count = len(arr) + if run_vertex_count > 0: + rp.draw(run_vertex_count, 1, run_first_vertex, 0) + + +# --------------------------------------------------------------------------- +# GPU upload helpers +# --------------------------------------------------------------------------- + + +def _batch_upload( + device: wgpu_t.GPUDevice, + arrays: list[np.ndarray], +) -> tuple[wgpu_t.GPUBuffer, list[int]]: + """Concatenate *arrays* into one bytes blob and upload as a VERTEX buffer.""" + import wgpu + + byte_offsets: list[int] = [] + parts: list[bytes] = [] + offset = 0 + for arr in arrays: + byte_offsets.append(offset) + b = arr.tobytes() + parts.append(b) + offset += len(b) + + buf = device.create_buffer_with_data( + data=b"".join(parts), + usage=wgpu.BufferUsage.VERTEX, + ) + return buf, byte_offsets + + +# --------------------------------------------------------------------------- +# VMobject cubic collector — geometry cache +# --------------------------------------------------------------------------- + + +def _collect_cubics( + vmobject: VMobject, +) -> tuple[np.ndarray, np.ndarray] | None: + """Return ``(fill_cubics, stroke_cubics)`` for the GPU compute shader. + + *fill_cubics* — ``(N, 4, 3)`` float32. All subpath cubics **plus** one + linear closing cubic per open subpath (required for correct winding- + number coverage during partial animations such as ``Create``). + + *stroke_cubics* — ``(M, 4, 3)`` float32. Only the actual subpath cubics, + no closing segment — the stroke should follow the visible curve only. + + Colors are NOT stored here; they are fetched fresh every frame in + ``collect_frame_data`` so that opacity animations work correctly. + + Returns ``None`` if the vmobject has no usable bezier curves. + """ + nppcc = vmobject.n_points_per_cubic_curve + + fill_cubics_list: list[np.ndarray] = [] + stroke_cubics_list: list[np.ndarray] = [] + + for subpath in vmobject.get_subpaths(): + n_curves = len(subpath) // nppcc + if n_curves == 0: + continue + pts = subpath[: n_curves * nppcc] + b0s = pts[0::nppcc].astype(np.float32) + h0s = pts[1::nppcc].astype(np.float32) + h1s = pts[2::nppcc].astype(np.float32) + b3s = pts[3::nppcc].astype(np.float32) + + cubics = np.stack([b0s, h0s, h1s, b3s], axis=1) # (n, 4, 3) + stroke_cubics_list.append(cubics) + fill_cubics_list.append(cubics) + + # Closing segment: linear cubic from the last anchor back to the first. + # Degree-elevation from a line (last→first) to a cubic: + # b0 = last, b1 = last + (first-last)/3, + # b2 = last + 2*(first-last)/3, b3 = first. + first = b0s[0] + last = b3s[-1] + if not np.allclose(first, last, atol=1e-6): + diff = first - last + closing = np.array( + [[last, last + diff * (1.0 / 3.0), last + diff * (2.0 / 3.0), first]], + dtype=np.float32, + ) + fill_cubics_list.append(closing) + + if not fill_cubics_list and not stroke_cubics_list: + return None + + fill_cubics = ( + np.concatenate(fill_cubics_list, axis=0) + if fill_cubics_list + else np.empty((0, 4, 3), dtype=np.float32) + ) + stroke_cubics = ( + np.concatenate(stroke_cubics_list, axis=0) + if stroke_cubics_list + else np.empty((0, 4, 3), dtype=np.float32) + ) + return fill_cubics, stroke_cubics + + +# --------------------------------------------------------------------------- +# Bounding-quad builder +# --------------------------------------------------------------------------- + + +def _build_fill_stroke_quad( + fill_cubics: np.ndarray, + stroke_cubics: np.ndarray, + fill_color: np.ndarray, + stroke_color: np.ndarray, + stroke_width: float, + fill_curve_start: int, + stroke_curve_start: int, + view_matrix: np.ndarray, + proj_matrix: np.ndarray, + fill_rule: int = 0, + fill_rgbas: np.ndarray | None = None, + gradient_start: np.ndarray | None = None, + gradient_end: np.ndarray | None = None, + stroke_rgbas: np.ndarray | None = None, + stroke_gradient_start: np.ndarray | None = None, + stroke_gradient_end: np.ndarray | None = None, +) -> np.ndarray: + """Build a ``_FILL_STROKE_DTYPE`` bounding quad (6 vertices) for one object. + + The bounding box is computed in NDC space (clip.xy / clip.w) from the + anchor points of both fill and stroke cubics, then mapped back to world + space at the average view-space Z. This is correct for both orthographic + (w = 1) and perspective projections. + + *stroke_half_ndc* is the stroke half-width in NDC units, computed from + the current projection matrix and average clip-w so that stroke width is + consistent across perspective depths. + + *fill_rgbas* — if provided and has more than one row, enables gradient fill. + *gradient_start* / *gradient_end* — world-space endpoints of the fill gradient axis. + *stroke_rgbas* — if provided and has more than one row, enables gradient stroke. + *stroke_gradient_start* / *stroke_gradient_end* — world-space endpoints of the stroke gradient axis. + """ + # Gather all anchor points (b0 and b3 of every cubic). + anchor_lists: list[np.ndarray] = [] + if len(fill_cubics) > 0: + anchor_lists.append(fill_cubics[:, 0]) + anchor_lists.append(fill_cubics[:, 3]) + if len(stroke_cubics) > 0: + anchor_lists.append(stroke_cubics[:, 0]) + anchor_lists.append(stroke_cubics[:, 3]) + + if not anchor_lists: + return np.empty(0, dtype=_FILL_STROKE_DTYPE) + + anchors = np.concatenate(anchor_lists, axis=0).astype(np.float32) # (N, 3) + + vm = view_matrix.astype(np.float32) + pm = proj_matrix.astype(np.float32) + R, t = vm[:3, :3], vm[:3, 3] + + pts_v = (R @ anchors.T).T + t # (N, 3) view space + avg_z_v = float(pts_v[:, 2].mean()) + + # Perspective divide → NDC. + ones = np.ones((len(pts_v), 1), dtype=np.float32) + clips = (pm @ np.hstack([pts_v, ones]).T).T # (N, 4) + w = clips[:, 3:4] + w_s = np.where(np.abs(w) > 1e-8, w, np.sign(w + 1e-38) * 1e-8) + ndcs = clips[:, :2] / w_s # (N, 2) NDC + + PAD = 0.05 + ndc_min = ndcs.min(axis=0) - PAD + ndc_max = ndcs.max(axis=0) + PAD + + # Stroke half-width in NDC. + # v_thickness = 0.004 * stroke_width (view-space, matching vmobject_stroke.wgsl) + # stroke_half_ndc = v_thickness * pm[0,0] / avg_clip_w + # where avg_clip_w = pm[3,2]*avg_z + pm[3,3] + avg_clip_w = float(pm[3, 2] * avg_z_v + pm[3, 3]) + avg_clip_w = avg_clip_w if abs(avg_clip_w) > 1e-8 else 1.0 + stroke_half_ndc = 0.0 + if stroke_width > 0.0 and float(stroke_color[3]) > 0.001: + stroke_half_ndc = float(0.004 * stroke_width * abs(pm[0, 0]) / abs(avg_clip_w)) + # Add stroke padding so the bounding quad covers the stroke edges. + ndc_min -= stroke_half_ndc * 2.0 + ndc_max += stroke_half_ndc * 2.0 + + # Invert NDC bounding corners to view space. + inv_px = 1.0 / (pm[0, 0] if abs(pm[0, 0]) > 1e-8 else 1.0) + inv_py = 1.0 / (pm[1, 1] if abs(pm[1, 1]) > 1e-8 else 1.0) + x0_v = (float(ndc_min[0]) * avg_clip_w - float(pm[0, 3])) * inv_px + x1_v = (float(ndc_max[0]) * avg_clip_w - float(pm[0, 3])) * inv_px + y0_v = (float(ndc_min[1]) * avg_clip_w - float(pm[1, 3])) * inv_py + y1_v = (float(ndc_max[1]) * avg_clip_w - float(pm[1, 3])) * inv_py + + corners_v = np.array( + [ + [x0_v, y0_v, avg_z_v], + [x1_v, y0_v, avg_z_v], + [x0_v, y1_v, avg_z_v], + [x1_v, y1_v, avg_z_v], + ], + dtype=np.float32, + ) + R_inv = R.T + t_inv = -(R_inv @ t) + corners_w = (R_inv @ corners_v.T).T + t_inv # (4, 3) world space + quad_pos = corners_w[[0, 1, 2, 1, 3, 2]] # (6, 3) two CCW triangles + + # ── Per-vertex fill colours (gradient support) ──────────────────────────── + # If fill_rgbas has >1 colour row, interpolate along the gradient axis. + # gradient_start / gradient_end are world-space endpoints of the axis. + if ( + fill_rgbas is not None + and fill_rgbas.shape[0] > 1 + and gradient_start is not None + and gradient_end is not None + ): + gs = np.asarray(gradient_start, dtype=np.float32) + ge = np.asarray(gradient_end, dtype=np.float32) + axis = ge - gs + axis_len2 = float(np.dot(axis, axis)) + if axis_len2 > 1e-12: + # Project each of the 4 corners onto the gradient axis → t ∈ [0, 1]. + t_corners = np.clip( + np.dot(corners_w - gs, axis) / axis_len2, 0.0, 1.0 + ) # (4,) + n_stops = fill_rgbas.shape[0] + # Interpolate: t_corners maps to colour stop indices. + idx_f = t_corners * (n_stops - 1) # float indices + idx_lo = np.floor(idx_f).astype(int).clip(0, n_stops - 2) + idx_hi = idx_lo + 1 + frac = (idx_f - idx_lo)[:, None] # (4, 1) + corner_colors = ( + fill_rgbas[idx_lo].astype(np.float32) * (1.0 - frac) + + fill_rgbas[idx_hi].astype(np.float32) * frac + ) # (4, 4) + # Map corners [0,1,2,3] → quad vertices [0,1,2,1,3,2]. + per_vertex_fill = corner_colors[[0, 1, 2, 1, 3, 2]] # (6, 4) + else: + per_vertex_fill = np.broadcast_to(fill_color, (6, 4)).copy() + else: + per_vertex_fill = np.broadcast_to(fill_color, (6, 4)).copy() + + # ── Per-vertex stroke colours (gradient support) ───────────────────────── + if ( + stroke_rgbas is not None + and stroke_rgbas.shape[0] > 1 + and stroke_gradient_start is not None + and stroke_gradient_end is not None + ): + sgs = np.asarray(stroke_gradient_start, dtype=np.float32) + sge = np.asarray(stroke_gradient_end, dtype=np.float32) + s_axis = sge - sgs + s_axis_len2 = float(np.dot(s_axis, s_axis)) + if s_axis_len2 > 1e-12: + t_corners = np.clip( + np.dot(corners_w - sgs, s_axis) / s_axis_len2, 0.0, 1.0 + ) # (4,) + n_stops = stroke_rgbas.shape[0] + idx_f = t_corners * (n_stops - 1) + idx_lo = np.floor(idx_f).astype(int).clip(0, n_stops - 2) + idx_hi = idx_lo + 1 + frac = (idx_f - idx_lo)[:, None] + corner_colors = ( + stroke_rgbas[idx_lo].astype(np.float32) * (1.0 - frac) + + stroke_rgbas[idx_hi].astype(np.float32) * frac + ) # (4, 4) + per_vertex_stroke = corner_colors[[0, 1, 2, 1, 3, 2]] # (6, 4) + else: + per_vertex_stroke = np.broadcast_to(stroke_color, (6, 4)).copy() + else: + per_vertex_stroke = np.broadcast_to(stroke_color, (6, 4)).copy() + + n_fill_quads = len(fill_cubics) * 4 # 4 quadratics per cubic + n_stroke_quads = len(stroke_cubics) * 4 + + verts = np.empty(6, dtype=_FILL_STROKE_DTYPE) + verts["in_pos"] = quad_pos + verts["in_fill_color"] = per_vertex_fill + verts["in_stroke_color"] = per_vertex_stroke + verts["stroke_half_ndc"] = stroke_half_ndc + verts["fill_curve_start"] = fill_curve_start + verts["n_fill_curves"] = n_fill_quads + verts["stroke_curve_start"] = stroke_curve_start + verts["n_stroke_curves"] = n_stroke_quads + verts["fill_rule"] = fill_rule + return verts + + +# --------------------------------------------------------------------------- +# Surface geometry collectors (unchanged from original) +# --------------------------------------------------------------------------- + + +def _surface_opacity_class(part: np.ndarray) -> str: + alphas = part["in_fill_color"][:, 3] + return "opaque" if float(alphas.min()) >= 0.99 else "oit" + + +def _collect_surface_geometry( + vmobject: VMobject, + view_matrix: np.ndarray, + proj_matrix: np.ndarray, + diffuse_strength: float = 0.8, + specular_strength: float = 0.9, + specular_exponent: float = 16.0, +) -> np.ndarray | None: + """Return a ``_SURFACE_COMBINED_DTYPE`` array for a shade_in_3d VMobject. + + Material parameters are passed from the parent :class:`~.Surface` so + that ``diffuse_strength``, ``specular_strength``, and + ``specular_exponent`` set on the parent are applied to every submobject + patch. + + Barycentric coordinates are assigned per triangle in the centroid fan: + centroid → bary = (1, 0, 0) (bary.x = 0 on outer edge) + anchor_i → bary = (0, 1, 0) + anchor_{i+1} → bary = (0, 0, 1) + + ``stroke_half_px`` is computed from the stroke width, projection matrix + and average clip-w of the surface anchors so that wireframe line width + is consistent across perspective depths. + """ + from manim import config + + fill_rgba = vmobject.get_fill_rgbas() + if fill_rgba.shape[0] == 0 or fill_rgba[0, 3] == 0: + return None + + fill_color = fill_rgba[0].astype(np.float32) + stroke_rgba = vmobject.get_stroke_rgbas() + stroke_color = ( + stroke_rgba[0].astype(np.float32) + if stroke_rgba.shape[0] > 0 + else np.zeros(4, dtype=np.float32) + ) + stroke_width = ( + float(vmobject.get_stroke_width()) if stroke_rgba.shape[0] > 0 else 0.0 + ) + + nppcc = vmobject.n_points_per_cubic_curve + + all_verts: list[np.ndarray] = [] + all_normals: list[np.ndarray] = [] + all_bary: list[np.ndarray] = [] + + for subpath in vmobject.get_subpaths(): + n_curves = len(subpath) // nppcc + if n_curves < 2: + continue + anchors = subpath[0::nppcc] + last = subpath[n_curves * nppcc - 1 : n_curves * nppcc] + if len(last) and not np.allclose(anchors[-1], last[0], atol=1e-6): + anchors = np.vstack([anchors, last]) + + n_pts = len(anchors) + if n_pts < 3: + continue + + centroid = anchors.mean(axis=0) + v0 = anchors[0] - centroid + v1 = anchors[1] - centroid + raw_normal = np.cross(v1, v0).astype(np.float64) + norm_len = np.linalg.norm(raw_normal) + normal = ( + (raw_normal / norm_len).astype(np.float32) + if norm_len > 1e-9 + else np.array([0.0, 0.0, 1.0], dtype=np.float32) + ) + + # Triangle fan: (centroid, anchor_i, anchor_{i+1}) + fan_verts = np.empty((n_pts * 3, 3), dtype=np.float32) + fan_verts[0::3] = centroid.astype(np.float32) + fan_verts[1::3] = anchors.astype(np.float32) + fan_verts[2::3] = np.roll(anchors, -1, axis=0).astype(np.float32) + + # Barycentric coords: centroid=(1,0,0), anchor_i=(0,1,0), next=(0,0,1) + bary_block = np.zeros((n_pts * 3, 3), dtype=np.float32) + bary_block[0::3] = [1.0, 0.0, 0.0] + bary_block[1::3] = [0.0, 1.0, 0.0] + bary_block[2::3] = [0.0, 0.0, 1.0] + + all_verts.append(fan_verts) + all_normals.append(np.tile(normal, (n_pts * 3, 1))) + all_bary.append(bary_block) + + if not all_verts: + return None + + verts = np.concatenate(all_verts, axis=0) + normals = np.concatenate(all_normals, axis=0) + bary = np.concatenate(all_bary, axis=0) + n_total = len(verts) + + # Compute stroke_half_px: half the wireframe line width in screen pixels. + # Formula matches _build_fill_stroke_quad: 0.004 * width * |pm[0,0]| / |avg_clip_w| + # then multiplied by pixel_width/2 to convert NDC to pixels. + stroke_half_px = 0.0 + if stroke_width > 0.0 and float(stroke_color[3]) > 0.001: + pm = proj_matrix.astype(np.float32) + vm = view_matrix.astype(np.float32) + R, t = vm[:3, :3], vm[:3, 3] + pts_v = (R @ verts.T).T + t # (N, 3) view space + avg_z_v = float(pts_v[:, 2].mean()) + avg_clip_w = float(pm[3, 2] * avg_z_v + pm[3, 3]) + avg_clip_w = avg_clip_w if abs(avg_clip_w) > 1e-8 else 1.0 + stroke_half_ndc = float(0.004 * stroke_width * abs(pm[0, 0]) / abs(avg_clip_w)) + stroke_half_px = stroke_half_ndc * config.pixel_width * 0.5 + + attrs = np.empty(n_total, dtype=_SURFACE_COMBINED_DTYPE) + attrs["in_vert"] = verts + attrs["in_normal"] = normals + attrs["in_fill_color"] = fill_color + attrs["in_stroke_color"] = stroke_color + attrs["in_bary"] = bary + attrs["stroke_half_px"] = stroke_half_px + attrs["diffuse_strength"] = float(diffuse_strength) + attrs["specular_strength"] = float(specular_strength) + attrs["specular_exponent"] = float(specular_exponent) + return attrs + + +def _smooth_surface_normals(surface_parts: list[np.ndarray]) -> None: + """Average normals at shared vertex positions (modifies in-place).""" + if not surface_parts: + return + + all_verts = np.concatenate([p["in_vert"] for p in surface_parts], axis=0) + all_norms = np.concatenate([p["in_normal"] for p in surface_parts], axis=0) + + PREC = 1e-5 + quantized = np.round(all_verts.astype(np.float64) / PREC).astype(np.int64) + _, inverse = np.unique(quantized, axis=0, return_inverse=True) + + n_unique = int(inverse.max()) + 1 + smooth = np.zeros((n_unique, 3), dtype=np.float64) + np.add.at(smooth, inverse, all_norms.astype(np.float64)) + + lengths = np.linalg.norm(smooth, axis=1, keepdims=True) + lengths = np.where(lengths < 1e-9, 1.0, lengths) + smooth = (smooth / lengths).astype(np.float32) + + idx = 0 + for part in surface_parts: + n = len(part) + part["in_normal"] = smooth[inverse[idx : idx + n]] + idx += n diff --git a/manim/scene/scene.py b/manim/scene/scene.py index b4ef54f38f..6e79eb5a37 100644 --- a/manim/scene/scene.py +++ b/manim/scene/scene.py @@ -50,6 +50,7 @@ from ..renderer.cairo_renderer import CairoRenderer from ..renderer.opengl_renderer import OpenGLCamera, OpenGLMobject, OpenGLRenderer from ..renderer.shader import Object3D +from ..renderer.webgpu.webgpu_renderer import WebGPURenderer from ..utils import opengl, space_ops from ..utils.exceptions import RerunSceneException from ..utils.family import extract_mobject_family_members @@ -167,7 +168,7 @@ def construct(self): def __init__( self, - renderer: CairoRenderer | OpenGLRenderer | None = None, + renderer: CairoRenderer | OpenGLRenderer | WebGPURenderer | None = None, camera_class: type[Camera] = Camera, always_update_mobjects: bool = False, random_seed: int | None = None, @@ -204,15 +205,20 @@ def __init__( if renderer is None: renderer = OpenGLRenderer() - if renderer is None: - self.renderer: CairoRenderer | OpenGLRenderer = CairoRenderer( - # TODO: Is it a suitable approach to make an instance of - # the self.camera_class here? - camera_class=self.camera_class, - skip_animations=self.skip_animations, - ) - else: - self.renderer = renderer + elif config.renderer == RendererType.WEBGPU: + if renderer is None: + renderer = WebGPURenderer() + + elif config.renderer == RendererType.CAIRO: + if renderer is None: + renderer = CairoRenderer( + # TODO: Is it a suitable approach to make an instance of + # the self.camera_class here? + camera_class=self.camera_class, + skip_animations=self.skip_animations, + ) + + self.renderer: CairoRenderer | OpenGLRenderer | WebGPURenderer = renderer self.renderer.init_scene(self) self.mobjects: list[Mobject] = [] @@ -451,8 +457,7 @@ def get_mobject_family_members(self) -> list[Mobject]: for mob in self.mobjects: family_members.extend(mob.get_family()) return family_members - else: - assert config.renderer == RendererType.CAIRO + elif config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: return extract_mobject_family_members( self.mobjects, use_z_index=self.renderer.camera.use_z_index, @@ -486,8 +491,7 @@ def add(self, *mobjects: Mobject | OpenGLMobject) -> Self: self.mobjects += new_mobjects # type: ignore[arg-type] self.remove(*new_meshes) # type: ignore[arg-type] self.meshes += new_meshes - else: - assert config.renderer == RendererType.CAIRO + elif config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: new_and_foreground_mobjects: list[Mobject] = [ *mobjects, # type: ignore[list-item] *self.foreground_mobjects, @@ -546,8 +550,7 @@ def lambda_function(mesh: Object3D) -> bool: filter(lambda_function, self.meshes), ) return self - else: - assert config.renderer == RendererType.CAIRO + elif config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: for list_name in "mobjects", "foreground_mobjects": self.restructure_mobjects(mobjects, list_name, False) return self @@ -1308,7 +1311,7 @@ def begin_animations(self) -> None: animation._setup_scene(self) animation.begin() - if config.renderer == RendererType.CAIRO: + if config.renderer in {RendererType.CAIRO, RendererType.WEBGPU}: # Paint all non-moving objects onto the screen, so they don't # have to be rendered every frame ( @@ -1361,7 +1364,9 @@ def play_internal(self, skip_rendering: bool = False) -> None: self.time_progression.close() def check_interactive_embed_is_valid(self) -> bool: - assert isinstance(self.renderer, OpenGLRenderer) + assert isinstance(self.renderer, OpenGLRenderer) or isinstance( + self.renderer, WebGPURenderer + ) if config["force_window"]: return True if self.skip_animation_preview: @@ -1387,7 +1392,42 @@ def check_interactive_embed_is_valid(self) -> bool: return True def interactive_embed(self) -> None: - """Like embed(), but allows for screen interaction.""" + """Like embed(), but allows for screen interaction. + + Drops into an IPython shell while the preview window stays alive and + responds to mouse / keyboard. Scene methods (``play``, ``wait``, + ``add``, ``remove``) are available without a ``self.`` prefix inside + the shell. + + Supported renderers: OpenGL, WebGPU (when ``-p`` / ``--preview`` is + active). Call this from inside :meth:`construct` after the animations + you want to have already played. + + Example + ------- + .. code-block:: python + + class MyScene(ThreeDScene): + def construct(self): + ax = ThreeDAxes() + self.add(ax) + self.interactive_embed() + """ + if config.renderer == RendererType.WEBGPU: + if not self.check_interactive_embed_is_valid(): + return + self.interactive_mode = True + from manim.renderer.webgpu.webgpu_interactive import ( + interactive_embed as _webgpu_embed, + ) + + currentframe: FrameType = inspect.currentframe() # type: ignore[assignment] + local_namespace = currentframe.f_back.f_locals # type: ignore[union-attr] + rerun = _webgpu_embed(self, self.renderer, local_namespace) + if rerun: + raise RerunSceneException + return + assert isinstance(self.camera, OpenGLCamera) assert isinstance(self.renderer, OpenGLRenderer) if not self.check_interactive_embed_is_valid(): diff --git a/manim/scene/scene_file_writer.py b/manim/scene/scene_file_writer.py index 17661c23aa..37232a22eb 100644 --- a/manim/scene/scene_file_writer.py +++ b/manim/scene/scene_file_writer.py @@ -54,6 +54,7 @@ from manim.renderer.cairo_renderer import CairoRenderer from manim.renderer.opengl_renderer import OpenGLRenderer + from manim.renderer.webgpu.webgpu_renderer import WebGPURenderer from manim.typing import PixelArray, StrPath @@ -227,7 +228,7 @@ class SceneFileWriter: def __init__( self, - renderer: CairoRenderer | OpenGLRenderer, + renderer: CairoRenderer | OpenGLRenderer | WebGPURenderer, scene_name: str, **kwargs: Any, ) -> None: @@ -552,7 +553,7 @@ def write_frame( else: frame = ( frame_or_renderer.get_frame() - if config.renderer == RendererType.OPENGL + if config.renderer in (RendererType.OPENGL, RendererType.WEBGPU) else frame_or_renderer ) @@ -575,7 +576,7 @@ def write_frame( else: image = ( frame_or_renderer.get_image() - if config.renderer == RendererType.OPENGL + if config.renderer in (RendererType.OPENGL, RendererType.WEBGPU) else Image.fromarray(frame_or_renderer) ) target_dir = self.image_file_path.parent / self.image_file_path.stem diff --git a/manim/scene/section.py b/manim/scene/section.py index 5b3463ec84..76f4f33ee1 100644 --- a/manim/scene/section.py +++ b/manim/scene/section.py @@ -57,6 +57,7 @@ class Section: :class:`.DefaultSectionType` :meth:`.CairoRenderer.update_skipping_status` :meth:`.OpenGLRenderer.update_skipping_status` + :meth:`.WebGPURenderer.update_skipping_status` """ def __init__( diff --git a/manim/scene/three_d_scene.py b/manim/scene/three_d_scene.py index 7d2337437e..2fd3e9d651 100644 --- a/manim/scene/three_d_scene.py +++ b/manim/scene/three_d_scene.py @@ -13,12 +13,14 @@ from manim.mobject.geometry.line import Line from manim.mobject.graphing.coordinate_systems import ThreeDAxes from manim.mobject.opengl.opengl_mobject import OpenGLMobject +from manim.mobject.three_d.light_source import AmbientLight, LightSource from manim.mobject.three_d.three_dimensions import Sphere from manim.mobject.value_tracker import ValueTracker from .. import config from ..animation.animation import Animation from ..animation.transform import Transform +from ..animation.updaters.update import UpdateFromAlphaFunc from ..camera.three_d_camera import ThreeDCamera from ..constants import DEGREES, RendererType from ..mobject.mobject import Mobject @@ -52,6 +54,32 @@ def __init__( ) super().__init__(camera_class=camera_class, **kwargs) + if config.renderer == RendererType.WEBGPU: + # Default ambient light — exactly one is kept at all times. + # WebGPU renderer reads self.mobjects to find LightSource instances. + self._ambient_light = AmbientLight(intensity=0.5) + self.add(self._ambient_light) + + def add(self, *mobjects): + """Override to enforce the single-ambient-light rule. + + Every scene can have only one ambient light setting. + If the caller adds a new :class:`~.AmbientLight`, the existing one is + removed first so only one ambient light is ever in the scene. + """ + if config.renderer == RendererType.WEBGPU: + for mob in mobjects: + if isinstance(mob, AmbientLight): + # Remove any existing AmbientLight before adding the new one. + existing = [m for m in self.mobjects if isinstance(m, AmbientLight)] + for old in existing: + super().remove(old) + self._ambient_light = mob + else: # do not allow LightSource to be added for Cairo or OpenGL renderer + mobjects = [mob for mob in mobjects if not isinstance(mob, LightSource)] + + return super().add(*mobjects) + def set_camera_orientation( self, phi: float | None = None, @@ -134,6 +162,15 @@ def begin_ambient_camera_rotation(self, rate: float = 0.02, about: str = "theta" } cam.add_updater(lambda m, dt: methods[about](rate * dt)) self.add(self.camera) + elif config.renderer == RendererType.WEBGPU: + cam = self.renderer.camera + methods = { + "theta": cam.increment_theta, + "phi": cam.increment_phi, + "gamma": cam.increment_gamma, + } + cam.add_updater(lambda m, dt: methods[about](rate * dt)) + self.add(cam) except Exception as e: raise ValueError("Invalid ambient rotation angle.") from e @@ -152,6 +189,8 @@ def stop_ambient_camera_rotation(self, about="theta"): self.remove(x) elif config.renderer == RendererType.OPENGL: self.camera.clear_updaters() + elif config.renderer == RendererType.WEBGPU: + self.renderer.camera.clear_updaters() except Exception as e: raise ValueError("Invalid ambient rotation angle.") from e @@ -176,6 +215,28 @@ def begin_3dillusion_camera_rotation( The azimutal angle the camera should move around. Defaults to the current theta angle. """ + if config.renderer == RendererType.WEBGPU: + cam = self.renderer.camera + if origin_theta is None: + origin_theta = cam.euler_angles[0] + if origin_phi is None: + origin_phi = cam.euler_angles[1] + + _theta_t = ValueTracker(0) + _phi_t = ValueTracker(0) + + def update_cam_illusion(m, dt): + _theta_t.increment_value(dt * rate) + _phi_t.increment_value(dt * rate) + m.set_euler_angles( + theta=origin_theta + 0.2 * np.sin(_theta_t.get_value()), + phi=origin_phi + 0.1 * np.cos(_phi_t.get_value()) - 0.1, + ) + + cam.add_updater(update_cam_illusion) + self.add(cam) + return + if origin_theta is None: origin_theta = self.renderer.camera.theta_tracker.get_value() if origin_phi is None: @@ -203,6 +264,10 @@ def update_phi(m, dt): def stop_3dillusion_camera_rotation(self): """This method stops all illusion camera rotations.""" + if config.renderer == RendererType.WEBGPU: + self.renderer.camera.clear_updaters() + self.remove(self.renderer.camera) + return self.renderer.camera.theta_tracker.clear_updaters() self.remove(self.renderer.camera.theta_tracker) self.renderer.camera.phi_tracker.clear_updaters() @@ -300,6 +365,35 @@ def move_camera( anims += [Transform(cam, cam2)] + elif config.renderer == RendererType.WEBGPU: + cam = self.renderer.camera + start_theta = cam.euler_angles[0] + start_phi = cam.euler_angles[1] + start_gamma = cam.euler_angles[2] + target_theta = theta if theta is not None else start_theta + target_phi = phi if phi is not None else start_phi + target_gamma = gamma if gamma is not None else start_gamma + + def update_cam(m, alpha): + m.set_euler_angles( + theta=start_theta + alpha * (target_theta - start_theta), + phi=start_phi + alpha * (target_phi - start_phi), + gamma=start_gamma + alpha * (target_gamma - start_gamma), + ) + + anims.append(UpdateFromAlphaFunc(cam, update_cam)) + + if focal_distance is not None: + start_fd = cam.focal_distance + anims.append( + UpdateFromAlphaFunc( + cam, + lambda m, a, _s=start_fd: setattr( + m, "focal_distance", _s + a * (focal_distance - _s) + ), + ) + ) + self.play(*anims + added_anims, **kwargs) # These lines are added to improve performance. If manim thinks that frame_center is moving, @@ -353,6 +447,9 @@ def add_fixed_orientation_mobjects(self, *mobjects: Mobject, **kwargs): mob: OpenGLMobject mob.fix_orientation() self.add(mob) + elif config.renderer == RendererType.WEBGPU: + self.renderer.camera.add_fixed_orientation_mobjects(*mobjects) + self.add(*mobjects) def add_fixed_in_frame_mobjects(self, *mobjects: Mobject): """ @@ -375,6 +472,9 @@ def add_fixed_in_frame_mobjects(self, *mobjects: Mobject): mob: OpenGLMobject mob.fix_in_frame() self.add(mob) + elif config.renderer == RendererType.WEBGPU: + self.renderer.camera.add_fixed_in_frame_mobjects(*mobjects) + self.add(*mobjects) def remove_fixed_orientation_mobjects(self, *mobjects: Mobject): """ @@ -395,6 +495,8 @@ def remove_fixed_orientation_mobjects(self, *mobjects: Mobject): mob: OpenGLMobject mob.unfix_orientation() self.remove(mob) + elif config.renderer == RendererType.WEBGPU: + self.renderer.camera.remove_fixed_orientation_mobjects(*mobjects) def remove_fixed_in_frame_mobjects(self, *mobjects: Mobject): """ @@ -414,6 +516,8 @@ def remove_fixed_in_frame_mobjects(self, *mobjects: Mobject): mob: OpenGLMobject mob.unfix_from_frame() self.remove(mob) + elif config.renderer == RendererType.WEBGPU: + self.renderer.camera.remove_fixed_in_frame_mobjects(*mobjects) ## def set_to_default_angled_camera_orientation(self, **kwargs): diff --git a/pyproject.toml b/pyproject.toml index 825d92c809..e98499d878 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,6 +49,9 @@ dependencies = [ "tqdm>=4.21.0", "typing-extensions>=4.12.0", "watchdog>=2.0.0", + "wgpu>=0.31.0", + "rendercanvas>=2.6.3", + "glfw>=2.9.0", ] diff --git a/tests/test_graphical_units/control_data/coordinate_system/plot_surface.npz b/tests/test_graphical_units/control_data/coordinate_system/plot_surface.npz index 83697c1d58..7a2ac508ac 100644 Binary files a/tests/test_graphical_units/control_data/coordinate_system/plot_surface.npz and b/tests/test_graphical_units/control_data/coordinate_system/plot_surface.npz differ diff --git a/tests/test_graphical_units/test_coordinate_systems.py b/tests/test_graphical_units/test_coordinate_systems.py index 299c7bcae6..fe2dcb1263 100644 --- a/tests/test_graphical_units/test_coordinate_systems.py +++ b/tests/test_graphical_units/test_coordinate_systems.py @@ -51,7 +51,7 @@ def param_trig(u, v): param_trig, u_range=(-5, 5), v_range=(-5, 5), - color=BLUE, + color=GREEN, ) scene.add(axes, trig_plane)