tgv2_admm.h

Go to the documentation of this file.

/*
 * Copyright 2025 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/index_type.h"
#include "num_collect/base/precondition.h"
#include "num_collect/linear/impl/operator_conjugate_gradient.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 tgv2_admm_tag =
    logging::log_tag_view("num_collect::regularization::tgv2_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 tgv2_admm : public iterative_regularized_solver_base<
                      tgv2_admm<Coeff, DerivativeMatrix, Data>, Data> {
public:
    using this_type = tgv2_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;

    tgv2_admm() : base_type(tgv2_admm_tag) {
        this->configure_child_algorithm_logger_if_exists(
            conjugate_gradient_solution_);
        this->configure_child_algorithm_logger_if_exists(conjugate_gradient_z_);
    }

    void compute(const Coeff& coeff,
        const DerivativeMatrix& first_derivative_matrix,
        const DerivativeMatrix& second_derivative_matrix, const Data& data) {
        coeff_ = &coeff;
        first_derivative_matrix_ = &first_derivative_matrix;
        second_derivative_matrix_ = &second_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_ != nullptr, this->logger(),
            "Coefficient matrix is not set.");
        NUM_COLLECT_PRECONDITION(first_derivative_matrix_ != nullptr,
            this->logger(), "First order derivative matrix is not set.");
        NUM_COLLECT_PRECONDITION(second_derivative_matrix_ != nullptr,
            this->logger(), "Second order derivative matrix is not set.");
        NUM_COLLECT_PRECONDITION(
            data_ != nullptr, this->logger(), "Data vector is not set.");
 
        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(
            first_derivative_matrix_->cols() == solution.rows(), this->logger(),
            "The number of columns in the first order derivative matrix must "
            "match the number of rows in solution vector.");
        NUM_COLLECT_PRECONDITION(second_derivative_matrix_->cols() ==
                first_derivative_matrix_->rows(),
            this->logger(),
            "The number of columns in the second order derivative matrix must "
            "match the number of rows in the first order derivative matrix.");
 
        iterations_ = 0;
 
        // Set variables to one of feasible solutions.
        z_ = data_type::Zero(first_derivative_matrix_->rows());
        s_ = (*first_derivative_matrix_) * solution;
        t_ = data_type::Zero(second_derivative_matrix_->rows());
        p_ = data_type::Zero(first_derivative_matrix_->rows());
        u_ = data_type::Zero(second_derivative_matrix_->rows());
 
        // Pre-allocate temporary vectors.
        temp_solution_ = data_type::Zero(solution.rows());
        temp_data_ = data_type::Zero(data_->rows());
        temp_z_ = data_type::Zero(z_.rows());
        temp_t_ = data_type::Zero(t_.rows());
        previous_solution_ = data_type::Zero(solution.rows());
        previous_z_ = data_type::Zero(z_.rows());
        previous_s_ = data_type::Zero(s_.rows());
        previous_t_ = data_type::Zero(t_.rows());
        p_update_ = data_type::Zero(p_.rows());
        u_update_ = data_type::Zero(u_.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_solution_.tolerance_rate(
            conjugate_gradient_tolerance_rate);
        conjugate_gradient_z_.tolerance_rate(conjugate_gradient_tolerance_rate);
    }

    void iterate(const scalar_type& param, data_type& solution) {
        update_rate_ = static_cast<scalar_type>(0);
 
        update_solution(param, solution);
        update_z(param, solution);
        update_s(param, solution);
        update_t(param, solution);
        update_p(param, solution);
        update_u(param, solution);
 
        ++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 (*first_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 =
            (*first_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) -> this_type& {
        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) -> this_type& {
        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:
    void update_solution(const scalar_type& param, data_type& solution) {
        (void)param;  // Not used for update of solution.
 
        temp_z_ = -p_ + constraint_coeff_ * z_ + constraint_coeff_ * s_;
        temp_solution_.noalias() =
            static_cast<scalar_type>(2) * (*coeff_).transpose() * (*data_);
        temp_solution_.noalias() +=
            (*first_derivative_matrix_).transpose() * temp_z_;
        previous_solution_ = solution;
        conjugate_gradient_solution_.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_z_.noalias() = (*first_derivative_matrix_) * target;
                result.noalias() += constraint_coeff_ *
                    (*first_derivative_matrix_).transpose() * temp_z_;
            },
            temp_solution_, solution);
        update_rate_ += (solution - previous_solution_).norm() /
            (solution.norm() + std::numeric_limits<scalar_type>::epsilon());
        residual_.noalias() = (*coeff_) * solution;
        residual_ -= (*data_);
    }

    void update_z(const scalar_type& param, const data_type& solution) {
        (void)param;  // Not used for update of z_.
 
        temp_z_ = p_;
        temp_z_.noalias() -= (*second_derivative_matrix_).transpose() * u_;
        temp_z_.noalias() +=
            constraint_coeff_ * (*first_derivative_matrix_) * solution;
        temp_z_.noalias() -= constraint_coeff_ * s_;
        temp_z_.noalias() +=
            constraint_coeff_ * (*second_derivative_matrix_).transpose() * t_;
        previous_z_ = z_;
        conjugate_gradient_z_.solve(
            [this](const data_type& target, data_type& result) {
                result.noalias() = constraint_coeff_ * target;
                temp_t_.noalias() = (*second_derivative_matrix_) * target;
                result.noalias() += constraint_coeff_ *
                    (*second_derivative_matrix_).transpose() * temp_t_;
            },
            temp_z_, z_);
        update_rate_ += (z_ - previous_z_).norm() /
            (z_.norm() + std::numeric_limits<scalar_type>::epsilon());
    }

    void update_s(const scalar_type& param, const data_type& solution) {
        previous_s_ = s_;
        s_.noalias() = (*first_derivative_matrix_) * solution;
        s_ += -z_ + p_ / constraint_coeff_;
        impl::apply_shrinkage_operator(s_, param / constraint_coeff_);
        update_rate_ += (s_ - previous_s_).norm() /
            (s_.norm() + std::numeric_limits<scalar_type>::epsilon());
    }

    void update_t(const scalar_type& param, const data_type& solution) {
        (void)solution;  // Not used for update of t_.
 
        previous_t_ = t_;
        t_.noalias() = (*second_derivative_matrix_) * z_;
        t_ += u_ / constraint_coeff_;
        impl::apply_shrinkage_operator(
            t_, param / constraint_coeff_ / second_derivative_ratio_);
        update_rate_ += (t_ - previous_t_).norm() /
            (t_.norm() + std::numeric_limits<scalar_type>::epsilon());
    }

    void update_p(const scalar_type& param, const data_type& solution) {
        (void)param;  // Not used for update of p_.
 
        p_update_.noalias() =
            constraint_coeff_ * (*first_derivative_matrix_) * solution;
        p_update_ += constraint_coeff_ * (-z_ - s_);
        p_ += p_update_;
        update_rate_ += p_update_.norm() /
            (p_.norm() + std::numeric_limits<scalar_type>::epsilon());
    }

    void update_u(const scalar_type& param, const data_type& solution) {
        (void)param;     // Not used for update of u_.
        (void)solution;  // Not used for update of u_.
 
        u_update_.noalias() =
            constraint_coeff_ * (*second_derivative_matrix_) * z_;
        u_update_ -= constraint_coeff_ * t_;
        u_ += u_update_;
        update_rate_ += u_update_.norm() /
            (u_.norm() + std::numeric_limits<scalar_type>::epsilon());
    }

    const coeff_type* coeff_{nullptr};

    const derivative_matrix_type* first_derivative_matrix_{nullptr};

    const derivative_matrix_type* second_derivative_matrix_{nullptr};

    const data_type* data_{nullptr};

    index_type iterations_{};

    data_type z_{};

    data_type s_{};

    data_type t_{};

    data_type p_{};

    data_type u_{};

    data_type temp_solution_{};

    data_type temp_data_{};

    data_type temp_z_{};

    data_type temp_t_{};

    data_type previous_solution_{};

    data_type previous_z_{};

    data_type previous_s_{};

    data_type previous_t_{};

    data_type p_update_{};

    data_type u_update_{};

    data_type residual_{};

    scalar_type update_rate_{};

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

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

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

    scalar_type second_derivative_ratio_{default_second_derivative_ratio};

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

    scalar_type constraint_coeff_{default_constraint_coeff};

    static constexpr index_type default_max_iterations = 100000;

    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