diff --git a/tornettools/stage.py b/tornettools/stage.py index 150561a..b05aa81 100644 --- a/tornettools/stage.py +++ b/tornettools/stage.py @@ -178,11 +178,34 @@ def stage_graph(args): # create new graph and copy only the nodes network = nx.classes.function.create_empty_copy(network) + if args.remove_nodes_without_country_code: + # prune nodes which don't contain the country_code key + network = remove_nodes_without_key(network_graph=network, key_to_check="country_code") + # it takes networkx a few minutes to read the atlas graph, so we save a smaller graph containing # only the atlas graph nodes so that the 'generate' step can read these nodes much quicker network_info_path = f"{args.prefix}/networkinfo_staging.gml" nx.readwrite.gml.write_gml(network, network_info_path) +def remove_nodes_without_key(network_graph, key_to_check): + nodes_to_remove = [] + for node in network_graph.nodes: + node_properties_dict = network_graph.nodes[node] + try: + node_properties_dict[key_to_check] + except KeyError: + nodes_to_remove.append(node) + + if len(nodes_to_remove) > 0: + network_graph.remove_nodes_from(nodes_to_remove) + logging.info( + f"Removed {len(nodes_to_remove)} nodes missing key {key_to_check}: " + f"{nodes_to_remove}" + ) + else: + logging.info(f"All nodes in network graph contain the required key {key_to_check}.") + return network_graph + def get_file_list(dir_path): file_paths = [] for root, _, filenames in os.walk(dir_path): diff --git a/tornettools/tornettools b/tornettools/tornettools index 7405947..5055cc5 100644 --- a/tornettools/tornettools +++ b/tornettools/tornettools @@ -221,6 +221,14 @@ def main(): help="""Compress output with xz.""", action="store_true", dest="do_compress", default=False) + + stage_parser.add_argument('--remove_nodes_without_country_code', + help="""Include only those nodes in the networkinfo_staging.gml + network graph that specify a country_code""", + action="store_true", + default=False, + dest="remove_nodes_without_country_code", + ) ############ # generate #