fix weight & consider short term schedule in d

src/fsrs_optimizer/fsrs_optimizer.py

@@ -630,10 +630,10 @@ def rating_counts(x):
             [button_usage_dict.get((1, i), 0) for i in range(1, 5)]
         )
         self.review_buttons = np.array(
-            [button_usage_dict.get((2, i), 0) for i in range(2, 5)]
+            [button_usage_dict.get((2, i), 0) for i in range(1, 5)]
         )
         self.first_rating_prob = self.learn_buttons / self.learn_buttons.sum()
-        self.review_rating_prob = self.review_buttons / self.review_buttons.sum()
+        self.review_rating_prob = self.review_buttons[1:] / self.review_buttons[1:].sum()
 
         df2 = (
             df1.groupby(by=["first_review_state", "first_review_rating"])[[1, 2, 3, 4]]
@@ -1313,8 +1313,8 @@ def find_optimal_retention(
         )
 
         weight = self.review_buttons / (50 + self.review_buttons)
-        self.review_costs = self.review_costs * weight[1:] + DEFAULT_REVIEW_COSTS * (
-            1 - weight[1:]
+        self.review_costs = self.review_costs * weight + DEFAULT_REVIEW_COSTS * (
+            1 - weight
         )
         self.forget_rating_offset = self.forget_rating_offset * weight[
             0
@@ -1328,7 +1328,7 @@ def find_optimal_retention(
             self.first_rating_prob * weight + DEFAULT_FIRST_RATING_PROB * (1 - weight)
         )
 
-        weight = sum(self.review_buttons) / (50 + sum(self.review_buttons))
+        weight = sum(self.review_buttons[1:]) / (50 + sum(self.review_buttons[1:]))
         self.review_rating_prob = (
             self.review_rating_prob * weight + DEFAULT_REVIEW_RATING_PROB * (1 - weight)
         )

src/fsrs_optimizer/fsrs_simulator.py

@@ -116,6 +116,12 @@ def stability_short_term(s, init_rating=None):
         new_s = s * np.exp(w[17] * (rating_offset + session_len * w[18]))
         return new_s
 
+    def init_d(rating):
+        new_d = w[4] - np.exp(w[5] * (rating - 1)) + 1
+        rating_offset = np.choose(rating - 1, first_rating_offset)
+        new_d -= w[6] * rating_offset
+        return np.clip(new_d, 1, 10)
+
     def next_d(d, rating):
         new_d = d - w[6] * (rating - 3)
         new_d = mean_reversion(w[4], new_d)
@@ -174,6 +180,9 @@ def mean_reversion(init, current):
             card_table[col["difficulty"]][true_review],
             card_table[col["rating"]][true_review],
         )
+        card_table[col["difficulty"]][true_review & forget] -= (
+            w[6] * forget_rating_offset
+        )
 
         need_learn = card_table[col["due"]] == learn_span
         card_table[col["cost"]][need_learn] = min(learn_costs)
@@ -200,12 +209,8 @@ def mean_reversion(init, current):
             card_table[col["stability"]][true_learn],
             init_rating=card_table[col["rating"]][true_learn].astype(int),
         )
-        card_table[col["difficulty"]][true_learn] = np.clip(
-            w[4]
-            - np.exp(w[5] * (card_table[col["rating"]][true_learn].astype(int) - 1))
-            + 1,
-            1,
-            10,
+        card_table[col["difficulty"]][true_learn] = init_d(
+            card_table[col["rating"]][true_learn].astype(int)
         )
 
         card_table[col["ivl"]][true_review | true_learn] = np.clip(