Add smoothing to jerk calculations and updated jerk thresholds.

Test: atest libinput_tests
Test: atest CtsInputTestCases
Test: atest MotionPredictorBenchmark MotionPredictorTest
Test: Using stylus in a drawing app and seeing the jerk logs.
Bug: 266747654
Bug: 353161308
Flag: com.android.input.flags.enable_prediction_pruning_via_jerk_thresholding
Change-Id: I3d6c47d94d66e5ff2b33474acbca72daca051242
diff --git a/libs/input/MotionPredictor.cpp b/libs/input/MotionPredictor.cpp
index 5b61d39..9204b95 100644
--- a/libs/input/MotionPredictor.cpp
+++ b/libs/input/MotionPredictor.cpp
@@ -75,6 +75,9 @@
 JerkTracker::JerkTracker(bool normalizedDt) : mNormalizedDt(normalizedDt) {}
 
 void JerkTracker::pushSample(int64_t timestamp, float xPos, float yPos) {
+    // If we previously had full samples, we have a previous jerk calculation
+    // to do weighted smoothing.
+    const bool applySmoothing = mTimestamps.size() == mTimestamps.capacity();
     mTimestamps.pushBack(timestamp);
     const int numSamples = mTimestamps.size();
 
@@ -115,6 +118,16 @@
         }
     }
 
+    if (numSamples == static_cast<int>(mTimestamps.capacity())) {
+        float newJerkMagnitude = std::hypot(newXDerivatives[3], newYDerivatives[3]);
+        ALOGD_IF(isDebug(), "raw jerk: %f", newJerkMagnitude);
+        if (applySmoothing) {
+            mJerkMagnitude = mJerkMagnitude + (mForgetFactor * (newJerkMagnitude - mJerkMagnitude));
+        } else {
+            mJerkMagnitude = newJerkMagnitude;
+        }
+    }
+
     std::swap(newXDerivatives, mXDerivatives);
     std::swap(newYDerivatives, mYDerivatives);
 }
@@ -125,11 +138,19 @@
 
 std::optional<float> JerkTracker::jerkMagnitude() const {
     if (mTimestamps.size() == mTimestamps.capacity()) {
-        return std::hypot(mXDerivatives[3], mYDerivatives[3]);
+        return mJerkMagnitude;
     }
     return std::nullopt;
 }
 
+void JerkTracker::setForgetFactor(float forgetFactor) {
+    mForgetFactor = forgetFactor;
+}
+
+float JerkTracker::getForgetFactor() const {
+    return mForgetFactor;
+}
+
 // --- MotionPredictor ---
 
 MotionPredictor::MotionPredictor(nsecs_t predictionTimestampOffsetNanos,
@@ -159,6 +180,7 @@
     if (!mModel) {
         mModel = TfLiteMotionPredictorModel::create();
         LOG_ALWAYS_FATAL_IF(!mModel);
+        mJerkTracker.setForgetFactor(mModel->config().jerkForgetFactor);
     }
 
     if (!mBuffers) {
@@ -357,4 +379,12 @@
     return true;
 }
 
+const TfLiteMotionPredictorModel::Config& MotionPredictor::getModelConfig() {
+    if (!mModel) {
+        mModel = TfLiteMotionPredictorModel::create();
+        LOG_ALWAYS_FATAL_IF(!mModel);
+    }
+    return mModel->config();
+}
+
 } // namespace android
diff --git a/libs/input/TfLiteMotionPredictor.cpp b/libs/input/TfLiteMotionPredictor.cpp
index b843a4b..b401c98 100644
--- a/libs/input/TfLiteMotionPredictor.cpp
+++ b/libs/input/TfLiteMotionPredictor.cpp
@@ -283,6 +283,7 @@
             .distanceNoiseFloor = parseXMLFloat(*configRoot, "distance-noise-floor"),
             .lowJerk = parseXMLFloat(*configRoot, "low-jerk"),
             .highJerk = parseXMLFloat(*configRoot, "high-jerk"),
+            .jerkForgetFactor = parseXMLFloat(*configRoot, "jerk-forget-factor"),
     };
 
     return std::unique_ptr<TfLiteMotionPredictorModel>(
diff --git a/libs/input/tests/MotionPredictor_test.cpp b/libs/input/tests/MotionPredictor_test.cpp
index d077760..5bd5794 100644
--- a/libs/input/tests/MotionPredictor_test.cpp
+++ b/libs/input/tests/MotionPredictor_test.cpp
@@ -88,6 +88,7 @@
 
 TEST(JerkTrackerTest, JerkCalculationNormalizedDtTrue) {
     JerkTracker jerkTracker(true);
+    jerkTracker.setForgetFactor(.5);
     jerkTracker.pushSample(/*timestamp=*/0, 20, 50);
     jerkTracker.pushSample(/*timestamp=*/1, 25, 53);
     jerkTracker.pushSample(/*timestamp=*/2, 30, 60);
@@ -118,11 +119,14 @@
      * y'':  3 -> -15
      * y''': -18
      */
-    EXPECT_FLOAT_EQ(jerkTracker.jerkMagnitude().value(), std::hypot(-50, -18));
+    const float newJerk = (1 - jerkTracker.getForgetFactor()) * std::hypot(10, -1) +
+            jerkTracker.getForgetFactor() * std::hypot(-50, -18);
+    EXPECT_FLOAT_EQ(jerkTracker.jerkMagnitude().value(), newJerk);
 }
 
 TEST(JerkTrackerTest, JerkCalculationNormalizedDtFalse) {
     JerkTracker jerkTracker(false);
+    jerkTracker.setForgetFactor(.5);
     jerkTracker.pushSample(/*timestamp=*/0, 20, 50);
     jerkTracker.pushSample(/*timestamp=*/10, 25, 53);
     jerkTracker.pushSample(/*timestamp=*/20, 30, 60);
@@ -153,7 +157,9 @@
      * y'':  .03 -> -.125 (delta above, divide by 10)
      * y''': -.0155 (delta above, divide by 10)
      */
-    EXPECT_FLOAT_EQ(jerkTracker.jerkMagnitude().value(), std::hypot(-.0375, -.0155));
+    const float newJerk = (1 - jerkTracker.getForgetFactor()) * std::hypot(.01, -.001) +
+            jerkTracker.getForgetFactor() * std::hypot(-.0375, -.0155);
+    EXPECT_FLOAT_EQ(jerkTracker.jerkMagnitude().value(), newJerk);
 }
 
 TEST(JerkTrackerTest, JerkCalculationAfterReset) {
@@ -291,15 +297,19 @@
     MotionPredictor predictor(/*predictionTimestampOffsetNanos=*/0,
                               []() { return true /*enable prediction*/; });
 
-    // Jerk is medium (1.05 normalized, which is halfway between LOW_JANK and HIGH_JANK)
-    predictor.record(getMotionEvent(DOWN, 0, 5.2, 20ms));
-    predictor.record(getMotionEvent(MOVE, 0, 11.5, 30ms));
-    predictor.record(getMotionEvent(MOVE, 0, 22, 40ms));
-    predictor.record(getMotionEvent(MOVE, 0, 37.75, 50ms));
-    predictor.record(getMotionEvent(MOVE, 0, 59.8, 60ms));
+    const float mediumJerk =
+            (predictor.getModelConfig().lowJerk + predictor.getModelConfig().highJerk) / 2;
+    const float a = 3; // initial acceleration
+    const float b = 4; // initial velocity
+    const float c = 5; // initial position
+    predictor.record(getMotionEvent(DOWN, 0, c, 20ms));
+    predictor.record(getMotionEvent(MOVE, 0, c + b, 30ms));
+    predictor.record(getMotionEvent(MOVE, 0, c + 2 * b + a, 40ms));
+    predictor.record(getMotionEvent(MOVE, 0, c + 3 * b + 3 * a + mediumJerk, 50ms));
+    predictor.record(getMotionEvent(MOVE, 0, c + 4 * b + 6 * a + 4 * mediumJerk, 60ms));
     std::unique_ptr<MotionEvent> predicted = predictor.predict(82 * NSEC_PER_MSEC);
     EXPECT_NE(nullptr, predicted);
-    // Halfway between LOW_JANK and HIGH_JANK means that half of the predictions
+    // Halfway between LOW_JERK and HIGH_JERK means that half of the predictions
     // will be pruned. If model prediction window is close enough to predict()
     // call time window, then half of the model predictions (5/2 -> 2) will be
     // ouputted.