pyhld/qgis_scripts/create_network_v3.py
2024-04-19 14:29:59 -05:00

95 lines
4.0 KiB
Python

# Import required libraries
import networkx as nx
import time
from qgis.PyQt.QtCore import QVariant
from qgis.core import QgsProject, QgsVectorLayer, QgsFeature, QgsField, QgsGeometry, QgsPoint
import heapq
# Get layers
edges_layer = QgsProject.instance().mapLayersByName('EDGES')[0]
nodes_layer = QgsProject.instance().mapLayersByName('NODES')[0]
home_points_layer = QgsProject.instance().mapLayersByName('HOME_POINTS')[0]
# Build the graph
G = nx.Graph()
nodes_data = {feature['id']: feature.geometry() for feature in nodes_layer.getFeatures()} # store node geometries by id
edges_data = {(feature['start_node'], feature['end_node']): (feature.geometry(), feature['length'], feature['cost']) for feature in edges_layer.getFeatures()} # store edge geometries by node pairs
for feature in nodes_layer.getFeatures():
G.add_node(feature['id'])
for feature in edges_layer.getFeatures():
start_node = feature['start_node'] # assuming edges layer has start_node & end_node fields
end_node = feature['end_node']
G.add_edge(start_node, end_node, weight=feature['cost'], type=feature['type'], length=feature['length'], cost=feature['cost'])
# Identify home nodes
home_nodes = {feature['drop_point_id'] for feature in home_points_layer.getFeatures() if feature['drop_point_id'] in G}
# Create a subgraph of G that only includes the home nodes and any nodes that connect them
subgraph = nx.Graph()
visited = set()
# Start at any 'home' node
start = next(iter(home_nodes))
visited.add(start)
counter = 0
total = len(home_nodes)
process_times = []
while home_nodes.difference(visited):
start_time = time.time()
# Find the nearest 'home' node that is not in our tree yet
next_node = min(
(node for node in home_nodes if node not in visited),
key=lambda node: nx.shortest_path_length(G, start, node, weight='cost')
)
# Create a subgraph excluding other home nodes except start and next_node
G_sub = G.copy()
remove_nodes = home_nodes - {start} - {next_node}
G_sub.remove_nodes_from(remove_nodes)
# Add the shortest path to this node to our tree
shortest_path = nx.shortest_path(G_sub, start, next_node, weight='cost')
subgraph.add_edges_from(zip(shortest_path, shortest_path[1:]))
visited.add(next_node)
start = next_node
end_time = time.time()
process_times.append(end_time - start_time)
counter += 1
avg_time = sum(process_times) / len(process_times)
remaining = (total - counter) * avg_time
hours, remainder = divmod(remaining, 3600)
minutes, seconds = divmod(remainder, 60)
print(f'Processing home node {counter} of {total}.\nEstimated time remaining: {int(hours)} hours, {int(minutes)} minutes, and {int(seconds)} seconds')
# Create a minimum spanning tree from the subgraph
mst = nx.minimum_spanning_tree(subgraph, weight='cost')
# Prepare the NETWORK layer
network_layer = QgsVectorLayer("LineString", "NETWORK", "memory")
network_prov = network_layer.dataProvider()
# Add 'id', 'type', 'length', and 'cost' fields to the NETWORK layer
network_prov.addAttributes([QgsField("id", QVariant.Int), QgsField("type", QVariant.String), QgsField("length", QVariant.Double), QgsField("cost", QVariant.Double)])
network_layer.updateFields()
# Add edges from the minimum spanning tree to the NETWORK layer
for edge in mst.edges():
edge_geometry, edge_length, edge_cost = edges_data[(edge[0], edge[1])] if (edge[0], edge[1]) in edges_data else edges_data[(edge[1], edge[0])] # fetch geometry, length, cost by node pair
edge_type = G.get_edge_data(edge[0], edge[1])['type'] # get edge 'type' from the original graph G
network_feature = QgsFeature()
network_feature.setGeometry(edge_geometry)
network_feature.setAttributes([edge[0], edge_type, edge_length, edge_cost]) # set the type, length, cost attribute from edge data
network_prov.addFeature(network_feature)
# Update NETWORK layer
network_layer.updateExtents()
# Add NETWORK layer to Layers panel
QgsProject.instance().addMapLayer(network_layer)
print("Done.")