2022-02-15 14:49:02 +08:00
|
|
|
from __future__ import annotations
|
|
|
|
|
2022-04-12 19:19:59 +08:00
|
|
|
import copy
|
2020-02-17 12:14:40 -08:00
|
|
|
import os
|
|
|
|
import re
|
2022-02-15 14:49:02 +08:00
|
|
|
|
2020-02-17 12:14:40 -08:00
|
|
|
import moderngl
|
2020-06-29 18:17:18 -07:00
|
|
|
import numpy as np
|
2020-02-17 12:14:40 -08:00
|
|
|
|
2022-12-27 21:39:20 -08:00
|
|
|
from manimlib.utils.iterables import resize_array
|
2023-01-16 14:18:35 -08:00
|
|
|
from manimlib.utils.shaders import get_shader_code_from_file
|
2020-02-17 12:14:40 -08:00
|
|
|
|
2022-04-12 19:19:59 +08:00
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
|
|
|
|
if TYPE_CHECKING:
|
2023-01-16 14:18:35 -08:00
|
|
|
from typing import List
|
2022-04-12 19:19:59 +08:00
|
|
|
|
|
|
|
|
2020-02-17 12:14:40 -08:00
|
|
|
# Mobjects that should be rendered with
|
|
|
|
# the same shader will be organized and
|
|
|
|
# clumped together based on keeping track
|
|
|
|
# of a dict holding all the relevant information
|
|
|
|
# to that shader
|
|
|
|
|
|
|
|
|
2020-06-29 18:17:18 -07:00
|
|
|
class ShaderWrapper(object):
|
2022-02-15 14:49:02 +08:00
|
|
|
def __init__(
|
|
|
|
self,
|
2022-12-27 14:53:55 -08:00
|
|
|
vert_data: np.ndarray,
|
2022-02-15 14:49:02 +08:00
|
|
|
vert_indices: np.ndarray | None = None,
|
|
|
|
shader_folder: str | None = None,
|
2023-01-16 11:50:48 -08:00
|
|
|
uniforms: dict[str, float | np.ndarray] | None = None, # A dictionary mapping names of uniform variables
|
2022-02-15 14:49:02 +08:00
|
|
|
texture_paths: dict[str, str] | None = None, # A dictionary mapping names to filepaths for textures.
|
|
|
|
depth_test: bool = False,
|
2023-01-13 21:42:34 -08:00
|
|
|
use_clip_plane: bool = False,
|
2022-02-15 14:49:02 +08:00
|
|
|
render_primitive: int = moderngl.TRIANGLE_STRIP,
|
|
|
|
):
|
2020-06-29 18:17:18 -07:00
|
|
|
self.vert_data = vert_data
|
|
|
|
self.vert_indices = vert_indices
|
|
|
|
self.vert_attributes = vert_data.dtype.names
|
2021-01-05 23:14:16 -08:00
|
|
|
self.shader_folder = shader_folder
|
2020-06-29 18:17:18 -07:00
|
|
|
self.uniforms = uniforms or dict()
|
|
|
|
self.texture_paths = texture_paths or dict()
|
|
|
|
self.depth_test = depth_test
|
2023-01-13 21:42:34 -08:00
|
|
|
self.use_clip_plane = use_clip_plane
|
2020-12-04 08:09:02 -08:00
|
|
|
self.render_primitive = str(render_primitive)
|
2021-01-09 18:52:54 -08:00
|
|
|
self.init_program_code()
|
|
|
|
self.refresh_id()
|
2020-06-29 18:17:18 -07:00
|
|
|
|
2022-04-23 18:51:03 -07:00
|
|
|
def __eq__(self, shader_wrapper: ShaderWrapper):
|
|
|
|
return all((
|
|
|
|
np.all(self.vert_data == shader_wrapper.vert_data),
|
|
|
|
np.all(self.vert_indices == shader_wrapper.vert_indices),
|
|
|
|
self.shader_folder == shader_wrapper.shader_folder,
|
|
|
|
all(
|
|
|
|
np.all(self.uniforms[key] == shader_wrapper.uniforms[key])
|
|
|
|
for key in self.uniforms
|
|
|
|
),
|
|
|
|
all(
|
|
|
|
self.texture_paths[key] == shader_wrapper.texture_paths[key]
|
|
|
|
for key in self.texture_paths
|
|
|
|
),
|
|
|
|
self.depth_test == shader_wrapper.depth_test,
|
|
|
|
self.render_primitive == shader_wrapper.render_primitive,
|
|
|
|
))
|
|
|
|
|
2020-06-29 18:17:18 -07:00
|
|
|
def copy(self):
|
|
|
|
result = copy.copy(self)
|
|
|
|
result.vert_data = np.array(self.vert_data)
|
|
|
|
if result.vert_indices is not None:
|
|
|
|
result.vert_indices = np.array(self.vert_indices)
|
|
|
|
if self.uniforms:
|
2022-04-24 13:23:30 -07:00
|
|
|
result.uniforms = {key: np.array(value) for key, value in self.uniforms.items()}
|
2020-06-29 18:17:18 -07:00
|
|
|
if self.texture_paths:
|
|
|
|
result.texture_paths = dict(self.texture_paths)
|
|
|
|
return result
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def is_valid(self) -> bool:
|
2020-06-29 18:17:18 -07:00
|
|
|
return all([
|
|
|
|
self.vert_data is not None,
|
2021-01-09 18:52:54 -08:00
|
|
|
self.program_code["vertex_shader"] is not None,
|
|
|
|
self.program_code["fragment_shader"] is not None,
|
2020-06-29 18:17:18 -07:00
|
|
|
])
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def get_id(self) -> str:
|
2020-06-29 18:17:18 -07:00
|
|
|
return self.id
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def get_program_id(self) -> int:
|
2020-06-29 18:17:18 -07:00
|
|
|
return self.program_id
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def create_id(self) -> str:
|
2020-06-29 18:17:18 -07:00
|
|
|
# A unique id for a shader
|
|
|
|
return "|".join(map(str, [
|
2021-01-09 18:52:54 -08:00
|
|
|
self.program_id,
|
2020-06-29 18:17:18 -07:00
|
|
|
self.uniforms,
|
|
|
|
self.texture_paths,
|
|
|
|
self.depth_test,
|
2020-12-04 08:09:02 -08:00
|
|
|
self.render_primitive,
|
2020-06-29 18:17:18 -07:00
|
|
|
]))
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def refresh_id(self) -> None:
|
2021-01-09 18:52:54 -08:00
|
|
|
self.program_id = self.create_program_id()
|
2020-06-29 18:17:18 -07:00
|
|
|
self.id = self.create_id()
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def create_program_id(self) -> int:
|
2021-01-09 18:52:54 -08:00
|
|
|
return hash("".join((
|
|
|
|
self.program_code[f"{name}_shader"] or ""
|
|
|
|
for name in ("vertex", "geometry", "fragment")
|
|
|
|
)))
|
2020-06-29 18:17:18 -07:00
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def init_program_code(self) -> None:
|
|
|
|
def get_code(name: str) -> str | None:
|
2021-01-05 23:14:16 -08:00
|
|
|
return get_shader_code_from_file(
|
|
|
|
os.path.join(self.shader_folder, f"{name}.glsl")
|
|
|
|
)
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
self.program_code: dict[str, str | None] = {
|
2021-01-05 23:14:16 -08:00
|
|
|
"vertex_shader": get_code("vert"),
|
|
|
|
"geometry_shader": get_code("geom"),
|
|
|
|
"fragment_shader": get_code("frag"),
|
2020-06-29 18:17:18 -07:00
|
|
|
}
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def get_program_code(self) -> dict[str, str | None]:
|
2021-01-09 18:52:54 -08:00
|
|
|
return self.program_code
|
|
|
|
|
2022-02-15 14:49:02 +08:00
|
|
|
def replace_code(self, old: str, new: str) -> None:
|
2021-01-09 18:52:54 -08:00
|
|
|
code_map = self.program_code
|
|
|
|
for (name, code) in code_map.items():
|
|
|
|
if code_map[name] is None:
|
|
|
|
continue
|
|
|
|
code_map[name] = re.sub(old, new, code_map[name])
|
|
|
|
self.refresh_id()
|
|
|
|
|
2022-12-27 21:39:20 -08:00
|
|
|
def combine_with(self, *shader_wrappers: ShaderWrapper) -> ShaderWrapper:
|
2022-12-29 12:04:35 -08:00
|
|
|
if len(shader_wrappers) > 0:
|
2023-01-15 20:27:19 -08:00
|
|
|
data_list = [self.vert_data, *(sw.vert_data for sw in shader_wrappers)]
|
|
|
|
if self.vert_indices is not None:
|
|
|
|
indices_list = [self.vert_indices, *(sw.vert_indices for sw in shader_wrappers)]
|
|
|
|
else:
|
|
|
|
indices_list = None
|
|
|
|
self.read_in(data_list, indices_list)
|
2022-12-27 21:39:20 -08:00
|
|
|
return self
|
|
|
|
|
2023-01-15 20:27:19 -08:00
|
|
|
def read_in(
|
|
|
|
self,
|
2023-01-16 08:59:50 -08:00
|
|
|
data_list: List[np.ndarray],
|
|
|
|
indices_list: List[np.ndarray] | None = None
|
2023-01-15 20:27:19 -08:00
|
|
|
) -> ShaderWrapper:
|
2022-12-27 21:39:20 -08:00
|
|
|
# Assume all are of the same type
|
2023-01-16 08:59:50 -08:00
|
|
|
total_len = sum(len(data) for data in data_list)
|
2022-12-27 21:39:20 -08:00
|
|
|
self.vert_data = resize_array(self.vert_data, total_len)
|
2023-01-15 20:27:19 -08:00
|
|
|
if total_len == 0:
|
|
|
|
return self
|
|
|
|
|
2023-01-16 13:26:34 -08:00
|
|
|
# Stack the data
|
|
|
|
np.concatenate(data_list, out=self.vert_data)
|
2022-12-27 21:39:20 -08:00
|
|
|
|
2023-01-16 13:26:34 -08:00
|
|
|
if indices_list is None or self.vert_indices is None:
|
|
|
|
return self
|
|
|
|
|
|
|
|
total_verts = sum(len(vi) for vi in indices_list)
|
|
|
|
self.vert_indices = resize_array(self.vert_indices, total_verts)
|
|
|
|
|
|
|
|
# Stack vert_indices, but adding the appropriate offset
|
|
|
|
# alogn the way
|
2022-12-27 21:39:20 -08:00
|
|
|
n_points = 0
|
|
|
|
n_verts = 0
|
2023-01-16 13:26:34 -08:00
|
|
|
for data, indices in zip(data_list, indices_list):
|
|
|
|
new_n_verts = n_verts + len(indices)
|
|
|
|
self.vert_indices[n_verts:new_n_verts] = indices + n_points
|
|
|
|
n_verts = new_n_verts
|
|
|
|
n_points += len(data)
|
2020-06-29 18:17:18 -07:00
|
|
|
return self
|