minor refactoring

Author: DillonZChen

include/feature_generation/feature_generators/ccwl.hpp

@@ -21,7 +21,7 @@ namespace feature_generation {
 
     CCWLFeatures(const std::string &filename);
 
-    Embedding embed(const std::shared_ptr<graph::Graph> &graph) override;
+    Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) override;
 
     void set_weights(const std::vector<double> &weights);
   };

include/feature_generation/feature_generators/iwl.hpp

@@ -28,13 +28,12 @@ namespace feature_generation {
 
     IWLFeatures(const std::string &filename);
 
-    Embedding embed(const std::shared_ptr<graph::Graph> &graph) override;
+    Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) override;
 
    protected:
     void collect_impl(const std::vector<graph::Graph> &graphs) override;
     void refine(const std::shared_ptr<graph::Graph> &graph,
                 std::vector<int> &colours,
-                std::vector<int> &colours_tmp,
                 int iteration);
   };
 }  // namespace feature_generation

include/feature_generation/feature_generators/kwl2.hpp

@@ -28,7 +28,7 @@ namespace feature_generation {
 
     KWL2Features(const std::string &filename);
 
-    Embedding embed(const std::shared_ptr<graph::Graph> &graph) override;
+    Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) override;
 
    protected:
     inline int get_initial_colour(int index,
@@ -39,7 +39,6 @@ namespace feature_generation {
     void collect_impl(const std::vector<graph::Graph> &graphs) override;
     void refine(const std::shared_ptr<graph::Graph> &graph,
                 std::vector<int> &colours,
-                std::vector<int> &colours_tmp,
                 int iteration);
   };
 }  // namespace feature_generation

include/feature_generation/feature_generators/lwl2.hpp

@@ -21,7 +21,7 @@ namespace feature_generation {
 
     LWL2Features(const std::string &filename);
 
-    Embedding embed(const std::shared_ptr<graph::Graph> &graph) override;
+    Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) override;
 
    protected:
     inline int get_initial_colour(int index,
@@ -33,7 +33,6 @@ namespace feature_generation {
     void refine(const std::shared_ptr<graph::Graph> &graph,
                 std::vector<std::set<int>> &pair_to_neighbours,
                 std::vector<int> &colours,
-                std::vector<int> &colours_tmp,
                 int iteration);
   };
 }  // namespace feature_generation

include/feature_generation/feature_generators/niwl.hpp

@@ -18,7 +18,7 @@ namespace feature_generation {
 
     NIWLFeatures(const std::string &filename);
 
-    Embedding embed(const std::shared_ptr<graph::Graph> &graph) override;
+    Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) override;
   };
 }  // namespace feature_generation
 

include/feature_generation/feature_generators/wl.hpp

@@ -25,14 +25,13 @@ namespace feature_generation {
 
     WLFeatures(const std::string &filename);
 
-    Embedding embed(const std::shared_ptr<graph::Graph> &graph) override;
+    Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) override;
 
    protected:
     void collect_impl(const std::vector<graph::Graph> &graphs) override;
     void refine(const std::shared_ptr<graph::Graph> &graph,
                 std::set<int> &nodes,
                 std::vector<int> &colours,
-                std::vector<int> &colours_tmp,
                 int iteration);
   };
 }  // namespace feature_generation

include/feature_generation/features.hpp

@@ -97,8 +97,9 @@ namespace feature_generation {
     // common init for initialisation and loading from file
     void initialise_variables();
 
-    // main collection body
+    // main virtual functions
     virtual void collect_impl(const std::vector<graph::Graph> &graphs) = 0;
+    virtual Embedding embed_impl(const std::shared_ptr<graph::Graph> &graph) = 0;
 
    public:
     Features(const std::string feature_name,
@@ -124,7 +125,8 @@ namespace feature_generation {
     std::vector<Embedding> embed_graphs(const std::vector<graph::Graph> &graphs);
     Embedding embed_graph(const graph::Graph &graph);
     Embedding embed_state(const planning::State &state);
-    virtual Embedding embed(const std::shared_ptr<graph::Graph> &graph) = 0;
+    Embedding embed(const std::shared_ptr<graph::Graph> &graph);
+
 
     void add_colour_to_x(int colour, int iteration, Embedding &x);
 

src/feature_generation/feature_generator_loader.cpp

@@ -28,8 +28,7 @@ std::shared_ptr<feature_generation::Features> load_feature_generator(const std::
   } else if (feature_name == "niwl") {
     feature_generator = std::make_shared<feature_generation::NIWLFeatures>(save_file);
   } else {
-    std::cout << "Feature name " << feature_name << " not recognised. Exiting." << std::endl;
-    exit(-1);
+    throw std::runtime_error("Feature name " + feature_name + " not recognised.");
   }
   std::cout << "Feature generator loaded!" << std::endl;
   return feature_generator;

src/feature_generation/feature_generators/ccwl.cpp

@@ -18,12 +18,7 @@ namespace feature_generation {
 
   CCWLFeatures::CCWLFeatures(const std::string &filename) : WLFeatures(filename) {}
 
-  Embedding CCWLFeatures::embed(const std::shared_ptr<graph::Graph> &graph) {
-    collecting = false;
-    if (!collected) {
-      throw std::runtime_error("CCWLFeatures::collect() must be called before embedding");
-    }
-
+  Embedding CCWLFeatures::embed_impl(const std::shared_ptr<graph::Graph> &graph) {
     // New additions to the WL algorithm are indicated with the [NUMERIC] comments.
     // We use a sum function for the pool operator as described in the ccWL algorithm.
     // To change this to max, we just need to replace += occurrences with std::max.
@@ -33,7 +28,6 @@ namespace feature_generation {
     Embedding x0(categorical_size * 2, 0);
     int n_nodes = graph->nodes.size();
     std::vector<int> colours(n_nodes);
-    std::vector<int> colours_tmp(n_nodes);
     std::set<int> nodes = graph->get_nodes_set();
 
     /* 2. Compute initial colours */
@@ -52,7 +46,7 @@ namespace feature_generation {
 
     /* 3. Main WL loop */
     for (int itr = 1; itr < iterations + 1; itr++) {
-      refine(graph, nodes, colours, colours_tmp, itr);
+      refine(graph, nodes, colours, itr);
       for (int node_i = 0; node_i < n_nodes; node_i++) {
         col = colours[node_i];
         is_seen_colour = (col != UNSEEN_COLOUR);  // prevent branch prediction

src/feature_generation/feature_generators/iwl.cpp

@@ -29,13 +29,14 @@ namespace feature_generation {
 
   void IWLFeatures::refine(const std::shared_ptr<graph::Graph> &graph,
                            std::vector<int> &colours,
-                           std::vector<int> &colours_tmp,
                            int iteration) {
     // memory for storing string and hashed int representation of colours
     std::vector<int> new_colour;
     std::vector<int> neighbour_vector;
     int new_colour_compressed;
 
+    std::vector<int> new_colours(colours.size(), UNSEEN_COLOUR);
+
     for (size_t u = 0; u < graph->nodes.size(); u++) {
       // skip unseen colours
       if (colours[u] == UNSEEN_COLOUR) {
@@ -65,16 +66,15 @@ namespace feature_generation {
       new_colour_compressed = get_colour_hash(new_colour, iteration);
 
     end_of_iteration:
-      colours_tmp[u] = new_colour_compressed;
+      new_colours[u] = new_colour_compressed;
     }
 
-    colours.swap(colours_tmp);
+    colours = new_colours;
   }
 
   void IWLFeatures::collect_impl(const std::vector<graph::Graph> &graphs) {
-    // intermediate graph colours during WL and extra memory for WL updates
+    // intermediate graph colours during WL
     std::vector<int> colours;
-    std::vector<int> colours_tmp;
 
     // init colours
     for (size_t graph_i = 0; graph_i < graphs.size(); graph_i++) {
@@ -84,7 +84,6 @@ namespace feature_generation {
       // individualisation for each node
       for (int node_i = 0; node_i < n_nodes; node_i++) {
         colours = std::vector<int>(n_nodes, 0);
-        colours_tmp = std::vector<int>(n_nodes, 0);
 
         for (int u = 0; u < n_nodes; u++) {
           std::vector<int> colour_key = {graph->nodes[u]};
@@ -97,30 +96,22 @@ namespace feature_generation {
 
         // main WL loop
         for (int iteration = 1; iteration < iterations + 1; iteration++) {
-          refine(graph, colours, colours_tmp, iteration);
+          refine(graph, colours, iteration);
         }
       }
     }
   }
 
-  Embedding IWLFeatures::embed(const std::shared_ptr<graph::Graph> &graph) {
-    collecting = false;
-    if (!collected) {
-      throw std::runtime_error("IWLFeatures::collect() must be called before embedding");
-    }
-
+  Embedding IWLFeatures::embed_impl(const std::shared_ptr<graph::Graph> &graph) {
     /* 1. Initialise embedding */
     Embedding x0(get_n_features(), 0);
-
-    /* 2. Set up memory for WL updates */
     int n_nodes = graph->nodes.size();
 
     /* Individualisation */
     for (int node_i = 0; node_i < n_nodes; node_i++) {
       std::vector<int> colours(n_nodes);
-      std::vector<int> colours_tmp(n_nodes);
 
-      /* 3. Compute initial colours */
+      /* 2. Compute initial colours */
       for (int u = 0; u < n_nodes; u++) {
         std::vector<int> colour_key = {graph->nodes[u]};
         if (u == node_i) {
@@ -130,9 +121,9 @@ namespace feature_generation {
         add_colour_to_x(col, 0, x0);
       }
 
-      /* 4. Main WL loop */
+      /* 3. Main WL loop */
       for (int itr = 1; itr < iterations + 1; itr++) {
-        refine(graph, colours, colours_tmp, itr);
+        refine(graph, colours, itr);
         for (const int col : colours) {
           add_colour_to_x(col, itr, x0);
         }

src/feature_generation/feature_generators/kwl2.cpp

@@ -35,14 +35,15 @@ namespace feature_generation {
 
   void KWL2Features::refine(const std::shared_ptr<graph::Graph> &graph,
                             std::vector<int> &colours,
-                            std::vector<int> &colours_tmp,
                             int iteration) {
     // memory for storing string and hashed int representation of colours
     std::vector<int> new_colour;
     std::vector<int> neighbour_vector;
     int new_colour_compressed, pair1, pair2, pair1_col, pair2_col;
     int n_nodes = graph->nodes.size();
 
+    std::vector<int> new_colours(colours.size(), UNSEEN_COLOUR);
+
     for (int u = 0; u < n_nodes; u++) {
       for (int v = 0; v < n_nodes; v++) {
         int index = kwl2_pair_to_index_map(n_nodes, u, v);
@@ -76,11 +77,11 @@ namespace feature_generation {
         new_colour_compressed = get_colour_hash(new_colour, iteration);
 
       end_of_iteration:
-        colours_tmp[index] = new_colour_compressed;
+        new_colours[index] = new_colour_compressed;
       }
     }
 
-    colours.swap(colours_tmp);
+    colours = new_colours;
   }
 
   std::vector<int> get_kwl2_pair_to_edge_label(std::shared_ptr<graph::Graph> graph) {
@@ -109,9 +110,8 @@ namespace feature_generation {
   }
 
   void KWL2Features::collect_impl(const std::vector<graph::Graph> &graphs) {
-    // intermediate graph colours during WL and extra memory for WL updates
+    // intermediate graph colours during WL
     std::vector<int> colours;
-    std::vector<int> colours_tmp;
 
     for (size_t graph_i = 0; graph_i < graphs.size(); graph_i++) {
       const auto graph = std::make_shared<graph::Graph>(graphs[graph_i]);
@@ -122,7 +122,6 @@ namespace feature_generation {
 
       // intermediate colours
       colours = std::vector<int>(n_pairs, 0);
-      colours_tmp = std::vector<int>(n_pairs, 0);
 
       std::vector<int> pair_to_edge_label = get_kwl2_pair_to_edge_label(graph);
 
@@ -137,29 +136,22 @@ namespace feature_generation {
 
       // main WL loop
       for (int iteration = 1; iteration < iterations + 1; iteration++) {
-        refine(graph, colours, colours_tmp, iteration);
+        refine(graph, colours, iteration);
       }
     }
   }
 
-  Embedding KWL2Features::embed(const std::shared_ptr<graph::Graph> &graph) {
-    collecting = false;
-    if (!collected) {
-      throw std::runtime_error("KWL2Features::collect() must be called before embedding");
-    }
-
+  Embedding KWL2Features::embed_impl(const std::shared_ptr<graph::Graph> &graph) {
     /* 1. Initialise embedding before pruning */
     Embedding x0(get_n_features(), 0);
 
-    /* 2. Set up memory for WL updates */
     int n_nodes = graph->nodes.size();
     int n_pairs = get_n_kwl2_pairs(n_nodes);
     std::vector<int> colours(n_pairs);
-    std::vector<int> colours_tmp(n_pairs);
 
     std::vector<int> pair_to_edge_label = get_kwl2_pair_to_edge_label(graph);
 
-    /* 3. Compute initial colours */
+    /* 2. Compute initial colours */
     for (int u = 0; u < n_nodes; u++) {
       for (int v = 0; v < n_nodes; v++) {
         int index = kwl2_pair_to_index_map(n_nodes, u, v);
@@ -169,9 +161,9 @@ namespace feature_generation {
       }
     }
 
-    /* 4. Main WL loop */
+    /* 3. Main WL loop */
     for (int itr = 1; itr < iterations + 1; itr++) {
-      refine(graph, colours, colours_tmp, itr);
+      refine(graph, colours, itr);
       for (const int col : colours) {
         add_colour_to_x(col, itr, x0);
       }