From a3af137b74d6e505e2dae32fc7cea2b2226f834c Mon Sep 17 00:00:00 2001
From: martacki <martha.frysztacki@kit.edu>
Date: Mon, 27 Jun 2022 14:18:47 +0200
Subject: [PATCH] clustering strats to configurables: move duplicate code to
 _helpers script & import

---
 scripts/_helpers.py         | 16 ++++++++++++++++
 scripts/cluster_network.py  |  8 ++------
 scripts/simplify_network.py | 14 +++-----------
 3 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/scripts/_helpers.py b/scripts/_helpers.py
index 6e47c053..3c3e6ac7 100644
--- a/scripts/_helpers.py
+++ b/scripts/_helpers.py
@@ -210,6 +210,22 @@ def progress_retrieve(url, file):
 
     urllib.request.urlretrieve(url, file, reporthook=dlProgress)
 
+def get_aggregation_strategies(aggregation_strategies):
+    # default aggregation strategies that cannot be defined in .yaml format must be specified within
+    # the function, otherwise (when defaults are passed in the function's definition) they get lost
+    # when custom values are specified in the config.
+
+    import numpy as np
+    from pypsa.networkclustering import _make_consense
+
+    bus_strategies = dict(country=_make_consense("Bus", "country"))
+    bus_strategies.update(aggregation_strategies.get("buses", {}))
+
+    generator_strategies = {'build_year': lambda x: 0, 'lifetime': lambda x: np.inf}
+    generator_strategies.update(aggregation_strategies.get("generators", {}))
+
+    return bus_strategies, generator_strategies
+
 
 def mock_snakemake(rulename, **wildcards):
     """
diff --git a/scripts/cluster_network.py b/scripts/cluster_network.py
index 93cd89cd..6186a00e 100644
--- a/scripts/cluster_network.py
+++ b/scripts/cluster_network.py
@@ -121,7 +121,7 @@ Exemplary unsolved network clustered to 37 nodes:
 """
 
 import logging
-from _helpers import configure_logging, update_p_nom_max
+from _helpers import configure_logging, update_p_nom_max, get_aggregation_strategies
 
 import pypsa
 import os
@@ -261,11 +261,7 @@ def clustering_for_n_clusters(n, n_clusters, custom_busmap=False, aggregate_carr
                               line_length_factor=1.25, aggregation_strategies=dict(), solver_name="cbc",
                               algorithm="kmeans", extended_link_costs=0, focus_weights=None):
 
-    bus_strategies = dict(country=_make_consense("Bus", "country"))
-    bus_strategies.update(aggregation_strategies.get("buses", {}))
-
-    generator_strategies = {'build_year': lambda x: 0, 'lifetime': lambda x: np.inf}
-    generator_strategies.update(aggregation_strategies.get("generators", {}))
+    bus_strategies, generator_strategies = get_aggregation_strategies(aggregation_strategies)
 
     if not isinstance(custom_busmap, pd.Series):
         busmap = busmap_for_n_clusters(n, n_clusters, solver_name, focus_weights, algorithm)
diff --git a/scripts/simplify_network.py b/scripts/simplify_network.py
index 0fda5c77..37e7e698 100644
--- a/scripts/simplify_network.py
+++ b/scripts/simplify_network.py
@@ -83,7 +83,7 @@ The rule :mod:`simplify_network` does up to four things:
 """
 
 import logging
-from _helpers import configure_logging, update_p_nom_max
+from _helpers import configure_logging, update_p_nom_max, get_aggregation_strategies
 
 from cluster_network import clustering_for_n_clusters, cluster_regions
 from add_electricity import load_costs
@@ -203,8 +203,7 @@ def _aggregate_and_move_components(n, busmap, connection_costs_to_bus, output,
 
     _adjust_capital_costs_using_connection_costs(n, connection_costs_to_bus, output)
 
-    generator_strategies = {'build_year': lambda x: 0, 'lifetime': lambda x: np.inf}
-    generator_strategies.update(aggregation_strategies.get("generators", {}))
+    _, generator_strategies = get_aggregation_strategies(aggregation_strategies)
 
     generators, generators_pnl = aggregategenerators(
         n, busmap, custom_strategies=generator_strategies
@@ -355,14 +354,7 @@ def aggregate_to_substations(n, aggregation_strategies=dict(), buses_i=None):
     busmap = n.buses.index.to_series()
     busmap.loc[buses_i] = dist.idxmin(1)
 
-    # default aggregation strategies that cannot be defined in .yaml format must be specified within
-    # the function, otherwise (when defaults are passed in the function's definition) they get lost
-    # in case custom values for different variables are specified in the config.
-    bus_strategies = dict(country=_make_consense("Bus", "country"))
-    bus_strategies.update(aggregation_strategies.get("buses", {}))
-
-    generator_strategies = {'build_year': lambda x: 0, 'lifetime': lambda x: np.inf}
-    generator_strategies.update(aggregation_strategies.get("generators", {}))
+    bus_strategies, generator_strategies = get_aggregation_strategies(aggregation_strategies)
 
     clustering = get_clustering_from_busmap(n, busmap,
                                             bus_strategies=bus_strategies,