From f6d07e0b61e1d9871283b1578018b93c8f57f221 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20C=2E=20Riven=C3=A6s?= <jan@rivenaes.net>
Date: Thu, 14 Sep 2023 18:29:07 +0200
Subject: [PATCH] WIP, seism atttrs v3

---
 src/xtgeo/surface/_regsurf_cube_window_v3.py | 37 +++++++++++++++-----
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/src/xtgeo/surface/_regsurf_cube_window_v3.py b/src/xtgeo/surface/_regsurf_cube_window_v3.py
index 7bc6b48c0..67748d2dc 100644
--- a/src/xtgeo/surface/_regsurf_cube_window_v3.py
+++ b/src/xtgeo/surface/_regsurf_cube_window_v3.py
@@ -3,7 +3,8 @@
 
 from __future__ import annotations
 
-from typing import Type
+import warnings
+from typing import Type, Union
 
 import numpy as np
 from scipy.interpolate import interp1d
@@ -11,6 +12,9 @@
 import xtgeo
 from xtgeo.common import XTGeoDialog
 
+warnings.filterwarnings(action="ignore", message="All-NaN slice encountered")
+warnings.filterwarnings(action="ignore", message="Mean of empty slice")
+
 xtg = XTGeoDialog()
 
 logger = xtg.functionlogger(__name__)
@@ -41,7 +45,7 @@
 # pylint: disable=W0613, W0612
 
 
-def _cut_cube_deadtraces(cube: Type[xtgeo.Cube]) -> Type[np.ndarray]:
+def _cut_cube_deadtraces(cube: xtgeo.Cube) -> np.ndarray:
     """Take the cube numpy values and filter away dead traces as np.nan."""
 
     logger.info("Assign dead traces")
@@ -58,10 +62,10 @@ def _cut_cube_deadtraces(cube: Type[xtgeo.Cube]) -> Type[np.ndarray]:
 def _upper_lower_surface(
     self,
     cube,
-    zsurf: Type[xtgeo.RegularSurface],
-    other: Type[xtgeo.RegularSurface],
-    other_position,
-    zrange,
+    zsurf: xtgeo.RegularSurface,
+    other: xtgeo.RegularSurface,
+    other_position: str,
+    zrange: float,
 ) -> list:
     """Return upper and lower surface, sampled to cube resolution."""
 
@@ -95,7 +99,7 @@ def _upper_lower_surface(
     return upper, lower
 
 
-def _create_depth_cube(cube):
+def _create_depth_cube(cube: xtgeo.Cube) -> np.ndarray:
     """Create a cube (np array) where values are cube depths; to be used as filter."""
     logger.info("Create a depth cube...")
     dcube = cube.values.copy()
@@ -103,6 +107,7 @@ def _create_depth_cube(cube):
     dcube[:, :, :] = darr
 
     logger.info("Created a depth cube starting from %s", np.mean(dcube))
+
     return dcube
 
 
@@ -162,7 +167,7 @@ def _filter_cube_values_upper_lower(cvalues, dvalues, upper, lower):
     return cvalues
 
 
-def _expand_attributes(attribute):
+def _expand_attributes(attribute: Union[str, list]) -> list:
     """The 'attribute' may be a name, 'all', or a list of attributes"""
     useattrs = None
     if isinstance(attribute, str):
@@ -173,6 +178,11 @@ def _expand_attributes(attribute):
     else:
         useattrs = attribute
 
+    if not all(item in ALLATTRS for item in useattrs):
+        raise ValueError(
+            f"One or more values are not a valid, input list is {useattrs}, "
+            f"allowed list is {ALLATTRS}"
+        )
     return useattrs
 
 
@@ -218,10 +228,12 @@ def _compute_stats(cref, attr, self, upper):
         use = cref.copy()
         use[cref >= 0] = np.nan
         values = np.nansum(use, axis=2)
+        values = np.ma.masked_greater_equal(values, 0.0)
     elif attr == "sumpos":
         use = cref.copy()
         use[cref < 0] = np.nan
         values = np.nansum(use, axis=2)
+        values = np.ma.masked_less_equal(values, 0.0)
     elif attr == "sumabs":
         use = cref.copy()
         use = np.abs(use)
@@ -231,7 +243,8 @@ def _compute_stats(cref, attr, self, upper):
 
     actual = self.copy()
     sampled = upper.copy()
-    sampled.values = values
+
+    sampled.values = np.ma.masked_invalid(values)
     actual.resample(sampled)
 
     logger.info("Compute stats... done")
@@ -267,7 +280,13 @@ def slice_cube_window(
     cref, dref = _refine_cubes_vertically(cvalues, dvalues, ndiv)
 
     cref = _filter_cube_values_upper_lower(cref, dref, upper, lower)
+
     cval = _filter_cube_values_upper_lower(cvalues, dvalues, upper, lower)
+    print(cval)
+    c = cube.copy()
+    c.values = np.nan_to_num(cval, nan=0)
+
+    c.to_file("/tmp/cval.segy")
 
     use_attrs = _expand_attributes(attribute)