From b13b449d589ed6cf20d29c6fb2db21e818733186 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20C=2E=20Riven=C3=A6s?= <jan@rivenaes.net>
Date: Tue, 19 Sep 2023 18:16:45 +0200
Subject: [PATCH] WIP

---
 src/xtgeo/surface/_regsurf_cube_window_v3.py |  8 +--
 src/xtgeo/surface/regular_surface.py         | 57 +++++++++++++-------
 2 files changed, 41 insertions(+), 24 deletions(-)

diff --git a/src/xtgeo/surface/_regsurf_cube_window_v3.py b/src/xtgeo/surface/_regsurf_cube_window_v3.py
index 67748d2dc..1ed42c900 100644
--- a/src/xtgeo/surface/_regsurf_cube_window_v3.py
+++ b/src/xtgeo/surface/_regsurf_cube_window_v3.py
@@ -192,14 +192,14 @@ def _compute_stats(cref, attr, self, upper):
 
     if attr == "mean":
         values = np.nanmean(cref, axis=2)
+    elif attr == "var":
+        values = np.nanvar(cref, axis=2)
     elif attr == "max":
         values = np.nanmax(cref, axis=2)
     elif attr == "min":
         values = np.nanmin(cref, axis=2)
     elif attr == "rms":
         values = np.sqrt(np.nanmean(np.square(cref), axis=2))
-    elif attr == "var":
-        values = np.var(cref, axis=2)
     elif attr == "maxneg":
         use = cref.copy()
         use[cref >= 0] = np.nan
@@ -228,12 +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)
+        values = np.ma.masked_greater_equal(values, 0.0)  # to make undefined map areas
     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)
+        values = np.ma.masked_less_equal(values, 0.0)  # to make undefined map areas
     elif attr == "sumabs":
         use = cref.copy()
         use = np.abs(use)
diff --git a/src/xtgeo/surface/regular_surface.py b/src/xtgeo/surface/regular_surface.py
index 2348772fc..3b18e8518 100644
--- a/src/xtgeo/surface/regular_surface.py
+++ b/src/xtgeo/surface/regular_surface.py
@@ -32,6 +32,7 @@
 # --------------------------------------------------------------------------------------
 
 # pylint: disable=too-many-public-methods
+from __future__ import annotations
 
 import functools
 import io
@@ -42,7 +43,7 @@
 from collections import OrderedDict
 from copy import deepcopy
 from types import FunctionType
-from typing import List, Optional, Tuple, Type, Union
+from typing import Dict, List, Literal, Optional, Tuple, Type, Union
 
 import deprecation
 import numpy as np
@@ -69,6 +70,22 @@
 xtg = xtgeo.common.XTGeoDialog()
 logger = xtg.functionlogger(__name__)
 
+# valid argumentts for seismic attributes
+ValidAttrs = Literal[
+    "all",
+    "max",
+    "min",
+    "rms",
+    "mean",
+    "var",
+    "maxpos",
+    "maxneg",
+    "sumpos",
+    "sumneg",
+    "meanabs",
+    "meanpos",
+    "meanneg",
+]
 
 # ======================================================================================
 # METHODS as wrappers to class init + import
@@ -2541,21 +2558,21 @@ def slice_cube(
 
     def slice_cube_window(
         self,
-        cube,
-        zsurf=None,
-        other=None,
-        other_position="below",
-        sampling="nearest",
-        mask=True,
-        zrange=None,
-        ndiv=None,
-        attribute="max",
-        maskthreshold=0.1,
-        snapxy=False,
-        showprogress=False,
-        deadtraces=True,
-        algorithm=2,
-    ):
+        cube: xtgeo.Cube,
+        zsurf: Optional[xtgeo.RegularSurface] = None,
+        other: Optional[xtgeo.RegularSurface] = None,
+        other_position: str = "below",
+        sampling: Literal["nearest", "cube", "trilinear"] = "nearest",
+        mask: bool = True,
+        zrange: Optional[float] = None,
+        ndiv: Optional[int] = None,
+        attribute: Union[List[ValidAttrs], ValidAttrs] = "max",
+        maskthreshold: float = 0.1,
+        snapxy: bool = False,
+        showprogress: bool = False,
+        deadtraces: bool = True,
+        algorithm: Literal[1, 2, 3] = 2,
+    ) -> Tuple(None, Dict[xtgeo.RegularSurface]):
         """Slice the cube within a vertical window and get the statistical attrubutes.
 
         The statistical attributes can be min, max etc. Attributes are:
@@ -2590,13 +2607,13 @@ def slice_cube_window(
         available.
 
         Args:
-            cube (Cube): Instance of a Cube()
-            zsurf (RegularSurface): Instance of a depth (or time) map, which
+            cube: Instance of a Cube() here
+            zsurf: Instance of a depth (or time) map, which
                 is the depth or time map (or...) that is used a slicer.
                 If None, then the surface instance itself is used a slice
                 criteria. Note that zsurf must have same map defs as the
                 surface instance.
-            other (RegularSurface): Instance of other surface if window is
+            other: Instance of other surface if window is
                 between surfaces instead of a static window. The zrange
                 input is then not applied.
             sampling (str): 'nearest'/'trilinear'/'cube' for nearest node (default),
@@ -2613,7 +2630,7 @@ def slice_cube_window(
                 means 'auto' sampling, using 0.5 of cube Z increment as basis. If
                 algorithm = 2 and sampling is 'cube', the cube Z increment
                 will be used.
-            attribute (str or list): The requested attribute(s), e.g.
+            attribute: The requested attribute(s), e.g.
                 'max' value. May also be a list of attributes, e.g.
                 ['min', 'rms', 'max']. By such, a dict of surface objects is
                 returned. Note 'all' will make a list of possible attributes