use super::validation::{invalid, matrix_size, nonnegative, positive};
use crate::{Error, Result as QuantumResult};
use burn::config::Config;
#[derive(Config, Debug)]
pub struct QnSpsaConfig {
#[config(default = 0.05)]
pub learning_rate: f64,
#[config(default = 0.1)]
pub perturbation: f64,
#[config(default = 0.602)]
pub learning_rate_decay: f64,
#[config(default = 0.101)]
pub perturbation_decay: f64,
#[config(default = 0.0)]
pub stability_constant: f64,
#[config(default = 0.01)]
pub regularization: f64,
#[config(default = 1.0)]
pub initial_metric_scale: f64,
#[config(default = 1)]
pub resamplings: usize,
#[config(default = 0)]
pub metric_delay: u64,
#[config(default = "None")]
pub blocking_tolerance: Option<f64>,
#[config(default = 0)]
pub seed: u64,
#[config(default = 65_536)]
pub max_metric_elements: usize,
}
impl QnSpsaConfig {
pub(super) fn validate(&self, parameters: usize) -> QuantumResult<usize> {
if !positive(self.learning_rate)
|| !positive(self.perturbation)
|| !positive(self.regularization)
|| !nonnegative(self.initial_metric_scale)
|| !nonnegative(self.learning_rate_decay)
|| !nonnegative(self.perturbation_decay)
|| !nonnegative(self.stability_constant)
|| self.blocking_tolerance.is_some_and(|x| !nonnegative(x))
|| self.resamplings == 0
{
return Err(invalid(
"invalid rate, perturbation, regularization, schedule, or resampling count",
));
}
self.resamplings
.checked_mul(4)
.and_then(|n| n.checked_add(2))
.ok_or(Error::Overflow)?;
matrix_size(parameters, self.max_metric_elements)
}
}