diff --git a/docs/changelog.md b/docs/changelog.md index 9409e0d9..1c1bbc08 100644 --- a/docs/changelog.md +++ b/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: diff --git a/src/extra/applications/grating.py b/src/extra/applications/grating.py index 92ccc264..6da170ae 100644 --- a/src/extra/applications/grating.py +++ b/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,6 +123,16 @@ 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 @@ -123,6 +140,7 @@ def setup(self, # 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), diff --git a/tests/conftest.py b/tests/conftest.py index c210cf84..3407b49b 100644 --- a/tests/conftest.py +++ b/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 diff --git a/tests/mockdata/camera.py b/tests/mockdata/camera.py index 5097c4bf..4ef7d1d3 100644 --- a/tests/mockdata/camera.py +++ b/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',) diff --git a/tests/test_applications_grating.py b/tests/test_applications_grating.py index 31b0184e..7ed1a9b4 100644 --- a/tests/test_applications_grating.py +++ b/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,11 +100,22 @@ 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() @@ -103,9 +123,15 @@ def test_reading_grating1d(mock_sqs_grating_calibration_run, tmp_path): 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)