aos/vision/math/vector.h

#ifndef AOS_VISION_MATH_VECTOR_H_
#define AOS_VISION_MATH_VECTOR_H_

#include <cmath>

#include "Eigen/Dense"

namespace aos::vision {

// Represents an n-dimensional vector of doubles with various convenient
// shortcuts for common operations.
//
// This includes overloads of various arithmetic operators for convenience.
// Multiplication by doubles is scalar multiplication and multiplication by
// other vectors is element-wise.
//
// Accessing elements which don't exist for a given size vector, mixing sizes of
// vectors in appropriately, etc are compile-time errors.
template <int size>
class Vector {
 public:
  Vector() : data_(::Eigen::Matrix<double, 1, size>::Zero()) {}

  Vector(double x, double y) { Set(x, y); }

  Vector(double x, double y, double z) { Set(x, y, z); }

  double Get(int index) const { return data_(index); }
  void Set(int index, double value) { data_(index) = value; }

  void Set(double x, double y) {
    static_assert(size == 2, "illegal size");
    data_(0) = x;
    data_(1) = y;
  }
  void Set(double x, double y, double z) {
    static_assert(size == 3, "illegal size");
    data_(0) = x;
    data_(1) = y;
    data_(2) = z;
  }
  void Set(double x, double y, double z, double w) {
    static_assert(size == 4, "illegal size");
    data_(0) = x;
    data_(1) = y;
    data_(2) = z;
    data_(3) = w;
  }

  double x() const {
    static_assert(size >= 1, "illegal size");
    return data_(0);
  }
  void x(double xX) {
    static_assert(size >= 1, "illegal size");
    data_(0) = xX;
  }

  double y() const {
    static_assert(size >= 2, "illegal size");
    return data_(1);
  }
  void y(double yY) {
    static_assert(size >= 2, "illegal size");
    data_(1) = yY;
  }

  double z() const {
    static_assert(size >= 3, "illegal size");
    return data_(2);
  }
  void z(double zZ) {
    static_assert(size >= 3, "illegal size");
    data_(2) = zZ;
  }

  double w() const {
    static_assert(size >= 4, "illegal size");
    return data_(3);
  }
  void w(double wW) {
    static_assert(size >= 4, "illegal size");
    data_(3) = wW;
  }

  // Fast part of length.
  double MagSqr() const { return data_.squaredNorm(); }

  // Length of the vector.
  double Mag() const { return data_.norm(); }

  // Get underlying data structure
  // TODO(austin): Everyone uses column not row vectors.
  ::Eigen::Matrix<double, 1, size> GetData() const { return data_; }

  // Set underlying data structure
  void SetData(const ::Eigen::Matrix<double, 1, size> &other) { data_ = other; }

  Vector<size> operator+(const Vector<size> &other) const {
    Vector<size> nv = *this;
    nv += other;
    return nv;
  }
  Vector<size> operator+=(const Vector<size> &other) {
    data_ += other.data_;
    return *this;
  }

  Vector<size> operator-(const Vector<size> &other) const {
    Vector<size> nv = *this;
    nv -= other;
    return nv;
  }
  Vector<size> operator-=(const Vector<size> &other) {
    data_ -= other.data_;
    return *this;
  }

  Vector<size> operator*(double other) {
    Vector<size> nv = *this;
    nv *= other;
    return nv;
  }
  Vector<size> operator*=(double other) {
    data_ *= other;
    return *this;
  }

  Vector<size> operator*(const Vector<size> &other) const {
    Vector<size> nv = *this;
    nv *= other;
    return nv;
  }
  Vector<size> operator*=(const Vector<size> &other) {
    for (int i = 0; i < size; i++) {
      Set(i, other.Get(i) * Get(i));
    }
    return *this;
  }

  bool operator==(const Vector<size> &other) const {
    for (int i = 0; i < size; i++) {
      if (Get(i) != other.Get(i)) {
        return false;
      }
    }
    return true;
  }

  double dot(const Vector<size> &other) const {
    return data_.dot(other.GetData());
  }

  Vector<size> cross(const Vector<size> &other) const {
    Vector<size> nv;
    nv.SetData(data_.cross(other.GetData()));
    return nv;
  }

  // Returns a vector in the same direction as this one with a magnitude of 1.
  Vector<size> Normalized() const {
    double mag = Mag();
    Vector<size> nv;
    for (int i = 0; i < size; i++) {
      nv.Set(i, Get(i) / mag);
    }
    return nv;
  }

  // Returns the angle between this vector and the 0 vector.
  // Only valid for 2-dimensional vectors.
  double AngleToZero() const {
    static_assert(size == 2, "illegal size");
    return ::std::atan2(y(), x());
  }

  // Return the angle between this and other.
  double AngleTo(const Vector<size> &other) const {
    // cos(theta) = u.dot(v) / (u.magnitude() * v.magnitude())
    return ::std::acos(dot(other) / (Mag() * other.Mag()));
  }

  // Returns the distance between this and other squared.
  double SquaredDistanceTo(const Vector<size> &other) const {
    Vector<size> tmp = *this - other;
    return tmp.MagSqr();
  }

  // Returns the distance between this and other.
  double DistanceTo(const Vector<size> &other) const {
    return std::sqrt(SquaredDistanceTo(other));
  }

 private:
  // The actual data.
  ::Eigen::Matrix<double, 1, size> data_;
};

// Returns the cross product of two points.
inline double PointsCrossProduct(const Vector<2> &a, const Vector<2> &b) {
  return a.x() * b.y() - a.y() * b.x();
}

// Rotates along the z, y, and then x axis (all radians).
inline Vector<3> Rotate(double rx, double ry, double rz, const Vector<3> vec) {
  ::Eigen::AngleAxis<double> ax(rx, ::Eigen::Vector3d(1.0, 0.0, 0.0));
  ::Eigen::AngleAxis<double> ay(ry, ::Eigen::Vector3d(0.0, 1.0, 0.0));
  ::Eigen::AngleAxis<double> az(rz, ::Eigen::Vector3d(0.0, 0.0, 1.0));
  Vector<3> result;
  result.SetData(ax * ay * az * vec.GetData());
  return result;
}

// scalar multiply
template <int Size>
inline Vector<Size> operator*(const double &lhs, Vector<Size> &rhs) {
  Vector<Size> nv;
  for (int i = 0; i < Size; i++) {
    nv.Set(i, lhs * rhs.Get(i));
  }
  return nv;
}

}  // namespace aos::vision

#endif  // AOS_VISION_MATH_VECTOR_H_