HB3AIntegrateDetectorPeaks.py

# Mantid Repository : https://github.com/mantidproject/mantid
#
# Copyright © 2020 ISIS Rutherford Appleton Laboratory UKRI,
#   NScD Oak Ridge National Laboratory, European Spallation Source,
#   Institut Laue - Langevin & CSNS, Institute of High Energy Physics, CAS
# SPDX - License - Identifier: GPL - 3.0 +
from mantid.api import (AlgorithmFactory, IPeaksWorkspaceProperty,
                        PythonAlgorithm, PropertyMode, ADSValidator,
                        WorkspaceGroup)
from mantid.kernel import (Direction, EnabledWhenProperty, IntArrayProperty, IntBoundedValidator,
                           Property, PropertyCriterion, IntArrayLengthValidator, StringArrayProperty,
                           StringListValidator)
from mantid.simpleapi import (mtd, IntegrateMDHistoWorkspace,
                              CreatePeaksWorkspace, DeleteWorkspace,
                              AnalysisDataService, SetGoniometer,
                              ConvertMDHistoToMatrixWorkspace,
                              ConvertToPointData, Fit,
                              CombinePeaksWorkspaces,
                              HB3AAdjustSampleNorm,
                              CentroidPeaksMD, RenameWorkspace)
import numpy as np


class HB3AIntegrateDetectorPeaks(PythonAlgorithm):

    def category(self):
        return "Crystal\\Integration"

    def seeAlso(self):
        return ["HB3AAdjustSampleNorm", "HB3AIntegratePeaks"]

    def name(self):
        return "HB3AIntegrateDetectorPeaks"

    def summary(self):
        return "Integrate four-circle rocking-scan peak in detector space from HB3A"

    def PyInit(self):
        self.declareProperty(StringArrayProperty("InputWorkspace", direction=Direction.Input, validator=ADSValidator()),
                             doc="Workspace or comma-separated workspace list containing input MDHisto scan data.")

        self.declareProperty("Method", direction=Direction.Input, defaultValue="Fitted",
                             validator=StringListValidator(["Counts", "CountsWithFitting", "Fitted"]),
                             doc="Integration method to use")

        self.declareProperty("NumBackgroundPts", direction=Direction.Input, defaultValue=3,
                             validator=IntBoundedValidator(lower=0),
                             doc="Number of background points from beginning and end of scan to use for background estimation")
        self.setPropertySettings("NumBackgroundPts", EnabledWhenProperty("Method", PropertyCriterion.IsEqualTo, "Counts"))

        self.declareProperty("WidthScale", direction=Direction.Input, defaultValue=2,
                             validator=IntBoundedValidator(lower=0, exclusive=True),
                             doc="Controls integration range (+/- WidthScale/2*FWHM) defined around motor positions "
                                 "for CountsWithFitting method")
        self.setPropertySettings("WidthScale",
                                 EnabledWhenProperty("Method", PropertyCriterion.IsEqualTo, "CountsWithFitting"))

        self.declareProperty(IntArrayProperty("LowerLeft", [128, 128], IntArrayLengthValidator(2),
                                              direction=Direction.Input), doc="Region of interest lower-left corner, in detector pixels")
        self.declareProperty(IntArrayProperty("UpperRight", [384, 384], IntArrayLengthValidator(2),
                                              direction=Direction.Input), doc="Region of interest upper-right corner, in detector pixels")

        self.declareProperty("StartX", Property.EMPTY_DBL,
                             doc="The start of the scan axis fitting range in degrees, either omega or chi axis.")
        self.declareProperty("EndX", Property.EMPTY_DBL,
                             doc="The end of the scan axis fitting range in degrees, either omega or chi axis.")

        self.declareProperty("ScaleFactor", 1.0, doc="scale the integrated intensity by this value")
        self.declareProperty("ChiSqMax", 10.0, doc="Fitting resulting in chi-sqaured higher than this won't be added to the output")
        self.declareProperty("SignalNoiseMin", 1.0,
                             doc="Minimum Signal/Noice ratio (Intensity/SigmaIntensity) of peak to be added to the output")
        self.declareProperty("ApplyLorentz", True, doc="If to apply Lorentz Correction to intensity")

        self.declareProperty("OutputFitResults", False, doc="This will output the fitting result workspace and a ROI workspace")
        self.setPropertySettings("OutputFitResults",
                                 EnabledWhenProperty("Method", PropertyCriterion.IsNotEqualTo, "Counts"))
        self.declareProperty("OptimizeQVector", True,
                             doc="This will convert the data to q and optimize the peak location using CentroidPeaksdMD")

        self.declareProperty(IPeaksWorkspaceProperty("OutputWorkspace", "", optional=PropertyMode.Mandatory, direction=Direction.Output),
                             doc="Output Peaks Workspace")

    def PyExec(self):
        input_workspaces = self._expand_groups()
        outWS = self.getPropertyValue("OutputWorkspace")
        CreatePeaksWorkspace(OutputType='LeanElasticPeak',
                             InstrumentWorkspace=input_workspaces[0],
                             NumberOfPeaks=0,
                             OutputWorkspace=outWS, EnableLogging=False)

        method = self.getProperty("Method").value
        n_bkgr_pts = self.getProperty("NumBackgroundPts").value
        n_fwhm = self.getProperty("WidthScale").value
        scale = self.getProperty("ScaleFactor").value
        chisqmax = self.getProperty("ChiSqMax").value
        signalNoiseMin = self.getProperty("SignalNoiseMin").value
        ll = self.getProperty("LowerLeft").value
        ur = self.getProperty("UpperRight").value
        startX = self.getProperty('StartX').value
        endX = self.getProperty('EndX').value
        use_lorentz = self.getProperty("ApplyLorentz").value
        optmize_q = self.getProperty("OptimizeQVector").value
        output_fit = self.getProperty("OutputFitResults").value

        if output_fit and method != "Counts":
            fit_results = WorkspaceGroup()
            AnalysisDataService.addOrReplace(outWS+"_fit_results", fit_results)

        for inWS in input_workspaces:
            tmp_inWS = '__tmp_' + inWS
            IntegrateMDHistoWorkspace(InputWorkspace=inWS,
                                      P1Bin=f'{ll[1]},{ur[1]}',
                                      P2Bin=f'{ll[0]},{ur[0]}',
                                      OutputWorkspace=tmp_inWS,
                                      EnableLogging=False)
            ConvertMDHistoToMatrixWorkspace(tmp_inWS, OutputWorkspace=tmp_inWS, EnableLogging=False)
            data = ConvertToPointData(tmp_inWS, OutputWorkspace=tmp_inWS, EnableLogging=False)

            run = mtd[inWS].getExperimentInfo(0).run()
            scan_log = 'omega' if np.isclose(run.getTimeAveragedStd('phi'), 0.0) else 'phi'
            scan_axis = run[scan_log].value
            scan_step = (scan_axis[-1] - scan_axis[0]) / (scan_axis.size - 1)
            data.setX(0, scan_axis)

            y = data.extractY().flatten()
            x = data.extractX().flatten()

            __tmp_pw = CreatePeaksWorkspace(OutputType='LeanElasticPeak',
                                            InstrumentWorkspace=inWS,
                                            NumberOfPeaks=0,
                                            EnableLogging=False)

            if method != "Counts":
                # fit against gaussian with flat background for both the Fitted and CountsWithFitting methods
                fit_result = self._fit_gaussian(inWS, data, x, y, startX, endX, output_fit)

                if fit_result and fit_result.OutputStatus == 'success' and fit_result.OutputChi2overDoF < chisqmax:
                    B, A, peak_centre, sigma, _ = fit_result.OutputParameters.toDict()['Value']
                    _, errA, _, errs, _ = fit_result.OutputParameters.toDict()['Error']

                    if method == "Fitted":
                        integrated_intensity = A * sigma * np.sqrt(2 * np.pi)

                        # Convert correlation back into covariance
                        cor_As = (fit_result.OutputNormalisedCovarianceMatrix.cell(1, 4) / 100
                                  * fit_result.OutputParameters.cell(1, 2) * fit_result.OutputParameters.cell(3, 2))
                        # σ^2 = 2π (A^2 σ_s^2 + σ_A^2 s^2 + 2 A s σ_As)
                        integrated_intensity_error = np.sqrt(
                            2 * np.pi * (A ** 2 * errs ** 2 + sigma ** 2 * errA ** 2 + 2 * A * sigma * cor_As))

                    elif method == "CountsWithFitting":
                        y = y[slice(np.searchsorted(x, peak_centre - 2.3548 * sigma * n_fwhm / 2),
                                    np.searchsorted(x, peak_centre + 2.3548 * sigma * n_fwhm / 2))]
                        # subtract out the fitted flat background
                        integrated_intensity = (y.sum() - B * y.size) * scan_step
                        integrated_intensity_error = np.sum(np.sqrt(y)) * scan_step

                    # update the goniometer position based on the fitted peak center
                    if scan_log == 'omega':
                        SetGoniometer(Workspace=__tmp_pw, Axis0=f'{peak_centre},0,1,0,-1', Axis1='chi,0,0,1,-1',
                                      Axis2='phi,0,1,0,-1',
                                      EnableLogging=False)
                    else:
                        SetGoniometer(Workspace=__tmp_pw, Axis0='omega,0,1,0,-1', Axis1='chi,0,0,1,-1',
                                      Axis2=f'{peak_centre},0,1,0,-1',
                                      EnableLogging=False)
                else:
                    self.log().warning("Failed to fit workspace {}: Output Status={}, ChiSq={}".format(inWS,
                                                                                                       fit_result.OutputStatus,
                                                                                                       fit_result.OutputChi2overDoF))
                    self._delete_tmp_workspaces(str(__tmp_pw), tmp_inWS)
                    continue
            else:
                integrated_intensity, integrated_intensity_error = self._counts_integration(data, n_bkgr_pts, scan_step)

                # set the goniometer position to use the average of the scan
                SetGoniometer(Workspace=__tmp_pw, Axis0='omega,0,1,0,-1', Axis1='chi,0,0,1,-1',
                              Axis2='phi,0,1,0,-1',
                              EnableLogging=False)

            integrated_intensity *= scale
            integrated_intensity_error *= scale

            peak = __tmp_pw.createPeakHKL([run['h'].getStatistics().median,
                                           run['k'].getStatistics().median,
                                           run['l'].getStatistics().median])
            peak.setWavelength(float(run['wavelength'].value))
            peak.setIntensity(integrated_intensity)
            peak.setSigmaIntensity(integrated_intensity_error)

            if integrated_intensity / integrated_intensity_error > signalNoiseMin:
                __tmp_pw.addPeak(peak)

                # correct q-vector using CentroidPeaksMD
                if optmize_q:
                    __tmp_q_ws = HB3AAdjustSampleNorm(InputWorkspaces=inWS, NormaliseBy='None', EnableLogging=False)
                    __tmp_pw = CentroidPeaksMD(__tmp_q_ws, __tmp_pw, EnableLogging=False)
                    DeleteWorkspace(__tmp_q_ws, EnableLogging=False)

                if use_lorentz:
                    # ILL Neutron Data Booklet, Second Edition, Section 2.9, Part 4.1, Equation 7
                    peak = __tmp_pw.getPeak(0)
                    lorentz = abs(np.sin(peak.getScattering() * np.cos(peak.getAzimuthal())))
                    peak.setIntensity(peak.getIntensity() * lorentz)
                    peak.setSigmaIntensity(peak.getSigmaIntensity() * lorentz)

                CombinePeaksWorkspaces(outWS, __tmp_pw, OutputWorkspace=outWS, EnableLogging=False)

                if output_fit and method != "Counts":
                    fit_results.addWorkspace(RenameWorkspace(tmp_inWS + '_Workspace', outWS + "_" + inWS + '_Workspace',
                                                             EnableLogging=False))
                    fit_results.addWorkspace(
                        RenameWorkspace(tmp_inWS + '_Parameters', outWS + "_" + inWS + '_Parameters',
                                        EnableLogging=False))
                    fit_results.addWorkspace(RenameWorkspace(tmp_inWS + '_NormalisedCovarianceMatrix',
                                                             outWS + "_" + inWS + '_NormalisedCovarianceMatrix',
                                                             EnableLogging=False))
                    fit_results.addWorkspace(IntegrateMDHistoWorkspace(InputWorkspace=inWS,
                                                                       P1Bin=f'{ll[1]},0,{ur[1]}',
                                                                       P2Bin=f'{ll[0]},0,{ur[0]}',
                                                                       P3Bin='0,{}'.format(
                                                                           mtd[inWS].getDimension(2).getNBins()),
                                                                       OutputWorkspace=outWS + "_" + inWS + "_ROI",
                                                                       EnableLogging=False))
            else:
                self.log().warning("Skipping peak from {} because Signal/Noise={:.3f} which is less than {}"
                                   .format(inWS, integrated_intensity/integrated_intensity_error, signalNoiseMin))

            self._delete_tmp_workspaces(str(__tmp_pw), tmp_inWS)

        self.setProperty("OutputWorkspace", mtd[outWS])

    def _delete_tmp_workspaces(self, tmp_peakws, tmp_inWS):
        for tmp_ws in (tmp_peakws, tmp_inWS, tmp_inWS + '_Workspace', tmp_inWS + '_Parameters',
                       tmp_inWS + '_NormalisedCovarianceMatrix'):
            if mtd.doesExist(tmp_ws):
                DeleteWorkspace(tmp_ws, EnableLogging=False)

    def _counts_integration(self, ws, n_background_pts, scan_step):
        """simple cuboid integration of detector counts with estimated background subtracted"""
        y = ws.extractY().flatten()
        avg_background = 0.0
        if n_background_pts != 0:
            avg_background = (y[:n_background_pts].sum() + y[-n_background_pts:].sum()) / (2 * n_background_pts)
        integrated_intensity = (y.sum() - avg_background * y.size) * scan_step
        integrated_intensity_error = np.sum(np.sqrt(y)) * scan_step
        return integrated_intensity, integrated_intensity_error

    def _fit_gaussian(self, inWS, ws, x, y, startX, endX, output_fit):
        """fits ws to Gaussian + Flat background"""
        function = f"name=FlatBackground, A0={np.nanmin(y)};" \
                   f"name=Gaussian, PeakCentre={x[np.nanargmax(y)]}, Height={np.nanmax(y) - np.nanmin(y)}, Sigma=0.25"
        constraints = f"f0.A0 > 0, f1.Height > 0, {x.min()} < f1.PeakCentre < {x.max()}"
        fit_result = ""
        try:
            fit_result = Fit(function, ws, Output=str(ws),
                             IgnoreInvalidData=True,
                             OutputParametersOnly=not output_fit,
                             Constraints=constraints,
                             StartX=startX, EndX=endX,
                             EnableLogging=False)
        except RuntimeError as e:
            self.log().warning("Failed to fit workspace {}: {}".format(inWS, e))
        return fit_result

    def _expand_groups(self):
        """expand workspace groups"""
        workspaces = self.getProperty("InputWorkspace").value
        input_workspaces = []
        for wsname in workspaces:
            wks = AnalysisDataService.retrieve(wsname)
            if isinstance(wks, WorkspaceGroup):
                input_workspaces.extend(wks.getNames())
            else:
                input_workspaces.append(wsname)
        return input_workspaces


AlgorithmFactory.subscribe(HB3AIntegrateDetectorPeaks)