src/regularisation.rs

/*!
Regularisation terms
*/

#[allow(unused_imports)] // Used by documentation.
use crate::fb::pointsource_fb_reg;
use crate::fb::InsertionConfig;
use crate::measures::{DeltaMeasure, DiscreteMeasure, Radon, RNDM};
#[allow(unused_imports)] // Used by documentation.
use crate::sliding_fb::pointsource_sliding_fb_reg;
use crate::types::*;
use alg_tools::instance::{Instance, Space};
use alg_tools::linops::Mapping;
use alg_tools::loc::Loc;
use alg_tools::norms::Norm;
use numeric_literals::replace_float_literals;
use serde::{Deserialize, Serialize};

use crate::subproblem::{
    l1squared_nonneg::l1squared_nonneg, l1squared_unconstrained::l1squared_unconstrained,
    nonneg::quadratic_nonneg, unconstrained::quadratic_unconstrained,
};
use alg_tools::bounds::{Bounds, MinMaxMapping};
use alg_tools::iterate::AlgIteratorFactory;
use alg_tools::nalgebra_support::ToNalgebraRealField;
use nalgebra::{DMatrix, DVector};

use std::cmp::Ordering::{Equal, Greater, Less};

/// The regularisation term $α\\|μ\\|\_{ℳ(Ω)} + δ_{≥ 0}(μ)$ for [`pointsource_fb_reg`] and other
/// algorithms.
///
/// The only member of the struct is the regularisation parameter α.
#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
pub struct NonnegRadonRegTerm<F: Float = f64>(pub F /* α */);

impl<'a, F: Float> NonnegRadonRegTerm<F> {
    /// Returns the regularisation parameter
    pub fn α(&self) -> F {
        let &NonnegRadonRegTerm(α) = self;
        α
    }
}

impl<'a, F: Float, const N: usize> Mapping<RNDM<N, F>> for NonnegRadonRegTerm<F> {
    type Codomain = F;

    fn apply<I>(&self, μ: I) -> F
    where
        I: Instance<RNDM<N, F>>,
    {
        self.α() * μ.eval(|x| x.norm(Radon))
    }
}

/// The regularisation term $α\|μ\|_{ℳ(Ω)}$ for [`pointsource_fb_reg`].
///
/// The only member of the struct is the regularisation parameter α.
#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
pub struct RadonRegTerm<F: Float = f64>(pub F /* α */);

impl<'a, F: Float> RadonRegTerm<F> {
    /// Returns the regularisation parameter
    pub fn α(&self) -> F {
        let &RadonRegTerm(α) = self;
        α
    }
}

impl<'a, F: Float, const N: usize> Mapping<RNDM<N, F>> for RadonRegTerm<F> {
    type Codomain = F;

    fn apply<I>(&self, μ: I) -> F
    where
        I: Instance<RNDM<N, F>>,
    {
        self.α() * μ.eval(|x| x.norm(Radon))
    }
}

/// Regularisation term configuration
#[derive(Clone, Copy, Eq, PartialEq, Serialize, Deserialize, Debug)]
pub enum Regularisation<F: Float = f64> {
    /// $α \\|μ\\|\_{ℳ(Ω)}$
    Radon(F),
    /// $α\\|μ\\|\_{ℳ(Ω)} + δ_{≥ 0}(μ)$
    NonnegRadon(F),
}

impl<'a, F: Float, const N: usize> Mapping<RNDM<N, F>> for Regularisation<F> {
    type Codomain = F;

    fn apply<I>(&self, μ: I) -> F
    where
        I: Instance<RNDM<N, F>>,
    {
        match *self {
            Self::Radon(α) => RadonRegTerm(α).apply(μ),
            Self::NonnegRadon(α) => NonnegRadonRegTerm(α).apply(μ),
        }
    }
}

/// Abstraction of regularisation terms.
pub trait RegTerm<Domain, F = f64>: Mapping<DiscreteMeasure<Domain, F>, Codomain = F>
where
    Domain: Space + Clone,
    F: Float + ToNalgebraRealField,
{
    /// Approximately solve the problem
    /// <div>$$
    ///     \min_{x ∈ ℝ^n} \frac{1}{2} x^⊤Ax - g^⊤ x + τ G(x)
    /// $$</div>
    /// for $G$ depending on the trait implementation.
    ///
    /// The parameter `mA` is $A$. An estimate for its opeator norm should be provided in
    /// `mA_normest`. The initial iterate and output is `x`. The current main tolerance is `ε`.
    ///
    /// Returns the number of iterations taken.
    fn solve_findim(
        &self,
        mA: &DMatrix<F::MixedType>,
        g: &DVector<F::MixedType>,
        τ: F,
        x: &mut DVector<F::MixedType>,
        mA_normest: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> usize;

    /// Approximately solve the problem
    /// <div>$$
    ///     \min_{x ∈ ℝ^n} \frac{1}{2} |x-y|_1^2 - g^⊤ x + τ G(x)
    /// $$</div>
    /// for $G$ depending on the trait implementation.
    ///
    /// Returns the number of iterations taken.
    fn solve_findim_l1squared(
        &self,
        y: &DVector<F::MixedType>,
        g: &DVector<F::MixedType>,
        τ: F,
        x: &mut DVector<F::MixedType>,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> usize;

    /// Find a point where `d` may violate the tolerance `ε`.
    ///
    /// If `skip_by_rough_check` is set, do not find the point if a rough check indicates that we
    /// are in bounds. `ε` is the current main tolerance and `τ` a scaling factor for the
    /// regulariser.
    ///
    /// Returns `None` if `d` is in bounds either based on the rough check, or a more precise check
    /// terminating early. Otherwise returns a possibly violating point, the value of `d` there,
    /// and a boolean indicating whether the found point is in bounds.
    fn find_tolerance_violation<M>(
        &self,
        d: &mut M,
        τ: F,
        ε: F,
        skip_by_rough_check: bool,
        config: &InsertionConfig<F>,
    ) -> Option<(Domain, F, bool)>
    where
        M: MinMaxMapping<Domain, F>,
    {
        self.find_tolerance_violation_slack(d, τ, ε, skip_by_rough_check, config, F::ZERO)
    }

    /// Find a point where `d` may violate the tolerance `ε`.
    ///
    /// This version includes a `slack` parameter to expand the tolerances.
    /// It is used for Radon-norm squared proximal term in [`crate::prox_penalty::radon_squared`].
    ///
    /// If `skip_by_rough_check` is set, do not find the point if a rough check indicates that we
    /// are in bounds. `ε` is the current main tolerance and `τ` a scaling factor for the
    /// regulariser.
    ///
    /// Returns `None` if `d` is in bounds either based on the rough check, or a more precise check
    /// terminating early. Otherwise returns a possibly violating point, the value of `d` there,
    /// and a boolean indicating whether the found point is in bounds.
    fn find_tolerance_violation_slack<M>(
        &self,
        d: &mut M,
        τ: F,
        ε: F,
        skip_by_rough_check: bool,
        config: &InsertionConfig<F>,
        slack: F,
    ) -> Option<(Domain, F, bool)>
    where
        M: MinMaxMapping<Domain, F>;

    /// Verify that `d` is in bounds `ε` for a merge candidate `μ`
    ///
    /// `ε` is the current main tolerance and `τ` a scaling factor for the regulariser.
    fn verify_merge_candidate<M>(
        &self,
        d: &mut M,
        μ: &DiscreteMeasure<Domain, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> bool
    where
        M: MinMaxMapping<Domain, F>;

    /// Verify that `d` is in bounds `ε` for a merge candidate `μ`
    ///
    /// This version is s used for Radon-norm squared proximal term in
    /// [`crate::prox_penalty::radon_squared`].
    /// The [measures][crate::measures::DiscreteMeasure] `μ` and `radon_μ` are supposed to have
    /// same coordinates at same agreeing indices.
    ///
    /// `ε` is the current main tolerance and `τ` a scaling factor for the regulariser.
    fn verify_merge_candidate_radonsq<M>(
        &self,
        d: &mut M,
        μ: &DiscreteMeasure<Domain, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
        radon_μ: &DiscreteMeasure<Domain, F>,
    ) -> bool
    where
        M: MinMaxMapping<Domain, F>;

    /// TODO: document this
    fn target_bounds(&self, τ: F, ε: F) -> Option<Bounds<F>>;

    /// Returns a scaling factor for the tolerance sequence.
    ///
    /// Typically this is the regularisation parameter.
    fn tolerance_scaling(&self) -> F;
}

/// Abstraction of regularisation terms for [`pointsource_sliding_fb_reg`].
pub trait SlidingRegTerm<Domain, F = f64>: RegTerm<Domain, F>
where
    Domain: Space + Clone,
    F: Float + ToNalgebraRealField,
{
    /// Calculate $τ[w(z) - w(y)]$ for some w in the subdifferential of the regularisation
    /// term, such that $-ε ≤ τw - d ≤ ε$.
    fn goodness<M>(
        &self,
        d: &mut M,
        μ: &DiscreteMeasure<Domain, F>,
        y: &Domain,
        z: &Domain,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> F
    where
        M: MinMaxMapping<Domain, F>;

    /// Convert bound on the regulariser to a bond on the Radon norm
    fn radon_norm_bound(&self, b: F) -> F;
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float + ToNalgebraRealField, const N: usize> RegTerm<Loc<N, F>, F>
    for NonnegRadonRegTerm<F>
{
    fn solve_findim(
        &self,
        mA: &DMatrix<F::MixedType>,
        g: &DVector<F::MixedType>,
        τ: F,
        x: &mut DVector<F::MixedType>,
        mA_normest: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> usize {
        let inner_tolerance = ε * config.inner.tolerance_mult;
        let inner_it = config.inner.iterator_options.stop_target(inner_tolerance);
        quadratic_nonneg(mA, g, τ * self.α(), x, mA_normest, &config.inner, inner_it)
    }

    fn solve_findim_l1squared(
        &self,
        y: &DVector<F::MixedType>,
        g: &DVector<F::MixedType>,
        τ: F,
        x: &mut DVector<F::MixedType>,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> usize {
        let inner_tolerance = ε * config.inner.tolerance_mult;
        let inner_it = config.inner.iterator_options.stop_target(inner_tolerance);
        l1squared_nonneg(y, g, τ * self.α(), 1.0, x, &config.inner, inner_it)
    }

    #[inline]
    fn find_tolerance_violation_slack<M>(
        &self,
        d: &mut M,
        τ: F,
        ε: F,
        skip_by_rough_check: bool,
        config: &InsertionConfig<F>,
        slack: F,
    ) -> Option<(Loc<N, F>, F, bool)>
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let τα = τ * self.α();
        let keep_above = -τα - slack - ε;
        let minimise_below = -τα - slack - ε * config.insertion_cutoff_factor;
        let refinement_tolerance = ε * config.refinement.tolerance_mult;

        // println!(
        //     "keep_above: {keep_above}, rough lower bound: {}, tolerance: {ε}, slack: {slack}, τα: {τα}",
        //     d.bounds().lower()
        // );

        // If preliminary check indicates that we are in bounds, and if it otherwise matches
        // the insertion strategy, skip insertion.
        if skip_by_rough_check && d.bounds().lower() >= keep_above {
            None
        } else {
            // If the rough check didn't indicate no insertion needed, find minimising point.
            let res = d.minimise_below(
                minimise_below,
                refinement_tolerance,
                config.refinement.max_steps,
            );

            res.map(|(ξ, v_ξ)| (ξ, v_ξ, v_ξ >= keep_above))
        }
    }

    fn verify_merge_candidate<M>(
        &self,
        d: &mut M,
        μ: &RNDM<N, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> bool
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let τα = τ * self.α();
        let refinement_tolerance = ε * config.refinement.tolerance_mult;
        let merge_tolerance = config.merge_tolerance_mult * ε;
        let keep_above = -τα - merge_tolerance;
        let keep_supp_below = -τα + merge_tolerance;
        let bnd = d.bounds();

        return (bnd.upper() <= keep_supp_below
            || μ
                .iter_spikes()
                .all(|&DeltaMeasure { α, ref x }| (α == 0.0) || d.apply(x) <= keep_supp_below))
            && (bnd.lower() >= keep_above
                || d.has_lower_bound(
                    keep_above,
                    refinement_tolerance,
                    config.refinement.max_steps,
                ));
    }

    fn verify_merge_candidate_radonsq<M>(
        &self,
        d: &mut M,
        μ: &RNDM<N, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
        radon_μ: &RNDM<N, F>,
    ) -> bool
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let τα = τ * self.α();
        let refinement_tolerance = ε * config.refinement.tolerance_mult;
        let merge_tolerance = config.merge_tolerance_mult * ε;
        let slack = radon_μ.norm(Radon);
        let bnd = d.bounds();

        return {
            μ.both_matching(radon_μ).all(|(α, rα, x)| {
                let v = -d.apply(x); // TODO: observe ad hoc negation here, after minus_τv
                                     // switch to τv.
                let (l1, u1) = match α.partial_cmp(&0.0).unwrap_or(Equal) {
                    Greater => (τα, τα),
                    _ => (F::NEG_INFINITY, τα),
                    // Less should not happen; treated as Equal
                };
                let (l2, u2) = match rα.partial_cmp(&0.0).unwrap_or(Equal) {
                    Greater => (slack, slack),
                    Equal => (-slack, slack),
                    Less => (-slack, -slack),
                };
                // TODO: both fail.
                (l1 + l2 - merge_tolerance <= v) && (v <= u1 + u2 + merge_tolerance)
            })
        } && {
            let keep_above = -τα - slack - merge_tolerance;
            bnd.lower() <= keep_above
                || d.has_lower_bound(
                    keep_above,
                    refinement_tolerance,
                    config.refinement.max_steps,
                )
        };
    }

    fn target_bounds(&self, τ: F, ε: F) -> Option<Bounds<F>> {
        let τα = τ * self.α();
        Some(Bounds(τα - ε, τα + ε))
    }

    fn tolerance_scaling(&self) -> F {
        self.α()
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float + ToNalgebraRealField, const N: usize> SlidingRegTerm<Loc<N, F>, F>
    for NonnegRadonRegTerm<F>
{
    fn goodness<M>(
        &self,
        d: &mut M,
        _μ: &RNDM<N, F>,
        y: &Loc<N, F>,
        z: &Loc<N, F>,
        τ: F,
        ε: F,
        _config: &InsertionConfig<F>,
    ) -> F
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let w = |x| 1.0.min((ε + d.apply(x)) / (τ * self.α()));
        w(z) - w(y)
    }

    fn radon_norm_bound(&self, b: F) -> F {
        b / self.α()
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float + ToNalgebraRealField, const N: usize> RegTerm<Loc<N, F>, F> for RadonRegTerm<F> {
    fn solve_findim(
        &self,
        mA: &DMatrix<F::MixedType>,
        g: &DVector<F::MixedType>,
        τ: F,
        x: &mut DVector<F::MixedType>,
        mA_normest: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> usize {
        let inner_tolerance = ε * config.inner.tolerance_mult;
        let inner_it = config.inner.iterator_options.stop_target(inner_tolerance);
        quadratic_unconstrained(mA, g, τ * self.α(), x, mA_normest, &config.inner, inner_it)
    }

    fn solve_findim_l1squared(
        &self,
        y: &DVector<F::MixedType>,
        g: &DVector<F::MixedType>,
        τ: F,
        x: &mut DVector<F::MixedType>,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> usize {
        let inner_tolerance = ε * config.inner.tolerance_mult;
        let inner_it = config.inner.iterator_options.stop_target(inner_tolerance);
        l1squared_unconstrained(y, g, τ * self.α(), 1.0, x, &config.inner, inner_it)
    }

    fn find_tolerance_violation_slack<M>(
        &self,
        d: &mut M,
        τ: F,
        ε: F,
        skip_by_rough_check: bool,
        config: &InsertionConfig<F>,
        slack: F,
    ) -> Option<(Loc<N, F>, F, bool)>
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let τα = τ * self.α();
        let keep_below = τα + slack + ε;
        let keep_above = -(τα + slack) - ε;
        let maximise_above = τα + slack + ε * config.insertion_cutoff_factor;
        let minimise_below = -(τα + slack) - ε * config.insertion_cutoff_factor;
        let refinement_tolerance = ε * config.refinement.tolerance_mult;

        // If preliminary check indicates that we are in bonds, and if it otherwise matches
        // the insertion strategy, skip insertion.
        if skip_by_rough_check && Bounds(keep_above, keep_below).superset(&d.bounds()) {
            None
        } else {
            // If the rough check didn't indicate no insertion needed, find maximising point.
            let mx = d.maximise_above(
                maximise_above,
                refinement_tolerance,
                config.refinement.max_steps,
            );
            let mi = d.minimise_below(
                minimise_below,
                refinement_tolerance,
                config.refinement.max_steps,
            );

            match (mx, mi) {
                (None, None) => None,
                (Some((ξ, v_ξ)), None) => Some((ξ, v_ξ, keep_below >= v_ξ)),
                (None, Some((ζ, v_ζ))) => Some((ζ, v_ζ, keep_above <= v_ζ)),
                (Some((ξ, v_ξ)), Some((ζ, v_ζ))) => {
                    if v_ξ - τα > τα - v_ζ {
                        Some((ξ, v_ξ, keep_below >= v_ξ))
                    } else {
                        Some((ζ, v_ζ, keep_above <= v_ζ))
                    }
                }
            }
        }
    }

    fn verify_merge_candidate<M>(
        &self,
        d: &mut M,
        μ: &RNDM<N, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
    ) -> bool
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let τα = τ * self.α();
        let refinement_tolerance = ε * config.refinement.tolerance_mult;
        let merge_tolerance = config.merge_tolerance_mult * ε;
        let keep_below = τα + merge_tolerance;
        let keep_above = -τα - merge_tolerance;
        let keep_supp_pos_above = τα - merge_tolerance;
        let keep_supp_neg_below = -τα + merge_tolerance;
        let bnd = d.bounds();

        return ((bnd.lower() >= keep_supp_pos_above && bnd.upper() <= keep_supp_neg_below)
            || μ
                .iter_spikes()
                .all(|&DeltaMeasure { α: β, ref x }| match β.partial_cmp(&0.0) {
                    Some(Greater) => d.apply(x) >= keep_supp_pos_above,
                    Some(Less) => d.apply(x) <= keep_supp_neg_below,
                    _ => true,
                }))
            && (bnd.upper() <= keep_below
                || d.has_upper_bound(
                    keep_below,
                    refinement_tolerance,
                    config.refinement.max_steps,
                ))
            && (bnd.lower() >= keep_above
                || d.has_lower_bound(
                    keep_above,
                    refinement_tolerance,
                    config.refinement.max_steps,
                ));
    }

    fn verify_merge_candidate_radonsq<M>(
        &self,
        d: &mut M,
        μ: &RNDM<N, F>,
        τ: F,
        ε: F,
        config: &InsertionConfig<F>,
        radon_μ: &RNDM<N, F>,
    ) -> bool
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let τα = τ * self.α();
        let refinement_tolerance = ε * config.refinement.tolerance_mult;
        let merge_tolerance = config.merge_tolerance_mult * ε;
        let slack = radon_μ.norm(Radon);
        let bnd = d.bounds();

        return {
            μ.both_matching(radon_μ).all(|(α, rα, x)| {
                let v = d.apply(x);
                let (l1, u1) = match α.partial_cmp(&0.0).unwrap_or(Equal) {
                    Greater => (τα, τα),
                    Equal => (-τα, τα),
                    Less => (-τα, -τα),
                };
                let (l2, u2) = match rα.partial_cmp(&0.0).unwrap_or(Equal) {
                    Greater => (slack, slack),
                    Equal => (-slack, slack),
                    Less => (-slack, -slack),
                };
                (l1 + l2 - merge_tolerance <= v) && (v <= u1 + u2 + merge_tolerance)
            })
        } && {
            let keep_below = τα + slack + merge_tolerance;
            bnd.upper() <= keep_below
                || d.has_upper_bound(
                    keep_below,
                    refinement_tolerance,
                    config.refinement.max_steps,
                )
        } && {
            let keep_above = -τα - slack - merge_tolerance;
            bnd.lower() >= keep_above
                || d.has_lower_bound(
                    keep_above,
                    refinement_tolerance,
                    config.refinement.max_steps,
                )
        };
    }

    fn target_bounds(&self, τ: F, ε: F) -> Option<Bounds<F>> {
        let τα = τ * self.α();
        Some(Bounds(-τα - ε, τα + ε))
    }

    fn tolerance_scaling(&self) -> F {
        self.α()
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F: Float + ToNalgebraRealField, const N: usize> SlidingRegTerm<Loc<N, F>, F>
    for RadonRegTerm<F>
{
    fn goodness<M>(
        &self,
        d: &mut M,
        _μ: &RNDM<N, F>,
        y: &Loc<N, F>,
        z: &Loc<N, F>,
        τ: F,
        ε: F,
        _config: &InsertionConfig<F>,
    ) -> F
    where
        M: MinMaxMapping<Loc<N, F>, F>,
    {
        let α = self.α();
        let w = |x| {
            let dx = d.apply(x);
            ((-ε + dx) / (τ * α)).max(1.0.min(ε + dx) / (τ * α))
        };
        w(z) - w(y)
    }

    fn radon_norm_bound(&self, b: F) -> F {
        b / self.α()
    }
}