diff --git a/phi/geom/__init__.py b/phi/geom/__init__.py index 14c8fd063..6a72e809d 100644 --- a/phi/geom/__init__.py +++ b/phi/geom/__init__.py @@ -13,8 +13,8 @@ # --- Low-level functions --- from ._geom import Geometry, GeometryException, Point, assert_same_rank, invert, sample_function -from ._functions import normal_from_slope, clip_length, cross -from ._transform import scale, rotate, rotation_matrix, rotation_angles, rotation_matrix_from_axis_and_angle, rotation_matrix_from_directions +from ._functions import normal_from_slope, clip_length, cross, rotation_matrix, rotation_angles, rotation_matrix_from_axis_and_angle, rotation_matrix_from_directions +from ._transform import scale, rotate # --- Geometry types --- from ._box import Box, BaseBox, Cuboid, bounding_box diff --git a/phi/geom/_box.py b/phi/geom/_box.py index e78805a6c..570175a04 100644 --- a/phi/geom/_box.py +++ b/phi/geom/_box.py @@ -75,7 +75,7 @@ def global_to_local(self, global_position: Tensor, scale=True, origin='lower') - assert origin in ['lower', 'center', 'upper'] origin_loc = getattr(self, origin) pos = global_position if math.always_close(origin_loc, 0) else global_position - origin_loc - pos = rotate(pos, self.rotation_matrix, invert=True) + pos = rotate(pos, self._rot_or_none, invert=True) if scale: pos /= (self.half_size if origin == 'center' else self.size) return pos @@ -84,7 +84,11 @@ def local_to_global(self, local_position, scale=True, origin='lower'): assert origin in ['lower', 'center', 'upper'] origin_loc = getattr(self, origin) pos = local_position * (self.half_size if origin == 'center' else self.size) if scale else local_position - return rotate(pos, self.rotation_matrix) + origin_loc + return rotate(pos, self._rot_or_none) + origin_loc + + @property + def _rot_or_none(self): + raise NotImplementedError def largest(self, dim: DimFilter) -> 'BaseBox': dim = self.shape.without('vector').only(dim) @@ -99,7 +103,7 @@ def smallest(self, dim: DimFilter) -> 'BaseBox': return Box(math.max(self.lower, dim), math.min(self.upper, dim)) def lies_inside(self, location: Tensor): - assert self.rotation_matrix is None, f"Please override lies_inside() for rotated boxes" + assert self._rot_or_none is None, f"Please override lies_inside() for rotated boxes" bool_inside = (location >= self.lower) & (location <= self.upper) bool_inside = math.all(bool_inside, 'vector') bool_inside = math.any(bool_inside, self.shape.instance - instance(location)) # union for instance dimensions @@ -152,8 +156,7 @@ def approximate_closest_surface(self, location: Tensor) -> Tuple[Tensor, Tensor, max_sgn_dist = math.max(sgn_surf_delta, 'vector') normal_axis = max_sgn_dist == sgn_surf_delta # ToDo only one if inside normal = math.vec_normalize(normal_axis * math.sign(loc_to_center)) - if self.rotation_matrix is not None: - normal = rotate(normal, self.rotation_matrix) + normal = rotate(normal, self._rot_or_none) surf_to_center = math.where(normal_axis, math.sign(loc_to_center) * self.half_size, loc_to_center) closest_to_center = math.clip(surf_to_center, -self.half_size, self.half_size) surface_pos = self.local_to_global(closest_to_center, scale=False, origin='center') @@ -190,7 +193,7 @@ def sample_uniform_surface(self, *shape: Shape) -> Tensor: return samples def corner_representation(self) -> 'Box': - assert self.rotation_matrix is None, f"corner_representation does not support rotations" + assert self._rot_or_none is None, f"corner_representation does not support rotations" return Box(self.lower, self.upper) box = corner_representation @@ -198,6 +201,8 @@ def corner_representation(self) -> 'Box': def center_representation(self, size_variable=True) -> 'Cuboid': return Cuboid(self.center, self.half_size, size_variable=size_variable) + cuboid = center_representation + def contains(self, other: 'BaseBox'): """ Tests if the other box lies fully inside this box. """ return np.all(other.lower >= self.lower) and np.all(other.upper <= self.upper) @@ -215,7 +220,7 @@ def boundary_faces(self) -> Dict[Any, Dict[str, slice]]: @property def faces(self) -> 'Geometry': - return Cuboid(self.face_centers, self._half_size, self._rotation_matrix, size_variable=False) + return Cuboid(self.face_centers, self.half_size, self._rot_or_none, size_variable=False) @property def face_centers(self) -> Tensor: @@ -224,7 +229,7 @@ def face_centers(self) -> Tensor: @property def face_normals(self) -> Tensor: unit_vectors = math.to_float(math.range(self.shape['vector']) == math.range(dual(**self.shape['vector'].untyped_dict))) - vectors = rotate(unit_vectors, self.rotation_matrix) + vectors = rotate(unit_vectors, self._rot_or_none) return vectors * math.vec(dual('side'), lower=-1, upper=1) @property @@ -238,7 +243,7 @@ def face_shape(self) -> Shape: return self.shape.without('vector') & dual(side='lower,upper') & dual(**self.shape['vector'].untyped_dict) @property - def corners(self): + def corners(self) -> Tensor: to_face = self.face_normals[{'~side': 'upper'}] * math.rename_dims(self.half_size, 'vector', dual) lower_upper = math.meshgrid(math.dual, **{dim: [-1, 1] for dim in self.vector.item_names}, stack_dim=dual('vector')) # (x=2, y=2, ... vector=x,y,...) to_corner = math.sum(lower_upper * to_face, '~vector') @@ -349,6 +354,10 @@ def __variable_attrs__(self): def __value_attrs__(self): return '_lower', '_upper' + @property + def _rot_or_none(self): + return None + @property def shape(self): if self._lower is None or self._upper is None: @@ -381,10 +390,10 @@ def rotation_matrix(self) -> Optional[Tensor]: @property def is_size_variable(self): - raise False + return False def at(self, center: Tensor) -> 'BaseBox': - return Cuboid(center, self.half_size, self.rotation_matrix) + return Cuboid(center, self.half_size, None) def shifted(self, delta, **delta_by_dim) -> 'Box': return Box(self.lower + delta, self.upper + delta) @@ -506,14 +515,18 @@ def upper(self): @property def rotation_matrix(self) -> Optional[Tensor]: - return self._rotation_matrix + return rotation_matrix(self._rotation_matrix, self.shape['vector'], none_to_unit=True) + + @property + def _rot_or_none(self): + return None if self._rotation_matrix is None else self.rotation_matrix @property def is_size_variable(self): return self._size_variable def at(self, center: Tensor) -> 'Cuboid': - return Cuboid(center, self.half_size, self.rotation_matrix, size_variable=self._size_variable) + return Cuboid(center, self.half_size, self._rotation_matrix, size_variable=self._size_variable) def rotated(self, angle) -> 'Cuboid': if self._rotation_matrix is None: diff --git a/phi/geom/_convert.py b/phi/geom/_convert.py index 9b3ab9e64..b4f30f81b 100644 --- a/phi/geom/_convert.py +++ b/phi/geom/_convert.py @@ -1,7 +1,7 @@ import numpy as np from phiml import math -from phiml.math import wrap, instance, batch, DimFilter, Tensor, spatial, pack_dims, dual, stack +from phiml.math import wrap, instance, batch, DimFilter, Tensor, spatial, pack_dims, dual, stack, to_int32, maximum from ._box import Cuboid, BaseBox from ._sphere import Sphere from ._functions import plane_sgn_dist @@ -9,6 +9,7 @@ from ._mesh import Mesh, mesh_from_numpy, mesh from ._sdf import SDF, numpy_sdf from ._sdf_grid import sample_sdf, SDFGrid +from ._spline_solid import SplineSolid def as_sdf(geo: Geometry, bounds=None, rel_margin=None, abs_margin=0., separate: DimFilter = None, method='auto') -> SDF: @@ -100,6 +101,12 @@ def surface_mesh(geo: Geometry, raise NotImplementedError("Only 3D SDF currently supported") if isinstance(geo, NoGeometry): return mesh_from_numpy([], [], element_rank=2) + # --- Determine resolution --- + if rel_dx is None and abs_dx is None: + rel_dx = 0.005 + rel_dx = None if rel_dx is None else rel_dx * geo.bounding_box().size.max + dx = math.minimum(rel_dx, abs_dx, allow_none=True) + # --- Check special cases --- if method == 'auto' and isinstance(geo, BaseBox): assert rel_dx is None and abs_dx is None, f"When method='auto', boxes will always use their corners as vertices. Leave rel_dx,abs_dx unspecified or pass 'lewiner' or 'lorensen' as method" vertices = pack_dims(geo.corners, dual, instance('vertices')) @@ -114,19 +121,18 @@ def surface_mesh(geo: Geometry, return mesh(vertices, faces, element_rank=2) elif method == 'auto' and isinstance(geo, Sphere): pass # ToDo analytic solution + elif method == 'auto' and isinstance(geo, SplineSolid): + return geo.surface_mesh() # ToDo resolution from dx + # --- Build mesh from SDF --- if isinstance(geo, SDFGrid): assert rel_dx is None and abs_dx is None, f"When creating a surface mesh from an SDF grid, rel_dx and abs_dx are determined from the grid and must be specified as None" sdf_grid = geo else: - if rel_dx is None and abs_dx is None: - rel_dx = 0.005 - rel_dx = None if rel_dx is None else rel_dx * geo.bounding_radius() * 2 - dx = math.minimum(rel_dx, abs_dx, allow_none=True) if isinstance(geo, SDF): sdf = geo else: sdf = as_sdf(geo, rel_margin=0, abs_margin=dx) - resolution = math.to_int32(math.round(sdf.bounds.size / dx)) + resolution = maximum(1, to_int32(math.round(sdf.bounds.size / dx))) resolution = spatial(**resolution.vector) sdf_grid = sample_sdf(sdf, sdf.bounds, resolution) from skimage.measure import marching_cubes diff --git a/phi/geom/_cylinder.py b/phi/geom/_cylinder.py index 312e6c754..b68c93cbb 100644 --- a/phi/geom/_cylinder.py +++ b/phi/geom/_cylinder.py @@ -7,7 +7,7 @@ from phiml.math._magic_ops import all_attributes from phiml.dataclasses import replace, sliceable from ._geom import Geometry -from ._transform import rotate, rotation_matrix, rotation_matrix_from_directions +from ._functions import rotate_vector as rotate, rotation_matrix, rotation_matrix_from_directions from ._sphere import Sphere diff --git a/phi/geom/_functions.py b/phi/geom/_functions.py index 47d9f9bd9..98bbb4586 100644 --- a/phi/geom/_functions.py +++ b/phi/geom/_functions.py @@ -1,13 +1,30 @@ -from typing import Sequence, Union +from typing import Sequence, Union, Optional from phiml import math -from phiml.math import Tensor, channel, Shape, vec_normalize, vec, sqrt, maximum, clip, vec_squared, vec_length, where, stack, dual, argmin, safe_div -from phiml.math._shape import parse_dim_order, DimFilter +from phiml.math import Tensor, channel, Shape, normalize, vec, sqrt, maximum, clip, vec_squared, norm, where, stack, dual, argmin, safe_div, arange, wrap, to_float, rename_dims +from phiml.math._shape import parse_dim_order, DimFilter, shape # No dependence on Geometry +def vec_normalize(x, epsilon=None, allow_infinite=False, allow_zero=False): + return normalize(x, 'vector', epsilon, allow_infinite=allow_infinite, allow_zero=allow_zero) + + +def vec_length(x, eps=None): + return norm(x, 'vector', eps) + + +def rotate_vector(x, rot, invert=False): + assert 'vector' in x.shape, f"vector must have exactly a channel dimension named 'vector'" + matrix = rotation_matrix(rot) + if invert: + matrix = rename_dims(matrix, '~vector,vector', matrix.shape['vector'] + matrix.shape['~vector']) + assert matrix.vector.dual.size == x.vector.size, f"Rotation matrix from {rot.shape} is {matrix.vector.dual.size}D but vector {x.shape} is {x.vector.size}D." + return math.dot(matrix, '~vector', x, 'vector') + + def cross(vec1: Tensor, vec2: Tensor) -> Tensor: """ Computes the cross product of two vectors in 2D. @@ -197,3 +214,175 @@ def distance_line_point(line_offset: Tensor, line_direction: Tensor, point: Tens return c +def closest_normal_vector(target: Tensor, normal: Tensor, eps=1e-10): + target_normal = normal * (target.vector @ normal.vector) + target_ortho = target - target_normal + return vec_normalize(target_ortho, 'vector', eps) + + + + +def rotation_matrix(x: float | math.Tensor | None, matrix_dim=channel('vector'), none_to_unit=False) -> Optional[Tensor]: + """ + Create a 2D or 3D rotation matrix from the corresponding angle(s). + + Args: + x: + 2D: scalar angle + 3D: Either vector pointing along the rotation axis with rotation angle as length or Euler angles. + Euler angles need to be laid out along a `angle` channel dimension with dimension names listing the spatial dimensions. + E.g. a 90° rotation about the z-axis is represented by `vec('angles', x=0, y=0, z=PI/2)`. + If a rotation matrix is passed for `angle`, it is returned without modification. + matrix_dim: Matrix dimension for 2D rotations. In 3D, the channel dimension of angle is used. + + Returns: + Matrix containing `matrix_dim` in primal and dual form as well as all non-channel dimensions of `x`. + """ + if x is None and not none_to_unit: + return None + elif x is None: + return to_float(arange(matrix_dim) == arange(matrix_dim.as_dual())) + if isinstance(x, Tensor) and x.dtype == object: # possibly None in matrices + return math.map(rotation_matrix, x, dims=object, matrix_dim=matrix_dim, none_to_unit=none_to_unit) + if isinstance(x, Tensor) and '~vector' in x.shape and 'vector' in x.shape.channel and x.shape.get_size('~vector') == x.shape.get_size('vector'): + return x # already a rotation matrix + elif 'angle' in shape(x) and shape(x).get_size('angle') == 3: # 3D Euler angles + assert channel(x).rank == 1 and channel(x).size == 3, f"x for 3D rotations needs to be a 3-vector but got {x}" + s1, s2, s3 = math.sin(x).angle # x, y, z + c1, c2, c3 = math.cos(x).angle + matrix_dim = matrix_dim.with_size(shape(x).get_item_names('angle')) + return wrap([[c3 * c2, c3 * s2 * s1 - s3 * c1, c3 * s2 * c1 + s3 * s1], + [s3 * c2, s3 * s2 * s1 + c3 * c1, s3 * s2 * c1 - c3 * s1], + [-s2, c2 * s1, c2 * c1]], matrix_dim, matrix_dim.as_dual()) # Rz * Ry * Rx (1. rotate about X by first angle) + elif 'vector' in shape(x) and shape(x).get_size('vector') == 3: # 3D axis + x + angle = vec_length(x) + s, c = math.sin(angle), math.cos(angle) + t = 1 - c + k1, k2, k3 = normalize(x, epsilon=1e-12).vector + matrix_dim = matrix_dim.with_size(shape(x).get_item_names('vector')) + return wrap([[c + k1**2 * t, k1 * k2 * t - k3 * s, k1 * k3 * t + k2 * s], + [k2 * k1 * t + k3 * s, c + k2**2 * t, k2 * k3 * t - k1 * s], + [k3 * k1 * t - k2 * s, k3 * k2 * t + k1 * s, c + k3**2 * t]], matrix_dim, matrix_dim.as_dual()) + else: # 2D rotation + sin = wrap(math.sin(x)) + cos = wrap(math.cos(x)) + return wrap([[cos, -sin], [sin, cos]], matrix_dim, matrix_dim.as_dual()) + + +def rotation_angles(rot: Tensor): + """ + Compute the scalar x in 2D or the Euler angles in 3D from a given rotation matrix. + This function returns one valid solution but often, there are multiple solutions. + + Args: + rot: Rotation matrix as created by `phi.math.rotation_matrix()`. + Must have exactly one channel and one dual dimension with equally-ordered elements. + + Returns: + Scalar x in 2D, Euler angles + """ + assert channel(rot).rank == 1 and dual(rot).rank == 1, f"Rotation matrix must have one channel and one dual dimension but got {rot.shape}" + if channel(rot).size == 2: + cos = rot[{channel: 0, dual: 0}] + sin = rot[{channel: 1, dual: 0}] + return math.arctan(sin, divide_by=cos) + elif channel(rot).size == 3: + a2 = -math.arcsin(rot[{channel: 2, dual: 0}]) # ToDo handle [2, 0] == 1 (i.e. cos_theta == 0) + cos2 = math.cos(a2) + a1 = math.arctan(rot[{channel: 2, dual: 1}] / cos2, divide_by=rot[{channel: 2, dual: 2}] / cos2) + a3 = math.arctan(rot[{channel: 1, dual: 0}] / cos2, divide_by=rot[{channel: 0, dual: 0}] / cos2) + regular_sol = stack([a1, a2, a3], channel(angle=channel(rot).item_names[0])) + # --- pole case cos(theta) == 1 --- + a3_pole = 0 # unconstrained + bottom_pole = rot[{channel: 2, dual: 0}] < 0 + a2_pole = math.where(bottom_pole, 1.57079632679, -1.57079632679) + a1_pole = math.where(bottom_pole, math.arctan(rot[{channel: 0, dual: 1}], divide_by=rot[{channel: 0, dual: 2}]), math.arctan(-rot[{channel: 0, dual: 1}], divide_by=-rot[{channel: 0, dual: 2}])) + pole_sol = stack([a1_pole, a2_pole, a3_pole], channel(regular_sol)) + return math.where(abs(rot[{channel: 2, dual: 0}]) >= 1, pole_sol, regular_sol) + else: + raise ValueError(f"") + + +def rotation_matrix_from_directions(source_dir: Tensor, target_dir: Tensor, vec_dim: str = 'vector', epsilon=None) -> Tensor: + """ + Computes a rotation matrix A, such that `target_dir = A @ source_dir` + + Args: + source_dir: Two or three-dimensional vector. `Tensor` with channel dim called 'vector'. + target_dir: Two or three-dimensional vector. `Tensor` with channel dim called 'vector'. + + Returns: + Rotation matrix as `Tensor` with 'vector' dim and its dual counterpart. + """ + if source_dir.vector.size == 3: + axis, angle = axis_angle_from_directions(source_dir, target_dir, vec_dim, epsilon=epsilon) + return rotation_matrix_from_axis_and_angle(axis, angle, is_axis_normalized=False, epsilon=epsilon) + raise NotImplementedError + + +def axis_angle_from_directions(source_dir: Tensor, target_dir: Tensor, vec_dim: str = 'vector', epsilon=None) -> tuple[Tensor, Tensor]: + if source_dir.vector.size == 3: + source_dir = normalize(source_dir, vec_dim, epsilon=epsilon) + target_dir = normalize(target_dir, vec_dim, epsilon=epsilon) + axis = cross(source_dir, target_dir) + lim = 1-epsilon if epsilon is not None else 1 + angle = math.arccos(math.clip(source_dir.vector @ target_dir.vector, -lim, lim)) + return axis, angle + raise NotImplementedError + + +def rotation_matrix_from_axis_and_angle(axis: Tensor, angle: float | Tensor, vec_dim='vector', is_axis_normalized=False, epsilon=1e-5) -> Tensor: + """ + Computes a rotation matrix that rotates by `angle` around `axis`. + + Args: + axis: 3D vector. `Tensor` with channel dim called 'vector'. + angle: Rotation angle. + is_axis_normalized: Whether `axis` has length 1. + epsilon: Minimum axis length. For shorter axes, the unit matrix is returned. + + Returns: + Rotation matrix as `Tensor` with 'vector' dim and its dual counterpart. + """ + if axis.vector.size == 3: # Rodrigues' rotation formula + axis = normalize(axis, vec_dim, epsilon=epsilon, allow_zero=False) if not is_axis_normalized else axis + kx, ky, kz = axis.vector + s = math.sin(angle) + c = 1 - math.cos(angle) + return wrap([ + (1 - c*(ky*ky+kz*kz), -kz*s + c*(kx*ky), ky*s + c*(kx*kz)), + ( kz*s + c*(kx*ky), 1 - c*(kx*kx+kz*kz), -kx*s + c*(ky * kz)), + ( -ky*s + c*(kx*kz), kx*s + c*(ky * kz), 1 - c*(kx*kx+ky*ky)), + ], axis.shape['vector'], axis.shape['vector'].as_dual()) + raise NotImplementedError + + +def matching_rotations(*matrices: Optional[Tensor]) -> Sequence[Tensor]: + """ + Replaces `None` rotations with unit matrices if any of `matrices` is not `None`. + """ + if all(m is None for m in matrices) or all(m is not None for m in matrices): + return matrices + some = [m for m in matrices if m is not None][0] + unit_matrix = arange(some.shape.primal) == arange(some.shape.dual) + return [unit_matrix if m is None else m for m in matrices] + + +def sample_helix(offset: Tensor, axis: Tensor, start: Tensor, end: Tensor, t: Tensor): + start -= offset + end -= offset + # --- Component along axis (h) --- + h_start = start.vector @ axis.vector + h_end = end.vector @ axis.vector + h = t * h_end + (1-t) * h_start + start = start - h_start * axis + end = end - h_end * axis + # --- Radius --- + r_start = math.norm(start, 'vector') + r_end = math.norm(end, 'vector') + r = t * r_end + (1-t) * r_start + # --- Angle --- + rot_axis, angle = axis_angle_from_directions(start, end) + a = angle * t + # --- Combine --- + return offset + h * axis + r/r_start * rotate_vector(start, rot_axis * a) diff --git a/phi/geom/_spline_solid.py b/phi/geom/_spline_solid.py new file mode 100644 index 000000000..6df77670d --- /dev/null +++ b/phi/geom/_spline_solid.py @@ -0,0 +1,215 @@ +from dataclasses import dataclass +from functools import cached_property +from typing import Union, Tuple, Dict + +from phiml import math +from phiml.dataclasses import replace +from phiml.math import Tensor, Shape, spatial, dual, dprod, stack, dmax, dmin, dmean, wrap, linspace, PI, where, concat, clip, dpack, \ + mean, arcsin +from phiml.math.extrapolation import ZERO_GRADIENT +from ._functions import vec_length, cross, vec_normalize, rotate_vector, sample_helix +from ._geom import Geometry +from ._mesh_builder import MeshBuilder + + +@dataclass +class SplineSolid(Geometry): + """ + Internal coordinates (u,v) are in the range [0, N] where N is the number of points along that axis. + """ + + points: Tensor + """2D spatial tensor of points""" + thickness: Tensor + """2D spatial tensor matching vertices""" + fillet: dict[str, Tensor] + """Roundness by vertex by boundary, such as 'u-', 'u+', 'v-', 'v+'. 1 puts a full cylinder at the edge, 0 is a sharp edge.""" + order: dict[str, int] + """Spline order along each axis, e.g. {'u': 1, 'v': 2}""" + + def __post_init__(self): + assert 'vector' in self.points.shape + for dim, o in self.order.items(): + assert f'{dim}-' in self.fillet and f'{dim}+' in self.fillet + assert dim in spatial(self.points) + assert o < self.points.shape.get_size(dim) + + @property + def shape(self) -> Shape: + return self.points.shape + + @cached_property + def center(self) -> Tensor: + lin_center = math.neighbor_mean(self.points, spatial, None) + return lin_center + + @cached_property + def volume(self) -> Tensor: + dx = math.spatial_gradient(self.points, difference='forward', stack_dim=dual('gradient'))[{d: slice(0, -1) for d in spatial(self.points)}] + return dprod(vec_length(dx)) + + @cached_property + def corner_shape(self) -> Shape: + return spatial(self.points).as_dual().with_sizes('lo,up') + (spatial(self.points) - 1) + + @cached_property + def corners(self) -> Tensor: + result = [self.points[{dim[1:]: slice(o, None if o else -1) for dim, o in offset.items()}] for offset in self.corner_shape.dual.meshgrid()] + return stack(result, self.corner_shape.dual) + + @cached_property + def corner_radii(self) -> Tensor: + result = [self.thickness[{dim: slice(o, None if o else -1) for dim, o in offset.items()}] for offset in self.corner_shape.dual.meshgrid()] + return stack(result, self.corner_shape.dual) + + @cached_property + def _central_point_tangents(self): + tangents = {dim.name[1:]: dmean(self.corners[{dim: 1}] - self.corners[{dim: 0}]) for dim in self.corner_shape.dual} + tangents = stack(tangents, '~tangents') + return math.neighbor_mean(tangents, spatial, ZERO_GRADIENT, extend_bounds=(1, 0)) + + @cached_property + def _central_point_normals(self): + t1, t2 = self._central_point_tangents.tangents.dual + return vec_normalize(cross(t1, t2)) + + @cached_property + def _surface_points(self): + front_back = wrap([-1, 1], dual(side='front,back')) + return self.points + front_back * .5 * self.thickness * self._central_point_normals + + @cached_property + def _surface_corners(self): + result = [self._surface_points[{dim[1:]: slice(o, None if o else -1) for dim, o in offset.items()}] for offset in self.corner_shape.dual.meshgrid()] + return stack(result, self.corner_shape.dual) + + @cached_property + def _surface_point_tangents(self): + tangents = {dim.name[1:]: mean(self._surface_corners[{dim: 1}] - self._surface_corners[{dim: 0}], self.corner_shape.dual) for dim in self.corner_shape.dual} + tangents = stack(tangents, '~tangents') + return math.neighbor_mean(tangents, spatial, ZERO_GRADIENT, extend_bounds=(1, 0)) + + @cached_property + def _surface_point_normals(self): + front_back = wrap([1, -1], dual(side='front,back')) + t1, t2 = self._surface_point_tangents.tangents.dual + return vec_normalize(cross(t1, t2)) * -front_back + + def surface_mesh(self, cyl_segments=5, corner_segments=2): # mesh can get self-intersections for larger corner_segments + front_back = wrap([1, -1], dual(side='front,back')) + u, v = spatial(self.points).names[::-1] + mb = MeshBuilder(2) + spline_idx = mb.new_quads('spline', self.points + .5 * self.thickness * self._surface_point_normals, flip=front_back > 0) + # --- Rounded edges --- + for edge_, fillet in self.fillet.items(): + edge, is_upper = edge_[:-1], edge_[-1] == '+' + other_edge = spatial(self.points).without(edge).name + fillet = clip(fillet, 1e-5, .99) + edge_points = self.points[{edge: -1 if is_upper else 0}] + edge_tangent_c = self._central_point_tangents[{'~tangents': other_edge, edge: -1 if is_upper else 0}] + edge_tangent_s = self._surface_point_tangents[{'~tangents': other_edge, edge: -1 if is_upper else 0}] + edge_tangent = edge_tangent_c # ToDo rotate to match thickness change along `edge` + edge_normals = self._surface_point_normals[{edge: -1 if is_upper else 0}] + edge_normals_c = self._central_point_normals[{edge: -1 if is_upper else 0}] + edge_thickness = self.thickness[{edge: -1 if is_upper else 0}] + cyl_center = edge_points + edge_normals * .5 * edge_thickness * (1-fillet) + incline_angle = arcsin(cross(edge_normals, edge_normals_c).vector @ edge_tangent_c.vector) * (-1 if is_upper else 1) + angle = linspace(0, PI/2 - incline_angle, spatial(cyl=cyl_segments+1)).cyl[1:] + cv = cyl_center + rotate_vector(.5 * fillet * edge_thickness * edge_normals, edge_tangent * angle * front_back * where(is_upper ^ (edge == v), 1, -1)) + e_idx = mb.add_vertices(edge_, cv) + mb.add_quads(concat([spline_idx[{edge: -1 if is_upper else 0}], e_idx], 'cyl', expand_values=True), flip=(front_back > 0) ^ is_upper ^ (edge == v)) + mb.add_quads(dpack(e_idx.cyl[-1], 'side:s'), flip=is_upper ^ (edge != v)) + # --- Rounded corners --- + for u_up in '-+': + for v_up in '-+': + u_verts = mb.vertices(f"{u}{u_up}")[{v: 0 if v_up == '-' else -1}] + v_verts = mb.vertices(f"{v}{v_up}")[{u: 0 if u_up == '-' else -1}] + normal = self._surface_point_normals[{u: 0 if u_up == '-' else -1, v: 0 if v_up == '-' else -1}] + corner = self.points[{u: 0 if u_up == '-' else -1, v: 0 if v_up == '-' else -1}] + t = linspace(0, 1, spatial(corner=corner_segments+1)).corner[1:-1] + corner_points_round = sample_helix(corner, normal, u_verts, v_verts, t) + corner_points_straight = t * u_verts + (1-t) * v_verts + fillet_u = self.fillet[f"{u}{u_up}"][{v: 0 if v_up == '-' else -1}] + fillet_v = self.fillet[f"{v}{v_up}"][{u: 0 if u_up == '-' else -1}] + corner_roundness = fillet_u * fillet_v + corner_points = (1-corner_roundness) * corner_points_straight + corner_roundness * corner_points_round + c_idx = mb.add_vertices(f"{u}{u_up},{v}{v_up}", corner_points) + # --- Triangles --- + u_idx = mb.vertex_indices(f"{u}{u_up}")[{v: 0 if v_up == '-' else -1}] + v_idx = mb.vertex_indices(f"{v}{v_up}")[{u: 0 if u_up == '-' else -1}] + c_idx = concat([u_idx, c_idx, v_idx], 'corner', expand_values=True) + mb.add_tris(spline_idx[{v: 0 if v_up == '-' else -1, u: 0 if u_up == '-' else -1}], c_idx.cyl[0], flip=(front_back > 0) ^ (u_up != v_up)) + # --- Quads --- + mb.add_quads(c_idx, (front_back > 0) ^ (u_up != v_up)) + mb.add_quads(dpack(c_idx.cyl[-1], 'side:s'), flip=u_up != v_up) + # mb.debug_show(normals=True) + return mb.build_mesh() + + def lies_inside(self, location: Tensor) -> Tensor: + return self.approximate_signed_distance(location) <= 0 + + def approximate_closest_surface(self, location: Tensor) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor]: + return self._surface_mesh.approximate_closest_surface(location) + # find closest segment (face) + # guess u,v from first-order + # fixed iter optimize + # boundary handling... + # normal to surface -> walk according to crease, put sphere + + @cached_property + def _surface_mesh(self): + return self.surface_mesh(5, 3) + + def approximate_signed_distance(self, location: Tensor) -> Tensor: + return self.approximate_closest_surface(location)[0] + + def sample_uniform(self, *shape: math.Shape) -> Tensor: + raise NotImplementedError + + def bounding_radius(self) -> Tensor: + return dmax(vec_length(self.corners - self.center) + self.corner_radii) + + def bounding_half_extent(self) -> Tensor: + return .5 * (dmax(self.corners) - dmin(self.corners)) + dmax(self.corner_radii) + + def at(self, center: Tensor) -> 'Geometry': + assert spatial(self.points) in center.shape, f"NURBSSolid.at() must be given new vertex positions" + return replace(self, vertices=center) + + def rotated(self, angle: Union[float, Tensor]) -> 'Geometry': + raise NotImplementedError + + def scaled(self, factor: Union[float, Tensor]) -> 'Geometry': + raise NotImplementedError + + @property + def faces(self) -> 'Geometry': + raise NotImplementedError + + @property + def face_centers(self) -> Tensor: + raise NotImplementedError + + @property + def face_areas(self) -> Tensor: + raise NotImplementedError + + @property + def face_normals(self) -> Tensor: + raise NotImplementedError + + @property + def boundary_elements(self) -> Dict[str, Dict[str, slice]]: + return {} + + @property + def boundary_faces(self) -> Dict[str, Dict[str, slice]]: + raise NotImplementedError # return outer edges + + +def spline(order: int, derivative: int, uv, points: Tensor) -> Tensor: + if order == 1 and derivative == 0: + # left_idx = minimum(to_int32(uv), points.shape.get_size() + left_idx = 0 + right_idx = left_idx + 1 + left, right = points diff --git a/phi/geom/_transform.py b/phi/geom/_transform.py index 9fdfca1a8..02f17a05b 100644 --- a/phi/geom/_transform.py +++ b/phi/geom/_transform.py @@ -1,10 +1,8 @@ -from typing import Optional, Sequence - -from phiml import math -from phiml.math import Tensor, channel, rename_dims, wrap, shape, normalize, dual, stack, length, arange, to_float, Shape +from phiml.math import Tensor +from phiml.math import Tensor +from ._functions import rotate_vector from ._geom import Geometry, GeometricType -from ._functions import cross def scale(obj: GeometricType, scale: float | Tensor, pivot: Tensor = None, dim='vector') -> GeometricType: @@ -57,155 +55,7 @@ def rotate(obj: GeometricType, rot: float | Tensor | None, invert=False, pivot: center = pivot + rotate(obj.center - pivot, rot) return obj.rotated(rot).at(center) elif isinstance(obj, Tensor): - assert 'vector' in obj.shape, f"vector must have exactly a channel dimension named 'vector'" - matrix = rotation_matrix(rot) - if invert: - matrix = rename_dims(matrix, '~vector,vector', matrix.shape['vector'] + matrix.shape['~vector']) - assert matrix.vector.dual.size == obj.vector.size, f"Rotation matrix from {rot.shape} is {matrix.vector.dual.size}D but vector {obj.shape} is {obj.vector.size}D." - return math.dot(matrix, '~vector', obj, 'vector') - - -def rotation_matrix(x: float | math.Tensor | None, matrix_dim=channel('vector'), none_to_unit=False) -> Optional[Tensor]: - """ - Create a 2D or 3D rotation matrix from the corresponding angle(s). - - Args: - x: - 2D: scalar angle - 3D: Either vector pointing along the rotation axis with rotation angle as length or Euler angles. - Euler angles need to be laid out along a `angle` channel dimension with dimension names listing the spatial dimensions. - E.g. a 90° rotation about the z-axis is represented by `vec('angles', x=0, y=0, z=PI/2)`. - If a rotation matrix is passed for `angle`, it is returned without modification. - matrix_dim: Matrix dimension for 2D rotations. In 3D, the channel dimension of angle is used. - - Returns: - Matrix containing `matrix_dim` in primal and dual form as well as all non-channel dimensions of `x`. - """ - if x is None and not none_to_unit: - return None - elif x is None: - return to_float(arange(matrix_dim) == arange(matrix_dim.as_dual())) - if isinstance(x, Tensor) and x.dtype == object: # possibly None in matrices - return math.map(rotation_matrix, x, dims=object, matrix_dim=matrix_dim, none_to_unit=none_to_unit) - if isinstance(x, Tensor) and '~vector' in x.shape and 'vector' in x.shape.channel and x.shape.get_size('~vector') == x.shape.get_size('vector'): - return x # already a rotation matrix - elif 'angle' in shape(x) and shape(x).get_size('angle') == 3: # 3D Euler angles - assert channel(x).rank == 1 and channel(x).size == 3, f"x for 3D rotations needs to be a 3-vector but got {x}" - s1, s2, s3 = math.sin(x).angle # x, y, z - c1, c2, c3 = math.cos(x).angle - matrix_dim = matrix_dim.with_size(shape(x).get_item_names('angle')) - return wrap([[c3 * c2, c3 * s2 * s1 - s3 * c1, c3 * s2 * c1 + s3 * s1], - [s3 * c2, s3 * s2 * s1 + c3 * c1, s3 * s2 * c1 - c3 * s1], - [-s2, c2 * s1, c2 * c1]], matrix_dim, matrix_dim.as_dual()) # Rz * Ry * Rx (1. rotate about X by first angle) - elif 'vector' in shape(x) and shape(x).get_size('vector') == 3: # 3D axis + x - angle = length(x) - s, c = math.sin(angle), math.cos(angle) - t = 1 - c - k1, k2, k3 = normalize(x, epsilon=1e-12).vector - matrix_dim = matrix_dim.with_size(shape(x).get_item_names('vector')) - return wrap([[c + k1**2 * t, k1 * k2 * t - k3 * s, k1 * k3 * t + k2 * s], - [k2 * k1 * t + k3 * s, c + k2**2 * t, k2 * k3 * t - k1 * s], - [k3 * k1 * t - k2 * s, k3 * k2 * t + k1 * s, c + k3**2 * t]], matrix_dim, matrix_dim.as_dual()) - else: # 2D rotation - sin = wrap(math.sin(x)) - cos = wrap(math.cos(x)) - return wrap([[cos, -sin], [sin, cos]], matrix_dim, matrix_dim.as_dual()) - - -def rotation_angles(rot: Tensor): - """ - Compute the scalar x in 2D or the Euler angles in 3D from a given rotation matrix. - This function returns one valid solution but often, there are multiple solutions. - - Args: - rot: Rotation matrix as created by `phi.math.rotation_matrix()`. - Must have exactly one channel and one dual dimension with equally-ordered elements. - - Returns: - Scalar x in 2D, Euler angles - """ - assert channel(rot).rank == 1 and dual(rot).rank == 1, f"Rotation matrix must have one channel and one dual dimension but got {rot.shape}" - if channel(rot).size == 2: - cos = rot[{channel: 0, dual: 0}] - sin = rot[{channel: 1, dual: 0}] - return math.arctan(sin, divide_by=cos) - elif channel(rot).size == 3: - a2 = -math.arcsin(rot[{channel: 2, dual: 0}]) # ToDo handle [2, 0] == 1 (i.e. cos_theta == 0) - cos2 = math.cos(a2) - a1 = math.arctan(rot[{channel: 2, dual: 1}] / cos2, divide_by=rot[{channel: 2, dual: 2}] / cos2) - a3 = math.arctan(rot[{channel: 1, dual: 0}] / cos2, divide_by=rot[{channel: 0, dual: 0}] / cos2) - regular_sol = stack([a1, a2, a3], channel(angle=channel(rot).item_names[0])) - # --- pole case cos(theta) == 1 --- - a3_pole = 0 # unconstrained - bottom_pole = rot[{channel: 2, dual: 0}] < 0 - a2_pole = math.where(bottom_pole, 1.57079632679, -1.57079632679) - a1_pole = math.where(bottom_pole, math.arctan(rot[{channel: 0, dual: 1}], divide_by=rot[{channel: 0, dual: 2}]), math.arctan(-rot[{channel: 0, dual: 1}], divide_by=-rot[{channel: 0, dual: 2}])) - pole_sol = stack([a1_pole, a2_pole, a3_pole], channel(regular_sol)) - return math.where(abs(rot[{channel: 2, dual: 0}]) >= 1, pole_sol, regular_sol) - else: - raise ValueError(f"") - - -def rotation_matrix_from_directions(source_dir: Tensor, target_dir: Tensor, vec_dim: str = 'vector', epsilon=None) -> Tensor: - """ - Computes a rotation matrix A, such that `target_dir = A @ source_dir` - - Args: - source_dir: Two or three-dimensional vector. `Tensor` with channel dim called 'vector'. - target_dir: Two or three-dimensional vector. `Tensor` with channel dim called 'vector'. - - Returns: - Rotation matrix as `Tensor` with 'vector' dim and its dual counterpart. - """ - if source_dir.vector.size == 3: - axis, angle = axis_angle_from_directions(source_dir, target_dir, vec_dim, epsilon=epsilon) - return rotation_matrix_from_axis_and_angle(axis, angle, is_axis_normalized=False, epsilon=epsilon) - raise NotImplementedError - - -def axis_angle_from_directions(source_dir: Tensor, target_dir: Tensor, vec_dim: str = 'vector', epsilon=None) -> tuple[Tensor, Tensor]: - if source_dir.vector.size == 3: - source_dir = normalize(source_dir, vec_dim, epsilon=epsilon) - target_dir = normalize(target_dir, vec_dim, epsilon=epsilon) - axis = cross(source_dir, target_dir) - lim = 1-epsilon if epsilon is not None else 1 - angle = math.arccos(math.clip(source_dir.vector @ target_dir.vector, -lim, lim)) - return axis, angle - raise NotImplementedError - - -def rotation_matrix_from_axis_and_angle(axis: Tensor, angle: float | Tensor, vec_dim='vector', is_axis_normalized=False, epsilon=1e-5) -> Tensor: - """ - Computes a rotation matrix that rotates by `angle` around `axis`. - - Args: - axis: 3D vector. `Tensor` with channel dim called 'vector'. - angle: Rotation angle. - is_axis_normalized: Whether `axis` has length 1. - epsilon: Minimum axis length. For shorter axes, the unit matrix is returned. - - Returns: - Rotation matrix as `Tensor` with 'vector' dim and its dual counterpart. - """ - if axis.vector.size == 3: # Rodrigues' rotation formula - axis = normalize(axis, vec_dim, epsilon=epsilon, allow_zero=False) if not is_axis_normalized else axis - kx, ky, kz = axis.vector - s = math.sin(angle) - c = 1 - math.cos(angle) - return wrap([ - (1 - c*(ky*ky+kz*kz), -kz*s + c*(kx*ky), ky*s + c*(kx*kz)), - ( kz*s + c*(kx*ky), 1 - c*(kx*kx+kz*kz), -kx*s + c*(ky * kz)), - ( -ky*s + c*(kx*kz), kx*s + c*(ky * kz), 1 - c*(kx*kx+ky*ky)), - ], axis.shape['vector'], axis.shape['vector'].as_dual()) - raise NotImplementedError - - -def matching_rotations(*matrices: Optional[Tensor]) -> Sequence[Tensor]: - """ - Replaces `None` rotations with unit matrices if any of `matrices` is not `None`. - """ - if all(m is None for m in matrices) or all(m is not None for m in matrices): - return matrices - some = [m for m in matrices if m is not None][0] - unit_matrix = arange(some.shape.primal) == arange(some.shape.dual) - return [unit_matrix if m is None else m for m in matrices] + if isinstance(pivot, Tensor): + return pivot + rotate_vector(obj - pivot, rot, invert=invert) + else: + return rotate_vector(obj, rot, invert=invert) diff --git a/phi/vis/_dash/_plotly_plots.py b/phi/vis/_dash/_plotly_plots.py index d25241887..5ee154c33 100644 --- a/phi/vis/_dash/_plotly_plots.py +++ b/phi/vis/_dash/_plotly_plots.py @@ -9,6 +9,8 @@ import numpy import numpy as np import plotly.graph_objs + +from phi.geom._spline_solid import SplineSolid from phiml.math._sparse import CompactSparseTensor from scipy.sparse import csr_matrix, coo_matrix @@ -16,12 +18,11 @@ import plotly.graph_objects as go from plotly.subplots import make_subplots from plotly.tools import DEFAULT_PLOTLY_COLORS -import plotly.io as pio from phiml.math import reshaped_numpy, dual, instance, non_dual, merge_shapes -from phi import math, field, geom +from phi import math, geom from phi.field import Field -from phi.geom import Sphere, BaseBox, Point, Box, SDF, SDFGrid, Cylinder +from phi.geom import Sphere, BaseBox, Point, Box, SDF, SDFGrid, Cylinder, Mesh from phi.geom._geom_ops import GeometryStack from phi.math import Tensor, spatial, channel, non_channel from phi.vis._dash.colormaps import COLORMAPS @@ -666,6 +667,25 @@ def plot_single_material(data: Field, color, alpha: float): math.map(plot_single_material, data, color, alpha, dims=channel(data.geometry) - 'vector', unwrap_scalars=False) +class SplineSolid3D(Recipe): + + def can_plot(self, data: Field, space: Box) -> bool: + return isinstance(data.geometry, SplineSolid) and data.spatial_rank == 3 + + def plot(self, + data: Field, + figure, + subplot, + space: Box, + min_val: float, + max_val: float, + show_color_bar: bool, + color: Tensor, + alpha: Tensor, + err: Tensor): + mesh: Mesh = data.geometry._surface_mesh + SurfaceMesh3D().plot(Field(mesh, math.NAN), figure, subplot, space, min_val, max_val, show_color_bar, color, alpha, err) + def _get_range(bounds: Box, index: int): lower = float(bounds.lower.vector[index]) @@ -948,5 +968,6 @@ def _build_subplot_title_annotations(subplot_titles, list_of_domains, title_edge Graph3D(), VectorCloud3D(), Object3D(), + SplineSolid3D(), Scatter3D(), ]) \ No newline at end of file