peak_fit.py

#!/usr/bin/python
# coding: utf8
# /*##########################################################################
#
# Copyright (c) 2015-2016 European Synchrotron Radiation Facility
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.
#
# ###########################################################################*/

from __future__ import absolute_import

__authors__ = ["D. Naudet"]
__date__ = "01/06/2016"
__license__ = "MIT"

import logging
import time
import ctypes
import multiprocessing as mp
from threading import Thread
import multiprocessing.sharedctypes as mp_sharedctypes
from silx.math.fit import snip1d

import numpy as np

from ... import config
from ...io import QSpaceH5
from ...io.FitH5 import BackgroundTypes
from ...fit import (GaussianFitter, GaussianResults,
                    CentroidFitter, CentroidResults,
                    SilxFitter, SilxResults)
from .sharedresults import FitTypes


_logger = logging.getLogger(__name__)


disp_times = False


def background_estimation(mode, data):
    """Estimates a background of mode kind for data

    :param BackgroundTypes mode: The kind of background to compute
    :param numpy.ndarray data: The data array to process
    :return: The estimated background as an array with same shape as input
    :rtype: numpy.ndarray
    :raises ValueError: In case mode is not known
    """
    # Background subtraction
    if mode == BackgroundTypes.CONSTANT:
        # Shift data so that smallest value is 0
        return np.ones_like(data) * np.nanmin(data)

    elif mode == BackgroundTypes.LINEAR:
        # Simple linear background
        return np.linspace(data[0], data[-1], num=len(data), endpoint=True)

    elif mode == BackgroundTypes.SNIP:
        # Using snip background
        return snip1d(data, snip_width=len(data))

    elif mode == BackgroundTypes.NONE:
        return np.zeros_like(data)

    else:
        raise ValueError("Unsupported background mode")


class PeakFitter(Thread):
    """
    :param str qspace_f: path to the HDF5 file containing the qspace cubes

    :param FitTypes fit_type:

    :param indices: indices of the cubes (in the input HDF5 dataset) for which
        the dim0/dim1/dim2 peaks coordinates will be computed. E.g : if the array
        [1, 2, 3] is provided, only those cubes will be fitted.
    :type indices: *optional* `array_like`

    :param Union[int,None] n_proc:
        Number of process to use. If None, the config value is used.

    :param BackgroundTypes background: The background subtraction to perform
    """

    READY, RUNNING, DONE, ERROR, CANCELED = __STATUSES = range(5)

    def __init__(self,
                 qspace_f,
                 fit_type=FitTypes.GAUSSIAN,
                 n_peaks=1,
                 indices=None,
                 n_proc=None,
                 roi_indices=None,
                 background=BackgroundTypes.NONE):
        super(PeakFitter, self).__init__()

        self.__results = None
        self.__thread = None
        self.__progress = 0
        self.__callback = None

        self.__status = self.READY

        self.__indices = None

        self.__qspace_f = qspace_f
        self.__fit_type = fit_type
        self.__background = background
        self.__n_peaks = n_peaks

        self.__n_proc = n_proc if n_proc else config.DEFAULT_PROCESS_NUMBER
        self.__shared_progress = mp_sharedctypes.RawArray(ctypes.c_int32,
                                                          self.__n_proc)

        if roi_indices is not None:
            self.__roi_indices = np.array(roi_indices[:])
        else:
            self.__roi_indices = None

        if fit_type not in FitTypes.ALLOWED:
            self.__set_status(self.ERROR)
            raise ValueError('Unknown fit type: {0}'.format(fit_type))

        if background not in BackgroundTypes.ALLOWED:
            self.__set_status(self.ERROR)
            raise ValueError('Unknown background type: {}'.format(background))

        try:
            with QSpaceH5.QSpaceH5(qspace_f) as qspace_h5:
                with qspace_h5.qspace_dset_ctx() as dset:
                    qdata_shape = dset.shape

                n_points = qdata_shape[0]

                if indices is None:
                    indices = list(range(n_points))
                else:
                    indices = indices[:]
        except IOError:
            self.__set_status(self.ERROR)
            raise

        self.__indices = np.array(indices)

    def __set_status(self, status):
        assert status in self.__STATUSES
        self.__status = status

    status = property(lambda self: self.__status)

    results = property(lambda self: self.__results)

    def peak_fit(self, blocking=True, callback=None):
        if self.__thread and self.__thread.is_alive():
            raise RuntimeError('A fit is already in progress.')

        self.__results = None

        if blocking:
            return self.__peak_fit()
        else:
            thread = self.__thread = Thread(target=self.__peak_fit)
            self.__callback = callback
            thread.start()

    def progress(self):
        return (100.0 *
                np.frombuffer(self.__shared_progress, dtype='int32').max() /
                (len(self.__indices) - 1))

    def __peak_fit(self):
        self.__set_status(self.RUNNING)

        t_total = time.time()

        progress = np.frombuffer(self.__shared_progress, dtype='int32')
        progress[:] = 0

        n_indices = len(self.__indices)

        try:
            with QSpaceH5.QSpaceH5(self.__qspace_f) as qspace_h5:
                with qspace_h5:
                    x_pos = qspace_h5.sample_x[self.__indices]
                    y_pos = qspace_h5.sample_y[self.__indices]
        except IOError:
            self.__set_status(self.ERROR)
            raise

        if self.__fit_type == FitTypes.GAUSSIAN:
            fit_class = GaussianFitter
            shared_results = GaussianResults(n_points=n_indices,
                                             n_peaks=max(1, self.__n_peaks))
        elif self.__fit_type == FitTypes.CENTROID:
            fit_class = CentroidFitter
            shared_results = CentroidResults(n_points=n_indices)
        elif self.__fit_type == FitTypes.SILX:
            fit_class = SilxFitter
            shared_results = SilxResults(n_points=n_indices,
                                         n_peaks=max(1, self.__n_peaks))
        else:
            raise RuntimeError('Unknown Fit Type')

        manager = mp.Manager()

        read_lock = manager.Lock()
        idx_queue = manager.Queue()

        pool = mp.Pool(self.__n_proc,
                       initializer=_init_thread,
                       initargs=(shared_results,
                                 self.__shared_progress,
                                 fit_class,
                                 (n_indices, 9),
                                 idx_queue,
                                 qspace_f,
                                 read_lock,
                                 self.__background))

        if disp_times:
            class MyTimes(object):
                def __init__(self):
                    self.t_read = 0.
                    self.t_mask = 0.
                    self.t_fit = 0.
                    self.t_write = 0.

                def update(self, arg):
                    (t_read_, t_mask_, t_fit_, t_write_) = arg
                    self.t_read += t_read_
                    self.t_mask += t_mask_
                    self.t_fit += t_fit_
                    self.t_write += t_write_

            res_times = MyTimes()
            callback = res_times.update
        else:
            callback = None

        # creating the processes
        res_list = []
        for th_idx in range(self.__n_proc):
            arg_list = (th_idx, self.__roi_indices)
            res = pool.apply_async(_fit_process,
                                   args=arg_list,
                                   callback=callback)
            res_list.append(res)

        # sending the image indices
        for i_fit, i_cube in enumerate(self.__indices):
            idx_queue.put([i_fit, i_cube])

        # sending the None value to let the threads know that they should return
        for th_idx in range(self.__n_proc):
            idx_queue.put(None)

        pool.close()
        pool.join()

        t_total = time.time() - t_total
        if disp_times:
            _logger.info('Total : {0}.'.format(t_total))
            _logger.info('Read {0}'.format(res_times.t_read))
            _logger.info('Mask {0}'.format(res_times.t_mask))
            _logger.info('Fit {0}'.format(res_times.t_fit))
            _logger.info('Write {0}'.format(res_times.t_write))

        with QSpaceH5.QSpaceH5(qspace_f) as qspace_h5:
            q_dim0, q_dim1, q_dim2 = qspace_h5.qspace_dimension_values

        if self.__roi_indices is not None:
            q_dim0 = q_dim0[self.__roi_indices[0][0]:self.__roi_indices[0][1]]
            q_dim1 = q_dim1[self.__roi_indices[1][0]:self.__roi_indices[1][1]]
            q_dim2 = q_dim2[self.__roi_indices[2][0]:self.__roi_indices[2][1]]

        fit_results = shared_results.fit_results(sample_x=x_pos,
                                                 sample_y=y_pos,
                                                 q_x=q_dim0,
                                                 q_y=q_dim1,
                                                 q_z=q_dim2,
                                                 background_mode=self.__background)

        self.__results = fit_results

        self.__set_status(self.DONE)

        if self.__callback:
            self.__callback()

        return fit_results


def _init_thread(shared_res_,
                 shared_prog_,
                 fit_class_,
                 result_shape_,
                 idx_queue_,
                 qspace_f_,
                 read_lock_,
                 background_):
    global shared_res, \
        shared_progress, \
        fit_class, \
        result_shape, \
        idx_queue, \
        qspace_f, \
        read_lock, \
        fit_background

    shared_res = shared_res_
    shared_progress = shared_prog_
    fit_class = fit_class_
    result_shape = result_shape_
    idx_queue = idx_queue_
    qspace_f = qspace_f_
    read_lock = read_lock_
    fit_background = background_


def _fit_process(th_idx, roiIndices=None):
    print('Thread {0} started.'.format(th_idx))
    try:
        t_read = 0.
        t_fit = 0.
        t_mask = 0.

        l_shared_res = shared_res.local_copy()
        progress = np.frombuffer(shared_progress, dtype='int32')

        qspace_h5 = QSpaceH5.QSpaceH5(qspace_f)

        # Put this in the main thread
        if roiIndices is not None:
            dim0Slice = slice(roiIndices[0][0], roiIndices[0][1], 1)
            dim1Slice = slice(roiIndices[1][0], roiIndices[1][1], 1)
            dim2Slice = slice(roiIndices[2][0], roiIndices[2][1], 1)

        # TODO : timeout to check if it has been canceled
        # read_lock.acquire()
        with qspace_h5:
            q_dim0, q_dim1, q_dim2 = qspace_h5.qspace_dimension_values

            with qspace_h5.qspace_dset_ctx() as dset:
                q_shape = dset.shape
                q_dtype = dset.dtype
            histo = qspace_h5.histo

            if roiIndices is not None:
                q_dim0 = q_dim0[dim0Slice]
                q_dim1 = q_dim1[dim1Slice]
                q_dim2 = q_dim2[dim2Slice]
                histo = histo[dim0Slice, dim1Slice, dim2Slice]

            mask = np.where(histo > 0)
            weights = histo[mask]

        # read_lock.release()
        read_cube = np.ascontiguousarray(np.zeros(q_shape[1:]),
                                         dtype=q_dtype)

        fitter = fit_class(q_dim0, q_dim1, q_dim2, l_shared_res)

        while True:
            # TODO : timeout
            data = idx_queue.get()

            if data is None:
                break

            i_fit, i_cube = data

            progress[th_idx] = i_fit

            if i_cube % 100 == 0:
                print('Processing cube {0}/{1}.'.format(i_fit, result_shape[0]))

            t0 = time.time()
            with qspace_h5.qspace_dset_ctx() as dset:
                dset.read_direct(read_cube,
                                 source_sel=np.s_[i_cube],
                                 dest_sel=None)
            t_read += time.time() - t0

            if roiIndices is not None:
                cube = read_cube[dim0Slice, dim1Slice, dim2Slice]
            else:
                cube = read_cube

            t0 = time.time()
            cube[mask] = cube[mask] / weights
            t_mask = time.time() - t0

            t0 = time.time()

            z_sum = cube.sum(axis=0).sum(axis=0)
            cube_sum_z = cube.sum(axis=2)
            y_sum = cube_sum_z.sum(axis=0)
            x_sum = cube_sum_z.sum(axis=1)

            if fit_background != BackgroundTypes.NONE:
                # Background subtraction
                for array in (z_sum, y_sum, x_sum):
                    array -= background_estimation(fit_background, array)

            fitter.fit(i_fit, i_cube, x_sum, y_sum, z_sum)

            t_fit += time.time() - t0

            t0 = time.time()

            t_write = time.time() - t0

    except Exception as ex:
        print('EX', ex)

    times = (t_read, t_mask, t_fit, t_write)
    if disp_times:
        print('Thread {0} done ({1}).'.format(th_idx, times))
    return times