Merge pull request #819 from HEXRD/grain-energy-y-correction

Add beam energy correction to grains

Author: psavery

hexrd/fitting/calibration/grain.py

@@ -5,6 +5,7 @@
 from hexrd import matrixutil as mutil
 from hexrd.rotations import angularDifference
 from hexrd.transforms import xfcapi
+from hexrd import xrdutil
 
 from .abstract_grain import AbstractGrainCalibrator
 from .lmfit_param_handling import DEFAULT_EULER_CONVENTION
@@ -81,6 +82,7 @@ def sxcal_obj_func(grain_params, instr, xyo_det, hkls_idx,
     bvec = instr.beam_vector
     chi = instr.chi
     tvec_s = instr.tvec
+    energy_correction = instr.energy_correction
 
     # right now just stuck on the end and assumed
     # to all be the same length... FIX THIS
@@ -127,9 +129,17 @@ def sxcal_obj_func(grain_params, instr, xyo_det, hkls_idx,
             ghat_s = mutil.unitVector(np.dot(vmat_s, np.dot(rmat_c, gvec_c)))
             ghat_c = np.dot(rmat_c.T, ghat_s)
 
+            # Apply an energy correction according to grain position
+            corrected_wavelength = xrdutil.apply_correction_to_wavelength(
+                wavelength,
+                energy_correction,
+                tvec_s,
+                tvec_c,
+            )
+
             match_omes, calc_omes_tmp = grainutil.matchOmegas(
                 xyo, ghkls.T,
-                chi, rmat_c, bmat, wavelength,
+                chi, rmat_c, bmat, corrected_wavelength,
                 vInv=vinv_s,
                 beamVec=bvec,
                 omePeriod=ome_period)

hexrd/fitting/calibration/lmfit_param_handling.py

@@ -51,6 +51,20 @@ def create_instr_params(instr, euler_convention=DEFAULT_EULER_CONVENTION,
         parms_list.append((
             names['beam_energy'], energy, False, energy - 0.2, energy + 0.2
         ))
+        if beam.get('energy_correction') is not None:
+            base_name = names['beam_energy_correction']
+            intercept = beam['energy_correction']['intercept']
+            slope = beam['energy_correction']['slope']
+            parms_list.append((
+                f'{base_name}_intercept',
+                intercept,
+                False,
+                -np.inf,
+                np.inf,
+            ))
+            parms_list.append((
+                f'{base_name}_slope', slope, False, -np.inf, np.inf
+            ))
 
     parms_list.append(('instr_chi', np.degrees(instr.chi),
                        False, np.degrees(instr.chi)-1,
@@ -210,6 +224,7 @@ def create_beam_param_names(instr: HEDMInstrument) -> dict[str, str]:
             'beam_polar': f'{prefix}beam_polar',
             'beam_azimuth': f'{prefix}beam_azimuth',
             'beam_energy': f'{prefix}beam_energy',
+            'beam_energy_correction': f'{prefix}beam_energy_correction',
         }
     return param_names
 
@@ -258,12 +273,21 @@ def update_instrument_from_params(
     # 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():
+        beam = instr.beam_dict[xrs_name]
         energy = params[param_names['beam_energy']].value
         azim = params[param_names['beam_azimuth']].value
         pola = params[param_names['beam_polar']].value
 
-        instr.beam_dict[xrs_name]['energy'] = energy
-        instr.beam_dict[xrs_name]['vector'] = calc_beam_vec(azim, pola)
+        beam['energy'] = energy
+        beam['vector'] = calc_beam_vec(azim, pola)
+
+        if beam.get('energy_correction'):
+            base_name = param_names['beam_energy_correction']
+            slope = params[f'{base_name}_slope'].value
+            intercept = params[f'{base_name}_intercept'].value
+
+            beam['energy_correction']['slope'] = slope
+            beam['energy_correction']['intercept'] = intercept
 
     # Trigger any needed updates from beam modifications
     instr.beam_dict_modified()

hexrd/fitting/grains.py

@@ -10,7 +10,10 @@
 from hexrd import constants
 from hexrd import rotations
 
-from hexrd.xrdutil import extract_detector_transformation
+from hexrd.xrdutil import (
+    apply_correction_to_wavelength,
+    extract_detector_transformation,
+)
 
 return_value_flag = None
 
@@ -185,6 +188,8 @@ def objFuncFitGrain(gFit, gFull, gFlag,
     """
     bVec = instrument.beam_vector
     eVec = instrument.eta_vector
+    tVec_s = instrument.tvec
+    energy_correction = instrument.energy_correction
 
     # fill out parameters
     gFull[gFlag] = gFit
@@ -195,6 +200,14 @@ def objFuncFitGrain(gFit, gFull, gFlag,
     vInv_s = gFull[6:]
     vMat_s = mutil.vecMVToSymm(vInv_s)  # NOTE: Inverse of V from F = V * R
 
+    # Apply an energy correction according to grain position
+    corrected_wavelength = apply_correction_to_wavelength(
+        wavelength,
+        energy_correction,
+        tVec_s,
+        tVec_c,
+    )
+
     # loop over instrument panels
     # CAVEAT: keeping track of key ordering in the "detectors" attribute of
     # instrument here because I am not sure if instatiating them using
@@ -265,7 +278,7 @@ def objFuncFitGrain(gFit, gFull, gFlag,
 
         # !!!: check that this operates on UNWARPED xy
         match_omes, calc_omes = matchOmegas(
-            meas_xyo, hkls, chi, rMat_c, bMat, wavelength,
+            meas_xyo, hkls, chi, rMat_c, bMat, corrected_wavelength,
             vInv=vInv_s, beamVec=bVec, etaVec=eVec,
             omePeriod=omePeriod)
 

hexrd/instrument/detector.py

@@ -1493,6 +1493,7 @@ def simulate_rotation_series(
         chi=0.0,
         tVec_s=ct.zeros_3,
         wavelength=None,
+        energy_correction=None,
     ):
         """
         Simulate a monochromatic rotation series for a list of grains.
@@ -1561,6 +1562,14 @@ def simulate_rotation_series(
             tVec_c = gparm[3:6]
             vInv_s = gparm[6:]
 
+            # Apply an energy correction according to grain position
+            corrected_wavelength = xrdutil.apply_correction_to_wavelength(
+                wavelength,
+                energy_correction,
+                tVec_s,
+                tVec_c,
+            )
+
             # All possible bragg conditions as vstacked [tth, eta, ome]
             # for each omega solution
             angList = np.vstack(
@@ -1569,7 +1578,7 @@ def simulate_rotation_series(
                     chi,
                     rMat_c,
                     bMat,
-                    wavelength,
+                    corrected_wavelength,
                     v_inv=vInv_s,
                     beam_vec=self.bvec,
                 )

hexrd/instrument/hedm_instrument.py

@@ -36,7 +36,7 @@
 import os
 from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
 from functools import partial
-from typing import Optional
+from typing import Optional, Union
 
 from tqdm import tqdm
 
@@ -604,6 +604,7 @@ def __init__(self, instrument_config=None,
                         beam['vector']['polar_angle'],
                     ),
                     'distance': beam.get('source_distance', np.inf),
+                    'energy_correction': beam.get('energy_correction', None),
                 }
 
             # Set the active beam name if not set already
@@ -811,6 +812,7 @@ def _create_default_beam(self):
             'energy': beam_energy_DFLT,
             'vector': beam_vec_DFLT.copy(),
             'distance': np.inf,
+            'energy_correction': None,
         }
 
         if self._active_beam_name is None:
@@ -889,6 +891,51 @@ def source_distance(self, x):
         self.active_beam['distance'] = x
         self.beam_dict_modified()
 
+    @property
+    def energy_correction(self) -> Union[dict,  None]:
+        """Energy correction dict appears as follows:
+
+            {
+                # The beam energy gradient center, along the specified
+                # axis, in millimeters.
+                'intercept': 0.0,
+
+                # The slope of the beam energy gradient along the
+                # specified axis, in eV/mm.
+                'slope': 0.0,
+
+                # The specified axis for the beam energy gradient,
+                # either 'x' or 'y'.
+                'axis': 'y',
+            }
+        """
+        return self.active_beam['energy_correction']
+
+    @energy_correction.setter
+    def energy_correction(self, v: Union[dict, None]):
+        if v is not None:
+            # First validate
+            keys = sorted(list(v))
+            default_keys = sorted(list(
+                self.create_default_energy_correction()
+            ))
+            if keys != default_keys:
+                msg = (
+                    f'Keys in energy correction dict, "{keys}", do not match '
+                    f'the required keys: "{default_keys}"'
+                )
+                raise RuntimeError(msg)
+
+        self.active_beam['energy_correction'] = v
+
+    @staticmethod
+    def create_default_energy_correction() -> dict[str, float]:
+        return {
+            'intercept': 0.0,  # in mm
+            'slope': 0.0,      # eV/mm
+            'axis': 'y',
+        }
+
     @property
     def eta_vector(self):
         return self._eta_vector
@@ -932,6 +979,11 @@ def write_config(self, file=None, style='yaml', calibration_dict={}):
             if beam['distance'] != np.inf:
                 beam_dict[beam_name]['source_distance'] = beam['distance']
 
+            if beam.get('energy_correction') is not None:
+                beam_dict[beam_name]['energy_correction'] = beam[
+                    'energy_correction'
+                ]
+
         if len(beam_dict) == 1:
             # Just write it out a single beam (classical way)
             beam_dict = next(iter(beam_dict.values()))
@@ -1508,7 +1560,8 @@ def simulate_rotation_series(self, plane_data, grain_param_list,
                 ome_ranges=ome_ranges,
                 ome_period=ome_period,
                 chi=self.chi, tVec_s=self.tvec,
-                wavelength=wavelength)
+                wavelength=wavelength,
+                energy_correction=self.energy_correction)
         return results
 
     def pull_spots(self, plane_data, grain_params,

hexrd/xrdutil/utils.py

@@ -941,6 +941,7 @@ def simulateGVecs(
     pixel_pitch: tuple[float] = (0.2, 0.2),
     distortion: DistortionABC = None,
     beam_vector: np.ndarray = constants.beam_vec,
+    energy_correction: Union[dict, None] = None,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """
     returns valid_ids, valid_hkl, valid_ang, valid_xy, ang_ps
@@ -988,10 +989,18 @@ def simulateGVecs(
     vInv_s = np.ascontiguousarray(grain_params[6:12])
     beam_vector = np.ascontiguousarray(beam_vector)
 
+    # Apply an energy correction according to grain position
+    corrected_wlen = apply_correction_to_wavelength(
+        wlen,
+        energy_correction,
+        tVec_s,
+        tVec_c,
+    )
+
     # first find valid G-vectors
     angList = np.vstack(
         xfcapi.oscill_angles_of_hkls(
-            full_hkls[:, 1:], chi, rMat_c, bMat, wlen, v_inv=vInv_s,
+            full_hkls[:, 1:], chi, rMat_c, bMat, corrected_wlen, v_inv=vInv_s,
             beam_vec=beam_vector
         )
     )
@@ -1514,3 +1523,49 @@ def extract_detector_transformation(
         chi = detector_params[6]
         tVec_s = np.ascontiguousarray(detector_params[7:10])
     return rMat_d, tVec_d, chi, tVec_s
+
+
+def apply_correction_to_wavelength(
+    wavelength: float,
+    energy_correction: Union[dict, None],
+    tvec_s: np.ndarray,
+    tvec_c: np.ndarray,
+) -> float:
+    """Apply an energy correction to the wavelength according to grain position
+
+    The energy correction dict appears as follows:
+
+        {
+            # The beam energy gradient center, in millimeters,
+            # along the specified axis
+            'intercept': 0.0,
+
+            # The slope of the beam energy gradient along the
+            # specified axis, in eV/mm.
+            'slope': 0.0,
+
+            # The specified axis for the beam energy gradient,
+            # either 'x' or 'y'.
+            'axis': 'y',
+        }
+
+    If the energy_correction dict is `None`, then no correction
+    is performed.
+    """
+    if not energy_correction:
+        # No correction
+        return wavelength
+
+    # 'c' here is the conversion factor between keV and angstrom
+    c = constants.keVToAngstrom(1)
+
+    ind = 1 if energy_correction['axis'] == 'y' else 0
+
+    # Correct wavelength according to grain position. Position is in mm.
+    position = tvec_c[ind] + tvec_s[ind] - energy_correction['intercept']
+
+    # The slope is in eV/mm. Convert to keV.
+    adjustment = position * energy_correction['slope'] / 1e3
+
+    # Convert to keV, apply the adjustment, and then convert back to wavelength
+    return c / (c / wavelength + adjustment)