Fixes #61, bug fix and change of fill_negative functions

Fixed buggy fill negatives function and replaced it by two new methods,
either filling from a maximum closeby channel, or from the weighted maximum. where the weights are given by a gaussian function

```
Fill negative channels with positive counts from neighboring channels.

The idea is that the negative counts are somehow connected to the (γ-ray)
resolution and should thus be filled from a channel within the resolution.

This implementation loops through the closest channels with positive number
of counts to fill the channle(s) with negative counts. Note that it can
happen that some bins will remain with negative counts (if not enough bins
with possitive counts are available within the window_size) .

The routine is performed for each Ex row independently.
```

ompy/__init__.py

@@ -25,7 +25,7 @@
 
     # Simply import all functions and classes from all files to make them
     # available at the package level
-    from .library import div0, fill_negative
+    from .library import div0, fill_negative_gauss, fill_negative_max
     from .spinfunctions import SpinFunctions
     from .abstractarray import AbstractArray
     from .matrix import Matrix

ompy/firstgeneration.py

@@ -32,7 +32,7 @@
 import logging
 import termtables as tt
 import numpy as np
-from typing import Tuple, Optional
+from typing import Tuple, Optional, Callable
 from .matrix import Matrix
 from .vector import Vector
 from .library import div0
@@ -56,7 +56,8 @@ class FirstGeneration:
         multiplicity_estimation (str): Selects which method should be used
             for the multiplicity estimation. Can be either "statistical",
             or "total". Default is "statistical".
-
+        window_size (int or np.ndarray): window_size for (fill and) remove
+            negatives on output. Defaults to 20.
         statistical_upper (float): Threshold for upper limit in
             `statistical` multiplicity estimation. Defaults to 430 keV.
         statistical_lower (float): Threshold for lower limit in
@@ -86,6 +87,7 @@ class FirstGeneration:
     def __init__(self):
         self.num_iterations: int = 10
         self.multiplicity_estimation: str = 'statistical'
+        self.window_size = 20
 
         self.statistical_upper: float = 430.0  # MAMA ThresSta
         self.statistical_lower: float = 200.0  # MAMA ThresTot
@@ -136,8 +138,7 @@ def apply(self, unfolded: Matrix) -> Matrix:
 
         final = Matrix(values=H, Eg=matrix.Eg, Ex=matrix.Ex)
         final.state = "firstgen"
-        final.fill_negative(window_size=10)
-        final.remove_negative()
+        self.remove_negative(final)
         return final
 
     def setup(self, matrix: Matrix) -> Tuple[Matrix, Matrix, Matrix]:
@@ -399,6 +400,16 @@ def allgen_from_primary(fg: Matrix,
 
         return ag
 
+    def remove_negative(self, matrix: Matrix):
+        """ (Fill and) remove negative counts
+
+        Wrapper for Matrix.fill_and_remove_negative()
+
+        Args:
+            matrix: Input matrix
+        """
+        matrix.fill_and_remove_negative(window_size=self.window_size)
+
 
 def normalize_rows(array: np.ndarray) -> np.ndarray:
     """ Normalize each row to unity """

ompy/library.py

@@ -27,9 +27,10 @@
 """
 import numpy as np
 from scipy.interpolate import interp1d, RectBivariateSpline
+from scipy.stats import truncnorm
 import matplotlib
 from itertools import product
-from typing import Optional, Tuple, Iterator, Any
+from typing import Optional, Tuple, Iterator, Any, Union
 import inspect
 import re
 
@@ -104,39 +105,141 @@ def E_array_from_calibration(a0: float, a1: float, *,
         raise ValueError("Either N or E_max must be given")
 
 
-def fill_negative(matrix, window_size: int):
+def fill_negative_max(array: np.ndarray,
+                      window_size: Union[int, np.array]) -> np.ndarray:
     """
-    Fill negative channels with positive counts from neighbouring channels
+    Fill negative channels with positive counts from neighboring channels.
 
-    The MAMA routine for this is very complicated. It seems to basically
-    use a sliding window along the Eg axis, given by the FWHM, to look for
-    neighbouring bins with lots of counts and then take some counts from there.
-    Can we do something similar in an easy way?
+    The idea is that the negative counts are somehow connected to the (γ-ray)
+    resolution and should thus be filled from a channel within the resolution.
+
+    This implementation loops through the closest channels with maximum number
+    of counts to fill the channle(s) with negative counts. Note that it can
+    happen that some bins will remain with negative counts (if not enough bins
+    with possitive counts are available within the window_size) .
+
+    The routine is performed for each Ex row independently.
+
+    Args:
+        array: Input array, ordered as [Ex, Eg]
+        window_size (Union[int, np.array]): Window_size, eg. FWHM. If `int`
+            `float`, the same FWHM will be applied for all `Eg` bins.
+            Otherwise, provide an array with the FWHM for each `Eg` bin.
+
+    Returns:
+        array with negative counts filled, where possible
     """
-    matrix_out = np.copy(matrix)
-    # Loop over rows:
-    for i_Ex in range(matrix.shape[0]):
-        for i_Eg in np.where(matrix[i_Ex, :] < 0)[0]:
-            # print("i_Ex = ", i_Ex, "i_Eg =", i_Eg)
-            # window_size = 4  # Start with a constant window size.
-            # TODO relate it to FWHM by energy arrays
-            i_Eg_low = max(0, i_Eg - window_size)
-            i_Eg_high = min(matrix.shape[1], i_Eg + window_size)
-            # Fill from the channel with the larges positive count
-            # in the neighbourhood
-            i_max = np.argmax(matrix[i_Ex, i_Eg_low:i_Eg_high])
-            # print("i_max =", i_max)
-            if matrix[i_Ex, i_max] <= 0:
-                pass
-            else:
-                positive = matrix[i_Ex, i_max]
-                negative = matrix[i_Ex, i_Eg]
-                fill = min(0, positive + negative)  # Don't fill more than to 0
-                rest = positive
-                # print("fill =", fill, "rest =", rest)
-                matrix_out[i_Ex, i_Eg] = fill
-                # matrix_out[i_Ex, i_max] = rest
-    return matrix_out
+    if isinstance(window_size, int):
+        window_size = np.full(array.shape[1], window_size)
+    else:
+        assert len(window_size) == array.shape[1], "Array length incompatible"
+        assert window_size.dtype == np.integer, "Check input"
+
+    array = np.copy(array)
+    N_Ex = array.shape[0]
+    N_Eg = array.shape[1]
+    for i_Ex in range(N_Ex):
+        row = array[i_Ex, :]
+        for i_Eg in np.where(row < 0)[0]:
+            window_size_Eg = window_size[i_Eg]
+            max_distance = window_size_Eg
+            max_distance = int(np.ceil((window_size_Eg - 1) / 2))
+            i_Eg_low = max(0, i_Eg - max_distance)
+            i_Eg_high = min(N_Eg, i_Eg + max_distance)
+            while row[i_Eg] < 0:
+                i_max = np.argmax(row[i_Eg_low:i_Eg_high + 1])
+                i_max = i_Eg_low + i_max
+                if row[i_max] <= 0:
+                    break
+                shuffle_counts(row, i_max, i_Eg)
+
+    return array
+
+
+def fill_negative_gauss(array: np.ndarray, Eg: np.ndarray,
+                        window_size: Union[int, float, np.array],
+                        n_trunc: float = 3) -> np.ndarray:
+    """
+    Fill negative channels with positive counts from weighted neighbor chnls.
+
+    The idea is that the negative counts are somehow connected to the (γ-ray)
+    resolution and should thus be filled from a channel within the resolution.
+
+    This implementation loops through channels with the maximum "weight", where
+    the weight is given by
+        weight = gaussian(i, loc=i, scale ~ window_size) * counts(i),
+    to fill the channle(s) with negative counts. Note that it can
+    happen that some bins will remain with negative counts (if not enough bins
+    with possitive counts are available within the window_size) .
+
+    The routine is performed for each Ex row independently.
+
+    Args:
+        array: Input array, ordered as [Ex, Eg]
+        Eg: Gamma-ray energies
+        window_size: FWHM for the gaussian. If `int` or
+            `float`, the same FWHM will be applied for all `Eg` bins.
+            Otherwise, provide an array with the FWHM for each `Eg` bin.
+        n_trun (float, optional): Truncate gaussian for faster calculation.
+            Defaults to 3.
+
+    Returns:
+        array with negative counts filled, where possible
+    """
+    if isinstance(window_size, (int, float)):
+        sigma = window_size/2.355  # convert FWHM to sigma
+        window_size = np.full(array.shape[1], window_size)
+    else:
+        assert len(window_size) == array.shape[1], "Array length incompatible"
+        sigma = window_size/2.355  # convert FWHM to sigma
+
+    # generate truncated gauss for each Eg bin, format [Eg-bin, gauss-values]
+    lower, upper = Eg - n_trunc*sigma, Eg + n_trunc*sigma
+    a = (lower - Eg) / sigma
+    b = (upper - Eg) / sigma
+    gauss = [truncnorm(a=a, b=b, loc=p, scale=sig).pdf(Eg)
+             for p, sig in zip(Eg, sigma)]
+    gauss = np.array(gauss)
+
+    array = np.copy(array)
+    N_Ex = array.shape[0]
+    for i_Ex in range(N_Ex):
+        row = array[i_Ex, :]
+        for i_Eg in np.nonzero(row < 0)[0]:
+            positives = np.where(row < 0, 0, row)
+            weights = positives * gauss[i_Eg, :]
+
+            for i_from in np.argsort(weights):
+                if row[i_from] < 0:
+                    break
+                shuffle_counts(row, i_from, i_Eg)
+                if row[i_Eg] >= 0:
+                    break
+
+    return array
+
+
+def shuffle_counts(row: np.ndarray, i_from: int, i_to: int):
+    """Shuffles counts in `row` from bin `i_from` to `i_to`
+
+    Transfers at maximum row[i_from] counts, so that row[i_from] cannot be
+    negative after the shuffling.
+
+    Note:
+        Assumes that row[i_from] > 0 and row[i_to] < 0.
+
+    Args:
+        row: Input array
+        i_from: Index of bin to take counts from
+        i_to: Index of bin to fill
+
+    """
+    positive = row[i_from]
+    negative = row[i_to]
+    fill = min(0, positive + negative)
+    rest = max(0, positive + negative)
+    row[i_to] = fill
+    row[i_from] = rest
 
 
 def cut_diagonal(matrix, Ex_array, Eg_array, E1, E2):

ompy/matrix.py

@@ -42,7 +42,7 @@
 from .filehandling import (mama_read, mama_write,
                            save_numpy_2D, load_numpy_2D, save_tar, load_tar,
                            filetype_from_suffix)
-from .library import div0, fill_negative, diagonal_elements
+from .library import div0, fill_negative_gauss, diagonal_elements
 from .matrixstate import MatrixState
 from .rebin import rebin_2D
 from .vector import Vector
@@ -752,16 +752,22 @@ def diagonal_elements(self) -> Iterator[Tuple[int, int]]:
         return diagonal_elements(self.values)
 
     def fill_negative(self, window_size: int):
-        """ Wrapper for :func:`ompy.fill_negative` """
-        self.values = fill_negative(self.values, window_size)
+        """ Wrapper for :func:`ompy.fill_negative_gauss` """
+        self.values = fill_negative_gauss(self.values, self.Eg, window_size)
 
     def remove_negative(self):
         """ Entries with negative values are set to 0 """
         self.values = np.where(self.values > 0, self.values, 0)
 
-    def fill_and_remove_negative(self, window_size: int):
+    def fill_and_remove_negative(self,
+                                 window_size: Tuple[int, np.ndarray] = 20):
         """ Combination of :meth:`ompy.Matrix.fill_negative` and
-        :meth:`ompy.Matrix.remove_negative` """
+        :meth:`ompy.Matrix.remove_negative`
+
+        Args:
+            window_size: See `fill_negative`. Defaults to 20 (arbitrary)!.
+            """
+
         self.fill_negative(window_size=window_size)
         self.remove_negative()
 

ompy/unfolder.py

@@ -30,7 +30,7 @@
 import pandas
 import logging
 import warnings
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Callable
 import termtables as tt
 from scipy.ndimage import gaussian_filter1d
 from copy import copy
@@ -78,8 +78,8 @@ class Unfolder:
              fluctuations. Defaults to 0.2
         minimum_iterations (int, optional): Minimum number of iterations.
             Defaults to 3.
-        window_size (float or int?, optional): window_size for fill negatives
-            on output. Defaults to 10.
+        window_size (int or np.ndarray): window_size for  (fill and) remove
+            negatives on output. Defaults to 20.
         use_compton_subtraction (bool, optional): Set usage of Compton
             subtraction method. Defaults to `True`.
         response_tab (DataFrame, optional): If `use_compton_subtraction=True`
@@ -91,7 +91,6 @@ class Unfolder:
         iscores (np.ndarray, optional): Selected iteration number in the
             unfolding. The unfolding and selection is performed independently
             for each `Ex` row, thus this is an array with `len(raw.Ex)`.
-
     TODO:
         - Unfolding for a single spectrum (currently needs to be mocked as a
             Matrix).
@@ -106,7 +105,7 @@ def __init__(self, num_iter: int = 33, response: Optional[Matrix] = None):
         self.num_iter = num_iter
         self.weight_fluctuation: float = 0.2
         self.minimum_iterations: int = 3
-        self.window_size = 10
+        self.window_size = 20
 
         self._R: Optional[Matrix] = response
         self.raw: Optional[Matrix] = None
@@ -116,9 +115,6 @@ def __init__(self, num_iter: int = 33, response: Optional[Matrix] = None):
         self.response_tab: Optional[pandas.DataFrame] = None
         self.FWHM_tweak_multiplier: Optional[dict[str, float]] = None
 
-        # remove negs. in final unfolded spectrum
-        self.remove_negatives: bool = True
-
         self.iscores: Optional[np.ndarray] = None
 
     def __call__(self, matrix: Matrix) -> Matrix:
@@ -217,8 +213,7 @@ def apply(self, raw: Matrix,
         unfolded = Matrix(unfolded, Eg=self.raw.Eg, Ex=self.raw.Ex)
         unfolded.state = "unfolded"
 
-        if self.remove_negatives:
-            unfolded.fill_and_remove_negative(window_size=self.window_size)
+        self.remove_negative(unfolded)
         return unfolded
 
     def step(self, unfolded: np.ndarray, folded: np.ndarray,
@@ -435,6 +430,16 @@ def compton_subtraction(self, unfolded: np.ndarray) -> np.ndarray:
 
         return unfolded
 
+    def remove_negative(self, matrix: Matrix):
+        """ (Fill and) remove negative counts
+
+        Wrapper for Matrix.fill_and_remove_negative()
+
+        Args:
+            matrix: Input matrix
+        """
+        matrix.fill_and_remove_negative(window_size=self.window_size)
+
 
 def shift(counts_in, E_array_in, energy_shift):
     """