ChessAnalysisPipeline 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

CHAP/utils/general.py

@@ -31,6 +31,52 @@ except ImportError:
 
 logger = getLogger(__name__)
 
+tiny = np.finfo(np.float64).resolution
+
+def gformat(val, length=11):
+    """
+    Format a number with '%g'-like format, while:
+      - the length of the output string will be of the requested length
+      - positive numbers will have a leading blank
+      - the precision will be as high as possible
+      - trailing zeros will not be trimmed
+    """
+    # Code taken from lmfit library
+    if val is None or isinstance(val, bool):
+        return f'{repr(val):>{length}s}'
+    try:
+        expon = int(np.log10(abs(val)))
+    except (OverflowError, ValueError):
+        expon = 0
+    except TypeError:
+        return f'{repr(val):>{length}s}'
+
+    length = max(length, 7)
+    form = 'e'
+    prec = length - 7
+    if abs(expon) > 99:
+        prec -= 1
+    elif ((expon > 0 and expon < (prec+4)) or
+          (expon <= 0 and -expon < (prec-1))):
+        form = 'f'
+        prec += 4
+        if expon > 0:
+            prec -= expon
+    return f'{val:{length}.{prec}{form}}'
+
+
+def getfloat_attr(obj, attr, length=11):
+    """Format an attribute of an object for printing."""
+    # Code taken from lmfit library
+    val = getattr(obj, attr, None)
+    if val is None:
+        return 'unknown'
+    if isinstance(val, int):
+        return f'{val}'
+    if isinstance(val, float):
+        return gformat(val, length=length).strip()
+    return repr(val)
+
 
 def depth_list(_list):
     """Return the depth of a list."""
@@ -86,6 +132,13 @@ def illegal_combination(
         raise ValueError(error_msg)
 
 
+def not_zero(value):
+    """Return value with a minimal absolute size of tiny,
+    preserving the sign.
+    """
+    return float(np.copysign(max(tiny, abs(value)), value))
+
+
 def test_ge_gt_le_lt(
         ge, gt, le, lt, func, location=None, raise_error=False, log=True):
     """
@@ -447,8 +500,8 @@ def index_nearest_down(a, value):
     return index
 
 
-def index_nearest_upp(a, value):
-    """Return index of nearest array value, rounded upp."""
+def index_nearest_up(a, value):
+    """Return index of nearest array value, rounded up."""
     a = np.asarray(a)
     if a.ndim > 1:
         raise ValueError(
@@ -477,14 +530,14 @@ def get_consecutive_int_range(a):
 
 
 def round_to_n(x, n=1):
-    """Round to a specific number of decimals."""
+    """Round to a specific number of sig figs ."""
     if x == 0.0:
         return 0
     return type(x)(round(x, n-1-int(np.floor(np.log10(abs(x))))))
 
 
 def round_up_to_n(x, n=1):
-    """Round up to a specific number of decimals."""
+    """Round up to a specific number of sig figs."""
     x_round = round_to_n(x, n)
     if abs(x/x_round) > 1.0:
         x_round += np.sign(x) * 10**(np.floor(np.log10(abs(x)))+1-n)
@@ -492,7 +545,7 @@ def round_up_to_n(x, n=1):
 
 
 def trunc_to_n(x, n=1):
-    """Truncate to a specific number of decimals."""
+    """Truncate to a specific number of sig figs."""
     x_round = round_to_n(x, n)
     if abs(x_round/x) > 1.0:
         x_round -= np.sign(x) * 10**(np.floor(np.log10(abs(x)))+1-n)
@@ -510,7 +563,9 @@ def almost_equal(a, b, sig_figs):
         f'b: {b}, {type(b)})')
 
 
-def string_to_list(s, split_on_dash=True, remove_duplicates=True, sort=True):
+def string_to_list(
+        s, split_on_dash=True, remove_duplicates=True, sort=True,
+        raise_error=False):
     """
     Return a list of numbers by splitting/expanding a string on any
     combination of commas, whitespaces, or dashes (when
@@ -524,8 +579,11 @@ def string_to_list(s, split_on_dash=True, remove_duplicates=True, sort=True):
         return []
     try:
         list1 = re_split(r'\s+,\s+|\s+,|,\s+|\s+|,', s.strip())
-    except (ValueError, TypeError, SyntaxError, MemoryError, RecursionError):
-        return None
+    except (ValueError, TypeError, SyntaxError, MemoryError,
+            RecursionError) as e:
+        if not raise_error:
+            return None
+        raise e
     if split_on_dash:
         try:
             l_of_i = []
@@ -540,8 +598,10 @@ def string_to_list(s, split_on_dash=True, remove_duplicates=True, sort=True):
                 else:
                     raise ValueError
         except (ValueError, TypeError, SyntaxError, MemoryError,
-                RecursionError):
-            return None
+                RecursionError) as e:
+            if not raise_error:
+                return None
+            raise e
     else:
         l_of_i = [literal_eval(x) for x in list1]
     if remove_duplicates:
@@ -551,6 +611,24 @@ def string_to_list(s, split_on_dash=True, remove_duplicates=True, sort=True):
     return l_of_i
 
 
+def list_to_string(a):
+    """Return a list of pairs of integers marking consecutive ranges
+    of integers in string notation."""
+    int_ranges = get_consecutive_int_range(a)
+    if not len(int_ranges):
+        return ''
+    if int_ranges[0][0] == int_ranges[0][1]:
+        s = f'{int_ranges[0][0]}'
+    else:
+        s = f'{int_ranges[0][0]}-{int_ranges[0][1]}'
+    for int_range in int_ranges[1:]:
+        if int_range[0] == int_range[1]:
+            s += f', {int_range[0]}'
+        else:
+            s += f', {int_range[0]}-{int_range[1]}'
+    return s
+
+
 def get_trailing_int(string):
     """Get the trailing integer in a string."""
     index_regex = re_compile(r'\d+$')
@@ -902,11 +980,255 @@ def file_exists_and_readable(f):
     return f
 
 
+def rolling_average(
+        y, x=None, dtype=None, start=0, end=None, width=None,
+        stride=None, num=None, average=True, mode='valid',
+        use_convolve=None):
+    """
+    Returns the rolling sum or average of an array over the last
+    dimension.
+    """
+    y = np.asarray(y)
+    y_shape = y.shape
+    if y.ndim == 1:
+        y = np.expand_dims(y, 0)
+    else:
+        y = y.reshape((np.prod(y.shape[0:-1]), y.shape[-1]))
+    if x is not None:
+        x = np.asarray(x)
+        if x.ndim != 1:
+            raise ValueError('Parameter "x" must be a 1D array-like')
+        if x.size != y.shape[1]:
+            raise ValueError(f'Dimensions of "x" and "y[1]" do not '
+                             f'match ({x.size} vs {y.shape[1]})')
+    if dtype is None:
+        if average:
+            dtype = y.dtype
+        else:
+            dtype = np.float32
+    if width is None and stride is None and num is None:
+        raise ValueError('Invalid input parameters, specify at least one of '
+                         '"width", "stride" or "num"')
+    if width is not None and not is_int(width, ge=1):
+        raise ValueError(f'Invalid "width" parameter ({width})')
+    if stride is not None and not is_int(stride, ge=1):
+        raise ValueError(f'Invalid "stride" parameter ({stride})')
+    if num is not None and not is_int(num, ge=1):
+        raise ValueError(f'Invalid "num" parameter ({num})')
+    if not isinstance(average, bool):
+        raise ValueError(f'Invalid "average" parameter ({average})')
+    if mode not in ('valid', 'full'):
+        raise ValueError(f'Invalid "mode" parameter ({mode})')
+    size = y.shape[1]
+    if size < 2:
+        raise ValueError(f'Invalid y[1] dimension ({size})')
+    if not is_int(start, ge=0, lt=size):
+        raise ValueError(f'Invalid "start" parameter ({start})')
+    if end is None:
+        end = size
+    elif not is_int(end, gt=start, le=size):
+        raise ValueError(f'Invalid "end" parameter ({end})')
+    if use_convolve is None:
+        if len(y_shape) ==1:
+            use_convolve = True
+        else:
+            use_convolve = False
+    if use_convolve and (start or end < size):
+        y = np.take(y, range(start, end), axis=1)
+        if x is not None:
+            x = x[start:end]
+        size = y.shape[1]
+    else:
+        size = end-start
+
+    if stride is None:
+        if width is None:
+            width = max(1, int(size/num))
+            stride = width
+        else:
+            width = min(width, size)
+            if num is None:
+                stride = width
+            else:
+                stride = max(1, int((size-width) / (num-1)))
+    else:
+        stride = min(stride, size-stride)
+        if width is None:
+            width = stride
+
+    if mode == 'valid':
+        num = 1 + max(0, int((size-width) / stride))
+    else:
+        num = int(size/stride)
+        if num*stride < size:
+            num += 1
+
+    if use_convolve:
+        n_start = 0
+        n_end = width
+        weight = np.empty((num))
+        for n in range(num):
+            n_num = n_end-n_start
+            weight[n] = n_num
+            n_start += stride
+            n_end = min(size, n_end+stride)
+
+        window = np.ones((width))
+        if x is not None:
+            if mode == 'valid':
+                rx = np.convolve(x, window)[width-1:1-width:stride]
+            else:
+                rx = np.convolve(x, window)[width-1::stride]
+            rx /= weight
+
+        ry = []
+        if mode == 'valid':
+            for i in range(y.shape[0]):
+                ry.append(np.convolve(y[i], window)[width-1:1-width:stride])
+        else:
+            for i in range(y.shape[0]):
+                ry.append(np.convolve(y[i], window)[width-1::stride])
+        ry = np.reshape(ry, (*y_shape[0:-1], num))
+        if len(y_shape) == 1:
+            ry = np.squeeze(ry)
+        if average:
+            ry = (np.asarray(ry).astype(np.float32)/weight).astype(dtype)
+        elif mode != 'valid':
+            weight = np.where(weight < width, width/weight, 1.0)
+            ry = (np.asarray(ry).astype(np.float32)*weight).astype(dtype)
+    else:
+        ry = np.zeros((num, y.shape[0]), dtype=y.dtype)
+        if x is not None:
+            rx = np.zeros(num, dtype=x.dtype)
+        n_start = start
+        n_end = n_start+width
+        for n in range(num):
+            y_sum = np.sum(y[:,n_start:n_end], 1)
+            n_num = n_end-n_start
+            if n_num < width:
+                y_sum *= width/n_num
+            ry[n] = y_sum
+            if x is not None:
+                rx[n] = np.sum(x[n_start:n_end])/n_num
+            n_start += stride
+            n_end = min(start+size, n_end+stride)
+        ry = np.reshape(ry.T, (*y_shape[0:-1], num))
+        if len(y_shape) == 1:
+            ry = np.squeeze(ry)
+        if average:
+            ry = (ry.astype(np.float32)/width).astype(dtype)
+
+    if x is None:
+        return ry
+    return ry, rx
+
+
+def baseline_arPLS(
+        y, mask=None, w=None, tol=1.e-8, lam=1.e6, max_iter=20,
+        full_output=False):
+    """Returns the smoothed baseline estimate of a spectrum.
+
+    Based on S.-J. Baek, A. Park, Y.-J Ahn, and J. Choo,
+    "Baseline correction using asymmetrically reweighted penalized
+    least squares smoothing", Analyst, 2015,140, 250-257
+
+    :param y: The spectrum.
+    :type y: array-like
+    :param mask: A mask to apply to the spectrum before baseline
+       construction, default to `None`.
+    :type mask: array-like, optional
+    :param w: The weights (allows restart for additional ieterations),
+        defaults to None.
+    :type w: numpy.array, optional
+    :param tol: The convergence tolerence, defaults to `1.e-8`.
+    :type tol: float, optional
+    :param lam: The &lambda (smoothness) parameter (the balance
+        between the residual of the data and the baseline and the
+        smoothness of the baseline). The suggested range is between
+        100 and 10^8, defaults to `10^6`.
+    :type lam: float, optional
+    :param max_iter: The maximum number of iterations,
+        defaults to `20`.
+    :type max_iter: int, optional
+    :param full_output: Whether or not to also output the baseline
+        corrected spectrum, the number of iterations and error in the
+        returned result, defaults to `False`.
+    :type full_output: bool, optional
+    :return: The smoothed baseline, with optionally the baseline
+        corrected spectrum, the weights, the number of iterations and
+        the error in the returned result.
+    :rtype: numpy.array [, numpy.array, int, float]
+    """
+    # With credit to: Daniel Casas-Orozco
+    # https://stackoverflow.com/questions/29156532/python-baseline-correction-library
+    # System modules
+    from sys import float_info
+
+    # Third party modules
+    from scipy.sparse import (
+        spdiags,
+        linalg,
+    )
+
+    if not is_num(tol, gt=0):
+        raise ValueError(f'Invalid tol parameter ({tol})')
+    if not is_num(lam, gt=0):
+        raise ValueError(f'Invalid lam parameter ({lam})')
+    if not is_int(max_iter, gt=0):
+        raise ValueError(f'Invalid max_iter parameter ({max_iter})')
+    if not isinstance(full_output, bool):
+        raise ValueError(f'Invalid full_output parameter ({max_iter})')
+    y = np.asarray(y)
+    if mask is not None:
+        mask = mask.astype(bool)
+        y_org = y
+        y = y[mask]
+    num = y.size
+
+    diag = np.ones((num-2))
+    D = spdiags([diag, -2*diag, diag], [0, -1, -2], num, num-2)
+
+    H = lam * D.dot(D.T)
+
+    if w is None:
+        w = np.ones(num)
+    W = spdiags(w, 0, num, num)
+
+    error = 1
+    num_iter = 0
+
+    exp_max = int(np.log(float_info.max))
+    while error > tol and num_iter < max_iter:
+        z = linalg.spsolve(W + H, W * y)
+        d = y - z
+        dn = d[d < 0]
+
+        m = np.mean(dn)
+        s = np.std(dn)
+
+        w_new = 1.0 / (1.0 + np.exp(
+            np.clip(2.0 * (d - (2.0*s - m))/s, None, exp_max)))
+        error = np.linalg.norm(w_new - w) / np.linalg.norm(w)
+        num_iter += 1
+        w = w_new
+        W.setdiag(w)
+
+    if mask is not None:
+        zz = np.zeros(y_org.size)
+        zz[mask] = z
+        z = zz
+        if full_output:
+            d = y_org - z
+    if full_output:
+        return z, d, w, num_iter, float(error)
+    return z
+
+
 def select_mask_1d(
         y, x=None, label=None, ref_data=[], preselected_index_ranges=None,
         preselected_mask=None, title=None, xlabel=None, ylabel=None,
         min_num_index_ranges=None, max_num_index_ranges=None,
-        interactive=True):
+        interactive=True, filename=None):
     """Display a lineplot and have the user select a mask.
 
     :param y: One-dimensional data array for which a mask will be
@@ -945,18 +1267,21 @@ def select_mask_1d(
         ranges, defaults to `None`.
     :type max_num_index_ranges: int, optional
     :param interactive: Show the plot and allow user interactions with
-        the matplotlib figure, defults to `True`.
+        the matplotlib figure, defaults to `True`.
     :type interactive: bool, optional
-    :return: A Matplotlib figure, a boolean mask array and the list of
-        selected index ranges.
-    :rtype: matplotlib.figure.Figure, numpy.ndarray,
-        list[tuple(int, int)]
+    :param filename: Save a .png of the plot to filename, defaults to
+        `None`, in which case the plot is not saved.
+    :type filename: str, optional
+    :return: A boolean mask array and the list of selected index
+        ranges.
+    :rtype: numpy.ndarray, list[tuple(int, int)]
     """
     # Third party modules
-    from matplotlib.patches import Patch
-    from matplotlib.widgets import Button, SpanSelector
+    if interactive or filename is not None:
+        from matplotlib.patches import Patch
+        from matplotlib.widgets import Button, SpanSelector
 
-    # local modules
+    # Local modules
     from CHAP.utils.general import index_nearest
 
     def change_fig_title(title):
@@ -1111,10 +1436,25 @@ def select_mask_1d(
             raise ValueError('Invalid parameter preselected_index_ranges '
                              f'({preselected_index_ranges})')
 
+    # Setup the preselected mask and index ranges if provided
+    if preselected_mask is not None:
+        preselected_index_ranges = update_index_ranges(
+            update_mask(
+                np.copy(np.asarray(preselected_mask, dtype=bool)),
+                preselected_index_ranges))
+
+    if not interactive and filename is None:
+
+        # Update the mask with the preselected index ranges
+        selected_mask = update_mask(len(x)*[False], preselected_index_ranges)
+
+        return selected_mask, preselected_index_ranges
+
     spans = []
     fig_title = []
     error_texts = []
 
+    # Setup the Matplotlib figure
     title_pos = (0.5, 0.95)
     title_props = {'fontsize': 'xx-large', 'horizontalalignment': 'center',
                    'verticalalignment': 'bottom'}
@@ -1138,12 +1478,7 @@ def select_mask_1d(
     ax.set_xlim(x[0], x[-1])
     fig.subplots_adjust(bottom=0.0, top=0.85)
 
-    # Setup the preselected mask and index ranges if provided
-    if preselected_mask is not None:
-        preselected_index_ranges = update_index_ranges(
-            update_mask(
-                np.copy(np.asarray(preselected_mask, dtype=bool)),
-                preselected_index_ranges))
+    # Add the preselected index ranges
     for min_, max_ in preselected_index_ranges:
         add_span(None, xrange_init=(x[min_], x[min(max_, num_data-1)]))
 
@@ -1158,7 +1493,8 @@ def select_mask_1d(
         fig.subplots_adjust(bottom=0.2)
 
         # Setup "Add span" button
-        add_span_btn = Button(plt.axes([0.15, 0.05, 0.15, 0.075]), 'Add span')
+        add_span_btn = Button(
+            plt.axes([0.15, 0.05, 0.15, 0.075]), 'Add span')
         add_span_cid = add_span_btn.on_clicked(add_span)
 
         # Setup "Reset" button
@@ -1188,15 +1524,18 @@ def select_mask_1d(
     # Update the mask with the currently selected index ranges
     selected_mask = update_mask(len(x)*[False], selected_index_ranges)
 
-    fig_title[0].set_in_layout(True)
-    fig.tight_layout(rect=(0, 0, 1, 0.95))
+    if filename is not None:
+        fig_title[0].set_in_layout(True)
+        fig.tight_layout(rect=(0, 0, 1, 0.95))
+        fig.savefig(filename)
+    plt.close()
 
-    return fig, selected_mask, selected_index_ranges
+    return selected_mask, selected_index_ranges
 
 
 def select_roi_1d(
         y, x=None, preselected_roi=None, title=None, xlabel=None, ylabel=None,
-        interactive=True):
+        interactive=True, filename=None):
     """Display a 2D plot and have the user select a single region
     of interest.
 
@@ -1217,8 +1556,11 @@ def select_roi_1d(
         defaults to `None`.
     :type ylabel: str, optional
     :param interactive: Show the plot and allow user interactions with
-        the matplotlib figure, defults to `True`.
+        the matplotlib figure, defaults to `True`.
     :type interactive: bool, optional
+    :param filename: Save a .png of the plot to filename, defaults to
+        `None`, in which case the plot is not saved.
+    :type filename: str, optional
     :return: The selected region of interest as array indices and a
         matplotlib figure.
     :rtype: matplotlib.figure.Figure, tuple(int, int)
@@ -1233,16 +1575,17 @@ def select_roi_1d(
                              f'({preselected_roi})')
         preselected_roi = [preselected_roi]
 
-    fig, mask, roi = select_mask_1d(
+    mask, roi = select_mask_1d(
         y, x=x, preselected_index_ranges=preselected_roi, title=title,
         xlabel=xlabel, ylabel=ylabel, min_num_index_ranges=1,
-        max_num_index_ranges=1, interactive=interactive)
+        max_num_index_ranges=1, interactive=interactive, filename=filename)
 
-    return fig, tuple(roi[0])
+    return tuple(roi[0])
 
 def select_roi_2d(
         a, preselected_roi=None, title=None, title_a=None,
-        row_label='row index', column_label='column index', interactive=True):
+        row_label='row index', column_label='column index', interactive=True,
+        filename=None):
     """Display a 2D image and have the user select a single rectangular
        region of interest.
 
@@ -1265,12 +1608,15 @@ def select_roi_2d(
     :param interactive: Show the plot and allow user interactions with
         the matplotlib figure, defaults to `True`.
     :type interactive: bool, optional
-    :return: The selected region of interest as array indices and a
-        matplotlib figure.
-    :rtype: matplotlib.figure.Figure, tuple(int, int, int, int)
+    :param filename: Save a .png of the plot to filename, defaults to
+        `None`, in which case the plot is not saved.
+    :type filename: str, optional
+    :return: The selected region of interest as array indices.
+    :rtype: tuple(int, int, int, int)
     """
     # Third party modules
-    from matplotlib.widgets import Button, RectangleSelector
+    if interactive or filename is not None:
+        from matplotlib.widgets import Button, RectangleSelector
 
     # Local modules
     from CHAP.utils.general import index_nearest
@@ -1327,9 +1673,6 @@ def select_roi_2d(
         change_fig_title(f'Selected ROI: {roi}')
         plt.close()
 
-    fig_title = []
-    subfig_title = []
-
     # Check inputs
     a = np.asarray(a)
     if a.ndim != 2:
@@ -1342,6 +1685,12 @@ def select_roi_2d(
     if title is None:
         title = 'Click and drag to select or adjust a region of interest (ROI)'
 
+    if not interactive and filename is None:
+        return preselected_roi
+
+    fig_title = []
+    subfig_title = []
+
     title_pos = (0.5, 0.95)
     title_props = {'fontsize': 'xx-large', 'horizontalalignment': 'center',
                    'verticalalignment': 'bottom'}
@@ -1370,8 +1719,9 @@ def select_roi_2d(
 
     if not interactive:
 
-        change_fig_title(
-            f'Selected ROI: {tuple(int(v) for v in preselected_roi)}')
+        if preselected_roi is not None:
+            change_fig_title(
+                f'Selected ROI: {tuple(int(v) for v in preselected_roi)}')
 
     else:
 
@@ -1400,20 +1750,26 @@ def select_roi_2d(
         reset_btn.ax.remove()
         confirm_btn.ax.remove()
 
-    fig_title[0].set_in_layout(True)
-    fig.tight_layout(rect=(0, 0, 1, 0.95))
+    if filename is not None:
+        if fig_title:
+            fig_title[0].set_in_layout(True)
+            fig.tight_layout(rect=(0, 0, 1, 0.95))
+        else:
+            fig.tight_layout(rect=(0, 0, 1, 1))
 
-    # Remove the handles before returning the figure
-    if interactive:
-        rects[0]._center_handle.set_visible(False)
-        rects[0]._corner_handles.set_visible(False)
-        rects[0]._edge_handles.set_visible(False)
+        # Remove the handles
+        if interactive:
+            rects[0]._center_handle.set_visible(False)
+            rects[0]._corner_handles.set_visible(False)
+            rects[0]._edge_handles.set_visible(False)
+        fig.savefig(filename)
+    plt.close()
 
     roi = tuple(int(v) for v in rects[0].extents)
     if roi[1]-roi[0] < 1 or roi[3]-roi[2] < 1:
         roi = None
 
-    return fig, roi
+    return roi
 
 
 def select_image_indices(
@@ -1918,9 +2274,11 @@ def nxcopy(
             name=nxobject.nxname)
     elif isinstance(nxobject, (NXlink, NXfield)):
         # The top level nxobject is a (linked) field: return a copy
+        attrs = nxobject.attrs
+        attrs.pop('target', None)
         nxobject_copy = NXfield(
             value=nxobject.nxdata, name=nxobject.nxname,
-            attrs=nxobject.attrs)
+            attrs=attrs)
         return nxobject_copy
     else:
         # Create a group with the same type/name as the nxobject