Include motor movements in the fit.

docs/changelog.md

@@ -76,6 +76,8 @@ Added:
 - [Grating1DCalibration][extra.applications.Grating1DCalibration] can use the calibration mask,
   removes the background subtraction (leaving it to the calibration DB),
   and use robust tools to improve the fit quality (!510).
+- [Grating1DCalibration][extra.applications.Grating1DCalibration] includes the motor position
+  in the calibration fit (!515).
 
 Fixed:
 

src/extra/applications/grating.py

@@ -74,13 +74,17 @@ def __init__(self, offset: Optional[int]=None, min_pixel: int=0, max_pixel: int=
                             "max_pixel",
                             "e0",
                             "slope",
+                            "slope_motor",
                             "energy_axis",
                             "calibration_energies",
                             "calibration_data",
+                            "calibration_motor",
                             "grating_source",
                             "grating_key",
                             "sources",
                             "grating_mask_key",
+                            "grating_motor_source",
+                            "grating_motor_key",
                             "sigma",
                             "_version",
                            ]
@@ -90,6 +94,7 @@ def setup(self,
               scan: Scan,
               pulses: XrayPulses,
               grating_mask: Optional[KeyData]=None,
+              grating_motor: Optional[KeyData]=None,
               ):
         """
         Setup calibration.
@@ -103,6 +108,8 @@ def setup(self,
                 Example: `XrayPulses(run)`
           grating_mask: Grating mask from the calibration system.
                    Example: `signal_run["SQS_EXP_GH2-2/CORR/RECEIVER:daqOutput", "data.mask"]`
+          grating_motor: KeyData corresponding to the motor position.
+               Example: `signal_run['SQS_DIAG3_SCAM/MOTOR/ST_AXIS_X', 'encoderPosition.value']`
         """
         self.grating_source = grating_signal.source
         self.grating_key = grating_signal.key
@@ -116,13 +123,24 @@ def setup(self,
                         self.grating_source,
                        ]
 
+        if grating_motor is not None:
+            self.grating_motor_source = grating_motor.source
+            self.grating_motor_key = grating_motor.key
+            self._grating_motor = grating_motor
+            self.sources += [self.grating_motor_source]
+        else:
+            self.grating_motor_source = ""
+            self.grating_motor_key = ""
+            self._grating_motor = None
+
         # create scan object
         self._scan = scan
         self.calibration_energies = self._scan.positions
 
         # outputs
         self.e0 = 0
         self.slope = 0
+        self.slope_motor = 0
         self.energy_axis = None
 
         if self.offset is None:
@@ -232,10 +250,18 @@ def load_data(self):
                      grating=self._grating_signal,
                      mask=self._grating_mask,
                      )
+        fn_motor = partial(calc_mean,
+                     scan=self._scan,
+                     grating=self._grating_motor,
+                     )
         energy_ids = np.arange(len(self.calibration_energies))
         # average data in each mono scan bin
         with ProcessPoolExecutor() as p:
             data = np.stack(list(p.map(fn, energy_ids)), axis=0)
+            if self.grating_motor_source != "":
+                data_motor = np.stack(list(p.map(fn_motor, energy_ids)), axis=0)
+            else:
+                data_motor = np.zeros((data.shape[0]))
         # skip offset and collect pulse data each pulse_period samples only
         self.calibration_data = data[:, self.offset::self.pulse_period, self.min_pixel:self.max_pixel]
         # apply mask
@@ -245,6 +271,7 @@ def load_data(self):
         if self.sigma > 0:
             self.calibration_data = gaussian_filter(np.nan_to_num(self.calibration_data), axes=-1, sigma=self.sigma)
         self.calibration_mask = np.ones(self.calibration_data.shape[0], dtype=bool)
+        self.calibration_motor = data_motor
 
     def mask_calibration_point(self, energy: float, mask: bool=False, tol: float=1.0):
         """
@@ -261,16 +288,23 @@ def mask_calibration_point(self, energy: float, mask: bool=False, tol: float=1.0
 
     def fit(self):
         """Fit line."""
+        #from scipy.stats import linregress
         mask = self.calibration_mask
         sample = np.arange(self.calibration_data.shape[-1])
         sample_mode = np.nanargmax(self.calibration_data, axis=-1)
-        x = sample_mode[mask]
+        motor_position = self.calibration_motor
+        #sample_mode = np.sum(self.calibration_data*sample, axis=-1)/np.sum(self.calibration_data, axis=-1)
+        #res = linregress(sample_mode[mask], self.calibration_energies[mask])
+        #self.slope = res.slope
+        #self.e0 = res.intercept
+        x = np.stack((sample_mode[mask], motor_position[mask]), axis=1)
         y = self.calibration_energies[mask]
         model = RANSACRegressor(estimator=LinearRegression(), random_state=42)
-        model.fit(x[:,np.newaxis], y[:, np.newaxis])
+        model.fit(x, y[:, np.newaxis])
         self.slope = model.estimator_.coef_[0,0]
+        self.slope_motor = model.estimator_.coef_[0,1]
         self.e0 = model.estimator_.intercept_[0]
-        self.energy_axis = self.e0 + self.slope*sample
+        self.energy_axis = self.e0 + self.slope*sample + self.slope_motor*motor_position.mean()
 
     def plot(self):
         """
@@ -280,24 +314,42 @@ def plot(self):
         plt.figure(figsize=(10, 8))
         sample = np.arange(self.calibration_data.shape[-1])
         sample_mode = np.nanargmax(self.calibration_data, axis=-1)
+        motor_position = self.calibration_motor
         plt.plot(sample, self.energy_axis, lw=2, label="Fit")
         plt.xlabel("Pixel")
         plt.ylabel("Energy [eV]")
+        plt.scatter(sample_mode, self.e0 + self.slope*sample_mode + self.slope_motor*motor_position,
+                    s=200, marker='o', facecolors='w', edgecolors="k", label="Prediction (w. motor)")
         plt.scatter(sample_mode, self.calibration_energies, s=200, marker='x', c='r', label="Data")
         plt.legend(frameon=False)
         plt.grid()
         plt.show()
 
-    def apply(self, run: DataCollection, load_all: bool=True) -> xr.Dataset:
+    def apply(self, run: DataCollection, load_all: bool=True, assume_motor=None) -> xr.Dataset:
         """
         Apply calibration to a new analysis run.
         It is assumed it contains the same settings.
 
         Args:
           run: Input run.
-          load_all: If True, load all data in memory at once. This is faster, but uses more memory.
+          load_all: If True, load all data in memory at once. This is faste, but uses more memory.
                     Disable if not enough memory is available.
+          assume_motor: Assume the grating motor is fixed at this position.
         """
+        if self.grating_motor_source != "":
+            if assume_motor is not None:
+                logging.info(f"Assuming motor position at {assume_motor}")
+                motor_position = assume_motor
+            else:
+                out_data_motor = run[self.grating_motor_source, self.grating_motor_key].xarray()
+                logging.info(f"Motor position being assumed fixed at {out_data_motor.mean()}. "
+                             f"I detect a rms variation of {out_data_motor.std()}. "
+                             f"This should be compatible with zero!")
+                motor_position = out_data_motor.mean().to_numpy()
+        else:
+            motor_position = 0.0
+        sample = np.arange(self.calibration_data.shape[-1])
+        energy = self.e0 + self.slope*sample + self.slope_motor*motor_position
         # do it per train to avoid memory overflow
         pulse_period = self.get_pulse_period(XrayPulses(run))
         if load_all:
@@ -327,7 +379,10 @@ def apply(self, run: DataCollection, load_all: bool=True) -> xr.Dataset:
                 trainId += [tid]
                 out_data += [d]
             out_data = np.stack(out_data, axis=0)
-        energy = self.energy_axis
+        #energy = self.energy_axis
+        if self.slope < 0:
+            energy = energy[::-1]
+            out_data = out_data[:, :, ::-1]
         out_data = xr.DataArray(data=out_data,
                             dims=('trainId', 'pulseIndex', 'energy'),
                             coords=dict(trainId=np.array(trainId),

tests/conftest.py

@@ -320,7 +320,8 @@ def mock_sqs_grating_calibration_directory():
         MonoMdl('SA3_XTD10_MONO/MDL/PHOTON_ENERGY', energy_data=energy),
         CameraWithData("SQS_DIAG3_BIU/CAM/CAM_6", data=data2d),
         GotthardIIWithData("SQS_EXP_GH2-2/CORR/RECEIVER", data=data_gh,
-                           mask=np.zeros_like(data_gh)),
+                           mask=np.zeros_like(data_gh),
+                           motor=np.zeros(data_gh.shape[0])),
               ]
 
     # for tests

tests/mockdata/camera.py

@@ -8,11 +8,13 @@ class GotthardIIWithData(DeviceBase):
     instrument_keys = [
         ("adc", "f4", (12,1000)),
         ("mask", "i4", (12,1000)),
+        ("motor", "f4", ()),
     ]
 
-    def __init__(self, *args, data, mask, **kwargs):
+    def __init__(self, *args, data, mask, motor, **kwargs):
         self.data = data
         self.mask = mask
+        self.motor = motor
         super().__init__(*args, **kwargs)
 
     def write_instrument(self, f):
@@ -22,6 +24,8 @@ def write_instrument(self, f):
         ds[:] = self.data
         ds = f[f'INSTRUMENT/{self.device_id}:daqOutput/data/mask']
         ds[:] = self.mask
+        ds = f[f'INSTRUMENT/{self.device_id}:daqOutput/data/motor']
+        ds[:] = self.motor
 
 class CameraWithData(DeviceBase):
     output_channels = ('daqOutput/data',)

tests/test_applications_grating.py

@@ -37,9 +37,14 @@ def test_grating_1d_fit():
     cal.calibration_data = data
     cal.calibration_mask = np.ones(data.shape[0], dtype=bool)
     cal.calibration_unc = np.zeros_like(data)
+    cal.calibration_motor = np.zeros((data.shape[0]))
 
     cal.fit()
 
+    import matplotlib as mpl
+    mpl.use('Agg')
+    cal.plot()
+
     assert np.isclose(cal.e0, true_offset, atol=1e-2, rtol=1e-2)
     assert np.isclose(cal.slope, true_slope, atol=1e-2, rtol=1e-2)
 
@@ -63,6 +68,10 @@ def test_grating_2d_fit():
 
     cal.fit()
 
+    import matplotlib as mpl
+    mpl.use('Agg')
+    cal.plot()
+
     assert np.isclose(cal.e0, true_offset, atol=1e-2, rtol=1e-2)
     assert np.isclose(cal.slope, true_slope, atol=1e-2, rtol=1e-2)
 
@@ -91,21 +100,38 @@ def test_reading_grating1d(mock_sqs_grating_calibration_run, tmp_path):
     final_photon_spectrometer = "SQS_EXP_GH2-2/CORR/RECEIVER:daqOutput"
 
     monochromator_scan = Scan(mock_sqs_grating_calibration_run[monochromator_energy, "actualEnergy"], resolution=1)
+
+    # fit without motor
     grating_calibration = Grating1DCalibration(min_pixel=0, max_pixel=1000)
     grating_calibration.setup(mock_sqs_grating_calibration_run[final_photon_spectrometer, "data.adc"],
                               monochromator_scan,
                               pulses=XrayPulses(mock_sqs_grating_calibration_run),
-                              grating_mask=mock_sqs_grating_calibration_run[final_photon_spectrometer, "data.mask"]
+                              grating_mask=mock_sqs_grating_calibration_run[final_photon_spectrometer, "data.mask"],
+                             )
+
+    # use motor information
+    grating_calibration = Grating1DCalibration(min_pixel=0, max_pixel=1000)
+    grating_calibration.setup(mock_sqs_grating_calibration_run[final_photon_spectrometer, "data.adc"],
+                              monochromator_scan,
+                              pulses=XrayPulses(mock_sqs_grating_calibration_run),
+                              grating_mask=mock_sqs_grating_calibration_run[final_photon_spectrometer, "data.mask"],
+                              grating_motor=mock_sqs_grating_calibration_run[final_photon_spectrometer, "data.motor"]
                              )
     d = tmp_path / "data"
     d.mkdir()
     fpath = str(d / "grating1d_test.h5")
     grating_calibration.to_file(fpath)
     grating_calibration = Grating1DCalibration.from_file(fpath)
 
-    calibrated = grating_calibration.apply(mock_sqs_grating_calibration_run.select_trains(np.s_[10:20]))
+    # with motor
+    calibrated = grating_calibration.apply(mock_sqs_grating_calibration_run.select_trains(np.s_[10:20]), assume_motor=0.0)
 
     assert np.isclose(grating_calibration.e0, 990.0, atol=1e-2, rtol=1e-2)
     assert np.isclose(grating_calibration.slope, 20.0/1000.0, atol=1e-2, rtol=1e-2)
 
+    # without motor
+    calibrated = grating_calibration.apply(mock_sqs_grating_calibration_run.select_trains(np.s_[10:20]))
+
+    assert np.isclose(grating_calibration.e0, 990.0, atol=1e-2, rtol=1e-2)
+    assert np.isclose(grating_calibration.slope, 20.0/1000.0, atol=1e-2, rtol=1e-2)