use crate::DEFAULT_PARAMETERS;
use crate::dataset::FSRSReview;
use crate::error::{FSRSError, Result};
use crate::inference::{FSRS5_DEFAULT_DECAY, MemoryState, Parameters};
use crate::parameter_clipper::clip_parameters;
use crate::simulation::{D_MAX, D_MIN, S_MAX, S_MIN};
#[cfg(test)]
use burn::{
module::{Module, Param},
tensor::{Shape, Tensor, TensorData, backend::Backend},
};
#[derive(Debug, Clone)]
#[allow(dead_code)]
pub struct ModelConfig {
pub freeze_initial_stability: bool,
pub initial_stability: Option<[f32; 4]>,
pub freeze_short_term_stability: bool,
pub num_relearning_steps: usize,
}
impl Default for ModelConfig {
fn default() -> Self {
Self {
freeze_initial_stability: false,
initial_stability: None,
freeze_short_term_stability: false,
num_relearning_steps: 1,
}
}
}
impl ModelConfig {
pub(crate) fn initial_parameters(&self) -> [f32; 21] {
let mut parameters: [f32; 21] = self
.initial_stability
.unwrap_or_else(|| DEFAULT_PARAMETERS[0..4].try_into().unwrap())
.into_iter()
.chain(DEFAULT_PARAMETERS[4..].iter().copied())
.collect::<Vec<_>>()
.try_into()
.unwrap();
if self.freeze_short_term_stability {
parameters[17] = 0.0;
parameters[18] = 0.0;
parameters[19] = 0.0;
}
parameters
}
#[cfg(test)]
pub fn init<B: Backend>(&self) -> Model<B> {
Model::new(self.clone())
}
}
#[derive(Debug, Clone)]
pub struct FSRS {
parameters: [f32; 21],
}
impl Default for FSRS {
fn default() -> Self {
Self::new(&[]).expect("Default parameters should be valid")
}
}
impl FSRS {
pub fn new(parameters: &Parameters) -> Result<Self> {
let parameters = check_and_fill_parameters(parameters)?;
let config = ModelConfig::default();
let parameters =
clip_parameters(¶meters, config.num_relearning_steps, Default::default());
Ok(Self {
parameters: parameters.try_into().unwrap(),
})
}
pub(crate) const fn parameters(&self) -> &[f32; 21] {
&self.parameters
}
#[inline]
pub(crate) fn power_forgetting_curve(&self, t: f32, s: f32) -> f32 {
power_forgetting_curve(&self.parameters, t, s)
}
#[inline]
pub(crate) fn next_interval_for_stability(
&self,
stability: f32,
desired_retention: f32,
) -> f32 {
next_interval(&self.parameters, stability, desired_retention)
}
#[inline]
pub(crate) fn init_stability(&self, rating: u32) -> f32 {
init_stability(&self.parameters, rating as usize)
}
#[inline]
pub(crate) fn init_difficulty(&self, rating: u32) -> f32 {
init_difficulty(&self.parameters, rating as usize)
}
#[inline]
#[cfg(test)]
pub(crate) fn next_difficulty(&self, difficulty: f32, rating: u32) -> f32 {
next_difficulty(&self.parameters, difficulty, rating as f32)
}
pub(crate) fn step(
&self,
delta_t: f32,
rating: u32,
state: MemoryState,
nth: usize,
) -> MemoryState {
step(&self.parameters, delta_t, rating as f32, state, nth)
}
pub(crate) fn forward_reviews(
&self,
reviews: &[FSRSReview],
starting_state: Option<MemoryState>,
) -> MemoryState {
let (mut state, start_index) = if let Some(state) = starting_state {
(state, 0)
} else if reviews.is_empty() {
(
MemoryState {
stability: 0.0,
difficulty: 0.0,
},
0,
)
} else {
let rating = reviews[0].rating;
if rating == 0 {
(
MemoryState {
stability: S_MIN,
difficulty: D_MIN,
},
1,
)
} else {
let rating = rating.clamp(1, 4);
(
MemoryState {
stability: self.init_stability(rating).clamp(S_MIN, S_MAX),
difficulty: self.init_difficulty(rating).clamp(D_MIN, D_MAX),
},
1,
)
}
};
for (index, review) in reviews.iter().enumerate().skip(start_index) {
state = self.step(review.delta_t as f32, review.rating, state, index);
}
state
}
}
#[inline]
pub(crate) fn power_forgetting_curve(w: &[f32], t: f32, s: f32) -> f32 {
let decay = -w[20];
let factor = (0.9f32.ln() / decay).exp() - 1.0;
(t / s * factor + 1.0).powf(decay)
}
#[inline]
pub(crate) fn next_interval(w: &[f32], stability: f32, desired_retention: f32) -> f32 {
let decay = -w[20];
let factor = (0.9f32.ln() / decay).exp() - 1.0;
stability / factor * (desired_retention.powf(1.0 / decay) - 1.0)
}
#[inline]
pub(crate) fn init_stability(w: &[f32], rating: usize) -> f32 {
w[rating.saturating_sub(1).min(3)]
}
#[inline]
pub(crate) fn init_difficulty(w: &[f32], rating: usize) -> f32 {
w[4] - (w[5] * rating.saturating_sub(1) as f32).exp() + 1.0
}
#[inline]
fn mean_reversion(w: &[f32], new_d: f32) -> f32 {
w[7] * (init_difficulty(w, 4) - new_d) + new_d
}
#[inline]
fn linear_damping(delta_d: f32, old_d: f32) -> f32 {
(10.0 - old_d) * delta_d / 9.0
}
#[inline]
pub(crate) fn next_difficulty(w: &[f32], difficulty: f32, rating: f32) -> f32 {
let delta_d = -w[6] * (rating - 3.0);
difficulty + linear_damping(delta_d, difficulty)
}
#[inline]
fn stability_after_success(w: &[f32], last_s: f32, last_d: f32, r: f32, rating: f32) -> f32 {
let hard_penalty = if rating == 2.0 { w[15] } else { 1.0 };
let easy_bonus = if rating == 4.0 { w[16] } else { 1.0 };
last_s
* (w[8].exp()
* (11.0 - last_d)
* last_s.powf(-w[9])
* (((1.0 - r) * w[10]).exp() - 1.0)
* hard_penalty
* easy_bonus
+ 1.0)
}
#[inline]
fn stability_after_failure(w: &[f32], last_s: f32, last_d: f32, r: f32) -> f32 {
let new_s = w[11]
* last_d.powf(-w[12])
* ((last_s + 1.0).powf(w[13]) - 1.0)
* ((1.0 - r) * w[14]).exp();
let new_s_min = last_s / (w[17] * w[18]).exp();
new_s.min(new_s_min)
}
#[inline]
fn stability_short_term(w: &[f32], last_s: f32, rating: f32) -> f32 {
let sinc = (w[17] * (rating - 3.0 + w[18])).exp() * last_s.powf(-w[19]);
last_s * if rating >= 2.0 { sinc.max(1.0) } else { sinc }
}
fn step(w: &[f32], delta_t: f32, rating: f32, state: MemoryState, nth: usize) -> MemoryState {
let last_s = state.stability.clamp(S_MIN, S_MAX);
let last_d = state.difficulty.clamp(D_MIN, D_MAX);
let retrievability = power_forgetting_curve(w, delta_t, last_s);
let stability_after_success =
stability_after_success(w, last_s, last_d, retrievability, rating);
let stability_after_failure = stability_after_failure(w, last_s, last_d, retrievability);
let stability_short_term = stability_short_term(w, last_s, rating);
let mut new_s = if rating == 1.0 {
stability_after_failure
} else {
stability_after_success
};
if delta_t == 0.0 {
new_s = stability_short_term;
}
let mut new_d = next_difficulty(w, last_d, rating);
new_d = mean_reversion(w, new_d).clamp(D_MIN, D_MAX);
if nth == 0 && state.stability == 0.0 {
let init_rating = (rating as u32).clamp(1, 4) as usize;
new_s = init_stability(w, init_rating);
new_d = init_difficulty(w, init_rating).clamp(D_MIN, D_MAX);
}
if rating == 0.0 {
new_s = last_s;
new_d = last_d;
}
MemoryState {
stability: new_s.clamp(S_MIN, S_MAX),
difficulty: new_d,
}
}
pub fn check_and_fill_parameters(parameters: &Parameters) -> Result<Vec<f32>, FSRSError> {
let parameters = match parameters.len() {
0 => DEFAULT_PARAMETERS.to_vec(),
17 => {
let mut parameters = parameters.to_vec();
parameters[4] = parameters[5].mul_add(2.0, parameters[4]);
parameters[5] = parameters[5].mul_add(3.0, 1.0).ln() / 3.0;
parameters[6] += 0.5;
parameters.extend_from_slice(&[0.0, 0.0, 0.0, FSRS5_DEFAULT_DECAY]);
parameters
}
19 => {
let mut parameters = parameters.to_vec();
parameters.extend_from_slice(&[0.0, FSRS5_DEFAULT_DECAY]);
parameters
}
21 => parameters.to_vec(),
_ => return Err(FSRSError::InvalidParameters),
};
if parameters.iter().any(|&w| !w.is_finite()) {
return Err(FSRSError::InvalidParameters);
}
Ok(parameters)
}
#[cfg(test)]
#[derive(Module, Debug)]
pub struct Model<B: Backend> {
pub w: Param<Tensor<B, 1>>,
}
#[cfg(test)]
pub(crate) trait Get<B: Backend, const N: usize> {
fn get(&self, n: usize) -> Tensor<B, N>;
}
#[cfg(test)]
impl<B: Backend, const N: usize> Get<B, N> for Tensor<B, N> {
fn get(&self, n: usize) -> Self {
self.clone().slice([n..(n + 1)])
}
}
#[cfg(test)]
impl<B: Backend> Model<B> {
pub fn new(config: ModelConfig) -> Self {
Self::new_with_device(config, &B::Device::default())
}
pub fn new_with_device(config: ModelConfig, device: &B::Device) -> Self {
Self {
w: Param::from_tensor(Tensor::from_floats(
TensorData::new(
config.initial_parameters().to_vec(),
Shape { dims: vec![21] },
),
device,
)),
}
}
pub fn power_forgetting_curve(&self, t: Tensor<B, 1>, s: Tensor<B, 1>) -> Tensor<B, 1> {
let decay = -self.w.get(20);
let factor = decay.clone().powi_scalar(-1).mul_scalar(0.9f32.ln()).exp() - 1.0;
(t / s * factor + 1.0).powf(decay)
}
#[allow(dead_code)]
pub fn next_interval(
&self,
stability: Tensor<B, 1>,
desired_retention: Tensor<B, 1>,
) -> Tensor<B, 1> {
let decay = -self.w.get(20);
let factor = decay.clone().powi_scalar(-1).mul_scalar(0.9f32.ln()).exp() - 1.0;
stability / factor * (desired_retention.powf(decay.powi_scalar(-1)) - 1.0)
}
fn stability_after_success(
&self,
last_s: Tensor<B, 1>,
last_d: Tensor<B, 1>,
r: Tensor<B, 1>,
rating: Tensor<B, 1>,
) -> Tensor<B, 1> {
let batch_size = rating.dims()[0];
let device = rating.device();
let hard_penalty = Tensor::ones([batch_size], &device)
.mask_where(rating.clone().equal_elem(2), self.w.get(15));
let easy_bonus =
Tensor::ones([batch_size], &device).mask_where(rating.equal_elem(4), self.w.get(16));
last_s.clone()
* (self.w.get(8).exp()
* (-last_d + 11)
* (last_s.powf(-self.w.get(9)))
* (((-r + 1) * self.w.get(10)).exp() - 1)
* hard_penalty
* easy_bonus
+ 1)
}
fn stability_after_failure(
&self,
last_s: Tensor<B, 1>,
last_d: Tensor<B, 1>,
r: Tensor<B, 1>,
) -> Tensor<B, 1> {
let new_s = self.w.get(11)
* last_d.powf(-self.w.get(12))
* ((last_s.clone() + 1).powf(self.w.get(13)) - 1)
* ((-r + 1) * self.w.get(14)).exp();
let new_s_min = last_s / (self.w.get(17) * self.w.get(18)).exp();
new_s
.clone()
.mask_where(new_s_min.clone().lower(new_s), new_s_min)
}
fn stability_short_term(&self, last_s: Tensor<B, 1>, rating: Tensor<B, 1>) -> Tensor<B, 1> {
let sinc = (self.w.get(17) * (rating.clone() - 3 + self.w.get(18))).exp()
* last_s.clone().powf(-self.w.get(19));
last_s
* sinc
.clone()
.mask_where(rating.greater_equal_elem(2), sinc.clamp_min(1.0))
}
fn mean_reversion(&self, new_d: Tensor<B, 1>) -> Tensor<B, 1> {
let device = new_d.device();
let rating = Tensor::from_floats([4.0], &device);
self.w.get(7) * (self.init_difficulty(rating) - new_d.clone()) + new_d
}
pub(crate) fn init_stability(&self, rating: Tensor<B, 1>) -> Tensor<B, 1> {
self.w.val().select(0, rating.int() - 1)
}
fn init_difficulty(&self, rating: Tensor<B, 1>) -> Tensor<B, 1> {
self.w.get(4) - (self.w.get(5) * (rating - 1)).exp() + 1
}
fn linear_damping(&self, delta_d: Tensor<B, 1>, old_d: Tensor<B, 1>) -> Tensor<B, 1> {
old_d.neg().add_scalar(10.0) * delta_d.div_scalar(9.0)
}
fn next_difficulty(&self, difficulty: Tensor<B, 1>, rating: Tensor<B, 1>) -> Tensor<B, 1> {
let delta_d = -self.w.get(6) * (rating - 3);
difficulty.clone() + self.linear_damping(delta_d, difficulty)
}
pub(crate) fn step(
&self,
delta_t: Tensor<B, 1>,
rating: Tensor<B, 1>,
state: MemoryStateTensors<B>,
nth: usize,
) -> MemoryStateTensors<B> {
let last_s = state.stability.clone().clamp(S_MIN, S_MAX);
let last_d = state.difficulty.clone().clamp(D_MIN, D_MAX);
let retrievability = self.power_forgetting_curve(delta_t.clone(), last_s.clone());
let stability_after_success = self.stability_after_success(
last_s.clone(),
last_d.clone(),
retrievability.clone(),
rating.clone(),
);
let stability_after_failure =
self.stability_after_failure(last_s.clone(), last_d.clone(), retrievability);
let stability_short_term = self.stability_short_term(last_s.clone(), rating.clone());
let mut new_s = stability_after_success
.mask_where(rating.clone().equal_elem(1), stability_after_failure);
new_s = new_s.mask_where(delta_t.equal_elem(0), stability_short_term);
let mut new_d = self.next_difficulty(last_d.clone(), rating.clone());
new_d = self.mean_reversion(new_d).clamp(D_MIN, D_MAX);
if nth == 0 {
let is_initial = state.stability.equal_elem(0.0);
let init_s = self.init_stability(rating.clone().clamp(1, 4));
let init_d = self
.init_difficulty(rating.clone().clamp(1, 4))
.clamp(D_MIN, D_MAX);
new_s = new_s.mask_where(is_initial.clone(), init_s);
new_d = new_d.mask_where(is_initial, init_d);
}
new_s = new_s.mask_where(rating.clone().equal_elem(0), last_s);
new_d = new_d.mask_where(rating.equal_elem(0), last_d);
MemoryStateTensors {
stability: new_s.clamp(S_MIN, S_MAX),
difficulty: new_d,
}
}
pub(crate) fn forward(
&self,
delta_ts: Tensor<B, 2>,
ratings: Tensor<B, 2>,
starting_state: Option<MemoryStateTensors<B>>,
) -> MemoryStateTensors<B> {
let [seq_len, batch_size] = delta_ts.dims();
let (mut state, start_index) = if let Some(state) = starting_state {
(state, 0)
} else if seq_len == 0 {
(MemoryStateTensors::zeros(batch_size), 0)
} else {
let rating = ratings.get(0).squeeze(0);
let initial_rating = rating.clone().clamp(1, 4);
let mut stability = self
.init_stability(initial_rating.clone())
.clamp(S_MIN, S_MAX);
let mut difficulty = self.init_difficulty(initial_rating).clamp(D_MIN, D_MAX);
let padding = rating.equal_elem(0);
let device = stability.device();
stability = stability.mask_where(
padding.clone(),
Tensor::ones([batch_size], &device).mul_scalar(S_MIN),
);
difficulty = difficulty.mask_where(
padding,
Tensor::ones([batch_size], &device).mul_scalar(D_MIN),
);
(
MemoryStateTensors {
stability,
difficulty,
},
1,
)
};
for i in start_index..seq_len {
let delta_t = delta_ts.get(i).squeeze(0);
let rating = ratings.get(i).squeeze(0);
state = self.step(delta_t, rating, state, i);
}
state
}
}
#[cfg(test)]
#[derive(Debug, Clone)]
pub(crate) struct MemoryStateTensors<B: Backend> {
pub stability: Tensor<B, 1>,
pub difficulty: Tensor<B, 1>,
}
#[cfg(test)]
impl<B: Backend> MemoryStateTensors<B> {
pub(crate) fn zeros(batch_size: usize) -> MemoryStateTensors<B> {
let device = B::Device::default();
MemoryStateTensors {
stability: Tensor::zeros([batch_size], &device),
difficulty: Tensor::zeros([batch_size], &device),
}
}
pub(crate) fn from_state(state: MemoryState) -> Self {
let device = B::Device::default();
Self {
stability: Tensor::from_floats([state.stability], &device),
difficulty: Tensor::from_floats([state.difficulty], &device),
}
}
}
#[cfg(test)]
impl<B: Backend> From<MemoryStateTensors<B>> for MemoryState {
fn from(m: MemoryStateTensors<B>) -> Self {
use burn::tensor::ElementConversion;
Self {
stability: m.stability.into_scalar().elem(),
difficulty: m.difficulty.into_scalar().elem(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::TestHelper;
use crate::test_helpers::{Model, Tensor};
use burn::backend::ndarray::NdArrayDevice;
use burn::tensor::TensorData;
static DEVICE: NdArrayDevice = NdArrayDevice::Cpu;
#[test]
fn test_w() {
let model = Model::new(ModelConfig::default());
assert_eq!(
model.w.val().to_data(),
TensorData::from(DEFAULT_PARAMETERS)
)
}
#[test]
fn test_convert_parameters() {
let fsrs4dot5_param = vec![
0.4, 0.6, 2.4, 5.8, 4.93, 0.94, 0.86, 0.01, 1.49, 0.14, 0.94, 2.18, 0.05, 0.34, 1.26,
0.29, 2.61,
];
let fsrs5_param = check_and_fill_parameters(&fsrs4dot5_param).unwrap();
assert_eq!(
fsrs5_param,
vec![
0.4, 0.6, 2.4, 5.8, 6.81, 0.44675013, 1.36, 0.01, 1.49, 0.14, 0.94, 2.18, 0.05,
0.34, 1.26, 0.29, 2.61, 0.0, 0.0, 0.0, 0.5
]
)
}
#[test]
fn test_power_forgetting_curve() {
let fsrs = FSRS::default();
let retrievability = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0]
.into_iter()
.zip([1.0, 2.0, 3.0, 4.0, 4.0, 2.0])
.map(|(delta_t, stability)| fsrs.power_forgetting_curve(delta_t, stability))
.collect::<Vec<_>>();
retrievability.assert_approx_eq([1.0, 0.9403443, 0.9253786, 0.9185229, 0.9, 0.8261359]);
}
#[test]
fn test_init_stability() {
let fsrs = FSRS::default();
let stability = [1, 2, 3, 4, 1, 2].map(|rating| fsrs.init_stability(rating));
assert_eq!(
stability,
[
DEFAULT_PARAMETERS[0],
DEFAULT_PARAMETERS[1],
DEFAULT_PARAMETERS[2],
DEFAULT_PARAMETERS[3],
DEFAULT_PARAMETERS[0],
DEFAULT_PARAMETERS[1]
]
)
}
#[test]
fn test_init_difficulty() {
let fsrs = FSRS::default();
let difficulty = [1, 2, 3, 4, 1, 2].map(|rating| fsrs.init_difficulty(rating));
assert_eq!(
difficulty,
[
DEFAULT_PARAMETERS[4],
DEFAULT_PARAMETERS[4] - DEFAULT_PARAMETERS[5].exp() + 1.0,
DEFAULT_PARAMETERS[4] - (2.0 * DEFAULT_PARAMETERS[5]).exp() + 1.0,
DEFAULT_PARAMETERS[4] - (3.0 * DEFAULT_PARAMETERS[5]).exp() + 1.0,
DEFAULT_PARAMETERS[4],
DEFAULT_PARAMETERS[4] - DEFAULT_PARAMETERS[5].exp() + 1.0,
]
)
}
#[test]
fn test_forward_matches_burn_oracle() {
let model = Model::new(ModelConfig::default());
let delta_ts = Tensor::from_floats(
[
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[1.0, 1.0, 1.0, 1.0, 2.0, 2.0],
],
&DEVICE,
);
let ratings = Tensor::from_floats(
[
[1.0, 2.0, 3.0, 4.0, 1.0, 2.0],
[1.0, 2.0, 3.0, 4.0, 1.0, 2.0],
],
&DEVICE,
);
let state = model.forward(delta_ts, ratings, None);
let burn_stability = state.stability.to_data().to_vec::<f32>().unwrap();
let burn_difficulty = state.difficulty.to_data().to_vec::<f32>().unwrap();
let fsrs = FSRS::default();
let scalar_states = [
[(1, 0), (1, 1)],
[(2, 0), (2, 1)],
[(3, 0), (3, 1)],
[(4, 0), (4, 1)],
[(1, 0), (1, 2)],
[(2, 0), (2, 2)],
]
.map(|reviews| FSRSReview {
rating: reviews[0].0,
delta_t: reviews[0].1,
})
.into_iter()
.zip([
FSRSReview {
rating: 1,
delta_t: 1,
},
FSRSReview {
rating: 2,
delta_t: 1,
},
FSRSReview {
rating: 3,
delta_t: 1,
},
FSRSReview {
rating: 4,
delta_t: 1,
},
FSRSReview {
rating: 1,
delta_t: 2,
},
FSRSReview {
rating: 2,
delta_t: 2,
},
])
.map(|(first, second)| fsrs.forward_reviews(&[first, second], None))
.collect::<Vec<_>>();
let scalar_stability = scalar_states
.iter()
.map(|state| state.stability)
.collect::<Vec<_>>();
let scalar_difficulty = scalar_states
.iter()
.map(|state| state.difficulty)
.collect::<Vec<_>>();
let burn_stability: [f32; 6] = burn_stability.try_into().unwrap();
let burn_difficulty: [f32; 6] = burn_difficulty.try_into().unwrap();
scalar_stability.assert_approx_eq(burn_stability);
scalar_difficulty.assert_approx_eq(burn_difficulty);
}
#[test]
fn test_next_difficulty() {
let fsrs = FSRS::default();
let next_difficulty = [1, 2, 3, 4].map(|rating| fsrs.next_difficulty(5.0, rating));
next_difficulty.assert_approx_eq([8.354889, 6.6774445, 5.0, 3.3225555]);
let next_difficulty = next_difficulty.map(|value| mean_reversion(fsrs.parameters(), value));
next_difficulty.assert_approx_eq([8.341763, 6.6659956, 4.990228, 3.3144615]);
}
#[test]
fn test_next_stability() {
let w = DEFAULT_PARAMETERS;
let s_recall = [1.0, 2.0, 3.0, 4.0]
.into_iter()
.zip([0.9, 0.8, 0.7, 0.6])
.map(|(rating, retrievability)| {
stability_after_success(&w, 5.0, rating, retrievability, rating)
})
.collect::<Vec<_>>();
s_recall.assert_approx_eq([25.602541, 28.226582, 58.656002, 127.226685]);
let s_forget = [1.0, 2.0, 3.0, 4.0]
.into_iter()
.zip([0.9, 0.8, 0.7, 0.6])
.map(|(difficulty, retrievability)| {
stability_after_failure(&w, 5.0, difficulty, retrievability)
})
.collect::<Vec<_>>();
s_forget.assert_approx_eq([1.0525396, 1.1894329, 1.3680838, 1.584989]);
let next_stability = [s_forget[0], s_recall[1], s_recall[2], s_recall[3]];
next_stability.assert_approx_eq([1.0525396, 28.226582, 58.656002, 127.226685]);
let next_stability =
[1.0, 2.0, 3.0, 4.0].map(|rating| stability_short_term(&w, 5.0, rating));
next_stability.assert_approx_eq([1.596818, 5.0, 5.0, 8.12961]);
}
#[test]
fn test_fsrs() {
FSRS::default()
.parameters()
.to_vec()
.assert_approx_eq(DEFAULT_PARAMETERS);
assert!(FSRS::new(&[]).is_ok());
assert!(FSRS::new(&[1.]).is_err());
assert!(FSRS::new(DEFAULT_PARAMETERS.as_slice()).is_ok());
assert!(FSRS::new(&DEFAULT_PARAMETERS[..17]).is_ok());
}
#[test]
fn scalar_step_matches_burn_oracle() {
let model = Model::new(ModelConfig::default());
let fsrs = FSRS::default();
let starting = MemoryState {
stability: 5.0,
difficulty: 6.0,
};
for (nth, state) in [
(
0,
MemoryState {
stability: 0.0,
difficulty: 0.0,
},
),
(1, starting),
] {
for delta_t in [0.0, 1.0, 21.0] {
for rating in 0..=4 {
let burn_state: MemoryState = model
.step(
Tensor::from_floats([delta_t], &DEVICE),
Tensor::from_floats([rating as f32], &DEVICE),
MemoryStateTensors::from_state(state),
nth,
)
.into();
let scalar_state = fsrs.step(delta_t, rating, state, nth);
[scalar_state.stability, scalar_state.difficulty]
.assert_approx_eq([burn_state.stability, burn_state.difficulty]);
}
}
}
}
}