fista.h

Go to the documentation of this file.

/*
 * Copyright 2021 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 <cmath>
#include <limits>
#include <utility>
 
#include "num_collect/base/exception.h"
#include "num_collect/base/index_type.h"
#include "num_collect/base/precondition.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/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 fista_tag =
    logging::log_tag_view("num_collect::regularization::fista");
 
template <typename Coeff, typename Data>
class fista
    : public iterative_regularized_solver_base<fista<Coeff, Data>, Data> {
public:
    using this_type = fista<Coeff, 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;
 
    fista() : base_type(fista_tag) {}
 
    void compute(const Coeff& coeff, const Data& data) {
        coeff_ = &coeff;
        data_ = &data;
        inv_max_eigen_ = static_cast<scalar_type>(1) /
            impl::approximate_max_eigen_aat(coeff);
        NUM_COLLECT_LOG_TRACE(
            this->logger(), "inv_max_eigen={}", inv_max_eigen_);
    }
 
    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.");
 
        iterations_ = 0;
        t_ = static_cast<scalar_type>(1);
        y_ = solution;
        update_ = std::numeric_limits<scalar_type>::infinity();
    }
 
    void iterate(const scalar_type& param, data_type& solution) {
        const scalar_type t_before = t_;
        using std::sqrt;
        t_ = static_cast<scalar_type>(0.5) *        // NOLINT
            (static_cast<scalar_type>(1) +          // NOLINT
                sqrt(static_cast<scalar_type>(1) +  // NOLINT
                    static_cast<scalar_type>(4)     // NOLINT
                        * t_before * t_before));
        const scalar_type coeff_update =
            (t_before - static_cast<scalar_type>(1)) / t_;
 
        const scalar_type twice_step = inv_max_eigen_;
        const scalar_type step = static_cast<scalar_type>(0.5) * twice_step;
        const scalar_type trunc_thresh = param * step;
 
        residual_ = -(*data_);
        auto update_sum2 = static_cast<scalar_type>(0);
#pragma omp parallel default(shared)
        {
            const index_type size = solution.size();
            data_type temp_res = Data::Zero(residual_.size());
#pragma omp for nowait
            for (index_type i = 0; i < size; ++i) {
                using std::abs;
                if (abs(y_(i)) > static_cast<scalar_type>(0)) {
                    temp_res += y_(i) * coeff_->col(i);
                }
            }
#pragma omp critical
            residual_ += temp_res;
#pragma omp barrier
 
#pragma omp for reduction(+ : update_sum2)
            for (index_type i = 0; i < size; ++i) {
                scalar_type cur_next_sol =
                    y_(i) - twice_step * coeff_->col(i).dot(residual_);
 
                if (cur_next_sol > trunc_thresh) {
                    cur_next_sol = cur_next_sol - trunc_thresh;
                } else if (cur_next_sol < -trunc_thresh) {
                    cur_next_sol = cur_next_sol + trunc_thresh;
                } else {
                    cur_next_sol = static_cast<scalar_type>(0);
                }
 
                const scalar_type current_update = cur_next_sol - solution(i);
                update_sum2 += current_update * current_update;
 
                y_(i) = cur_next_sol + coeff_update * current_update;
                solution(i) = cur_next_sol;
            }
        }
 
        update_ = sqrt(update_sum2);
        ++iterations_;
    }
 
    [[nodiscard]] auto is_stop_criteria_satisfied(
        const data_type& solution) const -> bool {
        return (iterations() > max_iterations()) ||
            (update() < tol_update_rate() * solution.norm());
    }
 
    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>(
            "Update", &this_type::update);
        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 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 scalar_type max_sol_est =
            (coeff_->transpose() * (*data_)).cwiseAbs().maxCoeff();
        NUM_COLLECT_LOG_TRACE(this->logger(), "max_sol_est={}", max_sol_est);
        constexpr auto tol_update_coeff_multiplier =
            static_cast<scalar_type>(10);
        return {max_sol_est *
                std::max(impl::weak_coeff_min_param<scalar_type>,
                    tol_update_coeff_multiplier * tol_update_rate_),
            max_sol_est * impl::weak_coeff_max_param<scalar_type>};
    }
 
    [[nodiscard]] auto iterations() const noexcept -> index_type {
        return iterations_;
    }
 
    [[nodiscard]] auto update() const noexcept -> scalar_type {
        return update_;
    }
 
    [[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) -> fista& {
        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) -> fista& {
        NUM_COLLECT_PRECONDITION(value > static_cast<scalar_type>(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 data_type* data_{nullptr};
 
    scalar_type inv_max_eigen_{};
 
    index_type iterations_{};
 
    scalar_type t_{};
 
    data_type y_{};
 
    data_type residual_{};
 
    scalar_type update_{};
 
    static constexpr index_type default_max_iterations = 1000;
 
    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};
};
 
}  // namespace num_collect::regularization