PDConvertReciprocalSpace.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.api import (
    PythonAlgorithm,
    AlgorithmFactory,
    WorkspaceUnitValidator,
    MatrixWorkspaceProperty,
    WorkspaceProperty,
    WorkspaceFactory,
)
from mantid.kernel import Direction, StringListValidator, logger

from pystog.converter import Converter

SQ = "S(Q)"
FQ = "F(Q)"
FKQ = "FK(Q)"
DCS = "DCS(Q)"


class PDConvertReciprocalSpace(PythonAlgorithm):
    """
    Convert between different reciprocal space functions
    """

    def category(self):
        """
        Return category
        """
        return "Diffraction\\Utility"

    def name(self):
        """
        Return name
        """
        return "PDConvertReciprocalSpace"

    def seeAlso(self):
        return ["PDConvertRealSpace"]

    def summary(self):
        return "Transforms a Workspace2D between different " + "reciprocal space functions."

    def PyInit(self):
        """
        Declare properties
        """
        allowed_units = WorkspaceUnitValidator("MomentumTransfer")
        self.declareProperty(
            MatrixWorkspaceProperty("InputWorkspace", "", direction=Direction.Input, validator=allowed_units),
            doc="Input workspace with units of momentum transfer",
        )
        functions = [SQ, FQ, FKQ, DCS]
        self.declareProperty("From", SQ, StringListValidator(functions), "Function type in the input workspace")
        self.declareProperty("To", SQ, StringListValidator(functions), "Function type in the output workspace")
        self.declareProperty(WorkspaceProperty("OutputWorkspace", "", direction=Direction.Output), doc="Output workspace")

    def validateInputs(self):
        result = dict()
        input_ws = self.getProperty("InputWorkspace").value
        from_quantity = self.getProperty("From").value
        to_quantity = self.getProperty("To").value

        if from_quantity != to_quantity:  # noop doesn't require material
            requireMaterial = True
            if from_quantity == SQ and to_quantity == FQ:
                requireMaterial = False
            elif from_quantity == FQ and to_quantity == SQ:
                requireMaterial = False
            if requireMaterial:
                if not input_ws.sample().getMaterial():
                    result["InputWorkspace"] = "Please run SetSample or SetSampleMaterial"
                else:
                    cohScattSqrd = input_ws.sample().getMaterial().cohScatterLengthSqrd()
                    if cohScattSqrd == 0.0:
                        if from_quantity in [SQ, FQ] and to_quantity in [DCS, FKQ]:
                            result["To"] = "Require non-zero coherent scattering length"
                        if to_quantity in [SQ, FQ] and from_quantity in [DCS, FKQ]:
                            result["From"] = "Require non-zero coherent scattering length"
        return result

    def PyExec(self):
        input_ws = self.getProperty("InputWorkspace").value
        output_ws_name = self.getProperty("OutputWorkspace").valueAsStr
        from_quantity = self.getProperty("From").value
        to_quantity = self.getProperty("To").value

        if input_ws.name() == output_ws_name:
            output_ws = input_ws
        else:
            output_ws = WorkspaceFactory.create(input_ws)

        self.setProperty("OutputWorkspace", output_ws)
        if from_quantity == to_quantity:
            logger.warning("The input and output functions are the same. Nothing to be done")
            return
        c = Converter()
        transformation = {
            SQ: {FQ: c.S_to_F, FKQ: c.S_to_FK, DCS: c.S_to_DCS},
            FQ: {SQ: c.F_to_S, FKQ: c.F_to_FK, DCS: c.F_to_DCS},
            FKQ: {SQ: c.FK_to_S, FQ: c.FK_to_F, DCS: c.FK_to_DCS},
            DCS: {SQ: c.DCS_to_S, FQ: c.DCS_to_F, FKQ: c.DCS_to_FK},
        }

        if input_ws.sample().getMaterial():
            sample_kwargs = {
                "<b_coh>^2": input_ws.sample().getMaterial().cohScatterLengthSqrd(),
                "<b_tot^2>": input_ws.sample().getMaterial().totalScatterLengthSqrd(),
                "rho": input_ws.sample().getMaterial().numberDensity,
            }
        else:
            sample_kwargs = dict()

        for sp_num in range(input_ws.getNumberHistograms()):
            x = input_ws.readX(sp_num)
            output_ws.setX(sp_num, x)
            y = input_ws.readY(sp_num)
            e = input_ws.readE(sp_num)
            if len(x) == len(y) + 1:
                x = 0.5 * (x[:-1] + x[1:])

            new_y, new_e = transformation[from_quantity][to_quantity](x, y, e, **sample_kwargs)
            output_ws.setY(sp_num, new_y)
            output_ws.setE(sp_num, new_e)


# Register algorithm with Mantid.
AlgorithmFactory.subscribe(PDConvertReciprocalSpace)