varpro 0.14.0

use crate::prelude::*;
use crate::util::Weights;
use nalgebra::Dyn;
use nalgebra::{ComplexField, DMatrix, MatrixView, Scalar, VectorView};

mod builder;

pub use builder::SeparableProblemBuilder;
pub use builder::SeparableProblemBuilderError;

/// trait describing the type of right hand side for the problem, meaning either
/// a single right hand side or multiple right hand sides. The latter implies
/// global fitting.
pub trait RhsType: std::fmt::Debug {}

/// This type indicates that the associated problem has a single (vector) right hand side.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct SingleRhs;

/// This type indicates that the associated problem has multiple right hand sides
/// and thus performs global fitting for the nonlinear parameters.
#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
pub struct MultiRhs;

impl RhsType for MultiRhs {}
impl RhsType for SingleRhs {}

/// This is the problem of fitting the separable model to data in a form that the
/// [levenberg_marquardt](https://crates.io/crates/levenberg-marquardt) crate can use it to
/// perform the least squares fit.
///
/// # Construction
///
/// Use the [`SeparableProblemBuilder`] to create an instance of a
/// levmar problem.
///
/// # Usage
///
/// After obtaining an instance of [`SeparableProblem`] we can pass it to the [`LevMarSolver`](crate::solvers::levmar::LevMarSolver)
/// which uses the Levenberg-Marquardt algorithm from the levenberg_marquardt crate for minimization.
/// Refer to the documentation of the [levenberg_marquardt](https://crates.io/crates/levenberg-marquardt)
/// crate for an overview of the algorithm. A usage example is provided in this crate's documentation as well.
///
/// # Right Hand Sides: Single vs Multiple
///
/// The problem is generic over the `Rhs` type parameter which indicates whether the
/// problem fits a single ([`SingleRhs`]) or multiple ([`MultiRhs`]) right
/// hand sides. This is decided during the building process. The underlying
/// math does not change, but the interface changes to use vectors for coefficients
/// and data in case of a single right hand side. For multiple right hand sides,
/// the coefficients and the data are matrices corresponding to columns of
/// coefficient vectors and data vectors respectively.
#[derive(Clone)]
#[allow(non_snake_case)]
pub struct SeparableProblem<Model, Rhs: RhsType>
where
    Model: SeparableNonlinearModel,
    Model::ScalarType: Scalar + ComplexField + Copy,
{
    /// The *weighted* data matrix to which to fit the model `$\boldsymbol{Y}_w$`.
    /// It is a matrix so it can accommodate multiple right hand sides. If
    /// the problem has only a single right hand side (SingleRhs), this is just
    /// a matrix with one column. The underlying math does not change in either case.
    /// **Attention:** The data matrix is weighted with the weights if some weights
    /// were provided (otherwise it is unweighted)
    pub(crate) Y_w: DMatrix<Model::ScalarType>,
    /// a reference to the separable model we are trying to fit to the data
    pub(crate) model: Model,
    /// the weights of the data. If none are given, the data is not weighted
    /// If weights were provided, the builder has checked that the weights have the
    /// correct dimension for the data
    pub(crate) weights: Weights<Model::ScalarType, Dyn>,
    phantom: std::marker::PhantomData<Rhs>,
}

impl<Model, Rhs: RhsType> std::fmt::Debug for SeparableProblem<Model, Rhs>
where
    Model: SeparableNonlinearModel,
    Model::ScalarType: Scalar + ComplexField + Copy,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SeparableProblem")
            .field("y_w", &self.Y_w)
            .field("model", &"/* omitted */")
            .field("weights", &self.weights)
            .finish()
    }
}

impl<Model> SeparableProblem<Model, MultiRhs>
where
    Model: SeparableNonlinearModel,
    Model::ScalarType: Scalar + ComplexField + Copy,
{
    /// The weighted data matrix `$\boldsymbol{Y}_w$` to which to fit the model. Note
    /// that the weights are already applied to the data matrix and this
    /// is not the original data vector.
    ///
    /// This method is for fitting multiple right hand sides, hence the data
    /// matrix is a matrix that contains the right hand sides as columns.
    pub fn weighted_data(&self) -> MatrixView<'_, Model::ScalarType, Dyn, Dyn> {
        self.Y_w.as_view()
    }
}

impl<Model> SeparableProblem<Model, SingleRhs>
where
    Model: SeparableNonlinearModel,
    Model::ScalarType: Scalar + ComplexField + Copy,
{
    /// The weighted data vector `$\vec{y}_w$` to which to fit the model. Note
    /// that the weights are already applied to the data vector and this
    /// is not the original data vector.
    ///
    /// This method is for fitting a single right hand side, hence the data
    /// is a single column vector.
    pub fn weighted_data(&self) -> VectorView<'_, Model::ScalarType, Dyn> {
        debug_assert_eq!(
            self.Y_w.ncols(),
            1,
            "data matrix must have exactly one column for single right hand side. This indicates a programming error in the library."
        );
        self.Y_w.as_view()
    }
}

impl<Model, Rhs: RhsType> SeparableProblem<Model, Rhs>
where
    Model: SeparableNonlinearModel,
    Model::ScalarType: Scalar + ComplexField + Copy,
{
    /// access the contained model immutably
    pub fn model(&self) -> &Model {
        &self.model
    }

    /// get the weights of the data for the fitting problem
    pub fn weights(&self) -> &Weights<Model::ScalarType, Dyn> {
        &self.weights
    }
}