mirror of
https://github.com/3b1b/manim.git
synced 2025-04-13 09:47:07 +00:00
045e1f6b8a
This means that surface shading will not necessarily work well for arbitrary transformations of the surface. But the existing solution was flimsy anyway, and caused annoying issues with singularity points.
408 lines
12 KiB
Python
408 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
|
|
from manimlib.constants import BLUE, BLUE_D, BLUE_E, GREY_A, BLACK
|
|
from manimlib.constants import IN, ORIGIN, OUT, RIGHT
|
|
from manimlib.constants import PI, TAU
|
|
from manimlib.mobject.mobject import Mobject
|
|
from manimlib.mobject.types.surface import SGroup
|
|
from manimlib.mobject.types.surface import Surface
|
|
from manimlib.mobject.types.vectorized_mobject import VGroup
|
|
from manimlib.mobject.types.vectorized_mobject import VMobject
|
|
from manimlib.mobject.geometry import Polygon
|
|
from manimlib.mobject.geometry import Square
|
|
from manimlib.utils.bezier import interpolate
|
|
from manimlib.utils.iterables import adjacent_pairs
|
|
from manimlib.utils.space_ops import compass_directions
|
|
from manimlib.utils.space_ops import get_norm
|
|
from manimlib.utils.space_ops import z_to_vector
|
|
|
|
from typing import TYPE_CHECKING
|
|
if TYPE_CHECKING:
|
|
from typing import Tuple, TypeVar
|
|
from manimlib.typing import ManimColor, Vect3, Sequence
|
|
|
|
T = TypeVar("T", bound=Mobject)
|
|
|
|
|
|
class SurfaceMesh(VGroup):
|
|
def __init__(
|
|
self,
|
|
uv_surface: Surface,
|
|
resolution: Tuple[int, int] = (21, 11),
|
|
stroke_width: float = 1,
|
|
stroke_color: ManimColor = GREY_A,
|
|
normal_nudge: float = 1e-2,
|
|
depth_test: bool = True,
|
|
joint_type: str = 'no_joint',
|
|
**kwargs
|
|
):
|
|
self.uv_surface = uv_surface
|
|
self.resolution = resolution
|
|
self.normal_nudge = normal_nudge
|
|
|
|
super().__init__(
|
|
stroke_color=stroke_color,
|
|
stroke_width=stroke_width,
|
|
depth_test=depth_test,
|
|
joint_type=joint_type,
|
|
**kwargs
|
|
)
|
|
|
|
def init_points(self) -> None:
|
|
uv_surface = self.uv_surface
|
|
|
|
full_nu, full_nv = uv_surface.resolution
|
|
part_nu, part_nv = self.resolution
|
|
# 'indices' are treated as floats. Later, there will be
|
|
# an interpolation between the floor and ceiling of these
|
|
# indices
|
|
u_indices = np.linspace(0, full_nu - 1, part_nu)
|
|
v_indices = np.linspace(0, full_nv - 1, part_nv)
|
|
|
|
points = uv_surface.get_points()
|
|
normals = uv_surface.get_unit_normals()
|
|
nudge = self.normal_nudge
|
|
nudged_points = points + nudge * normals
|
|
|
|
for ui in u_indices:
|
|
path = VMobject()
|
|
low_ui = full_nv * int(math.floor(ui))
|
|
high_ui = full_nv * int(math.ceil(ui))
|
|
path.set_points_smoothly(interpolate(
|
|
nudged_points[low_ui:low_ui + full_nv],
|
|
nudged_points[high_ui:high_ui + full_nv],
|
|
ui % 1
|
|
))
|
|
self.add(path)
|
|
for vi in v_indices:
|
|
path = VMobject()
|
|
path.set_points_smoothly(interpolate(
|
|
nudged_points[int(math.floor(vi))::full_nv],
|
|
nudged_points[int(math.ceil(vi))::full_nv],
|
|
vi % 1
|
|
))
|
|
self.add(path)
|
|
|
|
|
|
# 3D shapes
|
|
|
|
class Sphere(Surface):
|
|
def __init__(
|
|
self,
|
|
u_range: Tuple[float, float] = (0, TAU),
|
|
v_range: Tuple[float, float] = (0, PI),
|
|
resolution: Tuple[int, int] = (101, 51),
|
|
radius: float = 1.0,
|
|
**kwargs,
|
|
):
|
|
self.radius = radius
|
|
super().__init__(
|
|
u_range=u_range,
|
|
v_range=v_range,
|
|
resolution=resolution,
|
|
**kwargs
|
|
)
|
|
# Add bespoke normal specification to avoid issue at poles
|
|
self.data['d_normal_point'] = self.data['point'] * ((radius + self.normal_nudge) / radius)
|
|
|
|
def uv_func(self, u: float, v: float) -> np.ndarray:
|
|
return self.radius * np.array([
|
|
math.cos(u) * math.sin(v),
|
|
math.sin(u) * math.sin(v),
|
|
-math.cos(v)
|
|
])
|
|
|
|
|
|
class Torus(Surface):
|
|
def __init__(
|
|
self,
|
|
u_range: Tuple[float, float] = (0, TAU),
|
|
v_range: Tuple[float, float] = (0, TAU),
|
|
r1: float = 3.0,
|
|
r2: float = 1.0,
|
|
**kwargs,
|
|
):
|
|
self.r1 = r1
|
|
self.r2 = r2
|
|
super().__init__(
|
|
u_range=u_range,
|
|
v_range=v_range,
|
|
**kwargs,
|
|
)
|
|
|
|
def uv_func(self, u: float, v: float) -> np.ndarray:
|
|
P = np.array([math.cos(u), math.sin(u), 0])
|
|
return (self.r1 - self.r2 * math.cos(v)) * P - self.r2 * math.sin(v) * OUT
|
|
|
|
|
|
class Cylinder(Surface):
|
|
def __init__(
|
|
self,
|
|
u_range: Tuple[float, float] = (0, TAU),
|
|
v_range: Tuple[float, float] = (-1, 1),
|
|
resolution: Tuple[int, int] = (101, 11),
|
|
height: float = 2,
|
|
radius: float = 1,
|
|
axis: Vect3 = OUT,
|
|
**kwargs,
|
|
):
|
|
self.height = height
|
|
self.radius = radius
|
|
self.axis = axis
|
|
super().__init__(
|
|
u_range=u_range,
|
|
v_range=v_range,
|
|
resolution=resolution,
|
|
**kwargs
|
|
)
|
|
|
|
def init_points(self):
|
|
super().init_points()
|
|
self.scale(self.radius)
|
|
self.set_depth(self.height, stretch=True)
|
|
self.apply_matrix(z_to_vector(self.axis))
|
|
|
|
def uv_func(self, u: float, v: float) -> np.ndarray:
|
|
return np.array([np.cos(u), np.sin(u), v])
|
|
|
|
|
|
class Cone(Cylinder):
|
|
def __init__(
|
|
self,
|
|
u_range: Tuple[float, float] = (0, TAU),
|
|
v_range: Tuple[float, float] = (0, 1),
|
|
*args,
|
|
**kwargs,
|
|
):
|
|
super().__init__(u_range=u_range, v_range=v_range, *args, **kwargs)
|
|
|
|
def uv_func(self, u: float, v: float) -> np.ndarray:
|
|
return np.array([(1 - v) * np.cos(u), (1 - v) * np.sin(u), v])
|
|
|
|
|
|
class Line3D(Cylinder):
|
|
def __init__(
|
|
self,
|
|
start: Vect3,
|
|
end: Vect3,
|
|
width: float = 0.05,
|
|
resolution: Tuple[int, int] = (21, 25),
|
|
**kwargs
|
|
):
|
|
axis = end - start
|
|
super().__init__(
|
|
height=get_norm(axis),
|
|
radius=width / 2,
|
|
axis=axis,
|
|
resolution=resolution,
|
|
**kwargs
|
|
)
|
|
self.shift((start + end) / 2)
|
|
|
|
|
|
class Disk3D(Surface):
|
|
def __init__(
|
|
self,
|
|
radius: float = 1,
|
|
u_range: Tuple[float, float] = (0, 1),
|
|
v_range: Tuple[float, float] = (0, TAU),
|
|
resolution: Tuple[int, int] = (2, 100),
|
|
**kwargs
|
|
):
|
|
super().__init__(
|
|
u_range=u_range,
|
|
v_range=v_range,
|
|
resolution=resolution,
|
|
**kwargs,
|
|
)
|
|
self.scale(radius)
|
|
|
|
def uv_func(self, u: float, v: float) -> np.ndarray:
|
|
return np.array([
|
|
u * math.cos(v),
|
|
u * math.sin(v),
|
|
0
|
|
])
|
|
|
|
|
|
class Square3D(Surface):
|
|
def __init__(
|
|
self,
|
|
side_length: float = 2.0,
|
|
u_range: Tuple[float, float] = (-1, 1),
|
|
v_range: Tuple[float, float] = (-1, 1),
|
|
resolution: Tuple[int, int] = (2, 2),
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
u_range=u_range,
|
|
v_range=v_range,
|
|
resolution=resolution,
|
|
**kwargs
|
|
)
|
|
self.scale(side_length / 2)
|
|
|
|
def uv_func(self, u: float, v: float) -> np.ndarray:
|
|
return np.array([u, v, 0])
|
|
|
|
|
|
def square_to_cube_faces(square: T) -> list[T]:
|
|
radius = square.get_height() / 2
|
|
square.move_to(radius * OUT)
|
|
result = [square.copy()]
|
|
result.extend([
|
|
square.copy().rotate(PI / 2, axis=vect, about_point=ORIGIN)
|
|
for vect in compass_directions(4)
|
|
])
|
|
result.append(square.copy().rotate(PI, RIGHT, about_point=ORIGIN))
|
|
return result
|
|
|
|
|
|
class Cube(SGroup):
|
|
def __init__(
|
|
self,
|
|
color: ManimColor = BLUE,
|
|
opacity: float = 1,
|
|
shading: Tuple[float, float, float] = (0.1, 0.5, 0.1),
|
|
square_resolution: Tuple[int, int] = (2, 2),
|
|
side_length: float = 2,
|
|
**kwargs,
|
|
):
|
|
face = Square3D(
|
|
resolution=square_resolution,
|
|
side_length=side_length,
|
|
color=color,
|
|
opacity=opacity,
|
|
shading=shading,
|
|
)
|
|
super().__init__(*square_to_cube_faces(face), **kwargs)
|
|
|
|
|
|
class Prism(Cube):
|
|
def __init__(
|
|
self,
|
|
width: float = 3.0,
|
|
height: float = 2.0,
|
|
depth: float = 1.0,
|
|
**kwargs
|
|
):
|
|
super().__init__(**kwargs)
|
|
for dim, value in enumerate([width, height, depth]):
|
|
self.rescale_to_fit(value, dim, stretch=True)
|
|
|
|
|
|
class VGroup3D(VGroup):
|
|
def __init__(
|
|
self,
|
|
*vmobjects: VMobject,
|
|
depth_test: bool = True,
|
|
shading: Tuple[float, float, float] = (0.2, 0.2, 0.2),
|
|
joint_type: str = "no_joint",
|
|
**kwargs
|
|
):
|
|
super().__init__(*vmobjects, **kwargs)
|
|
self.set_shading(*shading)
|
|
self.set_joint_type(joint_type)
|
|
if depth_test:
|
|
self.apply_depth_test()
|
|
|
|
|
|
class VCube(VGroup3D):
|
|
def __init__(
|
|
self,
|
|
side_length: float = 2.0,
|
|
fill_color: ManimColor = BLUE_D,
|
|
fill_opacity: float = 1,
|
|
stroke_width: float = 0,
|
|
**kwargs
|
|
):
|
|
style = dict(
|
|
fill_color=fill_color,
|
|
fill_opacity=fill_opacity,
|
|
stroke_width=stroke_width,
|
|
**kwargs
|
|
)
|
|
face = Square(side_length=side_length, **style)
|
|
super().__init__(*square_to_cube_faces(face), **style)
|
|
|
|
|
|
class VPrism(VCube):
|
|
def __init__(
|
|
self,
|
|
width: float = 3.0,
|
|
height: float = 2.0,
|
|
depth: float = 1.0,
|
|
**kwargs
|
|
):
|
|
super().__init__(**kwargs)
|
|
for dim, value in enumerate([width, height, depth]):
|
|
self.rescale_to_fit(value, dim, stretch=True)
|
|
|
|
|
|
class Dodecahedron(VGroup3D):
|
|
def __init__(
|
|
self,
|
|
fill_color: ManimColor = BLUE_E,
|
|
fill_opacity: float = 1,
|
|
stroke_color: ManimColor = BLUE_E,
|
|
stroke_width: float = 1,
|
|
shading: Tuple[float, float, float] = (0.2, 0.2, 0.2),
|
|
**kwargs,
|
|
):
|
|
style = dict(
|
|
fill_color=fill_color,
|
|
fill_opacity=fill_opacity,
|
|
stroke_color=stroke_color,
|
|
stroke_width=stroke_width,
|
|
shading=shading,
|
|
**kwargs
|
|
)
|
|
|
|
# Start by creating two of the pentagons, meeting
|
|
# back to back on the positive x-axis
|
|
phi = (1 + math.sqrt(5)) / 2
|
|
x, y, z = np.identity(3)
|
|
pentagon1 = Polygon(
|
|
np.array([phi, 1 / phi, 0]),
|
|
np.array([1, 1, 1]),
|
|
np.array([1 / phi, 0, phi]),
|
|
np.array([1, -1, 1]),
|
|
np.array([phi, -1 / phi, 0]),
|
|
**style
|
|
)
|
|
pentagon2 = pentagon1.copy().stretch(-1, 2, about_point=ORIGIN)
|
|
pentagon2.reverse_points()
|
|
x_pair = VGroup(pentagon1, pentagon2)
|
|
z_pair = x_pair.copy().apply_matrix(np.array([z, -x, -y]).T)
|
|
y_pair = x_pair.copy().apply_matrix(np.array([y, z, x]).T)
|
|
|
|
pentagons = [*x_pair, *y_pair, *z_pair]
|
|
for pentagon in list(pentagons):
|
|
pc = pentagon.copy()
|
|
pc.apply_function(lambda p: -p)
|
|
pc.reverse_points()
|
|
pentagons.append(pc)
|
|
|
|
super().__init__(*pentagons, **style)
|
|
|
|
|
|
class Prismify(VGroup3D):
|
|
def __init__(self, vmobject, depth=1.0, direction=IN, **kwargs):
|
|
# At the moment, this assume stright edges
|
|
vect = depth * direction
|
|
pieces = [vmobject.copy()]
|
|
points = vmobject.get_anchors()
|
|
for p1, p2 in adjacent_pairs(points):
|
|
wall = VMobject()
|
|
wall.match_style(vmobject)
|
|
wall.set_points_as_corners([p1, p2, p2 + vect, p1 + vect])
|
|
pieces.append(wall)
|
|
top = vmobject.copy()
|
|
top.shift(vect)
|
|
top.reverse_points()
|
|
pieces.append(top)
|
|
super().__init__(*pieces, **kwargs)
|