/include/dgl/transform.h

/*!
 *  Copyright (c) 2019 by Contributors
 * \file dgl/transform.h
 * \brief DGL graph transformations
 */

#ifndef DGL_TRANSFORM_H_
#define DGL_TRANSFORM_H_

#include <vector>
#include <tuple>
#include <utility>
#include "base_heterograph.h"
#include "array.h"

namespace dgl {

namespace transform {

/*!
 * \brief Given a list of graphs, remove the common nodes that do not have inbound and
 * outbound edges.
 *
 * The graphs should have identical node ID space (i.e. should have the same set of nodes,
 * including types and IDs).
 *
 * \param graphs The list of graphs.
 * \param always_preserve The list of nodes to preserve regardless of whether the inbound
 *                        or outbound edges exist.
 *
 * \return A pair.  The first element is the list of compacted graphs, and the second
 * element is the mapping from the compacted graphs and the original graph.
 */
std::pair<std::vector<HeteroGraphPtr>, std::vector<IdArray>>
CompactGraphs(
    const std::vector<HeteroGraphPtr> &graphs,
    const std::vector<IdArray> &always_preserve);

/*!
 * \brief Convert a graph into a bipartite-structured graph for message passing.
 *
 * Specifically, we create one node type \c ntype_l on the "left" side and another
 * node type \c ntype_r on the "right" side for each node type \c ntype.  The nodes of
 * type \c ntype_r would contain the nodes designated by the caller, and node type
 * \c ntype_l would contain the nodes that has an edge connecting to one of the
 * designated nodes.
 *
 * The nodes of \c ntype_l would also contain the nodes in node type \c ntype_r.
 *
 * This function is often used for constructing a series of dependency graphs for
 * multi-layer message passing, where we first construct a series of frontier graphs
 * on the original node space, and run the following to get the bipartite graph needed
 * for message passing with each GNN layer:
 *
 * <code>
 *     bipartites = [None] * len(num_layers)
 *     for l in reversed(range(len(layers))):
 *         bipartites[l], seeds = to_bipartite(frontier[l], seeds)
 *     x = graph.ndata["h"][seeds]
 *     for g, layer in zip(bipartites, layers):
 *         x_src = x
 *         x_dst = x[:len(g.dsttype)]
 *         x = sageconv(g, (x_src, x_dst))
 *     output = x
 * </code>
 *
 * \param graph The graph.
 * \param rhs_nodes Designated nodes that would appear on the right side.
 * \param include_rhs_in_lhs If false, do not include the nodes of node type \c ntype_r
 *        in \c ntype_l.
 *
 * \return A triplet containing
 *         * The bipartite-structured graph,
 *         * The induced node from the left side for each graph,
 *         * The induced edges.
 *
 * \note If include_rhs_in_lhs is true, then for each node type \c ntype, the nodes
 *       in rhs_nodes[ntype] would always appear first in the nodes of type \c ntype_l
 *       in the new graph.
 */
std::tuple<HeteroGraphPtr, std::vector<IdArray>, std::vector<IdArray>>
ToBlock(HeteroGraphPtr graph, const std::vector<IdArray> &rhs_nodes, bool include_rhs_in_lhs);

/*!
 * \brief Convert a multigraph to a simple graph.
 *
 * \return A triplet of
 * * \c hg : The said simple graph.
 * * \c count : The array of edge occurrences per edge type.
 * * \c edge_map : The mapping from original edge IDs to new edge IDs per edge type.
 *
 * \note Example: consider a graph with the following edges
 *
 *     [(0, 1), (1, 3), (2, 2), (1, 3), (1, 4), (1, 4)]
 *
 * Then ToSimpleGraph(g) would yield the following elements:
 *
 * * The first element would be the simple graph itself with the following edges
 *
 *       [(0, 1), (1, 3), (1, 4), (2, 2)]
 *
 * * The second element is an array \c count.  \c count[i] stands for the number of edges
 *   connecting simple_g.src[i] and simple_g.dst[i] in the original graph.
 *
 *       count[0] = [1, 2, 2, 1]
 *
 * * One can find the mapping between edges from the original graph to the new simple
 *   graph.
 *
 *       edge_map[0] = [0, 1, 3, 1, 2, 2]
 */
std::tuple<HeteroGraphPtr, std::vector<IdArray>, std::vector<IdArray>>
ToSimpleGraph(const HeteroGraphPtr graph);

/*!
 * \brief Remove edges from a graph.
 *
 * \param graph The graph.
 * \param eids The edge IDs to remove per edge type.
 *
 * \return A pair of the graph with edges removed, as well as the edge ID mapping from
 *         the original graph to the new graph per edge type.
 */
std::pair<HeteroGraphPtr, std::vector<IdArray>>
RemoveEdges(const HeteroGraphPtr graph, const std::vector<IdArray> &eids);

};  // namespace transform

};  // namespace dgl

#endif  // DGL_TRANSFORM_H_