From 86eeb694007653b5fcef8e6024871ac819e4a8ae Mon Sep 17 00:00:00 2001
From: Expertium <83031600+Expertium@users.noreply.github.com>
Date: Wed, 24 Jul 2024 14:01:20 +0300
Subject: [PATCH] Increase sample size when "Days to simulate" is low.py (#122)

* Increase sample size when "Days to simulate" is low.py

This is to prevent results from being noisy

* Improve the sample size rule.py

* fix UnboundLocalError

* bump version

---------

Co-authored-by: Jarrett Ye <jarrett.ye@outlook.com>
---
 pyproject.toml                       |  2 +-
 src/fsrs_optimizer/fsrs_simulator.py | 10 +++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index bc10695..da1ce20 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "5.0.1"
+version = "5.0.2"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",
diff --git a/src/fsrs_optimizer/fsrs_simulator.py b/src/fsrs_optimizer/fsrs_simulator.py
index 5f1c18a..20d410c 100644
--- a/src/fsrs_optimizer/fsrs_simulator.py
+++ b/src/fsrs_optimizer/fsrs_simulator.py
@@ -33,9 +33,6 @@ def next_interval(s, r):
 ]
 col = {key: i for i, key in enumerate(columns)}
 
-SAMPLE_SIZE = 4
-
-
 DEFAULT_LEARN_COSTS = np.array([33.79, 24.3, 13.68, 6.5])
 DEFAULT_REVIEW_COSTS = np.array([23.0, 11.68, 7.33, 5.6])
 DEFAULT_FIRST_RATING_PROB = np.array([0.24, 0.094, 0.495, 0.171])
@@ -257,6 +254,13 @@ def sample(
     forget_session_len=DEFAULT_FORGET_SESSION_LEN,
 ):
     memorization = []
+    if learn_span < 100:
+        SAMPLE_SIZE = 16
+    elif learn_span < 365:
+        SAMPLE_SIZE = 8
+    else:
+        SAMPLE_SIZE = 4
+
     for i in range(SAMPLE_SIZE):
         _, _, _, memorized_cnt_per_day, cost_per_day = simulate(
             w,