Updated simulator to use marius

marius-team · Feb 22, 2024 · 57856a1 · 57856a1
1 parent 84f9198
commit 57856a1
Show file tree

Hide file tree

Showing 489 changed files with 84 additions and 924 deletions.
diff --git a/.clang-format b/.clang-format
diff --git a/.flake8 b/.flake8
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
diff --git a/.github/ISSUE_TEMPLATE/documentation-improvement.md b/.github/ISSUE_TEMPLATE/documentation-improvement.md
diff --git a/.github/ISSUE_TEMPLATE/feature_request.md b/.github/ISSUE_TEMPLATE/feature_request.md
diff --git a/.github/ISSUE_TEMPLATE/general-question.md b/.github/ISSUE_TEMPLATE/general-question.md
diff --git a/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md b/.github/PULL_REQUEST_TEMPLATE/pull_request_template.md
diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml
diff --git a/.github/workflows/db2graph_test_postgres.yml b/.github/workflows/db2graph_test_postgres.yml
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
diff --git a/.gitignore b/.gitignore
@@ -171,3 +171,4 @@ Thumbs.db
 
 
 # End of https://www.toptal.com/developers/gitignore/api/python
+simulator/datasets
diff --git a/.gitmodules b/.gitmodules
diff --git a/CMakeLists.txt b/CMakeLists.txt
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
diff --git a/LICENSE b/LICENSE
diff --git a/MANIFEST.in b/MANIFEST.in
diff --git a/README.md b/README.md
diff --git a/docs/.nojekyll b/docs/.nojekyll
diff --git a/docs/CMakeLists.txt b/docs/CMakeLists.txt
diff --git a/docs/Doxyfile b/docs/Doxyfile
diff --git a/docs/Doxyfile.in b/docs/Doxyfile.in
diff --git a/docs/README.md b/docs/README.md
diff --git a/docs/_static/css/marius_theme.css b/docs/_static/css/marius_theme.css
diff --git a/docs/_templates/layout.html b/docs/_templates/layout.html
diff --git a/docs/assets/configuration_lp.png b/docs/assets/configuration_lp.png
diff --git a/docs/assets/configuration_nc.png b/docs/assets/configuration_nc.png
diff --git a/docs/assets/marius_arch.png b/docs/assets/marius_arch.png
diff --git a/docs/assets/samples_dismult.png b/docs/assets/samples_dismult.png
diff --git a/docs/assets/samples_gs.png b/docs/assets/samples_gs.png
diff --git a/docs/conf.py b/docs/conf.py
diff --git a/docs/config_interface/configuration.rst b/docs/config_interface/configuration.rst
diff --git a/docs/config_interface/full_schema.rst b/docs/config_interface/full_schema.rst
diff --git a/docs/config_interface/index.rst b/docs/config_interface/index.rst
diff --git a/docs/config_interface/samples.rst b/docs/config_interface/samples.rst
diff --git a/docs/db2graph/db2graph.rst b/docs/db2graph/db2graph.rst
diff --git a/docs/examples/config/index.rst b/docs/examples/config/index.rst
diff --git a/docs/examples/config/lp_custom.rst b/docs/examples/config/lp_custom.rst
diff --git a/docs/examples/config/lp_fb15k237.rst b/docs/examples/config/lp_fb15k237.rst
diff --git a/docs/examples/config/lp_paleobiology.rst b/docs/examples/config/lp_paleobiology.rst
diff --git a/docs/examples/config/nc_custom.rst b/docs/examples/config/nc_custom.rst
diff --git a/docs/examples/config/nc_ogbn_arxiv.rst b/docs/examples/config/nc_ogbn_arxiv.rst
diff --git a/docs/examples/config/resume_training.rst b/docs/examples/config/resume_training.rst
diff --git a/docs/examples/index.rst b/docs/examples/index.rst
diff --git a/docs/examples/introduction.rst b/docs/examples/introduction.rst
diff --git a/docs/examples/prediction/command_line.rst b/docs/examples/prediction/command_line.rst
diff --git a/docs/examples/prediction/python.rst b/docs/examples/prediction/python.rst
diff --git a/docs/examples/preprocessing/command_line.rst b/docs/examples/preprocessing/command_line.rst
diff --git a/docs/examples/preprocessing/python.rst b/docs/examples/preprocessing/python.rst
diff --git a/docs/examples/python/index.rst b/docs/examples/python/index.rst
diff --git a/docs/examples/python/lp_custom.rst b/docs/examples/python/lp_custom.rst
diff --git a/docs/examples/python/lp_fb15k237.rst b/docs/examples/python/lp_fb15k237.rst
diff --git a/docs/examples/python/nc_ogbn_arxiv.rst b/docs/examples/python/nc_ogbn_arxiv.rst
diff --git a/docs/export_and_inference/index.rst b/docs/export_and_inference/index.rst
diff --git a/docs/export_and_inference/marius_postprocess.rst b/docs/export_and_inference/marius_postprocess.rst
diff --git a/docs/export_and_inference/marius_predict.rst b/docs/export_and_inference/marius_predict.rst
diff --git a/docs/favicon.ico b/docs/favicon.ico
diff --git a/docs/graph_learning/decoders.rst b/docs/graph_learning/decoders.rst
diff --git a/docs/graph_learning/downstream_tasks.rst b/docs/graph_learning/downstream_tasks.rst
diff --git a/docs/graph_learning/encoders.rst b/docs/graph_learning/encoders.rst
diff --git a/docs/graph_learning/index.rst b/docs/graph_learning/index.rst
diff --git a/docs/graph_learning/intro.rst b/docs/graph_learning/intro.rst
diff --git a/docs/graph_learning/learning_tasks.rst b/docs/graph_learning/learning_tasks.rst
diff --git a/docs/index.rst b/docs/index.rst
diff --git a/docs/introduction.rst b/docs/introduction.rst
diff --git a/docs/marius_logo.png b/docs/marius_logo.png
diff --git a/docs/marius_logo.svg b/docs/marius_logo.svg
diff --git a/docs/marius_logo_scaled.png b/docs/marius_logo_scaled.png
diff --git a/docs/preprocess_datasets/built_in.rst b/docs/preprocess_datasets/built_in.rst
diff --git a/docs/preprocess_datasets/command_line.rst b/docs/preprocess_datasets/command_line.rst
diff --git a/docs/preprocess_datasets/index.rst b/docs/preprocess_datasets/index.rst
diff --git a/docs/preprocess_datasets/python.rst b/docs/preprocess_datasets/python.rst
diff --git a/docs/python_api/configuration/index.rst b/docs/python_api/configuration/index.rst
diff --git a/docs/python_api/data/batch.rst b/docs/python_api/data/batch.rst
diff --git a/docs/python_api/data/dataloader.rst b/docs/python_api/data/dataloader.rst
diff --git a/docs/python_api/data/dense_graph.rst b/docs/python_api/data/dense_graph.rst
diff --git a/docs/python_api/data/graph.rst b/docs/python_api/data/graph.rst
diff --git a/docs/python_api/data/index.rst b/docs/python_api/data/index.rst
diff --git a/docs/python_api/data/samplers/edge.rst b/docs/python_api/data/samplers/edge.rst
diff --git a/docs/python_api/data/samplers/index.rst b/docs/python_api/data/samplers/index.rst
diff --git a/docs/python_api/data/samplers/negative.rst b/docs/python_api/data/samplers/negative.rst
diff --git a/docs/python_api/data/samplers/neighbor.rst b/docs/python_api/data/samplers/neighbor.rst
diff --git a/docs/python_api/index.rst b/docs/python_api/index.rst
diff --git a/docs/python_api/manager/index.rst b/docs/python_api/manager/index.rst
diff --git a/docs/python_api/nn/activation.rst b/docs/python_api/nn/activation.rst
diff --git a/docs/python_api/nn/decoders/decoder.rst b/docs/python_api/nn/decoders/decoder.rst
diff --git a/docs/python_api/nn/decoders/edge/comparators.rst b/docs/python_api/nn/decoders/edge/comparators.rst
diff --git a/docs/python_api/nn/decoders/edge/complex.rst b/docs/python_api/nn/decoders/edge/complex.rst
diff --git a/docs/python_api/nn/decoders/edge/distmult.rst b/docs/python_api/nn/decoders/edge/distmult.rst
diff --git a/docs/python_api/nn/decoders/edge/edge_decoder.rst b/docs/python_api/nn/decoders/edge/edge_decoder.rst
diff --git a/docs/python_api/nn/decoders/edge/index.rst b/docs/python_api/nn/decoders/edge/index.rst
diff --git a/docs/python_api/nn/decoders/edge/relation_operators.rst b/docs/python_api/nn/decoders/edge/relation_operators.rst
diff --git a/docs/python_api/nn/decoders/edge/transe.rst b/docs/python_api/nn/decoders/edge/transe.rst
diff --git a/docs/python_api/nn/decoders/index.rst b/docs/python_api/nn/decoders/index.rst
diff --git a/docs/python_api/nn/decoders/node/index.rst b/docs/python_api/nn/decoders/node/index.rst
diff --git a/docs/python_api/nn/decoders/node/node_decoder.rst b/docs/python_api/nn/decoders/node/node_decoder.rst
diff --git a/docs/python_api/nn/decoders/node/noop_node_decoder.rst b/docs/python_api/nn/decoders/node/noop_node_decoder.rst
diff --git a/docs/python_api/nn/encoders/general_encoder.rst b/docs/python_api/nn/encoders/general_encoder.rst
diff --git a/docs/python_api/nn/encoders/index.rst b/docs/python_api/nn/encoders/index.rst
diff --git a/docs/python_api/nn/index.rst b/docs/python_api/nn/index.rst
diff --git a/docs/python_api/nn/initialization.rst b/docs/python_api/nn/initialization.rst
diff --git a/docs/python_api/nn/layers/embedding.rst b/docs/python_api/nn/layers/embedding.rst
diff --git a/docs/python_api/nn/layers/feature.rst b/docs/python_api/nn/layers/feature.rst
diff --git a/docs/python_api/nn/layers/gnn.rst b/docs/python_api/nn/layers/gnn.rst
diff --git a/docs/python_api/nn/layers/index.rst b/docs/python_api/nn/layers/index.rst
diff --git a/docs/python_api/nn/layers/layer.rst b/docs/python_api/nn/layers/layer.rst
diff --git a/docs/python_api/nn/layers/reduction.rst b/docs/python_api/nn/layers/reduction.rst
diff --git a/docs/python_api/nn/loss.rst b/docs/python_api/nn/loss.rst
diff --git a/docs/python_api/nn/model.rst b/docs/python_api/nn/model.rst
diff --git a/docs/python_api/nn/optim.rst b/docs/python_api/nn/optim.rst
diff --git a/docs/python_api/pipeline/evaluator.rst b/docs/python_api/pipeline/evaluator.rst
diff --git a/docs/python_api/pipeline/graph_encoder.rst b/docs/python_api/pipeline/graph_encoder.rst
diff --git a/docs/python_api/pipeline/index.rst b/docs/python_api/pipeline/index.rst
diff --git a/docs/python_api/pipeline/trainer.rst b/docs/python_api/pipeline/trainer.rst
diff --git a/docs/python_api/reporting/index.rst b/docs/python_api/reporting/index.rst
diff --git a/docs/python_api/reporting/metrics.rst b/docs/python_api/reporting/metrics.rst
diff --git a/docs/python_api/reporting/reporters.rst b/docs/python_api/reporting/reporters.rst
diff --git a/docs/python_api/storage/graph_storage.rst b/docs/python_api/storage/graph_storage.rst
diff --git a/docs/python_api/storage/index.rst b/docs/python_api/storage/index.rst
diff --git a/docs/python_api/storage/storage.rst b/docs/python_api/storage/storage.rst
diff --git a/docs/python_api/tools/configuration/constants.rst b/docs/python_api/tools/configuration/constants.rst
diff --git a/docs/python_api/tools/configuration/datatypes.rst b/docs/python_api/tools/configuration/datatypes.rst
diff --git a/docs/python_api/tools/configuration/index.rst b/docs/python_api/tools/configuration/index.rst
diff --git a/docs/python_api/tools/configuration/marius_config.rst b/docs/python_api/tools/configuration/marius_config.rst
diff --git a/docs/python_api/tools/index.rst b/docs/python_api/tools/index.rst
diff --git a/docs/python_api/tools/preprocess/converters/index.rst b/docs/python_api/tools/preprocess/converters/index.rst
diff --git a/docs/python_api/tools/preprocess/datasets/index.rst b/docs/python_api/tools/preprocess/datasets/index.rst
diff --git a/docs/python_api/tools/preprocess/index.rst b/docs/python_api/tools/preprocess/index.rst
diff --git a/docs/python_api/tools/preprocess/partitioners/index.rst b/docs/python_api/tools/preprocess/partitioners/index.rst
diff --git a/docs/python_api/tools/preprocess/readers/index.rst b/docs/python_api/tools/preprocess/readers/index.rst
diff --git a/docs/python_api/tools/preprocess/writers/index.rst b/docs/python_api/tools/preprocess/writers/index.rst
diff --git a/docs/quickstart.rst b/docs/quickstart.rst
diff --git a/examples/configuration/custom_lp.yaml b/examples/configuration/custom_lp.yaml
diff --git a/examples/configuration/custom_nc.yaml b/examples/configuration/custom_nc.yaml
diff --git a/examples/configuration/fb15k_237.yaml b/examples/configuration/fb15k_237.yaml
diff --git a/examples/configuration/ogbn_arxiv.yaml b/examples/configuration/ogbn_arxiv.yaml
diff --git a/examples/configuration/sakila.yaml b/examples/configuration/sakila.yaml
diff --git a/examples/db2graph/dockerfile b/examples/db2graph/dockerfile
diff --git a/examples/db2graph/run.sh b/examples/db2graph/run.sh
diff --git a/examples/docker/README.md b/examples/docker/README.md
diff --git a/examples/docker/cpu_ubuntu/dockerfile b/examples/docker/cpu_ubuntu/dockerfile
diff --git a/examples/docker/gpu_ubuntu/dockerfile b/examples/docker/gpu_ubuntu/dockerfile
diff --git a/examples/preprocessing/custom_dataset.py b/examples/preprocessing/custom_dataset.py
diff --git a/examples/python/custom.py b/examples/python/custom.py
diff --git a/examples/python/custom_lp.py b/examples/python/custom_lp.py
diff --git a/examples/python/custom_nc_graphsage.py b/examples/python/custom_nc_graphsage.py
diff --git a/examples/python/fb15k_237.py b/examples/python/fb15k_237.py
diff --git a/examples/python/fb15k_237_gpu.py b/examples/python/fb15k_237_gpu.py
diff --git a/examples/python/ogbn_arxiv_nc.py b/examples/python/ogbn_arxiv_nc.py
diff --git a/pyproject.toml b/pyproject.toml
diff --git a/setup.cfg b/setup.cfg
diff --git a/setup.py b/setup.py
@@ -16,7 +16,7 @@ def __init__(self, name, sourcedir=""):
 class CMakeBuild(build_ext):
     def run(self):
         try:
-            _ = subprocess.check_output(["cmake", "--version"])
+            _ = subprocess.run(["cmake", "--version"])
         except OSError:
             raise RuntimeError(
                 "CMake must be installed to build the following extensions: "
@@ -69,8 +69,8 @@ def build_extension(self, ext):
 
         print(cmake_args)
 
-        subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=self.build_temp, env=env)
-        subprocess.check_call(["cmake", "--build", ".", "--target", "bindings"] + build_args, cwd=self.build_temp)
+        subprocess.run(["cmake", ext.sourcedir] + cmake_args, cwd=self.build_temp, env=env)
+        subprocess.run(["cmake", "--build", ".", "--target", "bindings"] + build_args, cwd=self.build_temp)
         print()  # Add an empty line for cleaner output
 
 

diff --git a/simulator/__init__.py b/simulator/__init__.py
diff --git a/simulator/configs/arvix_linear.json b/simulator/configs/arvix_linear.json
diff --git a/simulator/configs/arvix_linear_in_mem.json b/simulator/configs/arvix_linear_in_mem.json
diff --git a/simulator/configs/arvix_metis.json b/simulator/configs/arvix_metis.json
diff --git a/simulator/configs/arvix_neighbors.json b/simulator/configs/arvix_neighbors.json
diff --git a/simulator/configs/arvix_random.json b/simulator/configs/arvix_random.json
diff --git a/simulator/configs/papers.json b/simulator/configs/papers.json
diff --git a/simulator/configs/products_linear.json b/simulator/configs/products_linear.json
@@ -6,5 +6,6 @@
         "feature_dimension" : 100,
         "feature_size" : "float32"
     },
-    "sampling_depth" : 1
+    "sampling_depth" : 1,
+    "batch_size" : 11
 }
diff --git a/simulator/images/products_linear.png b/simulator/images/products_linear.png
diff --git a/simulator/main.py b/simulator/main.py
@@ -3,6 +3,7 @@
 import argparse
 import random
 import time
+import datetime
 
 from src.dataset_loader import *
 from src.features_loader import *
@@ -19,7 +20,7 @@ def read_arguments():
     parser.add_argument("--config_file", type=str, help="The config file containing the details for the simulation")
     parser.add_argument("--save_path", required=True, type=str, help="The path to save the resulting image to")
     parser.add_argument("--graph_title", required=True, type=str, help="The title of the saved graph")
-    parser.add_argument("--num_nodes", default = -1, type = int, help = "The number of nodes we want in our sample")
+    parser.add_argument("--sample_percentage", default = 50, type = int, help = "The percentage of batches to sample to generate CDF")
     parser.add_argument("--log_rate", type=int, default=20, help="Log rate of the nodes processed")
     return parser.parse_args()
 
@@ -30,37 +31,42 @@ def main():
 
     # Create the loaders
     data_loader = DatasetLoader(config)
-    print(data_loader.get_num_nodes(), data_loader.get_num_edges())
-    '''
     features_loader = get_featurizer(data_loader, config["features_stats"])
     sampler = SubgraphSampler(data_loader, features_loader, config)
-    print("Finished loading all objects")
 
-    # Perform sampling
-    nodes_to_sample = [i for i in range(data_loader.get_num_nodes())]
-    random.shuffle(nodes_to_sample)
-    if arguments.num_nodes > 0:
-        nodes_to_sample = nodes_to_sample[ : arguments.num_nodes]
-    log_rate = int(len(nodes_to_sample) / arguments.log_rate)
+    # Determine the sample nodes
+    sample_percentage = arguments.sample_percentage
+    all_nodes = data_loader.get_nodes_sorted_by_incoming()
+    batch_size = int(config["batch_size"])
+    all_nodes = all_nodes.reshape((-1, batch_size))
 
+    num_sample_nodes = int((sample_percentage * all_nodes.shape[0])/100.0) 
+    batches_to_keep = np.random.choice(all_nodes.shape[0], num_sample_nodes, replace = False)
+    sample_nodes = all_nodes[batches_to_keep, : ]
+    print("Processing", sample_nodes.shape[0], "batches with each batch having", batch_size, "nodes")
+
+    # Perform the sampling
+    log_rate = int(sample_nodes.shape[0]/arguments.log_rate)
     pages_loaded = []
-    for curr_node in nodes_to_sample:
-        num_pages_read = sampler.perform_sampling_for_node(curr_node)
-        if num_pages_read > 0:
-            pages_loaded.append(num_pages_read)
+    for idx, curr_batch in enumerate(sample_nodes):
+        average_pages_loaded = sampler.perform_sampling_for_nodes(curr_batch)
+        if average_pages_loaded > 0:
+            pages_loaded.append(average_pages_loaded)
 
-        if len(pages_loaded) > 0 and len(pages_loaded) % log_rate == 0:
-            percentage_finished = (100.0 * len(pages_loaded)) / len(nodes_to_sample)
+        if idx > 0 and idx % log_rate == 0:
+            percentage_finished = (100.0 * idx)/sample_nodes.shape[0]
             print("Finished processing", round(percentage_finished), "percent of nodes")
 
     # Get the arguments to log
-    vals_to_log = dict()
+    total_time = time.time() - start_time
+    vals_to_log = {
+        "Total Processing Time" : str(datetime.timedelta(seconds = int(total_time))),
+        "Batch Size" : batch_size,
+        "Sample Percentage" : sample_percentage,
+    }
+
     for curr_obj in [data_loader, features_loader, sampler]:
         vals_to_log.update(curr_obj.get_values_to_log())
-    
-    # Log the time taken
-    total_time = time.time() - start_time
-    print("Processed all", len(nodes_to_sample), "nodes in", total_time, "seconds")
 
     # Save the histogram
     os.makedirs(os.path.dirname(arguments.save_path), exist_ok=True)
@@ -73,7 +79,6 @@ def main():
         "values_to_log": vals_to_log,
     }
     visualize_results(visualize_arguments)
-    '''
 
 if __name__ == "__main__":
     main()
diff --git a/simulator/src/__init__.py b/simulator/src/__init__.py
diff --git a/simulator/src/dataset_loader.py b/simulator/src/dataset_loader.py
@@ -2,8 +2,10 @@
 import os
 import numpy as np
 import torch
+import traceback
 from collections import defaultdict
-import marius.storage
+from marius.data import Batch, DENSEGraph, MariusGraph
+from marius.data.samplers import LayeredNeighborSampler
 
 class DatasetLoader:
     SAVE_DIR = "datasets"
@@ -33,47 +35,40 @@ def load_dataset(self):
         edges_flaten_arr = np.frombuffer(edges_bytes, dtype=np.int32)
         edges_arr = edges_flaten_arr.reshape((-1, 2))
 
-        '''
         # Create the graph
-        self.edge_list = torch.from_numpy(edges_arr)
-        self.nodes = torch.unique(self.edge_list)
-        self.current_graph = MariusGraph(self.edge_list, self.edge_list[torch.argsort(edge_list[:, -1])], self.get_num_nodes())
-        self.sampler = LayeredNeighborSampler(full_graph, [-1 for _ in range(self.sampling_depth)])
-
-        # Neighbors cache
-        '''
+        self.edge_list = torch.tensor(edges_arr, dtype = torch.int64)
+        self.total_nodes = torch.max(self.edge_list).item() + 1
+        self.current_graph = MariusGraph(self.edge_list, self.edge_list[torch.argsort(self.edge_list[:, -1])], self.total_nodes)
+        self.sampler = LayeredNeighborSampler(self.current_graph, [-1 for _ in range(self.sampling_depth)])
 
     def get_num_nodes(self):
-        return self.nodes.shape(0)
-
-    def get_neigbhors_for_node(self, node_id, all_depths = False):
-        if node_id not in self.adjacency_map:
-            return []
-
-        return list(self.adjacency_map[node_id])
-
-    def get_incoming_neighbors(self, node_id):
-        if node_id not in self.num_incoming_edges:
-            return 0
-
-        return self.num_incoming_edges[node_id]
+        return self.total_nodes
+
+    def get_neigbhors_for_nodes(self, nodes):
+        # Get the neighbors for the nodes
+        try:
+            nodes_to_sample = torch.tensor(nodes, dtype = torch.int64)
+            sampled_nodes = self.sampler.getNeighbors(nodes_to_sample)
+            sampled_nodes.performMap()
+            return sampled_nodes.getNeighborIDs(True, True).numpy()
+        except:
+            return np.array([])
 
     def get_num_edges(self):
-        return self.edge_list.shape(0)
+        return self.edge_list.size(0)
+
+    def get_nodes_sorted_by_incoming(self):
+        return torch.argsort(torch.bincount(self.edge_list[ : , 1]), descending=True).numpy()
 
     def get_average_neighbors(self):
-        neighbors_count = []
-        for node_neighbors in self.adjacency_map.values():
-            neighbors_count.append(len(node_neighbors))
-
-        return np.mean(np.array(neighbors_count))
+        outgoing_nodes = self.edge_list[ : , 0]
+        outgoing_unique_nodes = torch.unique(outgoing_nodes)
+        return outgoing_nodes.size(0)/outgoing_unique_nodes.size(0)
 
     def get_average_incoming(self):
-        incoming_counts = []
-        for num_incoming in self.num_incoming_edges.values():
-            incoming_counts.append(num_incoming)
-
-        return np.mean(np.array(incoming_counts))
+        incoming_nodes = self.edge_list[ : , 1]
+        incoming_unique_nodes = torch.unique(incoming_nodes)
+        return incoming_nodes.size(0)/incoming_unique_nodes.size(0)
 
     def get_values_to_log(self):
         return {

diff --git a/simulator/src/features_loader.py b/simulator/src/features_loader.py
@@ -19,17 +19,19 @@ def __init__(self, data_loader, features_stat):
 
     def initialize(self):
         total_nodes = self.data_loader.get_num_nodes()
-        print("Total nodes of", total_nodes)
         self.total_pages = int(math.ceil(total_nodes / (1.0 * self.nodes_per_page)))
-        self.node_location_map = [i for i in range(self.total_pages)]
+        self.node_location_map = np.arange(total_nodes)
         if "feature_layout" in self.features_stat and self.features_stat["feature_layout"] == "random":
             random.shuffle(self.node_location_map)
+        self.node_to_page = (self.node_location_map/self.nodes_per_page).astype(int)
 
     def get_node_page(self, src_node, neighbor_node):
-        print("Querying for node", neighbor_node, "with total of", len(self.node_location_map), "nodes")
         start_node = int(self.node_location_map[neighbor_node] / self.nodes_per_page)
         curr_page_nodes = set(range(start_node, start_node + self.nodes_per_page))
         return curr_page_nodes
+
+    def num_pages_for_nodes(self, nodes):
+        return np.unique(self.node_to_page[nodes]).shape[0]
 
     def get_single_node_feature_size(self):
         return self.node_feature_size
Original file line number	Diff line number	Diff line change
Expand Up		@@ -171,3 +171,4 @@ Thumbs.db


		# End of https://www.toptal.com/developers/gitignore/api/python
		simulator/datasets