Program Listing for File multivariate_distribution_traits.hpp

↰ Return to documentation for file (include/beluga/random/multivariate_distribution_traits.hpp)

// Copyright 2024 Ekumen, Inc.
//
// 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 BELUGA_RANDOM_MULTIVARIATE_DISTRIBUTION_TRAITS_HPP
#define BELUGA_RANDOM_MULTIVARIATE_DISTRIBUTION_TRAITS_HPP

#include <Eigen/Core>
#include <sophus/se2.hpp>
#include <sophus/se3.hpp>
#include <sophus/so3.hpp>

namespace beluga {

template <class T, class Enable = void>
struct multivariate_distribution_traits;

template <class T>
struct multivariate_distribution_traits<T, std::enable_if_t<std::is_base_of_v<Eigen::EigenBase<T>, T>>> {
  static_assert(T::ColsAtCompileTime == 1 || T::RowsAtCompileTime == 1, "T should be a column or row vector");

  static constexpr int matrix_size = T::ColsAtCompileTime > T::RowsAtCompileTime  //
                                         ? T::ColsAtCompileTime
                                         : T::RowsAtCompileTime;

  using scalar_type = typename T::Scalar;

  using result_type = typename T::PlainMatrix;

  using vector_type = typename T::PlainMatrix;

  using covariance_type = typename Eigen::Matrix<scalar_type, matrix_size, matrix_size>;

  [[nodiscard]] static constexpr vector_type to_vector(const result_type& t) { return t; }

  [[nodiscard]] static constexpr result_type from_vector(const vector_type& v) { return v; }
};

template <class T>
struct multivariate_distribution_traits<T, std::enable_if_t<std::is_base_of_v<Sophus::SO2Base<T>, T>>> {
  using scalar_type = typename T::Scalar;

  using result_type = Sophus::SO2<scalar_type>;

  using vector_type = typename Eigen::Matrix<scalar_type, 1, 1>;

  using covariance_type = typename Eigen::Matrix<scalar_type, 1, 1>;

  [[nodiscard]] static constexpr vector_type to_vector(const result_type& t) { return vector_type{t.log()}; }

  [[nodiscard]] static constexpr result_type from_vector(const vector_type& v) {
    return Sophus::SO2<scalar_type>::exp(v.x());
  }
};

template <class T>
struct multivariate_distribution_traits<T, std::enable_if_t<std::is_base_of_v<Sophus::SE2Base<T>, T>>> {
  using scalar_type = typename T::Scalar;

  using result_type = Sophus::SE2<scalar_type>;

  using vector_type = typename Eigen::Matrix<scalar_type, 3, 1>;

  using covariance_type = typename Eigen::Matrix<scalar_type, 3, 3>;

  [[nodiscard]] static constexpr vector_type to_vector(const result_type& t) {
    vector_type v;
    v << t.translation(), t.so2().log();
    return v;
  }

  [[nodiscard]] static constexpr result_type from_vector(const vector_type& v) {
    return {Sophus::SO2<scalar_type>::exp(v.z()), v.head(2)};
  }
};

template <class T>
struct multivariate_distribution_traits<T, std::enable_if_t<std::is_base_of_v<Sophus::SO3Base<T>, T>>> {
  using scalar_type = typename T::Scalar;

  using result_type = Sophus::SO3<scalar_type>;

  using vector_type = typename Eigen::Matrix<scalar_type, 3, 1>;

  using covariance_type = typename Eigen::Matrix<scalar_type, 3, 3>;

  [[nodiscard]] static constexpr vector_type to_vector(const result_type& t) { return t.log(); }

  [[nodiscard]] static constexpr result_type from_vector(const vector_type& v) {
    // Projecting a Gaussian through the exponential map is not exactly correct.
    // Rotations in Lie groups don't have infinite support, making the term 'normal' distribution incorrect.
    // This approach follows the SO2/SE2 method, facing similar issues in 3D.
    // For small covariances, the resulting rotations in SO3 approximate a normal distribution.
    return Sophus::SO3<scalar_type>::exp(v);
  }
};

template <class T>
struct multivariate_distribution_traits<T, std::enable_if_t<std::is_base_of_v<Sophus::SE3Base<T>, T>>> {
  using scalar_type = typename T::Scalar;

  using result_type = Sophus::SE3<scalar_type>;

  using vector_type = typename Eigen::Matrix<scalar_type, 6, 1>;

  using covariance_type = typename Eigen::Matrix<scalar_type, 6, 6>;

  [[nodiscard]] static constexpr vector_type to_vector(const result_type& t) {
    vector_type v;
    v << t.translation(), t.so3().log();
    return v;
  }

  [[nodiscard]] static constexpr result_type from_vector(const vector_type& v) {
    // Projecting a Gaussian through the exponential map is not exactly correct.
    // Rotations in Lie groups don't have infinite support, making the term 'normal' distribution incorrect.
    // This approach follows the SO2/SE2 method, facing similar issues in 3D.
    // For small covariances, the resulting rotations in SO3 approximate a normal distribution.
    return result_type{Sophus::SO3<scalar_type>::exp(v.tail(3)), v.head(3)};
  }
};

}  // namespace beluga

#endif