From e8970433247f04050344f54397e4acbc57d5ad1d Mon Sep 17 00:00:00 2001
From: Runar Ask Johannessen <89020325+equinor-ruaj@users.noreply.github.com>
Date: Fri, 9 Jun 2023 13:33:43 +0200
Subject: [PATCH] add is_observation and is_prediction to metadata (#178)

* add is_observation and is_prediction to metadata

* do not expose isobservation and isprediction as explicit properties
---
 examples/explorer.ipynb                            |  2 ++
 src/fmu/sumo/explorer/objects/_child.py            |  4 ++--
 src/fmu/sumo/explorer/objects/_child_collection.py | 11 ++++++++---
 src/fmu/sumo/explorer/objects/cube.py              |  2 --
 src/fmu/sumo/explorer/objects/cube_collection.py   |  6 ++++++
 5 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/examples/explorer.ipynb b/examples/explorer.ipynb
index 85e9970e..55368b51 100644
--- a/examples/explorer.ipynb
+++ b/examples/explorer.ipynb
@@ -190,6 +190,8 @@
     "print(\"Tagname:\", cube.tagname)\n",
     "print(\"Iteration:\", cube.iteration)\n",
     "print(\"Realization:\", cube.realization)\n",
+    "print(\"is_observation:\", cube[\"data\"][\"is_observation\"])\n",
+    "print(\"is_prediction:\", cube[\"data\"][\"is_prediction\"])\n",
     "\n",
     "print(\"\\nOpenVDS data:\")\n",
     "openvds_handle = cube.openvds_handle\n",
diff --git a/src/fmu/sumo/explorer/objects/_child.py b/src/fmu/sumo/explorer/objects/_child.py
index 5c64b0b2..1645558b 100644
--- a/src/fmu/sumo/explorer/objects/_child.py
+++ b/src/fmu/sumo/explorer/objects/_child.py
@@ -27,12 +27,12 @@ def name(self) -> str:
     def tagname(self) -> str:
         """Object tagname"""
         return self._get_property(["data", "tagname"])
-    
+
     @property
     def stratigraphic(self) -> str:
         """Object stratigraphic"""
         return self._get_property(["data", "stratigraphic"])
-    
+
     @property
     def vertical_domain(self) -> str:
         """Object vertical_domain"""
diff --git a/src/fmu/sumo/explorer/objects/_child_collection.py b/src/fmu/sumo/explorer/objects/_child_collection.py
index fcb4925c..fa0a6b6f 100644
--- a/src/fmu/sumo/explorer/objects/_child_collection.py
+++ b/src/fmu/sumo/explorer/objects/_child_collection.py
@@ -25,6 +25,8 @@
     "access.asset",
     "masterdata.smda.field",
     "file.relative_path",
+    "data.is_observation",
+    "data.is_prediction"
 ]
 
 
@@ -71,18 +73,17 @@ def aggregations(self) -> List[str]:
     def stages(self) -> List[str]:
         """List of unique stages"""
         return self._get_field_values("fmu.context.stage.keyword")
-    
+
     @property
     def stratigraphic(self) -> List[str]:
         """List of unqiue object stratigraphic"""
         return self._get_field_values("data.stratigraphic")
-    
+
     @property
     def vertical_domain(self) -> List[str]:
         """List of unqiue object vertical domain"""
         return self._get_field_values("data.vertical_domain")
 
-
     def _init_query(self, doc_type: str, query: Dict = None) -> Dict:
         new_query = super()._init_query(doc_type, query)
         case_filter = {
@@ -108,6 +109,8 @@ def _add_filter(
         uuid: Union[str, List[str], bool] = None,
         stratigraphic: Union[str, List[str], bool] = None,
         vertical_domain: Union[str, List[str], bool] = None,
+        is_observation: bool = None,
+        is_prediction: bool = None
     ):
         must = []
         must_not = []
@@ -123,6 +126,8 @@ def _add_filter(
             "_id": uuid,
             "data.stratigraphic.keyword": stratigraphic,
             "data.vertical_domain.keyword": vertical_domain,
+            "data.is_observation.keyword": is_observation,
+            "data.is_prediction.keyword": is_prediction
         }
 
         for prop, value in prop_map.items():
diff --git a/src/fmu/sumo/explorer/objects/cube.py b/src/fmu/sumo/explorer/objects/cube.py
index e10ed051..fe699994 100644
--- a/src/fmu/sumo/explorer/objects/cube.py
+++ b/src/fmu/sumo/explorer/objects/cube.py
@@ -1,5 +1,4 @@
 """Module containing class for cube object"""
-import json
 import openvds
 from typing import Dict
 from sumo.wrapper import SumoClient
@@ -62,4 +61,3 @@ def interval(self) -> str:
             return (t0, t1)
 
         return None
-
diff --git a/src/fmu/sumo/explorer/objects/cube_collection.py b/src/fmu/sumo/explorer/objects/cube_collection.py
index aa521b76..c5eabf14 100644
--- a/src/fmu/sumo/explorer/objects/cube_collection.py
+++ b/src/fmu/sumo/explorer/objects/cube_collection.py
@@ -87,6 +87,8 @@ def filter(
         stage: Union[str, List[str], bool] = None,
         time: TimeFilter = None,
         uuid: Union[str, List[str], bool] = None,
+        is_observation: bool = None,
+        is_prediction: bool = None
     ) -> "CubeCollection":
         """Filter cubes
 
@@ -98,6 +100,8 @@ def filter(
             stage (Union[str, List[str], bool]): context/stage
             time (TimeFilter): time filter
             uuid (Union[str, List[str], bool]): cube object uuid
+            is_observation (bool): cube is_observation
+            is_prediction (bool): cube is_prediction
 
         Returns:
             CubeCollection: A filtered CubeCollection
@@ -110,6 +114,8 @@ def filter(
             stage=stage,
             time=time,
             uuid=uuid,
+            is_observation=is_observation,
+            is_prediction=is_prediction
         )
 
         return CubeCollection(self._sumo, self._case_uuid, query, self._pit)