ExtendedKalmanFilter.kt

package frc.team4069.saturn.lib.mathematics.statespace

import edu.wpi.first.wpilibj.math.Drake
import edu.wpi.first.wpilibj.math.StateSpaceUtils
import edu.wpi.first.wpiutil.math.Matrix
import edu.wpi.first.wpiutil.math.Nat
import edu.wpi.first.wpiutil.math.Num
import frc.team4069.saturn.lib.mathematics.matrix.Vector
import frc.team4069.saturn.lib.mathematics.matrix.eye
import frc.team4069.saturn.lib.mathematics.matrix.zeros
import frc.team4069.saturn.lib.mathematics.units.SIUnit
import frc.team4069.saturn.lib.mathematics.units.Second

class ExtendedKalmanFilter<States: Num, Inputs: Num, Outputs: Num>(
    val states: Nat<States>, val inputs: Nat<Inputs>, val outputs: Nat<Outputs>,
    val f: (x: Vector<States>, u: Vector<Inputs>) -> Vector<States>,
    val h: (x: Vector<States>, u: Vector<Inputs>) -> Vector<Outputs>,
    stateStdDevs: Vector<States>,
    measurementStdDevs: Vector<Outputs>,
    dt: SIUnit<Second>
) {

    var xHat: Vector<States> = zeros(states)
    private var P: Matrix<States, States>
    private val contQ: Matrix<States, States> = StateSpaceUtils.makeCovMatrix(states, stateStdDevs)
    private val contR: Matrix<Outputs, Outputs> = StateSpaceUtils.makeCovMatrix(outputs, measurementStdDevs)
    private var discR: Matrix<Outputs, Outputs>

    private val initP: Matrix<States, States>

    init {
        reset()

        val contA = numericalJacobianX(states, states, f, xHat, zeros(inputs))
        val C = numericalJacobianX(outputs, states, h, xHat, zeros(inputs))

        val discPair = StateSpaceUtils.discretizeAQTaylor(contA, contQ, dt.value)
        val discA = discPair.first
        val discQ = discPair.second

        discR = StateSpaceUtils.discretizeR(contR, dt.value)

        initP = if(StateSpaceUtils.isStabilizable(discA.transpose(), C.transpose())) {
            Matrix(Drake.discreteAlgebraicRiccatiEquation(discA.transpose(), C.transpose(), discQ, discR))
        } else {
            zeros(states, states)
        }

        P = initP
    }

    fun reset() {
        xHat = zeros(states)
        P = initP
    }

    fun predict(u: Vector<Inputs>, dt: SIUnit<Second>) {
        val contA = numericalJacobianX(states, states, f, xHat, u)

        val discPair = StateSpaceUtils.discretizeAQTaylor(contA, contQ, dt.value)
        val discA = discPair.first
        val discQ = discPair.second

        xHat = rungeKutta(f, xHat, u, dt)
        P = discA * P * discA.transpose() + discQ
        discR = StateSpaceUtils.discretizeR(contR, dt.value)
    }

    fun correct(u: Vector<Inputs>, y: Vector<Outputs>) {
        correct(outputs, u, y, h, discR)
    }

    fun <Rows: Num> correct(rows: Nat<Rows>, u: Vector<Inputs>, y: Vector<Rows>, h: (Vector<States>, Vector<Inputs>) -> Vector<Rows>, R: Matrix<Rows, Rows>) {
        val C = numericalJacobianX(rows, states, h, xHat, u)

        val S = C * P * C.transpose() + R
        val K = Matrix<States, Rows>(StateSpaceUtils.lltDecompose(S.transpose().storage).solve((C * P.transpose()).storage).transpose())
        xHat += K * (y - h(xHat, u))
        P = (eye(states) - K * C) * P
    }
}

NumericalJacobian.kt

package frc.team4069.saturn.lib.mathematics.statespace

import edu.wpi.first.wpiutil.math.Matrix
import edu.wpi.first.wpiutil.math.Nat
import edu.wpi.first.wpiutil.math.Num
import frc.team4069.saturn.lib.mathematics.matrix.Vector
import frc.team4069.saturn.lib.mathematics.matrix.set
import frc.team4069.saturn.lib.mathematics.matrix.zeros
import kotlin.jvm.functions.FunctionN

fun <Rows: Num, Cols: Num, States: Num> numericalJacobian(rows: Nat<Rows>, cols: Nat<Cols>, f: (Vector<Cols>) -> Vector<States>, x: Vector<Cols>,
                                                                              epsilon: Double = 1E-5): Matrix<Rows, Cols> {
    val result = zeros(rows, cols)

    for(i in 0 until cols.num) {
        val dxPlus = x.copy()
        val dxMinus = x.copy()
        dxPlus[i, 0] += epsilon
        dxMinus[i, 0] -= epsilon
        val dF = (f(dxPlus) - f(dxMinus)) / (2 * epsilon)

        result.storage.setColumn(i, 0, *dF.storage.ddrm.data)
    }

    return result
}

fun <Rows: Num, States: Num, Inputs: Num, D: Num> numericalJacobianX(rows: Nat<Rows>, states: Nat<States>, f: (Vector<States>, Vector<Inputs>) -> Vector<D>, x: Vector<States>,
                                                            u: Vector<Inputs>): Matrix<Rows, States> {
    return numericalJacobian(rows, states, { x: Vector<States> -> f(x, u) }, x)
}

fun <Rows: Num, States: Num, Inputs: Num> numericalJacobianU(rows: Nat<Rows>, inputs: Nat<Inputs>, f: (Vector<States>, Vector<Inputs>) -> Vector<States>, x: Vector<States>,
                                                             u: Vector<Inputs>): Matrix<Rows, Inputs> {
    return numericalJacobian(rows, inputs, { u: Vector<Inputs> -> f(x, u) }, u)
}

RungeKutta.kt

package frc.team4069.saturn.lib.mathematics.statespace

import edu.wpi.first.wpiutil.math.Num
import frc.team4069.saturn.lib.mathematics.matrix.Vector
import frc.team4069.saturn.lib.mathematics.units.SIUnit
import frc.team4069.saturn.lib.mathematics.units.Second

private operator fun <N: Num> Double.times(vec: Vector<N>) = vec.times(this)

/**
 * Performs 4th order Runge-Kutta integration of dx/dt = f(x, u) for dt.
 *
 * @param f  The function to integrate. It must take two arguments x and u.
 * @param x  The initial value of x.
 * @param u  The value u held constant over the integration period
 * @param dt The time over which to integrate
 */
fun <States: Num, Inputs: Num> rungeKutta(f: (Vector<States>, Vector<Inputs>) -> Vector<States>, x: Vector<States>, u: Vector<Inputs>, dt: SIUnit<Second>): Vector<States> {
    val halfDt = dt * 0.5

    val k1 = f(x, u)
    val k2 = f(x + k1 * halfDt.value, u)
    val k3 = f(x + k2 * halfDt.value, u)
    val k4 = f(x + k3 * dt.value, u)

    return x + dt.value / 6.0 * (k1 + 2.0 * k2 + 2.0 * k3 + k4)
}