refactoring

Author: hexqnt

src/models/arctangent_step.rs

@@ -1,6 +1,52 @@
 use super::common::{is_finite_non_negative, scale_and_mirror_upper_hessian};
 use ndarray::Array2;
 
+const PARAM_COUNT: usize = 4;
+
+#[derive(Clone, Copy)]
+struct Params<T> {
+    amplitude: T,
+    slope: T,
+    x0: T,
+    offset: T,
+}
+
+impl Params<f64> {
+    #[inline]
+    fn parse(param: &[f64]) -> Self {
+        let [amplitude, slope, x0, offset]: [f64; PARAM_COUNT] = param
+            .try_into()
+            .unwrap_or_else(|_| panic!("expected {} params", PARAM_COUNT));
+        Self {
+            amplitude,
+            slope,
+            x0,
+            offset,
+        }
+    }
+
+    #[inline]
+    fn value_at(self, x: f64) -> f64 {
+        self.amplitude * (self.slope * (x - self.x0)).atan() + self.offset
+    }
+
+    #[inline]
+    fn value_grad_at(self, x: f64, grad: &mut [f64]) -> f64 {
+        debug_assert_eq!(grad.len(), PARAM_COUNT);
+
+        let z = self.slope * (x - self.x0);
+        let atan_z = z.atan();
+        let inv_den = 1.0 / (1.0 + z * z);
+
+        grad[0] = atan_z;
+        grad[1] = self.amplitude * (x - self.x0) * inv_den;
+        grad[2] = -self.amplitude * self.slope * inv_den;
+        grad[3] = 1.0;
+
+        self.amplitude * atan_z + self.offset
+    }
+}
+
 /// Вычисляет арктангенс-ступень:
 /// `f(x) = amplitude * atan(slope * (x - x0)) + offset`,
 /// где:
@@ -10,31 +56,13 @@ use ndarray::Array2;
 /// - `offset` — вертикальный сдвиг.
 #[inline]
 pub(super) fn value_at(param: &[f64], x: f64) -> f64 {
-    let amplitude = param[0];
-    let slope = param[1];
-    let x0 = param[2];
-    let offset = param[3];
-    amplitude * (slope * (x - x0)).atan() + offset
+    Params::parse(param).value_at(x)
 }
 
+#[allow(dead_code)]
 #[inline]
 pub(super) fn value_grad_at(param: &[f64], x: f64, grad: &mut [f64]) -> f64 {
-    debug_assert_eq!(grad.len(), 4);
-
-    let amplitude = param[0];
-    let slope = param[1];
-    let x0 = param[2];
-    let offset = param[3];
-    let z = slope * (x - x0);
-    let atan_z = z.atan();
-    let inv_den = 1.0 / (1.0 + z * z);
-
-    grad[0] = atan_z;
-    grad[1] = amplitude * (x - x0) * inv_den;
-    grad[2] = -amplitude * slope * inv_den;
-    grad[3] = 1.0;
-
-    amplitude * atan_z + offset
+    Params::parse(param).value_grad_at(x, grad)
 }
 
 pub(super) fn add_value_grad(
@@ -45,18 +73,15 @@ pub(super) fn add_value_grad(
 ) {
     debug_assert_eq!(x_values.len(), value_first.len());
     debug_assert_eq!(gradient.len(), param.len());
+    let params = Params::parse(param);
 
-    let mut point_grad = [0.0; 4];
-    let mut index = 0;
-    while index < x_values.len() {
-        let upstream = value_first[index];
-        value_grad_at(param, x_values[index], &mut point_grad);
-
-        gradient[0] += upstream * point_grad[0];
-        gradient[1] += upstream * point_grad[1];
-        gradient[2] += upstream * point_grad[2];
-        gradient[3] += upstream * point_grad[3];
-        index += 1;
+    let mut point_grad = [0.0; PARAM_COUNT];
+    for (&x, &upstream) in x_values.iter().zip(value_first.iter()) {
+        params.value_grad_at(x, &mut point_grad);
+
+        for (gradient_value, point_grad_value) in gradient.iter_mut().zip(point_grad.iter()) {
+            *gradient_value += upstream * point_grad_value;
+        }
     }
 }
 
@@ -66,26 +91,24 @@ pub(super) fn add_value_grad_raw_hessian(
     value_first: &[f64],
     value_second: &[f64],
 ) -> Option<Array2<f64>> {
-    if param.len() != 4 {
+    if param.len() != PARAM_COUNT {
         return None;
     }
 
     let sample_count = x_values.len();
     let sample_scale = 1.0 / sample_count as f64;
-    let mut hessian = Array2::zeros((4, 4));
-    let amplitude = param[0];
-    let slope = param[1];
-    let x0 = param[2];
+    let mut hessian = Array2::zeros((PARAM_COUNT, PARAM_COUNT));
+    let params = Params::parse(param);
 
     let mut index = 0;
     while index < sample_count {
         let x = x_values[index];
-        let u = x - x0;
-        let z = slope * u;
+        let u = x - params.x0;
+        let z = params.slope * u;
         let atan_z = z.atan();
         let inv_den = 1.0 / (1.0 + z * z);
         let d2_shape_dz2 = -2.0 * z * inv_den * inv_den;
-        let model = value_at(param, x);
+        let model = params.value_at(x);
         if !model.is_finite() {
             return None;
         }
@@ -97,15 +120,15 @@ pub(super) fn add_value_grad_raw_hessian(
         }
 
         let jac_a = atan_z;
-        let jac_b = amplitude * inv_den * u;
-        let jac_c = -amplitude * inv_den * slope;
+        let jac_b = params.amplitude * inv_den * u;
+        let jac_c = -params.amplitude * inv_den * params.slope;
         let jac_d = 1.0;
 
         let d2_model_dadb = inv_den * u;
-        let d2_model_dadc = -inv_den * slope;
-        let d2_model_dbdb = amplitude * d2_shape_dz2 * u * u;
-        let d2_model_dbdc = amplitude * (d2_shape_dz2 * (-slope) * u - inv_den);
-        let d2_model_dcdc = amplitude * d2_shape_dz2 * slope * slope;
+        let d2_model_dadc = -inv_den * params.slope;
+        let d2_model_dbdb = params.amplitude * d2_shape_dz2 * u * u;
+        let d2_model_dbdc = params.amplitude * (d2_shape_dz2 * (-params.slope) * u - inv_den);
+        let d2_model_dcdc = params.amplitude * d2_shape_dz2 * params.slope * params.slope;
 
         hessian[[0, 0]] += value_second * jac_a * jac_a;
         hessian[[0, 1]] += value_second * jac_a * jac_b + value_first * d2_model_dadb;