PoldiMerge.py

# Mantid Repository : https://github.com/mantidproject/mantid
#
# Copyright © 2018 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 +
# pylint: disable=no-init,invalid-name
from mantid.kernel import StringArrayProperty, Direction
from mantid.simpleapi import *
from mantid.api import *

import numpy as np


class PoldiMerge(PythonAlgorithm):
    comparedPropertyNames = ["TablePositionX", "TablePositionY", "TablePositionZ"]
    comparedInstrumentParameters = [("detector", "two_theta"),
                                    ("chopper", "t0"),
                                    ("chopper", "t0_const")]
    outputWorkspaceName = None
    checkInstruments = True

    def category(self):
        return "SINQ\\Poldi"

    def name(self):
        return "PoldiMerge"

    def summary(self):
        return "PoldiMerge takes a list of workspace names and adds the counts, resulting in a new workspace."

    def PyInit(self):
        self.declareProperty(StringArrayProperty(name="WorkspaceNames",
                                                 direction=Direction.Input),
                             doc="List of Workspace names to merge.")

        self.declareProperty(WorkspaceProperty(name="OutputWorkspace",
                                               defaultValue="MergedPoldiWorkspaces",
                                               direction=Direction.Output),
                             doc="Workspace where all counts from the list workspaces have been added")

        self.declareProperty("CheckInstruments", True, "If checked, only workspaces with equal"
                                                       "instrument parameters are merged."
                                                       "Do not disable without a very good reason.")

    def PyExec(self):
        self.checkInstruments = self.getProperty("CheckInstruments").value

        workspaceNames = self.getProperty("WorkspaceNames").value
        self.outputWorkspaceName = self.getProperty("OutputWorkspace").valueAsStr

        self.log().information("Workspaces to merge: %i" % (len(workspaceNames)))

        workspaces = []

        for wsName in workspaceNames:
            if not AnalysisDataService.doesExist(wsName):
                raise KeyError("Not all strings in the input list are valid workspace names.")

            ws = AnalysisDataService.retrieve(wsName)
            workspaces += self.getWorkspacesRecursive(ws)

        workspaceCount = len(workspaces)

        for i in range(workspaceCount):
            currentWorkspace = workspaces[i]
            for j in range(i + 1, workspaceCount):
                try:
                    self.canMerge(currentWorkspace, workspaces[j])
                except RuntimeError as error:
                    self.handleError(error)

        self.setProperty("OutputWorkspace", MergeRuns(workspaceNames))

    def canMerge(self, leftWorkspace, rightWorkspace):
        if not self.timingsMatch(leftWorkspace.dataX(0), rightWorkspace.dataX(0)):
            raise RuntimeError("Timings don't match")

        # If this option is enabled, don't do any checks
        if not self.checkInstruments:
            self.log().warning('Instrument check has been disabled.')
            return True

        leftRun = leftWorkspace.getRun()
        rightRun = rightWorkspace.getRun()

        if not self.chopperSpeedsMatch(leftRun, rightRun):
            raise RuntimeError(
                "Chopper speeds do not match (" + '&'.join((leftWorkspace.name(), rightWorkspace.name())) + ")")

        return self.propertiesMatch(leftRun, rightRun) and self.instrumentsMatch(leftWorkspace, rightWorkspace)

    def timingsMatch(self, leftXData, rightXData):
        leftDeltaX = leftXData[1] - leftXData[0]
        rightDeltaX = rightXData[1] - rightXData[0]

        return abs(leftDeltaX - rightDeltaX) < 1e-4 and abs(rightXData[0] - leftXData[0]) < 1e-4

    def chopperSpeedsMatch(self, leftRun, rightRun):
        chopperSpeedLeft = self.makePlausibleChopperSpeed(self.getPropertyValue(leftRun.getProperty("ChopperSpeed")))
        chopperSpeedRight = self.makePlausibleChopperSpeed(self.getPropertyValue(rightRun.getProperty("ChopperSpeed")))

        return abs(chopperSpeedLeft - chopperSpeedRight) < 1e-4

    def makePlausibleChopperSpeed(self, chopperSpeed):
        # This is related to ticket #10090, where a new field in new data is used
        # when that ticket is finished, new data files will not need this
        # cleanup method anymore.
        return np.floor((chopperSpeed + 250.0) / 500.0) * 500.0

    def instrumentsMatch(self, leftWorkspace, rightWorkspace):
        leftInstrument = leftWorkspace.getInstrument()
        rightInstrument = rightWorkspace.getInstrument()

        return self.instrumentParametersMatch(leftInstrument, rightInstrument)

    def instrumentParametersMatch(self, leftInstrument, rightInstrument):
        if not leftInstrument.getDetector(0).getPos() == rightInstrument.getDetector(0).getPos():
            raise RuntimeError("Detector positions are not equal")

        for parameterTuple in self.comparedInstrumentParameters:
            leftValue = self.getParameterValue(leftInstrument, parameterTuple)
            rightValue = self.getParameterValue(rightInstrument, parameterTuple)

            if abs(leftValue - rightValue) > 1e-12:
                raise RuntimeError("Instrument parameter '%s'/'%s' does not match" % parameterTuple)

        return True

    def getParameterValue(self, instrument, parameterTuple):
        return instrument.getComponentByName(parameterTuple[0]).getNumberParameter(parameterTuple[1])[0]

    def propertiesMatch(self, leftRun, rightRun):
        for propertyName in self.comparedPropertyNames:
            if leftRun.hasProperty(propertyName) and rightRun.hasProperty(propertyName):
                leftProperty = leftRun.getProperty(propertyName)
                rightProperty = rightRun.getProperty(propertyName)
                if abs(self.getPropertyValue(leftProperty) - self.getPropertyValue(rightProperty)) > 5e-3:
                    raise RuntimeError("Property '%s' does not match" % (propertyName))
            else:
                self.log().warning('Property ' + propertyName + ' is not present in data - skipping comparison.')

        return True

    def getPropertyValue(self, runProperty):
        try:
            return runProperty.value[0]
        except TypeError:
            return runProperty.value

    def handleError(self, error):
        if AnalysisDataService.doesExist(self.outputWorkspaceName):
            AnalysisDataService.remove(self.outputWorkspaceName)

        raise RuntimeError("Workspaces can not be merged. %s. Aborting." % (str(error)))

    def getWorkspacesRecursive(self, workspace):
        returnList = []
        if isinstance(workspace, WorkspaceGroup):
            for i in range(workspace.getNumberOfEntries()):
                returnList += self.getWorkspacesRecursive(workspace.getItem(i))

        elif isinstance(workspace, MatrixWorkspace):
            returnList.append(workspace)

        else:
            raise RuntimeError("Can only merge MatrixWorkspaces, this is " + type(workspace))

        return returnList


AlgorithmFactory.subscribe(PoldiMerge)