Rotation3d.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 static org.wpilib.math.util.UnitConversions.rotationsToRadians;
import static org.wpilib.units.Units.Radians;
 
import io.avaje.jsonb.Json;
import java.util.Objects;
import org.wpilib.math.geometry.proto.Rotation3dProto;
import org.wpilib.math.geometry.struct.Rotation3dStruct;
import org.wpilib.math.interpolation.Interpolatable;
import org.wpilib.math.linalg.MatBuilder;
import org.wpilib.math.linalg.Matrix;
import org.wpilib.math.linalg.VecBuilder;
import org.wpilib.math.linalg.Vector;
import org.wpilib.math.numbers.N3;
import org.wpilib.math.util.MathSharedStore;
import org.wpilib.math.util.Nat;
import org.wpilib.units.measure.Angle;
import org.wpilib.util.protobuf.ProtobufSerializable;
import org.wpilib.util.struct.StructSerializable;

@Json
public final class Rotation3d
    implements Interpolatable<Rotation3d>, ProtobufSerializable, StructSerializable {
  public static final Rotation3d ZERO = new Rotation3d();
 
  @Json.Property("quaternion")
  private final Quaternion m_q;

  public Rotation3d() {
    m_q = new Quaternion();
  }

  @Json.Creator
  public Rotation3d(@Json.Alias("quaternion") Quaternion q) {
    m_q = q.normalize();
  }

  public Rotation3d(double roll, double pitch, double yaw) {
    // https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Euler_angles_to_quaternion_conversion
    double cr = Math.cos(roll * 0.5);
    double sr = Math.sin(roll * 0.5);
 
    double cp = Math.cos(pitch * 0.5);
    double sp = Math.sin(pitch * 0.5);
 
    double cy = Math.cos(yaw * 0.5);
    double sy = Math.sin(yaw * 0.5);
 
    m_q =
        new Quaternion(
            cr * cp * cy + sr * sp * sy,
            sr * cp * cy - cr * sp * sy,
            cr * sp * cy + sr * cp * sy,
            cr * cp * sy - sr * sp * cy);
  }

  public Rotation3d(Angle roll, Angle pitch, Angle yaw) {
    this(roll.in(Radians), pitch.in(Radians), yaw.in(Radians));
  }

  public Rotation3d(Vector<N3> rvec) {
    this(rvec, rvec.norm());
  }

  public Rotation3d(Vector<N3> axis, double angleRadians) {
    double norm = axis.norm();
    if (norm == 0.0) {
      m_q = new Quaternion();
      return;
    }
 
    // https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Definition
    var v = axis.times(1.0 / norm).times(Math.sin(angleRadians / 2.0));
    m_q = new Quaternion(Math.cos(angleRadians / 2.0), v.get(0, 0), v.get(1, 0), v.get(2, 0));
  }

  public Rotation3d(Vector<N3> axis, Angle angle) {
    this(axis, angle.in(Radians));
  }

  public Rotation3d(Matrix<N3, N3> rotationMatrix) {
    final var R = rotationMatrix;
 
    // Require that the rotation matrix is special orthogonal. This is true if
    // the matrix is orthogonal (RRᵀ = I) and normalized (determinant is 1).
    if (R.times(R.transpose()).minus(Matrix.eye(Nat.N3())).normF() > 1e-9) {
      var msg = "Rotation matrix isn't orthogonal\n\nR =\n" + R.getStorage().toString() + '\n';
      MathSharedStore.reportError(msg, Thread.currentThread().getStackTrace());
      throw new IllegalArgumentException(msg);
    }
    if (Math.abs(R.det() - 1.0) > 1e-9) {
      var msg =
          "Rotation matrix is orthogonal but not special orthogonal\n\nR =\n"
              + R.getStorage().toString()
              + '\n';
      MathSharedStore.reportError(msg, Thread.currentThread().getStackTrace());
      throw new IllegalArgumentException(msg);
    }
 
    // Turn rotation matrix into a quaternion
    // https://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToQuaternion/
    double trace = R.get(0, 0) + R.get(1, 1) + R.get(2, 2);
    double w;
    double x;
    double y;
    double z;
 
    if (trace > 0.0) {
      double s = 0.5 / Math.sqrt(trace + 1.0);
      w = 0.25 / s;
      x = (R.get(2, 1) - R.get(1, 2)) * s;
      y = (R.get(0, 2) - R.get(2, 0)) * s;
      z = (R.get(1, 0) - R.get(0, 1)) * s;
    } else {
      if (R.get(0, 0) > R.get(1, 1) && R.get(0, 0) > R.get(2, 2)) {
        double s = 2.0 * Math.sqrt(1.0 + R.get(0, 0) - R.get(1, 1) - R.get(2, 2));
        w = (R.get(2, 1) - R.get(1, 2)) / s;
        x = 0.25 * s;
        y = (R.get(0, 1) + R.get(1, 0)) / s;
        z = (R.get(0, 2) + R.get(2, 0)) / s;
      } else if (R.get(1, 1) > R.get(2, 2)) {
        double s = 2.0 * Math.sqrt(1.0 + R.get(1, 1) - R.get(0, 0) - R.get(2, 2));
        w = (R.get(0, 2) - R.get(2, 0)) / s;
        x = (R.get(0, 1) + R.get(1, 0)) / s;
        y = 0.25 * s;
        z = (R.get(1, 2) + R.get(2, 1)) / s;
      } else {
        double s = 2.0 * Math.sqrt(1.0 + R.get(2, 2) - R.get(0, 0) - R.get(1, 1));
        w = (R.get(1, 0) - R.get(0, 1)) / s;
        x = (R.get(0, 2) + R.get(2, 0)) / s;
        y = (R.get(1, 2) + R.get(2, 1)) / s;
        z = 0.25 * s;
      }
    }
 
    m_q = new Quaternion(w, x, y, z);
  }

  public Rotation3d(Vector<N3> initial, Vector<N3> last) {
    double dot = initial.dot(last);
    double normProduct = initial.norm() * last.norm();
    double dotNorm = dot / normProduct;
 
    if (dotNorm > 1.0 - 1E-9) {
      // If the dot product is 1, the two vectors point in the same direction so
      // there's no rotation. The default initialization of m_q will work.
      m_q = new Quaternion();
    } else if (dotNorm < -1.0 + 1E-9) {
      // If the dot product is -1, the two vectors are antiparallel, so a 180°
      // rotation is required. Any other vector can be used to generate an
      // orthogonal one.
 
      double x = Math.abs(initial.get(0, 0));
      double y = Math.abs(initial.get(1, 0));
      double z = Math.abs(initial.get(2, 0));
 
      // Find vector that is most orthogonal to initial vector
      Vector<N3> other;
      if (x < y) {
        if (x < z) {
          // Use x-axis
          other = VecBuilder.fill(1, 0, 0);
        } else {
          // Use z-axis
          other = VecBuilder.fill(0, 0, 1);
        }
      } else {
        if (y < z) {
          // Use y-axis
          other = VecBuilder.fill(0, 1, 0);
        } else {
          // Use z-axis
          other = VecBuilder.fill(0, 0, 1);
        }
      }
 
      var axis = Vector.cross(initial, other);
 
      double axisNorm = axis.norm();
      m_q =
          new Quaternion(
              0.0, axis.get(0, 0) / axisNorm, axis.get(1, 0) / axisNorm, axis.get(2, 0) / axisNorm);
    } else {
      var axis = Vector.cross(initial, last);
 
      // https://stackoverflow.com/a/11741520
      m_q =
          new Quaternion(normProduct + dot, axis.get(0, 0), axis.get(1, 0), axis.get(2, 0))
              .normalize();
    }
  }

  public Rotation3d(Rotation2d rotation) {
    this(0.0, 0.0, rotation.getRadians());
  }

  public static Rotation3d fromRadians(double roll, double pitch, double yaw) {
    return new Rotation3d(roll, pitch, yaw);
  }

  public static Rotation3d fromDegrees(double roll, double pitch, double yaw) {
    return new Rotation3d(Math.toRadians(roll), Math.toRadians(pitch), Math.toRadians(yaw));
  }

  public static Rotation3d fromRotations(double roll, double pitch, double yaw) {
    return new Rotation3d(
        rotationsToRadians(roll), rotationsToRadians(pitch), rotationsToRadians(yaw));
  }

  public Rotation3d inverse() {
    return new Rotation3d(m_q.inverse());
  }

  public Rotation3d times(double scalar) {
    return Rotation3d.ZERO.interpolate(this, scalar);
  }

  public Rotation3d div(double scalar) {
    return times(1.0 / scalar);
  }

  public Rotation3d rotateBy(Rotation3d other) {
    return new Rotation3d(other.m_q.times(m_q));
  }

  public Rotation3d relativeTo(Rotation3d other) {
    // To apply a quaternion intrinsically, we must right-multiply by that quaternion.
    // Therefore, "this_q relative to other_q" is the q such that other_q q = this_q:
    //
    //   other_q q = this_q
    //   q = other_q⁻¹ this_q
    return new Rotation3d(other.m_q.inverse().times(m_q));
  }

  public Rotation3d integrate(double rollRate, double pitchRate, double yawRate, double dt) {
    // qₖ₊₁ = qₖ exp(1/2 W dt) where W = 0 + ω_x î + ω_y ĵ + ω_z k̂
    //
    // https://math.stackexchange.com/a/2099673
    var W = new Quaternion(0.0, rollRate, pitchRate, yawRate);
    return new Rotation3d(m_q.times(W.times(dt / 2.0).exp()));
  }

  public Quaternion getQuaternion() {
    return m_q;
  }

  public double getX() {
    final var w = m_q.getW();
    final var x = m_q.getX();
    final var y = m_q.getY();
    final var z = m_q.getZ();
 
    // wpimath/docs/Quaternion.md
    final var cxcy = 1.0 - 2.0 * (x * x + y * y);
    final var sxcy = 2.0 * (w * x + y * z);
    final var cy_sq = cxcy * cxcy + sxcy * sxcy;
    if (cy_sq > 1e-20) {
      return Math.atan2(sxcy, cxcy);
    } else {
      return 0.0;
    }
  }

  public double getY() {
    final var w = m_q.getW();
    final var x = m_q.getX();
    final var y = m_q.getY();
    final var z = m_q.getZ();
 
    // https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles#Quaternion_to_Euler_angles_(in_3-2-1_sequence)_conversion
    double ratio = 2.0 * (w * y - z * x);
    if (Math.abs(ratio) >= 1.0) {
      return Math.copySign(Math.PI / 2.0, ratio);
    } else {
      return Math.asin(ratio);
    }
  }

  public double getZ() {
    final var w = m_q.getW();
    final var x = m_q.getX();
    final var y = m_q.getY();
    final var z = m_q.getZ();
 
    // wpimath/docs/Quaternion.md
    final var cycz = 1.0 - 2.0 * (y * y + z * z);
    final var cysz = 2.0 * (w * z + x * y);
    final var cy_sq = cycz * cycz + cysz * cysz;
    if (cy_sq > 1e-20) {
      return Math.atan2(cysz, cycz);
    } else {
      return Math.atan2(2.0 * w * z, w * w - z * z);
    }
  }

  public Angle getMeasureX() {
    return Radians.of(getX());
  }

  public Angle getMeasureY() {
    return Radians.of(getY());
  }

  public Angle getMeasureZ() {
    return Radians.of(getZ());
  }

  public Vector<N3> getAxis() {
    double norm =
        Math.sqrt(m_q.getX() * m_q.getX() + m_q.getY() * m_q.getY() + m_q.getZ() * m_q.getZ());
    if (norm == 0.0) {
      return VecBuilder.fill(0.0, 0.0, 0.0);
    } else {
      return VecBuilder.fill(m_q.getX() / norm, m_q.getY() / norm, m_q.getZ() / norm);
    }
  }

  public double getAngle() {
    double norm =
        Math.sqrt(m_q.getX() * m_q.getX() + m_q.getY() * m_q.getY() + m_q.getZ() * m_q.getZ());
    return 2.0 * Math.atan2(norm, m_q.getW());
  }

  public Matrix<N3, N3> toMatrix() {
    double w = m_q.getW();
    double x = m_q.getX();
    double y = m_q.getY();
    double z = m_q.getZ();
 
    // https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation#Quaternion-derived_rotation_matrix
    return MatBuilder.fill(
        Nat.N3(),
        Nat.N3(),
        1.0 - 2.0 * (y * y + z * z),
        2.0 * (x * y - w * z),
        2.0 * (x * z + w * y),
        2.0 * (x * y + w * z),
        1.0 - 2.0 * (x * x + z * z),
        2.0 * (y * z - w * x),
        2.0 * (x * z - w * y),
        2.0 * (y * z + w * x),
        1.0 - 2.0 * (x * x + y * y));
  }

  public Vector<N3> toVector() {
    return m_q.toRotationVector();
  }

  public Angle getMeasureAngle() {
    return Radians.of(getAngle());
  }

  public Rotation2d toRotation2d() {
    return new Rotation2d(getZ());
  }
 
  @Override
  public String toString() {
    return String.format("Rotation3d(%s)", m_q);
  }

  @Override
  public boolean equals(Object obj) {
    return obj instanceof Rotation3d other
        && Math.abs(Math.abs(m_q.dot(other.m_q)) - m_q.norm() * other.m_q.norm()) < 1e-9;
  }
 
  @Override
  public int hashCode() {
    return Objects.hash(m_q);
  }
 
  @Override
  public Rotation3d interpolate(Rotation3d endValue, double t) {
    // https://en.wikipedia.org/wiki/Slerp#Quaternion_Slerp
    //
    // slerp(q₀, q₁, t) = (q₁q₀⁻¹)ᵗq₀
    //
    // We negate the delta quaternion if necessary to take the shortest path
    var q0 = m_q;
    var q1 = endValue.m_q;
    var delta = q1.times(q0.inverse());
    if (delta.getW() < 0.0) {
      delta = new Quaternion(-delta.getW(), -delta.getX(), -delta.getY(), -delta.getZ());
    }
    return new Rotation3d(delta.pow(t).times(q0));
  }

  public static final Rotation3dProto proto = new Rotation3dProto();

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