he11lib/tests/test_geometry.py

import numpy as np
import pytest

from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration


def test_camera_model_tolerance_accepts_zero_and_positive():
    CameraModelTolerance(
        focal_length_px=0.0,
        position=(0.0, 0.0, 0.0),
        orientation_deg=(1.0, 2.0, 3.0),
        principal_point=(0.5, 0.5),
    )  # should not raise


def test_camera_model_tolerance_rejects_negative_scalar_field():
    with pytest.raises(ValueError, match="focal_length_px"):
        CameraModelTolerance(
            focal_length_px=-1.0,
            position=(0.0, 0.0, 0.0),
            orientation_deg=(0.0, 0.0, 0.0),
        )


def test_camera_model_tolerance_rejects_negative_tuple_component():
    with pytest.raises(ValueError, match="position"):
        CameraModelTolerance(
            focal_length_px=1.0,
            position=(0.0, -0.5, 0.0),
            orientation_deg=(0.0, 0.0, 0.0),
        )


def make_on_axis_camera(focal_length_px=2000.0, camera_z=-2.0):
    return CameraModel(
        focal_length_px=focal_length_px,
        position=(0.0, 0.0, camera_z),
        orientation_deg=(0.0, 0.0, 0.0),
    )


def make_tilted_camera():
    return CameraModel(
        focal_length_px=2000.0,
        position=(0.05, -0.03, -2.0),
        orientation_deg=(8.0, -5.0, 3.0),
    )


@pytest.mark.parametrize(
    "camera",
    [make_on_axis_camera(), make_tilted_camera()],
    ids=["on_axis", "tilted_off_center"],
)
@pytest.mark.parametrize("z", [0.3, 0.5, 0.8])
def test_projection_round_trip_recovers_pixel_grid(camera, z):
    image_shape = (41, 41)
    calib = GeometryCalibration(camera)

    x, y = calib.physical_coordinates(image_shape, z)
    row, col = calib.pixel_coordinates(x, y, z)

    rows, cols = image_shape
    row_idx = np.arange(rows) - rows // 2
    col_idx = np.arange(cols) - cols // 2
    expected_col, expected_row = np.meshgrid(col_idx, row_idx)

    np.testing.assert_allclose(row, expected_row, atol=1e-6)
    np.testing.assert_allclose(col, expected_col, atol=1e-6)


def test_keystone_regression_uniform_for_on_axis_camera():
    # A camera with zero orientation, centered on the beam axis, produces
    # uniform pixel spacing for evenly spaced physical points (no keystoning).
    camera = make_on_axis_camera()
    calib = GeometryCalibration(camera)
    z = 0.5

    xs = np.array([-0.02, -0.01, 0.0, 0.01, 0.02])
    ys = np.zeros_like(xs)
    _, col = calib.pixel_coordinates(xs, ys, z)

    spacings = np.diff(col)
    np.testing.assert_allclose(spacings, spacings[0], rtol=1e-6)


def test_keystone_regression_nonuniform_for_tilted_camera():
    # A tilted/off-axis camera produces non-uniform pixel spacing for the
    # same evenly spaced physical points -- genuine keystoning.
    camera = make_tilted_camera()
    calib = GeometryCalibration(camera)
    z = 0.5

    xs = np.array([-0.02, -0.01, 0.0, 0.01, 0.02])
    ys = np.zeros_like(xs)
    _, col = calib.pixel_coordinates(xs, ys, z)

    spacings = np.diff(col)
    assert not np.allclose(spacings, spacings[0], rtol=1e-3)


def test_pixel_coordinates_raises_when_point_behind_camera():
    camera = CameraModel(
        focal_length_px=2000.0,
        position=(0.0, 0.0, 10.0),
        orientation_deg=(0.0, 0.0, 0.0),
    )
    calib = GeometryCalibration(camera)

    with pytest.raises(ValueError):
        calib.pixel_coordinates(np.array([0.0]), np.array([0.0]), z=0.5)


def test_physical_coordinates_raises_when_plane_behind_camera():
    # Camera sits downstream of the target plane and looks further
    # downstream (boresight = +z world) -- the z=0.5 plane is behind it.
    camera = CameraModel(
        focal_length_px=2000.0,
        position=(0.0, 0.0, 10.0),
        orientation_deg=(0.0, 0.0, 0.0),
    )
    calib = GeometryCalibration(camera)

    with pytest.raises(ValueError):
        calib.physical_coordinates((21, 21), z=0.5)


def test_physical_coordinates_raises_when_edge_on():
    # Pitch=90 deg points the boresight along world -y, making the
    # z=const target plane edge-on (parallel to the view direction).
    camera = CameraModel(
        focal_length_px=2000.0,
        position=(0.0, 0.0, -2.0),
        orientation_deg=(0.0, 90.0, 0.0),
    )
    calib = GeometryCalibration(camera)

    with pytest.raises(ValueError):
        calib.physical_coordinates((41, 41), z=0.5)


def test_effective_pixel_scale_matches_on_axis_focal_length():
    focal_length_px = 2000.0
    camera_z = -2.0
    z = 0.5
    camera = make_on_axis_camera(focal_length_px=focal_length_px, camera_z=camera_z)
    calib = GeometryCalibration(camera)

    scale = calib.effective_pixel_scale((41, 41), z)
    expected = (z - camera_z) / focal_length_px
    assert scale == pytest.approx(expected, rel=1e-6)


def test_effective_pixel_scale_raises_for_image_too_small():
    camera = make_on_axis_camera()
    calib = GeometryCalibration(camera)
    z = 0.5

    with pytest.raises(ValueError, match="image_shape must be at least 2x2"):
        calib.effective_pixel_scale((1, 1), z)