nxs-analysis-tools 0.0.46__tar.gz → 0.0.47__tar.gz

{nxs_analysis_tools-0.0.46/src/nxs_analysis_tools.egg-info → nxs_analysis_tools-0.0.47}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nxs-analysis-tools
-Version: 0.0.46
+Version: 0.0.47
 Summary: Reduce and transform nexus format (.nxs) scattering data.
 Author-email: "Steven J. Gomez Alvarado" <stevenjgomez@ucsb.edu>
 License: MIT License

{nxs_analysis_tools-0.0.46 → nxs_analysis_tools-0.0.47}/pyproject.toml

@@ -6,7 +6,7 @@ build-backend = 'setuptools.build_meta'
 
 [project]
 name = 'nxs-analysis-tools'
-version = '0.0.46'
+version = '0.0.47'
 description = 'Reduce and transform nexus format (.nxs) scattering data.'
 readme = 'README.md'
 requires-python = '>=3.7'

{nxs_analysis_tools-0.0.46 → nxs_analysis_tools-0.0.47}/src/_meta/__init__.py

@@ -6,5 +6,5 @@ __author__ = 'Steven J. Gomez Alvarado'
 __email__ = 'stevenjgomez@ucsb.edu'
 __copyright__ = f"2023-2025, {__author__}"
 __license__ = 'MIT'
-__version__ = '0.0.46'
+__version__ = '0.0.47'
 __repo_url__ = 'https://github.com/stevenjgomez/nxs_analysis_tools'

{nxs_analysis_tools-0.0.46 → nxs_analysis_tools-0.0.47}/src/nxs_analysis_tools/datareduction.py

@@ -16,7 +16,7 @@ from scipy import ndimage
 # Specify items on which users are allowed to perform standalone imports
 __all__ = ['load_data', 'load_transform', 'plot_slice', 'Scissors',
            'reciprocal_lattice_params', 'rotate_data', 'rotate_data_2D'
-           'array_to_nxdata', 'Padder']
+                                                       'array_to_nxdata', 'Padder']
 
 
 def load_data(path, print_tree=True):
@@ -107,7 +107,7 @@ def array_to_nxdata(array, data_template, signal_name=None):
                   tuple(d[d.axes[i]] for i in range(len(d.axes))))
 
 
-def plot_slice(data, X=None, Y=None, transpose=False, vmin=None, vmax=None,
+def plot_slice(data, X=None, Y=None, sum_axis=None, transpose=False, vmin=None, vmax=None,
                skew_angle=90, ax=None, xlim=None, ylim=None,
                xticks=None, yticks=None, cbar=True, logscale=False,
                symlogscale=False, cmap='viridis', linthresh=1,
@@ -130,6 +130,10 @@ def plot_slice(data, X=None, Y=None, transpose=False, vmin=None, vmax=None,
         The Y axis values. If None, a default range from 0 to the number of
          rows in `data` is used.
 
+    sum_axis : int, optional
+        If the input data is 3D, this specifies the axis to sum over in order
+        to reduce the data to 2D for plotting. Required if `data` has three dimensions.
+
     transpose : bool, optional
         If True, transpose the dataset and its axes before plotting.
         Default is False.
@@ -199,7 +203,37 @@ def plot_slice(data, X=None, Y=None, transpose=False, vmin=None, vmax=None,
     p : :class:`matplotlib.collections.QuadMesh`
         The `matplotlib` QuadMesh object representing the plotted data.
     """
+    is_array = False
+    is_nxdata = False
+
     if isinstance(data, np.ndarray):
+        is_array = True
+    elif isinstance(data, (NXdata, NXfield)):
+        is_nxdata = True
+    else:
+        raise TypeError(f"Unexpected data type: {type(data)}. "
+                        f"Supported types are np.ndarray and NXdata.")
+
+    # If three-dimensional, demand sum_axis to reduce to two dimensions.
+    if is_array and len(data.shape) == 3:
+        assert sum_axis is not None, "sum_axis must be specified when data is a 3D array"
+
+        data = data.sum(axis=sum_axis)
+
+    if is_nxdata and len(data.shape) == 3:
+        assert sum_axis is not None, "sum_axis must be specified when data is a 3D array"
+
+        arr = data.nxsignal.nxdata
+        arr = arr.sum(axis=sum_axis)
+
+        # Create a 2D template from the original nxdata
+        slice_obj = [slice(None)] * len(data.shape)
+        slice_obj[sum_axis] = 0
+
+        # Use the 2D template to create a new nxdata
+        data = array_to_nxdata(arr, data[slice_obj])
+
+    if is_array:
         if X is None:
             X = NXfield(np.linspace(0, data.shape[0], data.shape[0]), name='x')
         if Y is None:
@@ -209,7 +243,7 @@ def plot_slice(data, X=None, Y=None, transpose=False, vmin=None, vmax=None,
             data = data.transpose()
         data = NXdata(NXfield(data, name='value'), (X, Y))
         data_arr = data[data.signal].nxdata.transpose()
-    elif isinstance(data, (NXdata, NXfield)):
+    elif is_nxdata:
         if X is None:
             X = data[data.axes[0]]
         if Y is None:
@@ -218,9 +252,6 @@ def plot_slice(data, X=None, Y=None, transpose=False, vmin=None, vmax=None,
             X, Y = Y, X
             data = data.transpose()
         data_arr = data[data.signal].nxdata.transpose()
-    else:
-        raise TypeError(f"Unexpected data type: {type(data)}. "
-                        f"Supported types are np.ndarray and NXdata.")
 
     # Display Markdown heading
     if mdheading is None:
@@ -559,22 +590,31 @@ class Scissors:
 
         return self.linecut
 
-    def highlight_integration_window(self, data=None, label=None, highlight_color='red', **kwargs):
+    def highlight_integration_window(self, data=None, width=None, height=None, label=None, highlight_color='red',
+                                     **kwargs):
         """
-        Plots integration window highlighted on the three principal cross sections of the first
-        temperature dataset.
+        Plots the integration window highlighted on the three principal 2D cross-sections of a 3D dataset.
 
         Parameters
         ----------
         data : array-like, optional
-            The 2D heatmap dataset to plot. If not provided, the dataset stored in `self.data` will
-            be used.
+            The 3D dataset to visualize. If not provided, uses `self.data`.
+        width : float, optional
+            Width of the visible x-axis range in each subplot. Used to zoom in on the integration region.
+        height : float, optional
+            Height of the visible y-axis range in each subplot. Used to zoom in on the integration region.
         label : str, optional
-            The label for the integration window plot.
+            Label for the rectangle patch marking the integration window, used in the legend.
         highlight_color : str, optional
-            The edge color used to highlight the integration window. Default is 'red'.
-        **kwargs : keyword arguments, optional
-            Additional keyword arguments to customize the plot.
+            Color of the rectangle edges highlighting the integration window. Default is 'red'.
+        **kwargs : dict, optional
+            Additional keyword arguments passed to `plot_slice` for customizing the plot (e.g., colormap, vmin, vmax).
+
+        Returns
+        -------
+        p1, p2, p3 : matplotlib.collections.QuadMesh
+            The plotted QuadMesh objects for the three cross-sections:
+            XY at fixed Z, XZ at fixed Y, and YZ at fixed X.
 
         """
         data = self.data if data is None else data
@@ -584,7 +624,7 @@ class Scissors:
         # Create a figure and subplots
         fig, axes = plt.subplots(1, 3, figsize=(15, 4))
 
-        # Plot cross section 1
+        # Plot cross-section 1
         slice_obj = [slice(None)] * data.ndim
         slice_obj[2] = center[2]
 
@@ -603,7 +643,12 @@ class Scissors:
         )
         ax.add_patch(rect_diffuse)
 
-        # Plot cross section 2
+        if 'xlim' not in kwargs and width is not None:
+            ax.set(xlim=(center[0] - width / 2, center[0] + width / 2))
+        if 'ylim' not in kwargs and height is not None:
+            ax.set(ylim=(center[1] - height / 2, center[1] + height / 2))
+
+        # Plot cross-section 2
         slice_obj = [slice(None)] * data.ndim
         slice_obj[1] = center[1]
 
@@ -622,7 +667,12 @@ class Scissors:
         )
         ax.add_patch(rect_diffuse)
 
-        # Plot cross section 3
+        if 'xlim' not in kwargs and width is not None:
+            ax.set(xlim=(center[0] - width / 2, center[0] + width / 2))
+        if 'ylim' not in kwargs and height is not None:
+            ax.set(ylim=(center[2] - height / 2, center[2] + height / 2))
+
+        # Plot cross-section 3
         slice_obj = [slice(None)] * data.ndim
         slice_obj[0] = center[0]
 
@@ -641,6 +691,11 @@ class Scissors:
         )
         ax.add_patch(rect_diffuse)
 
+        if 'xlim' not in kwargs and width is not None:
+            ax.set(xlim=(center[1] - width / 2, center[1] + width / 2))
+        if 'ylim' not in kwargs and height is not None:
+            ax.set(ylim=(center[2] - height / 2, center[2] + height / 2))
+
         # Adjust subplot padding
         fig.subplots_adjust(wspace=0.5)
 
@@ -665,7 +720,7 @@ class Scissors:
 
         fig, axes = plt.subplots(1, 3, figsize=(15, 4))
 
-        # Plot cross section 1
+        # Plot cross-section 1
         slice_obj = [slice(None)] * data.ndim
         slice_obj[2] = center[2]
         p1 = plot_slice(data[slice_obj],
@@ -685,7 +740,7 @@ class Scissors:
                         **kwargs)
         axes[1].set_aspect(len(data[data.axes[0]].nxdata) / len(data[data.axes[2]].nxdata))
 
-        # Plot cross section 3
+        # Plot cross-section 3
         slice_obj = [slice(None)] * data.ndim
         slice_obj[0] = center[0]
         p2 = plot_slice(data[slice_obj],
@@ -1015,7 +1070,7 @@ class Padder:
         self.data = data
 
         self.steps = tuple((data[axis].nxdata[1] - data[axis].nxdata[0])
-                            for axis in data.axes)
+                           for axis in data.axes)
 
         # Absolute value of the maximum value; assumes the domain of the input
         # is symmetric (eg, -H_min = H_max)

{nxs_analysis_tools-0.0.46 → nxs_analysis_tools-0.0.47/src/nxs_analysis_tools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nxs-analysis-tools
-Version: 0.0.46
+Version: 0.0.47
 Summary: Reduce and transform nexus format (.nxs) scattering data.
 Author-email: "Steven J. Gomez Alvarado" <stevenjgomez@ucsb.edu>
 License: MIT License

{nxs_analysis_tools-0.0.46 → nxs_analysis_tools-0.0.47}/src/nxs_analysis_tools.egg-info/SOURCES.txt

@@ -14,6 +14,7 @@ src/nxs_analysis_tools.egg-info/SOURCES.txt
 src/nxs_analysis_tools.egg-info/dependency_links.txt
 src/nxs_analysis_tools.egg-info/requires.txt
 src/nxs_analysis_tools.egg-info/top_level.txt
+tests/test_accurate_highlight.py
 tests/test_chess.py
 tests/test_chess_fitting.py
 tests/test_datareduction.py
@@ -22,4 +23,5 @@ tests/test_lmfit.py
 tests/test_mask_plotting.py
 tests/test_pairdistribution.py
 tests/test_plot_slice_with_ndarray.py
+tests/test_sum_axis.py
 tests/test_symmetrizer_rectangular_plane.py

nxs_analysis_tools-0.0.47/tests/test_accurate_highlight.py

@@ -0,0 +1,388 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.transforms import Affine2D
+from matplotlib.markers import MarkerStyle
+from matplotlib.ticker import MultipleLocator
+from matplotlib import colors
+from matplotlib import patches
+from IPython.display import display, Markdown
+from nexusformat.nexus import NXfield, NXdata, nxload, NeXusError, NXroot, NXentry, nxsave
+from scipy import ndimage
+from nxs_analysis_tools.datareduction import plot_slice, load_data
+
+
+class Scissors:
+    """
+    Scissors class provides functionality for reducing data to a 1D linecut using an integration
+    window.
+
+    Attributes
+    ----------
+    data : :class:`nexusformat.nexus.NXdata` or None
+        Input :class:`nexusformat.nexus.NXdata`.
+    center : tuple or None
+        Central coordinate around which to perform the linecut.
+    window : tuple or None
+        Extents of the window for integration along each axis.
+    axis : int or None
+        Axis along which to perform the integration.
+    integration_volume : :class:`nexusformat.nexus.NXdata` or None
+        Data array after applying the integration window.
+    integrated_axes : tuple or None
+        Indices of axes that were integrated.
+    linecut : :class:`nexusformat.nexus.NXdata` or None
+        1D linecut data after integration.
+    integration_window : tuple or None
+        Slice object representing the integration window in the data array.
+
+    Methods
+    -------
+    set_data(data)
+        Set the input :class:`nexusformat.nexus.NXdata`.
+    get_data()
+        Get the input :class:`nexusformat.nexus.NXdata`.
+    set_center(center)
+        Set the central coordinate for the linecut.
+    set_window(window, axis=None, verbose=False)
+        Set the extents of the integration window.
+    get_window()
+        Get the extents of the integration window.
+    cut_data(center=None, window=None, axis=None, verbose=False)
+        Reduce data to a 1D linecut using the integration window.
+    highlight_integration_window(data=None, label=None, highlight_color='red', **kwargs)
+        Plot the integration window highlighted on a 2D heatmap of the full dataset.
+    plot_integration_window(**kwargs)
+        Plot a 2D heatmap of the integration window data.
+    """
+
+    def __init__(self, data=None, center=None, window=None, axis=None):
+        """
+        Initializes a Scissors object.
+
+        Parameters
+        ----------
+        data : :class:`nexusformat.nexus.NXdata` or None, optional
+            Input NXdata. Default is None.
+        center : tuple or None, optional
+            Central coordinate around which to perform the linecut. Default is None.
+        window : tuple or None, optional
+            Extents of the window for integration along each axis. Default is None.
+        axis : int or None, optional
+            Axis along which to perform the integration. Default is None.
+        """
+
+        self.data = data
+        self.center = tuple(float(i) for i in center) if center is not None else None
+        self.window = tuple(float(i) for i in window) if window is not None else None
+        self.axis = axis
+
+        self.integration_volume = None
+        self.integrated_axes = None
+        self.linecut = None
+        self.integration_window = None
+
+    def set_data(self, data):
+        """
+        Set the input NXdata.
+
+        Parameters
+        ----------
+        data : :class:`nexusformat.nexus.NXdata`
+            Input data array.
+        """
+        self.data = data
+
+    def get_data(self):
+        """
+        Get the input data array.
+
+        Returns
+        -------
+        ndarray or None
+            Input data array.
+        """
+        return self.data
+
+    def set_center(self, center):
+        """
+        Set the central coordinate for the linecut.
+
+        Parameters
+        ----------
+        center : tuple
+            Central coordinate around which to perform the linecut.
+        """
+        self.center = tuple(float(i) for i in center) if center is not None else None
+
+    def set_window(self, window, axis=None, verbose=False):
+        """
+        Set the extents of the integration window.
+
+        Parameters
+        ----------
+        window : tuple
+            Extents of the window for integration along each axis.
+        axis : int or None, optional
+            The axis along which to perform the linecut. If not specified, the value from the
+            object's attribute will be used.
+        verbose : bool, optional
+            Enables printout of linecut axis and integrated axes. Default False.
+
+        """
+        self.window = tuple(float(i) for i in window) if window is not None else None
+
+        # Determine the axis for integration
+        self.axis = window.index(max(window)) if axis is None else axis
+
+        # Determine the integrated axes (axes other than the integration axis)
+        self.integrated_axes = tuple(i for i in range(self.data.ndim) if i != self.axis)
+
+        if verbose:
+            print("Linecut axis: " + str(self.data.axes[self.axis]))
+            print("Integrated axes: " + str([self.data.axes[axis]
+                                             for axis in self.integrated_axes]))
+
+    def get_window(self):
+        """
+        Get the extents of the integration window.
+
+        Returns
+        -------
+        tuple or None
+            Extents of the integration window.
+        """
+        return self.window
+
+    def cut_data(self, center=None, window=None, axis=None, verbose=False):
+        """
+        Reduces data to a 1D linecut with integration extents specified by the
+        window about a central coordinate.
+
+        Parameters
+        ----------
+        center : float or None, optional
+            Central coordinate for the linecut. If not specified, the value from the object's
+            attribute will be used.
+        window : tuple or None, optional
+            Integration window extents around the central coordinate. If not specified, the value
+            from the object's attribute will be used.
+        axis : int or None, optional
+            The axis along which to perform the linecut. If not specified, the value from the
+            object's attribute will be used.
+        verbose : bool
+            Enables printout of linecut axis and integrated axes. Default False.
+
+        Returns
+        -------
+        integrated_data : :class:`nexusformat.nexus.NXdata`
+            1D linecut data after integration.
+
+        """
+
+        # Extract necessary attributes from the object
+        data = self.data
+        center = center if center is not None else self.center
+        self.set_center(center)
+        window = window if window is not None else self.window
+        self.set_window(window, axis, verbose)
+
+        # Convert the center to a tuple of floats
+        center = tuple(float(c) for c in center)
+
+        # Calculate the start and stop indices for slicing the data
+        start = np.subtract(center, window)
+        stop = np.add(center, window)
+        slice_obj = tuple(slice(s, e) for s, e in zip(start, stop))
+        self.integration_window = slice_obj
+
+        # Perform the data cut
+        self.integration_volume = data[slice_obj]
+        self.integration_volume.nxname = data.nxname
+
+        # Perform integration along the integrated axes
+        integrated_data = np.sum(self.integration_volume[self.integration_volume.signal].nxdata,
+                                 axis=self.integrated_axes)
+
+        # Create an NXdata object for the linecut data
+        self.linecut = NXdata(NXfield(integrated_data, name=self.integration_volume.signal),
+                              self.integration_volume[self.integration_volume.axes[self.axis]])
+        self.linecut.nxname = self.integration_volume.nxname
+
+        return self.linecut
+
+    def highlight_integration_window(self, data=None, width=None, height=None, label=None, highlight_color='red', **kwargs):
+        """
+        Plots the integration window highlighted on the three principal 2D cross-sections of a 3D dataset.
+
+        Parameters
+        ----------
+        data : array-like, optional
+            The 3D dataset to visualize. If not provided, uses `self.data`.
+        width : float, optional
+            Width of the visible x-axis range in each subplot. Used to zoom in on the integration region.
+        height : float, optional
+            Height of the visible y-axis range in each subplot. Used to zoom in on the integration region.
+        label : str, optional
+            Label for the rectangle patch marking the integration window, used in the legend.
+        highlight_color : str, optional
+            Color of the rectangle edges highlighting the integration window. Default is 'red'.
+        **kwargs : dict, optional
+            Additional keyword arguments passed to `plot_slice` for customizing the plot (e.g., colormap, vmin, vmax).
+
+        Returns
+        -------
+        p1, p2, p3 : matplotlib.collections.QuadMesh
+            The plotted QuadMesh objects for the three cross-sections:
+            XY at fixed Z, XZ at fixed Y, and YZ at fixed X.
+
+        """
+        data = self.data if data is None else data
+        center = self.center
+        window = self.window
+
+        # Create a figure and subplots
+        fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+        # Plot cross-section 1
+        slice_obj = [slice(None)] * data.ndim
+        slice_obj[2] = center[2]
+
+        p1 = plot_slice(data[slice_obj],
+                        X=data[data.axes[0]],
+                        Y=data[data.axes[1]],
+                        ax=axes[0],
+                        **kwargs)
+        ax = axes[0]
+        rect_diffuse = patches.Rectangle(
+            (center[0] - window[0],
+             center[1] - window[1]),
+            2 * window[0], 2 * window[1],
+            linewidth=1, edgecolor=highlight_color,
+            facecolor='none', transform=p1.get_transform(), label=label,
+        )
+        ax.add_patch(rect_diffuse)
+
+        if 'xlim' not in kwargs and width is not None:
+            ax.set(xlim=(center[0]-width/2,center[0]+width/2))
+        if 'ylim' not in kwargs and height is not None:
+            ax.set(ylim=(center[1]-height/2,center[1]+height/2))
+
+        # Plot cross-section 2
+        slice_obj = [slice(None)] * data.ndim
+        slice_obj[1] = center[1]
+
+        p2 = plot_slice(data[slice_obj],
+                        X=data[data.axes[0]],
+                        Y=data[data.axes[2]],
+                        ax=axes[1],
+                        **kwargs)
+        ax = axes[1]
+        rect_diffuse = patches.Rectangle(
+            (center[0] - window[0],
+             center[2] - window[2]),
+            2 * window[0], 2 * window[2],
+            linewidth=1, edgecolor=highlight_color,
+            facecolor='none', transform=p2.get_transform(), label=label,
+        )
+        ax.add_patch(rect_diffuse)
+
+        if 'xlim' not in kwargs and width is not None:
+            ax.set(xlim=(center[0]-width/2,center[0]+width/2))
+        if 'ylim' not in kwargs and height is not None:
+            ax.set(ylim=(center[2]-height/2,center[2]+height/2))
+
+
+        # Plot cross-section 3
+        slice_obj = [slice(None)] * data.ndim
+        slice_obj[0] = center[0]
+
+        p3 = plot_slice(data[slice_obj],
+                        X=data[data.axes[1]],
+                        Y=data[data.axes[2]],
+                        ax=axes[2],
+                        **kwargs)
+        ax = axes[2]
+        rect_diffuse = patches.Rectangle(
+            (center[1] - window[1],
+             center[2] - window[2]),
+            2 * window[1], 2 * window[2],
+            linewidth=1, edgecolor=highlight_color,
+            facecolor='none', transform=p3.get_transform(), label=label,
+        )
+        ax.add_patch(rect_diffuse)
+
+        if 'xlim' not in kwargs and width is not None:
+            ax.set(xlim=(center[1]-width/2,center[1]+width/2))
+        if 'ylim' not in kwargs and height is not None:
+            ax.set(ylim=(center[2]-height/2,center[2]+height/2))
+
+
+        # Adjust subplot padding
+        fig.subplots_adjust(wspace=0.5)
+
+        if label is not None:
+            [ax.legend() for ax in axes]
+
+        plt.show()
+
+        return p1, p2, p3
+
+    def plot_integration_window(self, **kwargs):
+        """
+        Plots the three principal cross-sections of the integration volume on a single figure.
+
+        Parameters
+        ----------
+        **kwargs : keyword arguments, optional
+            Additional keyword arguments to customize the plot.
+        """
+        data = self.integration_volume
+        center = self.center
+
+        fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+        # Plot cross section 1
+        slice_obj = [slice(None)] * data.ndim
+        slice_obj[2] = center[2]
+        p1 = plot_slice(data[slice_obj],
+                        X=data[data.axes[0]],
+                        Y=data[data.axes[1]],
+                        ax=axes[0],
+                        **kwargs)
+        axes[0].set_aspect(len(data[data.axes[0]].nxdata) / len(data[data.axes[1]].nxdata))
+
+        # Plot cross section 2
+        slice_obj = [slice(None)] * data.ndim
+        slice_obj[1] = center[1]
+        p3 = plot_slice(data[slice_obj],
+                        X=data[data.axes[0]],
+                        Y=data[data.axes[2]],
+                        ax=axes[1],
+                        **kwargs)
+        axes[1].set_aspect(len(data[data.axes[0]].nxdata) / len(data[data.axes[2]].nxdata))
+
+        # Plot cross section 3
+        slice_obj = [slice(None)] * data.ndim
+        slice_obj[0] = center[0]
+        p2 = plot_slice(data[slice_obj],
+                        X=data[data.axes[1]],
+                        Y=data[data.axes[2]],
+                        ax=axes[2],
+                        **kwargs)
+        axes[2].set_aspect(len(data[data.axes[1]].nxdata) / len(data[data.axes[2]].nxdata))
+
+        # Adjust subplot padding
+        fig.subplots_adjust(wspace=0.3)
+
+        plt.show()
+
+        return p1, p2, p3
+
+
+data = load_data(r'C:\Users\steve\OneDrive\Documents\UCSB\Projects\RECd3P3\mini\15\3rot_hkli.nxs')
+s = Scissors(data)
+s.cut_data(center=(1,1,0), window=(0.5,0.2,0.2))
+print(s.highlight_integration_window(vmin=0, vmax=100, width=10, height=10))
+
+print(data.shape)

nxs_analysis_tools-0.0.47/tests/test_sum_axis.py

@@ -0,0 +1,299 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.transforms import Affine2D
+from matplotlib.markers import MarkerStyle
+from matplotlib.ticker import MultipleLocator
+from matplotlib import colors
+from matplotlib import patches
+from IPython.display import display, Markdown
+from nexusformat.nexus import NXfield, NXdata, nxload, NeXusError, NXroot, NXentry, nxsave
+from scipy import ndimage
+from nxs_analysis_tools.datareduction import load_data, array_to_nxdata, plot_slice
+
+# def plot_slice(data, X=None, Y=None, sum_axis=None, transpose=False, vmin=None, vmax=None,
+#                skew_angle=90, ax=None, xlim=None, ylim=None,
+#                xticks=None, yticks=None, cbar=True, logscale=False,
+#                symlogscale=False, cmap='viridis', linthresh=1,
+#                title=None, mdheading=None, cbartitle=None,
+#                **kwargs):
+#     """
+#     Plot a 2D slice of the provided dataset, with optional transformations
+#     and customizations.
+#
+#     Parameters
+#     ----------
+#     data : :class:`nexusformat.nexus.NXdata` or ndarray
+#         The dataset to plot. Can be an `NXdata` object or a `numpy` array.
+#
+#     X : NXfield, optional
+#         The X axis values. If None, a default range from 0 to the number of
+#          columns in `data` is used.
+#
+#     Y : NXfield, optional
+#         The Y axis values. If None, a default range from 0 to the number of
+#          rows in `data` is used.
+#
+#     sum_axis : int, optional
+#         If the input data is 3D, this specifies the axis to sum over in order
+#         to reduce the data to 2D for plotting. Required if `data` has three dimensions.
+#
+#     transpose : bool, optional
+#         If True, transpose the dataset and its axes before plotting.
+#         Default is False.
+#
+#     vmin : float, optional
+#         The minimum value for the color scale. If not provided, the minimum
+#          value of the dataset is used.
+#
+#     vmax : float, optional
+#         The maximum value for the color scale. If not provided, the maximum
+#          value of the dataset is used.
+#
+#     skew_angle : float, optional
+#         The angle in degrees to shear the plot. Default is 90 degrees (no skew).
+#
+#     ax : matplotlib.axes.Axes, optional
+#         The `matplotlib` axis to plot on. If None, a new figure and axis will
+#          be created.
+#
+#     xlim : tuple, optional
+#         The limits for the x-axis. If None, the limits are set automatically
+#          based on the data.
+#
+#     ylim : tuple, optional
+#         The limits for the y-axis. If None, the limits are set automatically
+#          based on the data.
+#
+#     xticks : float or list of float, optional
+#         The major tick interval or specific tick locations for the x-axis.
+#          Default is to use a minor tick interval of 1.
+#
+#     yticks : float or list of float, optional
+#         The major tick interval or specific tick locations for the y-axis.
+#         Default is to use a minor tick interval of 1.
+#
+#     cbar : bool, optional
+#         Whether to include a colorbar. Default is True.
+#
+#     logscale : bool, optional
+#         Whether to use a logarithmic color scale. Default is False.
+#
+#     symlogscale : bool, optional
+#         Whether to use a symmetrical logarithmic color scale. Default is False.
+#
+#     cmap : str or Colormap, optional
+#         The colormap to use for the plot. Default is 'viridis'.
+#
+#     linthresh : float, optional
+#         The linear threshold for symmetrical logarithmic scaling. Default is 1.
+#
+#     title : str, optional
+#         The title for the plot. If None, no title is set.
+#
+#     mdheading : str, optional
+#         A Markdown heading to display above the plot. If 'None' or not provided,
+#          no heading is displayed.
+#
+#     cbartitle : str, optional
+#         The title for the colorbar. If None, the colorbar label will be set to
+#          the name of the signal.
+#
+#     **kwargs
+#         Additional keyword arguments passed to `pcolormesh`.
+#
+#     Returns
+#     -------
+#     p : :class:`matplotlib.collections.QuadMesh`
+#         The `matplotlib` QuadMesh object representing the plotted data.
+#     """
+#     is_array = False
+#     is_nxdata = False
+#
+#     if isinstance(data, np.ndarray):
+#         is_array = True
+#     elif isinstance(data, (NXdata, NXfield)):
+#         is_nxdata = True
+#     else:
+#         raise TypeError(f"Unexpected data type: {type(data)}. "
+#                         f"Supported types are np.ndarray and NXdata.")
+#
+#     # If three-dimensional, demand sum_axis to reduce to two dimensions.
+#     if is_array and len(data.shape) == 3:
+#         assert sum_axis is not None, "sum_axis must be specified when data is a 3D array"
+#
+#         data = data.sum(axis=sum_axis)
+#
+#     if is_nxdata and len(data.shape) == 3:
+#         assert sum_axis is not None, "sum_axis must be specified when data is a 3D array"
+#
+#         arr = data.nxsignal.nxdata
+#         arr = arr.sum(axis=sum_axis)
+#
+#         # Create a 2D template from the original nxdata
+#         slice_obj = [slice(None)] * len(data.shape)
+#         slice_obj[sum_axis] = 0
+#
+#         # Use the 2D template to create a new nxdata
+#         data = array_to_nxdata(arr, data[slice_obj])
+#
+#     if is_array:
+#         if X is None:
+#             X = NXfield(np.linspace(0, data.shape[0], data.shape[0]), name='x')
+#         if Y is None:
+#             Y = NXfield(np.linspace(0, data.shape[1], data.shape[1]), name='y')
+#         if transpose:
+#             X, Y = Y, X
+#             data = data.transpose()
+#         data = NXdata(NXfield(data, name='value'), (X, Y))
+#         data_arr = data[data.signal].nxdata.transpose()
+#     elif is_nxdata:
+#         if X is None:
+#             X = data[data.axes[0]]
+#         if Y is None:
+#             Y = data[data.axes[1]]
+#         if transpose:
+#             X, Y = Y, X
+#             data = data.transpose()
+#         data_arr = data[data.signal].nxdata.transpose()
+#
+#     # Display Markdown heading
+#     if mdheading is None:
+#         pass
+#     elif mdheading == "None":
+#         display(Markdown('### Figure'))
+#     else:
+#         display(Markdown('### Figure - ' + mdheading))
+#
+#     # Inherit axes if user provides some
+#     if ax is not None:
+#         fig = ax.get_figure()
+#     # Otherwise set up some default axes
+#     else:
+#         fig = plt.figure()
+#         ax = fig.add_axes([0, 0, 1, 1])
+#
+#     # If limits not provided, use extrema
+#     if vmin is None:
+#         vmin = data_arr.min()
+#     if vmax is None:
+#         vmax = data_arr.max()
+#
+#     # Set norm (linear scale, logscale, or symlogscale)
+#     norm = colors.Normalize(vmin=vmin, vmax=vmax)  # Default: linear scale
+#
+#     if symlogscale:
+#         norm = colors.SymLogNorm(linthresh=linthresh, vmin=-1 * vmax, vmax=vmax)
+#     elif logscale:
+#         norm = colors.LogNorm(vmin=vmin, vmax=vmax)
+#
+#     # Plot data
+#     p = ax.pcolormesh(X.nxdata, Y.nxdata, data_arr, shading='auto', norm=norm, cmap=cmap, **kwargs)
+#
+#     ## Transform data to new coordinate system if necessary
+#     # Correct skew angle
+#     skew_angle_adj = 90 - skew_angle
+#     # Create blank 2D affine transformation
+#     t = Affine2D()
+#     # Scale y-axis to preserve norm while shearing
+#     t += Affine2D().scale(1, np.cos(skew_angle_adj * np.pi / 180))
+#     # Shear along x-axis
+#     t += Affine2D().skew_deg(skew_angle_adj, 0)
+#     # Return to original y-axis scaling
+#     t += Affine2D().scale(1, np.cos(skew_angle_adj * np.pi / 180)).inverted()
+#     ## Correct for x-displacement after shearing
+#     # If ylims provided, use those
+#     if ylim is not None:
+#         # Set ylims
+#         ax.set(ylim=ylim)
+#         ymin, ymax = ylim
+#     # Else, use current ylims
+#     else:
+#         ymin, ymax = ax.get_ylim()
+#     # Use ylims to calculate translation (necessary to display axes in correct position)
+#     p.set_transform(t
+#                     + Affine2D().translate(-ymin * np.sin(skew_angle_adj * np.pi / 180), 0)
+#                     + ax.transData)
+#
+#     # Set x limits
+#     if xlim is not None:
+#         xmin, xmax = xlim
+#     else:
+#         xmin, xmax = ax.get_xlim()
+#     if skew_angle <= 90:
+#         ax.set(xlim=(xmin, xmax + (ymax - ymin) / np.tan((90 - skew_angle_adj) * np.pi / 180)))
+#     else:
+#         ax.set(xlim=(xmin - (ymax - ymin) / np.tan((skew_angle_adj - 90) * np.pi / 180), xmax))
+#
+#     # Correct aspect ratio for the x/y axes after transformation
+#     ax.set(aspect=np.cos(skew_angle_adj * np.pi / 180))
+#
+#     # Add tick marks all around
+#     ax.tick_params(direction='in', top=True, right=True, which='both')
+#
+#     # Set tick locations
+#     if xticks is None:
+#         # Add default minor ticks
+#         ax.xaxis.set_minor_locator(MultipleLocator(1))
+#     else:
+#         # Otherwise use user provided values
+#         ax.xaxis.set_major_locator(MultipleLocator(xticks))
+#         ax.xaxis.set_minor_locator(MultipleLocator(1))
+#     if yticks is None:
+#         # Add default minor ticks
+#         ax.yaxis.set_minor_locator(MultipleLocator(1))
+#     else:
+#         # Otherwise use user provided values
+#         ax.yaxis.set_major_locator(MultipleLocator(yticks))
+#         ax.yaxis.set_minor_locator(MultipleLocator(1))
+#
+#     # Apply transform to tick marks
+#     for i in range(0, len(ax.xaxis.get_ticklines())):
+#         # Tick marker
+#         m = MarkerStyle(3)
+#         line = ax.xaxis.get_majorticklines()[i]
+#         if i % 2:
+#             # Top ticks (translation here makes their direction="in")
+#             m._transform.set(Affine2D().translate(0, -1) + Affine2D().skew_deg(skew_angle_adj, 0))
+#             # This first method shifts the top ticks horizontally to match the skew angle.
+#             # This does not look good in all cases.
+#             # line.set_transform(Affine2D().translate((ymax-ymin)*np.sin(skew_angle*np.pi/180),0) +
+#             #     line.get_transform())
+#             # This second method skews the tick marks in place and
+#             # can sometimes lead to them being misaligned.
+#             line.set_transform(line.get_transform())  # This does nothing
+#         else:
+#             # Bottom ticks
+#             m._transform.set(Affine2D().skew_deg(skew_angle_adj, 0))
+#
+#         line.set_marker(m)
+#
+#     for i in range(0, len(ax.xaxis.get_minorticklines())):
+#         m = MarkerStyle(2)
+#         line = ax.xaxis.get_minorticklines()[i]
+#         if i % 2:
+#             m._transform.set(Affine2D().translate(0, -1) + Affine2D().skew_deg(skew_angle_adj, 0))
+#         else:
+#             m._transform.set(Affine2D().skew_deg(skew_angle_adj, 0))
+#
+#         line.set_marker(m)
+#
+#     if cbar:
+#         colorbar = fig.colorbar(p)
+#         if cbartitle is None:
+#             colorbar.set_label(data.signal)
+#
+#     ax.set(
+#         xlabel=X.nxname,
+#         ylabel=Y.nxname,
+#     )
+#
+#     if title is not None:
+#         ax.set_title(title)
+#
+#     # Return the quadmesh object
+#     return p
+
+data = load_data(r'C:\Users\steve\OneDrive\Documents\UCSB\Projects\RECd3P3\mini\15\3rot_hkli.nxs')
+plot_slice(data[:,-0.2:0.2,:], vmin=0, vmax=100*10, sum_axis=1)
+plt.show()