HB3AIntegratePeaks.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, FileAction, FileProperty,
                        PythonAlgorithm, PropertyMode, ADSValidator,
                        WorkspaceGroup, WorkspaceProperty, MultipleExperimentInfos,
                        IPeaksWorkspace)
from mantid.kernel import Direction, FloatBoundedValidator, StringListValidator, StringArrayProperty
from mantid.simpleapi import (DeleteWorkspace, IntegratePeaksMD,
                              SaveHKL, SaveReflections,
                              CreatePeaksWorkspace, CopySample,
                              AnalysisDataService, FilterPeaks,
                              CombinePeaksWorkspaces, mtd)
import numpy as np


class HB3AIntegratePeaks(PythonAlgorithm):

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

    def seeAlso(self):
        return ["HB3AFindPeaks", "IntegratePeaksMD", "HB3AAdjustSampleNorm", "HB3APredictPeaks"]

    def name(self):
        return "HB3AIntegratePeaks"

    def summary(self):
        return 'Integrates peaks from the input MDEvent workspace and can optionally apply a Lorentz correction to ' \
               'the output peaks workspace; output can be saved in different formats.'

    def PyInit(self):
        self.declareProperty(StringArrayProperty("InputWorkspace", direction=Direction.Input, validator=ADSValidator()),
                             doc="Input MDEvent workspace to use for integration")

        self.declareProperty(StringArrayProperty("PeaksWorkspace", direction=Direction.Input, validator=ADSValidator()),
                             doc="Peaks workspace containing peaks to integrate")

        positive_val = FloatBoundedValidator(lower=0.0)
        self.declareProperty("PeakRadius", defaultValue=1.0, validator=positive_val,
                             doc="Fixed radius around each peak position in which to integrate"
                                 " (same units as input workspace) ")

        self.declareProperty("BackgroundInnerRadius", defaultValue=0.0, validator=positive_val,
                             doc="Inner radius used to evaluate the peak background")
        self.declareProperty("BackgroundOuterRadius", defaultValue=0.0, validator=positive_val,
                             doc="Outer radius used to evaluate the peak background")

        self.declareProperty("ApplyLorentz", defaultValue=True,
                             doc="Whether the Lorentz correction should be applied to the integrated peaks")

        self.declareProperty("RemoveZeroIntensity", defaultValue=True,
                             doc="If to remove peaks with 0 or less intensity from the output")

        formats = StringListValidator()
        formats.addAllowedValue("SHELX")
        formats.addAllowedValue("Fullprof")
        self.declareProperty("OutputFormat", defaultValue="SHELX", validator=formats,
                             doc="Save direction cosines in HKL, or the fullprof format")

        self.declareProperty(FileProperty(name="OutputFile", defaultValue="",
                                          direction=Direction.Input,
                                          action=FileAction.OptionalSave),
                             doc="Filepath to save the integrated peaks workspace in HKL format")

        self.declareProperty(WorkspaceProperty("OutputWorkspace", defaultValue="", direction=Direction.Output,
                                               optional=PropertyMode.Mandatory),
                             doc="Output peaks workspace (copy of input with updated peak intensities)")

    def validateInputs(self):
        issues = dict()

        # Make sure outer radius > inner radius
        inner_radius = self.getProperty("BackgroundInnerRadius").value
        outer_radius = self.getProperty("BackgroundOuterRadius").value

        if outer_radius < inner_radius:
            issues['BackgroundOuterRadius'] = "Outer radius should be >= to inner radius"

        input_workspaces, peak_workspaces = self._expand_groups()

        if len(input_workspaces) != len(peak_workspaces):
            issues["InputWorkspace"] = "The same number of PeaksWorkspace must be provided as InputWorkspace"

        for input_ws in input_workspaces:
            if not (isinstance(AnalysisDataService[input_ws], MultipleExperimentInfos)):
                issues["InputWorkspace"] = "Workspace need to be a MDEventWorkspace"
            elif AnalysisDataService[input_ws].getSpecialCoordinateSystem().name != "QSample":
                issues["InputWorkspace"] = "Input workspace expected to be in QSample, " \
                    "workspace is in '{}'".format(input_ws.getSpecialCoordinateSystem().name)
            elif AnalysisDataService[input_ws].getNumDims() != 3:
                issues["InputWorkspace"] = "Workspace has the wrong number of dimensions"

        for peak_ws in peak_workspaces:
            if not isinstance(AnalysisDataService[peak_ws], IPeaksWorkspace):
                issues["PeaksWorkspace"] = "Workspace need to be a PeaksWorkspace or LeanElasticPeaksWorkspace"

        return issues

    def PyExec(self):
        input_workspaces, peak_workspaces = self._expand_groups()
        output_workspace_name = self.getPropertyValue("OutputWorkspace")

        peak_radius = self.getProperty("PeakRadius").value
        inner_radius = self.getProperty("BackgroundInnerRadius").value
        outer_radius = self.getProperty("BackgroundOuterRadius").value

        remove_0_intensity = self.getProperty("RemoveZeroIntensity").value
        use_lorentz = self.getProperty("ApplyLorentz").value

        multi_ws = len(input_workspaces) > 1

        output_workspaces = []

        for input_ws, peak_ws in zip(input_workspaces, peak_workspaces):
            if multi_ws:
                peaks_ws_name = input_ws + '_' + output_workspace_name
                output_workspaces.append(peaks_ws_name)
            else:
                peaks_ws_name = output_workspace_name

            IntegratePeaksMD(InputWorkspace=input_ws,
                             PeakRadius=peak_radius,
                             BackgroundInnerRadius=inner_radius,
                             BackgroundOuterRadius=outer_radius,
                             PeaksWorkspace=peak_ws,
                             OutputWorkspace=peaks_ws_name)

        if multi_ws:
            peaks_ws_name = output_workspace_name
            CreatePeaksWorkspace(InstrumentWorkspace=input_workspaces[0],
                                 NumberOfPeaks=0,
                                 OutputWorkspace=peaks_ws_name,
                                 OutputType=mtd[peak_workspaces[0]].id().replace('sWorkspace',''))
            CopySample(InputWorkspace=output_workspaces[0],
                       OutputWorkspace=peaks_ws_name,
                       CopyName=False,
                       CopyMaterial=False,
                       CopyEnvironment=False,
                       CopyShape=False,
                       CopyLattice=True)
            for peak_ws in output_workspaces:
                CombinePeaksWorkspaces(peaks_ws_name, peak_ws, OutputWorkspace=peaks_ws_name)
                DeleteWorkspace(peak_ws)

        if use_lorentz:
            # Apply Lorentz correction:
            peaks = AnalysisDataService[peaks_ws_name]
            for p in range(peaks.getNumberPeaks()):
                peak = peaks.getPeak(p)
                lorentz = abs(np.sin(peak.getScattering() * np.cos(peak.getAzimuthal())))
                peak.setIntensity(peak.getIntensity() * lorentz)
                peak.setSigmaIntensity(peak.getSigmaIntensity() * lorentz)

        if remove_0_intensity:
            FilterPeaks(InputWorkspace=peaks_ws_name, OutputWorkspace=peaks_ws_name,
                        FilterVariable='Intensity', FilterValue=0, Operator='>')

        # Write output only if a file path was provided
        if not self.getProperty("OutputFile").isDefault:
            out_format = self.getProperty("OutputFormat").value
            filename = self.getProperty("OutputFile").value

            if out_format == "SHELX":
                SaveHKL(InputWorkspace=peaks_ws_name, Filename=filename, DirectionCosines=True, OutputWorkspace="__tmp")
                DeleteWorkspace("__tmp")
            elif out_format == "Fullprof":
                SaveReflections(InputWorkspace=peaks_ws_name, Filename=filename, Format="Fullprof")
            else:
                # This shouldn't happen
                RuntimeError("Invalid output format given")

        self.setProperty("OutputWorkspace", AnalysisDataService[peaks_ws_name])

    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)

        workspaces = self.getProperty("PeaksWorkspace").value
        peaks_workspaces = []
        for wsname in workspaces:
            wks = AnalysisDataService.retrieve(wsname)
            if isinstance(wks, WorkspaceGroup):
                peaks_workspaces.extend(wks.getNames())
            else:
                peaks_workspaces.append(wsname)

        return input_workspaces, peaks_workspaces


AlgorithmFactory.subscribe(HB3AIntegratePeaks)