From eb3f2a88507fd8673408c70fba740d0fa28d5616 Mon Sep 17 00:00:00 2001
From: Jakub Both <jakub.both@uib.no>
Date: Fri, 22 Nov 2024 12:28:38 +0100
Subject: [PATCH] MAINT: Convert print to logging, and add more structure to
 homogeneous vs heteroegneous analysis.

---
 src/darsia/restoration/averaging.py | 59 ++++++++++++++++++++++++++---
 1 file changed, 53 insertions(+), 6 deletions(-)

diff --git a/src/darsia/restoration/averaging.py b/src/darsia/restoration/averaging.py
index eaa4d881..ea288d5f 100644
--- a/src/darsia/restoration/averaging.py
+++ b/src/darsia/restoration/averaging.py
@@ -1,12 +1,15 @@
 """Module with tools for volume averaging."""
 
-from typing import Union, overload
+import logging
+from typing import Optional, Union, overload
 
 import numpy as np
 import scipy.ndimage
 
 import darsia
 
+logger = logging.getLogger(__name__)
+
 
 class REV:
     """Class defining an representative elementary volume."""
@@ -32,20 +35,25 @@ def __init__(self, size: Union[float, tuple[float]], img: darsia.Image) -> None:
 
 class VolumeAveraging:
 
-    def __init__(self, rev: REV, mask: darsia.Image) -> None:
+    def __init__(
+        self, rev: REV, mask: darsia.Image, labels: Optional[darsia.Image] = None
+    ) -> None:
         """Constructor.
 
         Args:
             rev (REV): representative elementary volume
             mask (Image): mask
+            labels (Image): labels; activating heterogeneous analysis, if not None
 
         """
         self.rev_size = rev.size
         """Size of the REV."""
         self.mask = mask
         """Mask."""
-        self.mean_pore_volume = scipy.ndimage.uniform_filter(
-            self.mask.astype(float).img, size=self.rev_size
+        self.labels = labels
+        """Labels."""
+        self.mean_pore_volume = self._heterogeneous_uniform_filter(
+            self.mask.astype(float).img
         )
         """Mean pore volume."""
         tol = 1e-12
@@ -53,7 +61,9 @@ def __init__(self, rev: REV, mask: darsia.Image) -> None:
         """Zero indices in the mean pore volume."""
 
         # User output
-        print(f"Number of zero indices: {len(self.zero_indices[0])}")
+        logger.info(
+            f"Number of zero indices in pore volume identification: {len(self.zero_indices[0])}"
+        )
 
     @overload  # type: ignore [override]
     def __call__(self, img: np.ndarray) -> np.ndarray: ...
@@ -83,19 +93,56 @@ def __call__(
     def _average_array(self, arr: np.ndarray) -> np.ndarray:
         """Application of volume averaging to numpy array.
 
+        Apply averaging to each channel of a 3D array.
+
         Args:
             arr (np.ndarray): array
 
         Returns:
             np.ndarray: volume averaged array
 
+        """
+        if arr.ndim == 2:
+            return self._average_array_single(arr)
+        elif arr.ndim == 3:
+            return np.stack(
+                [self._average_array_single(arr[..., i]) for i in range(arr.shape[-1])],
+                axis=-1,
+            )
+        else:
+            raise ValueError("Only 2D and 3D arrays are supported.")
+
+    def _average_array_single(self, arr: np.ndarray) -> np.ndarray:
+        """Application of volume averaging to two-dimensional numpy array.
+
+        Args:
+            arr (np.ndarray): two-dimensional array
+
+        Returns:
+            np.ndarray: two-dimensional volume averaged array
+
         """
         masked_data = np.multiply(arr, self.mask.img)
-        mean_masked_data = scipy.ndimage.uniform_filter(masked_data, size=self.rev_size)
+        mean_masked_data = self._heterogeneous_uniform_filter(masked_data)
         result = np.divide(mean_masked_data, self.mean_pore_volume)
         result[self.zero_indices] = 0
         return result
 
+    def _heterogeneous_uniform_filter(self, data: np.ndarray) -> np.ndarray:
+        """Application of a uniform filter to heterogeneous data.
+
+        Args:
+            data (np.ndarray): data
+
+        Returns:
+            np.ndarray: filtered data
+
+        """
+        if self.labels is None:
+            return scipy.ndimage.uniform_filter(data, size=self.rev_size)
+        else:
+            raise NotImplementedError("Heterogeneous analysis is not yet implemented.")
+
 
 def volume_average(img: darsia.Image, mask: darsia.Image, size: float) -> darsia.Image:
     """Fast-access function for volume averaging.