geometry.h

Go to the documentation of this file.

//
// Project NuriKit - Copyright 2023 SNU Compbio Lab.
// SPDX-License-Identifier: Apache-2.0
//
#ifndef NURI_CORE_GEOMETRY_H_
#define NURI_CORE_GEOMETRY_H_

#include <cmath>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>
 
#include <absl/log/absl_check.h>
#include <Eigen/Dense>
 
#include "nuri/eigen_config.h"
#include "nuri/meta.h"
 
namespace nuri {
namespace internal {
  using Array8i = Array<int, 8, 1>;
 
  class OCTreeNode {
  public:
    OCTreeNode(Array8i &&children, int begin, int nleaf)
        : children_(std::move(children)), begin_(begin), nleaf_(nleaf) { }
 
    const Array8i &children() const { return children_; }
 
    int child(int i) const { return children_[i]; }
 
    int operator[](int i) const { return child(i); }
 
    bool leaf() const { return nleaf_ <= 8; }
 
    int begin() const { return begin_; }
    int end() const { return begin_ + nleaf_; }
 
    int nleaf() const { return nleaf_; }
 
  private:
    Array8i children_;  // < 0 -> not exist, >= 0 -> index of child
    int begin_;
    int nleaf_;  // <= 8 -> leaf, > 8 -> internal node
  };
}  // namespace internal

class OCTree {
public:
  using Points = ConstRef<Matrix3Xd>;
 
  OCTree() = default;
 
  explicit OCTree(Points pts, int bucket_size = 32): bucket_size_(bucket_size) {
    rebuild_impl(pts);
  }
 
  void rebuild(Points pts, int bucket_size = -1) {
    if (bucket_size > 0)
      bucket_size_ = bucket_size;
 
    rebuild_impl(pts);
  }
 
  void find_neighbors_kd(const Vector3d &pt, int k, std::vector<int> &idxs,
                         std::vector<double> &distsq, double cutoff = -1) const;
 
  void find_neighbors_d(const Vector3d &pt, double cutoff,
                        std::vector<int> &idxs,
                        std::vector<double> &distsq) const;
 
  void find_neighbors_tree(const OCTree &oct, double cutoff,
                           std::vector<int> &self,
                           std::vector<int> &other) const;
 
  std::vector<std::vector<int>> find_neighbors_tree(const OCTree &oct,
                                                    double cutoff) const;
 
  void find_neighbors_self(double cutoff, std::vector<int> &left,
                           std::vector<int> &right) const;
 
  const Matrix3Xd &pts() const { return pts_; }
 
  const Vector3d &max() const { return max_; }
 
  const Vector3d &len() const { return len_; }

  void notify_transform(const Vector3d &new_max, const Vector3d &new_len);
 
  const std::vector<internal::OCTreeNode> &nodes() const { return nodes_; }
 
  int size() const { return static_cast<int>(nodes_.size()); }
 
  int root() const { return size() - 1; }
 
  int bucket_size() const { return bucket_size_; }
 
  const ArrayXi &idxs() const { return idxs_; }
 
  const internal::OCTreeNode &node(int i) const { return nodes_[i]; }
 
  const internal::OCTreeNode &operator[](int idx) const { return nodes_[idx]; }
 
private:
  void rebuild_impl(Points src);
 
  Matrix3Xd pts_;
  Vector3d max_;
  Vector3d len_;
  std::vector<internal::OCTreeNode> nodes_;
  ArrayXi idxs_;
 
  int bucket_size_ = 32;
};

class VoxelGrid {
public:
  using Points = ConstRef<Matrix3Xd>;
 
  VoxelGrid() = default;
 
  explicit VoxelGrid(Points pts, double cutoff): cutoff_(cutoff) {
    rebuild_impl(pts);
  }
 
  void rebuild(Points pts, double cutoff = -1.0) {
    if (cutoff > 0)
      cutoff_ = cutoff;
    rebuild_impl(pts);
  }

  void find_neighbors_d(const Vector3d &pt, std::vector<int> &idxs,
                        std::vector<double> &distsq) const;

  void find_neighbors_grid(const VoxelGrid &grid, std::vector<int> &self,
                           std::vector<int> &other) const;

  std::vector<std::vector<int>>
  find_neighbors_grid(const VoxelGrid &grid) const;

  void find_neighbors_self(std::vector<int> &left,
                           std::vector<int> &right) const;

  const Matrix3Xd &pts() const { return pts_; }
 
  double cutoff() const { return cutoff_; }
 
  const Vector3d &origin() const { return origin_; }
 
  const Array3i &dims() const { return dims_; }
 
  int num_cells() const { return static_cast<int>(cell_offset_.size()) - 1; }
 
  const ArrayXi &cell_offset() const { return cell_offset_; }
 
  const ArrayXi &cell_pts() const { return cell_pts_; }
 
private:
  void rebuild_impl(Points src);
 
  Matrix3Xd pts_;
  Vector3d origin_;
  Array3i dims_;
  ArrayXi cell_offset_;
  ArrayXi cell_pts_;
  double cutoff_;
  int max_occ_;
};
 
namespace constants {
  constexpr double kPi =
      3.1415926535897932384626433832795028841971693993751058209749445923078164;
 
  constexpr double kTwoPi =
      6.2831853071795864769252867665590057683943387987502116419498891846156328;
 
  // NOLINTBEGIN(*-identifier-naming)
  constexpr double kCos15 =
      0.9659258262890682867497431997288973676339048390084045504023430763;
  constexpr double kCos30 =
      0.8660254037844386467637231707529361834714026269051903140279034897;
  constexpr double kCos36 =
      0.8090169943749474241022934171828190588601545899028814310677243114;
  constexpr double kCos45 =
      0.7071067811865475244008443621048490392848359376884740365883398690;
  constexpr double kCos54 =
      0.5877852522924731291687059546390727685976524376431459910722724808;
  constexpr double kCos60 = 0.5;
  constexpr double kCos75 =
      0.2588190451025207623488988376240483283490689013199305138140032073;
  constexpr double kCos100 =
      -0.173648177666930348851716626769314796000375677184069387236241378;
  constexpr double kCos102 =
      -0.207911690817759337101742284405125166216584760627723836407181973;
  constexpr double kCos112 =
      -0.374606593415912035414963774501195131000158922253676174103440371;
  constexpr double kCos115 =
      -0.422618261740699436186978489647730181563129301194864623444415159;
  constexpr double kCos125 =
      -0.573576436351046096108031912826157864620433371450986351081027118;
  constexpr double kCos155 =
      -0.906307787036649963242552656754316983267712625175864680871298408;
  constexpr double kCos175 =
      -0.996194698091745532295010402473888046183562672645850974525442277;
  constexpr double kTan10_2 =
      0.0874886635259240052220186694349614581194542763681082291452366622;
  constexpr double kTan15_2 =
      0.1316524975873958534715264574097171035928141022232375735535653257;
  constexpr double kTan116_2 =
      1.6003345290410503553267330811833575255040718469227591484115002297;
  constexpr double kTan155_2 =
      4.5107085036620571342899391172547519686713241944553043587162345185;
  // NOLINTEND(*-identifier-naming)
}  // namespace constants
 
template <class DT, std::enable_if_t<std::is_floating_point_v<DT>, int> = 0>
constexpr DT deg2rad(DT deg) {
  return deg * static_cast<DT>(constants::kPi / 180);
}
 
template <class DT, std::enable_if_t<std::is_integral_v<DT>, int> = 0>
constexpr double deg2rad(DT deg) {
  return static_cast<double>(deg) * constants::kPi / 180;
}
 
template <class DT, std::enable_if_t<std::is_floating_point_v<DT>, int> = 0>
constexpr DT rad2deg(DT rad) {
  return rad * 180 / static_cast<DT>(constants::kPi);
}
 
template <class DT, std::enable_if_t<std::is_integral_v<DT>, int> = 0>
constexpr double rad2deg(DT rad) {
  return static_cast<double>(rad) * 180 / constants::kPi;
}
 
// NOLINTBEGIN(*-missing-std-forward)
 
template <class MatrixLike, class ArrayLike>
void pdistsq(ArrayLike &&distsq, const MatrixLike &m) {
  using DT = typename MatrixLike::Scalar;
  constexpr Eigen::Index rows = MatrixLike::RowsAtCompileTime;
 
  const Eigen::Index n = m.cols();
  ABSL_DCHECK(distsq.size() == n * (n - 1) / 2);
 
  Vector<DT, rows> v(m.rows());
  for (Eigen::Index i = 1, k = 0; i < n; ++i) {
    v = m.col(i);
    for (Eigen::Index j = 0; j < i; ++j, ++k) {
      distsq[k] = (v - m.col(j)).squaredNorm();
    }
  }
}
 
template <class MatrixLike>
auto pdistsq(const MatrixLike &m) {
  using DT = typename MatrixLike::Scalar;
 
  const Eigen::Index n = m.cols();
  ArrayX<DT> distsq(n * (n - 1) / 2);
  pdistsq(distsq, m);
  return distsq;
}
 
template <class MatrixLike, class ArrayLike>
void pdist(ArrayLike &&dist, const MatrixLike &m) {
  pdistsq(dist, m);
  dist = dist.sqrt();
}
 
template <class MatrixLike>
auto pdist(const MatrixLike &m) {
  auto ret = pdistsq(m);
  ret = ret.sqrt();
  return ret;
}
 
template <class AL, class DAL>
void to_square_form(DAL &&dists, const AL &pdists, Eigen::Index n) {
  ABSL_DCHECK(dists.rows() == n);
  ABSL_DCHECK(dists.cols() == n);
 
  dists.diagonal().setZero();
  for (Eigen::Index i = 1, k = 0; i < n; ++i)
    for (Eigen::Index j = 0; j < i; ++j, ++k)
      dists(i, j) = dists(j, i) = pdists[k];
}
 
template <class ArrayLike>
auto to_square_form(const ArrayLike &pdists, Eigen::Index n) {
  using DT = typename ArrayLike::Scalar;
 
  MatrixX<DT> dists(n, n);
  to_square_form(dists, pdists, n);
  return dists;
}
 
template <
    class ML1, class ML2, class DML,
    std::enable_if_t<internal::extract_if_enum_v(ML1::RowsAtCompileTime)
                         == internal::extract_if_enum_v(ML2::RowsAtCompileTime),
                     int> = 0>
void cdistsq(DML &&distsq, const ML1 &a, const ML2 &b) {
  ABSL_DCHECK(distsq.rows() == a.cols());
  ABSL_DCHECK(distsq.cols() == b.cols());
 
  for (Eigen::Index j = 0; j < b.cols(); ++j)
    distsq.col(j) = (a.colwise() - b.col(j)).colwise().squaredNorm();
}
 
template <class ML1, class ML2>
auto cdistsq(const ML1 &a, const ML2 &b) {
  using DT = typename ML1::Scalar;
 
  MatrixX<DT> distsq(a.cols(), b.cols());
  cdistsq(distsq, a, b);
  return distsq;
}
 
template <class ML1, class ML2, class DML>
void cdist(DML &&dist, const ML1 &a, const ML2 &b) {
  cdistsq(dist, a, b);
  dist = dist.cwiseSqrt();
}
 
template <class ML1, class ML2>
auto cdist(const ML1 &a, const ML2 &b) {
  auto ret = cdistsq(a, b);
  ret = ret.cwiseSqrt();
  return ret;
}
 
template <class ML1, class ML2>
auto msd(const ML1 &a, const ML2 &b) {
  return (a - b).colwise().squaredNorm().mean();
}
 
namespace internal {
  constexpr double safe_normalizer(double sqn, double eps = 1e-12) {
    return sqn > eps ? 1 / std::sqrt(sqn) : 0;
  }
 
  template <class VectorLike>
  inline void safe_normalize(VectorLike &&m, double eps = 1e-12) {
    m.array() *= safe_normalizer(m.squaredNorm(), eps);
  }
 
  template <class VectorLike>
  inline auto safe_normalized(VectorLike &&m, double eps = 1e-12) {
    using T = remove_cvref_t<VectorLike>;
    using Scalar = typename T::Scalar;
    constexpr auto size = T::SizeAtCompileTime;
    constexpr auto max_size = T::MaxSizeAtCompileTime;
 
    Matrix<Scalar, size, 1, 0, max_size, 1> ret = std::forward<VectorLike>(m);
    safe_normalize(ret, eps);
    return ret;
  }
 
  template <class MatrixLike>
  inline void safe_colwise_normalize(MatrixLike &&m, double eps = 1e-12) {
    using ArrayLike = decltype(m.colwise().squaredNorm().array());
    constexpr auto max_cols = ArrayLike::MaxColsAtCompileTime;
 
    if constexpr (max_cols != Eigen::Dynamic) {
      using T = remove_cvref_t<MatrixLike>;
      using Scalar = typename T::Scalar;
 
      constexpr auto cols = ArrayLike::ColsAtCompileTime;
 
      Array<Scalar, 1, cols, Eigen::RowMajor, 1, max_cols> norm =
          m.colwise().squaredNorm().array();
      m.array().rowwise() *= (norm > eps).select(norm.sqrt().inverse(), 0);
    } else {
      for (Eigen::Index i = 0; i < m.cols(); ++i)
        safe_normalize(m.col(i), eps);
    }
  }
 
  template <class MatrixLike>
  inline auto safe_colwise_normalized(MatrixLike &&m, double eps = 1e-12) {
    using T = remove_cvref_t<MatrixLike>;
    using Scalar = typename T::Scalar;
    constexpr auto rows = T::RowsAtCompileTime, cols = T::ColsAtCompileTime;
    constexpr auto max_rows = T::MaxRowsAtCompileTime,
                   max_cols = T::MaxColsAtCompileTime;
 
    Matrix<Scalar, rows, cols, 0, max_rows, max_cols> ret =
        std::forward<MatrixLike>(m);
    safe_colwise_normalize(ret, eps);
    return ret;
  }
}  // namespace internal
 
// NOLINTEND(*-missing-std-forward)

inline double cos_angle(const Vector3d &oa, const Vector3d &ob) {
  return oa.dot(ob)
         * internal::safe_normalizer(oa.squaredNorm() * ob.squaredNorm());
}

inline double cos_angle(const Vector3d &o, const Vector3d &a,
                        const Vector3d &b) {
  Vector3d oa = a - o, ob = b - o;
  return cos_angle(oa, ob);
}

template <class Scalar, class Indexer, int N, auto... Extra>
inline std::pair<Scalar, Scalar>
sum_tan2_half(const Matrix<Scalar, 3, N, 0, Extra...> &m, const Indexer &idxs) {
  double csum = (m + m(Eigen::all, idxs)).colwise().norm().sum(),
         ssum = (m - m(Eigen::all, idxs)).colwise().norm().sum();
  return std::make_pair(ssum, csum);
}

template <class Scalar, int N, auto... Extra>
inline std::pair<Scalar, Scalar>
sum_tan2_half(const Matrix<Scalar, 3, N, 0, Extra...> &m) {
  if constexpr (N == Eigen::Dynamic) {
    CyclicIndex<1> idxs(m.cols());
    return sum_tan2_half(m, idxs);
  }
 
  CyclicIndex<1, N> idxs;
  return sum_tan2_half(m, idxs);
}

inline double cos_dihedral(const Vector3d &axis, Vector3d v, Vector3d w) {
  v -= v.dot(axis) * axis;
  w -= w.dot(axis) * axis;
  return v.dot(w)
         * internal::safe_normalizer(v.squaredNorm() * w.squaredNorm());
}

inline double cos_dihedral(const Vector3d &a, const Vector3d &b,
                           const Vector3d &c, const Vector3d &d) {
  Vector3d axis = internal::safe_normalized(c - b);
  return cos_dihedral(axis, a - b, d - c);
}

template <class MatrixLike>
Vector4d fit_plane(const MatrixLike &pts, bool normalize = true) {
  Vector3d cntr = pts.rowwise().mean();
  MatrixXd m = pts.colwise() - cntr;
  auto svd = m.jacobiSvd(Eigen::ComputeThinU);
 
  Vector4d ret;
  ret.head<3>() = svd.matrixU().col(2);
  if (normalize)
    internal::safe_normalize(ret.head<3>());
  ret[3] = -ret.head<3>().dot(cntr);
  return ret;
}

template <class VectorLike>
Vector3d any_perpendicular(const VectorLike &v, bool normalize = true) {
  Vector3d w = { std::copysign(v[2], v[0]),  //
                 std::copysign(v[2], v[1]),
                 -std::copysign(v[0], v[2]) - std::copysign(v[1], v[2]) };
  if (normalize)
    internal::safe_normalize(w);
  return w;
}
 
enum class AlignMode : std::uint8_t {
  kMsdOnly = 0x1,
  kXformOnly = 0x2,
  kBoth = kMsdOnly | kXformOnly,
};

extern std::pair<Isometry3d, double> kabsch(ConstRef<Matrix3Xd> query,
                                            ConstRef<Matrix3Xd> templ,
                                            AlignMode mode = AlignMode::kBoth,
                                            bool reflection = false);

extern std::pair<Isometry3d, double>
qcp(const Eigen::Ref<const Matrix3Xd> &query,
    const Eigen::Ref<const Matrix3Xd> &templ, AlignMode mode = AlignMode::kBoth,
    bool reflection = false, double evalprec = 1e-11, double evecprec = 1e-6,
    int maxiter = 50);

extern std::pair<Isometry3d, double>
qcp_inplace(MutRef<Matrix3Xd> query, MutRef<Matrix3Xd> templ,
            AlignMode mode = AlignMode::kBoth, bool reflection = false,
            double evalprec = 1e-11, double evecprec = 1e-6, int maxiter = 50);

extern bool embed_distances_3d(Eigen::Ref<Matrix3Xd> pts, MatrixXd &dsqs);

extern bool embed_distances_4d(Eigen::Ref<Matrix4Xd> pts, MatrixXd &dsqs);

extern Matrix3Xd canonical_fibonacci_lattice(int npts);
}  // namespace nuri
 
#endif /* NURI_CORE_GEOMETRY_H_ */