Merge pull request #721 from HEXRD/multi-xrs

Add support for multiple x-ray sources

hexrd/fitting/calibration/laue.py

@@ -1,3 +1,6 @@
+import copy
+from typing import Optional
+
 import numpy as np
 from scipy import ndimage
 from scipy.integrate import nquad
@@ -7,6 +10,7 @@
 
 from hexrd import xrdutil
 from hexrd.constants import fwhm_to_sigma
+from hexrd.instrument import switch_xray_source
 from hexrd.rotations import angleAxisOfRotMat, RotMatEuler
 from hexrd.transforms import xfcapi
 from hexrd.utils.hkl import hkl_to_str, str_to_hkl
@@ -25,22 +29,45 @@ class LaueCalibrator(Calibrator):
     def __init__(self, instr, material, grain_params, default_refinements=None,
                  min_energy=5, max_energy=25, tth_distortion=None,
                  calibration_picks=None,
-                 euler_convention=DEFAULT_EULER_CONVENTION):
+                 euler_convention=DEFAULT_EULER_CONVENTION,
+                 xray_source: Optional[str] = None):
         self.instr = instr
         self.material = material
         self.grain_params = grain_params
         self.default_refinements = default_refinements
         self.energy_cutoffs = [min_energy, max_energy]
-        self.tth_distortion = tth_distortion
-
         self.euler_convention = euler_convention
+        self.xray_source = xray_source
 
         self.data_dict = None
         if calibration_picks is not None:
             self.calibration_picks = calibration_picks
 
+        self._tth_distortion = tth_distortion
+        self._update_tth_distortion_panels()
+
         self.param_names = []
 
+    @property
+    def tth_distortion(self):
+        return self._tth_distortion
+
+    @tth_distortion.setter
+    def tth_distortion(self, v):
+        self._tth_distortion = v
+        self._update_tth_distortion_panels()
+
+    def _update_tth_distortion_panels(self):
+        # Make sure the panels in the tth distortion are the same
+        # as those on the instrument, so their beam vectors get modified
+        # accordingly.
+        if self._tth_distortion is None:
+            return
+
+        self._tth_distortion = copy.deepcopy(self._tth_distortion)
+        for det_key, obj in self._tth_distortion.items():
+            obj.panel = self.instr.detectors[det_key]
+
     def create_lmfit_params(self, current_params):
         params = create_grain_params(
             self.material.name,
@@ -139,8 +166,6 @@ def autopick_points(self, raw_img_dict, tth_tol=5., eta_tol=5.,
                         use_blob_detection=True, blob_threshold=0.25,
                         fit_peaks=True, min_peak_int=1., fit_tth_tol=0.1):
         """
-
-
         Parameters
         ----------
         raw_img_dict : TYPE
@@ -167,6 +192,26 @@ def autopick_points(self, raw_img_dict, tth_tol=5., eta_tol=5.,
         None.
 
         """
+
+        with switch_xray_source(self.instr, self.xray_source):
+            return self._autopick_points(
+                raw_img_dict=raw_img_dict,
+                tth_tol=tth_tol,
+                eta_tol=eta_tol,
+                npdiv=npdiv,
+                do_smoothing=do_smoothing,
+                smoothing_sigma=smoothing_sigma,
+                use_blob_detection=use_blob_detection,
+                blob_threshold=blob_threshold,
+                fit_peaks=fit_peaks,
+                min_peak_int=min_peak_int,
+                fit_tth_tol=fit_tth_tol,
+            )
+
+    def _autopick_points(self, raw_img_dict, tth_tol=5., eta_tol=5.,
+                         npdiv=2, do_smoothing=True, smoothing_sigma=2,
+                         use_blob_detection=True, blob_threshold=0.25,
+                         fit_peaks=True, min_peak_int=1., fit_tth_tol=0.1):
         labelStructure = ndimage.generate_binary_structure(2, 1)
         rmat_s = np.eye(3)  # !!! forcing to identity
         omega = 0.  # !!! same ^^^
@@ -427,6 +472,10 @@ def _evaluate(self):
         return pick_hkls_dict, pick_xys_dict
 
     def residual(self):
+        with switch_xray_source(self.instr, self.xray_source):
+            return self._residual()
+
+    def _residual(self):
         # need this for laue obj
         pick_hkls_dict, pick_xys_dict = self._evaluate()
 
@@ -439,6 +488,10 @@ def residual(self):
         )
 
     def model(self):
+        with switch_xray_source(self.instr, self.xray_source):
+            return self._model()
+
+    def _model(self):
         # need this for laue obj
         pick_hkls_dict, pick_xys_dict = self._evaluate()
 

hexrd/fitting/calibration/lmfit_param_handling.py

@@ -4,6 +4,7 @@
 from hexrd.instrument import (
     calc_angles_from_beam_vec,
     calc_beam_vec,
+    HEDMInstrument,
 )
 from hexrd.rotations import (
     expMapOfQuat,
@@ -21,31 +22,23 @@
 def create_instr_params(instr, euler_convention=DEFAULT_EULER_CONVENTION):
     # add with tuples: (NAME VALUE VARY MIN  MAX  EXPR  BRUTE_STEP)
     parms_list = []
-    if instr.has_multi_beam:
-        for k, v in instr.multi_beam_dict.items():
-            azim, pol = calc_angles_from_beam_vec(v['beam_vector'])
-            pname = f'{k}_beam_polar'
-            aname = f'{k}_beam_azimuth'
-            parms_list.append((pname, pol, False, pol-1, pol+1))
-            parms_list.append((aname, azim, False, azim-1, azim+1))
-
-            bname = f'{k}_beam_energy'
-            beam_energy = v['beam_energy']
-            parms_list.append((bname,
-                               beam_energy,
-                               False,
-                               beam_energy-0.2,
-                               beam_energy+0.2))
-    else:
-        azim, pol = calc_angles_from_beam_vec(instr.beam_vector)
-        parms_list.append(('beam_polar', pol, False, pol-1, pol+1))
-        parms_list.append(('beam_azimuth', azim, False, azim-1, azim+1))
-
-        parms_list.append(('beam_energy',
-                           instr.beam_energy,
-                           False,
-                           instr.beam_energy-0.2,
-                           instr.beam_energy+0.2))
+
+    # This supports either single beam or multi-beam
+    beam_param_names = create_beam_param_names(instr)
+    for beam_name, beam in instr.beam_dict.items():
+        azim, pol = calc_angles_from_beam_vec(beam['vector'])
+        energy = beam['energy']
+
+        names = beam_param_names[beam_name]
+        parms_list.append((
+            names['beam_polar'], pol, False, pol - 1, pol + 1
+        ))
+        parms_list.append((
+            names['beam_azimuth'], azim, False, azim - 1, azim + 1
+        ))
+        parms_list.append((
+            names['beam_energy'], energy, False, energy - 0.2, energy + 0.2
+        ))
 
     parms_list.append(('instr_chi', np.degrees(instr.chi),
                        False, np.degrees(instr.chi)-1,
@@ -93,6 +86,18 @@ def create_instr_params(instr, euler_convention=DEFAULT_EULER_CONVENTION):
     return parms_list
 
 
+def create_beam_param_names(instr: HEDMInstrument) -> dict[str, str]:
+    param_names = {}
+    for k, v in instr.beam_dict.items():
+        prefix = f'{k}_' if instr.has_multi_beam else ''
+        param_names[k] = {
+            'beam_polar': f'{prefix}beam_polar',
+            'beam_azimuth': f'{prefix}beam_azimuth',
+            'beam_energy': f'{prefix}beam_energy',
+        }
+    return param_names
+
+
 def update_instrument_from_params(instr, params, euler_convention):
     """
     this function updates the instrument from the
@@ -107,11 +112,18 @@ def update_instrument_from_params(instr, params, euler_convention):
                f'Received: {params}')
         raise NotImplementedError(msg)
 
-    instr.beam_energy = params['beam_energy'].value
+    # This supports single XRS or multi XRS
+    beam_param_names = create_beam_param_names(instr)
+    for xrs_name, param_names in beam_param_names.items():
+        energy = params[param_names['beam_energy']].value
+        azim = params[param_names['beam_azimuth']].value
+        pola = params[param_names['beam_polar']].value
 
-    azim = params['beam_azimuth'].value
-    pola = params['beam_polar'].value
-    instr.beam_vector = calc_beam_vec(azim, pola)
+        instr.beam_dict[xrs_name]['energy'] = energy
+        instr.beam_dict[xrs_name]['vector'] = calc_beam_vec(azim, pola)
+
+    # Trigger any needed updates from beam modifications
+    instr.beam_dict_modified()
 
     chi = np.radians(params['instr_chi'].value)
     instr.chi = chi
@@ -150,23 +162,46 @@ def update_instrument_from_params(instr, params, euler_convention):
                     )
 
 
-def create_tth_parameters(meas_angles):
+def create_tth_parameters(
+    instr: HEDMInstrument,
+    meas_angles: dict[str, np.ndarray],
+) -> list[lmfit.Parameter]:
+
+    prefixes = tth_parameter_prefixes(instr)
+
     parms_list = []
-    for ii, tth in enumerate(meas_angles):
-        ds_ang = np.empty([0,])
-        for k, v in tth.items():
-            if v is not None:
-                ds_ang = np.concatenate((ds_ang, v[:, 0]))
-        if ds_ang.size != 0:
-            val = np.degrees(np.mean(ds_ang))
-            parms_list.append((f'DS_ring_{ii}',
+    for xray_source, angles in meas_angles.items():
+        prefix = prefixes[xray_source]
+        for ii, tth in enumerate(angles):
+            ds_ang = []
+            for k, v in tth.items():
+                if v is not None:
+                    ds_ang.append(v[:, 0])
+
+            if not ds_ang:
+                continue
+
+            val = np.degrees(np.mean(np.hstack(ds_ang)))
+
+            parms_list.append((f'{prefix}{ii}',
                                val,
                                True,
                                val-5.,
                                val+5.))
+
     return parms_list
 
 
+def tth_parameter_prefixes(instr: HEDMInstrument) -> dict[str, str]:
+    """Generate tth parameter prefixes according to beam names"""
+    prefix = 'DS_ring_'
+    beam_names = instr.beam_names
+    if len(beam_names) == 1:
+        return {beam_names[0]: prefix}
+
+    return {name: f'{name}_{prefix}' for name in beam_names}
+
+
 def create_material_params(material, refinements=None):
     # The refinements should be in reduced format
     refine_idx = 0

hexrd/fitting/calibration/powder.py

@@ -1,6 +1,10 @@
+import copy
+from typing import Optional
+
 import numpy as np
 
 from hexrd import matrixutil as mutil
+from hexrd.instrument import calc_angles_from_beam_vec, switch_xray_source
 from hexrd.utils.hkl import hkl_to_str, str_to_hkl
 
 from .calibrator import Calibrator
@@ -19,14 +23,16 @@ def __init__(self, instr, material, img_dict, default_refinements=None,
                  tth_tol=None, eta_tol=0.25,
                  fwhm_estimate=None, min_pk_sep=1e-3, min_ampl=0.,
                  pktype='pvoigt', bgtype='linear',
-                 tth_distortion=None, calibration_picks=None):
+                 tth_distortion=None, calibration_picks=None,
+                 xray_source: Optional[str] = None):
         assert list(instr.detectors.keys()) == list(img_dict.keys()), \
             "instrument and image dict must have the same keys"
 
         self.instr = instr
         self.material = material
         self.img_dict = img_dict
         self.default_refinements = default_refinements
+        self.xray_source = xray_source
 
         # for polar interpolation
         if tth_tol is not None:
@@ -40,9 +46,11 @@ def __init__(self, instr, material, img_dict, default_refinements=None,
         self.min_ampl = min_ampl
         self.pktype = pktype
         self.bgtype = bgtype
-        self.tth_distortion = tth_distortion
 
-        self.plane_data.wavelength = instr.beam_energy  # force
+        self._tth_distortion = tth_distortion
+        self._update_tth_distortion_panels()
+
+        self.plane_data.wavelength = instr.xrs_beam_energy(xray_source)
 
         self.param_names = []
 
@@ -51,17 +59,43 @@ def __init__(self, instr, material, img_dict, default_refinements=None,
             # container for calibration data
             self.calibration_picks = calibration_picks
 
+    @property
+    def tth_distortion(self):
+        return self._tth_distortion
+
+    @tth_distortion.setter
+    def tth_distortion(self, v):
+        self._tth_distortion = v
+        self._update_tth_distortion_panels()
+
+    def _update_tth_distortion_panels(self):
+        # Make sure the panels in the tth distortion are the same
+        # as those on the instrument, so their beam vectors get modified
+        # accordingly.
+        if self._tth_distortion is None:
+            return
+
+        self._tth_distortion = copy.deepcopy(self._tth_distortion)
+        for det_key, obj in self._tth_distortion.items():
+            obj.panel = self.instr.detectors[det_key]
+
     def create_lmfit_params(self, current_params):
         # There shouldn't be more than one calibrator for a given material, so
         # just assume we have a unique name...
         params = create_material_params(self.material,
                                         self.default_refinements)
 
+        # If multiple powder calibrators were used for the same material (such
+        # as in 2XRS), then don't add params again.
+        param_names = [x[0] for x in current_params]
+        params = [x for x in params if x[0] not in param_names]
+
         self.param_names = [x[0] for x in params]
         return params
 
     def update_from_lmfit_params(self, params_dict):
-        update_material_from_params(params_dict, self.material)
+        if self.param_names:
+            update_material_from_params(params_dict, self.material)
 
     @property
     def plane_data(self):
@@ -133,6 +167,12 @@ def autopick_points(self, fit_tth_tol=5., int_cutoff=1e-4):
 
         FIXME: can not yet handle tth ranges with multiple peaks!
         """
+        # If needed, change the x-ray source before proceeding.
+        # This does nothing for single x-ray sources.
+        with switch_xray_source(self.instr, self.xray_source):
+            return self._autopick_points(fit_tth_tol, int_cutoff)
+
+    def _autopick_points(self, fit_tth_tol=5., int_cutoff=1e-4):
         # ideal tth
         dsp_ideal = np.atleast_1d(self.plane_data.getPlaneSpacings())
         hkls_ref = self.plane_data.hkls.T
@@ -331,7 +371,13 @@ def _evaluate(self, output='residual'):
         return retval
 
     def residual(self):
-        return self._evaluate(output='residual')
+        # If needed, change the x-ray source before proceeding.
+        # This does nothing for single x-ray sources.
+        with switch_xray_source(self.instr, self.xray_source):
+            return self._evaluate(output='residual')
 
     def model(self):
-        return self._evaluate(output='model')
+        # If needed, change the x-ray source before proceeding.
+        # This does nothing for single x-ray sources.
+        with switch_xray_source(self.instr, self.xray_source):
+            return self._evaluate(output='model')