Quaternion.java

// Copyright (c) FIRST and other WPILib contributors.
// Open Source Software; you can modify and/or share it under the terms of
// the WPILib BSD license file in the root directory of this project.
 
package org.wpilib.math.geometry;
 
import io.avaje.jsonb.Json;
import java.util.Objects;
import org.wpilib.math.geometry.proto.QuaternionProto;
import org.wpilib.math.geometry.struct.QuaternionStruct;
import org.wpilib.math.linalg.VecBuilder;
import org.wpilib.math.linalg.Vector;
import org.wpilib.math.numbers.N3;
import org.wpilib.util.protobuf.ProtobufSerializable;
import org.wpilib.util.struct.StructSerializable;

@Json
public final class Quaternion implements ProtobufSerializable, StructSerializable {
  // Scalar r in versor form
  @Json.Ignore private final double m_w;
 
  // Vector v in versor form
  @Json.Ignore private final double m_x;
  @Json.Ignore private final double m_y;
  @Json.Ignore private final double m_z;

  public Quaternion() {
    m_w = 1.0;
    m_x = 0.0;
    m_y = 0.0;
    m_z = 0.0;
  }

  @Json.Creator
  public Quaternion(
      @Json.Alias("W") double w,
      @Json.Alias("X") double x,
      @Json.Alias("Y") double y,
      @Json.Alias("Z") double z) {
    m_w = w;
    m_x = x;
    m_y = y;
    m_z = z;
  }

  public Quaternion plus(Quaternion other) {
    return new Quaternion(
        getW() + other.getW(), getX() + other.getX(), getY() + other.getY(), getZ() + other.getZ());
  }

  public Quaternion minus(Quaternion other) {
    return new Quaternion(
        getW() - other.getW(), getX() - other.getX(), getY() - other.getY(), getZ() - other.getZ());
  }

  public Quaternion divide(double scalar) {
    return new Quaternion(getW() / scalar, getX() / scalar, getY() / scalar, getZ() / scalar);
  }

  public Quaternion times(double scalar) {
    return new Quaternion(getW() * scalar, getX() * scalar, getY() * scalar, getZ() * scalar);
  }

  public Quaternion times(Quaternion other) {
    // https://en.wikipedia.org/wiki/Quaternion#Scalar_and_vector_parts
    final var r1 = m_w;
    final var r2 = other.m_w;
 
    // v₁ ⋅ v₂
    double dot = m_x * other.m_x + m_y * other.m_y + m_z * other.m_z;
 
    // v₁ x v₂
    double cross_x = m_y * other.m_z - other.m_y * m_z;
    double cross_y = other.m_x * m_z - m_x * other.m_z;
    double cross_z = m_x * other.m_y - other.m_x * m_y;
 
    return new Quaternion(
        // r = r₁r₂ − v₁ ⋅ v₂
        r1 * r2 - dot,
        // v = r₁v₂ + r₂v₁ + v₁ x v₂
        r1 * other.m_x + r2 * m_x + cross_x,
        r1 * other.m_y + r2 * m_y + cross_y,
        r1 * other.m_z + r2 * m_z + cross_z);
  }
 
  @Override
  public String toString() {
    return String.format("Quaternion(%s, %s, %s, %s)", getW(), getX(), getY(), getZ());
  }

  @Override
  public boolean equals(Object obj) {
    return obj instanceof Quaternion other
        && Math.abs(dot(other) - norm() * other.norm()) < 1e-9
        && Math.abs(norm() - other.norm()) < 1e-9;
  }
 
  @Override
  public int hashCode() {
    return Objects.hash(m_w, m_x, m_y, m_z);
  }

  public Quaternion conjugate() {
    return new Quaternion(getW(), -getX(), -getY(), -getZ());
  }

  public double dot(final Quaternion other) {
    return getW() * other.getW()
        + getX() * other.getX()
        + getY() * other.getY()
        + getZ() * other.getZ();
  }

  public Quaternion inverse() {
    var norm = norm();
    return conjugate().divide(norm * norm);
  }

  public double norm() {
    return Math.sqrt(dot(this));
  }

  public Quaternion normalize() {
    double norm = norm();
    if (norm == 0.0) {
      return new Quaternion();
    } else {
      return new Quaternion(getW() / norm, getX() / norm, getY() / norm, getZ() / norm);
    }
  }

  public Quaternion pow(double t) {
    // q^t = e^(ln(q^t)) = e^(t * ln(q))
    return this.log().times(t).exp();
  }

  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);
  }

  public Quaternion log() {
    var norm = norm();
    if (norm == 0.0) {
      return new Quaternion(0.0, 0.0, 0.0, 0.0);
    }
 
    var scalar = Math.log(norm);
 
    var v_norm = Math.sqrt(getX() * getX() + getY() * getY() + getZ() * getZ());
 
    var w = getW();
    var s_norm = w / 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 && w != 0.0) {
      // Taylor series expansion of atan2(y/x)/y at y = 0:
      //
      //   1/x - 1/3 y²/x³ + O(y⁴)
      v_scalar = 1.0 / w - 1.0 / 3.0 * v_norm * v_norm / (w * w * w);
    } else if (v_norm == 0.0) {
      v_scalar = 0.0;
    } else {
      v_scalar = Math.atan2(v_norm, w) / v_norm;
    }
 
    return new Quaternion(scalar, v_scalar * getX(), v_scalar * getY(), v_scalar * getZ());
  }

  @Json.Property(value = "W")
  public double getW() {
    return m_w;
  }

  @Json.Property(value = "X")
  public double getX() {
    return m_x;
  }

  @Json.Property(value = "Y")
  public double getY() {
    return m_y;
  }

  @Json.Property(value = "Z")
  public double getZ() {
    return m_z;
  }

  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));
  }

  public Vector<N3> toRotationVector() {
    // See equation (31) in "Integrating Generic Sensor Fusion Algorithms with
    // Sound State Representation through Encapsulation of Manifolds"
    //
    // https://arxiv.org/pdf/1107.1119.pdf
    double norm = Math.sqrt(getX() * getX() + getY() * getY() + getZ() * getZ());
 
    double coeff;
    if (norm < 1e-9) {
      coeff = 2.0 / getW() - 2.0 / 3.0 * norm * norm / (getW() * getW() * getW());
    } else {
      if (getW() < 0.0) {
        coeff = 2.0 * Math.atan2(-norm, -getW()) / norm;
      } else {
        coeff = 2.0 * Math.atan2(norm, getW()) / norm;
      }
    }
 
    return VecBuilder.fill(coeff * getX(), coeff * getY(), coeff * getZ());
  }

  public static final QuaternionProto proto = new QuaternionProto();

  public static final QuaternionStruct struct = new QuaternionStruct();
}