python/dpt/geometry.py

1"""Right-handed rigid geometry in millimetres, with explicit object-to-world maps.23These immutable host contracts import without Warp. Rotations are row-major4matrices acting on column vectors; an acquisition origin is its first pixel5centre. Small pose transfers belong at the optimisation control boundary.6"""78from __future__ import annotations910import math11from dataclasses import dataclass12from typing import cast1314from dpt.contracts import ContractError, finite_tuple1516type Vector3 = tuple[float, float, float]17type Matrix3 = tuple[float, float, float, float, float, float, float, float, float]18IDENTITY: Matrix3 = (1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)192021def vector3(value: Vector3, name: str) -> Vector3:22    return cast(Vector3, finite_tuple(value, name, minimum=None, length=3))232425def dot(a: Vector3, b: Vector3) -> float:26    return sum(x * y for x, y in zip(a, b, strict=True))272829def cross(a: Vector3, b: Vector3) -> Vector3:30    return (a[1] * b[2] - a[2] * b[1], a[2] * b[0] - a[0] * b[2], a[0] * b[1] - a[1] * b[0])313233def matvec(matrix: Matrix3, vector: Vector3) -> Vector3:34    return cast(35        Vector3, tuple(sum(matrix[3 * i + j] * vector[j] for j in range(3)) for i in range(3))36    )373839def transpose(matrix: Matrix3) -> Matrix3:40    return cast(Matrix3, tuple(matrix[3 * j + i] for i in range(3) for j in range(3)))414243def matmul(left: Matrix3, right: Matrix3) -> Matrix3:44    return cast(45        Matrix3,46        tuple(47            sum(left[3 * i + k] * right[3 * k + j] for k in range(3))48            for i in range(3)49            for j in range(3)50        ),51    )525354def rotation_matrix(value: Matrix3, name: str = "rotation") -> Matrix3:55    value = cast(Matrix3, finite_tuple(value, name, minimum=None, length=9))56    product = matmul(transpose(value), value)57    if max(abs(a - b) for a, b in zip(product, IDENTITY, strict=True)) > 1e-10:58        raise ContractError(f"{name} must be orthonormal to 1e-10")59    columns = [(value[i], value[i + 3], value[i + 6]) for i in range(3)]60    if dot(cross(columns[0], columns[1]), columns[2]) < 0.0:61        raise ContractError(f"{name} must be right-handed")62    return cast(Matrix3, tuple(float(x) for x in value))636465# region book:geometry-rigid-transform66@dataclass(frozen=True, slots=True)67class RigidTransform:68    """T_WO maps object points into world coordinates; translation is in mm."""6970    rotation: Matrix3 = IDENTITY71    translation_mm: Vector3 = (0.0, 0.0, 0.0)7273    def __post_init__(self) -> None:74        object.__setattr__(self, "rotation", rotation_matrix(self.rotation))75        object.__setattr__(self, "translation_mm", vector3(self.translation_mm, "translation"))7677    def point(self, point_mm: Vector3) -> Vector3:78        rotated = matvec(self.rotation, point_mm)79        return cast(Vector3, tuple(rotated[i] + self.translation_mm[i] for i in range(3)))8081    def direction(self, direction: Vector3) -> Vector3:82        return matvec(self.rotation, direction)8384    def inverse(self) -> RigidTransform:85        inverse_rotation = transpose(self.rotation)86        translation = matvec(inverse_rotation, self.translation_mm)87        return RigidTransform(inverse_rotation, cast(Vector3, tuple(-x for x in translation)))8889    def compose(self, right: RigidTransform) -> RigidTransform:90        """Return self @ right; the right-hand map is applied first."""91        return RigidTransform(92            matmul(self.rotation, right.rotation), self.point(right.translation_mm)93        )9495    def packed(self) -> tuple[float, ...]:96        """Twelve float64 device values: row-major rotation, then translation."""97        return (*self.rotation, *self.translation_mm)9899100# endregion book:geometry-rigid-transform101102103@dataclass(frozen=True, slots=True)104class DetectorGeometry:105    """Finite source-to-pixel segments; u increases columns and v increases rows."""106107    source_mm: Vector3108    origin_mm: Vector3109    u: Vector3110    v: Vector3111    spacing_mm: tuple[float, float]112    shape: tuple[int, int]113114    def __post_init__(self) -> None:115        for name in ("source_mm", "origin_mm", "u", "v"):116            object.__setattr__(self, name, vector3(getattr(self, name), name))117        if abs(dot(self.u, self.u) - 1.0) > 1e-10 or abs(dot(self.v, self.v) - 1.0) > 1e-10:118            raise ContractError("detector basis vectors must have unit length")119        if abs(dot(self.u, self.v)) > 1e-10:120            raise ContractError("detector basis vectors must be perpendicular")121        spacing = finite_tuple(self.spacing_mm, "pixel spacing", positive=True, length=2)122        if len(self.shape) != 2 or any(type(x) is not int or x <= 0 for x in self.shape):123            raise ContractError("detector shape is (height, width), both positive integers")124        object.__setattr__(self, "shape", tuple(self.shape))125        object.__setattr__(self, "spacing_mm", spacing)126        offset: Vector3 = cast(127            Vector3, tuple(self.source_mm[i] - self.origin_mm[i] for i in range(3))128        )129        if abs(dot(offset, cross(self.u, self.v))) <= 1e-12:130            raise ContractError("source must lie outside the detector plane")131        if self.pixels > 2**31 - 1:132            raise ContractError("detector exceeds signed 32-bit indexing")133134    @property135    def pixels(self) -> int:136        return self.shape[0] * self.shape[1]137138    def pixel_centre(self, row: int, column: int) -> Vector3:139        if type(row) is not int or type(column) is not int:140            raise ContractError("pixel indices must be integers")141        if not 0 <= row < self.shape[0] or not 0 <= column < self.shape[1]:142            raise ContractError("pixel index outside detector")143        return cast(144            Vector3,145            tuple(146                self.origin_mm[i]147                + column * self.spacing_mm[0] * self.u[i]148                + row * self.spacing_mm[1] * self.v[i]149                for i in range(3)150            ),151        )152153154def _skew(vector: Vector3) -> Matrix3:155    x, y, z = vector156    return (0.0, -z, y, z, 0.0, -x, -y, x, 0.0)157158159def _coefficients(theta2: float) -> tuple[float, float, float, float, float]:160    # Series avoid cancellation in both the exponential and its derivative.161    if not math.isfinite(theta2):162        raise ContractError("rotation norm exceeds the finite geometry range")163    if theta2 < 1e-2:164        a = 1 - theta2 / 6 + theta2**2 / 120 - theta2**3 / 5040165        b = 0.5 - theta2 / 24 + theta2**2 / 720 - theta2**3 / 40320166        c = 1 / 6 - theta2 / 120 + theta2**2 / 5040 - theta2**3 / 362880167        db = -1 / 12 + theta2 / 180 - theta2**2 / 6720 + theta2**3 / 453600168        dc = -1 / 60 + theta2 / 1260 - theta2**2 / 60480 + theta2**3 / 4989600169        return a, b, c, db, dc170    theta = math.sqrt(theta2)171    sin, cos = math.sin(theta), math.cos(theta)172    return (173        sin / theta,174        (1 - cos) / theta2,175        (theta - sin) / (theta2 * theta),176        (theta * sin - 2 * (1 - cos)) / theta2**2,177        (3 * sin - theta * (2 + cos)) / (theta2**2 * theta),178    )179180181# region book:geometry-right-se3-update182def compose_pose(anchor: RigidTransform, increment: tuple[float, ...]) -> RigidTransform:183    """Apply anchor @ exp(xi^); xi=(tx,ty,tz,rx,ry,rz), mm and radians.184185    Translation is the Lie algebra coordinate, not a separately added world186    displacement. A solver keeps anchor fixed until it discards curvature history.187    """188    increment = finite_tuple(increment, "pose increment", minimum=None, length=6)189    translation: Vector3 = (increment[0], increment[1], increment[2])190    rotation: Vector3 = (increment[3], increment[4], increment[5])191    omega = _skew(rotation)192    omega2 = matmul(omega, omega)193    a, b, c, _, _ = _coefficients(dot(rotation, rotation))194    matrix: Matrix3 = cast(195        Matrix3, tuple(IDENTITY[i] + a * omega[i] + b * omega2[i] for i in range(9))196    )197    velocity: Matrix3 = cast(198        Matrix3, tuple(IDENTITY[i] + b * omega[i] + c * omega2[i] for i in range(9))199    )200    return anchor.compose(RigidTransform(matrix, matvec(velocity, translation)))201202203# endregion book:geometry-right-se3-update204205206def right_jacobian_se3(increment: tuple[float, ...]) -> tuple[tuple[float, ...], ...]:207    """Map fixed-anchor chart perturbations to local right perturbations.208209    exp(xi+d)^ ~= exp(xi^) exp((J_r(xi)d)^). This is an analytic210    derivative of Rodrigues' formula, including translation/rotation coupling.211    """212    transform = compose_pose(RigidTransform(), increment)213    rho: Vector3 = (increment[0], increment[1], increment[2])214    rotation: Vector3 = (increment[3], increment[4], increment[5])215    omega = _skew(rotation)216    omega2 = matmul(omega, omega)217    _, b, c, db, dc = _coefficients(dot(rotation, rotation))218    rt = transpose(transform.rotation)219    velocity: Matrix3 = cast(220        Matrix3, tuple(IDENTITY[i] + b * omega[i] + c * omega2[i] for i in range(9))221    )222    translation_block = matmul(rt, velocity)223    rotation_block = tuple(IDENTITY[i] - b * omega[i] + c * omega2[i] for i in range(9))224    coupling: list[Vector3] = []225    for axis in range(3):226        basis: Vector3 = cast(Vector3, tuple(float(j == axis) for j in range(3)))227        derivative = _skew(basis)228        left, right = matmul(derivative, omega), matmul(omega, derivative)229        dv: Matrix3 = cast(230            Matrix3,231            tuple(232                db * rotation[axis] * omega[i]233                + b * derivative[i]234                + dc * rotation[axis] * omega2[i]235                + c * (left[i] + right[i])236                for i in range(9)237            ),238        )239        coupling.append(matvec(rt, matvec(dv, rho)))240    return tuple(241        tuple(242            translation_block[3 * i + j]243            if i < 3 and j < 3244            else coupling[j - 3][i]245            if i < 3246            else rotation_block[3 * (i - 3) + (j - 3)]247            if j >= 3248            else 0.0249            for j in range(6)250        )251        for i in range(6)252    )253254255def chart_gradient(256    increment: tuple[float, ...], right_gradient: tuple[float, ...]257) -> tuple[float, ...]:258    right_gradient = finite_tuple(right_gradient, "right tangent gradient", minimum=None, length=6)259    jacobian = right_jacobian_se3(increment)260    return tuple(sum(jacobian[i][j] * right_gradient[i] for i in range(6)) for j in range(6))261