varpro 0.14.0

use crate::prelude::*;
use crate::problem::SeparableProblem;
use crate::util::Weights;
use nalgebra::{ComplexField, DMatrix, Dyn, OMatrix, OVector, RealField, Scalar};

use num_traits::{float::TotalOrder, Float, Zero};
use std::ops::Mul;
use thiserror::Error as ThisError;

use super::{MultiRhs, RhsType, SingleRhs};

/// Errors pertaining to use errors of the `SeparableProblemBuilder`
#[derive(Debug, Clone, ThisError, PartialEq, Eq)]
#[allow(missing_docs)]
pub enum SeparableProblemBuilderError {
    /// the data for y variable was not given to the builder
    #[error("Right hand side(s) not provided")]
    YDataMissing,

    /// x and y vector have different lengths
    #[error(
        "Vectors x and y must have same lengths. Given x length = {} and y length = {}",
        x_length,
        y_length
    )]
    InvalidLengthOfData { x_length: usize, y_length: usize },

    /// the provided x and y vectors must not have zero elements
    #[error("x or y must have nonzero number of elements.")]
    ZeroLengthVector,

    /// the model has a different number of parameters than the provided initial guesses
    #[error(
        "Initial guess vector must have same length as parameters. Model has {} parameters and {} initial guesses were provided.",
        model_count,
        provided_count
    )]
    InvalidParameterCount {
        model_count: usize,
        provided_count: usize,
    },

    /// y vector and weights have different lengths
    #[error("The weights must have the same length as the data y.")]
    InvalidLengthOfWeights,
}

/// A builder structure to create a [SeparableProblem](super::SeparableProblem), which can be used for
/// fitting a separable model to data.
/// # Example
/// The following code shows how to create an unweighted least squares problem to fit the separable model
/// `$\vec{f}(\vec{x},\vec{\alpha},\vec{c})$` given by `model` to data `$\vec{y}$` (given as `y`) using the
/// independent variable `$\vec{x}$` (given as `x`). Furthermore we need an initial guess `params`
/// for the nonlinear model parameters `$\vec{\alpha}$`
/// ```rust
/// # use nalgebra::{DVector, Scalar};
/// # use varpro::model::SeparableModel;
/// # use varpro::problem::*;
/// # fn model(model : SeparableModel<f64>,y: DVector<f64>) {
///   let problem = SeparableProblemBuilder::new(model)
///                 .observations(y)
///                 .build()
///                 .unwrap();
/// # }
/// ```
///
/// # Building a Model
///
/// A new builder is constructed with the [new](SeparableProblemBuilder::new) constructor. It must be filled with
/// content using the methods described in the following. After all mandatory fields have been filled,
/// the [build](SeparableProblemBuilder::build) method can be called. This returns a [Result](std::result::Result)
/// type that contains the finished model iff all mandatory fields have been set with valid values. Otherwise
/// it contains an error variant.
///
/// ## Multiple Right Hand Sides
///
/// We can also construct a problem with multiple right hand sides, using the
/// [mrhs](SeparableProblemBuilder::mrhs) constructor, see [SeparableProblem] for
/// additional details.
#[derive(Clone)]
#[allow(non_snake_case)]
pub struct SeparableProblemBuilder<Model, Rhs: RhsType>
where
    Model::ScalarType: Scalar + ComplexField + Copy,
    <Model::ScalarType as ComplexField>::RealField: Float,
    Model: SeparableNonlinearModel,
{
    /// Required: the data `$\vec{y}(\vec{x})$` that we want to fit
    Y: Option<DMatrix<Model::ScalarType>>,
    /// Required: the model to be fitted to the data
    separable_model: Model,
    /// Optional: weights to be applied to the data for weightes least squares
    /// if no weights are given, the problem is unweighted, i.e. the same as if
    /// all weights were 1.
    /// Must have the same length as x and y.
    weights: Weights<Model::ScalarType, Dyn>,
    phantom: std::marker::PhantomData<Rhs>,
}

impl<Model> SeparableProblemBuilder<Model, SingleRhs>
where
    Model::ScalarType: Scalar + ComplexField + Zero + Copy,
    <Model::ScalarType as ComplexField>::RealField: Float,
    <<Model as SeparableNonlinearModel>::ScalarType as ComplexField>::RealField:
        Mul<Model::ScalarType, Output = Model::ScalarType> + Float,
    Model: SeparableNonlinearModel,
{
    /// Create a new builder based on the given model for a problem
    /// with a **single right hand side**. This is the standard use case,
    /// where the data is a vector that is fitted by the model.
    pub fn new(model: Model) -> Self {
        Self {
            Y: None,
            separable_model: model,
            weights: Weights::default(),
            phantom: Default::default(),
        }
    }
}

impl<Model> SeparableProblemBuilder<Model, SingleRhs>
where
    Model::ScalarType: Scalar + ComplexField + Zero + Copy,
    <Model::ScalarType as ComplexField>::RealField: Float,
    <<Model as SeparableNonlinearModel>::ScalarType as ComplexField>::RealField:
        Mul<Model::ScalarType, Output = Model::ScalarType> + Float,
    Model: SeparableNonlinearModel,
{
    /// **Mandatory**: Set the data which we want to fit: since this
    /// is called on a model builder for problems with **single right hand sides**,
    /// this is a column vector `$\vec{y}=\vec{y}(\vec{x})$` containing the
    /// values we want to fit with the model.
    ///
    /// The length of `$\vec{y}` and the output dimension of the model must be
    /// the same.
    pub fn observations(self, observed: OVector<Model::ScalarType, Dyn>) -> Self {
        let nrows = observed.nrows();
        Self {
            Y: Some(observed.reshape_generic(Dyn(nrows), Dyn(1))),
            ..self
        }
    }
}

impl<Model> SeparableProblemBuilder<Model, MultiRhs>
where
    Model::ScalarType: Scalar + ComplexField + Zero + Copy,
    <Model::ScalarType as ComplexField>::RealField: Float,
    <<Model as SeparableNonlinearModel>::ScalarType as ComplexField>::RealField:
        Mul<Model::ScalarType, Output = Model::ScalarType> + Float,
    Model: SeparableNonlinearModel,
{
    /// Create a new builder based on the given model
    /// for a problem with **multiple right hand sides** and perform a global
    /// fit. This uses single threaded calculations.
    ///
    /// That means the observations are expected to be a matrix, where the
    /// columns correspond to the individual observations.
    ///
    /// For a set of observations `$\vec{y}_1,\dots,\vec{y}_S$` (column vectors) we
    /// now have to pass a _matrix_ `$Y$` of observations, rather than a single
    /// vector to the builder. As explained above, the resulting matrix would look
    /// like this.
    ///
    /// ```math
    /// \boldsymbol{Y}=\left(\begin{matrix}
    ///  \vert &  & \vert \\
    ///  \vec{y}_1 &  \dots & \vec{y}_S \\
    ///  \vert &  & \vert \\
    /// \end{matrix}\right)
    /// ```
    ///
    /// The nonlinear parameters will be optimized across all the observations
    /// globally, but the best linear coefficients are calculated for each observation
    /// individually. Hence, the latter also become a matrix `$C$`, where the columns
    /// correspond to the linear coefficients of the observation in the same column.
    ///
    /// ```math
    /// \boldsymbol{C}=\left(\begin{matrix}
    ///  \vert &  & \vert \\
    ///  \vec{c}_1 &  \dots & \vec{c}_S \\
    ///  \vert &  & \vert \\
    /// \end{matrix}\right)
    /// ```
    ///
    /// The (column) vector of linear coefficients `$\vec{c}_j$` is for the observation
    /// `$\vec{y}_j$` in the same column.
    pub fn mrhs(model: Model) -> Self {
        Self {
            Y: None,
            separable_model: model,
            weights: Weights::default(),
            phantom: Default::default(),
        }
    }
}

impl<Model> SeparableProblemBuilder<Model, MultiRhs>
where
    Model::ScalarType: Scalar + ComplexField + Zero + Copy,
    <Model::ScalarType as ComplexField>::RealField: Float,
    <<Model as SeparableNonlinearModel>::ScalarType as ComplexField>::RealField:
        Mul<Model::ScalarType, Output = Model::ScalarType> + Float,
    Model: SeparableNonlinearModel,
{
    /// **Mandatory**: Set the data which we want to fit: This is either a single vector
    /// `$\vec{y}=\vec{y}(\vec{x})$` or a matrix `$\boldsymbol{Y}$` of multiple
    /// vectors. In the former case this corresponds to fitting a single right hand side,
    /// in the latter case, this corresponds to global fitting of a problem with
    /// multiple right hand sides.
    /// The length of `$\vec{x}$` and the number of _rows_ in the data must
    /// be the same.
    pub fn observations(self, observed: OMatrix<Model::ScalarType, Dyn, Dyn>) -> Self {
        Self {
            Y: Some(observed),
            ..self
        }
    }
}

impl<Model, Rhs: RhsType> SeparableProblemBuilder<Model, Rhs>
where
    Model::ScalarType: Scalar + ComplexField + Zero + Copy,
    <Model::ScalarType as ComplexField>::RealField: Float,
    <<Model as SeparableNonlinearModel>::ScalarType as ComplexField>::RealField:
        Mul<Model::ScalarType, Output = Model::ScalarType> + Float,
    Model: SeparableNonlinearModel,
{
    /// **Optional** Add diagonal weights to the problem (meaning data points are statistically
    /// independent). If this is not given, the problem is unweighted, i.e. each data point has
    /// unit weight.
    ///
    /// **Note** The weighted residual is calculated as `$||W(\vec{y}-\vec{f}(\vec{\alpha})||^2$`, so
    /// to make weights that have a statistical meaning, the diagonal elements of the weight matrix should be
    /// set to `w_{jj} = 1/\sigma_j` where `$\sigma_j$` is the (estimated) standard deviation associated with
    /// data point `$y_j$`.
    pub fn weights(self, weights: OVector<Model::ScalarType, Dyn>) -> Self {
        Self {
            weights: Weights::diagonal(weights),
            ..self
        }
    }

    /// build the least squares problem from the builder.
    /// # Prerequisites
    /// * All mandatory parameters have been set (see individual builder methods for details)
    /// * `$\vec{x}$` and `$\vec{y}$` have the same number of elements
    /// * `$\vec{x}$` and `$\vec{y}$` have a nonzero number of elements
    /// * the length of the initial guesses vector is the same as the number of model parameters
    /// # Returns
    /// If all prerequisites are fulfilled, returns a [SeparableProblem](super::SeparableProblem) with the given
    /// content and the parameters set to the initial guess. Otherwise returns an error variant.
    #[allow(non_snake_case)]
    pub fn build(self) -> Result<SeparableProblem<Model, Rhs>, SeparableProblemBuilderError>
    where
        Model::ScalarType: Float + RealField + TotalOrder,
    {
        // and assign the defaults to the values we don't have
        let Y = self.Y.ok_or(SeparableProblemBuilderError::YDataMissing)?;
        let model = self.separable_model;
        let weights = self.weights;

        // now do some sanity checks for the values and return
        // an error if they do not pass the test
        let x_len: usize = model.output_len();
        if x_len == 0 || Y.is_empty() {
            return Err(SeparableProblemBuilderError::ZeroLengthVector);
        }

        if x_len != Y.nrows() {
            return Err(SeparableProblemBuilderError::InvalidLengthOfData {
                x_length: x_len,
                y_length: Y.nrows(),
            });
        }

        if !weights.is_size_correct_for_data_length(Y.nrows()) {
            return Err(SeparableProblemBuilderError::InvalidLengthOfWeights);
        }

        let Y_w = &weights * Y;

        Ok(SeparableProblem {
            Y_w,
            model,
            weights,
            phantom: Default::default(),
        })
    }
}
// make available for testing and doc tests
#[cfg(any(test, doctest))]
mod test;