tv_admm.h

Go to the documentation of this file.

/*
 * Copyright 2024 MusicScience37 (Kenta Kabashima)
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <algorithm>
#include <limits>
#include <utility>
 
#include "num_collect/base/concepts/dense_matrix.h"
#include "num_collect/base/concepts/dense_vector.h"
#include "num_collect/base/concepts/sparse_matrix.h"
#include "num_collect/base/exception.h"
#include "num_collect/base/index_type.h"
#include "num_collect/base/precondition.h"
#include "num_collect/linear/impl/operator_conjugate_gradient.h"
#include "num_collect/logging/iterations/iteration_logger.h"
#include "num_collect/logging/log_tag_view.h"
#include "num_collect/logging/logging_macros.h"
#include "num_collect/regularization/impl/apply_shrinkage_operator.h"
#include "num_collect/regularization/impl/approximate_max_eigen_aat.h"
#include "num_collect/regularization/impl/weak_coeff_param.h"  // IWYU pragma: keep
#include "num_collect/regularization/iterative_regularized_solver_base.h"
 
namespace num_collect::regularization {

constexpr auto tv_admm_tag =
    logging::log_tag_view("num_collect::regularization::tv_admm");

template <typename Coeff, typename DerivativeMatrix,
    base::concepts::dense_vector Data>
    requires((base::concepts::sparse_matrix<Coeff> ||
                 base::concepts::dense_matrix<Coeff>) &&
        (base::concepts::sparse_matrix<DerivativeMatrix> ||
            base::concepts::dense_matrix<DerivativeMatrix>))
class tv_admm : public iterative_regularized_solver_base<
                    tv_admm<Coeff, DerivativeMatrix, Data>, Data> {
public:
    using this_type = tv_admm<Coeff, DerivativeMatrix, Data>;

    using base_type = iterative_regularized_solver_base<this_type, Data>;
 
    using typename base_type::data_type;
    using typename base_type::scalar_type;

    using coeff_type = Coeff;

    using derivative_matrix_type = DerivativeMatrix;

    tv_admm() : base_type(tv_admm_tag) {
        this->configure_child_algorithm_logger_if_exists(conjugate_gradient_);
    }

    void compute(const Coeff& coeff, const DerivativeMatrix& derivative_matrix,
        const Data& data) {
        coeff_ = &coeff;
        derivative_matrix_ = &derivative_matrix;
        data_ = &data;
        // Sizes will be checked in init.
    }

    void init(const scalar_type& param, data_type& solution) {
        (void)param;
 
        NUM_COLLECT_PRECONDITION(coeff_->rows() == data_->rows(),
            this->logger(),
            "Coefficient matrix and data vector must have the same number of "
            "rows.");
        NUM_COLLECT_PRECONDITION(coeff_->cols() == solution.rows(),
            this->logger(),
            "The number of columns in the coefficient matrix must match the "
            "number of rows in solution vector.");
        NUM_COLLECT_PRECONDITION(data_->cols() == solution.cols(),
            this->logger(),
            "Data and solution must have the same number of columns.");
        NUM_COLLECT_PRECONDITION(derivative_matrix_->cols() == solution.rows(),
            this->logger(),
            "The number of columns in the derivative matrix must match the "
            "number of rows in solution vector.");
 
        iterations_ = 0;
        derivative_ = (*derivative_matrix_) * solution;
        lagrange_multiplier_ = data_type::Zero(derivative_matrix_->rows());
        temp_solution_ = solution;
        temp_data_ = data_type::Zero(data_->rows());
        temp_derivative_ = data_type::Zero(derivative_matrix_->rows());
        residual_ = (*coeff_) * solution - (*data_);
        update_rate_ = std::numeric_limits<scalar_type>::infinity();
 
        const scalar_type conjugate_gradient_tolerance_rate =
            rate_of_cg_tol_rate_to_tol_update_rate_ * tol_update_rate_;
        conjugate_gradient_.tolerance_rate(conjugate_gradient_tolerance_rate);
    }

    void iterate(const scalar_type& param, data_type& solution) {
        // Update solution.
        temp_solution_.noalias() =
            static_cast<scalar_type>(2) * (*coeff_).transpose() * (*data_);
        temp_solution_.noalias() -=
            (*derivative_matrix_).transpose() * lagrange_multiplier_;
        temp_solution_.noalias() += derivative_constraint_coeff_ *
            (*derivative_matrix_).transpose() * derivative_;
        previous_solution_ = solution;
        conjugate_gradient_.solve(
            [this](const data_type& target, data_type& result) {
                temp_data_.noalias() = (*coeff_) * target;
                result.noalias() = static_cast<scalar_type>(2) *
                    (*coeff_).transpose() * temp_data_;
                temp_derivative_.noalias() = (*derivative_matrix_) * target;
                result.noalias() += derivative_constraint_coeff_ *
                    (*derivative_matrix_).transpose() * temp_derivative_;
            },
            temp_solution_, solution);
        update_rate_ = (solution - previous_solution_).norm() /
            (solution.norm() + std::numeric_limits<scalar_type>::epsilon());
        residual_.noalias() = (*coeff_) * solution;
        residual_ -= (*data_);
 
        // Update derivative.
        previous_derivative_ = derivative_;
        derivative_.noalias() = (*derivative_matrix_) * solution;
        derivative_ += lagrange_multiplier_ / derivative_constraint_coeff_;
        const scalar_type derivative_shrinkage_threshold =
            param / derivative_constraint_coeff_;
        impl::apply_shrinkage_operator(
            derivative_, derivative_shrinkage_threshold);
        update_rate_ += (derivative_ - previous_derivative_).norm() /
            (derivative_.norm() + std::numeric_limits<scalar_type>::epsilon());
 
        // Update lagrange multiplier.
        lagrange_multiplier_update_.noalias() =
            derivative_constraint_coeff_ * (*derivative_matrix_) * solution;
        lagrange_multiplier_update_ -=
            derivative_constraint_coeff_ * derivative_;
        lagrange_multiplier_ += lagrange_multiplier_update_;
        update_rate_ += lagrange_multiplier_update_.norm() /
            (lagrange_multiplier_.norm() +
                std::numeric_limits<scalar_type>::epsilon());
 
        ++iterations_;
    }

    [[nodiscard]] auto is_stop_criteria_satisfied(
        const data_type& solution) const -> bool {
        (void)solution;
        return (iterations() > max_iterations()) ||
            (update_rate() < tol_update_rate());
    }

    void configure_iteration_logger(
        logging::iterations::iteration_logger<this_type>& iteration_logger)
        const {
        iteration_logger.template append<index_type>(
            "Iter.", &this_type::iterations);
        iteration_logger.template append<scalar_type>(
            "UpdateRate", &this_type::update_rate);
        iteration_logger.template append<scalar_type>(
            "Res.Rate", &this_type::residual_norm_rate);
    }

    [[nodiscard]] auto residual_norm(const data_type& solution) const
        -> scalar_type {
        return ((*coeff_) * solution - (*data_)).squaredNorm();
    }

    [[nodiscard]] auto regularization_term(const data_type& solution) const
        -> scalar_type {
        return (*derivative_matrix_ * solution).template lpNorm<1>();
    }

    void change_data(const data_type& data) { data_ = &data; }

    void calculate_data_for(const data_type& solution, data_type& data) const {
        data = (*coeff_) * solution;
    }

    [[nodiscard]] auto data_size() const -> index_type { return data_->size(); }

    [[nodiscard]] auto param_search_region() const
        -> std::pair<scalar_type, scalar_type> {
        const data_type approx_order_of_solution = coeff_->transpose() *
            (*data_) / impl::approximate_max_eigen_aat(*coeff_);
        const scalar_type approx_order_of_param =
            (*derivative_matrix_ * approx_order_of_solution)
                .cwiseAbs()
                .maxCoeff();
        NUM_COLLECT_LOG_TRACE(
            this->logger(), "approx_order_of_param={}", approx_order_of_param);
        constexpr auto tol_update_coeff_multiplier =
            static_cast<scalar_type>(10);
        return {approx_order_of_param *
                std::max(impl::weak_coeff_min_param<scalar_type>,
                    tol_update_coeff_multiplier * tol_update_rate_),
            approx_order_of_param * impl::weak_coeff_max_param<scalar_type>};
    }

    [[nodiscard]] auto iterations() const noexcept -> index_type {
        return iterations_;
    }

    [[nodiscard]] auto update_rate() const noexcept -> scalar_type {
        return update_rate_;
    }

    [[nodiscard]] auto residual_norm_rate() const -> scalar_type {
        return residual_.squaredNorm() / data_->squaredNorm();
    }

    [[nodiscard]] auto max_iterations() const -> index_type {
        return max_iterations_;
    }

    auto max_iterations(index_type value) -> tv_admm& {
        NUM_COLLECT_PRECONDITION(value > 0, this->logger(),
            "Maximum number of iterations must be a positive integer.");
        max_iterations_ = value;
        return *this;
    }

    [[nodiscard]] auto tol_update_rate() const -> scalar_type {
        return tol_update_rate_;
    }

    auto tol_update_rate(scalar_type value) -> tv_admm& {
        NUM_COLLECT_PRECONDITION(value > 0, this->logger(),
            "Tolerance of update rate of the solution must be a positive "
            "value.");
        tol_update_rate_ = value;
        return *this;
    }
 
private:
    const coeff_type* coeff_{nullptr};

    const derivative_matrix_type* derivative_matrix_{nullptr};

    const data_type* data_{nullptr};

    index_type iterations_{};

    data_type derivative_{};

    data_type lagrange_multiplier_{};

    data_type temp_solution_{};

    data_type temp_data_{};

    data_type temp_derivative_{};

    data_type previous_solution_{};

    data_type previous_derivative_{};

    data_type lagrange_multiplier_update_{};

    data_type residual_{};

    scalar_type update_rate_{};

    linear::impl::operator_conjugate_gradient<data_type> conjugate_gradient_{};

    static constexpr auto default_derivative_constraint_coeff =
        static_cast<scalar_type>(1);

    scalar_type derivative_constraint_coeff_{
        default_derivative_constraint_coeff};

    static constexpr index_type default_max_iterations = 10000;

    index_type max_iterations_{default_max_iterations};

    static constexpr auto default_tol_update_rate =
        static_cast<scalar_type>(1e-4);

    scalar_type tol_update_rate_{default_tol_update_rate};

    static constexpr auto default_rate_of_cg_tol_rate_to_tol_update_rate =
        static_cast<scalar_type>(1e-2);

    scalar_type rate_of_cg_tol_rate_to_tol_update_rate_{
        default_rate_of_cg_tol_rate_to_tol_update_rate};
};
 
}  // namespace num_collect::regularization