lu_rosenbrock_equation_solver.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 <optional>
 
#include <Eigen/Core>  // IWYU pragma: keep
#include <Eigen/LU>
 
#include "num_collect/base/index_type.h"
#include "num_collect/ode/concepts/mass_problem.h"
#include "num_collect/ode/concepts/multi_variate_differentiable_problem.h"
#include "num_collect/ode/concepts/time_differentiable_problem.h"
#include "num_collect/ode/evaluation_type.h"
 
namespace num_collect::ode::rosenbrock {
 
template <concepts::multi_variate_differentiable_problem Problem>
class lu_rosenbrock_equation_solver {
public:
    using problem_type = Problem;
 
    using variable_type = typename problem_type::variable_type;
 
    using scalar_type = typename problem_type::scalar_type;
 
    using jacobian_type = typename problem_type::jacobian_type;
 
    using lu_solver_type = Eigen::PartialPivLU<jacobian_type>;
 
    static constexpr bool use_time_derivative =
        concepts::time_differentiable_problem<problem_type>;
 
    static constexpr bool use_mass = concepts::mass_problem<problem_type>;
 
    explicit lu_rosenbrock_equation_solver(
        const scalar_type& inverted_jacobian_coeff)
        : inverted_jacobian_coeff_(inverted_jacobian_coeff) {}
 
    void evaluate_and_update_jacobian(problem_type& problem,
        const scalar_type& time, const scalar_type& step_size,
        const variable_type& variable) {
        problem.evaluate_on(time, variable,
            evaluation_type{.diff_coeff = true,
                .jacobian = true,
                .time_derivative = use_time_derivative,
                .mass = use_mass});
        jacobian_ = problem.jacobian();
        if constexpr (use_time_derivative) {
            time_derivative_ = problem.time_derivative();
        }
 
        const index_type variable_size = variable.size();
        if constexpr (use_mass) {
            lu_.compute(problem.mass() -
                step_size * inverted_jacobian_coeff_ * jacobian_);
        } else {
            lu_.compute(jacobian_type::Identity(variable_size, variable_size) -
                step_size * inverted_jacobian_coeff_ * jacobian_);
        }
 
        // TODO: Check that condition number is not so large.
    }
 
    void apply_jacobian(const variable_type& target, variable_type& result) {
        result = jacobian_ * target;
    }
 
    void add_time_derivative_term(const scalar_type& step_size,
        const scalar_type& coeff, variable_type& target) {
        if constexpr (use_time_derivative) {
            if (time_derivative_) {
                target += step_size * coeff * (*time_derivative_);
            }
        }
    }
 
    void solve(const variable_type& rhs, variable_type& result) {
        result = lu_.solve(rhs);
    }
 
private:
    jacobian_type jacobian_{};
 
    std::optional<variable_type> time_derivative_{};
 
    lu_solver_type lu_{};
 
    scalar_type inverted_jacobian_coeff_{static_cast<scalar_type>(1)};
};
 
}  // namespace num_collect::ode::rosenbrock