Move processing of STOCK

etsap-TIMES · Mar 7, 2024 · 112ab74 · 112ab74
1 parent 6611d63
commit 112ab74
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 57 deletions.
diff --git a/xl2times/__main__.py b/xl2times/__main__.py
@@ -119,7 +119,6 @@ def convert_xl_to_times(
         transforms.include_tables_source,
         transforms.merge_tables,
         transforms.complete_processes,
-        transforms.apply_more_fixups,
         transforms.process_units,
         transforms.complete_commodity_groups,
         transforms.process_wildcards,

diff --git a/xl2times/transforms.py b/xl2times/transforms.py
@@ -2729,63 +2729,43 @@ def apply_final_fixup(
             axis=1,
         )
 
-    tables[datatypes.Tag.fi_t] = df
-
-    return tables
-
+    # Handle STOCK specified for a single year
+    i = (df["attribute"] == "STOCK") & df["process"].notna()
+    # Temporary solution to include only processes defined in BASE
+    i_vt = i & (df["source_filename"].str.contains("VT_", case=False))
+    if any(i):
+        extra_rows = []
+        for region in df[i]["region"].unique():
+            i_reg = i & (df["region"] == region)
+            for process in df[(i_reg & i_vt)]["process"].unique():
+                i_reg_prc = i_reg & (df["process"] == process)
+                if any(i_reg_prc):
+                    extra_rows.append(["NCAP_BND", region, process, "UP", 0, 2])
+                # TODO: TIMES already handles this. Drop?
+                if len(df[i_reg_prc]["year"].unique()) == 1:
+                    year = df[i_reg_prc]["year"].unique()[0]
+                    i_attr = (
+                        df["attribute"].isin({"NCAP_TLIFE", "LIFE"})
+                        & (df["region"] == region)
+                        & (df["process"] == process)
+                    )
+                    if any(i_attr):
+                        lifetime = df[i_attr]["value"].unique()[-1]
+                    else:
+                        lifetime = 30
+                    extra_rows.append(
+                        ["STOCK", region, process, "", year + lifetime, 0]
+                    )
+        if len(extra_rows) > 0:
+            cols = ["attribute", "region", "process", "limtype", "year", "value"]
+            df = pd.concat(
+                [
+                    df,
+                    pd.DataFrame(extra_rows, columns=cols),
+                ]
+            )
 
-def apply_more_fixups(
-    config: datatypes.Config,
-    tables: Dict[str, DataFrame],
-    model: datatypes.TimesModel,
-) -> Dict[str, DataFrame]:
-    # TODO: This should only be applied to processes introduced in BASE
-    df = tables.get(datatypes.Tag.fi_t)
-    if df is not None:
-        index = df["attribute"] == "STOCK"
-        # Temporary solution to include only processes defined in BASE
-        i_vt = index & (df["source_filename"].str.contains("VT_", case=False))
-        if any(index):
-            extra_rows = []
-            for region in df[index]["region"].unique():
-                i_reg = index & (df["region"] == region)
-                for process in df[(i_reg & i_vt)]["process"].unique():
-                    i_reg_prc = i_reg & (df["process"] == process)
-                    if any(i_reg_prc):
-                        extra_rows.append(["NCAP_BND", region, process, "UP", 0, 2])
-                    # TODO: TIMES already handles this. Drop?
-                    if len(df[i_reg_prc]["year"].unique()) == 1:
-                        year = df[i_reg_prc]["year"].unique()[0]
-                        i_attr = (
-                            df["attribute"].isin({"NCAP_TLIFE", "LIFE"})
-                            & (df["region"] == region)
-                            & (df["process"] == process)
-                        )
-                        if any(i_attr):
-                            lifetime = df[i_attr]["value"].unique()[-1]
-                        else:
-                            lifetime = 30
-                        extra_rows.append(
-                            ["STOCK", region, process, "", year + lifetime, 0]
-                        )
-            if len(extra_rows) > 0:
-                df = pd.concat(
-                    [
-                        df,
-                        pd.DataFrame(
-                            extra_rows,
-                            columns=[
-                                "attribute",
-                                "region",
-                                "process",
-                                "limtype",
-                                "year",
-                                "value",
-                            ],
-                        ),
-                    ]
-                )
-        tables[datatypes.Tag.fi_t] = df
+    tables[datatypes.Tag.fi_t] = df
 
     return tables