Add method to BaseSX to plot result of IntegratePeaksMD and save in pdf

docs/source/release/v6.10.0/Diffraction/Single_Crystal/New_features/36947.rst

@@ -0,0 +1 @@
+- Add method ``plot_integrated_peaks_MD`` to ``BaseSX`` to plot result of IntegratePeaksMD and save in pdf

scripts/Diffraction/single_crystal/base_sx.py

@@ -9,11 +9,15 @@
 from enum import Enum
 import mantid.simpleapi as mantid
 from mantid.api import FunctionFactory, AnalysisDataService as ADS
-from mantid.kernel import logger
+from mantid.kernel import logger, SpecialCoordinateSystem
 from FindGoniometerFromUB import getSignMaxAbsValInCol
 from mantid.geometry import CrystalStructure, SpaceGroupFactory, ReflectionGenerator, ReflectionConditionFilter
 from os import path
-
+from json import loads as json_loads
+from matplotlib.pyplot import subplots, close
+from matplotlib.patches import Circle
+from matplotlib.colors import LogNorm
+from matplotlib.backends.backend_pdf import PdfPages
 from abc import ABC, abstractmethod
 
 
@@ -614,6 +618,155 @@ def remove_non_integrated_peaks(peaks):
             EnableLogging=False,
         )
 
+    @staticmethod
+    def plot_integrated_peaks_MD(wsMD, peaks, filename, nbins_max=21, extent=1.5, log_norm=True):
+        """
+        :param wsMD:  MD workspace to plot
+        :param peaks: integrated peaks using IntegratePeaksMD
+        :param filename: filename to store pdf output
+        :param nbins: number of bins along major radius of ellipsoid
+        :param extent: extent in units of largest outer background radius
+        :param log_norm: use log normalisation in colorscale
+        """
+        wsMD = BaseSX.retrieve(wsMD)
+        peaks = BaseSX.retrieve(peaks)
+
+        # find appropriate getter for peak centre given MD frame
+        frame = wsMD.getSpecialCoordinateSystem()
+        frame_to_peak_centre_attr = "getQLabFrame"
+        if frame == SpecialCoordinateSystem.QSample:
+            frame_to_peak_centre_attr = "getQSampleFrame"
+        elif frame == SpecialCoordinateSystem.HKL:
+            frame_to_peak_centre_attr = "getHKL"
+
+        # loop over peaks and plot
+        try:
+            with PdfPages(filename) as pdf:
+                ws_cut = None
+                for ipk, pk in enumerate(peaks):
+                    peak_shape = pk.getPeakShape()
+                    if peak_shape.shapeName().lower() == "none":
+                        continue
+                    ws_cut, radii, bg_inner_radii, bg_outer_radii, box_lengths, imax = BaseSX._bin_MD_around_peak(
+                        wsMD, pk, peak_shape, nbins_max, extent, frame_to_peak_centre_attr
+                    )
+
+                    fig, axes = subplots(1, 3, figsize=(12, 4), subplot_kw={"projection": "mantid"})
+                    for iax, ax in enumerate(axes):
+                        dims = list(range(3))
+                        dims.pop(iax)
+                        # plot slice
+                        im = ax.imshow(ws_cut.getSignalArray().sum(axis=iax)[:, ::-1].T)  # reshape to match sliceviewer
+                        im.set_extent([-1, 1, -1, 1])  # so that ellipsoid is a circle
+                        if log_norm:
+                            im.set_norm(LogNorm())
+                        # plot peak position
+                        ax.plot(0, 0, "xr")
+                        # plot peak representation (circle given extents)
+                        patch = Circle((0, 0), radii[imax] / box_lengths[imax], facecolor="none", edgecolor="r", ls="--")
+                        ax.add_patch(patch)
+                        # add background shell
+                        if np.all(bg_outer_radii > 1e-10):
+                            patch = Circle(
+                                (0, 0), bg_outer_radii[imax] / box_lengths[imax], facecolor="none", edgecolor=3 * [0.7]
+                            )  # outer radius
+                            ax.add_patch(patch)
+                            patch = Circle(
+                                (0, 0), bg_inner_radii[imax] / box_lengths[imax], facecolor="none", edgecolor=3 * [0.7], ls="--"
+                            )  # inner radius
+                            ax.add_patch(patch)
+                        # format axes
+                        ax.set_xlim(-1, 1)
+                        ax.set_ylim(-1, 1)
+                        ax.set_aspect("equal")
+                        ax.set_xlabel(ws_cut.getDimension(dims[0]).name)
+                        ax.set_ylabel(ws_cut.getDimension(dims[1]).name)
+                    fig.suptitle(
+                        f"{ipk} ({','.join(str(np.round(pk.getHKL(), 2))[1:-1].split())})"
+                        f"  $I/\\sigma$={np.round(pk.getIntensityOverSigma(), 2)}\n"
+                        rf"$\lambda$={np.round(pk.getWavelength(), 2)} $\AA$; "
+                        rf"$2\theta={np.round(np.degrees(pk.getScattering()), 1)}^\circ$; "
+                        rf"d={np.round(pk.getDSpacing(), 2)} $\AA$"
+                    )
+                    fig.tight_layout()
+                    pdf.savefig(fig)
+                    close(fig)
+
+                if ws_cut is not None:
+                    mantid.DeleteWorkspace(ws_cut)
+        except OSError:
+            raise RuntimeError(
+                f"OutputFile ({filename}) could not be opened - please check it is not open by "
+                f"another programme and that the user has permission to write to that directory."
+            )
+
+    @staticmethod
+    def _bin_MD_around_peak(wsMD, pk, peak_shape, nbins_max, extent, frame_to_peak_centre_attr):
+        shape_info = json_loads(peak_shape.toJSON())
+        if peak_shape.shapeName().lower() == "spherical":
+            BaseSX.convert_spherical_representation_to_ellipsoid(shape_info)
+        # get radii
+        radii = np.array([shape_info[f"radius{iax}"] for iax in range(3)])
+        bg_inner_radii = np.array([shape_info[f"background_inner_radius{iax}"] for iax in range(3)])
+        bg_outer_radii = np.array([shape_info[f"background_outer_radius{iax}"] for iax in range(3)])
+        isort = np.argsort(radii)
+        imax = isort[-1]
+        # eignevectors of ellipsoid
+        evecs = np.zeros((3, 3))
+        for iax in range(len(radii)):
+            evec = np.array([float(elem) for elem in shape_info[f"direction{iax}"].split()])
+            evecs[:, iax] = evec / np.linalg.norm(evec)
+        # center and translation
+        translation = np.array([shape_info[f"translation{iax}"] for iax in range(3)])
+        cen = getattr(pk, frame_to_peak_centre_attr)()
+        cen = np.matmul(evecs.T, np.array(cen) + translation)
+        # get extents
+        box_lengths = bg_outer_radii if np.all(bg_outer_radii > 1e-10) else radii
+        box_lengths = box_lengths * extent
+        extents = np.vstack((cen - box_lengths, cen + box_lengths))
+        # get nbins along each axis
+        nbins = np.array([int(nbins_max * radii[iax] / radii[imax]) for iax in range(len(radii))])
+        # ensure minimum of 3 bins inside radius along each dim
+        min_nbins_in_radius = np.min(nbins * radii / box_lengths)
+        if min_nbins_in_radius < 3:
+            nbins = nbins * 3 / min_nbins_in_radius
+        # call BinMD
+        ws_cut = mantid.BinMD(
+            InputWorkspace=wsMD,
+            OutputWorkspace="__ws_cut",
+            AxisAligned=False,
+            BasisVector0=r"Q$_0$,unit," + ",".join(np.array2string(evecs[:, 0], precision=6).strip("[]").split()),
+            BasisVector1=r"Q$_1$,unit," + ",".join(np.array2string(evecs[:, 1], precision=6).strip("[]").split()),
+            BasisVector2=r"Q$_2$,unit," + ",".join(np.array2string(evecs[:, 2], precision=6).strip("[]").split()),
+            OutputExtents=extents.flatten(order="F"),
+            OutputBins=nbins.astype(int),
+            EnableLogging=False,
+        )
+        return ws_cut, radii, bg_inner_radii, bg_outer_radii, box_lengths, imax
+
+    @staticmethod
+    def _convert_spherical_representation_to_ellipsoid(shape_info):
+        # copied from mantidqt.widgets.sliceviewer.peaksviewer.representation.ellipsoid - can't import here though
+        # convert shape_info dict from sphere to ellipsoid for plotting
+        for key in ["radius", "background_inner_radius", "background_outer_radius"]:
+            if key in shape_info:
+                shape_info[f"{key}{0}"] = shape_info.pop(key)
+            else:
+                shape_info[f"{key}{0}"] = 0.0  # null value
+            for idim in [1, 2]:
+                shape_info[f"{key}{idim}"] = shape_info[f"{key}{0}"]
+        # add axes along basis vecs of frame and set 0 translation
+        shape_info.update(
+            {
+                "direction0": "1 0 0",
+                "direction1": "0 1 0",
+                "direction2": "0 0 1",
+                "translation0": 0.0,
+                "translation1": 0.0,
+                "translation2": 0.0,
+            }
+        )
+
     @staticmethod
     def retrieve(ws):
         if isinstance(ws, str):

scripts/Diffraction/single_crystal/test/test_base_sx.py

@@ -5,14 +5,19 @@
     LoadEmptyInstrument,
     AnalysisDataService,
     CreatePeaksWorkspace,
+    CreateMDWorkspace,
     SetUB,
     IndexPeaks,
     CreateSampleWorkspace,
     SetSample,
     TransformHKL,
+    IntegratePeaksMD,
+    FakeMDEventData,
 )
 from mantid.dataobjects import Workspace2D
 from Diffraction.single_crystal.base_sx import BaseSX
+import tempfile
+import shutil
 
 base_sx_path = "Diffraction.single_crystal.base_sx"
 
@@ -23,10 +28,12 @@ def setUpClass(cls):
         cls.ws = LoadEmptyInstrument(InstrumentName="SXD", OutputWorkspace="empty")
         axis = cls.ws.getAxis(0)
         axis.setUnit("TOF")
+        cls._test_dir = tempfile.mkdtemp()
 
     @classmethod
     def tearDownClass(cls):
         AnalysisDataService.clear()
+        shutil.rmtree(cls._test_dir)
 
     def test_retrieve(self):
         self.assertTrue(isinstance(BaseSX.retrieve("empty"), Workspace2D))  # ADS name of self.ws is "empty"
@@ -151,6 +158,39 @@ def test_make_UB_consistent(self):
             allclose(peaks_ref.sample().getOrientedLattice().getUB().tolist(), peaks.sample().getOrientedLattice().getUB().tolist())
         )
 
+    def test_plot_integrated_peaks_MD(self):
+        # make fake dataset
+        ws = CreateMDWorkspace(
+            Dimensions="3",
+            Extents="-5,5,-5,5,-5,5",
+            Names="H,K,L",
+            Units="r.l.u.,r.l.u.,r.l.u.",
+            Frames="HKL,HKL,HKL",
+            SplitInto="2",
+            SplitThreshold="50",
+        )
+        # generate fake data at (0,1,0)
+        FakeMDEventData(ws, EllipsoidParams="1e+03,0,1,0,1,0,0,0,1,0,0,0,1,0.01,0.04,0.02,1", RandomSeed="3873875")
+        # create peak at (0,1,0) and integrate
+        peaks = self._make_peaks_HKL(hs=[0], ks=[1], ls=[0], wsname="peaks_md")
+        peaks_int = IntegratePeaksMD(
+            InputWorkspace=ws,
+            PeakRadius=0.6,
+            BackgroundInnerRadius=0.6,
+            BackgroundOuterRadius=0.8,
+            PeaksWorkspace=peaks,
+            OutputWorkspace="peaks_int_sphere",
+            Ellipsoid=False,
+            UseOnePercentBackgroundCorrection=False,
+        )
+
+        # plot
+        out_file = path.join(self._test_dir, "out_plot_MD.pdf")
+        BaseSX.plot_integrated_peaks_MD(ws, peaks_int, out_file, nbins_max=21, extent=1.5, log_norm=True)
+
+        # check output file saved
+        self.assertTrue(path.exists(out_file))
+
     # --- helper funcs ---
 
     def _make_peaks_HKL(self, hs=None, ks=[0], ls=[1], wsname="peaks_hkl"):