nxs-analysis-tools 0.0.35__py3-none-any.whl → 0.0.36__py3-none-any.whl

_meta/__init__.py

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

nxs_analysis_tools/__init__.py

@@ -4,7 +4,8 @@ Reduce and transform nexus format (.nxs) scattering data.
 
 import numpy as np
 from _meta import __author__, __copyright__, __license__, __version__
-from .datareduction import load_data, load_transform, plot_slice, reciprocal_lattice_params, Scissors, rotate_data
+from .datareduction import load_data, load_transform, plot_slice,\
+    reciprocal_lattice_params, Scissors, rotate_data
 from .chess import TempDependence
 
 # What to import when running "from nxs_analysis_tools import *"

nxs_analysis_tools/chess.py

@@ -4,11 +4,12 @@ This module provides classes and functions for analyzing scattering datasets col
 plotting linecuts.
 """
 import os
+import re
+
 import matplotlib.pyplot as plt
 import matplotlib as mpl
 import numpy as np
 from IPython.display import display, Markdown
-from nexusformat.nexus import nxload
 from nxs_analysis_tools import load_data, Scissors
 from nxs_analysis_tools.fitting import LinecutModel
 from nxs_analysis_tools.datareduction import load_transform, reciprocal_lattice_params
@@ -16,14 +17,93 @@ from nxs_analysis_tools.datareduction import load_transform, reciprocal_lattice_
 
 class TempDependence:
     """
-    Class for analyzing scattering datasets collected at CHESS (ID4B) with temperature dependence.
+    A class for analyzing temperature-dependent scattering datasets collected at CHESS (ID4B).
+
+    The `TempDependence` class facilitates the loading, processing, and analysis of scattering
+    data across different temperatures. It includes methods for handling datasets, setting
+    lattice parameters, performing linecuts, modeling the data, and visualizing the results.
+
+    Attributes
+    ----------
+    sample_directory : str
+        Path to the directory containing the datasets.
+    xlabel : str
+        Label for the x-axis of plots, determined by the axis of the linecuts.
+    datasets : dict
+        Dictionary storing datasets keyed by temperature.
+    temperatures : list of str
+        List of temperatures for which data is available.
+    scissors : dict
+        Dictionary of Scissors objects, one for each temperature, used for data manipulation and
+        linecut operations.
+    linecuts : dict
+        Dictionary storing the linecut data for each temperature.
+    linecutmodels : dict
+        Dictionary of LinecutModel objects, one for each temperature, used for fitting the linecuts.
+    a, b, c, al, be, ga : float or None
+        Lattice parameters (a, b, c, alpha, beta, gamma) of the crystal.
+    a_star, b_star, c_star, al_star, be_star, ga_star : float or None
+        Reciprocal lattice parameters (a*, b*, c*, alpha*, beta*, gamma*).
+
+    Methods
+    -------
+    set_temperatures(temperatures):
+        Set the list of temperatures for the datasets.
+    set_sample_directory(path):
+        Set the directory path where the datasets are located.
+    initialize():
+        Initialize Scissors and LinecutModel objects for each temperature.
+    set_data(temperature, data):
+        Set the dataset for a specific temperature.
+    load_transforms(temperatures_list=None):
+        Load transform datasets (from nxrefine) based on temperature.
+    load_datasets(file_ending='hkli.nxs', temperatures_list=None):
+        Load datasets (CHESS format) from the specified folder.
+    get_sample_directory():
+        Get the folder path where the datasets are located.
+    clear_datasets():
+        Clear the datasets stored in the TempDependence instance.
+    set_Lattice_params(lattice_params):
+        Set lattice parameters and calculate reciprocal lattice parameters.
+    set_window(window, verbose=False):
+        Set the extents of the integration window for each temperature.
+    set_center(center):
+        Set the central coordinate for the linecut for each temperature.
+    cut_data(center=None, window=None, axis=None, verbose=False):
+        Perform data cutting for each temperature dataset.
+    plot_linecuts(vertical_offset=0, **kwargs):
+        Plot the linecuts obtained from data cutting.
+    plot_linecuts_heatmap(ax=None, **kwargs):
+        Plot a heatmap of the linecuts obtained from data cutting.
+    highlight_integration_window(temperature=None, **kwargs):
+        Display the integration window plot for a specific temperature.
+    plot_integration_window(temperature=None, **kwargs):
+        Plot the integration window cross-sections for a specific temperature.
+    set_model_components(model_components):
+        Set the model components for all line cut models.
+    set_param_hint(*args, **kwargs):
+        Set parameter hints for all line cut models.
+    make_params():
+        Create parameters for all line cut models.
+    guess():
+        Make initial parameter guesses for all line cut models.
+    print_initial_params():
+        Print the initial parameter values for all line cut models.
+    plot_initial_guess():
+        Plot the initial guess for all line cut models.
+    fit(verbose=False):
+        Fit the line cut models for each temperature.
+    plot_fit(mdheadings=False, **kwargs):
+        Plot the fit results for each temperature.
+    print_fit_report():
+        Print the fit report for each temperature.
     """
 
     def __init__(self):
         """
-        Initialize TempDependence class.
+        Initialize the TempDependence class with default values.
         """
-        self.nxs_files_list = []
+
         self.sample_directory = ''
         self.xlabel = ''
         self.datasets = {}
@@ -31,63 +111,107 @@ class TempDependence:
         self.scissors = {}
         self.linecuts = {}
         self.linecutmodels = {}
+        self.a, self.b, self.c, self.al, self.be, self.ga, \
+            self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star \
+            = [None] * 12
 
     def set_temperatures(self, temperatures):
+        """
+        Set the list of temperatures for the datasets.
+
+        Parameters
+        ----------
+        temperatures : list
+            List of temperatures to set.
+        """
         self.temperatures = temperatures
 
     def set_sample_directory(self, path):
+        """
+        Set the directory path where the datasets are located.
+
+        Parameters
+        ----------
+        path : str
+            Path to the sample directory.
+        """
         self.sample_directory = os.path.normpath(path)
 
     def initialize(self):
+        """
+        Initialize Scissors and LinecutModel objects for each temperature.
+        """
         for temperature in self.temperatures:
             self.scissors[temperature] = Scissors()
             self.scissors[temperature] = LinecutModel()
 
     def set_data(self, temperature, data):
+        """
+        Set the dataset for a specific temperature.
+
+        Parameters
+        ----------
+        temperature : str
+            Temperature for which to set the data.
+        data : object
+            The dataset to be set.
+        """
         self.datasets[temperature] = data
 
     def load_transforms(self, temperatures_list=None):
+        """
+        Load transform datasets (from nxrefine) based on temperature.
+
+        Parameters
+        ----------
+        temperatures_list : list of int or None, optional
+            List of temperatures to load. If None, all available temperatures are loaded.
+        """
         # Convert all temperatures to strings
         if temperatures_list:
             temperatures_list = [str(t) for t in temperatures_list]
 
-        # Search all files in the sample directory
+        # Identify files to load
+        items_to_load = []
+        # Search for nxrefine .nxs files
         for item in os.listdir(self.sample_directory):
-
-            # Find the nxs files
-            if item.endswith('.nxs'):
-
-                # Load the nxs files
-                path = os.path.join(self.sample_directory, item)
-                g = nxload(path)
-
-                # Extract the sample temperature
-                temperature = str(int(np.round(g.entry.sample.temperature.nxdata)))
-
-                # Load all samples (or a subset if temperatures_list is provided)
+            pattern = r'_(\d+)\.nxs'
+            match = re.search(pattern, item)
+            if match:
+                print(f'Found {item}')
+                # Identify temperature
+                temperature = match.group(1)
+                # print(f'Temperature = {temperature}')
                 if (temperatures_list is None) or (temperature in temperatures_list):
+                    # Prepare file to be loaded
                     self.temperatures.append(temperature)
-                    self.datasets[temperature] = load_transform(path)
-                    # Initialize scissors object at each temperature
-                    self.scissors[temperature] = Scissors()
-                    self.scissors[temperature].set_data(self.datasets[temperature])
+                    items_to_load.append(item)
+                    # print(f'Preparing to load {temperature} K data: {item}')
 
-                    # Initialize linecutmodel object at each temperature
-                    self.linecutmodels[temperature] = LinecutModel()
-
-        # Convert to int temporarily to sort temperatures list
+        # Convert all temperatures to int temporarily to sort temperatures list before loading
         self.temperatures = [int(t) for t in self.temperatures]
         self.temperatures.sort()
         self.temperatures = [str(t) for t in self.temperatures]
 
+        for i, item in enumerate(items_to_load):
+            path = os.path.join(self.sample_directory, item)
+
+            # Save dataset
+            self.datasets[self.temperatures[i]] = load_transform(path)
+
+            # Initialize scissors object
+            self.scissors[self.temperatures[i]] = Scissors()
+            self.scissors[self.temperatures[i]].set_data(self.datasets[self.temperatures[i]])
+
+            # Initialize linecutmodel object
+            self.linecutmodels[self.temperatures[i]] = LinecutModel()
+
     def load_datasets(self, file_ending='hkli.nxs', temperatures_list=None):
         """
-        Load scattering datasets from the specified folder.
+        Load datasets (CHESS format) from the specified folder.
 
         Parameters
         ----------
-        folder : str
-            The path to the folder where the datasets are located.
         file_ending : str, optional
             The file extension of the datasets to be loaded. The default is 'hkli.nxs'.
         temperatures_list : list of int or None, optional
@@ -136,8 +260,8 @@ class TempDependence:
 
         Returns
         -------
-            str:
-                The folder path.
+        str
+            The folder path.
         """
         return self.sample_directory
 
@@ -147,26 +271,39 @@ class TempDependence:
         """
         self.datasets = {}
 
-    def set_Lattice_params(self, lattice_params):
+    def set_lattice_params(self, lattice_params):
+        """
+        Set lattice parameters and calculate reciprocal lattice parameters.
+
+        Parameters
+        ----------
+        lattice_params : tuple
+            Tuple containing lattice parameters (a, b, c, al, be, ga).
+        """
         self.a, self.b, self.c, self.al, self.be, self.ga = lattice_params
-        self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star = reciprocal_lattice_params(lattice_params)
-    def set_window(self, window):
+        self.a_star, self.b_star, self.c_star, \
+            self.al_star, self.be_star, self.ga_star = reciprocal_lattice_params(lattice_params)
+
+    def set_window(self, window, verbose=False):
         """
-        Set the extents of the integration window.
+        Set the extents of the integration window for each temperature.
 
         Parameters
         ----------
         window : tuple
             Extents of the window for integration along each axis.
+        verbose : bool, optional
+            Enables printout of linecut axis and integrated axes. Default is False.
         """
         for T in self.temperatures:
-            print("----------------------------------")
-            print("T = " + T + " K")
-            self.scissors[T].set_window(window)
+            if verbose:
+                print("----------------------------------")
+                print("T = " + T + " K")
+            self.scissors[T].set_window(window, verbose)
 
     def set_center(self, center):
         """
-        Set the central coordinate for the linecut.
+        Set the central coordinate for the linecut for each temperature.
 
         Parameters
         ----------
@@ -176,39 +313,46 @@ class TempDependence:
         for T in self.temperatures:
             self.scissors[T].set_center(center)
 
-    def cut_data(self, center=None, window=None, axis=None):
+    def cut_data(self, center=None, window=None, axis=None, verbose=False):
         """
         Perform data cutting for each temperature dataset.
 
         Parameters
         ----------
-        center : tuple
+        center : tuple, optional
             The center point for cutting the data.
-        window : tuple
+            Defaults to the first temperature's center if None.
+        window : tuple, optional
             The window size for cutting the data.
+            Defaults to the first temperature's window if None.
         axis : int or None, optional
-            The axis along which to perform the cutting. If None, cutting is performed along the
-            longest axis in `window`. The default is None.
+            The axis along which to perform the cutting.
+            Defaults to the longest axis in `window` if None.
+        verbose : bool, optional
+            Enables printout of linecut progress. Default is False.
 
         Returns
         -------
-        list
-            A list of linecuts obtained from the cutting operation.
+        dict
+            A dictionary of linecuts obtained from the cutting operation.
         """
 
         center = center if center is not None else self.scissors[self.temperatures[0]].center
         window = window if window is not None else self.scissors[self.temperatures[0]].window
+        axis = axis if axis is not None else self.scissors[self.temperatures[0]].axis
 
         for T in self.temperatures:
-            print("-------------------------------")
-            print("Cutting T = " + T + " K data...")
-            self.scissors[T].cut_data(center, window, axis)
+            if verbose:
+                print("-------------------------------")
+                print("Cutting T = " + T + " K data...")
+            self.scissors[T].cut_data(center, window, axis, verbose)
             self.linecuts[T] = self.scissors[T].linecut
             self.linecutmodels[T].set_data(self.linecuts[T])
 
         xlabel_components = [self.linecuts[self.temperatures[0]].axes
                              if i == self.scissors[self.temperatures[0]].axis
-                             else str(c) for i, c in enumerate(self.scissors[self.temperatures[0]].center)]
+                             else str(c) for i, c in
+                             enumerate(self.scissors[self.temperatures[0]].center)]
         self.xlabel = ' '.join(xlabel_components)
 
         return self.linecuts
@@ -250,10 +394,53 @@ class TempDependence:
 
         return fig, ax
 
+    def plot_linecuts_heatmap(self, ax=None, **kwargs):
+        """
+        Plot the linecuts obtained from data cutting.
+
+        Parameters
+        ----------
+        ax : matplotlib.axes.Axes, optional
+            The axes on which to plot the heatmap. If None, a new figure and axes
+            are created. The default is None.
+        **kwargs
+            Additional keyword arguments to be passed to the `pcolormesh` function.
+
+        Returns
+        -------
+        QuadMesh
+            The plotted heatmap object.
+        """
+
+        # Retrieve linecut data for the first temperature and extract x-axis data
+        cut = self.linecuts[self.temperatures[0]]
+        x = cut[cut.axes].nxdata
+
+        # Convert the list of temperatures to a NumPy array for the y-axis
+        y = np.array([int(t) for t in self.temperatures])
+
+        # Collect counts from each temperature and ensure they are numpy arrays
+        v = [self.linecuts[T].counts.nxdata for T in self.temperatures]
+
+        # Convert list of arrays to a 2D array for the heatmap
+        v_2d = np.array(v)
+
+        # Create the grid for the heatmap
+        X, Y = np.meshgrid(x, y)
+
+        # Plot using pcolormesh
+        if ax is None:
+            _, ax = plt.subplots()
+        p = ax.pcolormesh(X, Y, v_2d, **kwargs)
+        plt.colorbar(p, label='counts')
+        ax.set(xlabel=cut.axes, ylabel=r'$T$ (K)')
+
+        return p
+
     def highlight_integration_window(self, temperature=None, **kwargs):
         """
-        Displays the integration window plot for a specific temperature, or for the first temperature if
-        none is provided.
+        Displays the integration window plot for a specific temperature,
+        or for the first temperature if none is provided.
 
         Parameters
         ----------
@@ -266,8 +453,10 @@ class TempDependence:
         """
 
         if temperature is not None:
-            p = self.scissors[self.temperatures[0]].highlight_integration_window(data=self.datasets[temperature],
-                                                                                 **kwargs)
+            p = self.scissors[
+                self.temperatures[0]].highlight_integration_window(
+                data=self.datasets[temperature], **kwargs
+            )
         else:
             p = self.scissors[self.temperatures[0]].highlight_integration_window(
                 data=self.datasets[self.temperatures[0]], **kwargs
@@ -277,15 +466,18 @@ class TempDependence:
 
     def plot_integration_window(self, temperature=None, **kwargs):
         """
-        Plots the three principal cross-sections of the integration volume on a single figure for a specific
-        temperature, or for the first temperature if none is provided.
+        Plots the three principal cross-sections of the integration volume on
+        a single figure for a specific temperature, or for the first temperature
+        if none is provided.
 
         Parameters
         ----------
         temperature : str, optional
-            The temperature at which to plot the integration volume. If provided, the plot
-            will be generated using the dataset corresponding to the specified temperature. If not
-            provided, the integration window plots will be generated for the first temperature.
+            The temperature at which to plot the integration volume. If provided,
+            the plot will be generated using the dataset corresponding to the
+            specified temperature. If not provided, the integration window plots
+            will be generated for the first temperature.
+
         **kwargs : keyword arguments, optional
             Additional keyword arguments to customize the plot.
         """
@@ -301,9 +493,9 @@ class TempDependence:
         """
         Set the model components for all line cut models.
 
-        This method sets the same model components for all line cut models in the analysis.
-        It iterates over each line cut model and calls their respective `set_model_components` method
-        with the provided `model_components`.
+        This method sets the same model components for all line cut models in the
+        analysis. It iterates over each line cut model and calls their respective
+        `set_model_components` method with the provided `model_components`.
 
         Parameters
         ----------
@@ -311,7 +503,8 @@ class TempDependence:
             The model components to set for all line cut models.
 
         """
-        [linecutmodel.set_model_components(model_components) for linecutmodel in self.linecutmodels.values()]
+        [linecutmodel.set_model_components(model_components) for
+         linecutmodel in self.linecutmodels.values()]
 
     def set_param_hint(self, *args, **kwargs):
         """
@@ -329,7 +522,8 @@ class TempDependence:
             Arbitrary keyword arguments.
 
         """
-        [linecutmodel.set_param_hint(*args, **kwargs) for linecutmodel in self.linecutmodels.values()]
+        [linecutmodel.set_param_hint(*args, **kwargs)
+         for linecutmodel in self.linecutmodels.values()]
 
     def make_params(self):
         """
@@ -354,8 +548,9 @@ class TempDependence:
         """
         Print the initial parameter values for all line cut models.
 
-        This method prints the initial parameter values for all line cut models in the analysis.
-        It iterates over each line cut model and calls their respective `print_initial_params` method.
+        This method prints the initial parameter values for all line cut models
+        in the analysis. It iterates over each line cut model and calls their
+        respective `print_initial_params` method.
 
         """
         [linecutmodel.print_initial_params() for linecutmodel in self.linecutmodels.values()]
@@ -369,22 +564,29 @@ class TempDependence:
 
         """
         for T, linecutmodel in self.linecutmodels.items():
-            fig, ax = plt.subplots()
+            _, ax = plt.subplots()
             ax.set(title=T + ' K')
             linecutmodel.plot_initial_guess()
 
-    def fit(self):
+    def fit(self, verbose=False):
         """
         Fit the line cut models.
 
         This method fits the line cut models for each temperature in the analysis.
         It iterates over each line cut model, performs the fit, and prints the fitting progress.
 
+        Parameters
+        ----------
+       verbose : bool, optional
+            Enables printout of fitting progress. Default False.
+
         """
         for T, linecutmodel in self.linecutmodels.items():
-            print(f"Fitting {T} K  data...")
+            if verbose:
+                print(f"Fitting {T} K  data...")
             linecutmodel.fit()
-            print("Done.")
+            if verbose:
+                print("Done.")
         print("Fits completed.")
 
     def plot_fit(self, mdheadings=False, **kwargs):
@@ -401,7 +603,9 @@ class TempDependence:
             if mdheadings:
                 display(Markdown(f"### {T} K Fit Results"))
             # Plot fit
-            linecutmodel.plot_fit(xlabel=self.xlabel, ylabel=self.datasets[self.temperatures[0]].signal, title=f"{T} K",
+            linecutmodel.plot_fit(xlabel=self.xlabel,
+                                  ylabel=self.datasets[self.temperatures[0]].signal,
+                                  title=f"{T} K",
                                   **kwargs)
 
     def print_fit_report(self):