Fix/S0 dataset initialization fallback (#156)

open-spaced-repetition · Jan 2, 2025 · 333651f · 333651f
1 parent 4a48b43
commit 333651f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "5.6.2"
+version = "5.6.3"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",

diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -1002,6 +1002,15 @@ def pretrain(self, dataset=None, verbose=True):
             )
         else:
             self.dataset = dataset
+            self.dataset["first_rating"] = self.dataset["r_history"].map(
+                lambda x: x[0] if len(x) > 0 else ""
+            )
+            self.S0_dataset_group = (
+                self.dataset[self.dataset["i"] == 2]
+                .groupby(by=["first_rating", "delta_t"], group_keys=False)
+                .agg({"y": ["mean", "count"]})
+                .reset_index()
+            )
         self.dataset = self.dataset[
             (self.dataset["i"] > 1) & (self.dataset["delta_t"] > 0)
         ]