differentiate.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 <cstddef>
#include <memory>
#include <utility>  // IWYU pragma: keep
#include <vector>
 
#include <Eigen/Core>
 
#include "num_collect/auto_diff/backward/graph/node_differentiator.h"
#include "num_collect/auto_diff/backward/variable.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/util/safe_cast.h"
 
namespace num_collect::auto_diff::backward {

template <typename Scalar>
[[nodiscard]] inline auto differentiate(
    const variable<Scalar>& func_value, const variable<Scalar>& arg) -> Scalar {
    graph::node_differentiator<Scalar> diff;
    diff.compute(func_value.node());
    return diff.coeff(arg.node());
}
 
namespace impl {

template <typename ArgType>
class differentiate_scalar_with_matrix_functor {
public:
    using variable_type = typename ArgType::Scalar;

    using scalar_type = typename variable_type::scalar_type;

    using result_type = Eigen::Matrix<scalar_type, ArgType::RowsAtCompileTime,
        ArgType::ColsAtCompileTime, Eigen::ColMajor,
        ArgType::MaxRowsAtCompileTime, ArgType::MaxColsAtCompileTime>;

    differentiate_scalar_with_matrix_functor(const ArgType& arg,
        std::shared_ptr<const graph::node_differentiator<scalar_type>> diff)
        : arg_(arg), diff_(std::move(diff)) {}

    auto operator()(index_type row, index_type col) const -> scalar_type {
        return diff_->coeff(arg_(row, col).node());
    }
 
private:
    const ArgType& arg_;

    std::shared_ptr<const graph::node_differentiator<scalar_type>> diff_;
};
 
}  // namespace impl

template <typename Scalar, typename ArgDerived>
[[nodiscard]] inline auto differentiate(const variable<Scalar>& func_value,
    const Eigen::MatrixBase<ArgDerived>& arg)
    -> Eigen::CwiseNullaryOp<
        impl::differentiate_scalar_with_matrix_functor<ArgDerived>,
        typename impl::differentiate_scalar_with_matrix_functor<
            ArgDerived>::result_type> {
    static_assert(
        std::is_same_v<typename ArgDerived::Scalar, variable<Scalar>>);
 
    auto diff = std::make_shared<graph::node_differentiator<Scalar>>();
    diff->compute(func_value.node());
 
    using result_type = typename impl::differentiate_scalar_with_matrix_functor<
        ArgDerived>::result_type;
    return result_type::NullaryExpr(arg.rows(), arg.cols(),
        impl::differentiate_scalar_with_matrix_functor<ArgDerived>(
            arg.derived(), std::move(diff)));
}
 
namespace impl {

template <typename FuncValType, typename ArgType>
class differentiate_vector_with_vector_functor {
public:
    using variable_type = typename FuncValType::Scalar;
 
    static_assert(std::is_same_v<typename ArgType::Scalar, variable_type>);

    using scalar_type = typename variable_type::scalar_type;

    using result_type =
        Eigen::Matrix<scalar_type, FuncValType::RowsAtCompileTime,
            ArgType::RowsAtCompileTime, Eigen::ColMajor,
            FuncValType::MaxRowsAtCompileTime, ArgType::MaxRowsAtCompileTime>;

    differentiate_vector_with_vector_functor(const ArgType& arg,
        std::shared_ptr<std::vector<graph::node_differentiator<scalar_type>>>
            diff)
        : arg_(arg), diff_(std::move(diff)) {}

    auto operator()(index_type row, index_type col) const -> scalar_type {
        return diff_->at(util::safe_cast<std::size_t>(row))
            .coeff(arg_(col, 0).node());
    }
 
private:
    const ArgType& arg_;

    std::shared_ptr<std::vector<graph::node_differentiator<scalar_type>>> diff_;
};
 
}  // namespace impl

template <typename FuncValType, typename ArgType>
[[nodiscard]] inline auto differentiate(
    const Eigen::MatrixBase<FuncValType>& func_value,
    const Eigen::MatrixBase<ArgType>& arg)
    -> Eigen::CwiseNullaryOp<
        impl::differentiate_vector_with_vector_functor<FuncValType, ArgType>,
        typename impl::differentiate_vector_with_vector_functor<FuncValType,
            ArgType>::result_type> {
    using variable_type = typename FuncValType::Scalar;
    using scalar_type = typename variable_type::scalar_type;
 
    NUM_COLLECT_PRECONDITION(func_value.cols() == 1 && arg.cols() == 1,
        "differentiate function requires vectors as arguments.");
 
    auto diff = std::make_shared<
        std::vector<graph::node_differentiator<scalar_type>>>();
    diff->resize(util::safe_cast<std::size_t>(func_value.size()));
    for (index_type i = 0; i < func_value.size(); ++i) {
        diff->operator[](util::safe_cast<std::size_t>(i))
            .compute(func_value(i, 0).node());
    }
 
    using result_type =
        typename impl::differentiate_vector_with_vector_functor<FuncValType,
            ArgType>::result_type;
    return result_type::NullaryExpr(func_value.size(), arg.size(),
        impl::differentiate_vector_with_vector_functor<FuncValType, ArgType>(
            arg.derived(), std::move(diff)));
}
 
}  // namespace num_collect::auto_diff::backward