Feat/linear damping && w[7] >= 0.001 (#143)

L-M-Sherlock · web-flow · commit 403235260980 · 2024-10-22T10:46:20.000Z
* Feat/linear damping

* update default parameters

* update formula of simulator

* add more tests

* add test for loss and grad

* update ParameterClipper

* bump version
diff --git a/pyproject.toml b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "FSRS-Optimizer"
-version = "5.2.1"
+version = "5.2.2"
 readme = "README.md"
 dependencies = [
     "matplotlib>=3.7.0",
diff --git a/src/fsrs_optimizer/fsrs_optimizer.py b/src/fsrs_optimizer/fsrs_optimizer.py
@@ -5,7 +5,7 @@
 import numpy as np
 import os
 import math
-from typing import List, Optional
+from typing import List, Optional, Tuple
 from datetime import timedelta, datetime
 from collections import defaultdict
 import statsmodels.api as sm  # type: ignore
@@ -42,25 +42,25 @@
 Relearning = 3
 
 DEFAULT_PARAMETER = [
-    0.4072,
-    1.1829,
-    3.1262,
-    15.4722,
-    7.2102,
-    0.5316,
-    1.0651,
-    0.0234,
-    1.616,
-    0.1544,
-    1.0824,
-    1.9813,
-    0.0953,
-    0.2975,
-    2.2042,
-    0.2407,
-    2.9466,
-    0.5034,
-    0.6567,
+    0.40255,
+    1.18385,
+    3.173,
+    15.69105,
+    7.1949,
+    0.5345,
+    1.4604,
+    0.0046,
+    1.54575,
+    0.1192,
+    1.01925,
+    1.9395,
+    0.11,
+    0.29605,
+    2.2698,
+    0.2315,
+    2.9898,
+    0.51655,
+    0.6621,
 ]
 
 S_MIN = 0.01
@@ -105,8 +105,12 @@ def init_d(self, rating: Tensor) -> Tensor:
         new_d = self.w[4] - torch.exp(self.w[5] * (rating - 1)) + 1
         return new_d
 
+    def linear_damping(self, delta_d: Tensor, old_d: Tensor) -> Tensor:
+        return delta_d * (10 - old_d) / 9
+
     def next_d(self, state: Tensor, rating: Tensor) -> Tensor:
-        new_d = state[:, 1] - self.w[6] * (rating - 3)
+        delta_d = -self.w[6] * (rating - 3)
+        new_d = state[:, 1] + self.linear_damping(delta_d, state[:, 1])
         new_d = self.mean_reversion(self.init_d(4), new_d)
         return new_d
 
@@ -151,7 +155,9 @@ def step(self, X: Tensor, state: Tensor) -> Tensor:
         new_s = new_s.clamp(S_MIN, 36500)
         return torch.stack([new_s, new_d], dim=1)
 
-    def forward(self, inputs: Tensor, state: Optional[Tensor] = None) -> Tensor:
+    def forward(
+        self, inputs: Tensor, state: Optional[Tensor] = None
+    ) -> Tuple[Tensor, Tensor]:
         """
         :param inputs: shape[seq_len, batch_size, 2]
         """
@@ -179,16 +185,16 @@ def __call__(self, module):
             w[2] = w[2].clamp(S_MIN, 100)
             w[3] = w[3].clamp(S_MIN, 100)
             w[4] = w[4].clamp(1, 10)
-            w[5] = w[5].clamp(0.01, 4)
-            w[6] = w[6].clamp(0.01, 4)
-            w[7] = w[7].clamp(0, 0.75)
+            w[5] = w[5].clamp(0.001, 4)
+            w[6] = w[6].clamp(0.001, 4)
+            w[7] = w[7].clamp(0.001, 0.75)
             w[8] = w[8].clamp(0, 4.5)
             w[9] = w[9].clamp(0, 0.8)
-            w[10] = w[10].clamp(0.01, 3.5)
-            w[11] = w[11].clamp(0.1, 5)
-            w[12] = w[12].clamp(0.01, 0.25)
-            w[13] = w[13].clamp(0.01, 0.9)
-            w[14] = w[14].clamp(0.01, 4)
+            w[10] = w[10].clamp(0.001, 3.5)
+            w[11] = w[11].clamp(0.001, 5)
+            w[12] = w[12].clamp(0.001, 0.25)
+            w[13] = w[13].clamp(0.001, 0.9)
+            w[14] = w[14].clamp(0, 4)
             w[15] = w[15].clamp(0, 1)
             w[16] = w[16].clamp(1, 6)
             w[17] = w[17].clamp(0, 2)
@@ -2075,22 +2081,3 @@ def wrap_short_term_ratings(r_history, t_history):
     else:
         result.pop()
     return "".join(result)
-
-
-if __name__ == "__main__":
-    model = FSRS(DEFAULT_PARAMETER)
-    stability = torch.tensor([5.0] * 4)
-    difficulty = torch.tensor([1.0, 2.0, 3.0, 4.0])
-    retention = torch.tensor([0.9, 0.8, 0.7, 0.6])
-    rating = torch.tensor([1, 2, 3, 4])
-    state = torch.stack([stability, difficulty]).unsqueeze(0)
-    s_recall = model.stability_after_success(state, retention, rating)
-    print(s_recall)
-    s_forget = model.stability_after_failure(state, retention)
-    print(s_forget)
-
-    retentions = torch.tensor([0.1, 0.2, 0.3, 0.4])
-    labels = torch.tensor([0.0, 1.0, 0.0, 1.0])
-    loss_fn = nn.BCELoss()
-    loss = loss_fn(retentions, labels)
-    print(loss)
diff --git a/src/fsrs_optimizer/fsrs_simulator.py b/src/fsrs_optimizer/fsrs_simulator.py
@@ -127,8 +127,12 @@ def init_d_with_short_term(rating):
         new_d = init_d(rating) - w[6] * rating_offset
         return np.clip(new_d, 1, 10)
 
+    def linear_damping(delta_d, old_d):
+        return delta_d * (10 - old_d) / 9
+
     def next_d(d, rating):
-        new_d = d - w[6] * (rating - 3)
+        delta_d = -w[6] * (rating - 3)
+        new_d = d + linear_damping(delta_d, d)
         new_d = mean_reversion(init_d(4), new_d)
         return np.clip(new_d, 1, 10)
 
diff --git a/tests/model_test.py b/tests/model_test.py
@@ -0,0 +1,137 @@
+from src.fsrs_optimizer import *
+
+
+class Test_Model:
+    def test_next_stability(self):
+        model = FSRS(DEFAULT_PARAMETER)
+        stability = torch.tensor([5.0] * 4)
+        difficulty = torch.tensor([1.0, 2.0, 3.0, 4.0])
+        retention = torch.tensor([0.9, 0.8, 0.7, 0.6])
+        rating = torch.tensor([1, 2, 3, 4])
+        state = torch.stack([stability, difficulty]).unsqueeze(0)
+        s_recall = model.stability_after_success(state, retention, rating)
+        assert torch.allclose(
+            s_recall, torch.tensor([25.7761, 14.1219, 60.4044, 208.9760]), atol=1e-4
+        )
+        s_forget = model.stability_after_failure(state, retention)
+        assert torch.allclose(
+            s_forget, torch.tensor([1.7029, 1.9799, 2.3760, 2.8885]), atol=1e-4
+        )
+        s_short_term = model.stability_short_term(state, rating)
+        assert torch.allclose(
+            s_short_term, torch.tensor([2.5051, 4.1992, 7.0389, 11.7988]), atol=1e-4
+        )
+
+    def test_next_difficulty(self):
+        model = FSRS(DEFAULT_PARAMETER)
+        stability = torch.tensor([5.0] * 4)
+        difficulty = torch.tensor([5.0] * 4)
+        rating = torch.tensor([1, 2, 3, 4])
+        state = torch.stack([stability, difficulty]).unsqueeze(0)
+        d_recall = model.next_d(state, rating)
+        assert torch.allclose(
+            d_recall,
+            torch.tensor([6.6070, 5.7994, 4.9918, 4.1842]),
+            atol=1e-4,
+        )
+
+    def test_power_forgetting_curve(self):
+        delta_t = torch.tensor([0, 1, 2, 3, 4, 5])
+        stability = torch.tensor([1, 2, 3, 4, 4, 2])
+        retention = power_forgetting_curve(delta_t, stability)
+        assert torch.allclose(
+            retention,
+            torch.tensor([1.0, 0.946059, 0.9299294, 0.9221679, 0.90000004, 0.79394597]),
+            atol=1e-4,
+        )
+
+    def test_forward(self):
+        model = FSRS(DEFAULT_PARAMETER)
+        delta_ts = torch.tensor(
+            [
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                [1.0, 1.0, 1.0, 1.0, 2.0, 2.0],
+            ]
+        )
+        ratings = torch.tensor(
+            [
+                [1.0, 2.0, 3.0, 4.0, 1.0, 2.0],
+                [1.0, 2.0, 3.0, 4.0, 1.0, 2.0],
+            ]
+        )
+        inputs = torch.stack([delta_ts, ratings], dim=2)
+        _, state = model.forward(inputs)
+        stability = state[:, 0]
+        difficulty = state[:, 1]
+        assert torch.allclose(
+            stability,
+            torch.tensor([0.2619, 1.7073, 5.8691, 25.0123, 0.3403, 2.1482]),
+            atol=1e-4,
+        )
+        assert torch.allclose(
+            difficulty,
+            torch.tensor([8.0827, 7.0405, 5.2729, 2.1301, 8.0827, 7.0405]),
+            atol=1e-4,
+        )
+
+    def test_loss_and_grad(self):
+        model = FSRS(DEFAULT_PARAMETER)
+        loss_fn = nn.BCELoss(reduction="none")
+        t_histories = torch.tensor(
+            [
+                [0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+                [0.0, 1.0, 1.0, 3.0],
+                [1.0, 3.0, 3.0, 5.0],
+                [3.0, 6.0, 6.0, 12.0],
+            ]
+        )
+        r_histories = torch.tensor(
+            [
+                [1.0, 2.0, 3.0, 4.0],
+                [3.0, 4.0, 2.0, 4.0],
+                [1.0, 4.0, 4.0, 3.0],
+                [4.0, 3.0, 3.0, 3.0],
+                [3.0, 1.0, 3.0, 3.0],
+                [2.0, 3.0, 3.0, 4.0],
+            ]
+        )
+        delta_ts = torch.tensor([4.0, 11.0, 12.0, 23.0])
+        labels = torch.tensor([1, 1, 1, 0], dtype=torch.float32, requires_grad=False)
+        inputs = torch.stack([t_histories, r_histories], dim=2)
+        seq_lens = inputs.shape[0]
+        real_batch_size = inputs.shape[1]
+        outputs, _ = model.forward(inputs)
+        stabilities = outputs[seq_lens - 1, torch.arange(real_batch_size), 0]
+        retentions = power_forgetting_curve(delta_ts, stabilities)
+        loss = loss_fn(retentions, labels).sum()
+        assert round(loss.item(), 4) == 4.4467
+        loss.backward()
+        assert torch.allclose(
+            model.w.grad,
+            torch.tensor(
+                [
+                    -0.0583,
+                    -0.0068,
+                    -0.0026,
+                    0.0105,
+                    -0.0513,
+                    1.3643,
+                    0.0837,
+                    -0.9502,
+                    0.5345,
+                    -2.8929,
+                    0.5142,
+                    -0.0131,
+                    0.0419,
+                    -0.1183,
+                    -0.0009,
+                    -0.1445,
+                    0.2024,
+                    0.2141,
+                    0.0323,
+                ]
+            ),
+            atol=1e-4,
+        )
diff --git a/tests/simulator_test.py b/tests/simulator_test.py
@@ -10,7 +10,7 @@ def test_simulate(self):
             memorized_cnt_per_day,
             cost_per_day,
         ) = simulate(w=DEFAULT_PARAMETER, request_retention=0.9)
-        assert memorized_cnt_per_day[-1] == 5918.574208243532
+        assert memorized_cnt_per_day[-1] == 5875.025236206539
 
     def test_optimal_retention(self):
         default_params = {
@@ -24,4 +24,4 @@ def test_optimal_retention(self):
             "loss_aversion": 2.5,
         }
         r = optimal_retention(**default_params)
-        assert r == 0.8346739534878145
+        assert r == 0.8263932
