graph.h

Go to the documentation of this file.

/* 
* Copyright (C) 2020-2026 MEmilio
*
* Authors: Daniel Abele, Martin J. Kuehn
*
* Contact: Martin J. Kuehn <Martin.Kuehn@DLR.de>
*
* 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.
*/
#ifndef GRAPH_H
#define GRAPH_H
 
#include "memilio/epidemiology/age_group.h"
#include "memilio/epidemiology/damping.h"
#include "memilio/geography/regions.h"
#include "memilio/utils/date.h"
#include "memilio/utils/parameter_distributions.h"
#include "memilio/utils/stl_util.h"
#include "memilio/utils/uncertain_value.h"
 
#include <algorithm>
#include <iostream>
#include <concepts>
 
namespace mio
{
 
struct OutEdgeBase {
    OutEdgeBase(size_t start)
        : start_node_idx(start)
    {
    }
    size_t start_node_idx;
};
struct InEdgeBase {
    InEdgeBase(size_t end)
        : end_node_idx(end)
    {
    }
    size_t end_node_idx;
};
struct EdgeBase : public OutEdgeBase, public InEdgeBase {
    EdgeBase(size_t start, size_t end)
        : OutEdgeBase(start)
        , InEdgeBase(end)
    {
    }
};
 
template <class NodePropertyT>
struct Node {
    template <class... Args>
    Node(int node_id, Args&&... args)
        : id{node_id}
        , property(std::forward<Args>(args)...)
    {
    }
    int id;
    NodePropertyT property;
};
 
template <class EdgePropertyT>
struct Edge : public EdgeBase {
    template <class... Args>
    Edge(size_t start, size_t end, Args&&... args)
        : EdgeBase{start, end}
        , property(std::forward<Args>(args)...)
    {
    }
 
    EdgePropertyT property;
};
 
template <std::equality_comparable T>
bool operator==(const Node<T>& n1, const Node<T>& n2)
{
    return n1.id == n2.id && n1.property == n2.property;
}
template <std::equality_comparable T>
bool operator!=(const Node<T>& n1, const Node<T>& n2)
{
    return !(n1 == n2);
}
 
template <std::equality_comparable T>
bool operator==(const Edge<T>& e1, const Edge<T>& e2)
{
    return e1.start_node_idx == e2.start_node_idx && e1.end_node_idx == e2.end_node_idx && e1.property == e2.property;
}
template <std::equality_comparable T>
bool operator!=(const Edge<T>& e1, const Edge<T>& e2)
{
    return !(e1 == e2);
}
 
template <class T>
    requires HasOstreamOperator<T>
std::ostream& operator<<(std::ostream& os, const Edge<T>& e)
{
    os << e.start_node_idx << " > " << e.end_node_idx << " : " << e.property;
    return os;
}
 
template <class T>
    requires(!HasOstreamOperator<T>)
std::ostream& operator<<(std::ostream& os, const Edge<T>& e)
{
    os << e.start_node_idx << " > " << e.end_node_idx;
    return os;
}
 
template <class NodePropertyT, class EdgePropertyT>
class Graph
    //ensure correct std::is_copy_constructible; it's not correct by default because the nodes and edges are stored in std::vector
    : not_copyable_if_t<!(std::is_copy_constructible_v<NodePropertyT> && std::is_copy_constructible_v<EdgePropertyT>)>
{
public:
    using NodeProperty = NodePropertyT;
    using EdgeProperty = EdgePropertyT;
 
    Graph(const std::vector<int>& node_ids, const std::vector<NodePropertyT>& node_properties)
    {
        assert(node_ids.size() == node_properties.size());
 
        for (auto i = size_t(0); i < node_ids.size(); ++i) {
            add_node(node_ids[i], node_properties[i]);
        }
    }
 
    Graph(std::vector<NodePropertyT>& node_properties)
    {
        for (auto i = size_t(0); i < node_properties.size(); ++i) {
            add_node(i, node_properties[i]);
        }
    }
 
    template <class... Args>
        requires std::constructible_from<NodePropertyT, Args...>
    Graph(const std::vector<int>& node_ids, Args&&... node_args)
    {
        for (int id : node_ids) {
            add_node(id, std::forward<Args>(node_args)...);
        }
    }
 
    template <class... Args>
        requires std::constructible_from<NodePropertyT, Args...>
    Graph(const int number_of_nodes, Args&&... args)
    {
        for (int id = 0; id < number_of_nodes; ++id) {
            add_node(id, std::forward<Args>(args)...);
        }
    }
 
    Graph() = default;
 
    Graph(std::vector<Node<NodePropertyT>>&& nodes, std::vector<Edge<EdgePropertyT>>&& edges)
        : m_nodes(std::move(nodes))
        , m_edges(std::move(edges))
    {
    }
 
    template <class... Args>
        requires std::constructible_from<NodePropertyT, Args...>
    void add_node(int id, Args&&... args)
    {
        m_nodes.emplace_back(id, std::forward<Args>(args)...);
    }
 
    template <class... Args>
        requires std::constructible_from<EdgePropertyT, Args...>
    void add_edge(size_t start_node_idx, size_t end_node_idx, Args&&... args)
    {
        assert(m_nodes.size() > start_node_idx && m_nodes.size() > end_node_idx);
        insert_sorted_replace(m_edges, Edge<EdgePropertyT>(start_node_idx, end_node_idx, std::forward<Args>(args)...),
                              [](auto&& e1, auto&& e2) {
                                  return e1.start_node_idx == e2.start_node_idx ? e1.end_node_idx < e2.end_node_idx
                                                                                : e1.start_node_idx < e2.start_node_idx;
                              });
    }
 
    auto nodes()
    {
        return Range(m_nodes);
    }
 
    auto nodes() const
    {
        return Range(m_nodes);
    };
 
    auto edges()
    {
        return Range(m_edges);
    }
 
    auto edges() const
    {
        return Range(m_edges);
    }
 
    auto out_edges(size_t node_idx)
    {
        return out_edges(begin(m_edges), end(m_edges), node_idx);
    }
 
    auto out_edges(size_t node_idx) const
    {
        return out_edges(begin(m_edges), end(m_edges), node_idx);
    }
 
private:
    template <typename Iter>
    static auto out_edges(Iter b, Iter e, size_t idx)
    {
        return Range(std::equal_range(b, e, OutEdgeBase(idx), [](auto&& e1, auto&& e2) {
            return e1.start_node_idx < e2.start_node_idx;
        }));
    }
 
private:
    std::vector<Node<NodePropertyT>> m_nodes;
    std::vector<Edge<EdgePropertyT>> m_edges;
}; // namespace mio
 
template <typename FP, class TestAndTrace, class ContactPattern, class Model, class MobilityParams, class Parameters,
          class ReadFunction, class NodeIdFunction>
IOResult<void> set_nodes(const Parameters& params, Date start_date, Date end_date,
                         const std::filesystem::path& data_dir, const std::string& population_data_path,
                         bool is_node_for_county, Graph<Model, MobilityParams>& params_graph, ReadFunction&& read_func,
                         NodeIdFunction&& node_func, const std::vector<FP>& scaling_factor_inf, FP scaling_factor_icu,
                         FP tnt_capacity_factor, int num_days = 0, bool export_time_series = false,
                         bool rki_age_groups = true)
 
{
    BOOST_OUTCOME_TRY(auto&& node_ids, node_func(population_data_path, is_node_for_county, rki_age_groups));
    std::vector<Model> nodes(node_ids.size(), Model(int(size_t(params.get_num_groups()))));
    for (auto& node : nodes) {
        node.parameters = params;
    }
 
    BOOST_OUTCOME_TRY(read_func(nodes, start_date, node_ids, scaling_factor_inf, scaling_factor_icu, data_dir.string(),
                                num_days, export_time_series));
 
    for (size_t node_idx = 0; node_idx < nodes.size(); ++node_idx) {
 
        auto tnt_capacity = nodes[node_idx].populations.get_total() * tnt_capacity_factor;
 
        //local parameters
        auto& tnt_value = nodes[node_idx].parameters.template get<TestAndTrace>();
        tnt_value       = UncertainValue<FP>(0.5 * (1.2 * tnt_capacity + 0.8 * tnt_capacity));
        tnt_value.set_distribution(mio::ParameterDistributionUniform(0.8 * tnt_capacity, 1.2 * tnt_capacity));
 
        auto id = 0;
        if (is_node_for_county) {
            id = int(regions::CountyId(node_ids[node_idx]));
        }
        else {
            id = int(regions::DistrictId(node_ids[node_idx]));
        }
        //holiday periods
        auto holiday_periods = regions::get_holidays(regions::get_state_id(id), start_date, end_date);
        auto& contacts       = nodes[node_idx].parameters.template get<ContactPattern>();
        contacts.get_school_holidays() =
            std::vector<std::pair<mio::SimulationTime<FP>, mio::SimulationTime<FP>>>(holiday_periods.size());
        std::transform(
            holiday_periods.begin(), holiday_periods.end(), contacts.get_school_holidays().begin(), [=](auto& period) {
                return std::make_pair(mio::SimulationTime<FP>(mio::get_offset_in_days(period.first, start_date)),
                                      mio::SimulationTime<FP>(mio::get_offset_in_days(period.second, start_date)));
            });
 
        //uncertainty in populations
        for (auto i = mio::AgeGroup(0); i < params.get_num_groups(); i++) {
            for (auto j = Index<typename Model::Compartments>(0); j < Model::Compartments::Count; ++j) {
                auto& compartment_value = nodes[node_idx].populations[{i, j}];
                compartment_value =
                    UncertainValue<FP>(0.5 * (1.1 * compartment_value.value() + 0.9 * compartment_value.value()));
                compartment_value.set_distribution(mio::ParameterDistributionUniform(0.9 * compartment_value.value(),
                                                                                     1.1 * compartment_value.value()));
            }
        }
 
        params_graph.add_node(node_ids[node_idx], nodes[node_idx]);
    }
    return success();
}
 
template <typename FP, class ContactLocation, class Model, class MobilityParams, class MobilityCoefficientGroup,
          class InfectionState, class ReadFunction>
IOResult<void> set_edges(const std::filesystem::path& mobility_data_file, Graph<Model, MobilityParams>& params_graph,
                         std::initializer_list<InfectionState>& mobile_compartments, size_t contact_locations_size,
                         ReadFunction&& read_func, std::vector<FP> commuting_weights,
                         std::vector<std::vector<size_t>> indices_of_saved_edges = {})
{
    // mobility between nodes
    BOOST_OUTCOME_TRY(auto&& mobility_data_commuter, read_func(mobility_data_file.string()));
    if (mobility_data_commuter.rows() != Eigen::Index(params_graph.nodes().size()) ||
        mobility_data_commuter.cols() != Eigen::Index(params_graph.nodes().size())) {
        return mio::failure(mio::StatusCode::InvalidValue,
                            "Mobility matrices do not have the correct size. You may need to run "
                            "transformMobilitydata.py from pycode memilio epidata package.");
    }
 
    for (size_t county_idx_i = 0; county_idx_i < params_graph.nodes().size(); ++county_idx_i) {
        for (size_t county_idx_j = 0; county_idx_j < params_graph.nodes().size(); ++county_idx_j) {
            auto& populations = params_graph.nodes()[county_idx_i].property.populations;
            // mobility coefficients have the same number of components as the contact matrices.
            // so that the same NPIs/dampings can be used for both (e.g. more home office => fewer commuters)
            auto mobility_coeffs = MobilityCoefficientGroup(contact_locations_size, populations.numel());
            auto num_age_groups  = (size_t)params_graph.nodes()[county_idx_i].property.parameters.get_num_groups();
            commuting_weights =
                (commuting_weights.size() == 0 ? std::vector<FP>(num_age_groups, 1.0) : commuting_weights);
            //commuters
            FP working_population = 0.0;
            for (auto age = AgeGroup(0); age < populations.template size<mio::AgeGroup>(); ++age) {
                working_population += populations.get_group_total(age) * commuting_weights[size_t(age)];
            }
            auto commuter_coeff_ij = mobility_data_commuter(county_idx_i, county_idx_j) /
                                     working_population; //data is absolute numbers, we need relative
            for (auto age = AgeGroup(0); age < populations.template size<mio::AgeGroup>(); ++age) {
                for (auto compartment : mobile_compartments) {
                    auto coeff_index = populations.get_flat_index({age, compartment});
                    mobility_coeffs[size_t(ContactLocation::Work)].get_baseline()[coeff_index] =
                        commuter_coeff_ij * commuting_weights[size_t(age)];
                }
            }
            //only add edges with mobility above thresholds for performance
            //thresholds are chosen empirically so that more than 99% of mobility is covered, approx. 1/3 of the edges
            if (commuter_coeff_ij > 4e-5) {
                params_graph.add_edge(county_idx_i, county_idx_j, std::move(mobility_coeffs), indices_of_saved_edges);
            }
        }
    }
 
    return success();
}
 
template <class T>
    requires(!HasOstreamOperator<T>)
void print_graph_object(std::ostream& os, size_t idx, const T&)
{
    os << idx;
}
 
template <class T>
    requires HasOstreamOperator<T>
void print_graph_object(std::ostream& os, size_t idx, const T& o)
{
    os << idx << " [" << o << "]";
}
 
template <class Graph>
void print_graph(std::ostream& os, const Graph& g)
{
    auto nodes = g.nodes();
    for (size_t i = 0; i < nodes.size(); ++i) {
        os << "NODE ";
        print_graph_object(os, i, nodes[i]);
        os << '\n';
    }
 
    auto edges = g.edges();
    for (size_t i = 0; i < edges.size(); ++i) {
        auto& e = edges[i];
        os << "EDGE ";
        print_graph_object(os, i, e.property);
        os << " FROM NODE ";
        print_graph_object(os, e.start_node_idx, nodes[e.start_node_idx]);
        os << " TO ";
        print_graph_object(os, e.end_node_idx, nodes[e.end_node_idx]);
        os << '\n';
    }
}
 
} // namespace mio
 
#endif //GRAPH_H