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 org.wpilib.math.autodiff;
006
007import java.util.function.BiFunction;
008
009/** Numerical integration utilities. */
010public final class NumericalIntegration {
011  private NumericalIntegration() {
012    // Utility class.
013  }
014
015  /**
016   * Performs 4th order Runge-Kutta integration of dx/dt = f(x, u) for dt.
017   *
018   * @param f The function to integrate. It must take two arguments x and u.
019   * @param x The initial value of x.
020   * @param u The value u held constant over the integration period.
021   * @param dt The time over which to integrate.
022   * @return the integration of dx/dt = f(x, u) for dt.
023   */
024  public static VariableMatrix rk4(
025      BiFunction<VariableMatrix, VariableMatrix, VariableMatrix> f,
026      VariableBlock x,
027      VariableBlock u,
028      double dt) {
029    return rk4(f, new VariableMatrix(x), new VariableMatrix(u), dt);
030  }
031
032  /**
033   * Performs 4th order Runge-Kutta integration of dx/dt = f(x, u) for dt.
034   *
035   * @param f The function to integrate. It must take two arguments x and u.
036   * @param x The initial value of x.
037   * @param u The value u held constant over the integration period.
038   * @param dt The time over which to integrate.
039   * @return the integration of dx/dt = f(x, u) for dt.
040   */
041  public static VariableMatrix rk4(
042      BiFunction<VariableMatrix, VariableMatrix, VariableMatrix> f,
043      VariableBlock x,
044      VariableMatrix u,
045      double dt) {
046    return rk4(f, new VariableMatrix(x), u, dt);
047  }
048
049  /**
050   * Performs 4th order Runge-Kutta integration of dx/dt = f(x, u) for dt.
051   *
052   * @param f The function to integrate. It must take two arguments x and u.
053   * @param x The initial value of x.
054   * @param u The value u held constant over the integration period.
055   * @param dt The time over which to integrate.
056   * @return the integration of dx/dt = f(x, u) for dt.
057   */
058  public static VariableMatrix rk4(
059      BiFunction<VariableMatrix, VariableMatrix, VariableMatrix> f,
060      VariableMatrix x,
061      VariableBlock u,
062      double dt) {
063    return rk4(f, x, new VariableMatrix(u), dt);
064  }
065
066  /**
067   * Performs 4th order Runge-Kutta integration of dx/dt = f(x, u) for dt.
068   *
069   * @param f The function to integrate. It must take two arguments x and u.
070   * @param x The initial value of x.
071   * @param u The value u held constant over the integration period.
072   * @param dt The time over which to integrate.
073   * @return the integration of dx/dt = f(x, u) for dt.
074   */
075  public static VariableMatrix rk4(
076      BiFunction<VariableMatrix, VariableMatrix, VariableMatrix> f,
077      VariableMatrix x,
078      VariableMatrix u,
079      double dt) {
080    var h = dt;
081
082    var k1 = f.apply(x, u);
083    var k2 = f.apply(x.plus(k1.times(h * 0.5)), u);
084    var k3 = f.apply(x.plus(k2.times(h * 0.5)), u);
085    var k4 = f.apply(x.plus(k3.times(h)), u);
086
087    return x.plus(k1.plus(k2.times(2.0)).plus(k3.times(2.0)).plus(k4).times(h / 6.0));
088  }
089}