Ensure regions and zones are always strings in dataframe (#143)

equinor · Jun 21, 2021 · 6cba81b · 6cba81b
1 parent d151e7e
commit 6cba81b
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 6 deletions.
diff --git a/src/fmu/tools/rms/volumetrics.py b/src/fmu/tools/rms/volumetrics.py
@@ -140,14 +140,17 @@ def myregionrenamer(s):
     if zonerenamer:
         vol_df["ZONE"] = vol_df["ZONE"].apply(zonerenamer)
 
-    # Remove the Totals rows in case they are present.
-    #  (todo: do this for all columns that are not not of numeric type)
-    checkfortotals = ["ZONE", "REGION", "LICENSE", "FACIES"]
+    index_columns = ["ZONE", "REGION", "LICENSE", "FACIES"]
+    present_index_columns = list(set(index_columns).intersection(vol_df.columns))
 
+    # Index columns should always be of string datatype:
+    vol_df[present_index_columns] = vol_df[present_index_columns].astype(str)
+
+    # Remove the Totals rows in case they are present, signified by the
+    # magic value "Totals" in any of the index columns:
     totalsrows = pd.Series([False] * len(vol_df))
-    for col in checkfortotals:
-        if col in vol_df.columns:
-            totalsrows = totalsrows | (vol_df[col] == "Totals")
+    for col in present_index_columns:
+        totalsrows = totalsrows | (vol_df[col] == "Totals")
     vol_df = vol_df[~totalsrows].reset_index(drop=True)
 
     if outfile:

diff --git a/tests/rms/test_rmsvolumetrics.py b/tests/rms/test_rmsvolumetrics.py
@@ -43,6 +43,25 @@
             # This "fails" as there was not a double space between Zone and Bulk
             pd.DataFrame([{"Zone Bulk": "Upper 1"}]),
         ),
+        (
+            # Integer zone, returned as string datatype
+            "Zone  Bulk\n1  1.0",
+            "oil.txt",
+            pd.DataFrame([{"ZONE": "1", "BULK_OIL": 1.0}]),
+        ),
+        (
+            # Integer Region, returned as string datatype
+            "Region index  Bulk\n1  1.0",
+            "oil.txt",
+            pd.DataFrame([{"REGION": "1", "BULK_OIL": 1.0}]),
+        ),
+        (
+            # Floating point Region, returned as string datatype.
+            # User is probably up for trouble if this is used.
+            "Region index  Bulk\n1.0  1.0",
+            "oil.txt",
+            pd.DataFrame([{"REGION": "1.0", "BULK_OIL": 1.0}]),
+        ),
         (
             # Two spaces:
             "Zone  Bulk\nUpper  1.0",