diff --git a/manimlib/camera/camera.py b/manimlib/camera/camera.py index 59b43f7b..647f28a9 100644 --- a/manimlib/camera/camera.py +++ b/manimlib/camera/camera.py @@ -1,25 +1,21 @@ from __future__ import annotations -import math +import moderngl +from colour import Color +import OpenGL.GL as gl + import itertools as it -import moderngl import numpy as np +from scipy.spatial.transform import Rotation from PIL import Image -import OpenGL.GL as gl -from colour import Color from manimlib.constants import * from manimlib.mobject.mobject import Mobject from manimlib.mobject.mobject import Point from manimlib.utils.config_ops import digest_config from manimlib.utils.simple_functions import fdiv -from manimlib.utils.simple_functions import clip -from manimlib.utils.space_ops import angle_of_vector -from manimlib.utils.space_ops import rotation_matrix_transpose_from_quaternion -from manimlib.utils.space_ops import rotation_matrix_transpose -from manimlib.utils.space_ops import quaternion_from_angle_axis -from manimlib.utils.space_ops import quaternion_mult +from manimlib.utils.space_ops import normalize from typing import TYPE_CHECKING @@ -31,15 +27,13 @@ class CameraFrame(Mobject): CONFIG = { "frame_shape": (FRAME_WIDTH, FRAME_HEIGHT), "center_point": ORIGIN, - # Theta, phi, gamma - "euler_angles": [0, 0, 0], "focal_distance": 2, } - def init_data(self) -> None: - super().init_data() - self.data["euler_angles"] = np.array(self.euler_angles, dtype=float) - self.refresh_rotation_matrix() + def init_uniforms(self) -> None: + super().init_uniforms() + # As a quaternion + self.uniforms["orientation"] = Rotation.identity().as_quat() def init_points(self) -> None: self.set_points([ORIGIN, LEFT, RIGHT, DOWN, UP]) @@ -47,42 +41,29 @@ class CameraFrame(Mobject): self.set_height(self.frame_shape[1], stretch=True) self.move_to(self.center_point) + def set_orientation(self, rotation: Rotation): + self.uniforms["orientation"][:] = rotation.as_quat() + return self + + def get_orientation(self): + return Rotation.from_quat(self.uniforms["orientation"]) + def to_default_state(self): self.center() self.set_height(FRAME_HEIGHT) self.set_width(FRAME_WIDTH) - self.set_euler_angles(0, 0, 0) + self.set_orientation(Rotation.identity()) return self - def get_euler_angles(self) -> np.ndarray: - return self.data["euler_angles"] + def get_euler_angles(self): + return self.get_orientation().as_euler("xzy") - def get_inverse_camera_rotation_matrix(self) -> list[list[float]]: - return self.inverse_camera_rotation_matrix - - def refresh_rotation_matrix(self) -> None: - # Rotate based on camera orientation - theta, phi, gamma = self.get_euler_angles() - quat = quaternion_mult( - quaternion_from_angle_axis(theta, OUT, axis_normalized=True), - quaternion_from_angle_axis(phi, RIGHT, axis_normalized=True), - quaternion_from_angle_axis(gamma, OUT, axis_normalized=True), - ) - self.inverse_camera_rotation_matrix = rotation_matrix_transpose_from_quaternion(quat) + def get_inverse_camera_rotation_matrix(self): + return self.get_orientation().as_matrix().T def rotate(self, angle: float, axis: np.ndarray = OUT, **kwargs): - curr_rot_T = self.get_inverse_camera_rotation_matrix() - added_rot_T = rotation_matrix_transpose(angle, axis) - new_rot_T = np.dot(curr_rot_T, added_rot_T) - Fz = new_rot_T[2] - phi = np.arccos(clip(Fz[2], -1, 1)) - theta = angle_of_vector(Fz[:2]) + PI / 2 - partial_rot_T = np.dot( - rotation_matrix_transpose(phi, RIGHT), - rotation_matrix_transpose(theta, OUT), - ) - gamma = angle_of_vector(np.dot(partial_rot_T, new_rot_T.T)[:, 0]) - self.set_euler_angles(theta, phi, gamma) + rot = Rotation.from_rotvec(angle * normalize(axis)) + self.set_orientation(rot * self.get_orientation()) return self def set_euler_angles( @@ -92,13 +73,11 @@ class CameraFrame(Mobject): gamma: float | None = None, units: float = RADIANS ): - if theta is not None: - self.data["euler_angles"][0] = theta * units - if phi is not None: - self.data["euler_angles"][1] = phi * units - if gamma is not None: - self.data["euler_angles"][2] = gamma * units - self.refresh_rotation_matrix() + eulers = self.get_euler_angles() # phi, theta, gamma + for i, var in enumerate([phi, theta, gamma]): + if var is not None: + eulers[i] = var * units + self.set_orientation(Rotation.from_euler('xzy', eulers)) return self def reorient( @@ -124,32 +103,18 @@ class CameraFrame(Mobject): return self.set_euler_angles(gamma=gamma) def increment_theta(self, dtheta: float): - self.data["euler_angles"][0] += dtheta - self.refresh_rotation_matrix() + self.rotate(dtheta, OUT) return self def increment_phi(self, dphi: float): - phi = self.data["euler_angles"][1] - new_phi = clip(phi + dphi, 0, PI) - self.data["euler_angles"][1] = new_phi - self.refresh_rotation_matrix() + self.rotate(dphi, self.get_inverse_camera_rotation_matrix()[0]) return self def increment_gamma(self, dgamma: float): - self.data["euler_angles"][2] += dgamma - self.refresh_rotation_matrix() + self.rotate(dgamma, self.get_inverse_camera_rotation_matrix()[2]) return self - def get_theta(self) -> float: - return self.data["euler_angles"][0] - - def get_phi(self) -> float: - return self.data["euler_angles"][1] - - def get_gamma(self) -> float: - return self.data["euler_angles"][2] - - def get_shape(self) -> tuple[float, float]: + def get_shape(self): return (self.get_width(), self.get_height()) def get_center(self) -> np.ndarray: @@ -167,19 +132,10 @@ class CameraFrame(Mobject): def get_focal_distance(self) -> float: return self.focal_distance * self.get_height() - def get_implied_camera_location(self) -> tuple[float, float, float]: - theta, phi, gamma = self.get_euler_angles() + def get_implied_camera_location(self) -> np.ndarray: + to_camera = self.get_inverse_camera_rotation_matrix()[2] dist = self.get_focal_distance() - x, y, z = self.get_center() - return ( - x + dist * math.sin(theta) * math.sin(phi), - y - dist * math.cos(theta) * math.sin(phi), - z + dist * math.cos(phi) - ) - - def interpolate(self, *args, **kwargs): - super().interpolate(*args, **kwargs) - self.refresh_rotation_matrix() + return self.get_center() + dist * to_camera class Camera(object): @@ -474,7 +430,7 @@ class Camera(object): shader[name].value = tid for name, value in it.chain(self.perspective_uniforms.items(), shader_wrapper.uniforms.items()): try: - if isinstance(value, np.ndarray): + if isinstance(value, np.ndarray) and value.ndim > 0: value = tuple(value) shader[name].value = value except KeyError: diff --git a/manimlib/mobject/coordinate_systems.py b/manimlib/mobject/coordinate_systems.py index 81f8540a..2d21d9d4 100644 --- a/manimlib/mobject/coordinate_systems.py +++ b/manimlib/mobject/coordinate_systems.py @@ -408,10 +408,10 @@ class Axes(VGroup, CoordinateSystem): def coords_to_point(self, *coords: float) -> np.ndarray: origin = self.x_axis.number_to_point(0) - result = origin.copy() - for axis, coord in zip(self.get_axes(), coords): - result += (axis.number_to_point(coord) - origin) - return result + return origin + sum( + axis.number_to_point(coord) - origin + for axis, coord in zip(self.get_axes(), coords) + ) def point_to_coords(self, point: np.ndarray) -> tuple[float, ...]: return tuple([ diff --git a/manimlib/mobject/geometry.py b/manimlib/mobject/geometry.py index b856238c..f555ee52 100644 --- a/manimlib/mobject/geometry.py +++ b/manimlib/mobject/geometry.py @@ -647,7 +647,7 @@ class Elbow(VMobject): class Arrow(Line): CONFIG = { - "stroke_color": GREY_A, + "color": GREY_A, "stroke_width": 5, "tip_width_ratio": 4, "width_to_tip_len": 0.0075, diff --git a/manimlib/mobject/mobject.py b/manimlib/mobject/mobject.py index 6506dda4..17fe9ad0 100644 --- a/manimlib/mobject/mobject.py +++ b/manimlib/mobject/mobject.py @@ -383,7 +383,7 @@ class Mobject(object): h_buff: float | None = None, v_buff: float | None = None, buff_ratio: float | None = None, - h_buff_ratio: float =0.5, + h_buff_ratio: float = 0.5, v_buff_ratio: float = 0.5, aligned_edge: np.ndarray = ORIGIN, fill_rows_first: bool = True @@ -1043,31 +1043,22 @@ class Mobject(object): name: str = "rgbas", recurse: bool = True ): + max_len = 0 if color is not None: rgbs = np.array([color_to_rgb(c) for c in listify(color)]) + max_len = len(rgbs) if opacity is not None: - opacities = listify(opacity) + opacities = np.array(listify(opacity)) + max_len = max(max_len, len(opacities)) - # Color only - if color is not None and opacity is None: - for mob in self.get_family(recurse): - mob.data[name] = resize_array(mob.data[name], len(rgbs)) - mob.data[name][:, :3] = rgbs - - # Opacity only - if color is None and opacity is not None: - for mob in self.get_family(recurse): - mob.data[name] = resize_array(mob.data[name], len(opacities)) - mob.data[name][:, 3] = opacities - - # Color and opacity - if color is not None and opacity is not None: - rgbas = np.array([ - [*rgb, o] - for rgb, o in zip(*make_even(rgbs, opacities)) - ]) - for mob in self.get_family(recurse): - mob.data[name] = rgbas.copy() + for mob in self.get_family(recurse): + if max_len > len(mob.data[name]): + mob.data[name] = resize_array(mob.data[name], max_len) + size = len(mob.data[name]) + if color is not None: + mob.data[name][:, :3] = resize_array(rgbs, size) + if opacity is not None: + mob.data[name][:, 3] = resize_array(opacities, size) return self def set_color(self, color: ManimColor, opacity: float | None = None, recurse: bool = True): diff --git a/manimlib/mobject/number_line.py b/manimlib/mobject/number_line.py index f4e2a496..13b6a13b 100644 --- a/manimlib/mobject/number_line.py +++ b/manimlib/mobject/number_line.py @@ -104,8 +104,8 @@ class NumberLine(Line): def get_tick_marks(self) -> VGroup: return self.ticks - def number_to_point(self, number: float) -> np.ndarray: - alpha = float(number - self.x_min) / (self.x_max - self.x_min) + def number_to_point(self, number: float | np.ndarray) -> np.ndarray: + alpha = (number - self.x_min) / (self.x_max - self.x_min) return interpolate(self.get_start(), self.get_end(), alpha) def point_to_number(self, point: np.ndarray) -> float: @@ -159,7 +159,7 @@ class NumberLine(Line): def add_numbers( self, x_values: Iterable[float] | None = None, - excluding: Iterable[float] | None =None, + excluding: Iterable[float] | None = None, font_size: int = 24, **kwargs ) -> VGroup: diff --git a/manimlib/mobject/svg/drawings.py b/manimlib/mobject/svg/drawings.py index 41e9e907..6eee3428 100644 --- a/manimlib/mobject/svg/drawings.py +++ b/manimlib/mobject/svg/drawings.py @@ -1,6 +1,7 @@ from manimlib.animation.animation import Animation from manimlib.animation.rotation import Rotating from manimlib.constants import * +from manimlib.mobject.boolean_ops import Difference from manimlib.mobject.geometry import Arc from manimlib.mobject.geometry import Circle from manimlib.mobject.geometry import Line @@ -12,6 +13,7 @@ from manimlib.mobject.svg.svg_mobject import SVGMobject from manimlib.mobject.svg.tex_mobject import Tex from manimlib.mobject.svg.tex_mobject import TexText from manimlib.mobject.three_dimensions import Cube +from manimlib.mobject.three_dimensions import Prismify from manimlib.mobject.types.vectorized_mobject import VGroup from manimlib.mobject.types.vectorized_mobject import VMobject from manimlib.utils.config_ops import digest_config @@ -19,6 +21,7 @@ from manimlib.utils.rate_functions import linear from manimlib.utils.space_ops import angle_of_vector from manimlib.utils.space_ops import complex_to_R3 from manimlib.utils.space_ops import rotate_vector +from manimlib.utils.space_ops import midpoint class Checkmark(TexText): @@ -433,3 +436,84 @@ class VectorizedEarth(SVGMobject): ) circle.replace(self) self.add_to_back(circle) + + +class Piano(VGroup): + n_white_keys = 52 + black_pattern = [0, 2, 3, 5, 6] + white_keys_per_octave = 7 + white_key_dims = (0.15, 1.0) + black_key_dims = (0.1, 0.66) + key_buff = 0.02 + white_key_color = WHITE + black_key_color = GREY_E + total_width = 13 + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.add_white_keys() + self.add_black_keys() + self.sort_keys() + self[:-1].reverse_points() + self.set_width(self.total_width) + + def add_white_keys(self): + key = Rectangle(*self.white_key_dims) + key.set_fill(self.white_key_color, 1) + key.set_stroke(width=0) + self.white_keys = key.get_grid(1, self.n_white_keys, buff=self.key_buff) + self.add(*self.white_keys) + + def add_black_keys(self): + key = Rectangle(*self.black_key_dims) + key.set_fill(self.black_key_color, 1) + key.set_stroke(width=0) + + self.black_keys = VGroup() + for i in range(len(self.white_keys) - 1): + if i % self.white_keys_per_octave not in self.black_pattern: + continue + wk1 = self.white_keys[i] + wk2 = self.white_keys[i + 1] + bk = key.copy() + bk.move_to(midpoint(wk1.get_top(), wk2.get_top()), UP) + big_bk = bk.copy() + big_bk.stretch((bk.get_width() + self.key_buff) / bk.get_width(), 0) + big_bk.stretch((bk.get_height() + self.key_buff) / bk.get_height(), 1) + big_bk.move_to(bk, UP) + for wk in wk1, wk2: + wk.become(Difference(wk, big_bk).match_style(wk)) + self.black_keys.add(bk) + self.add(*self.black_keys) + + def sort_keys(self): + self.sort(lambda p: p[0]) + + +class Piano3D(VGroup): + CONFIG = { + "depth_test": True, + "reflectiveness": 1.0, + "stroke_width": 0.25, + "stroke_color": BLACK, + "key_depth": 0.1, + "black_key_shift": 0.05, + } + piano_2d_config = { + "white_key_color": GREY_A, + "key_buff": 0.001 + } + + def __init__(self, **kwargs): + digest_config(self, kwargs) + piano_2d = Piano(**self.piano_2d_config) + super().__init__(*( + Prismify(key, self.key_depth) + for key in piano_2d + )) + self.set_stroke(self.stroke_color, self.stroke_width) + self.apply_depth_test() + # Elevate black keys + for i, key in enumerate(self): + if piano_2d[i] in piano_2d.black_keys: + key.shift(self.black_key_shift * OUT) diff --git a/manimlib/mobject/three_dimensions.py b/manimlib/mobject/three_dimensions.py index 5d39fe7a..329732b1 100644 --- a/manimlib/mobject/three_dimensions.py +++ b/manimlib/mobject/three_dimensions.py @@ -11,6 +11,7 @@ from manimlib.mobject.geometry import Square from manimlib.mobject.geometry import Polygon from manimlib.utils.bezier import interpolate from manimlib.utils.config_ops import digest_config +from manimlib.utils.iterables import adjacent_pairs from manimlib.utils.space_ops import get_norm from manimlib.utils.space_ops import z_to_vector from manimlib.utils.space_ops import compass_directions @@ -280,3 +281,22 @@ class Prism(Cube): Cube.init_points(self) for dim, value in enumerate(self.dimensions): self.rescale_to_fit(value, dim, stretch=True) + + +class Prismify(VGroup): + CONFIG = { + "apply_depth_test": True + } + + def __init__(self, vmobject, depth=1.0, direction=IN, **kwargs): + # At the moment, this assume stright edges + super().__init__(**kwargs) + vect = depth * direction + self.add(vmobject.copy()) + points = vmobject.get_points()[::vmobject.n_points_per_curve] + for p1, p2 in adjacent_pairs(points): + wall = VMobject() + wall.match_style(vmobject) + wall.set_points_as_corners([p1, p2, p2 + vect, p1 + vect]) + self.add(wall) + self.add(vmobject.copy().shift(vect).reverse_points()) diff --git a/manimlib/mobject/types/dot_cloud.py b/manimlib/mobject/types/dot_cloud.py index d44bdc25..3e48aa8f 100644 --- a/manimlib/mobject/types/dot_cloud.py +++ b/manimlib/mobject/types/dot_cloud.py @@ -141,9 +141,13 @@ class TrueDot(DotCloud): super().__init__(points=[center], **kwargs) -class GlowDot(TrueDot): +class GlowDots(DotCloud): CONFIG = { "glow_factor": 2, "radius": DEFAULT_GLOW_DOT_RADIUS, "color": YELLOW, } + + +class GlowDot(GlowDots, TrueDot): + pass diff --git a/manimlib/scene/scene.py b/manimlib/scene/scene.py index ed9233d5..3b649c96 100644 --- a/manimlib/scene/scene.py +++ b/manimlib/scene/scene.py @@ -48,6 +48,7 @@ class Scene(object): "preview": True, "presenter_mode": False, "linger_after_completion": True, + "pan_sensitivity": 3, } def __init__(self, **kwargs): @@ -606,8 +607,8 @@ class Scene(object): frame = self.camera.frame if self.window.is_key_pressed(ord("d")): - frame.increment_theta(-d_point[0]) - frame.increment_phi(d_point[1]) + frame.increment_theta(-self.pan_sensitivity * d_point[0]) + frame.increment_phi(self.pan_sensitivity * d_point[1]) elif self.window.is_key_pressed(ord("s")): shift = -d_point shift[0] *= frame.get_width() / 2 diff --git a/manimlib/utils/bezier.py b/manimlib/utils/bezier.py index fa33cc55..16ec3e10 100644 --- a/manimlib/utils/bezier.py +++ b/manimlib/utils/bezier.py @@ -82,7 +82,13 @@ def partial_quadratic_bezier_points( def interpolate(start: T, end: T, alpha: float) -> T: try: - return (1 - alpha) * start + alpha * end + if isinstance(alpha, float): + return (1 - alpha) * start + alpha * end + # Otherwise, assume alpha is a list or array, and return + # an appropriated shaped array of all corresponding + # interpolations + result = np.outer(1 - alpha, start) + np.outer(alpha, end) + return result.reshape((*np.shape(alpha), *np.shape(start))) except TypeError: log.debug(f"`start` parameter with type `{type(start)}` and dtype `{start.dtype}`") log.debug(f"`end` parameter with type `{type(end)}` and dtype `{end.dtype}`") diff --git a/manimlib/utils/space_ops.py b/manimlib/utils/space_ops.py index 210ceae7..4dde08d5 100644 --- a/manimlib/utils/space_ops.py +++ b/manimlib/utils/space_ops.py @@ -8,6 +8,7 @@ from typing import Callable, Iterable, Sequence import numpy as np import numpy.typing as npt from mapbox_earcut import triangulate_float32 as earcut +from scipy.spatial.transform import Rotation from manimlib.constants import RIGHT from manimlib.constants import DOWN @@ -26,26 +27,36 @@ def cross(v1: np.ndarray, v2: np.ndarray) -> list[np.ndarray]: ] -def get_norm(vect: np.ndarray) -> np.flaoting: +def get_norm(vect: Iterable) -> float: return sum((x**2 for x in vect))**0.5 -# Quaternions -# TODO, implement quaternion type +def normalize(vect: np.ndarray, fall_back: np.ndarray | None = None) -> np.ndarray: + norm = get_norm(vect) + if norm > 0: + return np.array(vect) / norm + elif fall_back is not None: + return fall_back + else: + return np.zeros(len(vect)) + + +# Operations related to rotation def quaternion_mult(*quats: Sequence[float]) -> list[float]: + # Real part is last entry, which is bizzare, but fits scipy Rotation convention if len(quats) == 0: - return [1, 0, 0, 0] + return [0, 0, 0, 1] result = quats[0] for next_quat in quats[1:]: - w1, x1, y1, z1 = result - w2, x2, y2, z2 = next_quat + x1, y1, z1, w1 = result + x2, y2, z2, w2 = next_quat result = [ - w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2, w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2, + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, ] return result @@ -53,99 +64,55 @@ def quaternion_mult(*quats: Sequence[float]) -> list[float]: def quaternion_from_angle_axis( angle: float, axis: np.ndarray, - axis_normalized: bool = False ) -> list[float]: - if not axis_normalized: - axis = normalize(axis) - return [math.cos(angle / 2), *(math.sin(angle / 2) * axis)] + return Rotation.from_rotvec(angle * normalize(axis)).as_quat() -def angle_axis_from_quaternion( - quaternion: Sequence[float] -) -> tuple[float, np.ndarray]: - axis = normalize( - quaternion[1:], - fall_back=[1, 0, 0] - ) - angle = 2 * np.arccos(quaternion[0]) - if angle > TAU / 2: - angle = TAU - angle - return angle, axis +def angle_axis_from_quaternion(quat: Sequence[float]) -> tuple[float, np.ndarray]: + rot_vec = Rotation.from_quat(quat).as_rotvec() + norm = get_norm(rot_vec) + return norm, rot_vec / norm def quaternion_conjugate(quaternion: Iterable) -> list: result = list(quaternion) - for i in range(1, len(result)): + for i in range(3): result[i] *= -1 return result def rotate_vector( - vector: Iterable, - angle: float, + vector: Iterable, + angle: float, axis: np.ndarray = OUT ) -> np.ndarray | list[float]: - if len(vector) == 2: - # Use complex numbers...because why not - z = complex(*vector) * np.exp(complex(0, angle)) - result = [z.real, z.imag] - elif len(vector) == 3: - # Use quaternions...because why not - quat = quaternion_from_angle_axis(angle, axis) - quat_inv = quaternion_conjugate(quat) - product = quaternion_mult(quat, [0, *vector], quat_inv) - result = product[1:] - else: - raise Exception("vector must be of dimension 2 or 3") - - if isinstance(vector, np.ndarray): - return np.array(result) - return result + rot = Rotation.from_rotvec(angle * normalize(axis)) + return np.dot(vector, rot.as_matrix().T) -def thick_diagonal(dim: int, thickness: int = 2) -> np.ndarray: - row_indices = np.arange(dim).repeat(dim).reshape((dim, dim)) - col_indices = np.transpose(row_indices) - return (np.abs(row_indices - col_indices) < thickness).astype('uint8') +def rotate_vector_2d(vector: Iterable, angle: float): + # Use complex numbers...because why not + z = complex(*vector) * np.exp(complex(0, angle)) + return np.array([z.real, z.imag]) -def rotation_matrix_transpose_from_quaternion(quat: Iterable) -> list[list[float]]: - quat_inv = quaternion_conjugate(quat) - return [ - quaternion_mult(quat, [0, *basis], quat_inv)[1:] - for basis in [ - [1, 0, 0], - [0, 1, 0], - [0, 0, 1], - ] - ] +def rotation_matrix_transpose_from_quaternion(quat: Iterable) -> np.ndarray: + return Rotation.from_quat(quat).as_matrix() def rotation_matrix_from_quaternion(quat: Iterable) -> np.ndarray: return np.transpose(rotation_matrix_transpose_from_quaternion(quat)) -def rotation_matrix_transpose(angle: float, axis: np.ndarray) -> list[list[float]]: - if axis[0] == 0 and axis[1] == 0: - # axis = [0, 0, z] case is common enough it's worth - # having a shortcut - sgn = 1 if axis[2] > 0 else -1 - cos_a = math.cos(angle) - sin_a = math.sin(angle) * sgn - return [ - [cos_a, sin_a, 0], - [-sin_a, cos_a, 0], - [0, 0, 1], - ] - quat = quaternion_from_angle_axis(angle, axis) - return rotation_matrix_transpose_from_quaternion(quat) - - def rotation_matrix(angle: float, axis: np.ndarray) -> np.ndarray: """ Rotation in R^3 about a specified axis of rotation. """ - return np.transpose(rotation_matrix_transpose(angle, axis)) + return Rotation.from_rotvec(angle * normalize(axis)).as_matrix() + + +def rotation_matrix_transpose(angle: float, axis: np.ndarray) -> np.ndarray: + return rotation_matrix(angle, axis).T def rotation_about_z(angle: float) -> list[list[float]]: @@ -156,30 +123,19 @@ def rotation_about_z(angle: float) -> list[list[float]]: ] -def z_to_vector(vector: np.ndarray) -> np.ndarray: - """ - Returns some matrix in SO(3) which takes the z-axis to the - (normalized) vector provided as an argument - """ - axis = cross(OUT, vector) - if get_norm(axis) == 0: - if vector[2] > 0: - return np.identity(3) - else: - return rotation_matrix(PI, RIGHT) - angle = np.arccos(np.dot(OUT, normalize(vector))) - return rotation_matrix(angle, axis=axis) - - def rotation_between_vectors(v1, v2) -> np.ndarray: if np.all(np.isclose(v1, v2)): return np.identity(3) return rotation_matrix( angle=angle_between_vectors(v1, v2), - axis=normalize(np.cross(v1, v2)) + axis=np.cross(v1, v2) ) +def z_to_vector(vector: np.ndarray) -> np.ndarray: + return rotation_between_vectors(OUT, vector) + + def angle_of_vector(vector: Sequence[float]) -> float: """ Returns polar coordinate theta when vector is project on xy plane @@ -192,7 +148,10 @@ def angle_between_vectors(v1: np.ndarray, v2: np.ndarray) -> float: Returns the angle between two 3D vectors. This angle will always be btw 0 and pi """ - return math.acos(clip(np.dot(normalize(v1), normalize(v2)), -1, 1)) + n1 = get_norm(v1) + n2 = get_norm(v2) + cos_angle = np.dot(v1, v2) / (n1 * n2) + return math.acos(clip(cos_angle, -1, 1)) def project_along_vector(point: np.ndarray, vector: np.ndarray) -> np.ndarray: @@ -200,19 +159,6 @@ def project_along_vector(point: np.ndarray, vector: np.ndarray) -> np.ndarray: return np.dot(point, matrix.T) -def normalize( - vect: np.ndarray, - fall_back: np.ndarray | None = None -) -> np.ndarray: - norm = get_norm(vect) - if norm > 0: - return np.array(vect) / norm - elif fall_back is not None: - return fall_back - else: - return np.zeros(len(vect)) - - def normalize_along_axis( array: np.ndarray, axis: np.ndarray, @@ -227,7 +173,7 @@ def normalize_along_axis( def get_unit_normal( v1: np.ndarray, v2: np.ndarray, - tol: float=1e-6 + tol: float = 1e-6 ) -> np.ndarray: v1 = normalize(v1) v2 = normalize(v2) @@ -246,6 +192,12 @@ def get_unit_normal( ### +def thick_diagonal(dim: int, thickness: int = 2) -> np.ndarray: + row_indices = np.arange(dim).repeat(dim).reshape((dim, dim)) + col_indices = np.transpose(row_indices) + return (np.abs(row_indices - col_indices) < thickness).astype('uint8') + + def compass_directions(n: int = 4, start_vect: np.ndarray = RIGHT) -> np.ndarray: angle = TAU / n return np.array([