[wpimath] Expand Quaternion class with additional operators (#5600)

Co-authored-by: Tyler Veness <calcmogul@gmail.com>

wpimath/src/main/java/edu/wpi/first/math/geometry/Quaternion.java

@@ -52,6 +52,48 @@ public class Quaternion {
     m_z = z;
   }
 
+  /**
+   * Adds another quaternion to this quaternion entrywise.
+   *
+   * @param other The other quaternion.
+   * @return The quaternion sum.
+   */
+  public Quaternion plus(Quaternion other) {
+    return new Quaternion(
+        getW() + other.getW(), getX() + other.getX(), getY() + other.getY(), getZ() + other.getZ());
+  }
+
+  /**
+   * Subtracts another quaternion from this quaternion entrywise.
+   *
+   * @param other The other quaternion.
+   * @return The quaternion difference.
+   */
+  public Quaternion minus(Quaternion other) {
+    return new Quaternion(
+        getW() - other.getW(), getX() - other.getX(), getY() - other.getY(), getZ() - other.getZ());
+  }
+
+  /**
+   * Divides by a scalar.
+   *
+   * @param scalar The value to scale each component by.
+   * @return The scaled quaternion.
+   */
+  public Quaternion divide(double scalar) {
+    return new Quaternion(getW() / scalar, getX() / scalar, getY() / scalar, getZ() / scalar);
+  }
+
+  /**
+   * Multiplies with a scalar.
+   *
+   * @param scalar The value to scale each component by.
+   * @return The scaled quaternion.
+   */
+  public Quaternion times(double scalar) {
+    return new Quaternion(getW() * scalar, getX() * scalar, getY() * scalar, getZ() * scalar);
+  }
+
   /**
    * Multiply with another quaternion.
    *
@@ -96,12 +138,8 @@ public class Quaternion {
     if (obj instanceof Quaternion) {
       var other = (Quaternion) obj;
 
-      return Math.abs(
-              getW() * other.getW()
-                  + getX() * other.getX()
-                  + getY() * other.getY()
-                  + getZ() * other.getZ())
-          > 1.0 - 1E-9;
+      return Math.abs(dot(other) - norm() * other.norm()) < 1e-9
+          && Math.abs(norm() - other.norm()) < 1e-9;
     }
     return false;
   }
@@ -111,13 +149,45 @@ public class Quaternion {
     return Objects.hash(m_w, m_x, m_y, m_z);
   }
 
+  /**
+   * Returns the conjugate of the quaternion.
+   *
+   * @return The conjugate quaternion.
+   */
+  public Quaternion conjugate() {
+    return new Quaternion(getW(), -getX(), -getY(), -getZ());
+  }
+
+  /**
+   * Returns the elementwise product of two quaternions.
+   *
+   * @param other The other quaternion.
+   * @return The dot product of two quaternions.
+   */
+  public double dot(final Quaternion other) {
+    return getW() * other.getW()
+        + getX() * other.getX()
+        + getY() * other.getY()
+        + getZ() * other.getZ();
+  }
+
   /**
    * Returns the inverse of the quaternion.
    *
    * @return The inverse quaternion.
    */
   public Quaternion inverse() {
-    return new Quaternion(getW(), -getX(), -getY(), -getZ());
+    var norm = norm();
+    return conjugate().divide(norm * norm);
+  }
+
+  /**
+   * Calculates the L2 norm of the quaternion.
+   *
+   * @return The L2 norm.
+   */
+  public double norm() {
+    return Math.sqrt(dot(this));
   }
 
   /**
@@ -126,7 +196,7 @@ public class Quaternion {
    * @return The normalized quaternion.
    */
   public Quaternion normalize() {
-    double norm = Math.sqrt(getW() * getW() + getX() * getX() + getY() * getY() + getZ() * getZ());
+    double norm = norm();
     if (norm == 0.0) {
       return new Quaternion();
     } else {
@@ -134,6 +204,104 @@ public class Quaternion {
     }
   }
 
+  /**
+   * Rational power of a quaternion.
+   *
+   * @param t the power to raise this quaternion to.
+   * @return The quaternion power
+   */
+  public Quaternion pow(double t) {
+    // q^t = e^(ln(q^t)) = e^(t * ln(q))
+    return this.log().times(t).exp();
+  }
+
+  /**
+   * Matrix exponential of a quaternion.
+   *
+   * @param adjustment the "Twist" that will be applied to this quaternion.
+   * @return The quaternion product of exp(adjustment) * this
+   */
+  public Quaternion exp(Quaternion adjustment) {
+    return adjustment.exp().times(this);
+  }
+
+  /**
+   * Matrix exponential of a quaternion.
+   *
+   * <p>source: wpimath/algorithms.md
+   *
+   * <p>If this quaternion is in 𝖘𝖔(3) and you are looking for an element of SO(3), use {@link
+   * fromRotationVector}
+   *
+   * @return The Matrix exponential of this quaternion.
+   */
+  public Quaternion exp() {
+    var scalar = Math.exp(getW());
+
+    var axial_magnitude = Math.sqrt(getX() * getX() + getY() * getY() + getZ() * getZ());
+    var cosine = Math.cos(axial_magnitude);
+
+    double axial_scalar;
+
+    if (axial_magnitude < 1e-9) {
+      // Taylor series of sin(θ) / θ near θ = 0: 1 − θ²/6 + θ⁴/120 + O(n⁶)
+      var axial_magnitude_sq = axial_magnitude * axial_magnitude;
+      var axial_magnitude_sq_sq = axial_magnitude_sq * axial_magnitude_sq;
+      axial_scalar = 1.0 - axial_magnitude_sq / 6.0 + axial_magnitude_sq_sq / 120.0;
+    } else {
+      axial_scalar = Math.sin(axial_magnitude) / axial_magnitude;
+    }
+
+    return new Quaternion(
+        cosine * scalar,
+        getX() * axial_scalar * scalar,
+        getY() * axial_scalar * scalar,
+        getZ() * axial_scalar * scalar);
+  }
+
+  /**
+   * Log operator of a quaternion.
+   *
+   * @param end The quaternion to map this quaternion onto.
+   * @return The "Twist" that maps this quaternion to the argument.
+   */
+  public Quaternion log(Quaternion end) {
+    return end.times(this.inverse()).log();
+  }
+
+  /**
+   * The Log operator of a general quaternion.
+   *
+   * <p>source: wpimath/algorithms.md
+   *
+   * <p>If this quaternion is in SO(3) and you are looking for an element of 𝖘𝖔(3), use {@link
+   * toRotationVector}
+   *
+   * @return The logarithm of this quaternion.
+   */
+  public Quaternion log() {
+    var scalar = Math.log(norm());
+
+    var v_norm = Math.sqrt(getX() * getX() + getY() * getY() + getZ() * getZ());
+
+    var s_norm = getW() / norm();
+
+    if (Math.abs(s_norm + 1) < 1e-9) {
+      return new Quaternion(scalar, -Math.PI, 0, 0);
+    }
+
+    double v_scalar;
+
+    if (v_norm < 1e-9) {
+      // Taylor series expansion of atan2(y / x) / y around y = 0 => 1/x - y²/3*x³ + O(y⁴)
+      v_scalar = 1.0 / getW() - 1.0 / 3.0 * v_norm * v_norm / (getW() * getW() * getW());
+    } else {
+      v_scalar = Math.atan2(v_norm, getW()) / v_norm;
+    }
+
+    return new Quaternion(scalar, v_scalar * getX(), v_scalar * getY(), v_scalar * getZ());
+  }
+
   /**
    * Returns W component of the quaternion.
    *
@@ -174,6 +342,37 @@ public class Quaternion {
     return m_z;
   }
 
+  /**
+   * Returns the quaternion representation of this rotation vector.
+   *
+   * <p>This is also the exp operator of 𝖘𝖔(3).
+   *
+   * <p>source: wpimath/algorithms.md
+   *
+   * @param rvec The rotation vector.
+   * @return The quaternion representation of this rotation vector.
+   */
+  public static Quaternion fromRotationVector(Vector<N3> rvec) {
+    double theta = rvec.norm();
+
+    double cos = Math.cos(theta / 2);
+
+    double axial_scalar;
+
+    if (theta < 1e-9) {
+      // taylor series expansion of sin(θ/2) / θ = 1/2 - θ²/48 + O(θ⁴)
+      axial_scalar = 1.0 / 2.0 - theta * theta / 48.0;
+    } else {
+      axial_scalar = Math.sin(theta / 2) / theta;
+    }
+
+    return new Quaternion(
+        cos,
+        axial_scalar * rvec.get(0, 0),
+        axial_scalar * rvec.get(1, 0),
+        axial_scalar * rvec.get(2, 0));
+  }
+
   /**
    * Returns the rotation vector representation of this quaternion.
    *

wpimath/src/main/java/edu/wpi/first/math/geometry/Rotation3d.java

@@ -420,7 +420,7 @@ public class Rotation3d implements Interpolatable<Rotation3d> {
   public boolean equals(Object obj) {
     if (obj instanceof Rotation3d) {
       var other = (Rotation3d) obj;
-      return m_q.equals(other.m_q);
+      return Math.abs(Math.abs(m_q.dot(other.m_q)) - m_q.norm() * other.m_q.norm()) < 1e-9;
     }
     return false;
   }