move append_bus_shapes to build_bus_regions, apply it where possible

scripts/build_bus_regions.py

@@ -109,6 +109,34 @@ def voronoi_partition_pts(points, outline):
     return polygons
 
 
+def append_bus_shapes(n, shapes, type):
+    """
+    Append shapes to the network. If shapes with the same component and type
+    already exist, they will be removed.
+
+    Parameters:
+        n (pypsa.Network): The network to which the shapes will be appended.
+        shapes (geopandas.GeoDataFrame): The shapes to be appended.
+        **kwargs: Additional keyword arguments used in `n.madd`.
+
+    Returns:
+        None
+    """
+    remove = n.shapes.query("component == 'Bus' and type == @type").index
+    n.mremove("Shape", remove)
+
+    offset = n.shapes.index.astype(int).max() + 1 if not n.shapes.empty else 0
+    shapes = shapes.rename(lambda x: int(x) + offset)
+    n.madd(
+        "Shape",
+        shapes.index,
+        geometry=shapes.geometry,
+        idx=shapes.name,
+        component="Bus",
+        type=type,
+    )
+
+
 if __name__ == "__main__":
     if "snakemake" not in globals():
         from _helpers import mock_snakemake
@@ -119,7 +147,8 @@ if __name__ == "__main__":
 
     countries = snakemake.params.countries
 
-    n = pypsa.Network(snakemake.input.base_network)
+    base_network = snakemake.input.base_network
+    n = pypsa.Network(base_network)
 
     country_shapes = gpd.read_file(snakemake.input.country_shapes).set_index("name")[
         "geometry"
@@ -175,32 +204,15 @@ if __name__ == "__main__":
 
     shapes = pd.concat(onshore_regions, ignore_index=True)
     shapes.to_file(snakemake.output.regions_onshore)
-
+    append_bus_shapes(n, shapes, "onshore")
-    offset = n.shapes.index.astype(int).max() + 1 if not n.shapes.empty else 0
-    shapes.index = shapes.index.astype(int) + offset
-    n.madd(
-        "Shape",
-        shapes.index,
-        geometry=shapes.geometry,
-        idx=shapes.name,
-        component="Bus",
-        type="onshore",
-    )
 
     if offshore_regions:
         shapes = pd.concat(offshore_regions, ignore_index=True)
         shapes.to_file(snakemake.output.regions_offshore)
-
+        append_bus_shapes(n, shapes, "offshore")
-        offset = n.shapes.index.astype(int).max() + 1 if not n.shapes.empty else 0
-        shapes.index = shapes.index.astype(int) + offset
-        n.madd(
-            "Shape",
-            shapes.index,
-            geometry=shapes.geometry,
-            idx=shapes.name,
-            component="Bus",
-            type="offshore",
-        )
 
     else:
         offshore_shapes.to_frame().to_file(snakemake.output.regions_offshore)
+
+    # save network with shapes
+    n.export_to_netcdf(base_network)

scripts/cluster_network.py

@@ -135,6 +135,7 @@ import pypsa
 import seaborn as sns
 from _helpers import configure_logging, set_scenario_config, update_p_nom_max
 from add_electricity import load_costs
+from build_bus_regions import append_bus_shapes
 from packaging.version import Version, parse
 from pypsa.clustering.spatial import (
     busmap_by_greedy_modularity,
@@ -447,33 +448,6 @@ def cluster_regions(busmaps, regions):
     return regions_c.reset_index()
 
 
-def append_bus_shapes(n, shapes, type):
-    """
-    Append shapes to the network.
-
-    Parameters:
-        n (pypsa.Network): The network to which the shapes will be appended.
-        shapes (geopandas.GeoDataFrame): The shapes to be appended.
-        **kwargs: Additional keyword arguments used in `n.madd`.
-
-    Returns:
-        None
-    """
-    remove = n.shapes.query("component == 'Bus' and type == @which").index
-    n.mremove("Shape", remove)
-
-    offset = n.shapes.index.astype(int).max() + 1 if not n.shapes.empty else 0
-    shapes.index = shapes.index.astype(int) + offset
-    n.madd(
-        "Shape",
-        shapes.index,
-        geometry=shapes.geometry,
-        idx=shapes.name,
-        component="Bus",
-        type=type,
-    )
-
-
 def plot_busmap_for_n_clusters(n, n_clusters, solver_name="scip", fn=None):
     busmap = busmap_for_n_clusters(n, n_clusters, solver_name)
     cs = busmap.unique()
@@ -589,8 +563,8 @@ if __name__ == "__main__":
     for which in ["regions_onshore", "regions_offshore"]:
         regions = gpd.read_file(snakemake.input[which])
         clustered_regions = cluster_regions((clustering.busmap,), regions)
-        append_bus_shapes(nc, clustered_regions, type=which.split("_")[1])
         clustered_regions.to_file(snakemake.output[which])
+        append_bus_shapes(nc, clustered_regions, type=which.split("_")[1])
 
     nc.meta = dict(snakemake.config, **dict(wildcards=dict(snakemake.wildcards)))
     nc.export_to_netcdf(snakemake.output.network)

scripts/simplify_network.py

@@ -88,12 +88,14 @@ The rule :mod:`simplify_network` does up to four things:
 import logging
 from functools import reduce
 
+import geopandas as gpd
 import numpy as np
 import pandas as pd
 import pypsa
 import scipy as sp
 from _helpers import configure_logging, set_scenario_config, update_p_nom_max
 from add_electricity import load_costs
+from build_bus_regions import append_bus_shapes
 from cluster_network import cluster_regions, clustering_for_n_clusters
 from pypsa.clustering.spatial import (
     aggregateoneport,
@@ -610,6 +612,7 @@ if __name__ == "__main__":
     n.lines.drop(remove, axis=1, errors="ignore", inplace=True)
 
     if snakemake.wildcards.simpl:
+        shapes = n.shapes
         n, cluster_map = cluster(
             n,
             int(snakemake.wildcards.simpl),
@@ -619,15 +622,19 @@ if __name__ == "__main__":
             params.simplify_network["feature"],
             params.aggregation_strategies,
         )
+        n.shapes = shapes
         busmaps.append(cluster_map)
 
     update_p_nom_max(n)
 
-    n.meta = dict(snakemake.config, **dict(wildcards=dict(snakemake.wildcards)))
-    n.export_to_netcdf(snakemake.output.network)
-
     busmap_s = reduce(lambda x, y: x.map(y), busmaps[1:], busmaps[0])
     busmap_s.to_csv(snakemake.output.busmap)
 
     for which in ["regions_onshore", "regions_offshore"]:
-        cluster_regions(n, busmaps, which, snakemake.input, snakemake.output)
+        regions = gpd.read_file(snakemake.input[which])
+        clustered_regions = cluster_regions(busmaps, regions)
+        clustered_regions.to_file(snakemake.output[which])
+        append_bus_shapes(n, clustered_regions, type=which.split("_")[1])
+
+    n.meta = dict(snakemake.config, **dict(wildcards=dict(snakemake.wildcards)))
+    n.export_to_netcdf(snakemake.output.network)