import json from importlib.resources import files import numpy as np import pytest from pydantic import ValidationError from aarecommon.math.beam_center import BeamCenterFromDetectorStage, fit_beam_centre_model from aarecommon.models.beam_centre import BeamCentre, BeamCentreMeasurements BEAM_CENTRES_JSON = files("aarecommon.config") / "beamline_configs" / "beam_centres.json" # Ground truth used to synthesise noiseless measurements: # beam_px = TRUE_MATRIX @ (det_z_mm, det_y_mm) + TRUE_OFFSET # Rows index the pixel axes (x, y), columns the stage axes (z, y). TRUE_MATRIX = np.array([[0.35, -1.20], [-0.80, 26.50]]) TRUE_OFFSET = np.array([2100.0, 550.0]) def synth(det_z_mm, det_y_mm): """Noiseless (beam_x_px, beam_y_px) for the ground-truth model above.""" stage = np.column_stack([np.ravel(det_z_mm), np.ravel(det_y_mm)]).astype(float) beam = stage @ TRUE_MATRIX.T + TRUE_OFFSET return beam[:, 0], beam[:, 1] @pytest.fixture def stage_grid(): """A 2D sweep of both stage axes -- enough rank to pin down the affine model.""" det_z, det_y = np.meshgrid([170.0, 250.0, 320.0, 400.0], [62.0, 100.0, 120.0]) return det_z.ravel(), det_y.ravel() @pytest.fixture def synthetic_fit(stage_grid): det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) return fit_beam_centre_model(det_z_mm, det_y_mm, beam_x_px, beam_y_px) @pytest.fixture(scope="module") def beam_centres_config(): return json.loads(BEAM_CENTRES_JSON.read_text()) @pytest.fixture(scope="module") def x10sa_measured(beam_centres_config): return beam_centres_config["x10sa"]["measured"] # -------------------------------------------------------------------------------------- # fit_beam_centre_model # -------------------------------------------------------------------------------------- def test_fit_recovers_known_matrix(synthetic_fit): np.testing.assert_allclose(synthetic_fit.matrix_px_per_mm, TRUE_MATRIX, atol=1e-9) def test_fit_recovers_known_offset(synthetic_fit, stage_grid): # offset_px is referenced to stage_mean_mm, not to the origin, so undo the centring # before comparing against the absolute ground-truth offset. det_z_mm, det_y_mm = stage_grid expected = TRUE_OFFSET + TRUE_MATRIX @ np.array([det_z_mm.mean(), det_y_mm.mean()]) np.testing.assert_allclose(synthetic_fit.offset_px, expected, atol=1e-9) def test_offset_px_is_the_prediction_at_the_stage_mean(synthetic_fit): beam_x_px, beam_y_px = synthetic_fit.predict(*synthetic_fit.stage_mean_mm) np.testing.assert_allclose([beam_x_px[0], beam_y_px[0]], synthetic_fit.offset_px, atol=1e-9) def test_stage_mean_is_mean_of_inputs(synthetic_fit, stage_grid): det_z_mm, det_y_mm = stage_grid np.testing.assert_allclose(synthetic_fit.stage_mean_mm, [det_z_mm.mean(), det_y_mm.mean()]) def test_noiseless_fit_has_zero_residuals(synthetic_fit, stage_grid): assert synthetic_fit.residuals_px.shape == (stage_grid[0].size, 2) np.testing.assert_allclose(synthetic_fit.residuals_px, 0.0, atol=1e-9) def test_fit_accepts_lists_not_just_arrays(stage_grid): det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) from_lists = fit_beam_centre_model( det_z_mm.tolist(), det_y_mm.tolist(), beam_x_px.tolist(), beam_y_px.tolist() ) np.testing.assert_allclose(from_lists.matrix_px_per_mm, TRUE_MATRIX, atol=1e-9) def test_fit_ravels_nested_inputs(stage_grid): """2D stage inputs are flattened rather than rejected.""" det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) nested = fit_beam_centre_model( det_z_mm.reshape(3, 4), det_y_mm.reshape(3, 4), beam_x_px, beam_y_px ) np.testing.assert_allclose(nested.matrix_px_per_mm, TRUE_MATRIX, atol=1e-9) @pytest.mark.parametrize("bad_value", [np.nan, np.inf, -np.inf]) @pytest.mark.parametrize("column", ["det_z_mm", "det_y_mm", "beam_x_px", "beam_y_px"]) def test_non_finite_samples_are_dropped(stage_grid, column, bad_value): """A poisoned row in any of the four columns must not contaminate the fit.""" det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) columns = { "det_z_mm": det_z_mm.copy(), "det_y_mm": det_y_mm.copy(), "beam_x_px": beam_x_px.copy(), "beam_y_px": beam_y_px.copy(), } columns[column][0] = bad_value fit = fit_beam_centre_model(**columns) np.testing.assert_allclose(fit.matrix_px_per_mm, TRUE_MATRIX, atol=1e-9) # The dropped row is gone from the residuals and from the stage mean. assert fit.residuals_px.shape == (det_z_mm.size - 1, 2) np.testing.assert_allclose(fit.stage_mean_mm, [det_z_mm[1:].mean(), det_y_mm[1:].mean()]) def test_fit_rejects_mismatched_lengths(): with pytest.raises(ValueError, match="same length"): fit_beam_centre_model([170.0, 250.0, 320.0], [62.0, 62.0, 62.0], [1.0, 2.0], [3.0, 4.0]) def test_fit_rejects_too_few_samples(): with pytest.raises(ValueError, match="at least 3 finite samples"): fit_beam_centre_model([170.0, 250.0], [62.0, 100.0], [2090.0, 2095.0], [2200.0, 3250.0]) def test_fit_rejects_too_few_samples_after_dropping_non_finite(stage_grid): """12 samples, but only 2 survive the finiteness filter.""" det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) beam_x_px = beam_x_px.copy() beam_x_px[2:] = np.nan with pytest.raises(ValueError, match="at least 3 finite samples"): fit_beam_centre_model(det_z_mm, det_y_mm, beam_x_px, beam_y_px) def test_fit_is_least_squares_on_noisy_data(stage_grid): """Residuals of a proper LSQ solution are orthogonal to every design column.""" det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) noise = np.linspace(-3.0, 3.0, det_z_mm.size) fit = fit_beam_centre_model(det_z_mm, det_y_mm, beam_x_px + noise, beam_y_px - noise) design = BeamCenterFromDetectorStage.design_matrix(fit.stage_mean_mm, det_z_mm, det_y_mm) np.testing.assert_allclose(design.T @ fit.residuals_px, 0.0, atol=1e-9) def test_fit_is_insensitive_to_sample_order(stage_grid): det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) order = np.arange(det_z_mm.size)[::-1] forward = fit_beam_centre_model(det_z_mm, det_y_mm, beam_x_px, beam_y_px) reversed_ = fit_beam_centre_model( det_z_mm[order], det_y_mm[order], beam_x_px[order], beam_y_px[order] ) np.testing.assert_allclose(forward.matrix_px_per_mm, reversed_.matrix_px_per_mm, atol=1e-9) np.testing.assert_allclose(forward.offset_px, reversed_.offset_px, atol=1e-9) # -------------------------------------------------------------------------------------- # BeamCenterFromDetectorStage.design_matrix / .predict # -------------------------------------------------------------------------------------- def test_design_matrix_is_stage_centred_with_intercept(): det_z_mm = np.array([170.0, 250.0, 400.0]) det_y_mm = np.array([62.0, 100.0, 120.0]) design = BeamCenterFromDetectorStage.design_matrix(np.array([250.0, 100.0]), det_z_mm, det_y_mm) np.testing.assert_allclose(design, [[-80.0, -38.0, 1.0], [0.0, 0.0, 1.0], [150.0, 20.0, 1.0]]) def test_predict_round_trips_the_fitted_measurements(synthetic_fit, stage_grid): det_z_mm, det_y_mm = stage_grid expected_x, expected_y = synth(det_z_mm, det_y_mm) beam_x_px, beam_y_px = synthetic_fit.predict(det_z_mm, det_y_mm) np.testing.assert_allclose(beam_x_px, expected_x, atol=1e-9) np.testing.assert_allclose(beam_y_px, expected_y, atol=1e-9) def test_predict_extrapolates_beyond_the_measured_range(synthetic_fit): """The model is affine, so points outside the calibration grid stay exact.""" expected_x, expected_y = synth([600.0], [20.0]) beam_x_px, beam_y_px = synthetic_fit.predict(600.0, 20.0) np.testing.assert_allclose(beam_x_px, expected_x, atol=1e-9) np.testing.assert_allclose(beam_y_px, expected_y, atol=1e-9) def test_predict_on_scalars_returns_length_one_arrays(synthetic_fit): beam_x_px, beam_y_px = synthetic_fit.predict(250.0, 100.0) assert beam_x_px.shape == (1,) assert beam_y_px.shape == (1,) def test_predict_broadcasts_scalar_against_array(synthetic_fit): det_z_mm = np.array([170.0, 250.0, 400.0]) beam_x_px, beam_y_px = synthetic_fit.predict(det_z_mm, 100.0) expected_x, expected_y = synth(det_z_mm, np.full(3, 100.0)) np.testing.assert_allclose(beam_x_px, expected_x, atol=1e-9) np.testing.assert_allclose(beam_y_px, expected_y, atol=1e-9) def test_predict_flattens_2d_stage_inputs(synthetic_fit): """Output is always 1D, one entry per input point, regardless of input shape.""" det_z_mm = np.array([[170.0, 250.0], [320.0, 400.0]]) det_y_mm = np.array([[62.0, 100.0], [100.0, 120.0]]) beam_x_px, beam_y_px = synthetic_fit.predict(det_z_mm, det_y_mm) expected_x, expected_y = synth(det_z_mm, det_y_mm) assert beam_x_px.shape == (4,) np.testing.assert_allclose(beam_x_px, expected_x, atol=1e-9) np.testing.assert_allclose(beam_y_px, expected_y, atol=1e-9) def test_predict_accepts_lists(synthetic_fit): beam_x_px, beam_y_px = synthetic_fit.predict([170.0, 400.0], [62.0, 120.0]) expected_x, expected_y = synth([170.0, 400.0], [62.0, 120.0]) np.testing.assert_allclose(beam_x_px, expected_x, atol=1e-9) np.testing.assert_allclose(beam_y_px, expected_y, atol=1e-9) # -------------------------------------------------------------------------------------- # BeamCentreMeasurements / BeamCentre # -------------------------------------------------------------------------------------- def test_measurements_accept_canonical_field_names(stage_grid): det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) measurements = BeamCentreMeasurements( det_z_mm=det_z_mm.tolist(), det_y_mm=det_y_mm.tolist(), beam_x_px=beam_x_px.tolist(), beam_y_px=beam_y_px.tolist(), ) assert measurements.det_z_mm == det_z_mm.tolist() def test_measurements_accept_json_aliases(x10sa_measured): measurements = BeamCentreMeasurements.model_validate(x10sa_measured) assert measurements.det_z_mm == x10sa_measured["det_z"] assert measurements.beam_y_px == x10sa_measured["beam_y"] def test_measurements_reject_missing_column(x10sa_measured): incomplete = {k: v for k, v in x10sa_measured.items() if k != "beam_y"} with pytest.raises(ValidationError): BeamCentreMeasurements.model_validate(incomplete) def test_beam_centre_fits_model_from_measurements(stage_grid): det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) beam_centre = BeamCentre.model_validate( { "measured": { "det_z_mm": det_z_mm.tolist(), "det_y_mm": det_y_mm.tolist(), "beam_x_px": beam_x_px.tolist(), "beam_y_px": beam_y_px.tolist(), } } ) assert isinstance(beam_centre.model, BeamCenterFromDetectorStage) np.testing.assert_allclose(beam_centre.model.matrix_px_per_mm, TRUE_MATRIX, atol=1e-9) def test_beam_centre_keeps_an_explicitly_supplied_model(synthetic_fit, stage_grid): """A caller-provided model must be used verbatim, not silently refitted.""" det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) beam_centre = BeamCentre.model_validate( { "measured": { "det_z_mm": det_z_mm.tolist(), "det_y_mm": det_y_mm.tolist(), "beam_x_px": (beam_x_px + 1000.0).tolist(), "beam_y_px": beam_y_px.tolist(), }, "model": synthetic_fit, } ) np.testing.assert_allclose(beam_centre.model.offset_px, synthetic_fit.offset_px) def test_beam_centre_revalidation_is_idempotent(stage_grid): det_z_mm, det_y_mm = stage_grid beam_x_px, beam_y_px = synth(det_z_mm, det_y_mm) beam_centre = BeamCentre.model_validate( { "measured": { "det_z_mm": det_z_mm.tolist(), "det_y_mm": det_y_mm.tolist(), "beam_x_px": beam_x_px.tolist(), "beam_y_px": beam_y_px.tolist(), } } ) again = BeamCentre.model_validate(beam_centre) np.testing.assert_allclose(again.model.offset_px, beam_centre.model.offset_px) def test_beam_centre_propagates_fit_errors(): with pytest.raises(ValueError, match="same length"): BeamCentre.model_validate( { "measured": { "det_z_mm": [170.0, 250.0, 320.0], "det_y_mm": [62.0, 62.0, 62.0], "beam_x_px": [2090.0, 2095.0], "beam_y_px": [2200.0, 2210.0], } } ) # -------------------------------------------------------------------------------------- # beam_centres.json # -------------------------------------------------------------------------------------- def test_config_file_is_importable_package_data(): """Guards the pyproject package-data glob -- a missing json entry breaks the wheel.""" assert BEAM_CENTRES_JSON.is_file() @pytest.mark.parametrize("beamline", ["x10sa", "x06da"]) def test_config_has_measurements(beam_centres_config, beamline): assert beamline in beam_centres_config assert set(beam_centres_config[beamline]["measured"]) == {"det_z", "det_y", "beam_x", "beam_y"} def test_config_columns_are_equal_length_and_finite(beam_centres_config): for beamline, entry in beam_centres_config.items(): columns = entry["measured"] lengths = {name: len(values) for name, values in columns.items()} assert len(set(lengths.values())) == 1, f"{beamline} has ragged columns: {lengths}" for name, values in columns.items(): assert np.isfinite(values).all(), f"{beamline}.{name} has non-finite entries" def test_config_has_enough_samples_to_fit(beam_centres_config): for beamline, entry in beam_centres_config.items(): assert len(entry["measured"]["det_z"]) >= 3, f"{beamline} cannot be fitted" def test_config_sweeps_both_stage_axes(beam_centres_config): """A single-axis sweep leaves the fit rank-deficient and the matrix meaningless.""" for beamline, entry in beam_centres_config.items(): columns = entry["measured"] assert len(set(columns["det_z"])) > 1, f"{beamline} does not sweep det_z" assert len(set(columns["det_y"])) > 1, f"{beamline} does not sweep det_y" @pytest.fixture(scope="module") def x10sa_beam_centre(x10sa_measured): return BeamCentre.model_validate({"measured": x10sa_measured}) def test_x10sa_fit_matches_a_direct_fit(x10sa_beam_centre, x10sa_measured): direct = fit_beam_centre_model( det_z_mm=x10sa_measured["det_z"], det_y_mm=x10sa_measured["det_y"], beam_x_px=x10sa_measured["beam_x"], beam_y_px=x10sa_measured["beam_y"], ) np.testing.assert_allclose(x10sa_beam_centre.model.matrix_px_per_mm, direct.matrix_px_per_mm) np.testing.assert_allclose(x10sa_beam_centre.model.offset_px, direct.offset_px) def test_x10sa_beam_y_tracks_det_y(x10sa_beam_centre): """det_y moves the detector across the beam: ~26.5 px/mm on the y pixel axis.""" d_beam_y_d_det_y = x10sa_beam_centre.model.matrix_px_per_mm[1, 1] assert d_beam_y_d_det_y == pytest.approx(26.5, abs=1.0) def test_x10sa_is_nearly_decoupled_along_the_beam(x10sa_beam_centre): """det_z is along the beam, so it barely moves the pierce point in either axis.""" matrix = x10sa_beam_centre.model.matrix_px_per_mm assert abs(matrix[0, 0]) < 0.5 assert abs(matrix[1, 0]) < 0.5 def test_x10sa_offset_sits_on_the_detector(x10sa_beam_centre): """offset_px is the pierce point at the mean stage position -- must be plausible.""" beam_x_px, beam_y_px = x10sa_beam_centre.model.offset_px assert 1500.0 < beam_x_px < 2500.0 assert 2500.0 < beam_y_px < 3500.0 def test_x10sa_residuals_are_within_measurement_scatter(x10sa_beam_centre): """The affine model should explain the real data to a few tens of pixels.""" rms_px = np.sqrt((x10sa_beam_centre.model.residuals_px**2).mean(axis=0)) assert rms_px[0] < 30.0, f"beam_x residual too large: {rms_px[0]:.1f} px" assert rms_px[1] < 30.0, f"beam_y residual too large: {rms_px[1]:.1f} px" def test_x10sa_predict_reproduces_the_measurements(x10sa_beam_centre, x10sa_measured): beam_x_px, beam_y_px = x10sa_beam_centre.model.predict( x10sa_measured["det_z"], x10sa_measured["det_y"] ) np.testing.assert_allclose(beam_x_px, x10sa_measured["beam_x"], atol=60.0) np.testing.assert_allclose(beam_y_px, x10sa_measured["beam_y"], atol=60.0) @pytest.fixture(scope="module") def x06da_measured(beam_centres_config): return beam_centres_config["x06da"]["measured"] @pytest.fixture(scope="module") def x06da_beam_centre(x06da_measured): return BeamCentre.model_validate({"measured": x06da_measured}) def test_x06da_beam_y_tracks_det_y(x06da_beam_centre): """~5.35 px/mm -- a coarser pixel pitch than x10sa, hence the smaller gradient.""" d_beam_y_d_det_y = x06da_beam_centre.model.matrix_px_per_mm[1, 1] assert d_beam_y_d_det_y == pytest.approx(5.35, abs=0.2) def test_x06da_is_nearly_decoupled_along_the_beam(x06da_beam_centre): matrix = x06da_beam_centre.model.matrix_px_per_mm assert abs(matrix[0, 0]) < 0.1 assert abs(matrix[1, 0]) < 0.1 def test_x06da_offset_sits_on_the_detector(x06da_beam_centre): beam_x_px, beam_y_px = x06da_beam_centre.model.offset_px assert 500.0 < beam_x_px < 1000.0 assert 800.0 < beam_y_px < 1500.0 def test_x06da_residuals_are_small(x06da_beam_centre): """x06da was measured cleanly -- the affine model holds to ~1 px.""" rms_px = np.sqrt((x06da_beam_centre.model.residuals_px**2).mean(axis=0)) assert rms_px[0] < 2.0, f"beam_x residual too large: {rms_px[0]:.2f} px" assert rms_px[1] < 2.0, f"beam_y residual too large: {rms_px[1]:.2f} px" def test_x06da_predict_reproduces_the_measurements(x06da_beam_centre, x06da_measured): beam_x_px, beam_y_px = x06da_beam_centre.model.predict( x06da_measured["det_z"], x06da_measured["det_y"] ) np.testing.assert_allclose(beam_x_px, x06da_measured["beam_x"], atol=5.0) np.testing.assert_allclose(beam_y_px, x06da_measured["beam_y"], atol=5.0)