def conjugate(dual_quaternion, name=None): """Computes the conjugate of a dual quaternion. Note: In the following, A1 to An are optional batch dimensions. Args: dual_quaternion: A tensor of shape `[A1, ..., An, 8]`, where the last dimension represents a normalized dual quaternion. name: A name for this op that defaults to "dual_quaternion_conjugate". Returns: A tensor of shape `[A1, ..., An, 8]`, where the last dimension represents a normalized dual quaternion. Raises: ValueError: If the shape of `dual_quaternion` is not supported. """ with tf.compat.v1.name_scope(name, "dual_quaternion_conjugate", [dual_quaternion]): dual_quaternion = tf.convert_to_tensor(value=dual_quaternion) shape.check_static(tensor=dual_quaternion, tensor_name="dual_quaternion", has_dim_equals=(-1, 8)) quaternion_real, quaternion_dual = tf.split(dual_quaternion, (4, 4), axis=-1) quaternion_real = asserts.assert_normalized(quaternion_real) return tf.concat((quaternion.conjugate(quaternion_real), quaternion.conjugate(quaternion_dual)), axis=-1)
def from_quaternion(quaternion, name=None): """Converts a quaternion to an axis-angle representation. Note: In the following, A1 to An are optional batch dimensions. Args: quaternion: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. name: A name for this op that defaults to "axis_angle_from_quaternion". Returns: Tuple of two tensors of shape `[A1, ..., An, 3]` and `[A1, ..., An, 1]`, where the first tensor represents the axis, and the second represents the angle. The resulting axis is a normalized vector. Raises: ValueError: If the shape of `quaternion` is not supported. """ with tf.compat.v1.name_scope(name, "axis_angle_from_quaternion", [quaternion]): quaternion = tf.convert_to_tensor(value=quaternion) shape.check_static(tensor=quaternion, tensor_name="quaternion", has_dim_equals=(-1, 4)) quaternion = asserts.assert_normalized(quaternion) # This prevents zero norm xyz and zero w, and is differentiable. quaternion += asserts.select_eps_for_addition(quaternion.dtype) xyz, w = tf.split(quaternion, (3, 1), axis=-1) norm = tf.norm(tensor=xyz, axis=-1, keepdims=True) angle = 2.0 * tf.atan2(norm, tf.abs(w)) axis = safe_ops.safe_unsigned_div(safe_ops.nonzero_sign(w) * xyz, norm) return axis, angle
def relative_angle(quaternion1, quaternion2, name=None): r"""Computes the unsigned relative rotation angle between 2 unit quaternions. Given two normalized quanternions $$\mathbf{q}_1$$ and $$\mathbf{q}_2$$, the relative angle is computed as $$\theta = 2\arccos(\mathbf{q}_1^T\mathbf{q}_2)$$. Note: In the following, A1 to An are optional batch dimensions. Args: quaternion1: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. quaternion2: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. name: A name for this op that defaults to "quaternion_relative_angle". Returns: A tensor of shape `[A1, ..., An, 1]` where the last dimension represents rotation angles in the range [0.0, pi]. Raises: ValueError: If the shape of `quaternion1` or `quaternion2` is not supported. """ with (tf.compat.v1.name_scope(name, "quaternion_relative_angle", [quaternion1, quaternion2])): quaternion1 = tf.convert_to_tensor(value=quaternion1) quaternion2 = tf.convert_to_tensor(value=quaternion2) shape.check_static(tensor=quaternion1, tensor_name="quaternion1", has_dim_equals=(-1, 4)) shape.check_static(tensor=quaternion2, tensor_name="quaternion2", has_dim_equals=(-1, 4)) quaternion1 = asserts.assert_normalized(quaternion1) quaternion2 = asserts.assert_normalized(quaternion2) dot_product = vector.dot(quaternion1, quaternion2, keepdims=False) # Ensure dot product is in range [-1. 1]. eps_dot_prod = 4.0 * asserts.select_eps_for_addition(dot_product.dtype) dot_product = safe_ops.safe_shrink(dot_product, -1.0, 1.0, False, eps=eps_dot_prod) return 2.0 * tf.acos(tf.abs(dot_product))
def from_axis_angle(axis, angle, name="rotation_matrix_3d_from_axis_angle"): """Convert an axis-angle representation to a rotation matrix. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents a normalized axis. name: A name for this op that defaults to "rotation_matrix_3d_from_axis_angle". Returns: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represents a 3d rotation matrix. Raises: ValueError: If the shape of `axis` or `angle` is not supported. """ with tf.name_scope(name): axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static( tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions( tensors=(axis, angle), tensor_names=("axis", "angle"), last_axes=-2, broadcast_compatible=True) axis = asserts.assert_normalized(axis) sin_axis = tf.sin(angle) * axis cos_angle = tf.cos(angle) cos1_axis = (1.0 - cos_angle) * axis _, axis_y, axis_z = tf.unstack(axis, axis=-1) cos1_axis_x, cos1_axis_y, _ = tf.unstack(cos1_axis, axis=-1) sin_axis_x, sin_axis_y, sin_axis_z = tf.unstack(sin_axis, axis=-1) tmp = cos1_axis_x * axis_y m01 = tmp - sin_axis_z m10 = tmp + sin_axis_z tmp = cos1_axis_x * axis_z m02 = tmp + sin_axis_y m20 = tmp - sin_axis_y tmp = cos1_axis_y * axis_z m12 = tmp - sin_axis_x m21 = tmp + sin_axis_x diag = cos1_axis * axis + cos_angle diag_x, diag_y, diag_z = tf.unstack(diag, axis=-1) matrix = tf.stack((diag_x, m01, m02, m10, diag_y, m12, m20, m21, diag_z), axis=-1) # pyformat: disable output_shape = tf.concat((tf.shape(input=axis)[:-1], (3, 3)), axis=-1) return tf.reshape(matrix, shape=output_shape)
def test_assert_normalized_exception_raised(self, dtype): """Checks that assert_normalized raises exceptions for invalid input.""" vector = _pick_random_vector() + 10.0 vector = tf.convert_to_tensor(value=vector, dtype=dtype) vector = tf.abs(vector) with self.assertRaises(tf.errors.InvalidArgumentError): self.evaluate(asserts.assert_normalized(vector))
def rotate(point, axis, angle, name=None): r"""Rotates a 3d point using an axis-angle by applying the Rodrigues' formula. Rotates a vector $$\mathbf{v} \in {\mathbb{R}^3}$$ into a vector $$\mathbf{v}' \in {\mathbb{R}^3}$$ using the Rodrigues' rotation formula: $$\mathbf{v}'=\mathbf{v}\cos(\theta)+(\mathbf{a}\times\mathbf{v})\sin(\theta) +\mathbf{a}(\mathbf{a}\cdot\mathbf{v})(1-\cos(\theta)).$$ Note: In the following, A1 to An are optional batch dimensions. Args: point: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point to rotate. axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents an angle. name: A name for this op that defaults to "axis_angle_rotate". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. Raises: ValueError: If `point`, `axis`, or `angle` are of different shape or if their respective shape is not supported. """ with tf.compat.v1.name_scope(name, "axis_angle_rotate", [point, axis, angle]): point = tf.convert_to_tensor(value=point) axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=point, tensor_name="point", has_dim_equals=(-1, 3)) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static(tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions(tensors=(point, axis, angle), tensor_names=("point", "axis", "angle"), last_axes=-2, broadcast_compatible=True) axis = asserts.assert_normalized(axis) cos_angle = tf.cos(angle) axis_dot_point = vector.dot(axis, point) return point * cos_angle + vector.cross(axis, point) * tf.sin( angle) + axis * axis_dot_point * (1.0 - cos_angle)
def from_quaternion( quaternion: type_alias.TensorLike, name: str = "rotation_matrix_3d_from_quaternion") -> tf.Tensor: """Convert a quaternion to a rotation matrix. Note: In the following, A1 to An are optional batch dimensions. Args: quaternion: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. name: A name for this op that defaults to "rotation_matrix_3d_from_quaternion". Returns: A tensor of shape `[A1, ..., An, 3, 3]`, where the last two dimensions represent a 3d rotation matrix. Raises: ValueError: If the shape of `quaternion` is not supported. """ with tf.name_scope(name): quaternion = tf.convert_to_tensor(value=quaternion) shape.check_static(tensor=quaternion, tensor_name="quaternion", has_dim_equals=(-1, 4)) quaternion = asserts.assert_normalized(quaternion) x, y, z, w = tf.unstack(quaternion, axis=-1) tx = 2.0 * x ty = 2.0 * y tz = 2.0 * z twx = tx * w twy = ty * w twz = tz * w txx = tx * x txy = ty * x txz = tz * x tyy = ty * y tyz = tz * y tzz = tz * z matrix = tf.stack( (1.0 - (tyy + tzz), txy - twz, txz + twy, txy + twz, 1.0 - (txx + tzz), tyz - twx, txz - twy, tyz + twx, 1.0 - (txx + tyy)), axis=-1) # pyformat: disable output_shape = tf.concat((tf.shape(input=quaternion)[:-1], (3, 3)), axis=-1) return tf.reshape(matrix, shape=output_shape)
def distance_to_ray(point: type_alias.TensorLike, origin: type_alias.TensorLike, direction: type_alias.TensorLike, keepdims: bool = True, name: str = "point_distance_to_ray" ) -> tf.Tensor: """Computes the distance from a M-d point to a M-d ray. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: point: A tensor of shape `[A1, ..., An, M]`. origin: A tensor of shape `[A1, ..., An, M]`. direction: A tensor of shape `[A1, ..., An, M]`. The last dimension must be normalized. keepdims: A `bool`, whether to keep the last dimension with length 1 or to remove it. name: A name for this op. Defaults to "point_distance_to_ray". Returns: A tensor of shape `[A1, ..., An, 1]` containing the distance from each point to the corresponding ray. Raises: ValueError: If the shape of `point`, `origin`, or 'direction' is not supported. """ with tf.name_scope(name): point = tf.convert_to_tensor(value=point) origin = tf.convert_to_tensor(value=origin) direction = tf.convert_to_tensor(value=direction) shape.compare_dimensions((point, origin, direction), -1, ("point", "origin", "direction")) shape.compare_batch_dimensions( tensors=(point, origin, direction), last_axes=-2, broadcast_compatible=True) direction = asserts.assert_normalized(direction) vec = point - origin dot = vector.dot(vec, direction) vec -= dot * direction return tf.norm(tensor=vec, axis=-1, keepdims=keepdims)
def rotate(point: type_alias.TensorLike, quaternion: type_alias.TensorLike, name: str = "quaternion_rotate") -> tf.Tensor: """Rotates a point using a quaternion. Note: In the following, A1 to An are optional batch dimensions. Args: point: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. quaternion: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. name: A name for this op that defaults to "quaternion_rotate". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a 3d point. Raises: ValueError: If the shape of `point` or `quaternion` is not supported. """ with tf.name_scope(name): point = tf.convert_to_tensor(value=point) quaternion = tf.convert_to_tensor(value=quaternion) shape.check_static(tensor=point, tensor_name="point", has_dim_equals=(-1, 3)) shape.check_static(tensor=quaternion, tensor_name="quaternion", has_dim_equals=(-1, 4)) shape.compare_batch_dimensions(tensors=(point, quaternion), last_axes=-2, broadcast_compatible=True) quaternion = asserts.assert_normalized(quaternion) padding = [[0, 0] for _ in range(point.shape.ndims)] padding[-1][-1] = 1 point = tf.pad(tensor=point, paddings=padding, mode="CONSTANT") point = multiply(quaternion, point) point = multiply(point, conjugate(quaternion)) xyz, _ = tf.split(point, (3, 1), axis=-1) return xyz
def from_axis_angle(axis: type_alias.TensorLike, angle: type_alias.TensorLike, name: str = "quaternion_from_axis_angle") -> tf.Tensor: """Converts an axis-angle representation to a quaternion. Note: In the following, A1 to An are optional batch dimensions. Args: axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents an angle. name: A name for this op that defaults to "quaternion_from_axis_angle". Returns: A tensor of shape `[A1, ..., An, 4]`, where the last dimension represents a normalized quaternion. Raises: ValueError: If the shape of `axis` or `angle` is not supported. """ with tf.name_scope(name): axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static(tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions(tensors=(axis, angle), last_axes=-2, broadcast_compatible=True) axis = asserts.assert_normalized(axis) half_angle = 0.5 * angle w = tf.cos(half_angle) xyz = tf.sin(half_angle) * axis return tf.concat((xyz, w), axis=-1)
def reflect(vector: TensorLike, normal: TensorLike, axis: int = -1, name: str = "vector_reflect") -> TensorLike: r"""Computes the reflection direction for an incident vector. For an incident vector \\(\mathbf{v}\\) and normal $$\mathbf{n}$$ this function computes the reflected vector as \\(\mathbf{r} = \mathbf{v} - 2(\mathbf{n}^T\mathbf{v})\mathbf{n}\\). Note: In the following, A1 to An are optional batch dimensions, which should be broadcast compatible. Args: vector: A tensor of shape `[A1, ..., Ai, ..., An]`, where the dimension i = axis represents a vector. normal: A tensor of shape `[A1, ..., Ai, ..., An]`, where the dimension i = axis represents a normal around which the vector needs to be reflected. The normal vector needs to be normalized. axis: The dimension along which to compute the reflection. name: A name for this op which defaults to "vector_reflect". Returns: A tensor of shape `[A1, ..., Ai, ..., An]`, where the dimension i = axis represents a reflected vector. """ with tf.name_scope(name): vector = tf.convert_to_tensor(value=vector) normal = tf.convert_to_tensor(value=normal) shape.compare_dimensions(tensors=(vector, normal), axes=axis, tensor_names=("vector", "normal")) shape.compare_batch_dimensions(tensors=(vector, normal), last_axes=-1, broadcast_compatible=True) normal = asserts.assert_normalized(normal, axis=axis) dot_product = dot(vector, normal, axis=axis) return vector - 2.0 * dot_product * normal
def project_to_ray(point: type_alias.TensorLike, origin: type_alias.TensorLike, direction: type_alias.TensorLike, name: str = "point_project_to_ray" ) -> tf.Tensor: """Computes the projection of a M-d point on a M-d ray. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Args: point: A tensor of shape `[A1, ..., An, M]`. origin: A tensor of shape `[A1, ..., An, M]`. direction: A tensor of shape `[A1, ..., An, M]`. The last dimension must be normalized. name: A name for this op. Defaults to "point_project_to_ray". Returns: A tensor of shape `[A1, ..., An, M]` containing the projected point. Raises: ValueError: If the shape of `point`, `origin`, or 'direction' is not supported. """ with tf.name_scope(name): point = tf.convert_to_tensor(value=point) origin = tf.convert_to_tensor(value=origin) direction = tf.convert_to_tensor(value=direction) shape.compare_dimensions((point, origin, direction), -1, ("point", "origin", "direction")) shape.compare_batch_dimensions( tensors=(point, origin, direction), last_axes=-2, broadcast_compatible=True) direction = asserts.assert_normalized(direction) vec = point - origin dot = vector.dot(vec, direction) return origin + dot * direction
def inverse(axis, angle, name=None): """Computes the axis-angle that is the inverse of the input axis-angle. Note: In the following, A1 to An are optional batch dimensions. Args: axis: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized axis. angle: A tensor of shape `[A1, ..., An, 1]` where the last dimension represents an angle. name: A name for this op that defaults to "axis_angle_inverse". Returns: A tuple of two tensors, respectively of shape `[A1, ..., An, 3]` and `[A1, ..., An, 1]`, where the first tensor represents the axis, and the second represents the angle. The resulting axis is a normalized vector. Raises: ValueError: If the shape of `axis` or `angle` is not supported. """ with tf.compat.v1.name_scope(name, "axis_angle_inverse", [axis, angle]): axis = tf.convert_to_tensor(value=axis) angle = tf.convert_to_tensor(value=angle) shape.check_static(tensor=axis, tensor_name="axis", has_dim_equals=(-1, 3)) shape.check_static(tensor=angle, tensor_name="angle", has_dim_equals=(-1, 1)) shape.compare_batch_dimensions(tensors=(axis, angle), tensor_names=("axis", "angle"), last_axes=-2, broadcast_compatible=True) axis = asserts.assert_normalized(axis) return axis, -angle
def interpolate_attributes(attribute, barycentric, name=None): """Interpolates attributes using barycentric weights. Note: In the following, A1 to An are optional batch dimensions. Args: attribute: A tensor of shape `[A1, ..., An, 3, B]`, where the last dimension stores a per-vertex `B`-dimensional attribute. barycentric: A tensor of shape `[A1, ..., An, 3]`, where the last dimension contains barycentric coordinates. name: A name for this op. Defaults to 'interpolate_attributes'. Returns: A tensor of shape `[A1, ..., An, B]`, containing interpolated attributes. """ with tf.compat.v1.name_scope(name, "interpolate_attributes", (attribute, barycentric)): attribute = tf.convert_to_tensor(value=attribute) barycentric = tf.convert_to_tensor(value=barycentric) shape.check_static(tensor=attribute, tensor_name="attribute", has_dim_equals=(-2, 3)) shape.check_static(tensor=barycentric, tensor_name="barycentric", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions(tensors=(attribute, barycentric), last_axes=(-2, -1), tensor_names=("attribute", "barycentric"), broadcast_compatible=True) barycentric = asserts.assert_normalized(barycentric, order=1) return tf.reduce_sum( input_tensor=tf.expand_dims(barycentric, axis=-1) * attribute, axis=-2)
def brdf(direction_incoming_light: type_alias.TensorLike, direction_outgoing_light: type_alias.TensorLike, surface_normal: type_alias.TensorLike, albedo: type_alias.TensorLike, name: str = "lambertian_brdf") -> tf.Tensor: """Evaluates the brdf of a Lambertian surface. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Note: The gradient of this function is not smooth when the dot product of the normal with any light is 0.0. Args: direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized incoming light vector. direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized outgoing light vector. surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized surface normal. albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents albedo with values in [0,1]. name: A name for this op. Defaults to "lambertian_brdf". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the amount of reflected light in any outgoing direction. Raises: ValueError: if the shape of `direction_incoming_light`, `direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not supported. InvalidArgumentError: if at least one element of `albedo` is outside of [0,1]. """ with tf.name_scope(name): direction_incoming_light = tf.convert_to_tensor( value=direction_incoming_light) direction_outgoing_light = tf.convert_to_tensor( value=direction_outgoing_light) surface_normal = tf.convert_to_tensor(value=surface_normal) albedo = tf.convert_to_tensor(value=albedo) shape.check_static(tensor=direction_incoming_light, tensor_name="direction_incoming_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=direction_outgoing_light, tensor_name="direction_outgoing_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=surface_normal, tensor_name="surface_normal", has_dim_equals=(-1, 3)) shape.check_static(tensor=albedo, tensor_name="albedo", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(direction_incoming_light, direction_outgoing_light, surface_normal, albedo), tensor_names=("direction_incoming_light", "direction_outgoing_light", "surface_normal", "albedo"), last_axes=-2, broadcast_compatible=True) direction_incoming_light = asserts.assert_normalized( direction_incoming_light) direction_outgoing_light = asserts.assert_normalized( direction_outgoing_light) surface_normal = asserts.assert_normalized(surface_normal) albedo = asserts.assert_all_in_range(albedo, 0.0, 1.0, open_bounds=False) # Checks whether the incoming or outgoing light point behind the surface. dot_incoming_light_surface_normal = vector.dot( -direction_incoming_light, surface_normal) dot_outgoing_light_surface_normal = vector.dot( direction_outgoing_light, surface_normal) min_dot = tf.minimum(dot_incoming_light_surface_normal, dot_outgoing_light_surface_normal) common_shape = shape.get_broadcasted_shape(min_dot.shape, albedo.shape) d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim ) common_shape = [d_val(dim) for dim in common_shape] condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape) albedo = tf.broadcast_to(albedo, common_shape) return tf.where(condition, albedo / math.pi, tf.zeros_like(albedo))
def quaternion_weights( quaternion1: type_alias.TensorLike, quaternion2: type_alias.TensorLike, percent: Union[type_alias.Float, type_alias.TensorLike], eps: Optional[type_alias.Float] = None, name: str = "quaternion_weights") -> Tuple[tf.Tensor, tf.Tensor]: """Calculates slerp weights for two normalized quaternions. Given a percent and two normalized quaternions, this function returns the slerp weights. It can also produce extrapolation weights when percent is outside of the [0, 1] range. It reduces to lerp when input quaternions are almost parallel or anti-parallel. Input quaternions are assumed to be normalized. The tf.graphics debug flag TFG_ADD_ASSERTS_TO_GRAPH defined in tfg_flags.py can be set to add assertions to the graph that check whether the inputs are normalized, and whether Inf or Nan values are produced. Note: In the following, A1 to An are optional batch dimensions. Args: quaternion1: A tensor of shape `[A1, ... , An, 4]` storing normalized quaternions in its last dimension. quaternion2: A tensor of shape `[A1, ... , An, 4]` storing normalized quaternions in its last dimension. percent: A `float` or a tensor with a shape broadcastable to the shape `[A1, ... , An]`. eps: A `float` used to make operations safe. When left as None, the function automatically picks the best epsilon based on the dtype and the operation. name: A name for this op. Defaults to "quaternion_weights". Raises: ValueError: If the shapes of quaternions do not match, if the last dimensions of quaternions are not 4, or if percent is neither a float, nor a tensor with last dimension 1. Returns: Two tensors of shape `[A1, ... , An, 1]` each, which are the two slerp weights for each quaternion. """ with tf.name_scope(name): quaternion1 = tf.convert_to_tensor(value=quaternion1) quaternion2 = tf.convert_to_tensor(value=quaternion2) percent = tf.convert_to_tensor(value=percent, dtype=quaternion1.dtype) if percent.shape.ndims == 0: percent = tf.expand_dims(percent, axis=0) shape.check_static( tensor=quaternion1, tensor_name="quaternion1", has_dim_equals=(-1, 4)) shape.check_static( tensor=quaternion2, tensor_name="quaternion2", has_dim_equals=(-1, 4)) shape.compare_batch_dimensions( tensors=(quaternion1, quaternion2, percent), last_axes=(-2, -2, -1), broadcast_compatible=True, tensor_names=("quaternion1", "quaternion2", "percent")) quaternion1 = asserts.assert_normalized(quaternion1) quaternion2 = asserts.assert_normalized(quaternion2) dot_product = _safe_dot(quaternion1, quaternion2, eps) # Take the shorter path theta = tf.acos(tf.abs(dot_product)) # safe_sinpx_div_sinx returns p for very small x, which means slerp reduces # to lerp automatically. scale1 = safe_ops.safe_sinpx_div_sinx(theta, 1.0 - percent, eps) scale2 = safe_ops.safe_sinpx_div_sinx(theta, percent, eps) # Flip the sign of scale1 if quaternions are in different hemispheres. # tf.sign can make scale1 zero if quaternions are orthogonal. scale1 *= safe_ops.nonzero_sign(dot_product) return scale1, scale2
def brdf(direction_incoming_light: type_alias.TensorLike, direction_outgoing_light: type_alias.TensorLike, surface_normal: type_alias.TensorLike, shininess: type_alias.TensorLike, albedo: type_alias.TensorLike, brdf_normalization: bool = True, name: str = "phong_brdf") -> tf.Tensor: """Evaluates the specular brdf of the Phong model. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. Note: The gradient of this function is not smooth when the dot product of the normal with any light is 0.0. Args: direction_incoming_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized incoming light vector. direction_outgoing_light: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized outgoing light vector. surface_normal: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents a normalized surface normal. shininess: A tensor of shape `[A1, ..., An, 1]`, where the last dimension represents a non-negative shininess coefficient. albedo: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents albedo with values in [0,1]. brdf_normalization: A `bool` indicating whether normalization should be applied to enforce the energy conservation property of BRDFs. Note that `brdf_normalization` must be set to False in order to use the original Blinn specular model. name: A name for this op. Defaults to "phong_brdf". Returns: A tensor of shape `[A1, ..., An, 3]`, where the last dimension represents the amount of light reflected in the outgoing light direction. Raises: ValueError: if the shape of `direction_incoming_light`, `direction_outgoing_light`, `surface_normal`, `shininess` or `albedo` is not supported. InvalidArgumentError: if not all of shininess values are non-negative, or if at least one element of `albedo` is outside of [0,1]. """ with tf.name_scope(name): direction_incoming_light = tf.convert_to_tensor( value=direction_incoming_light) direction_outgoing_light = tf.convert_to_tensor( value=direction_outgoing_light) surface_normal = tf.convert_to_tensor(value=surface_normal) shininess = tf.convert_to_tensor(value=shininess) albedo = tf.convert_to_tensor(value=albedo) shape.check_static(tensor=direction_incoming_light, tensor_name="direction_incoming_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=direction_outgoing_light, tensor_name="direction_outgoing_light", has_dim_equals=(-1, 3)) shape.check_static(tensor=surface_normal, tensor_name="surface_normal", has_dim_equals=(-1, 3)) shape.check_static(tensor=shininess, tensor_name="shininess", has_dim_equals=(-1, 1)) shape.check_static(tensor=albedo, tensor_name="albedo", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(direction_incoming_light, direction_outgoing_light, surface_normal, shininess, albedo), tensor_names=("direction_incoming_light", "direction_outgoing_light", "surface_normal", "shininess", "albedo"), last_axes=-2, broadcast_compatible=True) direction_incoming_light = asserts.assert_normalized( direction_incoming_light) direction_outgoing_light = asserts.assert_normalized( direction_outgoing_light) surface_normal = asserts.assert_normalized(surface_normal) albedo = asserts.assert_all_in_range(albedo, 0.0, 1.0, open_bounds=False) shininess = asserts.assert_all_above(shininess, 0.0, open_bound=False) # Checks whether the incoming or outgoing light point behind the surface. dot_incoming_light_surface_normal = vector.dot( -direction_incoming_light, surface_normal) dot_outgoing_light_surface_normal = vector.dot( direction_outgoing_light, surface_normal) min_dot = tf.minimum(dot_incoming_light_surface_normal, dot_outgoing_light_surface_normal) perfect_reflection_direction = vector.reflect(direction_incoming_light, surface_normal) perfect_reflection_direction = tf.math.l2_normalize( perfect_reflection_direction, axis=-1) cos_alpha = vector.dot(perfect_reflection_direction, direction_outgoing_light, axis=-1) cos_alpha = tf.maximum(cos_alpha, tf.zeros_like(cos_alpha)) phong_model = albedo * tf.pow(cos_alpha, shininess) if brdf_normalization: phong_model *= _brdf_normalization_factor(shininess) common_shape = shape.get_broadcasted_shape(min_dot.shape, phong_model.shape) d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim ) common_shape = [d_val(dim) for dim in common_shape] condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape) phong_model = tf.broadcast_to(phong_model, common_shape) return tf.where(condition, phong_model, tf.zeros_like(phong_model))
def intersection_ray_sphere(sphere_center, sphere_radius, ray, point_on_ray, name=None): """Finds positions and surface normals where the sphere and the ray intersect. Note: In the following, A1 to An are optional batch dimensions. Args: sphere_center: A tensor of shape `[3]` representing the 3d sphere center. sphere_radius: A tensor of shape `[1]` containing a strictly positive value defining the radius of the sphere. ray: A tensor of shape `[A1, ..., An, 3]` containing normalized 3D vectors. point_on_ray: A tensor of shape `[A1, ..., An, 3]`. name: A name for this op. The default value of None means "ray_intersection_ray_sphere". Returns: A tensor of shape `[2, A1, ..., An, 3]` containing the position of the intersections, and a tensor of shape `[2, A1, ..., An, 3]` the associated surface normals at that point. Both tensors contain NaNs when there is no intersections. The first dimension of the returned tensor provides access to the first and second intersections of the ray with the sphere. Raises: ValueError: if the shape of `sphere_center`, `sphere_radius`, `ray` or `point_on_ray` is not supported. tf.errors.InvalidArgumentError: If `ray` is not normalized. """ with tf.compat.v1.name_scope( name, "ray_intersection_ray_sphere", [sphere_center, sphere_radius, ray, point_on_ray]): sphere_center = tf.convert_to_tensor(value=sphere_center) sphere_radius = tf.convert_to_tensor(value=sphere_radius) ray = tf.convert_to_tensor(value=ray) point_on_ray = tf.convert_to_tensor(value=point_on_ray) shape.check_static( tensor=sphere_center, tensor_name="sphere_center", has_rank=1, has_dim_equals=(0, 3)) shape.check_static( tensor=sphere_radius, tensor_name="sphere_radius", has_rank=1, has_dim_equals=(0, 1)) shape.check_static(tensor=ray, tensor_name="ray", has_dim_equals=(-1, 3)) shape.check_static( tensor=point_on_ray, tensor_name="point_on_ray", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(ray, point_on_ray), last_axes=(-2, -2), broadcast_compatible=False) sphere_radius = asserts.assert_all_above( sphere_radius, 0.0, open_bound=True) ray = asserts.assert_normalized(ray) vector_sphere_center_to_point_on_ray = sphere_center - point_on_ray distance_sphere_center_to_point_on_ray = tf.norm( tensor=vector_sphere_center_to_point_on_ray, axis=-1, keepdims=True) distance_projection_sphere_center_on_ray = vector.dot( vector_sphere_center_to_point_on_ray, ray) closest_distance_sphere_center_to_ray = tf.sqrt( tf.square(distance_sphere_center_to_point_on_ray) - tf.pow(distance_projection_sphere_center_on_ray, 2)) half_secant_length = tf.sqrt( tf.square(sphere_radius) - tf.square(closest_distance_sphere_center_to_ray)) distances = tf.stack( (distance_projection_sphere_center_on_ray - half_secant_length, distance_projection_sphere_center_on_ray + half_secant_length), axis=0) intersections_points = distances * ray + point_on_ray normals = tf.math.l2_normalize( intersections_points - sphere_center, axis=-1) return intersections_points, normals
def estimate_radiance(point_light_radiance, point_light_position, surface_point_position, surface_point_normal, observation_point, brdf, name="estimate_radiance", reflected_light_fall_off=False): """Estimates the spectral radiance of a point light reflected from the surface point towards the observation point. Note: In the following, A1 to An are optional batch dimensions, which must be broadcast compatible. B1 to Bm are optional batch dimensions for the lights, which must be broadcast compatible. Note: In case the light or the observation point are located behind the surface the function will return 0. Note: The gradient of this function is not smooth when the dot product of the normal with the light-to-surface or surface-to-observation vectors is 0. Args: point_light_radiance: A tensor of shape '[B1, ..., Bm, K]', where the last axis represents the radiance of the point light at a specific wave length. point_light_position: A tensor of shape `[B1, ..., Bm, 3]`, where the last axis represents the position of the point light. surface_point_position: A tensor of shape `[A1, ..., An, 3]`, where the last axis represents the position of the surface point. surface_point_normal: A tensor of shape `[A1, ..., An, 3]`, where the last axis represents the normalized surface normal at the given surface point. observation_point: A tensor of shape `[A1, ..., An, 3]`, where the last axis represents the observation point. brdf: The BRDF of the surface as a function of: incoming_light_direction - The incoming light direction as the last axis of a tensor with shape `[A1, ..., An, 3]`. outgoing_light_direction - The outgoing light direction as the last axis of a tensor with shape `[A1, ..., An, 3]`. surface_point_normal - The surface normal as the last axis of a tensor with shape `[A1, ..., An, 3]`. Note - The BRDF should return a tensor of size '[A1, ..., An, K]' where the last axis represents the amount of reflected light in each wave length. name: A name for this op. Defaults to "estimate_radiance". reflected_light_fall_off: A boolean specifying whether or not to include the fall off of the light reflected from the surface towards the observation point in the calculation. Defaults to False. Returns: A tensor of shape `[A1, ..., An, B1, ..., Bm, K]`, where the last axis represents the amount of light received at the observation point after being reflected from the given surface point. Raises: ValueError: if the shape of `point_light_position`, `surface_point_position`, `surface_point_normal`, or `observation_point` is not supported. InvalidArgumentError: if 'surface_point_normal' is not normalized. """ with tf.name_scope(name): point_light_radiance = tf.convert_to_tensor(value=point_light_radiance) point_light_position = tf.convert_to_tensor(value=point_light_position) surface_point_position = tf.convert_to_tensor(value=surface_point_position) surface_point_normal = tf.convert_to_tensor(value=surface_point_normal) observation_point = tf.convert_to_tensor(value=observation_point) shape.check_static( tensor=point_light_position, tensor_name="point_light_position", has_dim_equals=(-1, 3)) shape.check_static( tensor=surface_point_position, tensor_name="surface_point_position", has_dim_equals=(-1, 3)) shape.check_static( tensor=surface_point_normal, tensor_name="surface_point_normal", has_dim_equals=(-1, 3)) shape.check_static( tensor=observation_point, tensor_name="observation_point", has_dim_equals=(-1, 3)) shape.compare_batch_dimensions( tensors=(surface_point_position, surface_point_normal, observation_point), tensor_names=("surface_point_position", "surface_point_normal", "observation_point"), last_axes=-2, broadcast_compatible=True) shape.compare_batch_dimensions( tensors=(point_light_radiance, point_light_position), tensor_names=("point_light_radiance", "point_light_position"), last_axes=-2, broadcast_compatible=True) surface_point_normal = asserts.assert_normalized(surface_point_normal) # Get the number of lights dimensions (B1,...,Bm). lights_num_dimensions = max( len(point_light_radiance.shape), len(point_light_position.shape)) - 1 # Reshape the other parameters so they can be broadcasted to the output of # shape [A1,...,An, B1,...,Bm, K]. surface_point_position = tf.reshape( surface_point_position, surface_point_position.shape[:-1] + (1,) * lights_num_dimensions + (3,)) surface_point_normal = tf.reshape( surface_point_normal, surface_point_normal.shape[:-1] + (1,) * lights_num_dimensions + (3,)) observation_point = tf.reshape( observation_point, observation_point.shape[:-1] + (1,) * lights_num_dimensions + (3,)) light_to_surface_point = surface_point_position - point_light_position distance_light_surface_point = tf.norm( tensor=light_to_surface_point, axis=-1, keepdims=True) incoming_light_direction = tf.math.l2_normalize( light_to_surface_point, axis=-1) surface_to_observation_point = observation_point - surface_point_position outgoing_light_direction = tf.math.l2_normalize( surface_to_observation_point, axis=-1) brdf_value = brdf(incoming_light_direction, outgoing_light_direction, surface_point_normal) incoming_light_dot_surface_normal = vector.dot(-incoming_light_direction, surface_point_normal) outgoing_light_dot_surface_normal = vector.dot(outgoing_light_direction, surface_point_normal) estimated_radiance = (point_light_radiance * \ brdf_value * incoming_light_dot_surface_normal) / \ (4. * math.pi * tf.math.square(distance_light_surface_point)) if reflected_light_fall_off: distance_surface_observation_point = tf.norm( tensor=surface_to_observation_point, axis=-1, keepdims=True) estimated_radiance = estimated_radiance / \ tf.math.square(distance_surface_observation_point) # Create a condition for checking whether the light or observation point are # behind the surface. min_dot = tf.minimum(incoming_light_dot_surface_normal, outgoing_light_dot_surface_normal) common_shape = shape.get_broadcasted_shape(min_dot.shape, estimated_radiance.shape) d_val = lambda dim: 1 if dim is None else tf.compat.dimension_value(dim) common_shape = [d_val(dim) for dim in common_shape] condition = tf.broadcast_to(tf.greater_equal(min_dot, 0.0), common_shape) return tf.where(condition, estimated_radiance, tf.zeros_like(estimated_radiance))