Fix/S0 dataset initialization fallback

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",

src/fsrs_optimizer/fsrs_optimizer.py

@@ -519,6 +519,7 @@ def loss(stability):
 class Optimizer:
     float_delta_t: bool = False
     enable_short_term: bool = True
+    S0_dataset_group = None
 
     def __init__(
         self, float_delta_t: bool = False, enable_short_term: bool = True
@@ -1007,6 +1008,16 @@ def pretrain(self, dataset=None, verbose=True):
         ]
         if self.dataset.empty:
             raise ValueError("Training data is inadequate.")
+        if self.S0_dataset_group is None:
+            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()
+            )
         rating_stability = {}
         rating_count = {}
         average_recall = self.dataset["y"].mean()