Add features

Add features to Experiment and NexusScan, plus other improvements.

experiment.py
 - add Experiment.load_xas
 - add str, repr

nexus_reader.py
 - add HdfMap.generate_ids
 - add xas_scan
 - add instrument model

scan_plot_manager.py
 - add histogram

Author: DanPorter

docs/run_local_docs.py

@@ -1,5 +1,6 @@
 """
 Run local instance of docs for testing
+Run make_docs.py first
 """
 
 from mkdocs.commands import serve

mmg_toolbox/nexus/nexus_object.py mmg_toolbox/nexus/instrument_model.pymmg_toolbox/nexus/nexus_object.py renamed to mmg_toolbox/nexus/instrument_model.py

@@ -1,14 +1,26 @@
+"""
+Instrument Model from NeXus files
+
+Makes use of NeXus classes and NXtransformations to build a structure of the instrument from a scan file.
+"""
+
 import h5py
 import numpy as np
-from matplotlib import pyplot as plt
 
 from mmg_toolbox.nexus import nexus_names as nn
 from mmg_toolbox.nexus.nexus_functions import nx_find, get_dataset_value
 from mmg_toolbox.nexus.nexus_transformations import nx_direction, nx_transformations_max_size, \
     nx_transformations_matrix, nx_transform_vector
 from mmg_toolbox.utils.rotations import norm_vector, transform_by_t_matrix
 from mmg_toolbox.utils.xray_utils import wavevector, photon_energy, photon_wavelength, bmatrix
+from mmg_toolbox.plotting.matplotlib import Axes3D
 
+# types
+Shape = tuple[int, int, int]  #  (n, i, j) == (frame, slow_pixel, fast_pixel) #TODO: should this be fast, slow?
+Pixel = tuple[int, float, float]  # (n, i, j) pixel coordinates
+
+#TODO: add polarisation
+#TODO: improve docs
 
 class NXBeam:
     """
@@ -31,18 +43,10 @@ def energy_wavelength(self):
         :return: incident_energy, incident_wavelength
         """
         if nn.NX_WL in self.beam:
-            dataset = self.beam[nn.NX_WL]
-            units = dataset.attrs.get(nn.NX_UNITS, b'nm').decode()
-            wl = dataset[()]
-            if units == 'nm':
-                wl = 10 * wl  # wavelength in Angstroms
+            wl = get_dataset_value(nn.NX_WL, self.beam, units='A')
             return photon_energy(wl), wl
         elif nn.NX_EN in self.beam:
-            dataset = self.beam[nn.NX_WL]
-            units = dataset.attrs.get(nn.NX_UNITS, b'ev').decode()
-            en = dataset[()]
-            if units.lower() == 'ev':
-                en = en / 1000.  # wavelength in keV
+            en = get_dataset_value(nn.NX_EN, self.beam, units='keV')
             return en, photon_wavelength(en)
         else:
             raise KeyError(f"{self.beam} contains no '{nn.NX_WL}' or '{nn.NX_EN}'")
@@ -75,14 +79,15 @@ def __init__(self, path: str, hdf_file: h5py.File):
     def __repr__(self):
         return f"NXSsample({self.sample})"
 
-    def hkl2q(self, hkl: tuple[float, float, float] | np.ndarray):
+    def hkl2q(self, hkl: tuple[float, float, float] | np.ndarray, index: int = 0) -> np.ndarray:
         """
-        Returns wavecector direction for given hkl
-        :param hkl: Miller indices, in units of reciprocal lattice vectors
+        Returns wavevector direction for given hkl in lab space
+        :param hkl: Miller indices (h, k, l), in units of reciprocal lattice vectors
+        :param index: index of scan
         :return: Q position in inverse Angstroms
         """
         hkl = np.reshape(hkl, (-1, 3))
-        z = self.transforms[0][:3, :3]
+        z = self.transforms[index][:3, :3]
         ub = 2 * np.pi * self.ub_matrix
         return np.dot(z, np.dot(ub, hkl.T)).T
 
@@ -121,7 +126,7 @@ def __init__(self, path: str, hdf_file: h5py.File):
     def __repr__(self):
         return f"NXDetectorModule({self.module})"
 
-    def shape(self):
+    def shape(self) -> Shape:
         """
         Return scan shape of module
             (n, i, j)
@@ -132,26 +137,26 @@ def shape(self):
         """
         return self.size, self.data_size[0], self.data_size[1]
 
-    def pixel_wavevector(self, point: tuple[int, int, int], wavelength_a) -> np.ndarray:
+    def pixel_wavevector(self, point: Pixel, wavelength_a) -> np.ndarray:
         """
-        Return wavevector of pixel
+        Return wavevector of pixel in lab frame
         :param point: (n, i, j) == (frame, slow_axis_pixel, fast_axis_pixel)
         :param wavelength_a: wavelength in Angstrom
-        :return: [dx, dy, dz] unit vector
+        :return: [x, y, z] in inverse Angstrom
         """
         return wavevector(wavelength_a) * self.pixel_direction(point)
 
-    def pixel_direction(self, point: tuple[int, int, int]) -> np.ndarray:
+    def pixel_direction(self, point: Pixel) -> np.ndarray:
         """
-        Return direction of pixel
+        Return direction of pixel in lab frame
         :param point: (n, i, j) == (frame, slow_axis_pixel, fast_axis_pixel)
         :return: [dx, dy, dz] unit vector
         """
         return norm_vector(self.pixel_position(point))
 
-    def pixel_position(self, point: tuple[int, int, int]) -> np.ndarray:
+    def pixel_position(self, point: Pixel) -> np.ndarray:
         """
-        Return position of pixel (n, i, j)
+        Return position of pixel (n, i, j) in lab frame
             n = frame in scan
             i = pixel along slow axis
             j = pixel along fast axis
@@ -167,6 +172,7 @@ def pixel_position(self, point: tuple[int, int, int]) -> np.ndarray:
         return np.squeeze(ii * slow_direction + jj * fast_direction + module_origin)
 
     def corners(self, frame: int) -> np.ndarray:
+        """return corners of the detector module at this frame in scan"""
         shape = self.shape()
         corners = np.vstack([
             self.pixel_position((frame, 0, 0)),  # module origin
@@ -200,9 +206,9 @@ def __repr__(self):
         return f"NXDetector({self.detector}) with {len(self.modules)} modules"
 
 
-class NXScan:
+class NXInstrumentModel:
     """
-    NXScan object
+    NXInstrumentModel object
     """
 
     def __init__(self, hdf_file: h5py.File):
@@ -228,26 +234,47 @@ def __init__(self, hdf_file: h5py.File):
 
         sample_obj = nx_find(self.entry, nn.NX_SAMPLE)
         self.sample = NXSsample(sample_obj.name, hdf_file)
-        beam_obj = nx_find(self.sample_obj, nn.NX_BEAM)
+        beam_obj = nx_find(sample_obj, nn.NX_BEAM)
         self.beam = NXBeam(beam_obj.name, hdf_file)
 
     def __repr__(self):
-        return f"NXScan({self.file})"
+        return f"NXInstrumentModel({self.file})"
 
-    def shape(self):
-        detector_module = self.detectors[0].modules[0]
+    def _first_detector(self):
+        return self.detectors[0].modules[0]
+
+    def shape(self) -> Shape:
+        """
+        Return scan shape from the first detector module
+            (n, i, j)
+        Where:
+            n = frames in scan
+            i = pixels along slow axis
+            j = pixels along fast axis
+        """
+        detector_module = self._first_detector()
         return detector_module.shape()
 
-    def detector_q(self, point: tuple[int, int, int] = (0, 0, 0)):
+    def detector_q(self, point: Pixel = (0, 0, 0)) -> np.ndarray:
+        """
+        return wavevector transfer, kf-ki, of first detector module, in lab frame
+        :param point: (n, i, j) == (frame, slow_axis_pixel, fast_axis_pixel)
+        :return: [x, y, z] in inverse Angstrom
+        """
         wavelength = self.beam.wl
         ki = self.beam.incident_wavevector
-        detector_module = self.detectors[0].modules[0]
+        detector_module = self._first_detector()
         kf = detector_module.pixel_wavevector(point, wavelength)
         return kf - ki
 
-    def hkl(self, point: tuple[int, int, int] = (0, 0, 0)):
+    def hkl(self, point: Pixel = (0, 0, 0)):
+        """
+        return wavevector transfer, Q, of pixel coordinate
+        :param point: (n, i, j) == (frame, slow_axis_pixel, fast_axis_pixel)
+        :return: [x, y, z] in inverse Angstrom
+        """
         q = self.detector_q(point)
-        z = self.sample.transforms[-1][:3, :3]  #TODO: should this chain all transforms?
+        z = self.sample.transforms[point[0]][:3, :3]
         ub = 2 * np.pi * self.sample.ub_matrix
 
         inv_ub = np.linalg.inv(ub)
@@ -264,98 +291,89 @@ def hkl2q(self, hkl: tuple[float, float, float] | np.ndarray):
         """
         return self.sample.hkl2q(hkl)
 
-    def plot_instrument(self, figsize=[16, 6], dpi=100):
-        fig = plt.figure(figsize=figsize, dpi=dpi)
-        ax = fig.add_subplot(projection='3d')
-
+    def plot_instrument(self, axes: Axes3D):
         instrument_name = get_dataset_value('name', self.instrument, 'no name')
         max_distance = max([np.linalg.norm(position) for position in self.component_positions.values()])
         max_position = max_distance * self.beam.direction
 
-        ax.plot([-max_position[0], 0], [-max_position[2], 0], [-max_position[1], 0], 'k-')  # beam
+        axes.plot([-max_position[0], 0], [-max_position[2], 0], [-max_position[1], 0], 'k-')  # beam
         beam_cont = np.linalg.norm(self.detectors[0].position) * self.beam.direction
-        ax.plot([0, beam_cont[0]], [0, beam_cont[2]], [0, beam_cont[1]], 'k:')  # continued beam
+        axes.plot([0, beam_cont[0]], [0, beam_cont[2]], [0, beam_cont[1]], 'k:')  # continued beam
         # detectors
         for detector in self.detectors:
             pos = detector.position
-            ax.plot([0, pos[0]], [0, pos[2]], [0, pos[1]], 'k-')  # scattered beam
+            axes.plot([0, pos[0]], [0, pos[2]], [0, pos[1]], 'k-')  # scattered beam
         # components
         for component, position in self.component_positions.items():
-            ax.plot(position[0], position[2], position[1], 'r+')
-            ax.text(position[0], position[2], position[1], s=component)
+            axes.plot(position[0], position[2], position[1], 'r+')
+            axes.text(position[0], position[2], position[1], s=component)
 
-        ax.set_xlabel('X [mm]')
-        ax.set_ylabel('Z [mm]')
-        ax.set_zlabel('Y [mm]')
-        ax.set_title(f"Instrument: {instrument_name}")
+        axes.set_xlabel('X [mm]')
+        axes.set_ylabel('Z [mm]')
+        axes.set_zlabel('Y [mm]')
+        axes.set_title(f"Instrument: {instrument_name}")
         # ax.set_aspect('equalxz')
-        fig.show()
-
-    def plot_wavevectors(self, figsize=[16, 6], dpi=100):
-        fig = plt.figure(figsize=figsize, dpi=dpi)
-        ax = fig.add_subplot(projection='3d')
 
-        pixel_centre = tuple([i // 2 for i in self.shape()])
+    def plot_wavevectors(self, axes: Axes3D):
+        frames, ii, jj = self.shape()
+        pixel_centre = (frames // 2, ii // 2, jj // 2)
         ki = self.beam.incident_wavevector
         detector_module = self.detectors[0].modules[0]
         kf = detector_module.pixel_wavevector(pixel_centre, self.beam.wl)
         q = kf - ki
+        hkl = self.hkl(pixel_centre)
 
-        ax.plot([-ki[0], 0], [-ki[2], 0], [-ki[1], 0], '-k')
-        ax.plot([0, kf[0]], [0, kf[2]], [0, kf[1]], '-k')
-        ax.plot([0, q[0]], [0, q[2]], [0, q[1]], '-r')
+        axes.plot([-ki[0], 0], [-ki[2], 0], [-ki[1], 0], '-k')
+        axes.plot([0, kf[0]], [0, kf[2]], [0, kf[1]], '-k')
+        axes.plot([0, q[0]], [0, q[2]], [0, q[1]], '-r')
 
-        shape = self.shape()
         wl = self.beam.wl
-        for frame in range(shape[0]):
+        for frame in range(frames):
             corners = np.vstack([
                 detector_module.pixel_wavevector((frame, 0, 0), wl),  # module origin
-                detector_module.pixel_wavevector((frame, shape[1], 0), wl),  # module origin + slow pixels
-                detector_module.pixel_wavevector((frame, shape[1], shape[2]), wl),  # o + slow + fast
-                detector_module.pixel_wavevector((frame, 0, shape[2]), wl),  # o + fast
+                detector_module.pixel_wavevector((frame, ii, 0), wl),  # module origin + slow pixels
+                detector_module.pixel_wavevector((frame, ii, jj), wl),  # o + slow + fast
+                detector_module.pixel_wavevector((frame, 0, jj), wl),  # o + fast
                 detector_module.pixel_wavevector((frame, 0, 0), wl),  # module origin
             ])
-            ax.plot(corners[:, 0], corners[:, 2], corners[:, 1], '-k')
+            axes.plot(corners[:, 0], corners[:, 2], corners[:, 1], '-k')
             corners_q = corners - ki
-            ax.plot(corners_q[:, 0], corners_q[:, 2], corners_q[:, 1], '-r')
+            axes.plot(corners_q[:, 0], corners_q[:, 2], corners_q[:, 1], '-r')
 
         # plot Reciprocal lattice
         astar, bstar, cstar = self.hkl2q(np.eye(3))
-        ax.plot([0, astar[0]], [0, astar[2]], [0, astar[1]], '-g')
-        ax.plot([0, bstar[0]], [0, bstar[2]], [0, bstar[1]], '-g')
-        ax.plot([0, cstar[0]], [0, cstar[2]], [0, cstar[1]], '-g')
-        ax.text(astar[0], astar[2], astar[1], s='a*')
-        ax.text(bstar[0], bstar[2], bstar[1], s='b*')
-        ax.text(cstar[0], cstar[2], cstar[1], s='c*')
-
-        ax.set_xlabel('X')
-        ax.set_ylabel('Z')
-        ax.set_zlabel('Y')
-        ax.set_title(f"Wavevectors\nHKL: {self.hkl(pixel_centre)}")
-        ax.set_aspect('equal')
-        fig.show()
-
-    def plot_hkl(self, figsize=[16, 6], dpi=100):
-        fig = plt.figure(figsize=figsize, dpi=dpi)
-        ax = fig.add_subplot(projection='3d')
-
-        shape = self.shape()
-        pixel_centre = tuple([i // 2 for i in shape])
-        for frame in range(shape[0]):
+        axes.plot([0, astar[0]], [0, astar[2]], [0, astar[1]], '-g')
+        axes.plot([0, bstar[0]], [0, bstar[2]], [0, bstar[1]], '-g')
+        axes.plot([0, cstar[0]], [0, cstar[2]], [0, cstar[1]], '-g')
+        axes.text(astar[0], astar[2], astar[1], s='a*')
+        axes.text(bstar[0], bstar[2], bstar[1], s='b*')
+        axes.text(cstar[0], cstar[2], cstar[1], s='c*')
+
+        axes.set_xlabel('X')
+        axes.set_ylabel('Z')
+        axes.set_zlabel('Y')
+        axes.set_title(f"Wavevectors\nHKL: {hkl}")
+        axes.set_aspect('equal')
+
+    def plot_hkl(self, axes: Axes3D):
+        frames, ii, jj = self.shape()
+        pixel_centre = (frames // 2, ii // 2, jj // 2)
+        hkl = self.hkl(pixel_centre)
+        for frame in range(frames):
             corners = np.vstack([
                 self.hkl((frame, 0, 0)),  # module origin
-                self.hkl((frame, shape[1], 0)),  # module origin + slow pixels
-                self.hkl((frame, shape[1], shape[2])),  # o + slow + fast
-                self.hkl((frame, 0, shape[2])),  # o + fast
+                self.hkl((frame, ii, 0)),  # module origin + slow pixels
+                self.hkl((frame, ii, jj)),  # o + slow + fast
+                self.hkl((frame, 0, jj)),  # o + fast
                 self.hkl((frame, 0, 0)),  # module origin
             ])
-            ax.plot(corners[:, 0], corners[:, 2], corners[:, 1], '-r')
+            axes.plot(corners[:, 0], corners[:, 2], corners[:, 1], '-r')
         origin = self.hkl((0, 0, 0))
-        ax.plot(origin[0], origin[2], origin[1], '+k')
-
-        ax.set_xlabel('H')
-        ax.set_ylabel('L')
-        ax.set_zlabel('K')
-        ax.set_title(f"HKL: {self.hkl(pixel_centre)}")
-        ax.set_aspect('equal')
-        fig.show()
+        axes.plot(origin[0], origin[2], origin[1], '+k')
+
+        axes.set_xlabel('H')
+        axes.set_ylabel('L')
+        axes.set_zlabel('K')
+        axes.set_title(f"HKL: {hkl}")
+        axes.set_aspect('equal')
+

mmg_toolbox/nexus/nexus_functions.py

@@ -7,7 +7,7 @@
 import h5py
 
 from mmg_toolbox.nexus import nexus_names as nn
-from mmg_toolbox.nexus.nexus_names import METERS
+from mmg_toolbox.utils.units import METERS, unit_converter
 from mmg_toolbox.utils.xray_utils import photon_energy, photon_wavelength
 
 
@@ -196,20 +196,25 @@ def get_metadata(group: h5py.Group, *name_paths_default: tuple[str, tuple, str])
     return metadata
 
 
-def get_dataset_value(path: str, group: h5py.Group | h5py.File, default: np.ndarray) -> np.ndarray:
+def get_dataset_value(path: str, group: h5py.Group,
+                      default: str | float | np.ndarray | None = None, units: str = '') -> np.ndarray | None:
     """
     Get value from dataset in group, or return default
     :param path: hdf path of dataset in group
     :param group: hdf group
     :param default: returned if path doesn't exist
-    :return: value or default
+    :param units: converts to given units if units in attrs
+    :return: value or default as ndarray, or None if default is None
     """
     if path in group:
         dataset = group[path]
         if np.issubdtype(dataset, np.number):
-            return np.squeeze(dataset[...])
+            data = np.squeeze(dataset[...])
+            if units and 'units' in dataset.attrs:
+                return unit_converter(data, bytes2str(dataset.attrs['units']), units)
+            return data
         return dataset.asstr()[...]
-    return default
+    return None if default is None else np.asarray(default)
 
 
 def nx_beam_energy(beam: h5py.Group) -> tuple[float, float]:

mmg_toolbox/nexus/nexus_names.py

@@ -52,11 +52,6 @@
 ENTRY_CLASSES = ['NXentry', 'NXsubentry']
 XAS_DEFINITIONS = ['NXxas', 'NXxas_new']
 SEARCH_ATTRS = (NX_CLASS, 'local_name')  # DLS attribute 'local_name' helps match metadata to scan commands
-METERS = {  # conversion to meters
-    'km': 1e3, 'm': 1, 'cm': 0.1, 'mm': 1e-3,
-    'um': 1e-6, 'μm': 1e-6, 'nm': 1e-9,
-    'A': 1e-10, 'Å': 1e-10
-}
 
 # Polarisation field names inside NeXus groups
 # See https://manual.nexusformat.org/classes/base_classes/NXbeam.html#nxbeam