Epigraph of sq. Eucl. norm: improvements

About:
- algorithmic improvements for better num. stability
- better documentation

Author: alphaville

src/constraints/epigraph_squared_norm.rs

@@ -3,84 +3,157 @@ use crate::matrix_operations;
 use super::Constraint;
 
 #[derive(Copy, Clone, Default)]
-/// The epigraph of the squared Eucliden norm is a set of the form
-/// $X = \\{x = (z, t) \in \mathbb{R}^{n}\times \mathbb{R} {}:{} \\|z\\|^2 \leq t \\}.$
+/// The epigraph of the squared Euclidean norm, that is,
+/// $$
+/// X = \{x = (z, t) \in \mathbb{R}^{n}\times \mathbb{R} : \|z\|_2^2 \le t \}.
+/// $$
+///
+/// A point is represented by a slice `x` whose last entry is the scalar
+/// component `t`, while the preceding entries form the vector component `z`.
 pub struct EpigraphSquaredNorm {}
 
 impl EpigraphSquaredNorm {
     /// Create a new instance of the epigraph of the squared norm.
     ///
     /// Note that you do not need to specify the dimension.
+    #[must_use]
     pub fn new() -> Self {
         EpigraphSquaredNorm {}
     }
 }
 
 impl Constraint for EpigraphSquaredNorm {
-    ///Project on the epigraph of the squared Euclidean norm.
+    /// Project on the epigraph of the squared Euclidean norm.
     ///
-    /// The projection is computed as detailed
-    /// [here](https://mathematix.wordpress.com/2017/05/02/projection-on-the-epigraph-of-the-squared-euclidean-norm/).
+    /// Let the input be represented as $(z,t)$, where `z` is the vector formed
+    /// by the first `x.len() - 1` entries of `x`, and `t` is the last entry.
+    /// This method computes the Euclidean projection of $(z,t)$ onto
+    ///
+    /// $$
+    /// \operatorname{epi}\|\cdot\|_2^2
+    /// =
+    /// \{(u,s) \in \mathbb{R}^n \times \mathbb{R} : \|u\|_2^2 \le s\}.
+    /// $$
+    ///
+    /// If the point is already feasible, that is, if
+    ///
+    /// $$
+    /// \|z\|_2^2 \le t,
+    /// $$
+    ///
+    /// then the input is left unchanged.
+    ///
+    /// Otherwise, the projection is computed using the methodology described
+    /// [here](https://mathematix.wordpress.com/2017/05/02/projection-on-the-epigraph-of-the-squared-euclidean-norm/):
+    ///
+    /// 1. Eliminate the projected vector variable to obtain a cubic equation
+    ///    in the projected scalar variable.
+    /// 2. Compute the real roots of that cubic.
+    /// 3. Select the admissible root that is consistent with the projection
+    ///    formula.
+    /// 4. Refine the selected root with a few Newton iterations.
+    /// 5. Recover the projected vector component using the refined root.
     ///
     /// ## Arguments
-    /// - `x`: The given vector $x$ is updated with the projection on the set
+    ///
+    /// - `x`: The given vector `x` is updated with the projection on the set
+    ///
+    /// ## Panics
+    ///
+    /// Panics if:
+    ///
+    /// - `x.len() < 2`,
+    /// - no admissible real root is found,
+    /// - the Newton derivative becomes too small,
+    /// - the final scaling factor is numerically singular.
     ///
     /// ## Example
     ///
     /// ```rust
     /// use optimization_engine::constraints::*;
     ///
     /// let epi = EpigraphSquaredNorm::new();
-    /// let mut x = [1., 2., 3., 4.];    
+    ///
+    /// // Here, z = [1., 2., 3.] and t = 4.
+    /// let mut x = [1., 2., 3., 4.];
+    ///
     /// epi.project(&mut x);
     /// ```
     fn project(&self, x: &mut [f64]) {
+        assert!(
+            x.len() >= 2,
+            "EpigraphSquaredNorm::project requires x.len() >= 2"
+        );
+
         let nx = x.len() - 1;
-        assert!(nx > 0, "x must have a length of at least 2");
-        let z: &[f64] = &x[..nx];
-        let t: f64 = x[nx];
+        let z = &x[..nx];
+        let t = x[nx];
         let norm_z_sq = matrix_operations::norm2_squared(z);
+
+        // Already feasible
         if norm_z_sq <= t {
             return;
         }
 
-        let theta = 1. - 2. * t;
-        let a3 = 4.;
-        let a2 = 4. * theta;
+        // Cubic:
+        // 4 r^3 + 4 theta r^2 + theta^2 r - ||z||^2 = 0
+        let theta = 1.0 - 2.0 * t;
+        let a3 = 4.0;
+        let a2 = 4.0 * theta;
         let a1 = theta * theta;
         let a0 = -norm_z_sq;
 
         let cubic_poly_roots = roots::find_roots_cubic(a3, a2, a1, a0);
-        let mut right_root = f64::NAN;
-        let mut scaling = f64::NAN;
-
-        // Find right root
-        cubic_poly_roots.as_ref().iter().for_each(|ri| {
-            if *ri > 0. {
-                let denom = 1. + 2. * (*ri - t);
-                if ((norm_z_sq / (denom * denom)) - *ri).abs() < 1e-6 {
-                    right_root = *ri;
-                    scaling = denom;
+
+        let root_tol = 1e-6;
+        let mut right_root: Option<f64> = None;
+
+        // Pick the first admissible real root
+        for &ri in cubic_poly_roots.as_ref().iter() {
+            let denom = 1.0 + 2.0 * (ri - t);
+
+            // We need a valid scaling and consistency with ||z_proj||^2 = ri
+            if denom > 0.0 {
+                let candidate_norm_sq = norm_z_sq / (denom * denom);
+                if (candidate_norm_sq - ri).abs() <= root_tol {
+                    right_root = Some(ri);
+                    break;
                 }
             }
-        });
+        }
+
+        let mut zsol =
+            right_root.expect("EpigraphSquaredNorm::project: no admissible real root found");
 
-        // Refinement of root with Newton-Raphson
-        let mut refinement_error = 1.;
+        // Newton refinement
         let newton_max_iters: usize = 5;
         let newton_eps = 1e-14;
-        let mut zsol = right_root;
-        let mut iter = 0;
-        while refinement_error > newton_eps && iter < newton_max_iters {
+
+        for _ in 0..newton_max_iters {
             let zsol_sq = zsol * zsol;
             let zsol_cb = zsol_sq * zsol;
+
             let p_z = a3 * zsol_cb + a2 * zsol_sq + a1 * zsol + a0;
-            let dp_z = 3. * a3 * zsol_sq + 2. * a2 * zsol + a1;
+            if p_z.abs() <= newton_eps {
+                break;
+            }
+
+            let dp_z = 3.0 * a3 * zsol_sq + 2.0 * a2 * zsol + a1;
+            assert!(
+                dp_z.abs() > 1e-15,
+                "EpigraphSquaredNorm::project: Newton derivative too small"
+            );
+
             zsol -= p_z / dp_z;
-            refinement_error = p_z.abs();
-            iter += 1;
         }
-        right_root = zsol;
+
+        let right_root = zsol;
+        let scaling = 1.0 + 2.0 * (right_root - t);
+
+        assert!(
+            scaling.abs() > 1e-15,
+            "EpigraphSquaredNorm::project: scaling factor too small"
+        );
 
         // Projection
         for xi in x.iter_mut().take(nx) {
@@ -89,7 +162,7 @@ impl Constraint for EpigraphSquaredNorm {
         x[nx] = right_root;
     }
 
-    /// This is a convex set, so this function returns `True`
+    /// This is a convex set, so this function returns `true`.
     fn is_convex(&self) -> bool {
         true
     }

src/constraints/mod.rs

@@ -22,7 +22,6 @@ mod rectangle;
 mod simplex;
 mod soc;
 mod sphere2;
-mod squared_norm_epigraph;
 mod zero;
 
 pub use affine_space::AffineSpace;
@@ -39,7 +38,6 @@ pub use rectangle::Rectangle;
 pub use simplex::Simplex;
 pub use soc::SecondOrderCone;
 pub use sphere2::Sphere2;
-pub use squared_norm_epigraph::EpiSqNorm;
 pub use zero::Zero;
 
 /// A set which can be used as a constraint