Fix/set mean reversion target to d0(4) (#205)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "fsrs"
-version = "1.1.3"
+version = "1.1.4"
 authors = ["Open Spaced Repetition"]
 categories = ["algorithms", "science"]
 edition = "2021"

src/inference.rs

@@ -22,8 +22,8 @@ pub type Parameters = [f32];
 use itertools::izip;
 
 pub static DEFAULT_PARAMETERS: [f32; 19] = [
-    0.4197, 1.1869, 3.0412, 15.2441, 7.1434, 0.6477, 1.0007, 0.0674, 1.6597, 0.1712, 1.1178,
-    2.0225, 0.0904, 0.3025, 2.1214, 0.2498, 2.9466, 0.4891, 0.6468,
+    0.4072, 1.1829, 3.1262, 15.4722, 7.2102, 0.5316, 1.0651, 0.0234, 1.616, 0.1544, 1.0824, 1.9813,
+    0.0953, 0.2975, 2.2042, 0.2407, 2.9466, 0.5034, 0.6567,
 ];
 
 fn infer<B: Backend>(
@@ -377,9 +377,25 @@ mod tests {
     };
 
     static PARAMETERS: &[f32] = &[
-        0.72466177, 1.6790825, 4.562257, 10.0608635, 7.7002444, 0.912309, 1.0909119, 0.03472257,
-        1.4395499, 0.1712, 0.8977034, 2.1090207, 0.0904, 0.3025, 2.3506782, 0.23486121, 3.1349943,
-        0.18253358, 0.13885707,
+        0.6845422,
+        1.6790825,
+        4.7349424,
+        10.042885,
+        7.4410233,
+        0.64219797,
+        1.071918,
+        0.0025195254,
+        1.432437,
+        0.1544,
+        0.8692766,
+        2.0696752,
+        0.0953,
+        0.2975,
+        2.4691248,
+        0.19542035,
+        3.201072,
+        0.18046261,
+        0.121442534,
     ];
 
     #[test]
@@ -427,8 +443,8 @@ mod tests {
         assert_eq!(
             fsrs.memory_state(item, None).unwrap(),
             MemoryState {
-                stability: 29.448196,
-                difficulty: 7.7002444
+                stability: 31.722975,
+                difficulty: 7.382128
             }
         );
 
@@ -445,8 +461,8 @@ mod tests {
             .good
             .memory,
             MemoryState {
-                stability: 40.669125,
-                difficulty: 7.0292006
+                stability: 40.874535,
+                difficulty: 6.9913807
             }
         );
         Ok(())
@@ -478,26 +494,26 @@ mod tests {
         let metrics = fsrs.evaluate(items.clone(), |_| true).unwrap();
 
         Data::from([metrics.log_loss, metrics.rmse_bins])
-            .assert_approx_eq(&Data::from([0.202137, 0.025207]), 5);
+            .assert_approx_eq(&Data::from([0.210983, 0.037216]), 5);
 
         let fsrs = FSRS::new(Some(&[]))?;
         let metrics = fsrs.evaluate(items.clone(), |_| true).unwrap();
 
         Data::from([metrics.log_loss, metrics.rmse_bins])
-            .assert_approx_eq(&Data::from([0.216539, 0.039671]), 5);
+            .assert_approx_eq(&Data::from([0.217689, 0.039710]), 5);
 
         let fsrs = FSRS::new(Some(PARAMETERS))?;
         let metrics = fsrs.evaluate(items.clone(), |_| true).unwrap();
 
         Data::from([metrics.log_loss, metrics.rmse_bins])
-            .assert_approx_eq(&Data::from([0.202754, 0.026365]), 5);
+            .assert_approx_eq(&Data::from([0.203235, 0.026295]), 5);
 
         let (self_by_other, other_by_self) = fsrs
             .universal_metrics(items.clone(), &DEFAULT_PARAMETERS, |_| true)
             .unwrap();
 
         Data::from([self_by_other, other_by_self])
-            .assert_approx_eq(&Data::from([0.015230, 0.032233]), 5);
+            .assert_approx_eq(&Data::from([0.014476, 0.031874]), 5);
 
         Ok(())
     }
@@ -531,31 +547,31 @@ mod tests {
             NextStates {
                 again: ItemState {
                     memory: MemoryState {
-                        stability: 3.0271175,
-                        difficulty: 9.80631
+                        stability: 2.969144,
+                        difficulty: 9.520562
                     },
                     interval: 3
                 },
                 hard: ItemState {
                     memory: MemoryState {
-                        stability: 16.725859,
-                        difficulty: 8.753278
+                        stability: 17.091442,
+                        difficulty: 8.4513445
                     },
                     interval: 17
                 },
                 good: ItemState {
                     memory: MemoryState {
-                        stability: 29.448196,
-                        difficulty: 7.7002444
+                        stability: 31.722975,
+                        difficulty: 7.382128
                     },
-                    interval: 29
+                    interval: 32
                 },
                 easy: ItemState {
                     memory: MemoryState {
-                        stability: 64.947784,
-                        difficulty: 6.647212
+                        stability: 71.75015,
+                        difficulty: 6.3129106
                     },
-                    interval: 65
+                    interval: 72
                 }
             }
         );
@@ -581,13 +597,13 @@ mod tests {
         let fsrs = FSRS::new(Some(&[]))?;
         let memory_state = fsrs.memory_state_from_sm2(2.5, 10.0, 0.9).unwrap();
         Data::from([memory_state.stability, memory_state.difficulty])
-            .assert_approx_eq(&Data::from([9.999996, 7.422087]), 5);
+            .assert_approx_eq(&Data::from([9.999996, 7.279789]), 5);
         let memory_state = fsrs.memory_state_from_sm2(2.5, 10.0, 0.8).unwrap();
         Data::from([memory_state.stability, memory_state.difficulty])
-            .assert_approx_eq(&Data::from([4.170096, 9.545_82]), 5);
+            .assert_approx_eq(&Data::from([4.170096, 9.462736]), 5);
         let memory_state = fsrs.memory_state_from_sm2(2.5, 10.0, 0.95).unwrap();
         Data::from([memory_state.stability, memory_state.difficulty])
-            .assert_approx_eq(&Data::from([21.712555, 2.593589]), 5);
+            .assert_approx_eq(&Data::from([21.712555, 2.380210]), 5);
         let memory_state = fsrs.memory_state_from_sm2(1.3, 20.0, 0.9).unwrap();
         Data::from([memory_state.stability, memory_state.difficulty])
             .assert_approx_eq(&Data::from([19.999992, 10.0]), 5);

src/model.rs

@@ -92,7 +92,8 @@ impl<B: Backend> Model<B> {
     }
 
     fn mean_reversion(&self, new_d: Tensor<B, 1>) -> Tensor<B, 1> {
-        self.w.get(7) * (self.w.get(4) - new_d.clone()) + new_d
+        let rating = Tensor::from_floats([4.0], &B::Device::default());
+        self.w.get(7) * (self.init_difficulty(rating) - new_d.clone()) + new_d
     }
 
     pub(crate) fn init_stability(&self, rating: Tensor<B, 1>) -> Tensor<B, 1> {
@@ -358,7 +359,7 @@ mod tests {
         next_difficulty.clone().backward();
         assert_eq!(
             next_difficulty.to_data(),
-            Data::from([7.0109706, 6.077718, 5.144465, 4.211212])
+            Data::from([7.040172, 5.999995, 4.959819, 3.9196422])
         )
     }
 
@@ -379,24 +380,24 @@ mod tests {
         s_recall.clone().backward();
         assert_eq!(
             s_recall.to_data(),
-            Data::from([28.603035, 16.240442, 68.610886, 237.08693])
+            Data::from([27.43741, 15.276875, 65.24019, 224.3506])
         );
         let s_forget = model.stability_after_failure(stability.clone(), difficulty, retention);
         s_forget.clone().backward();
         assert_eq!(
             s_forget.to_data(),
-            Data::from([1.7989675, 2.089014, 2.4897401, 2.9990985])
+            Data::from([1.7390966, 2.029377, 2.433932, 2.9520853])
         );
         let next_stability = s_recall.mask_where(rating.clone().equal_elem(1), s_forget);
         next_stability.clone().backward();
         assert_eq!(
             next_stability.to_data(),
-            Data::from([1.7989675, 16.240442, 68.610886, 237.08693])
+            Data::from([1.7390966, 15.276875, 65.24019, 224.3506])
         );
         let next_stability = model.stability_short_term(stability, rating);
         assert_eq!(
             next_stability.to_data(),
-            Data::from([2.5794802, 4.206739, 6.8605514, 11.188516])
+            Data::from([2.542685, 4.2064567, 6.958895, 11.512355])
         )
     }
 

src/optimal_retention.rs

@@ -118,14 +118,18 @@ fn stability_short_term(w: &[f32], s: f32, rating_offset: f32, session_len: f32)
     s * (w[17] * (rating_offset + session_len * w[18])).exp()
 }
 
-fn init_d(w: &[f32], rating: usize, rating_offset: f32) -> f32 {
-    let new_d = w[4] - (w[5] * (rating - 1) as f32).exp() + 1.0 - w[6] * rating_offset;
+fn init_d(w: &[f32], rating: usize) -> f32 {
+    w[4] - (w[5] * (rating - 1) as f32).exp() + 1.0
+}
+
+fn init_d_with_short_term(w: &[f32], rating: usize, rating_offset: f32) -> f32 {
+    let new_d = init_d(w, rating) - w[6] * rating_offset;
     new_d.clamp(1.0, 10.0)
 }
 
 fn next_d(w: &[f32], d: f32, rating: usize) -> f32 {
     let new_d = d - w[6] * (rating as f32 - 3.0);
-    mean_reversion(w, w[4], new_d).clamp(1.0, 10.0)
+    mean_reversion(w, init_d(w, 4), new_d).clamp(1.0, 10.0)
 }
 
 fn mean_reversion(w: &[f32], init: f32, current: f32) -> f32 {
@@ -394,7 +398,7 @@ pub fn simulate(
                 first_rating_offsets[rating - 1],
                 first_session_lens[rating - 1],
             );
-            *new_diff = init_d(w, rating, first_rating_offsets[rating - 1]);
+            *new_diff = init_d_with_short_term(w, rating, first_rating_offsets[rating - 1]);
         });
         let old_interval = card_table.slice(s![Column::Interval, ..]);
         let mut new_interval = old_interval.to_owned();
@@ -1004,7 +1008,7 @@ mod tests {
             simulate(&config, &DEFAULT_PARAMETERS, 0.9, None, None);
         assert_eq!(
             memorized_cnt_per_day[memorized_cnt_per_day.len() - 1],
-            7020.523
+            6521.068
         )
     }
 
@@ -1048,8 +1052,8 @@ mod tests {
         assert_eq!(
             results.1.to_vec(),
             vec![
-                0, 15, 16, 17, 70, 73, 77, 82, 75, 87, 86, 111, 113, 110, 105, 112, 124, 131, 127,
-                119, 122, 163, 145, 150, 171, 150, 136, 163, 167, 156
+                0, 15, 16, 27, 69, 66, 81, 83, 75, 85, 90, 102, 107, 105, 122, 120, 141, 117, 132,
+                153, 153, 154, 152, 152, 149, 157, 166, 165, 184, 174
             ]
         );
         assert_eq!(
@@ -1071,7 +1075,7 @@ mod tests {
             ..Default::default()
         };
         let optimal_retention = fsrs.optimal_retention(&config, &[], |_v| true).unwrap();
-        assert_eq!(optimal_retention, 0.8263932);
+        assert_eq!(optimal_retention, 0.7921062);
         assert!(fsrs.optimal_retention(&config, &[1.], |_v| true).is_err());
         Ok(())
     }
@@ -1091,7 +1095,7 @@ mod tests {
         let optimal_retention = fsrs
             .optimal_retention(&config, &DEFAULT_PARAMETERS[..17], |_v| true)
             .unwrap();
-        assert_eq!(optimal_retention, 0.81812924);
+        assert_eq!(optimal_retention, 0.8430037);
         Ok(())
     }
 

src/pre_training.rs

@@ -296,9 +296,9 @@ mod tests {
         let count = Array1::from(vec![435.0, 97.0, 63.0, 38.0, 28.0]);
         let default_s0 = DEFAULT_PARAMETERS[0] as f64;
         let actual = loss(&delta_t, &recall, &count, 1.017056, default_s0);
-        assert_eq!(actual, 280.7489989949864);
+        assert_eq!(actual, 280.7497802442413);
         let actual = loss(&delta_t, &recall, &count, 1.017011, default_s0);
-        assert_eq!(actual, 280.74866497463466);
+        assert_eq!(actual, 280.7494462238896);
     }
 
     #[test]
@@ -363,6 +363,6 @@ mod tests {
         let mut rating_stability = HashMap::from([(2, 0.35)]);
         let rating_count = HashMap::from([(2, 1)]);
         let actual = smooth_and_fill(&mut rating_stability, &rating_count).unwrap();
-        assert_eq!(actual, [0.123763576, 0.34999996, 0.8968067, 4.495269]);
+        assert_eq!(actual, [0.12048356, 0.35, 0.9249894, 4.577961]);
     }
 }