general_spline_equation_solver.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 <cmath>
#include <limits>
 
#include <Eigen/Core>
#include <Eigen/Eigenvalues>
#include <Eigen/QR>
#include <fmt/core.h>  // IWYU pragma: keep
 
#include "num_collect/base/concepts/real_scalar.h"
#include "num_collect/base/exception.h"
#include "num_collect/base/index_type.h"
#include "num_collect/base/precondition.h"
#include "num_collect/logging/logging_macros.h"
#include "num_collect/rbf/kernel_matrix_type.h"
 
namespace num_collect::rbf::impl {

template <typename KernelValue, typename FunctionValue,
    kernel_matrix_type KernelMatrixType, bool UsesGlobalLengthParameter>
class general_spline_equation_solver;

template <base::concepts::real_scalar KernelValue, typename FunctionValue>
class general_spline_equation_solver<KernelValue, FunctionValue,
    kernel_matrix_type::dense, true> {
public:
    using kernel_matrix_type = Eigen::MatrixX<KernelValue>;

    using additional_matrix_type = Eigen::MatrixX<KernelValue>;

    using vector_type = Eigen::VectorX<FunctionValue>;

    using scalar_type = KernelValue;

    general_spline_equation_solver() = default;

    void compute(const kernel_matrix_type& kernel_matrix,
        const additional_matrix_type& additional_matrix,
        const vector_type& data) {
        num_variables_ = kernel_matrix.rows();
        NUM_COLLECT_PRECONDITION(kernel_matrix.cols() == num_variables_,
            "Kernel matrix must be a square matrix.");
        NUM_COLLECT_PRECONDITION(additional_matrix.rows() == num_variables_,
            "Matrix of additional terms must have the same number of rows as "
            "the kernel matrix.");
        num_additional_terms_ = additional_matrix.cols();
        NUM_COLLECT_PRECONDITION(num_variables_ > num_additional_terms_,
            "The number of variables must be larger than the number of "
            "additional terms.");
        kernel_subspace_dimensions_ = num_variables_ - num_additional_terms_;
 
        qr_decomposition_.compute(additional_matrix);
        if (qr_decomposition_.rank() != additional_matrix.cols()) {
            NUM_COLLECT_LOG_AND_THROW(algorithm_failure,
                "The matrix of additional terms must have full "
                "column rank. (columns={}, rand={})",
                additional_matrix.cols(), qr_decomposition_.rank());
        }
        q_matrix_ = qr_decomposition_.householderQ();
 
        transformed_kernel_matrix_ =
            q_matrix_.rightCols(kernel_subspace_dimensions_).transpose() *
            kernel_matrix * q_matrix_.rightCols(kernel_subspace_dimensions_);
 
        eigen_value_decomposition_.compute(
            transformed_kernel_matrix_, Eigen::ComputeEigenvectors);
        data_transformation_matrix_ =
            eigen_value_decomposition_.eigenvectors().transpose() *
            q_matrix_.rightCols(kernel_subspace_dimensions_).transpose();
        spectre_ = data_transformation_matrix_ * data;
 
        kernel_matrix_ = &kernel_matrix;
        data_ = &data;
    }

    void solve(vector_type& kernel_coefficients,
        vector_type& additional_coefficients, scalar_type reg_param) const {
        NUM_COLLECT_PRECONDITION(kernel_matrix_ != nullptr && data_ != nullptr,
            "compute() must be called before solve().");
 
        reg_param = correct_reg_param_if_needed(reg_param);
 
        kernel_coefficients = data_transformation_matrix_.transpose() *
            (eigen_value_decomposition_.eigenvalues().array() + reg_param)
                .inverse()
                .matrix()
                .asDiagonal() *
            spectre_;
 
        additional_coefficients = qr_decomposition_.solve(
            (*data_) - (*kernel_matrix_) * kernel_coefficients);
    }

    [[nodiscard]] auto calc_mle_objective(scalar_type reg_param) const
        -> scalar_type {
        reg_param = correct_reg_param_if_needed(reg_param);
 
        constexpr scalar_type limit = std::numeric_limits<scalar_type>::max() *
            static_cast<scalar_type>(1e-20);
        if (eigen_value_decomposition_.eigenvalues()(0) + reg_param <=
            static_cast<scalar_type>(0)) {
            return limit;
        }
 
        using std::log;
        const scalar_type value =
            static_cast<scalar_type>(kernel_subspace_dimensions_) *
                log(calc_reg_term(reg_param)) +
            calc_log_determinant(reg_param);
        if (value < limit) {
            return value;
        }
        return limit;
    }
 
private:
    [[nodiscard]] auto calc_reg_term(const scalar_type& reg_param) const
        -> scalar_type {
        return (spectre_.array().abs2().rowwise().sum() /
            (eigen_value_decomposition_.eigenvalues().array() + reg_param))
            .sum();
    }

    [[nodiscard]] auto calc_log_determinant(const scalar_type& reg_param) const
        -> scalar_type {
        return (eigen_value_decomposition_.eigenvalues().array() + reg_param)
            .log()
            .sum();
    }

    [[nodiscard]] auto correct_reg_param_if_needed(
        const scalar_type& reg_param) const noexcept -> scalar_type {
        const scalar_type smallest_eigenvalue =
            eigen_value_decomposition_.eigenvalues()(0);
        const scalar_type largest_eigenvalue =
            eigen_value_decomposition_.eigenvalues()(
                eigen_value_decomposition_.eigenvalues().size() - 1);
        const scalar_type eigenvalue_safe_limit =
            largest_eigenvalue * std::numeric_limits<scalar_type>::epsilon();
        const scalar_type reg_param_safe_limit =
            eigenvalue_safe_limit - smallest_eigenvalue;
 
        if (reg_param < reg_param_safe_limit) {
            return reg_param_safe_limit;
        }
        return reg_param;
    }

    Eigen::SelfAdjointEigenSolver<kernel_matrix_type>
        eigen_value_decomposition_{};

    Eigen::ColPivHouseholderQR<additional_matrix_type> qr_decomposition_{};

    kernel_matrix_type q_matrix_{};

    kernel_matrix_type transformed_kernel_matrix_{};

    kernel_matrix_type data_transformation_matrix_{};

    vector_type spectre_{};

    const kernel_matrix_type* kernel_matrix_{nullptr};

    const vector_type* data_{nullptr};

    index_type num_variables_{};

    index_type num_additional_terms_{};

    index_type kernel_subspace_dimensions_{};
};
 
}  // namespace num_collect::rbf::impl