node_differentiator.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 <queue>
#include <unordered_map>
 
#include "num_collect/auto_diff/backward/graph/node.h"
#include "num_collect/base/concepts/real_scalar.h"
#include "num_collect/base/index_type.h"
#include "num_collect/util/assert.h"
 
namespace num_collect::auto_diff::backward::graph {
 
template <base::concepts::real_scalar Scalar>
class node_differentiator {
public:
    using scalar_type = Scalar;
 
    node_differentiator() = default;
 
    void compute(const node_ptr<scalar_type>& top_node) {
        NUM_COLLECT_ASSERT(top_node);
        list_nodes(top_node);
        compute_coeffs(top_node);
    }
 
    [[nodiscard]] auto coeff(const node_ptr<scalar_type>& node) const
        -> scalar_type {
        if (const auto iter = info_dict_.find(node); iter != info_dict_.end()) {
            return iter->second.diff;
        }
        return static_cast<scalar_type>(0);
    }
 
private:
    struct node_info {
        scalar_type diff{static_cast<scalar_type>(0)};
 
        index_type ref_count{0};
    };
 
    void list_nodes(const node_ptr<scalar_type>& top_node) {
        info_dict_.clear();
        node_queue_ = std::queue<node_ptr<scalar_type>>();
 
        info_dict_.try_emplace(top_node);
        node_queue_.push(top_node);
        while (!node_queue_.empty()) {
            const auto& node = node_queue_.front();
            for (const auto& child_node : node->children()) {
                const auto [iter, inserted] =
                    info_dict_.try_emplace(child_node.node());
                if (inserted) {
                    node_queue_.push(child_node.node());
                }
                ++iter->second.ref_count;
            }
            node_queue_.pop();
        }
    }
 
    void compute_coeffs(const node_ptr<scalar_type>& top_node) {
        info_dict_[top_node].diff = static_cast<scalar_type>(1);
        node_queue_.push(top_node);
        while (!node_queue_.empty()) {
            const auto& node = node_queue_.front();
            const auto& info = info_dict_[node];
            for (const auto& child_node : node->children()) {
                auto& child_info = info_dict_[child_node.node()];
                child_info.diff += info.diff * child_node.sensitivity();
                --child_info.ref_count;
                if (child_info.ref_count == 0) {
                    node_queue_.push(child_node.node());
                }
            }
            node_queue_.pop();
        }
    }
 
    std::unordered_map<node_ptr<scalar_type>, node_info> info_dict_{};
 
    std::queue<node_ptr<scalar_type>> node_queue_{};
};
 
}  // namespace num_collect::auto_diff::backward::graph