From f19af22b0969ff84abd611edfc711464ec94817f Mon Sep 17 00:00:00 2001 From: Louis Gregg Date: Tue, 20 Oct 2026 20:54:25 +0200 Subject: [PATCH 1/2] A check to ensure the country_code key exists in all nodes in networkinfo_staging.gml. Does not check for other node keys (yet) --- tornettools/stage.py | 23 +++++++++++++++++++++++ tornettools/tornettools | 8 ++++++++ 2 files changed, 31 insertions(+) diff --git a/tornettools/stage.py b/tornettools/stage.py index 150561a..f041742 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 (occurs once, node "284") + network = remove_nodes_without_country_code(network_graph=network) + # 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_country_code(network_graph, key_to_check = "country_code"): + 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 nodes_to_remove: + 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 # From 7f553f9bf8253a62e6837df36a966d2619226a2d Mon Sep 17 00:00:00 2001 From: Louis Gregg Date: Thu, 22 Oct 2026 17:06:46 +0200 Subject: [PATCH 2/2] function renamed to remove_nodes_without_key and key_to_check arg now mandatory --- tornettools/stage.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tornettools/stage.py b/tornettools/stage.py index f041742..b05aa81 100644 --- a/tornettools/stage.py +++ b/tornettools/stage.py @@ -179,15 +179,15 @@ def stage_graph(args): 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 (occurs once, node "284") - network = remove_nodes_without_country_code(network_graph=network) + # 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_country_code(network_graph, key_to_check = "country_code"): +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] @@ -196,7 +196,7 @@ def remove_nodes_without_country_code(network_graph, key_to_check = "country_cod except KeyError: nodes_to_remove.append(node) - if nodes_to_remove: + 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}: "