From 72b75fd9feeeb137d3ad671419a8e6a574d2a041 Mon Sep 17 00:00:00 2001
From: virio-andreyana <114650479+virio-andreyana@users.noreply.github.com>
Date: Fri, 28 Apr 2023 03:43:20 +0200
Subject: [PATCH] merge add_nice_carrier_names with add_missing_carrier to
 become add_missing_carrier_with_nice_names(n, config)

---
 scripts/add_electricity.py        | 33 ++++++++++++++-----------------
 scripts/add_existing_baseyear.py  |  4 ++--
 scripts/add_extra_components.py   |  7 ++-----
 scripts/prepare_sector_network.py |  4 ++--
 4 files changed, 21 insertions(+), 27 deletions(-)

diff --git a/scripts/add_electricity.py b/scripts/add_electricity.py
index 348e49af..456cc48c 100755
--- a/scripts/add_electricity.py
+++ b/scripts/add_electricity.py
@@ -121,13 +121,26 @@ def calculate_annuity(n, r):
         return 1 / n
 
 
-def add_missing_carrier(n):
+def add_missing_carriers_with_nice_names(n,config):
     components = [n.buses, n.generators, n.lines, n.links, n.storage_units, n.stores]
     for c in components:
         missing_carrier = np.setdiff1d(c.carrier.unique(), n.carriers.index)
         if len(missing_carrier):
             n.madd("Carrier", missing_carrier)
 
+    carrier_i = n.carriers.index
+    nice_names = (
+        pd.Series(config["plotting"]["nice_names"])
+        .reindex(carrier_i)
+        .fillna(carrier_i.to_series().str.title())
+    )
+    n.carriers["nice_name"] = nice_names
+    colors = pd.Series(config["plotting"]["tech_colors"]).reindex(carrier_i)
+    if colors.isna().any():
+        missing_i = list(colors.index[colors.isna()])
+        logger.warning(f"tech_colors for carriers {missing_i} not defined in config.")
+    n.carriers["color"] = colors
+
 
 def _add_missing_carriers_from_costs(n, costs, carriers):
     missing_carriers = pd.Index(carriers).difference(n.carriers.index)
@@ -694,21 +707,6 @@ def estimate_renewable_capacities(n, config):
             )
 
 
-def add_nice_carrier_names(n, config):
-    carrier_i = n.carriers.index
-    nice_names = (
-        pd.Series(config["plotting"]["nice_names"])
-        .reindex(carrier_i)
-        .fillna(carrier_i.to_series().str.title())
-    )
-    n.carriers["nice_name"] = nice_names
-    colors = pd.Series(config["plotting"]["tech_colors"]).reindex(carrier_i)
-    if colors.isna().any():
-        missing_i = list(colors.index[colors.isna()])
-        logger.warning(f"tech_colors for carriers {missing_i} not defined in config.")
-    n.carriers["color"] = colors
-
-
 if __name__ == "__main__":
     if "snakemake" not in globals():
         from _helpers import mock_snakemake
@@ -841,8 +839,7 @@ if __name__ == "__main__":
 
     update_p_nom_max(n)
 
-    add_missing_carrier(n)
-    add_nice_carrier_names(n, snakemake.config)
+    add_missing_carriers_with_nice_names(n, snakemake.config)
 
     n.meta = snakemake.config
     n.export_to_netcdf(snakemake.output[0])
diff --git a/scripts/add_existing_baseyear.py b/scripts/add_existing_baseyear.py
index a8a07322..1435e660 100644
--- a/scripts/add_existing_baseyear.py
+++ b/scripts/add_existing_baseyear.py
@@ -22,7 +22,7 @@ import numpy as np
 import pypsa
 import xarray as xr
 from _helpers import override_component_attrs, update_config_with_sector_opts
-from add_electricity import add_missing_carrier
+from add_electricity import add_missing_carriers_with_nice_names
 from prepare_sector_network import cluster_heat_buses, define_spatial, prepare_costs
 
 cc = coco.CountryConverter()
@@ -668,6 +668,6 @@ if __name__ == "__main__":
 
     n.meta = dict(snakemake.config, **dict(wildcards=dict(snakemake.wildcards)))
 
-    add_missing_carrier(n)
+    add_missing_carriers_with_nice_names(n, snakemake.config)
 
     n.export_to_netcdf(snakemake.output[0])
diff --git a/scripts/add_extra_components.py b/scripts/add_extra_components.py
index 88606305..da47b7b7 100644
--- a/scripts/add_extra_components.py
+++ b/scripts/add_extra_components.py
@@ -58,8 +58,7 @@ import pypsa
 from _helpers import configure_logging
 from add_electricity import (
     _add_missing_carriers_from_costs,
-    add_missing_carrier,
-    add_nice_carrier_names,
+    add_missing_carriers_with_nice_names,
     load_costs,
 )
 
@@ -253,9 +252,7 @@ if __name__ == "__main__":
     attach_stores(n, costs, elec_config)
     attach_hydrogen_pipelines(n, costs, elec_config)
 
-    add_missing_carrier(n)
-
-    add_nice_carrier_names(n, snakemake.config)
+    add_missing_carriers_with_nice_names(n, snakemake.config)
 
     n.meta = dict(snakemake.config, **dict(wildcards=dict(snakemake.wildcards)))
     n.export_to_netcdf(snakemake.output[0])
diff --git a/scripts/prepare_sector_network.py b/scripts/prepare_sector_network.py
index a0d2f472..d952cb69 100644
--- a/scripts/prepare_sector_network.py
+++ b/scripts/prepare_sector_network.py
@@ -22,7 +22,7 @@ from _helpers import (
     override_component_attrs,
     update_config_with_sector_opts,
 )
-from add_electricity import add_missing_carrier
+from add_electricity import add_missing_carriers_with_nice_names
 from build_energy_totals import build_co2_totals, build_eea_co2, build_eurostat_co2
 from networkx.algorithms import complement
 from networkx.algorithms.connectivity.edge_augmentation import k_edge_augmentation
@@ -3400,6 +3400,6 @@ if __name__ == "__main__":
 
     n.meta = dict(snakemake.config, **dict(wildcards=dict(snakemake.wildcards)))
 
-    add_missing_carrier(n)
+    add_missing_carriers_with_nice_names(n, snakemake.config)
 
     n.export_to_netcdf(snakemake.output[0])