ConvertWANDSCDtoQ.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 +
from mantid.api import (PythonAlgorithm, AlgorithmFactory,
                        PropertyMode, WorkspaceProperty, Progress,
                        IMDHistoWorkspaceProperty, mtd)
from mantid.kernel import Direction, FloatArrayProperty, FloatArrayLengthValidator, StringListValidator, FloatBoundedValidator
from mantid import config
from mantid import logger
import numpy as np


class ConvertWANDSCDtoQ(PythonAlgorithm):

    def category(self):
        return 'DataHandling\\Nexus'

    def seeAlso(self):
        return [ "LoadWANDSCD", "ConvertHFIRSCDtoMDE" ]

    def name(self):
        return 'ConvertWANDSCDtoQ'

    def summary(self):
        return 'Convert the output of LoadWANDSCD to Q or HKL'

    def PyInit(self):
        self.declareProperty(IMDHistoWorkspaceProperty("InputWorkspace", "",
                                                       optional=PropertyMode.Mandatory,
                                                       direction=Direction.Input),
                             "Input Workspace")
        self.declareProperty(IMDHistoWorkspaceProperty("NormalisationWorkspace", "",
                                                       optional=PropertyMode.Optional,
                                                       direction=Direction.Input),
                             "Workspace to use for normalisation")
        self.declareProperty(WorkspaceProperty("UBWorkspace", "",
                                               optional=PropertyMode.Optional,
                                               direction=Direction.Input),
                             "Workspace containing the UB matrix to use")
        self.declareProperty("Wavelength", 1.488, validator=FloatBoundedValidator(0.0), doc="Wavelength to set the workspace")
        self.declareProperty("S1Offset", 0., doc="Offset to apply (in degrees) to the s1 of the input workspace")
        self.declareProperty('NormaliseBy', 'Monitor', StringListValidator(['None', 'Time', 'Monitor']),
                             "Normalise to monitor, time or None.")
        self.declareProperty('Frame', 'Q_sample', StringListValidator(['Q_sample', 'HKL']),
                             "Selects Q-dimensions of the output workspace")
        self.declareProperty(FloatArrayProperty("Uproj", [1,0,0], FloatArrayLengthValidator(3), direction=Direction.Input),
                             "Defines the first projection vector of the target Q coordinate system in HKL mode")
        self.declareProperty(FloatArrayProperty("Vproj", [0,1,0], FloatArrayLengthValidator(3), direction=Direction.Input),
                             "Defines the second projection vector of the target Q coordinate system in HKL mode")
        self.declareProperty(FloatArrayProperty("Wproj", [0,0,1], FloatArrayLengthValidator(3), direction=Direction.Input),
                             "Defines the third projection vector of the target Q coordinate system in HKL mode")
        self.declareProperty(FloatArrayProperty("BinningDim0", [-8.02,8.02,401], FloatArrayLengthValidator(3), direction=Direction.Input),
                             "Binning parameters for the 0th dimension. Enter it as a"
                             "comma-separated list of values with the"
                             "format: 'minimum,maximum,number_of_bins'.")
        self.declareProperty(FloatArrayProperty("BinningDim1", [-0.82,0.82,41], FloatArrayLengthValidator(3), direction=Direction.Input),
                             "Binning parameters for the 1st dimension. Enter it as a"
                             "comma-separated list of values with the"
                             "format: 'minimum,maximum,number_of_bins'.")
        self.declareProperty(FloatArrayProperty("BinningDim2", [-8.02,8.02,401], FloatArrayLengthValidator(3), direction=Direction.Input),
                             "Binning parameters for the 2nd dimension. Enter it as a"
                             "comma-separated list of values with the"
                             "format: 'minimum,maximum,number_of_bins'.")
        self.declareProperty('KeepTemporaryWorkspaces', False,
                             "If True the normalization and data workspaces in addition to the normalized data will be outputted")
        self.declareProperty("ObliquityParallaxCoefficient", 1.0, validator=FloatBoundedValidator(0.0),
                             doc="Geometrical correction for shift in vertical beam position due to wide beam.")
        self.declareProperty(WorkspaceProperty("OutputWorkspace", "",
                                               optional=PropertyMode.Mandatory,
                                               direction=Direction.Output),
                             "Output Workspace")

    def validateInputs(self):  # noqa C901
        issues = dict()

        inWS = self.getProperty("InputWorkspace").value
        instrument = inWS.getExperimentInfo(0).getInstrument().getName()

        if inWS.getNumDims() != 3:
            issues["InputWorkspace"] = "InputWorkspace has wrong number of dimensions, need 3"
            return issues

        d0 = inWS.getDimension(0)
        d1 = inWS.getDimension(1)
        d2 = inWS.getDimension(2)
        number_of_runs = d2.getNBins()

        if (d0.name != 'y' or d1.name != 'x' or d2.name != 'scanIndex'):
            issues["InputWorkspace"] = "InputWorkspace has wrong dimensions"
            return issues

        if inWS.getNumExperimentInfo() == 0:
            issues["InputWorkspace"] = "InputWorkspace is missing ExperimentInfo"
            return issues

        # Check that all logs are there and are of correct length
        run = inWS.getExperimentInfo(0).run()

        # Check number of goniometers
        if run.getNumGoniometers() != number_of_runs:
            issues["InputWorkspace"] = "goniometers not set correctly, did you run SetGoniometer with Average=False"

        if instrument == "HB3A":
            for prop in ['monitor', 'time']:
                if run.hasProperty(prop):
                    p = run.getProperty(prop).value
                    if np.size(p) != number_of_runs:
                        issues["InputWorkspace"] = "log {} is of incorrect length".format(prop)
                else:
                    issues["InputWorkspace"] = "missing log {}".format(prop)
        else:
            for prop in ['duration', 'monitor_count']:
                if run.hasProperty(prop):
                    p = run.getProperty(prop).value
                    if np.size(p) != number_of_runs:
                        issues["InputWorkspace"] = "log {} is of incorrect length".format(prop)
                else:
                    issues["InputWorkspace"] = "missing log {}".format(prop)

        for prop in ['azimuthal', 'twotheta']:
            if run.hasProperty(prop):
                p = run.getProperty(prop).value
                if np.size(p) != d0.getNBins()*d1.getNBins():
                    issues["InputWorkspace"] = "log {} is of incorrect length".format(prop)

        normWS = self.getProperty("NormalisationWorkspace").value

        if normWS:
            nd0 = normWS.getDimension(0)
            nd1 = normWS.getDimension(1)
            nd2 = normWS.getDimension(2)
            if (nd0.name != d0.name or nd0.getNBins() != d0.getNBins()
               or nd1.name != d1.name or nd1.getNBins() != d1.getNBins()
               or nd2.name != d2.name):
                issues["NormalisationWorkspace"] = "NormalisationWorkspace dimensions are not compatible with InputWorkspace"

        ubWS = self.getProperty("UBWorkspace").value
        if ubWS:
            try:
                sample = ubWS.sample()
            except AttributeError:
                sample = ubWS.getExperimentInfo(0).sample()
            if not sample.hasOrientedLattice():
                issues["UBWorkspace"] = "UBWorkspace does not has an OrientedLattice"
        else:
            if self.getProperty("Frame").value == 'HKL':
                if not inWS.getExperimentInfo(0).sample().hasOrientedLattice():
                    issues["Frame"] = "HKL selected but neither an UBWorkspace workspace was provided or " \
                                      "the InputWorkspace has an OrientedLattice"

        return issues

    def PyExec(self):  # noqa C901
        inWS = self.getProperty("InputWorkspace").value
        normWS = self.getProperty("NormalisationWorkspace").value
        _norm = bool(normWS)

        instrument = inWS.getExperimentInfo(0).getInstrument().getName()

        dim0_min, dim0_max, dim0_bins = self.getProperty('BinningDim0').value
        dim1_min, dim1_max, dim1_bins = self.getProperty('BinningDim1').value
        dim2_min, dim2_max, dim2_bins = self.getProperty('BinningDim2').value
        dim0_bins = int(dim0_bins)
        dim1_bins = int(dim1_bins)
        dim2_bins = int(dim2_bins)
        dim0_bin_size = (dim0_max-dim0_min)/dim0_bins
        dim1_bin_size = (dim1_max-dim1_min)/dim1_bins
        dim2_bin_size = (dim2_max-dim2_min)/dim2_bins

        data_array = inWS.getSignalArray() # getSignalArray returns a F_CONTIGUOUS view of the signal array

        number_of_runs = data_array.shape[2]

        progress = Progress(self, 0.0, 1.0, number_of_runs+4)

        normaliseBy = self.getProperty("NormaliseBy").value
        if normaliseBy == "Monitor":
            if instrument == "HB3A":
                scale = np.asarray(inWS.getExperimentInfo(0).run().getProperty('monitor').value)
            else:
                scale = np.asarray(inWS.getExperimentInfo(0).run().getProperty('monitor_count').value)
        elif normaliseBy == "Time":
            if instrument == "HB3A":
                scale = np.asarray(inWS.getExperimentInfo(0).run().getProperty('time').value)
            else:
                scale = np.asarray(inWS.getExperimentInfo(0).run().getProperty('duration').value)
        else:
            scale = np.ones(number_of_runs)

        if _norm:
            if normaliseBy == "Monitor":
                if instrument == "HB3A":
                    norm_scale = np.sum(normWS.getExperimentInfo(0).run().getProperty('monitor').value)
                else:
                    norm_scale = np.sum(normWS.getExperimentInfo(0).run().getProperty('monitor_count').value)
            elif normaliseBy == "Time":
                if instrument == "HB3A":
                    norm_scale = np.sum(normWS.getExperimentInfo(0).run().getProperty('time').value)
                else:
                    norm_scale = np.sum(normWS.getExperimentInfo(0).run().getProperty('duration').value)
            else:
                norm_scale = 1.
            norm_array = normWS.getSignalArray().sum(axis=2)

        W = np.eye(3)
        UBW = np.eye(3)
        if self.getProperty("Frame").value == 'HKL':
            W[:,0] = self.getProperty('Uproj').value
            W[:,1] = self.getProperty('Vproj').value
            W[:,2] = self.getProperty('Wproj').value
            ubWS = self.getProperty("UBWorkspace").value
            if ubWS:
                try:
                    ol = ubWS.sample().getOrientedLattice()
                except AttributeError:
                    ol = ubWS.getExperimentInfo(0).sample().getOrientedLattice()
                logger.notice("Using UB matrix from {} with {}".format(ubWS.name(), ol))
            else:
                ol = inWS.getExperimentInfo(0).sample().getOrientedLattice()
                logger.notice("Using UB matrix from {} with {}".format(inWS.name(), ol))
            UB = ol.getUB()
            UBW = np.dot(UB, W)
            char_dict = {0:'0', 1:'{1}', -1:'-{1}'}
            chars=['H','K','L']
            names = ['['+','.join(char_dict.get(j, '{0}{1}')
                                  .format(j,chars[np.argmax(np.abs(W[:,i]))]) for j in W[:,i])+']' for i in range(3)]
            units = 'in {:.3f} A^-1,in {:.3f} A^-1,in {:.3f} A^-1'.format(ol.qFromHKL(W[0]).norm(),
                                                                          ol.qFromHKL(W[1]).norm(),
                                                                          ol.qFromHKL(W[2]).norm())
            frames = 'HKL,HKL,HKL'
            k = 1/self.getProperty("Wavelength").value # Not 2pi/wavelength to save dividing by 2pi later
        else:
            names = 'Q_sample_x,Q_sample_y,Q_sample_z'
            units = 'Angstrom^-1,Angstrom^-1,Angstrom^-1'
            frames = 'QSample,QSample,QSample'
            k = 2*np.pi/self.getProperty("Wavelength").value

        progress.report('Calculating Qlab for each pixel')
        if inWS.getExperimentInfo(0).run().hasProperty('twotheta'):
            polar = np.array(inWS.getExperimentInfo(0).run().getProperty('twotheta').value)
        else:
            di = inWS.getExperimentInfo(0).detectorInfo()
            polar = np.array([di.twoTheta(i) for i in range(di.size()) if not di.isMonitor(i)])
            if inWS.getExperimentInfo(0).getInstrument().getName() == 'HB3A':
                polar = polar.reshape(512*3, 512).T.flatten()

        if inWS.getExperimentInfo(0).run().hasProperty('azimuthal'):
            azim = np.array(inWS.getExperimentInfo(0).run().getProperty('azimuthal').value)
        else:
            di = inWS.getExperimentInfo(0).detectorInfo()
            azim = np.array([di.azimuthal(i) for i in range(di.size()) if not di.isMonitor(i)])
            if inWS.getExperimentInfo(0).getInstrument().getName() == 'HB3A':
                azim = azim.reshape(512*3, 512).T.flatten()

        # check convention to determine the sign
        if config['Q.convention'] == 'Crystallography':
            k *= -1.0

        cop = self.getProperty('ObliquityParallaxCoefficient').value

        qlab = np.vstack((np.sin(polar)*np.cos(azim),
                          np.sin(polar)*np.sin(azim)*cop,
                          np.cos(polar) - 1)).T * -k # Kf - Ki(0,0,1)

        progress.report('Calculating Q volume')

        output = np.zeros((dim0_bins+2, dim1_bins+2, dim2_bins+2))
        outputr = output.ravel()

        output_scale = np.zeros_like(output)
        output_scaler = output_scale.ravel()

        if _norm:
            output_norm = np.zeros_like(output)
            output_normr = output_norm.ravel()
            output_norm_scale = np.zeros_like(output)
            output_norm_scaler = output_norm_scale.ravel()

        bin_size = np.array([[dim0_bin_size],
                             [dim1_bin_size],
                             [dim2_bin_size]])

        offset = np.array([[dim0_min/dim0_bin_size],
                           [dim1_min/dim1_bin_size],
                           [dim2_min/dim2_bin_size]])-0.5

        assert not data_array[:,:,0].flags.owndata
        assert not data_array[:,:,0].ravel('F').flags.owndata
        assert data_array[:,:,0].flags.fnc

        s1offset = np.deg2rad(self.getProperty("S1Offset").value)
        s1offset = np.array([[ np.cos(s1offset), 0, np.sin(s1offset)],
                             [             0,    1,                0],
                             [-np.sin(s1offset), 0, np.cos(s1offset)]])

        for n in range(number_of_runs):
            R = inWS.getExperimentInfo(0).run().getGoniometer(n).getR()
            R = np.dot(s1offset, R)
            RUBW = np.dot(R,UBW)
            q = np.round(np.dot(np.linalg.inv(RUBW),qlab.T)/bin_size-offset).astype(np.int)
            q_index = np.ravel_multi_index(q, (dim0_bins+2, dim1_bins+2, dim2_bins+2), mode='clip')
            q_uniq, inverse = np.unique(q_index, return_inverse=True)
            outputr[q_uniq] += np.bincount(inverse, data_array[:,:,n].ravel('F'))
            output_scaler[q_uniq] += np.bincount(inverse)*scale[n]
            if _norm:
                output_normr[q_uniq] += np.bincount(inverse, norm_array.ravel('F'))
                output_norm_scaler[q_uniq] += np.bincount(inverse)

            progress.report()

        if _norm:
            output *= output_norm_scale*norm_scale
            output_norm *= output_scale
        else:
            output_norm = output_scale

        if self.getProperty('KeepTemporaryWorkspaces').value:
            # Create data workspace
            progress.report('Creating data MDHistoWorkspace')
            createWS_alg = self.createChildAlgorithm("CreateMDHistoWorkspace", enableLogging=False)
            createWS_alg.setProperty("SignalInput", output[1:-1,1:-1,1:-1].ravel('F'))
            createWS_alg.setProperty("ErrorInput", np.sqrt(output[1:-1,1:-1,1:-1].ravel('F')))
            createWS_alg.setProperty("Dimensionality", 3)
            createWS_alg.setProperty("Extents", '{},{},{},{},{},{}'.format(dim0_min,dim0_max,dim1_min,dim1_max,dim2_min,dim2_max))
            createWS_alg.setProperty("NumberOfBins", '{},{},{}'.format(dim0_bins,dim1_bins,dim2_bins))
            createWS_alg.setProperty("Names", names)
            createWS_alg.setProperty("Units", units)
            createWS_alg.setProperty("Frames", frames)
            createWS_alg.execute()
            outWS_data = createWS_alg.getProperty("OutputWorkspace").value
            mtd.addOrReplace(self.getPropertyValue("OutputWorkspace")+'_data', outWS_data)

            # Create normalisation workspace
            progress.report('Creating norm MDHistoWorkspace')
            createWS_alg = self.createChildAlgorithm("CreateMDHistoWorkspace", enableLogging=False)
            createWS_alg.setProperty("SignalInput", output_norm[1:-1,1:-1,1:-1].ravel('F'))
            createWS_alg.setProperty("ErrorInput", np.sqrt(output_norm[1:-1,1:-1,1:-1].ravel('F')))
            createWS_alg.setProperty("Dimensionality", 3)
            createWS_alg.setProperty("Extents", '{},{},{},{},{},{}'.format(dim0_min,dim0_max,dim1_min,dim1_max,dim2_min,dim2_max))
            createWS_alg.setProperty("NumberOfBins", '{},{},{}'.format(dim0_bins,dim1_bins,dim2_bins))
            createWS_alg.setProperty("Names", names)
            createWS_alg.setProperty("Units", units)
            createWS_alg.setProperty("Frames", frames)
            createWS_alg.execute()
            mtd.addOrReplace(self.getPropertyValue("OutputWorkspace")+'_normalization', createWS_alg.getProperty("OutputWorkspace").value)

        old_settings = np.seterr(divide='ignore', invalid='ignore') # Ignore RuntimeWarning: invalid value encountered in true_divide
        output /= output_norm # We often divide by zero here and we get NaN's, this is desired behaviour
        np.seterr(**old_settings)

        progress.report('Creating MDHistoWorkspace')
        createWS_alg = self.createChildAlgorithm("CreateMDHistoWorkspace", enableLogging=False)
        createWS_alg.setProperty("SignalInput", output[1:-1,1:-1,1:-1].ravel('F'))
        createWS_alg.setProperty("ErrorInput", np.sqrt(output[1:-1,1:-1,1:-1].ravel('F')))
        createWS_alg.setProperty("Dimensionality", 3)
        createWS_alg.setProperty("Extents", '{},{},{},{},{},{}'.format(dim0_min,dim0_max,dim1_min,dim1_max,dim2_min,dim2_max))
        createWS_alg.setProperty("NumberOfBins", '{},{},{}'.format(dim0_bins,dim1_bins,dim2_bins))
        createWS_alg.setProperty("Names", names)
        createWS_alg.setProperty("Units", units)
        createWS_alg.setProperty("Frames", frames)
        createWS_alg.execute()
        outWS = createWS_alg.getProperty("OutputWorkspace").value

        # Copy experiment infos
        if inWS.getNumExperimentInfo() > 0:
            outWS.copyExperimentInfos(inWS)

        outWS.getExperimentInfo(0).run().addProperty('RUBW_MATRIX', list(UBW.flatten()), True)
        outWS.getExperimentInfo(0).run().addProperty('W_MATRIX', list(W.flatten()), True)
        outWS.getExperimentInfo(0).run().addProperty('wavelength', self.getProperty("Wavelength").value, True)
        try:
            if outWS.getExperimentInfo(0).sample().hasOrientedLattice():
                outWS.getExperimentInfo(0).sample().getOrientedLattice().setUB(UB)
        except NameError:
            pass

        if self.getProperty('KeepTemporaryWorkspaces').value:
            outWS_data.copyExperimentInfos(outWS)

        progress.report()
        self.setProperty("OutputWorkspace", outWS)


AlgorithmFactory.subscribe(ConvertWANDSCDtoQ)