001// Copyright (c) FIRST and other WPILib contributors.
002// Open Source Software; you can modify and/or share it under the terms of
003// the WPILib BSD license file in the root directory of this project.
004
005package edu.wpi.first.math.estimator;
006
007import edu.wpi.first.math.MathUtil;
008import edu.wpi.first.math.Matrix;
009import edu.wpi.first.math.Num;
010import edu.wpi.first.math.numbers.N1;
011import java.util.function.BiFunction;
012import org.ejml.simple.SimpleMatrix;
013
014public final class AngleStatistics {
015  private AngleStatistics() {
016    // Utility class
017  }
018
019  /**
020   * Subtracts a and b while normalizing the resulting value in the selected row as if it were an
021   * angle.
022   *
023   * @param <S> Number of rows in vector.
024   * @param a A vector to subtract from.
025   * @param b A vector to subtract with.
026   * @param angleStateIdx The row containing angles to be normalized.
027   * @return Difference of two vectors with angle at the given index normalized.
028   */
029  public static <S extends Num> Matrix<S, N1> angleResidual(
030      Matrix<S, N1> a, Matrix<S, N1> b, int angleStateIdx) {
031    Matrix<S, N1> ret = a.minus(b);
032    ret.set(angleStateIdx, 0, MathUtil.angleModulus(ret.get(angleStateIdx, 0)));
033
034    return ret;
035  }
036
037  /**
038   * Returns a function that subtracts two vectors while normalizing the resulting value in the
039   * selected row as if it were an angle.
040   *
041   * @param <S> Number of rows in vector.
042   * @param angleStateIdx The row containing angles to be normalized.
043   * @return Function returning difference of two vectors with angle at the given index normalized.
044   */
045  public static <S extends Num>
046      BiFunction<Matrix<S, N1>, Matrix<S, N1>, Matrix<S, N1>> angleResidual(int angleStateIdx) {
047    return (a, b) -> angleResidual(a, b, angleStateIdx);
048  }
049
050  /**
051   * Adds a and b while normalizing the resulting value in the selected row as an angle.
052   *
053   * @param <S> Number of rows in vector.
054   * @param a A vector to add with.
055   * @param b A vector to add with.
056   * @param angleStateIdx The row containing angles to be normalized.
057   * @return Sum of two vectors with angle at the given index normalized.
058   */
059  public static <S extends Num> Matrix<S, N1> angleAdd(
060      Matrix<S, N1> a, Matrix<S, N1> b, int angleStateIdx) {
061    Matrix<S, N1> ret = a.plus(b);
062    ret.set(angleStateIdx, 0, MathUtil.angleModulus(ret.get(angleStateIdx, 0)));
063
064    return ret;
065  }
066
067  /**
068   * Returns a function that adds two vectors while normalizing the resulting value in the selected
069   * row as an angle.
070   *
071   * @param <S> Number of rows in vector.
072   * @param angleStateIdx The row containing angles to be normalized.
073   * @return Function returning of two vectors with angle at the given index normalized.
074   */
075  public static <S extends Num> BiFunction<Matrix<S, N1>, Matrix<S, N1>, Matrix<S, N1>> angleAdd(
076      int angleStateIdx) {
077    return (a, b) -> angleAdd(a, b, angleStateIdx);
078  }
079
080  /**
081   * Computes the mean of sigmas with the weights Wm while computing a special angle mean for a
082   * select row.
083   *
084   * @param <S> Number of rows in sigma point matrix.
085   * @param sigmas Sigma points.
086   * @param Wm Weights for the mean.
087   * @param angleStateIdx The row containing the angles.
088   * @return Mean of sigma points.
089   */
090  @SuppressWarnings("checkstyle:ParameterName")
091  public static <S extends Num> Matrix<S, N1> angleMean(
092      Matrix<S, ?> sigmas, Matrix<?, N1> Wm, int angleStateIdx) {
093    double[] angleSigmas = sigmas.extractRowVector(angleStateIdx).getData();
094    Matrix<N1, ?> sinAngleSigmas = new Matrix<>(new SimpleMatrix(1, sigmas.getNumCols()));
095    Matrix<N1, ?> cosAngleSigmas = new Matrix<>(new SimpleMatrix(1, sigmas.getNumCols()));
096    for (int i = 0; i < angleSigmas.length; i++) {
097      sinAngleSigmas.set(0, i, Math.sin(angleSigmas[i]));
098      cosAngleSigmas.set(0, i, Math.cos(angleSigmas[i]));
099    }
100
101    double sumSin = sinAngleSigmas.times(Matrix.changeBoundsUnchecked(Wm)).elementSum();
102    double sumCos = cosAngleSigmas.times(Matrix.changeBoundsUnchecked(Wm)).elementSum();
103
104    Matrix<S, N1> ret = sigmas.times(Matrix.changeBoundsUnchecked(Wm));
105    ret.set(angleStateIdx, 0, Math.atan2(sumSin, sumCos));
106
107    return ret;
108  }
109
110  /**
111   * Returns a function that computes the mean of sigmas with the weights Wm while computing a
112   * special angle mean for a select row.
113   *
114   * @param <S> Number of rows in sigma point matrix.
115   * @param angleStateIdx The row containing the angles.
116   * @return Function returning mean of sigma points.
117   */
118  @SuppressWarnings("LambdaParameterName")
119  public static <S extends Num> BiFunction<Matrix<S, ?>, Matrix<?, N1>, Matrix<S, N1>> angleMean(
120      int angleStateIdx) {
121    return (sigmas, Wm) -> angleMean(sigmas, Wm, angleStateIdx);
122  }
123}