From 9947ac629e36b78338e18b9e80332856cbb78ac8 Mon Sep 17 00:00:00 2001
From: romainsacchi <romain@Romains-iMac.psi.ch>
Date: Tue, 5 Mar 2024 15:40:29 +0100
Subject: [PATCH] Add locations

---
 pathways/pathways.py | 19 ++++++++++++++++---
 pathways/utils.py    |  6 ++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/pathways/pathways.py b/pathways/pathways.py
index ae4ffa9..4a36ac0 100644
--- a/pathways/pathways.py
+++ b/pathways/pathways.py
@@ -179,6 +179,19 @@ def fetch_indices(mapping, regions, variables, A_index, geo):
     return vars_idx
 
 
+def fetch_inventories_locations(A_index: Dict[str, Tuple[str, str, str]]) -> List[str]:
+    """
+    Fetch the locations of the inventories.
+    :param A_index: Dictionary with the indices of the activities in the technosphere matrix.
+    :return: List of locations.
+    """
+    locations = []
+    for act in A_index:
+        if act[1] not in locations:
+            locations.append(act[1])
+    return locations
+
+
 def generate_A_indices(A_index, reverse_classifications, lca_results_coords):
     # Generate a list of activity indices for each activity category
     acts_idx = []
@@ -549,9 +562,8 @@ def calculate(
             methods = None
 
         # Set default values if arguments are not provided
-        if methods is None:
-            if characterization is True:
-                methods = get_lcia_method_names()
+        if methods is None and characterization is True:
+            methods = get_lcia_method_names()
         if models is None:
             models = self.scenarios.coords["model"].values
             models = [m.lower() for m in models]
@@ -614,6 +626,7 @@ def calculate(
 
                     # Create xarray for storing LCA results if not already present
                     if self.lca_results is None:
+                        locations = fetch_inventories_locations(A_index)
                         self.lca_results = create_lca_results_array(
                             methods,
                             B_index,
diff --git a/pathways/utils.py b/pathways/utils.py
index 45fe8c2..f922357 100644
--- a/pathways/utils.py
+++ b/pathways/utils.py
@@ -163,6 +163,7 @@ def create_lca_results_array(
     B_indices: Dict,
     years: List[int],
     regions: List[str],
+    locations: List[str],
     models: List[str],
     scenarios: List[str],
     classifications: dict,
@@ -180,6 +181,8 @@ def create_lca_results_array(
     :type years: List[int]
     :param regions: List of regions to consider in the LCA results.
     :type regions: List[str]
+    :param locations: List of locations to consider in the LCA results.
+    :type locations: List[str]
     :param models: List of models to consider in the LCA results.
     :type models: List[str]
     :param scenarios: List of scenarios to consider in the LCA results.
@@ -189,12 +192,14 @@ def create_lca_results_array(
     :rtype: xr.DataArray
     """
 
+
     # Define the coordinates for the xarray DataArray
     coords = {
         "act_category": list(set(classifications.values())),
         "variable": list(mapping.keys()),
         "year": years,
         "region": regions,
+        "location": locations,
         "model": models,
         "scenario": scenarios,
     }
@@ -204,6 +209,7 @@ def create_lca_results_array(
         len(coords["variable"]),
         len(years),
         len(regions),
+        len(locations),
         len(models),
         len(scenarios),
     )