Program Listing for File so3.hpp

            

            
/*
* SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef HOLOSCAN_POSE_TREE_MATH_SO3_HPP
#define HOLOSCAN_POSE_TREE_MATH_SO3_HPP

#include <cmath>

#include "holoscan/pose_tree/math/so2.hpp"
#include "holoscan/pose_tree/math/types.hpp"

namespace holoscan {

template <typename K>
class SO3 {
 public:
  using Scalar = K;
  static constexpr int kDimension = 3;

  SO3() {}

  static SO3 identity() { return SO3(Quaternion<K>::Identity()); }

  static SO3 from_scaled_axis(const Vector3<K>& axis_angle) {
    const K norm = axis_angle.norm();
    if (is_almost_zero(norm)) {
      return SO3::identity();
    } else {
      return SO3(Quaternion<K>(Eigen::AngleAxis<K>(norm, axis_angle / norm)));
    }
  }

  static SO3 from_axis_angle(const Vector3<K>& axis, K angle) {
    return SO3(Quaternion<K>(Eigen::AngleAxis<K>(angle, axis.normalized())));
  }

  static SO3 from_angle_axis(K angle, const Vector3<K>& axis) {
    return SO3(Quaternion<K>(Eigen::AngleAxis<K>(angle, axis.normalized())));
  }

  static SO3 from_quaternion(const Quaternion<K>& quaternion) {
    return SO3(quaternion.normalized());
  }

  static SO3 from_normalized_quaternion(const Quaternion<K>& quaternion) { return SO3(quaternion); }

  static SO3 from_so2_xy(const SO2<K>& rotation) {
    return from_axis_angle({K(0), K(0), K(1)}, rotation.angle());
  }

  static SO3 from_matrix(const Matrix<K, 3, 3>& matrix) {
    return SO3(Quaternion<K>(matrix).normalized());
  }

  Vector3<K> axis() const { return quaternion_.coeffs().head(3).normalized(); }

  K angle() const {
    return K(2) * std::atan2(quaternion_.coeffs().head(3).norm(), quaternion_.coeffs().w());
  }

  const Quaternion<K>& quaternion() const { return quaternion_; }

  Matrix3<K> matrix() const { return quaternion_.toRotationMatrix(); }

  Vector3<K> euler_angles_rpy() const {
    Vector3d euler_angles = quaternion_.toRotationMatrix().eulerAngles(0, 1, 2);
    // Make sure the roll is in range [-Pi/2, Pi/2]
    if (std::abs(euler_angles[0]) > M_PI * K(0.5)) {
      euler_angles[0] -= M_PI;
      euler_angles[1] = M_PI - euler_angles[1];
      euler_angles[2] += M_PI;
      // TODO(bbutin): make sure we are in the range [-Pi, Pi]
    }
    return euler_angles;
  }

  SO3 inverse() const {
    // conjugate == inverse iff quaternion_.norm() == 1
    return SO3(quaternion_.conjugate());
  }

  template <typename S, typename std::enable_if_t<!std::is_same<S, K>::value, int> = 0>
  SO3<S> cast() const {
    // We need to re-normalize in the new type
    return SO3<S>::from_quaternion(quaternion().template cast<S>());
  }

  template <typename S, typename std::enable_if_t<std::is_same<S, K>::value, int> = 0>
  const SO3& cast() const {
    // Nothing to do as the type does not change
    return *this;
  }

  SO2<K> to_so2_xy() const {
    // 2D rotation matrix:
    // cos(a) -sin(a)
    // sin(a) cos(a)
    // Quaternion to 3D rotation matrix:
    // 1 - 2*(qy^2 + qz^2) 2*(qx*qy - qz*qw) ...
    // 2*(qx*qy + qz*qw) 1 - 2*(qx^2 + qz^2) ...
    // It follows (modulo re-normalization):
    // cos(a) = 1 - (qx^2 + qy^2 + 2*qz^2)
    // sin(a) = 2*qz*qw
    // These formulas correspond to the half-angle formulas for sin/cos.
    const K qx = quaternion_.x();
    const K qy = quaternion_.y();
    const K qz = quaternion_.z();
    const K qw = quaternion_.w();
    const K cos_a = K(1) - (qx * qx + qy * qy + K(2) * qz * qz);
    const K sin_a = K(2) * qz * qw;
    return SO2<K>::from_direction(cos_a, sin_a);
  }

  friend SO3 operator*(const SO3& lhs, const SO3& rhs) {
    return from_quaternion(lhs.quaternion_ * rhs.quaternion_);
  }

  friend Vector3<K> operator*(const SO3& rot, const Vector3<K>& vec) {
    // TODO: faster implementation
    return (rot.quaternion_ * Quaternion<K>(K(0), vec.x(), vec.y(), vec.z()) *
            rot.quaternion_.conjugate())
        .coeffs()
        .head(3);
  }

  static SO3 from_euler_angles_rpy(K roll_angle, K pitch_angle, K yaw_angle) {
    SO3 roll = SO3::from_angle_axis(roll_angle, {K(1), K(0), K(0)});
    SO3 pitch = SO3::from_angle_axis(pitch_angle, {K(0), K(1), K(0)});
    SO3 yaw = SO3::from_angle_axis(yaw_angle, {K(0), K(0), K(1)});
    return roll * pitch * yaw;
  }

  static SO3 from_euler_angles_rpy(const Vector3d& roll_pitch_yaw) {
    return from_euler_angles_rpy(roll_pitch_yaw[0], roll_pitch_yaw[1], roll_pitch_yaw[2]);
  }

  Matrix<K, 3, 4> vector_rotation_jacobian(const Vector3<K>& n) const {
    // Ref: https://www.weizmann.ac.il/sci-tea/benari/sites/sci-tea.benari/files/uploads/
    // softwareAndLearningMaterials/quaternion-tutorial-2-0-1.pdf
    // Rotation Matrix in Quaternion (R):
    // [w^2 + x^2 - y^2 - z^2 2xy - 2wz 2wy + 2xz]
    // [2wz + 2xy w^2 - x^2 + y^2 - z^2 2yz - 2wx]
    // [2xz - 2wy 2wx + 2yz w^2 - x^2 - y^2 + z^2]
    const K qx = quaternion_.x();
    const K qy = quaternion_.y();
    const K qz = quaternion_.z();
    const K qw = quaternion_.w();
    Matrix<K, 3, 4> result;
    result << K(2) * (qw * n[0] - qz * n[1] + qy * n[2]),
        K(2) * (qx * n[0] + qy * n[1] + qz * n[2]), K(-2) * (qy * n[0] - qx * n[1] - qw * n[2]),
        K(-2) * (qz * n[0] + qw * n[1] - qx * n[2]),

        K(2) * (qz * n[0] + qw * n[1] - qx * n[2]), K(2) * (qy * n[0] - qx * n[1] - qw * n[2]),
        K(2) * (qx * n[0] + qy * n[1] + qz * n[2]), K(2) * (qw * n[0] - qz * n[1] + qy * n[2]),

        K(-2) * (qy * n[0] - qx * n[1] - qw * n[2]), K(2) * (qz * n[0] + qw * n[1] - qx * n[2]),
        K(-2) * (qw * n[0] - qz * n[1] + qy * n[2]), K(2) * (qx * n[0] + qy * n[1] + qz * n[2]);
    return result;
  }

 private:
  explicit SO3(const Quaternion<K>& quaternion) : quaternion_(quaternion) {}

  Quaternion<K> quaternion_;
};

using SO3d = SO3<double>;
using SO3f = SO3<float>;

}  // namespace holoscan

#endif/* HOLOSCAN_POSE_TREE_MATH_SO3_HPP */