Fix/columns not found for two-button user (#126)

* Fix/Columns not found for two-button users

* format

* bump version

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "5.0.3"
+version = "5.0.4"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",

src/fsrs_optimizer/fsrs_optimizer.py

@@ -576,7 +576,8 @@ def flags2str(flags: List[int]) -> str:
 
     def extract_simulation_config(self, df):
         def rating_counts(x):
-            tmp = defaultdict(int, x.value_counts().to_dict())
+            tmp = {1: 0, 2: 0, 3: 0, 4: 0}
+            tmp.update(x.value_counts().to_dict())
             first = x.iloc[0]
             tmp[first] -= 1
             return tmp
@@ -628,8 +629,12 @@ def rating_counts(x):
                 .to_dict()
             ),
         )
-        self.learn_buttons = np.array([button_usage_dict[(1, i)] for i in range(1, 5)])
-        self.review_buttons = np.array([button_usage_dict[(2, i)] for i in range(1, 5)])
+        self.learn_buttons = (
+            np.array([button_usage_dict[(1, i)] for i in range(1, 5)]) + 1
+        )
+        self.review_buttons = (
+            np.array([button_usage_dict[(2, i)] for i in range(1, 5)]) + 1
+        )
         self.first_rating_prob = self.learn_buttons / self.learn_buttons.sum()
         self.review_rating_prob = (
             self.review_buttons[1:] / self.review_buttons[1:].sum()