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

46 lines
1.3 KiB
Python
Executable File

#!/usr/bin/env python3
import geopandas as gpd
import networkx as nx
from scipy.spatial import distance_matrix
from shapely.geometry import LineString
import numpy as np
import pandas as pd
# Load the poles from a shapefile
poles_gdf = gpd.read_file('poles.shp')
# Extract the points coordinates for the distance matrix
points = np.array(list(zip(poles_gdf.geometry.x, poles_gdf.geometry.y)))
# Create a complete distance matrix
dist_matrix = distance_matrix(points, points)
# Create a complete graph from the distance matrix
G = nx.complete_graph(len(points))
# Add edges to the graph with distances as weights
for i, node in enumerate(G.nodes()):
for j, neighbor in enumerate(G.nodes()):
if i != j: # skip if same node
G[node][neighbor]['weight'] = dist_matrix[i][j]
# Compute the minimum spanning tree of the complete graph
mst = nx.minimum_spanning_tree(G)
# Generate the lines for the MST
mst_lines = []
for edge in mst.edges(data=False):
point1 = points[edge[0]]
point2 = points[edge[1]]
line = LineString([point1, point2])
mst_lines.append(line)
# Create a GeoDataFrame from the lines
lines_gdf = gpd.GeoDataFrame(geometry=mst_lines, crs=poles_gdf.crs)
# Save the lines to a new shapefile
lines_gdf.to_file('pole_lines.shp')
print("Pole lines have been saved to pole_lines.shp.")