review: use int for realization map

equinor · Jan 14, 2025 · b784231 · b784231
1 parent 4ebd579
commit b784231
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/src/everest/api/everest_data_api.py b/src/everest/api/everest_data_api.py
@@ -226,7 +226,7 @@ def summary_values(self, batches=None, keys=None):
                 )
                 realizations = pl.Series(
                     "realization",
-                    [realization_map.get(str(sim)) for sim in summary["simulation"]],
+                    [realization_map.get(int(sim)) for sim in summary["simulation"]],
                 )
                 realizations = realizations.cast(pl.Int64, strict=False)
                 summary = summary.with_columns(realizations)

diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py
@@ -73,7 +73,7 @@ class EverestStorageDataFrames:
     realization_weights: polars.DataFrame | None = None
 
     @property
-    def simulation_to_geo_realization_map(self) -> dict[str, str]:
+    def simulation_to_geo_realization_map(self) -> dict[int, int]:
         """
         Mapping from simulation ID to geo-realization
         """
@@ -91,8 +91,7 @@ def simulation_to_geo_realization_map(self) -> dict[str, str]:
 
         mapping = {}
         for d in dummy_df.select("realization", "simulation_id").to_dicts():
-            # Currently we work with str, but should maybe not be done in future
-            mapping[str(d["simulation_id"])] = str(d["realization"])
+            mapping[int(d["simulation_id"])] = int(d["realization"])
 
         return mapping