RModel_GNN.hxx

Go to the documentation of this file.

#ifndef TMVA_SOFIE_RMODEL_GNN
#define TMVA_SOFIE_RMODEL_GNN
 
#include <ctime>
 
#include "TMVA/RModel_Base.hxx"
#include "TMVA/RModel.hxx"
#include "TMVA/RFunction.hxx"
 
namespace TMVA {
namespace Experimental {
namespace SOFIE {
 
class RFunction_Update;
class RFunction_Aggregate;
 
struct GNN_Init {
    // updation blocks
    std::unique_ptr<RFunction_Update> edges_update_block;
    std::unique_ptr<RFunction_Update> nodes_update_block;
    std::unique_ptr<RFunction_Update> globals_update_block;
 
    // aggregation blocks
    std::unique_ptr<RFunction_Aggregate> edge_node_agg_block;
    std::unique_ptr<RFunction_Aggregate> edge_global_agg_block;
    std::unique_ptr<RFunction_Aggregate> node_global_agg_block;
 
    int num_nodes;
    std::vector<std::pair<int,int>> edges;
 
    std::size_t num_node_features;
    std::size_t num_edge_features;
    std::size_t num_global_features;
 
    std::string filename;
 
    ~GNN_Init() {
        edges_update_block.reset();
        nodes_update_block.reset();
        globals_update_block.reset();
 
        edge_node_agg_block.reset();
        edge_global_agg_block.reset();
        node_global_agg_block.reset();
    }
 
    template <typename T>
    void createUpdateFunction(T& updateFunction) {
        switch(updateFunction.GetFunctionTarget()) {
        case FunctionTarget::EDGES: {
            edges_update_block.reset(new T(updateFunction));
            break;
        }
        case FunctionTarget::NODES: {
            nodes_update_block.reset(new T(updateFunction));
            break;
        }
        case FunctionTarget::GLOBALS: {
            globals_update_block.reset(new T(updateFunction));
            break;
        }
        default: {
            throw std::runtime_error("TMVA SOFIE: Invalid Update function supplied for creating GNN function block.");
        }
        }
    }
 
    template <typename T>
    void createAggregateFunction(T& aggFunction, FunctionRelation relation) {
        switch(relation) {
        case FunctionRelation::NODES_EDGES : {
            edge_node_agg_block.reset(new T(aggFunction));
            break;
        }
        case FunctionRelation::NODES_GLOBALS: {
            node_global_agg_block.reset(new T(aggFunction));
            break;
        }
        case FunctionRelation::EDGES_GLOBALS: {
            edge_global_agg_block.reset(new T(aggFunction));
            break;
        }
        default: {
            throw std::runtime_error("TMVA SOFIE: Invalid Aggregate function supplied for creating GNN function block.");
        }
        }
    }
 
};
 
class RModel_GNN: public RModel_GNNBase {
 
private:
 
    // updation function for edges, nodes & global attributes
    std::unique_ptr<RFunction_Update> edges_update_block;
    std::unique_ptr<RFunction_Update> nodes_update_block;
    std::unique_ptr<RFunction_Update> globals_update_block;
 
    // aggregation function for edges, nodes & global attributes
    std::unique_ptr<RFunction_Aggregate> edge_node_agg_block;
    std::unique_ptr<RFunction_Aggregate> edge_global_agg_block;
    std::unique_ptr<RFunction_Aggregate> node_global_agg_block;
 
    int num_nodes;
    int num_edges;
    std::vector<int> senders;              // contains node indices
    std::vector<int> receivers;            // contains node indices
 
    std::size_t num_node_features;
    std::size_t num_edge_features;
    std::size_t num_global_features;
 
public:
 
    //explicit move ctor/assn
    RModel_GNN(RModel_GNN&& other);
 
    RModel_GNN& operator=(RModel_GNN&& other);
 
    //disallow copy
    RModel_GNN(const RModel_GNN& other) = delete;
    RModel_GNN& operator=(const RModel_GNN& other) = delete;
 
    RModel_GNN(GNN_Init& graph_input_struct);
    RModel_GNN() {}
 
    void Generate();
 
    ~RModel_GNN() {}
//    ClassDef(RModel_GNN,1);
};
 
}//SOFIE
}//Experimental
}//TMVA
 
#endif //TMVA_SOFIE_RMODEL_GNN