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.interpolation;
006
007import java.util.TreeMap;
008import org.wpilib.math.linalg.Matrix;
009import org.wpilib.math.util.Num;
010
011/**
012 * Interpolating Tree Maps are used to get values at points that are not defined by making a guess
013 * from points that are defined. This uses linear interpolation.
014 *
015 * @param <K> Key type.
016 * @param <R> Number of matrix rows.
017 * @param <C> Number of matrix columns.
018 */
019public class InterpolatingMatrixTreeMap<K extends Number, R extends Num, C extends Num> {
020  private final TreeMap<K, Matrix<R, C>> m_map = new TreeMap<>();
021
022  /** Default constructor. */
023  public InterpolatingMatrixTreeMap() {}
024
025  /**
026   * Inserts a key-value pair.
027   *
028   * @param key The key.
029   * @param value The value.
030   */
031  public void put(K key, Matrix<R, C> value) {
032    m_map.put(key, value);
033  }
034
035  /**
036   * Returns the value associated with a given key.
037   *
038   * <p>If there's no matching key, the value returned will be a linear interpolation between the
039   * keys before and after the provided one.
040   *
041   * @param key The key.
042   * @return The value associated with the given key.
043   */
044  public Matrix<R, C> get(K key) {
045    Matrix<R, C> val = m_map.get(key);
046    if (val == null) {
047      K ceilingKey = m_map.ceilingKey(key);
048      K floorKey = m_map.floorKey(key);
049
050      if (ceilingKey == null && floorKey == null) {
051        return null;
052      }
053      if (ceilingKey == null) {
054        return m_map.get(floorKey);
055      }
056      if (floorKey == null) {
057        return m_map.get(ceilingKey);
058      }
059      Matrix<R, C> floor = m_map.get(floorKey);
060      Matrix<R, C> ceiling = m_map.get(ceilingKey);
061
062      return interpolate(floor, ceiling, inverseInterpolate(ceilingKey, key, floorKey));
063    } else {
064      return val;
065    }
066  }
067
068  /**
069   * Return the value interpolated between val1 and val2 by the interpolant d.
070   *
071   * @param val1 The lower part of the interpolation range.
072   * @param val2 The upper part of the interpolation range.
073   * @param d The interpolant in the range [0, 1].
074   * @return The interpolated value.
075   */
076  public Matrix<R, C> interpolate(Matrix<R, C> val1, Matrix<R, C> val2, double d) {
077    var dydx = val2.minus(val1);
078    return dydx.times(d).plus(val1);
079  }
080
081  /**
082   * Return where within interpolation range [0, 1] q is between down and up.
083   *
084   * @param up Upper part of interpolation range.
085   * @param q Query.
086   * @param down Lower part of interpolation range.
087   * @return Interpolant in range [0, 1].
088   */
089  public double inverseInterpolate(K up, K q, K down) {
090    double upperToLower = up.doubleValue() - down.doubleValue();
091    if (upperToLower <= 0) {
092      return 0.0;
093    }
094    double queryToLower = q.doubleValue() - down.doubleValue();
095    if (queryToLower <= 0) {
096      return 0.0;
097    }
098    return queryToLower / upperToLower;
099  }
100}