123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393 |
- 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] = (1e-5, PI - 1e-5),
- 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
- )
- 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 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)
|