Merge pull request #766 from HEXRD/group-relative-constraints

Add group ("Detector") relative constraints

Author: psavery

hexrd/fitting/calibration/lmfit_param_handling.py

@@ -6,6 +6,7 @@
 from hexrd.instrument import (
     calc_angles_from_beam_vec,
     calc_beam_vec,
+    Detector,
     HEDMInstrument,
 )
 from hexrd.rotations import (
@@ -26,6 +27,8 @@
 # First is the axes_order, second is extrinsic
 DEFAULT_EULER_CONVENTION = ('zxz', False)
 
+EULER_CONVENTION_TYPES = dict | tuple | None
+
 
 def create_instr_params(instr, euler_convention=DEFAULT_EULER_CONVENTION,
                         relative_constraints=None):
@@ -66,8 +69,12 @@ def create_instr_params(instr, euler_convention=DEFAULT_EULER_CONVENTION,
             parms_list,
         )
     elif relative_constraints.type == RelativeConstraintsType.group:
-        # This should be implemented soon
-        raise NotImplementedError(relative_constraints.type)
+        add_group_constrained_detector_parameters(
+            instr,
+            euler_convention,
+            parms_list,
+            relative_constraints,
+        )
     elif relative_constraints.type == RelativeConstraintsType.system:
         add_system_constrained_detector_parameters(
             instr,
@@ -120,12 +127,14 @@ def add_unconstrained_detector_parameters(instr, euler_convention, parms_list):
                                -np.inf, np.inf))
 
 
-def add_system_constrained_detector_parameters(
-        instr, euler_convention,
-        parms_list, relative_constraints: RelativeConstraints):
-    system_params = relative_constraints.params
-    system_tvec = system_params['translation']
-    system_tilt = system_params['tilt']
+def _add_constrained_detector_parameters(
+    euler_convention: EULER_CONVENTION_TYPES,
+    parms_list: list[tuple],
+    prefix: str,
+    constraint_params: dict,
+):
+    tvec = constraint_params['translation']
+    tilt = constraint_params['tilt']
 
     if euler_convention is not None:
         # Convert the tilt to the specified Euler convention
@@ -136,30 +145,63 @@ def add_system_constrained_detector_parameters(
             extrinsic=normalized[1],
         )
 
-        rme.rmat = _tilt_to_rmat(system_tilt, None)
-        system_tilt = np.degrees(rme.angles)
+        rme.rmat = _tilt_to_rmat(tilt, None)
+        tilt = np.degrees(rme.angles)
 
     tvec_names = [
-        'system_tvec_x',
-        'system_tvec_y',
-        'system_tvec_z',
+        f'{prefix}_tvec_x',
+        f'{prefix}_tvec_y',
+        f'{prefix}_tvec_z',
     ]
     tvec_deltas = [1, 1, 1]
 
-    tilt_names = param_names_euler_convention('system', euler_convention)
+    tilt_names = param_names_euler_convention(prefix, euler_convention)
     tilt_deltas = [2, 2, 2]
 
     for i, name in enumerate(tvec_names):
-        value = system_tvec[i]
+        value = tvec[i]
         delta = tvec_deltas[i]
         parms_list.append((name, value, True, value - delta, value + delta))
 
     for i, name in enumerate(tilt_names):
-        value = system_tilt[i]
+        value = tilt[i]
         delta = tilt_deltas[i]
         parms_list.append((name, value, True, value - delta, value + delta))
 
 
+def add_system_constrained_detector_parameters(
+    instr: HEDMInstrument,
+    euler_convention: EULER_CONVENTION_TYPES,
+    parms_list: list[tuple],
+    relative_constraints: RelativeConstraints,
+):
+    prefix = 'system'
+    constraint_params = relative_constraints.params
+    _add_constrained_detector_parameters(
+        euler_convention,
+        parms_list,
+        prefix,
+        constraint_params,
+    )
+
+
+def add_group_constrained_detector_parameters(
+    instr: HEDMInstrument,
+    euler_convention: EULER_CONVENTION_TYPES,
+    parms_list: list[tuple],
+    relative_constraints: RelativeConstraints,
+):
+    for group in instr.detector_groups:
+        prefix = group.replace('-', '_')
+        constraint_params = relative_constraints.params[group]
+        _add_constrained_detector_parameters(
+            euler_convention,
+            parms_list,
+            prefix,
+            constraint_params,
+        )
+
+
 def create_beam_param_names(instr: HEDMInstrument) -> dict[str, str]:
     param_names = {}
     for k, v in instr.beam_dict.items():
@@ -220,8 +262,12 @@ def update_instrument_from_params(
             euler_convention,
         )
     elif relative_constraints.type == RelativeConstraintsType.group:
-        # This should be implemented soon
-        raise NotImplementedError(relative_constraints.type)
+        update_group_constrained_detector_parameters(
+            instr,
+            params,
+            euler_convention,
+            relative_constraints,
+        )
     elif relative_constraints.type == RelativeConstraintsType.system:
         update_system_constrained_detector_parameters(
             instr,
@@ -263,60 +309,109 @@ def update_unconstrained_detector_parameters(instr, params, euler_convention):
                     )
 
 
-def update_system_constrained_detector_parameters(
-        instr, params, euler_convention,
-        relative_constraints: RelativeConstraints):
-    system_params = relative_constraints.params
-    system_tvec = system_params['translation']
-    system_tilt = system_params['tilt']
+def _update_constrained_detector_parameters(
+    detectors: list[Detector],
+    params: dict,
+    rotation_center: np.ndarray,
+    euler_convention: EULER_CONVENTION_TYPES,
+    prefix: str,
+    constraint_params: dict,
+
+):
+    tvec = constraint_params['translation']
+    tilt = constraint_params['tilt']
 
     tvec_names = [
-        'system_tvec_x',
-        'system_tvec_y',
-        'system_tvec_z',
+        f'{prefix}_tvec_x',
+        f'{prefix}_tvec_y',
+        f'{prefix}_tvec_z',
     ]
-    tilt_names = param_names_euler_convention('system', euler_convention)
+    tilt_names = param_names_euler_convention(prefix, euler_convention)
 
     # Just like the detectors, we will apply tilt first and then translation
     # Only apply these transforms if they were marked "Vary".
 
     if any(params[x].vary for x in tilt_names):
-        # Get the center of rotation (depending on the settings)
-        rotation_center = relative_constraints.center_of_rotation(instr)
-
         # Find the change in tilt, create an rmat, then apply to detector tilts
         # and translations.
-        new_system_tilt = np.array([params[x].value for x in tilt_names])
+        new_tilt = np.array([params[x].value for x in tilt_names])
 
-        # The old system tilt was in the None convention
-        old_rmat = _tilt_to_rmat(system_tilt, None)
-        new_rmat = _tilt_to_rmat(new_system_tilt, euler_convention)
+        # The old tilt was in the None convention
+        old_rmat = _tilt_to_rmat(tilt, None)
+        new_rmat = _tilt_to_rmat(new_tilt, euler_convention)
 
         # Compute the rmat used to convert from old to new
         rmat_diff = new_rmat @ old_rmat.T
 
         # Rotate each detector using the rmat_diff
-        for panel in instr.detectors.values():
+        for panel in detectors:
             panel.tilt = _rmat_to_tilt(rmat_diff @ panel.rmat)
 
             # Also rotate the detectors about the rotation center
             panel.tvec = (
                 rmat_diff @ (panel.tvec - rotation_center) + rotation_center
             )
 
-        # Update the system tilt
-        system_tilt[:] = _rmat_to_tilt(new_rmat)
+        # Update the tilt
+        tilt[:] = _rmat_to_tilt(new_rmat)
 
     if any(params[x].vary for x in tvec_names):
         # Find the change in center and shift all tvecs
-        new_system_tvec = np.array([params[x].value for x in tvec_names])
+        new_tvec = np.array([params[x].value for x in tvec_names])
 
-        diff = new_system_tvec - system_tvec
-        for panel in instr.detectors.values():
+        diff = new_tvec - tvec
+        for panel in detectors:
             panel.tvec += diff
 
-        # Update the system tvec
-        system_tvec[:] = new_system_tvec
+        # Update the tvec
+        tvec[:] = new_tvec
+
+
+def update_system_constrained_detector_parameters(
+    instr: HEDMInstrument,
+    params: dict,
+    euler_convention: EULER_CONVENTION_TYPES,
+    relative_constraints: RelativeConstraints,
+):
+    detectors = list(instr.detectors.values())
+
+    # Get the center of rotation (depending on the settings)
+    rotation_center = relative_constraints.center_of_rotation(instr)
+    prefix = 'system'
+    constraint_params = relative_constraints.params
+
+    _update_constrained_detector_parameters(
+        detectors,
+        params,
+        rotation_center,
+        euler_convention,
+        prefix,
+        constraint_params,
+    )
+
+
+def update_group_constrained_detector_parameters(
+    instr: HEDMInstrument,
+    params: dict,
+    euler_convention: EULER_CONVENTION_TYPES,
+    relative_constraints: RelativeConstraints,
+):
+    for group in instr.detector_groups:
+        detectors = list(instr.detectors_in_group(group).values())
+
+        # Get the center of rotation (depending on the settings)
+        rotation_center = relative_constraints.center_of_rotation(instr, group)
+        prefix = group.replace('-', '_')
+        constraint_params = relative_constraints.params[group]
+
+        _update_constrained_detector_parameters(
+            detectors,
+            params,
+            rotation_center,
+            euler_convention,
+            prefix,
+            constraint_params,
+        )
 
 
 def _tilt_to_rmat(tilt: np.ndarray,

hexrd/fitting/calibration/relative_constraints.py

@@ -21,6 +21,9 @@ class RotationCenter(Enum):
     # Rotate about the mean center of all the detectors
     instrument_mean_center = 'InstrumentMeanCenter'
 
+    # Rotate about the mean center of the group
+    group_mean_center = 'GroupMeanCenter'
+
     # Rotate about lab origin, which is (0, 0, 0)
     lab_origin = 'Origin'
 
@@ -75,11 +78,7 @@ class RelativeConstraintsGroup(RelativeConstraints):
     type = RelativeConstraintsType.group
 
     def __init__(self, instr: HEDMInstrument):
-        self._groups = []
-        for panel in instr.detectors.values():
-            if panel.group is not None and panel.group not in self._groups:
-                self._groups.append(panel.group)
-
+        self._groups = instr.detector_groups
         self.reset()
 
     def reset(self):
@@ -96,7 +95,7 @@ def reset_params(self):
             }
 
     def reset_rotation_center(self):
-        self._rotation_center = RotationCenter.instrument_mean_center
+        self._rotation_center = RotationCenter.group_mean_center
 
     @property
     def params(self) -> dict:
@@ -110,6 +109,20 @@ def rotation_center(self):
     def rotation_center(self, v: RotationCenter):
         self._rotation_center = v
 
+    def center_of_rotation(
+        self,
+        instr: HEDMInstrument,
+        group: str,
+    ) -> np.ndarray:
+        if self.rotation_center == RotationCenter.instrument_mean_center:
+            return instr.mean_detector_center
+        elif self.rotation_center == RotationCenter.lab_origin:
+            return np.array([0.0, 0.0, 0.0])
+        elif self.rotation_center == RotationCenter.group_mean_center:
+            return instr.mean_group_center(group)
+
+        raise NotImplementedError(self.rotation_center)
+
 
 class RelativeConstraintsSystem(RelativeConstraints):
     type = RelativeConstraintsType.system
@@ -160,7 +173,7 @@ def create_relative_constraints(type: RelativeConstraintsType,
     }
 
     kwargs = {}
-    if type == 'System':
+    if type.value == 'Group':
         kwargs['instr'] = instr
 
     return types[type.value](**kwargs)

hexrd/instrument/hedm_instrument.py

@@ -80,6 +80,7 @@
 from .cylindrical_detector import CylindricalDetector
 from .detector import (
     beam_energy_DFLT,
+    Detector,
     max_workers_DFLT,
 )
 from .planar_detector import PlanarDetector
@@ -797,21 +798,25 @@ def mean_detector_center(self) -> np.ndarray:
         centers = np.array([panel.tvec for panel in self.detectors.values()])
         return centers.sum(axis=0) / len(centers)
 
+    def mean_group_center(self, group: str) -> np.ndarray:
+        """Return the mean center for detectors belonging to a group"""
+        centers = np.array([
+            x.tvec for x in self.detectors_in_group(group).values()
+        ])
+        return centers.sum(axis=0) / len(centers)
+
     @property
-    def mean_group_centers(self) -> dict[str, np.ndarray]:
-        """Return the mean center for every group of detectors"""
-        centers = {}
+    def detector_groups(self) -> list[str]:
+        groups = []
         for panel in self.detectors.values():
-            if panel.group is None:
-                # Skip over panels without groups
-                continue
-
-            if panel.group not in centers:
-                centers[panel.group] = []
+            group = panel.group
+            if group is not None and group not in groups:
+                groups.append(group)
 
-            centers[panel.group].append(panel.tvec)
+        return groups
 
-        return {k: v.sum(axis=0) / len(v) for k, v in centers.items()}
+    def detectors_in_group(self, group: str) -> dict[str, Detector]:
+        return {k: v for k, v in self.detectors.items() if v.group == group}
 
     # properties for physical size of rectangular detector
     @property