he11lib/examples/full_pipeline_example.py

"""End-to-end demonstration of the he11lib reconstruction pipeline.

Simulates a gyrotron beam that is mostly the LG_00 fundamental mode with a
small admixture of LG_01, viewed by a thermal camera at four distances from
the output window through a tilted, off-axis pinhole `CameraModel`. The
camera has an unknown transverse offset/pointing (two independent tilt
angles) and adds sensor noise; the target also has some thermal-diffusion
blur that we correct for. The nominal camera pose and each plane's nominal
`z` are deliberately offset from the (unknown-to-the-reconstructor) ground
truth, simulating realistic calibration/measurement error, and are jointly
refined by the fit within their tolerances. We then reconstruct the mode
purity, beam center/pointing, camera pose, and per-plane z, and plot the
diagnostics.

Run with:

    python examples/full_pipeline_example.py
"""

from __future__ import annotations

import matplotlib.pyplot as plt

from he11lib import (
    BeamReconstructor,
    CameraModel,
    CameraModelTolerance,
    DiffusionDeconvolver,
    GeometryCalibration,
    LGBasis,
    SyntheticBeamGenerator,
    plot_center_trace,
    plot_mode_purity,
    plot_residuals,
)

# --- Known reference beam parameters (from the gyrotron/mode-converter design) ---
W0 = 5e-3          # reference waist radius, meters
Z0 = 0.5           # reference waist location, meters from the output window
WAVELENGTH = 1.76e-3  # radiation wavelength, meters (e.g. a 170 GHz gyrotron)

# --- Ground truth for the synthetic beam (unknown to the reconstructor) ---
TRUE_COEFFICIENTS = {(0, 0): 0.95 + 0j, (0, 1): 0.25 + 0.05j}
TRUE_CENTER = (0.4e-3, -0.3e-3)              # beam offset from the optical axis
TRUE_POINTING_HORIZONTAL_DEG = 0.15          # beam pointing (horizontal tilt)
TRUE_POINTING_VERTICAL_DEG = -0.08           # beam pointing (vertical tilt)
IMAGE_SHAPE = (81, 81)

# A camera positioned well upstream of the target planes and mildly tilted,
# so true perspective projection (keystoning) is in play but the frame
# still comfortably contains the beam at every z below. Calibration only
# gives us a nominal estimate of this pose -- it's refined jointly with
# everything else, within CAMERA_TOLERANCE, because of mechanical vibration
# between calibration and measurement.
PIXEL_SCALE = 4e-4  # meters/pixel, used only to size FOCAL_LENGTH_PX below
CAMERA_DISTANCE = 5.0  # meters upstream of the output window
FOCAL_LENGTH_PX = (CAMERA_DISTANCE + Z0) / PIXEL_SCALE

# NOTE on the small magnitudes below: at CAMERA_DISTANCE + Z0 ~= 5.5 m with
# an 81x81 frame at PIXEL_SCALE, the imaged footprint is only ~3.2 cm wide.
# A pinhole camera's boresight sweeps laterally by
# (CAMERA_DISTANCE + Z0) * tan(angle) at the target plane, so even a
# fraction of a degree of yaw/pitch error translates to centimeters of
# parallax there -- degrees-scale tilt (as in a naive "mildly tilted"
# choice) would sweep the boresight tens of centimeters away from the beam,
# off the edge of the frame entirely. Small angles here (~0.01-0.03 deg)
# still exercise true keystoning/off-axis projection while keeping the beam
# comfortably inside the frame at every z. Similarly, CAMERA_TOLERANCE's
# yaw/pitch bounds are kept tight: yaw/pitch changes are nearly degenerate
# with the beam's own two pointing angles (both produce an z-dependent
# transverse drift), so a loose (e.g. several-degree) bound lets the
# optimizer trade pointing off against yaw/pitch and converge to a very
# wrong split of the two; roll doesn't share that degeneracy and is given
# more room.
TRUE_CAMERA = CameraModel(
    focal_length_px=FOCAL_LENGTH_PX,
    position=(0.002, -0.003, -CAMERA_DISTANCE),
    orientation_deg=(0.03, -0.02, 0.01),
)
# Nominal (calibrated) camera pose, deliberately offset from TRUE_CAMERA
# within CAMERA_TOLERANCE, standing in for real calibration error.
NOMINAL_CAMERA = CameraModel(
    focal_length_px=FOCAL_LENGTH_PX * 1.01,
    position=(0.0015, -0.0025, -CAMERA_DISTANCE),
    orientation_deg=(0.028, -0.018, 0.005),
)
CAMERA_TOLERANCE = CameraModelTolerance(
    focal_length_px=FOCAL_LENGTH_PX * 0.05,
    position=(0.005, 0.005, 0.02),
    orientation_deg=(0.01, 0.01, 0.02),
)

# Measurement plane distances, meters. Kept within roughly +/-2 Rayleigh
# ranges of z0 so the (widening) beam stays well within the camera frame --
# planes much farther out would be clipped by the finite frame, which
# degrades the fit.
Z_LIST = [0.4, 0.45, 0.55, 0.6]
# Each plane's z is only known to a nominal precision (e.g. a translation
# stage's readout); offset the nominal value from the true z used to
# render the plane, and let the fit recover the true z within Z_TOLERANCE.
NOMINAL_Z_OFFSETS = {0.4: 0.003, 0.45: -0.002, 0.55: 0.004, 0.6: -0.003}
Z_TOLERANCE = 0.01

# --- Target thermal-diffusion blur (known target material properties) ---
THERMAL_DIFFUSIVITY = 1e-6  # m^2/s
DWELL_TIME = 0.2            # s


def main() -> None:
    basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
    generator = SyntheticBeamGenerator(basis=basis, camera=TRUE_CAMERA)

    planes = generator.generate(
        coefficients=TRUE_COEFFICIENTS,
        z_list=Z_LIST,
        image_shape=IMAGE_SHAPE,
        center=TRUE_CENTER,
        pointing_angle_horizontal_deg=TRUE_POINTING_HORIZONTAL_DEG,
        pointing_angle_vertical_deg=TRUE_POINTING_VERTICAL_DEG,
        z_tolerance=Z_TOLERANCE,
        nominal_z_offsets=NOMINAL_Z_OFFSETS,
        noise_std=2e-4,
        seed=42,
    )

    # Apply the same thermal-diffusion blur a real target would exhibit,
    # using the nominal (not true) camera to compute the pixel scale --
    # exactly what BeamReconstructor itself does internally.
    blur_deconvolver = DiffusionDeconvolver(
        thermal_diffusivity=THERMAL_DIFFUSIVITY, dwell_time=DWELL_TIME
    )
    nominal_calibration = GeometryCalibration(NOMINAL_CAMERA)
    for plane in planes:
        pixel_scale = nominal_calibration.effective_pixel_scale(plane.flux.shape, plane.z)
        plane.flux = blur_deconvolver.blur(plane.flux, pixel_scale)

    # The ground truth only has order-0 and order-1 content, so a max_order
    # of 1 is enough for automatic mode-set growth to find it; growing much
    # further would start fitting deconvolution/noise artifacts as spurious
    # higher-order modes.
    reconstructor = BeamReconstructor(
        w0=W0,
        z0=Z0,
        wavelength=WAVELENGTH,
        camera=NOMINAL_CAMERA,
        camera_tolerance=CAMERA_TOLERANCE,
        max_order=1,
        deconvolver=blur_deconvolver,
    )
    result = reconstructor.reconstruct(planes)

    print("Mode purity table (power fraction, phase [rad]):")
    for mode, (fraction, phase) in sorted(
        result.purity.items(), key=lambda item: -item[1][0]
    ):
        print(f"  LG_{mode[0]},{mode[1]}: {fraction:6.3%}  (phase {phase:+.3f} rad)")

    print(
        "\nFitted pointing angles: "
        f"horizontal={result.pointing_angle_horizontal_deg:.4f} deg, "
        f"vertical={result.pointing_angle_vertical_deg:.4f} deg"
    )
    print("Fitted beam center per plane (m):")
    for plane, (cx, cy) in zip(planes, result.centers):
        print(f"  z={plane.z:.2f} m -> ({cx:.3e}, {cy:.3e})")

    print("\nFitted camera geometry:")
    for key, value in result.geometry.items():
        print(f"  {key}: {value:.6g}")

    print(f"\nUsed phase-retrieval fallback: {result.used_phase_retrieval}")

    plot_mode_purity(result)
    plot_center_trace(planes, result)
    plot_residuals(planes, result)
    plt.show()


if __name__ == "__main__":
    main()