diff --git a/docs/source/release/v6.10.0/Diffraction/Single_Crystal/New_features/36947.rst b/docs/source/release/v6.10.0/Diffraction/Single_Crystal/New_features/36947.rst new file mode 100644 index 000000000000..3eb064040ada --- /dev/null +++ b/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 \ No newline at end of file diff --git a/scripts/Diffraction/single_crystal/base_sx.py b/scripts/Diffraction/single_crystal/base_sx.py index 7c530cc6c602..809b2dfe2829 100644 --- a/scripts/Diffraction/single_crystal/base_sx.py +++ b/scripts/Diffraction/single_crystal/base_sx.py @@ -4,16 +4,20 @@ # NScD Oak Ridge National Laboratory, European Spallation Source # & Institut Laue - Langevin # SPDX - License - Identifier: GPL - 3.0 + -from typing import Sequence +from typing import Sequence, Optional import numpy as np from enum import Enum import mantid.simpleapi as mantid -from mantid.api import FunctionFactory, AnalysisDataService as ADS -from mantid.kernel import logger +from mantid.api import FunctionFactory, AnalysisDataService as ADS, IMDEventWorkspace, IPeaksWorkspace, IPeak +from mantid.kernel import logger, SpecialCoordinateSystem from FindGoniometerFromUB import getSignMaxAbsValInCol -from mantid.geometry import CrystalStructure, SpaceGroupFactory, ReflectionGenerator, ReflectionConditionFilter +from mantid.geometry import CrystalStructure, SpaceGroupFactory, ReflectionGenerator, ReflectionConditionFilter, PeakShape 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,194 @@ def remove_non_integrated_peaks(peaks): EnableLogging=False, ) + @staticmethod + def plot_integrated_peaks_MD( + wsMD: IMDEventWorkspace, + peaks: IPeaksWorkspace, + filename: str, + nbins_max: Optional[int] = 21, + extent: Optional[float] = 1.5, + log_norm: Optional[bool] = True, + ): + """ + Function to plot summary of IntegratePeaksMD integration comprising 3 colorfill plots per peak saved in a pdf. + :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: + 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 not np.allclose(bg_outer_radii, 0.0): + 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) + 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: IMDEventWorkspace, pk: IPeak, peak_shape: PeakShape, nbins_max: int, extent: float, frame_to_peak_centre_attr: str + ): + """ + Bin MD workspace in peak region with projection axes given by the axes of the ellipsoid shape + :param wsMD: MD workspace (with 3 dims) + :param pk: Peak object that has been integrated + :param peak_shape: shape of integrated peak (spherical, ellipsoid or none) + :param nbins_max: number of bins along the major axis of the ellipsoid + :param extent: extent of output MD workspace along each dimension in units of the outer background radius + :param frame_to_peak_centre_attr: name of attribute to return peak centre in appropriate frame + :return ws_cut: MD workspace in peak region + :return radii: radii of the ellipsoid peak region + :return bg_inner_radii: inner background radii of the ellipsoid peak region + :return bg_outer_radii: outer background radii of the ellipsoid peak region + :return box_lengths: length of each dimension in the MD workspace + :return: imax: index of dimension with largest radius + """ + ndims = wsMD.getNumDims() + radii, bg_inner_radii, bg_outer_radii, evecs, translation = BaseSX._get_ellipsoid_params_from_peak(peak_shape, ndims) + imax = np.argmax(radii) + # calc center in frame of ellipsoid axes + 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 not np.allclose(bg_outer_radii, 0.0) 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(ndims)]) + # ensure minimum of 3 bins inside radius along each dim + min_nbins_in_radius = 3 + nbins_in_radius = np.min(nbins * radii / box_lengths) # along most coarsely binned dimension + if nbins_in_radius < min_nbins_in_radius: + nbins = nbins * min_nbins_in_radius / nbins_in_radius + # call BinMD + ws_cut = mantid.BinMD( + InputWorkspace=wsMD, + 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, + StoreInADS=False, + ) + return ws_cut, radii, bg_inner_radii, bg_outer_radii, box_lengths, imax + + @staticmethod + def _get_ellipsoid_params_from_peak(peak_shape: PeakShape, ndims: int): + """ + Extract ellipsoid parameters (eigenvectors, radii etc.) from PeakShape object + :param peak_shape: PeakShape object of a integrated peak + :param ndims: number of dimensions in the MD workspace + :return radii: array of ellipsoid radii (3 sigma) defining peak region + :return bg_inner_radii: array of inner radii for ellipsoid background shell + :return bg_outer_radii: array of outer radii for ellipsoid background shell + :return evecs: ndims x ndims array of eignevectors - each column corresponds to an ellipsoid axis + :return translation: translation of the ellipsoid center in the coordinates/frame of the MD workspace integrated + """ + 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(ndims)]) + bg_inner_radii = np.array([shape_info[f"background_inner_radius{iax}"] for iax in range(ndims)]) + bg_outer_radii = np.array([shape_info[f"background_outer_radius{iax}"] for iax in range(ndims)]) + # eignevectors of ellipsoid + evecs = np.zeros((ndims, ndims)) + for iax in range(ndims): + evec = np.array([float(elem) for elem in shape_info[f"direction{iax}"].split()]) + evecs[:, iax] = evec / np.linalg.norm(evec) + # get translation in frame of wsMD + translation = np.array([shape_info[f"translation{iax}"] for iax in range(ndims)]) + return radii, bg_inner_radii, bg_outer_radii, evecs, translation + + @staticmethod + def _convert_spherical_representation_to_ellipsoid(shape_info: dict): + """ + Update shape_info dict of a spherical peak shape to have same keys/fields as an ellipsoid + :param shape_info: dictionary of JSON shape representation + """ + # 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"]: + shape_info[f"{key}{0}"] = shape_info.pop(key) if key in shape_info else 0.0 + 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): diff --git a/scripts/Diffraction/single_crystal/test/test_base_sx.py b/scripts/Diffraction/single_crystal/test/test_base_sx.py index 605728f90258..6b2b56b851a2 100644 --- a/scripts/Diffraction/single_crystal/test/test_base_sx.py +++ b/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"):