format & fix workload graph

open-spaced-repetition · Jul 10, 2024 · dbd7214 · dbd7214
1 parent 6450519
commit dbd7214
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 28 deletions.
diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -569,7 +569,7 @@ def flags2str(flags: List[int]) -> str:
         df.to_csv("revlog.csv", index=False)
         tqdm.write("revlog.csv saved.")
 
-    def extract_simulation_config(self, df, verbose: bool = True):
+    def extract_simulation_config(self, df):
         def rating_counts(x):
             tmp = x.value_counts().to_dict()
             first = x.iloc[0]
@@ -633,7 +633,9 @@ def rating_counts(x):
             [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[1:] / self.review_buttons[1:].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]]
@@ -651,18 +653,6 @@ def rating_counts(x):
         self.forget_rating_offset = rating_offset_dict.get((2, 1), 0)
         self.forget_session_len = session_len_dict.get((2, 1), 0)
 
-        if verbose:
-            print("Learn costs: ", self.learn_costs)
-            print("Review costs: ", self.review_costs)
-            print("Learn buttons: ", self.learn_buttons)
-            print("Review buttons: ", self.review_buttons)
-            print("First rating prob: ", self.first_rating_prob)
-            print("Review rating prob: ", self.review_rating_prob)
-            print("First rating offset: ", self.first_rating_offset)
-            print("First session len: ", self.first_session_len)
-            print("Forget rating offset: ", self.forget_rating_offset)
-            print("Forget session len: ", self.forget_session_len)
-
     def create_time_series(
         self,
         timezone: str,
@@ -1299,6 +1289,17 @@ def find_optimal_retention(
         verbose=True,
     ):
         """should not be called before predict_memory_states"""
+        if verbose:
+            print("Learn costs: ", self.learn_costs)
+            print("Review costs: ", self.review_costs)
+            print("Learn buttons: ", self.learn_buttons)
+            print("Review buttons: ", self.review_buttons)
+            print("First rating prob: ", self.first_rating_prob)
+            print("Review rating prob: ", self.review_rating_prob)
+            print("First rating offset: ", self.first_rating_offset)
+            print("First session len: ", self.first_session_len)
+            print("Forget rating offset: ", self.forget_rating_offset)
+            print("Forget session len: ", self.forget_session_len)
 
         weight = self.learn_buttons / (50 + self.learn_buttons)
         self.learn_costs = self.learn_costs * weight + DEFAULT_LEARN_COSTS * (
@@ -1415,7 +1416,8 @@ def moving_average(data, window_size=365 // 20):
         ax.legend()
         ax.grid(True)
 
-        simulate_config["max_cost_perday"] = 1800
+        simulate_config["deck_size"] = 20000
+        simulate_config["max_cost_perday"] = 1200
         simulate_config["learn_limit_perday"] = math.inf
         fig6 = workload_graph(simulate_config)
 

diff --git a/src/fsrs_optimizer/fsrs_simulator.py b/src/fsrs_optimizer/fsrs_simulator.py
@@ -70,6 +70,9 @@ def simulate(
     card_table[col["due"]] = learn_span
     card_table[col["difficulty"]] = 1e-10
     card_table[col["stability"]] = 1e-10
+    card_table[col["rating"]] = np.random.choice(
+        [1, 2, 3, 4], deck_size, p=first_rating_prob
+    )
     card_table[col["rating"]] = card_table[col["rating"]].astype(int)
 
     review_cnt_per_day = np.zeros(learn_span)
@@ -185,22 +188,15 @@ def mean_reversion(init, current):
         )
 
         need_learn = card_table[col["due"]] == learn_span
-        card_table[col["cost"]][need_learn] = min(learn_costs)
-        plan_learn = (
+        card_table[col["cost"]][need_learn] = np.choose(
+            card_table[col["rating"]][need_learn].astype(int) - 1,
+            learn_costs,
+        )
+        true_learn = (
             need_learn
             & (np.cumsum(card_table[col["cost"]]) <= max_cost_perday)
             & (np.cumsum(need_learn) <= learn_limit_perday)
         )
-        card_table[col["rating"]][plan_learn] = np.random.choice(
-            [1, 2, 3, 4], np.sum(plan_learn), p=first_rating_prob
-        )
-        card_table[col["cost"]][plan_learn] = np.choose(
-            card_table[col["rating"]][plan_learn].astype(int) - 1,
-            learn_costs,
-        )
-        true_learn = plan_learn & (
-            np.cumsum(card_table[col["cost"]]) <= max_cost_perday
-        )
         card_table[col["last_date"]][true_learn] = today
         card_table[col["stability"]][true_learn] = np.choose(
             card_table[col["rating"]][true_learn].astype(int) - 1, w[:4]
@@ -668,7 +664,7 @@ def workload_graph(default_params):
             0.4891,
             0.6468,
         ],
-        "deck_size": 10000,
+        "deck_size": 20000,
         "learn_span": 365,
         "max_cost_perday": 1800,
         "learn_limit_perday": math.inf,