Compare commits

...

14 Commits

Author SHA1 Message Date
Martino Ferrari 965469fe25 Merge branch 'worktree-camera-geometry-redesign' 2026-07-03 14:27:01 +02:00
Martino Ferrari 83ca084945 Fix final-review findings: pixel-scale size guard, DRY camera field names
effective_pixel_scale now raises a clear ValueError for image shapes
too small for its finite-difference indexing, instead of an IndexError.
ModalFitter.fit()'s geometry dict now reuses CAMERA_FIELD_NAMES from
geometry.py instead of a duplicated hardcoded tuple.

Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
2026-07-03 14:04:24 +02:00
Martino Ferrari 396e605d56 Update CLAUDE.md for the CameraModel geometry redesign
Refreshes the module responsibilities to describe CameraModel/
CameraModelTolerance, z_tolerance, and the two pointing angles, and adds
the new shared-geometry-underdetermined pitfall.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 13:09:21 +02:00
Martino Ferrari 1f2557a0e4 Update full_pipeline_example.py for the CameraModel geometry redesign
Demonstrates the new CameraModel/CameraModelTolerance API, two
independent beam pointing angles, and per-plane z refinement with a
deliberately offset nominal camera pose and z values.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 13:04:19 +02:00
Martino Ferrari c6b824660d Update docs/api.md for the CameraModel geometry redesign
Documents CameraModel/CameraModelTolerance, the rewritten
GeometryCalibration, z_tolerance, the two pointing angles, and every
downstream signature change (ModalFitter, SyntheticBeamGenerator,
PhaseRetriever, BeamReconstructor) introduced by the redesign.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 12:33:00 +02:00
Martino Ferrari 57dc7d743e Wire CameraModel/CameraModelTolerance through BeamReconstructor
BeamReconstructor now requires camera/camera_tolerance, threading them
into ModalFitter.fit_auto and PhaseRetriever.retrieve, and uses
GeometryCalibration.effective_pixel_scale for deconvolution instead of
the removed MeasurementPlane.pixel_scale.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 12:26:18 +02:00
Martino Ferrari 1d399ccc9b Thread CameraModel through PhaseRetriever.retrieve
retrieve() now takes a CameraModel directly instead of the removed
pixel_scale/viewing_angle_deg override kwargs, matching the rest of the
pipeline's shared-camera convention.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 12:19:34 +02:00
Martino Ferrari b81964fbad Update fit_auto for CameraModel and warn on underdetermined geometry fits
fit_auto now threads camera/camera_tolerance through to fit and _bic
(whose parameter count must include any free camera/z unknowns). Emits a
UserWarning, not an error, when free camera+z geometry parameters exceed
the number of measurement planes -- a new documented degeneracy pitfall.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 12:14:07 +02:00
Martino Ferrari ef57ec81e4 Rewrite ModalFitter.fit around the CameraModel tolerance mechanism
The optimizer's parameter vector is now built dynamically: LG
coefficients, per-plane center, and both pointing angles stay always
free; each CameraModel field and each plane's z join the fit (bounded to
its +/- tolerance) only when its paired tolerance is nonzero, and are
otherwise substituted as fixed constants.

Also seeds the first mode's coefficient with a small imaginary offset
(1.0 + 0.05j instead of 1.0 + 0j) rather than exactly on the real axis:
for a single mode, intensity depends only on |c| (phase is unobservable),
so Im=0 sits exactly on that flat/degenerate valley, aligned with a
coordinate axis, giving a zero-gradient Jacobian column that destabilized
trf's trust-region step and prevented pointing-angle recovery from a
default (0, 0) initial guess.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 12:04:47 +02:00
Martino Ferrari c2ae95e01f Rewrite SyntheticBeamGenerator around CameraModel and 2D beam pointing
Renders each plane via true pinhole projection through a shared
CameraModel instead of the old cosine-compression formula, adds
independent horizontal/vertical pointing drift, and supports generating
a deliberately-offset nominal z (vs. true z) per plane for tolerance-
recovery testing in later tasks.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 11:43:20 +02:00
Martino Ferrari dffca62f81 Replace per-plane pixel_scale/viewing_angle with z_tolerance; split pointing angle
MeasurementPlane now carries a z_tolerance (uniform tolerance mechanism)
instead of the removed pixel_scale/viewing_angle_deg fields.
ReconstructionResult.pointing_angle_deg becomes horizontal/vertical fields
to match the beam's two independent tilt degrees of freedom.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 11:39:16 +02:00
Martino Ferrari 4f65c2ce4f Replace cosine-compression geometry model with a full pinhole CameraModel
GeometryCalibration now performs true perspective forward/inverse
projection (with genuine keystoning) around a shared CameraModel, paired
with a CameraModelTolerance that will drive ModalFitter's per-field
fixed/refined behavior in a later task.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 11:32:31 +02:00
Martino Ferrari fabb3d4efc Add implementation plan for camera geometry & measurement uncertainty redesign
Covers CameraModel/CameraModelTolerance, tolerance-unified refinement of
camera pose/intrinsics and per-plane z, 2D beam pointing, and updates to
every affected module, tests, docs, and the example script.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 11:30:27 +02:00
Martino Ferrari 18d6276ab9 Add design spec for full camera geometry & measurement uncertainty
Generalizes camera modeling from a single pixel-scale/viewing-angle pair
to a full pinhole camera (3D orientation + position + intrinsics),
projected via true perspective (not just uniform cosine compression),
plus 2D beam pointing and per-plane z uncertainty. Every nominal
geometry value is now paired with an explicit tolerance that determines
whether it's held fixed or refined jointly with the mode fit.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-07-03 09:00:22 +02:00
20 changed files with 4889 additions and 421 deletions
+48 -21
View File
@@ -35,46 +35,63 @@ This project is not (yet) a git repository.
## Architecture ## Architecture
Data flows through the pipeline as a list of `MeasurementPlane` (one per imaging Data flows through the pipeline as a list of `MeasurementPlane` (one per imaging
distance `z`), each holding a raw 2D `flux` array plus optionally-known distance `z`), each holding a raw 2D `flux` array plus a nominal `z` and `z_tolerance`.
`pixel_scale`/`viewing_angle_deg`. Everything downstream is keyed off `LGBasis`, which Camera pose/intrinsics are a single shared `CameraModel`/`CameraModelTolerance` for the
whole reconstruction, not per-plane. Everything downstream is keyed off `LGBasis`, which
defines the mode basis relative to a known waist `w0`/`z0`/`wavelength`. defines the mode basis relative to a known waist `w0`/`z0`/`wavelength`.
Module responsibilities (`he11lib/`): Module responsibilities (`he11lib/`):
- **`data.py`** — `MeasurementPlane`, `ReconstructionResult` (the shared input/output - **`data.py`** — `MeasurementPlane` (flux, nominal `z`, `z_tolerance`), `ReconstructionResult`
types) and `validate_planes` (>=3 planes, matching shapes, distinct `z`). (the shared input/output types) and `validate_planes` (>=3 planes, matching shapes,
distinct `z`).
- **`modes.py`** — `LGBasis`: closed-form paraxial LG fields, beam radius `w(z)`, - **`modes.py`** — `LGBasis`: closed-form paraxial LG fields, beam radius `w(z)`,
Gouy phase, inverse radius of curvature, and projection of a measured field onto a Gouy phase, inverse radius of curvature, and projection of a measured field onto a
candidate mode set. This is the analytic ground truth all fitting is checked against. candidate mode set. This is the analytic ground truth all fitting is checked against.
- **`geometry.py`** — `GeometryCalibration`: resolves a plane's pixel-to-physical - **`geometry.py`** — `CameraModel`/`CameraModelTolerance` (a nominal pinhole camera
coordinate grid, deferring to known `pixel_scale`/`viewing_angle_deg` on the plane pose/intrinsics and its paired per-field refinement bound) and `GeometryCalibration`:
over any override passed in. resolves pixel<->physical coordinates via true pinhole forward/inverse projection
(ray-plane intersection), producing genuine keystoning for tilted/off-axis poses
rather than a uniform affine correction. A tolerance of `0` on a `CameraModel` field
means it's trusted exactly; `>0` means it's refined within `[nominal-tolerance,
nominal+tolerance]` by `ModalFitter` — the same mechanism applies to
`MeasurementPlane.z`/`z_tolerance`.
- **`noise.py`** — `NoiseEstimator`: automatic per-image noise-std estimation - **`noise.py`** — `NoiseEstimator`: automatic per-image noise-std estimation
(Laplacian method) and per-pixel weights for noise-weighted least squares. (Laplacian method) and per-pixel weights for noise-weighted least squares.
- **`deconvolution.py`** — `DiffusionDeconvolver`: optional forward blur / Wiener - **`deconvolution.py`** — `DiffusionDeconvolver`: optional forward blur / Wiener
deconvolution for thermal-diffusion blur in the absorbing target. The blur kernel is deconvolution for thermal-diffusion blur in the absorbing target. The blur kernel is
isotropic in pixel space, so it's only exact when `viewing_angle_deg == 0` (an isotropic in pixel space; callers use `GeometryCalibration.effective_pixel_scale`
oblique view makes x/y pixel scales differ) — an accepted approximation, not a bug. (a finite-difference approximation at the frame center) as a single figure, so it's
only exact for an on-axis, untilted camera — an accepted approximation, not a bug.
- **`synthetic.py`** — `SyntheticBeamGenerator`: forward model that produces - **`synthetic.py`** — `SyntheticBeamGenerator`: forward model that produces
`MeasurementPlane`s from known ground-truth coefficients/center/pointing/geometry. `MeasurementPlane`s from a known ground-truth `CameraModel`, coefficients, center,
Used throughout the test suite and examples to validate the pipeline end-to-end. and two pointing angles, rendering each plane at its own true `z` (which may
deliberately differ from the plane's nominal `z`, via `nominal_z_offsets`). Used
throughout the test suite and examples to validate the pipeline end-to-end.
- **`fitting.py`** — `ModalFitter` (`fit`, `fit_auto`) and `generate_mode_shells`: the - **`fitting.py`** — `ModalFitter` (`fit`, `fit_auto`) and `generate_mode_shells`: the
core joint nonlinear least-squares fit (complex LG coefficients + beam core joint nonlinear least-squares fit via `scipy.optimize.least_squares`. Complex LG
center/pointing + unknown geometry) via `scipy.optimize.least_squares`. `fit_auto` coefficients, per-plane beam center, and the two pointing angles (horizontal/vertical)
grows the candidate mode set shell-by-shell (by order `2p + |l|`), stopping via a BIC are always free; each `CameraModel` field and each plane's `z` is additionally free
improvement threshold, capped at `max_order` (emits `UserWarning`, doesn't raise, if (bounded by its tolerance) only when its paired tolerance is nonzero, otherwise held
still improving at the cap). fixed as a constant. `fit_auto` grows the candidate mode set shell-by-shell (by order
`2p + |l|`), stopping via a BIC improvement threshold, capped at `max_order` (emits
`UserWarning`, doesn't raise, if still improving at the cap, or if the free
camera+`z` parameter count is large relative to the number of planes — see the
degeneracy pitfall below).
- **`phase_retrieval.py`** — `propagate_angular_spectrum` (FFT-based paraxial - **`phase_retrieval.py`** — `propagate_angular_spectrum` (FFT-based paraxial
free-space propagation) and `PhaseRetriever` (multi-plane Gerchberg-Saxton), the free-space propagation) and `PhaseRetriever` (multi-plane Gerchberg-Saxton), the
fallback reconstruction path for when a finite mode basis doesn't fit well. fallback reconstruction path for when a finite mode basis doesn't fit well. Takes a
- **`reconstruct.py`** — `BeamReconstructor`: the orchestrator. Pipeline order: shared `CameraModel` (not per-plane pixel scale) to derive its common physical grid.
validate planes → optional deconvolution (requires known `pixel_scale` per plane) → - **`reconstruct.py`** — `BeamReconstructor`: the orchestrator, now constructed with a
required `camera`/`camera_tolerance`. Pipeline order: validate planes → optional
deconvolution (using `GeometryCalibration(camera).effective_pixel_scale`) →
`ModalFitter.fit_auto` → optional `PhaseRetriever` fallback (forced via `ModalFitter.fit_auto` → optional `PhaseRetriever` fallback (forced via
`force_phase_retrieval`, or triggered automatically when the noise-weighted RMS `force_phase_retrieval`, or triggered automatically when the noise-weighted RMS
residual exceeds `phase_retrieval_residual_threshold`). The fallback path projects residual exceeds `phase_retrieval_residual_threshold`). The fallback path projects
the recovered field onto all modes up to `max_order` and produces a the recovered field onto all modes up to `max_order` and produces a
`ReconstructionResult` with `used_phase_retrieval=True`, empty `residuals`, and NaN `ReconstructionResult` with `used_phase_retrieval=True`, empty `residuals`, empty
`coefficient_uncertainty` (no fit covariance available from phase retrieval). `geometry`, NaN pointing angles, and NaN `coefficient_uncertainty` (no fit covariance
available from phase retrieval).
- **`plotting.py`** — diagnostic figures (`plot_mode_purity`, `plot_center_trace`, - **`plotting.py`** — diagnostic figures (`plot_mode_purity`, `plot_center_trace`,
`plot_residuals`); each returns a `Figure` rather than calling `plt.show()`. `plot_residuals`); each returns a `Figure` rather than calling `plt.show()`.
@@ -102,3 +119,13 @@ debugging time in this project's history:
2-mode ground truth). When demonstrating growth with deconvolution or noise, set 2-mode ground truth). When demonstrating growth with deconvolution or noise, set
`max_order` close to the true expected mode content rather than generously high, `max_order` close to the true expected mode content rather than generously high,
unless the test specifically targets growth behavior itself. unless the test specifically targets growth behavior itself.
3. **Shared camera/`z` geometry can be underdetermined with few planes.** With only
3-10 planes, adding the ~7-9 shared `CameraModel` unknowns (whichever have nonzero
`CameraModelTolerance`) plus one `z` correction per plane (for nonzero
`z_tolerance`) can be practically underdetermined even though each plane
contributes many pixels of data, because those unknowns are *global* and only
weakly constrained by subtle keystone differences between planes. `fit_auto`/
`BeamReconstructor` emit a `UserWarning` (not an error) when the free-parameter
count is large relative to the number of planes — if you see it, tighten
`CameraModelTolerance`/`z_tolerance` toward values you actually trust rather than
leaving them generously wide.
+132 -48
View File
@@ -15,11 +15,33 @@ Every class/function below is exported from the top-level `he11lib` package
## Quick start ## Quick start
```python ```python
from he11lib import BeamReconstructor, MeasurementPlane from he11lib import (
BeamReconstructor,
CameraModel,
CameraModelTolerance,
MeasurementPlane,
)
# planes: a list of >=3 MeasurementPlane objects built from your own # planes: a list of >=3 MeasurementPlane objects built from your own
# flux arrays (see MeasurementPlane below). # flux arrays (see MeasurementPlane below).
reconstructor = BeamReconstructor(w0=5e-3, z0=0.5, wavelength=1.76e-3)
# Nominal camera pose/intrinsics from calibration; every field here is
# refined jointly with the mode fit because its tolerance is nonzero.
camera = CameraModel(
focal_length_px=2000.0,
position=(0.0, 0.0, -2.0),
orientation_deg=(0.0, 0.0, 0.0),
)
camera_tolerance = CameraModelTolerance(
focal_length_px=20.0,
position=(0.01, 0.01, 0.05),
orientation_deg=(2.0, 2.0, 2.0),
)
reconstructor = BeamReconstructor(
w0=5e-3, z0=0.5, wavelength=1.76e-3,
camera=camera, camera_tolerance=camera_tolerance,
)
result = reconstructor.reconstruct(planes) result = reconstructor.reconstruct(planes)
for mode, (power_fraction, phase_rad) in result.purity.items(): for mode, (power_fraction, phase_rad) in result.purity.items():
@@ -28,19 +50,22 @@ for mode, (power_fraction, phase_rad) in result.purity.items():
## `data` — `MeasurementPlane`, `ReconstructionResult` ## `data` — `MeasurementPlane`, `ReconstructionResult`
### `MeasurementPlane(flux, z, pixel_scale=None, viewing_angle_deg=None, label=None)` ### `MeasurementPlane(flux, z, z_tolerance=0.0, label=None)`
One measurement: a 2D flux array plus its acquisition metadata. One measurement: a 2D flux array plus its acquisition metadata.
- `flux` — 2D `np.ndarray` of flux values. Dead-pixel correction, background - `flux` — 2D `np.ndarray` of flux values. Dead-pixel correction, background
subtraction, and saturation clipping are assumed already handled upstream. subtraction, and saturation clipping are assumed already handled upstream.
- `z` — nominal distance from the output window, in meters. Must be `> 0`. - `z` — nominal distance from the output window, in meters. Must be `> 0`.
- `pixel_scale` — known meters/pixel, or `None` if unknown (then jointly - `z_tolerance``+/-` bound, in meters, around the nominal `z` within
fit by `ModalFitter`/`BeamReconstructor`). which the true distance is jointly refined by `ModalFitter`. Must be
- `viewing_angle_deg` — known camera viewing angle relative to the beam `>= 0`; `0` (the default) means `z` is trusted exactly and held fixed.
axis, in degrees, or `None` if unknown (also jointly fit).
- `label` — optional human-readable identifier. - `label` — optional human-readable identifier.
Per-plane camera geometry (`pixel_scale`/`viewing_angle_deg`) no longer
lives on `MeasurementPlane` — camera pose/intrinsics are a single shared
`CameraModel` for the whole reconstruction (see `geometry` below).
### `validate_planes(planes)` ### `validate_planes(planes)`
Raises `ValueError` if there are fewer than 3 planes, planes have Raises `ValueError` if there are fewer than 3 planes, planes have
@@ -58,9 +83,14 @@ and `BeamReconstructor.reconstruct`):
- `purity: dict[(p, l), (power_fraction, phase_rad)]` - `purity: dict[(p, l), (power_fraction, phase_rad)]`
- `reconstructed_field: np.ndarray` — reconstructed complex field. - `reconstructed_field: np.ndarray` — reconstructed complex field.
- `centers: list[(x, y)]` — fitted beam transverse center per plane, meters. - `centers: list[(x, y)]` — fitted beam transverse center per plane, meters.
- `pointing_angle_deg: float` — fitted shared beam pointing angle (tilt). - `pointing_angle_horizontal_deg`, `pointing_angle_vertical_deg: float`
- `geometry: dict[str, float]` — geometry parameters used or fitted (keys fitted shared beam pointing (tilt) angles, independent horizontal and
`pixel_scale_{i}`, `viewing_angle_deg_{i}` per plane index `i`). vertical.
- `geometry: dict[str, float]` — geometry parameters used or fitted: the 9
`CameraModel` field names from `he11lib.geometry.CAMERA_FIELD_NAMES`
(`focal_length_px`, `position_x`, `position_y`, `position_z`, `yaw_deg`,
`pitch_deg`, `roll_deg`, `principal_point_x`, `principal_point_y`), plus
`z_{i}` per plane index `i` (that plane's fitted/held distance).
- `residuals: list[np.ndarray]` — per-plane (measured modeled) flux maps. - `residuals: list[np.ndarray]` — per-plane (measured modeled) flux maps.
Empty when `used_phase_retrieval` is `True`. Empty when `used_phase_retrieval` is `True`.
- `coefficient_uncertainty: dict[(p, l), float]` — 1-sigma uncertainty on - `coefficient_uncertainty: dict[(p, l), float]` — 1-sigma uncertainty on
@@ -87,17 +117,55 @@ waist radius `w0` (m), waist location `z0` (m), and radiation `wavelength`
onto each `(p, l)` in `modes`, returning `dict[(p, l), complex]` onto each `(p, l)` in `modes`, returning `dict[(p, l), complex]`
coefficients (Riemann-sum inner product; `dx` is the grid spacing). coefficients (Riemann-sum inner product; `dx` is the grid spacing).
## `geometry` — `GeometryCalibration` ## `geometry` — `CameraModel`, `CameraModelTolerance`, `GeometryCalibration`
`GeometryCalibration(plane)` wraps a single `MeasurementPlane` and resolves ### `CameraModel(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))`
its pixel-to-physical-coordinate mapping.
- `pixel_scale_known` / `viewing_angle_known``bool` properties. A nominal pinhole camera pose/intrinsics shared across every plane in one
- `physical_coordinates(pixel_scale=None, viewing_angle_deg=None)` reconstruction. Always a point estimate — never trusted as exact by
returns `(x, y)` physical coordinate grids matching the plane's flux itself; trust is expressed via the paired `CameraModelTolerance`.
shape. Values known on the `MeasurementPlane` take precedence over the
`override` arguments; raises `ValueError` if a value is neither known nor - `focal_length_px` — focal length in pixel units.
overridden. - `position``(x, y, z)` camera position in the beam-axis world frame,
meters; `z=0` is the output window.
- `orientation_deg``(yaw, pitch, roll)`, degrees. All-zero means the
boresight is normal to every `z=const` target plane with no in-plane
rotation.
- `principal_point``(px, px)` offset from the frame center.
### `CameraModelTolerance(focal_length_px, position, orientation_deg, principal_point=(0.0, 0.0))`
Per-field `+/-` refinement bound, same shape as `CameraModel`. Every field
must be `>= 0` (raises `ValueError` otherwise). A field's tolerance of `0`
holds that `CameraModel` field fixed at its nominal value during fitting;
`> 0` lets `ModalFitter` refine it within `[nominal - tolerance, nominal +
tolerance]`.
### `GeometryCalibration(camera)`
Wraps a `CameraModel` and resolves pixel <-> physical coordinate mappings
via true pinhole projection (not a uniform affine/cosine approximation).
- `pixel_coordinates(x, y, z) -> (row, col)` — forward-projects physical
`(x, y)` at depth `z` to pixel coordinates. Raises `ValueError` if the
point is behind the camera (`Z_cam <= 0`).
- `physical_coordinates(image_shape, z) -> (x, y)` — inverse-projects every
pixel in a frame of `image_shape` to physical `(x, y)` on the `z=const`
plane, via ray-plane intersection (this is what produces genuine
keystoning — non-uniform spacing across the frame — for tilted/off-axis
poses). Raises `ValueError` if the plane is edge-on to or behind the
camera.
- `effective_pixel_scale(image_shape, z) -> float` — a single isotropic
meters/pixel figure (finite-difference approximation at the frame
center), for callers like `DiffusionDeconvolver` that assume one
isotropic pixel-space kernel.
### `CAMERA_FIELD_NAMES`, `camera_to_values`, `tolerance_to_values`, `camera_from_values`
Module-level helpers used internally by `ModalFitter` to flatten/unflatten
`CameraModel`/`CameraModelTolerance` into the optimizer's parameter vector.
Not usually needed by application code, but exported for advanced use
(e.g. inspecting `CAMERA_FIELD_NAMES` to interpret `ReconstructionResult.geometry` keys).
## `noise` — `NoiseEstimator` ## `noise` — `NoiseEstimator`
@@ -121,25 +189,30 @@ pass a `deconvolver` to `BeamReconstructor`.
- `deconvolve(image, pixel_scale, noise_to_signal_ratio=1e-3)` — regularized - `deconvolve(image, pixel_scale, noise_to_signal_ratio=1e-3)` — regularized
(Wiener) removal of the blur. (Wiener) removal of the blur.
Note: the blur/deconvolution kernel is isotropic in pixel space. If a Note: the blur/deconvolution kernel is isotropic in pixel space. A tilted
plane has a nonzero `viewing_angle_deg`, its `x` and `y` pixel axes have or off-axis `CameraModel` produces a pixel scale that varies across the
different physical scales (see `SyntheticBeamGenerator` below), so frame and between `x`/`y` (keystoning), so `deconvolve` uses
deconvolution is only exact for `viewing_angle_deg == 0`; at oblique `GeometryCalibration.effective_pixel_scale` — a single isotropic
angles it is an approximation. approximation evaluated at the frame center. This is exact only for an
on-axis, untilted camera; at oblique poses it is an accepted
approximation (see `CLAUDE.md`).
## `synthetic` — `SyntheticBeamGenerator` ## `synthetic` — `SyntheticBeamGenerator`
`SyntheticBeamGenerator(basis, image_shape, pixel_scale)` — forward model `SyntheticBeamGenerator(basis, camera)` — forward model used to validate
used to validate the pipeline against known ground truth, and to evaluate the pipeline against known ground truth, and to evaluate experimental
experimental design (e.g. "would these distances separate my modes?"). design. `camera` is the ground-truth `CameraModel` (position/orientation/
`pixel_scale` is the physical pixel size, in meters, along the non-tilted intrinsics) used to render each plane via true perspective projection.
`y` axis; the `x` axis is compressed by `1/cos(viewing_angle_deg)` to model
an oblique camera view.
- `generate(coefficients, z_list, *, center=(0, 0), pointing_angle_deg=0.0, viewing_angle_deg=0.0, noise_std=0.0, seed=None)` - `generate(coefficients, z_list, image_shape, *, center=(0.0, 0.0), pointing_angle_horizontal_deg=0.0, pointing_angle_vertical_deg=0.0, z_tolerance=0.0, nominal_z_offsets=None, noise_std=0.0, seed=None) -> list[MeasurementPlane]`
— returns one `MeasurementPlane` per `z` in `z_list`. The beam's — returns one `MeasurementPlane` per (true) `z` in `z_list`. The beam's
transverse center drifts linearly with `z` according to transverse center drifts linearly with `z` according to the two
`pointing_angle_deg`, starting from `center` at `z0`. independent pointing angles, starting from `center` at the basis's
`z0`. `nominal_z_offsets`, if given, maps a true `z` to an offset
applied to that plane's *nominal* `z` — letting a reconstruction be
tested against a deliberately-offset nominal input while the plane's
flux is still rendered at the true `z`. Every resulting plane shares
`z_tolerance`.
## `fitting` — `ModalFitter`, `generate_mode_shells` ## `fitting` — `ModalFitter`, `generate_mode_shells`
@@ -152,16 +225,22 @@ Groups candidate `LG_{p,l}` modes into shells of increasing order
### `ModalFitter(basis, noise_estimator=None)` ### `ModalFitter(basis, noise_estimator=None)`
Core reconstruction path: a joint nonlinear least-squares fit of complex LG Core reconstruction path: a joint nonlinear least-squares fit of complex LG
coefficients, beam center/pointing, and (if unknown) geometry. coefficients, beam center/pointing, and any nonzero-tolerance camera/`z`
geometry.
- `fit(planes, modes, initial_coefficients=None, initial_center=(0.0, 0.0), initial_tilt_deg=(0.0, 0.0), initial_pixel_scale=None, initial_viewing_angle_deg=0.0) -> ReconstructionResult` - `fit(planes, modes, camera, camera_tolerance, initial_coefficients=None, initial_center=(0.0, 0.0), initial_pointing_deg=(0.0, 0.0)) -> ReconstructionResult`
— fits exactly the given candidate `modes`. — fits exactly the given candidate `modes`. Every `CameraModel` field
- `fit_auto(planes, max_order=4, bic_improvement_threshold=10.0) -> ReconstructionResult` with a nonzero `camera_tolerance` entry, and every plane whose
`z_tolerance` is nonzero, is refined within `[nominal - tolerance,
nominal + tolerance]`; zero-tolerance fields are held fixed at their
nominal value.
- `fit_auto(planes, camera, camera_tolerance, max_order=4, bic_improvement_threshold=10.0) -> ReconstructionResult`
— starts from `LG_00` and grows the candidate mode set shell-by-shell — starts from `LG_00` and grows the candidate mode set shell-by-shell
(via `generate_mode_shells`), stopping once BIC no longer improves by (via `generate_mode_shells`), stopping once BIC no longer improves by
more than `bic_improvement_threshold`, capped at `max_order`. Emits a more than `bic_improvement_threshold`, capped at `max_order`. Emits a
`UserWarning` (does not raise) if the cap is reached while the fit is `UserWarning` (does not raise) if the cap is reached while the fit is
still improving. still improving, or if the number of free camera+`z` parameters is large
relative to the number of planes (see `CLAUDE.md`'s degeneracy pitfall).
## `phase_retrieval` — `PhaseRetriever`, `propagate_angular_spectrum` ## `phase_retrieval` — `PhaseRetriever`, `propagate_angular_spectrum`
@@ -176,11 +255,12 @@ model implicitly assumed by `LGBasis`'s closed-form paraxial modes.
### `PhaseRetriever(wavelength)` ### `PhaseRetriever(wavelength)`
- `retrieve(planes, pixel_scale=None, viewing_angle_deg=None, max_iterations=200) -> PhaseRetrievalResult` - `retrieve(planes, camera, max_iterations=200) -> PhaseRetrievalResult`
— multi-plane Gerchberg-Saxton phase retrieval: propagates a trial — multi-plane Gerchberg-Saxton phase retrieval: propagates a trial
complex field back and forth between planes, enforcing the measured complex field back and forth between planes, enforcing the measured
amplitude (`sqrt(flux)`) at each plane, without assuming a finite mode amplitude (`sqrt(flux)`) at each plane, without assuming a finite mode
basis. basis. All planes are propagated on one common physical grid, derived
from `camera` at the smallest-`z` plane's depth.
### `PhaseRetrievalResult` ### `PhaseRetrievalResult`
@@ -192,23 +272,27 @@ table, as `BeamReconstructor` does internally for its fallback path.
## `reconstruct` — `BeamReconstructor` ## `reconstruct` — `BeamReconstructor`
`BeamReconstructor(w0, z0, wavelength, max_order=4, noise_estimator=None, deconvolver=None, force_phase_retrieval=False, phase_retrieval_residual_threshold=None)` `BeamReconstructor(w0, z0, wavelength, camera, camera_tolerance, max_order=4, noise_estimator=None, deconvolver=None, force_phase_retrieval=False, phase_retrieval_residual_threshold=None)`
High-level orchestrator wiring together the full pipeline: optional High-level orchestrator wiring together the full pipeline: optional
diffusion deblurring → `ModalFitter.fit_auto` → optional diffusion deblurring → `ModalFitter.fit_auto` → optional
`PhaseRetriever` fallback. `PhaseRetriever` fallback. `camera`/`camera_tolerance` are the nominal
shared `CameraModel` and its per-field refinement bounds for this
reconstruction.
- `reconstruct(planes) -> ReconstructionResult` - `reconstruct(planes) -> ReconstructionResult`
1. Validates `planes` (see `validate_planes`). 1. Validates `planes` (see `validate_planes`).
2. If `deconvolver` is set, deblurs each plane (raises `ValueError` if a 2. If `deconvolver` is set, deblurs each plane using
plane's `pixel_scale` isn't known). `GeometryCalibration(camera).effective_pixel_scale(plane.flux.shape, plane.z)`.
3. Runs `ModalFitter(basis, noise_estimator).fit_auto(planes, max_order)`. 3. Runs `ModalFitter(basis, noise_estimator).fit_auto(planes, camera, camera_tolerance, max_order)`.
4. Runs the `PhaseRetriever` fallback instead, projecting its recovered 4. Runs the `PhaseRetriever` fallback instead, projecting its recovered
field onto all modes up to `max_order`, if `force_phase_retrieval` is field onto all modes up to `max_order`, if `force_phase_retrieval` is
`True`, or if `phase_retrieval_residual_threshold` is set and the `True`, or if `phase_retrieval_residual_threshold` is set and the
modal fit's noise-weighted RMS residual exceeds it. In that case modal fit's noise-weighted RMS residual exceeds it. In that case
`result.residuals` is empty and `coefficient_uncertainty` is `NaN` `result.residuals` is empty, `coefficient_uncertainty` is `NaN` per
per mode (phase retrieval doesn't produce a fit covariance). mode, `geometry` is empty, and both pointing-angle fields are `NaN`
(phase retrieval doesn't fit geometry/pointing or produce a fit
covariance).
## `plotting` — diagnostic visualizations ## `plotting` — diagnostic visualizations
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,222 @@
# he11lib — Full Camera Geometry & Measurement Uncertainty Redesign
**Date:** 2026-07-03
**Status:** Approved for planning
**Supersedes:** parts of `docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md`
relating to `MeasurementPlane.pixel_scale`/`viewing_angle_deg`, `GeometryCalibration`,
and the single-scalar beam pointing angle. All other decisions in the original spec
(mode basis, noise, deconvolution, phase-retrieval fallback, package layout) are
unchanged and still apply.
## Purpose
The original design modeled camera geometry as a single, isotropic pixel scale plus
one viewing-angle tilt (a uniform, affine correction across the whole frame), and
treated each `MeasurementPlane`'s `z` distance as exact. Real measurement setups are
more demanding:
- The camera's orientation relative to the beam axis has a full 3 rotational degrees
of freedom (yaw, pitch, roll), not one.
- The camera is close/wide-angle enough that true perspective projection (keystoning
that varies across the frame) matters, not just a uniform compression factor.
- A physical calibration step (e.g. fiducial markers) gives nominal values for camera
position, orientation, and intrinsics — but mechanical vibration means none of these
can be trusted as exact; all must be refined jointly with everything else.
- The beam's pointing/tilt is two-dimensional (independent horizontal and vertical
tilt), not a single scalar angle.
- Each plane's `z` distance (from a translation stage or tape measure) is also only
known to a nominal precision and should be refined, not trusted exactly.
## Scope of this change
- Replace the single pixel-scale/viewing-angle model with a full pinhole camera model
shared across all planes in one reconstruction.
- Generalize beam pointing from one angle to two (horizontal, vertical).
- Model each plane's `z` as a nominal value with its own refinable uncertainty.
- Unify "trusted/fixed" vs. "uncertain/refined" behind a single tolerance mechanism
(see below), replacing the old `None`-means-unknown convention for geometry.
- Out of scope: lens distortion (radial/tangential), rolling-shutter effects, and
multi-camera setups. These are not part of the current measurement setup and would
need a separate design if they become relevant.
## Architecture
### Data model changes (`data.py`)
- `MeasurementPlane` drops `pixel_scale` and `viewing_angle_deg` (per-plane geometry
no longer makes sense once the camera pose is a single shared physical setup for
the whole reconstruction). It gains `z_tolerance: float`, the ± bound (in meters)
around the nominal `z` within which the true distance is refined. `z_tolerance` must
be `>= 0`; `0` means `z` is trusted exactly and held fixed.
- `ReconstructionResult.pointing_angle_deg: float` becomes two fields:
`pointing_angle_horizontal_deg: float` and `pointing_angle_vertical_deg: float`.
- `ReconstructionResult.geometry` gains entries for the fitted `CameraModel` fields
(see below) alongside the existing per-plane fitted `z` values.
### `CameraModel` and `CameraModelTolerance` (`geometry.py`)
```python
@dataclass
class CameraModel:
focal_length_px: float
position: tuple[float, float, float] # (x, y, z) in the beam-axis frame,
# z=0 at the output window
orientation_deg: tuple[float, float, float] # (yaw, pitch, roll); all-zero =
# boresight normal to the target
# plane, no in-plane rotation
principal_point: tuple[float, float] = (0.0, 0.0) # (px, px) offset from frame center
@dataclass
class CameraModelTolerance:
focal_length_px: float
position: tuple[float, float, float]
orientation_deg: tuple[float, float, float]
principal_point: tuple[float, float] = (0.0, 0.0) # (px, px)
```
`CameraModel` always represents a nominal point estimate (from calibration, or an
assumed default), never a value trusted as exact by itself — trust/uncertainty is
expressed entirely through the paired `CameraModelTolerance` (see "Tolerance
mechanism" below).
`GeometryCalibration` is rewritten around these two types. Given a `CameraModel` and a
target plane's `z`, it provides:
- **Forward projection** — physical `(X, Y)` at that `z` to pixel `(row, col)`, via
standard pinhole projection: rotate/translate the world point into the camera
frame using `orientation_deg`/`position`, then perspective-divide using
`focal_length_px`/`principal_point`.
- **Inverse projection** — pixel `(row, col)` to physical `(X, Y)`, via rayplane
intersection: cast a ray from the camera through each pixel and intersect it with
the known `z =` const target plane. This is what produces genuine keystoning (the
correction varies across the frame), replacing the old uniform
`x = col * scale / cos(angle)` formula.
- Raises `ValueError` for degenerate poses where the target plane is edge-on to or
behind the camera (no valid intersection).
### Tolerance mechanism (applies to `CameraModel`/`CameraModelTolerance` and
`MeasurementPlane.z`/`z_tolerance`)
Every nominal geometry value is paired with a tolerance that determines how
`ModalFitter` treats it:
- **`tolerance == 0`**: held fixed at the nominal value. Excluded entirely from the
optimizer's parameter vector and substituted as a constant into the model function.
This recovers the old "fully known, don't fit it" behavior.
- **`tolerance > 0`**: included in the fit, bounded to
`[nominal - tolerance, nominal + tolerance]` via `scipy.optimize.least_squares`'
`bounds=`.
There is no "fully unbounded" mode for these parameters — if a value is genuinely
unconstrained, its tolerance should be set generously wide rather than infinite,
since an unbounded 7-9 parameter homography fit from 3-10 planes has little chance of
converging usefully without some bound.
### `ModalFitter` changes (`fitting.py`)
The optimizer's parameter vector is now built dynamically per reconstruction:
- Always free (unchanged in kind, generalized in the pointing case): complex LG mode
coefficients, per-plane beam transverse center `(x, y)`, and the two beam pointing
angles (`pointing_angle_horizontal_deg`, `pointing_angle_vertical_deg`).
- Conditionally free (new): any `CameraModel` field whose paired
`CameraModelTolerance` entry is nonzero, and any plane's `z` whose `z_tolerance` is
nonzero — each bounded as described above.
`BeamReconstructor.__init__` gains required `camera: CameraModel` and
`camera_tolerance: CameraModelTolerance` parameters (alongside existing `w0`, `z0`,
`wavelength`), replacing the old optional per-plane pixel-scale/viewing-angle
handling.
### `SyntheticBeamGenerator` changes (`synthetic.py`)
Takes an exact ground-truth `CameraModel` (position/orientation/intrinsics) and the
two beam pointing angles, and generates each plane at its own exact/ground-truth `z`
— which may deliberately differ from the nominal `z` given to the resulting
`MeasurementPlane`, so tests can verify the fit recovers the true `z` despite a
deliberately offset nominal input.
## Data flow (updated)
1. Build a list of `MeasurementPlane` (flux array + nominal `z` + `z_tolerance` +
label), plus a nominal `CameraModel` + `CameraModelTolerance` for the whole
reconstruction.
2. `GeometryCalibration` resolves the pixel↔physical mapping per plane from the
(possibly-still-being-refined) `CameraModel` and that plane's (possibly-being-
refined) `z`.
3. `NoiseEstimator` computes per-plane noise weights (unchanged).
4. `DiffusionDeconvolver` optionally deblurs each plane (unchanged; still assumes an
isotropic pixel-space kernel — noted as an existing approximation, not addressed
by this redesign).
5. `ModalFitter` runs the joint noise-weighted nonlinear least-squares fit over LG
coefficients + per-plane center + 2 pointing angles + any nonzero-tolerance camera
and `z` parameters, growing the mode set automatically as before.
6. Phase-retrieval fallback and `ReconstructionResult` assembly proceed as in the
original design, with `pointing_angle_deg` replaced by the horizontal/vertical
pair and `geometry` extended to include the fitted `CameraModel` fields and
per-plane fitted `z`.
## Testing strategy
In addition to the original design's synthetic-ground-truth-recovery approach:
- **`CameraModel` projection round-trip**: for several poses (on-axis, tilted,
off-center) and several `z` values, physical→pixel→physical recovers the original
point.
- **Keystone regression**: projecting a symmetric grid through a tilted/off-axis
`CameraModel` produces non-uniform spacing across the frame (distinguishing it from
the old uniform cosine-compression model).
- **Degenerate pose**: a pose placing the target plane edge-on to or behind the
camera raises `ValueError`.
- **Tolerance semantics**: a `tolerance=0` field stays exactly at its (deliberately
wrong) nominal value rather than being corrected; a `tolerance>0` field recovers a
ground truth offset from nominal but within its band; a ground truth placed outside
a deliberately too-tight band is clipped to the bound rather than escaping it.
- **End-to-end**: the full pipeline recovers mode purity, both pointing angles, the
camera pose, and per-plane `z`, when `SyntheticBeamGenerator`'s ground truth is
offset from the nominal inputs (within their tolerances) — simulating realistic
calibration/measurement error rather than assuming perfect nominal values.
## Error handling
- `CameraModelTolerance` fields and `MeasurementPlane.z_tolerance` must be `>= 0`;
`ValueError` otherwise. Validated at construction, consistent with the existing
"validate only at boundaries" approach.
- Degenerate camera geometry (target plane edge-on to or behind the camera) raises
`ValueError` from the projection code rather than producing NaNs.
- **New documented pitfall**: with only 3-10 planes, adding ~7-9 shared camera
unknowns plus one `z` correction per plane can be practically underdetermined even
though each plane contributes many pixels of data, because those unknowns are
*global* and only weakly constrained by subtle keystone differences between planes.
`fit_auto`/`BeamReconstructor` emit a `UserWarning` (not an error, consistent with
existing warn-don't-raise style) when the free-parameter count is large relative to
the number of planes, prompting the user toward tighter tolerances rather than
silently returning an ill-conditioned fit.
## Migration impact
This is a breaking change to the public API, acceptable pre-1.0 with no external
users yet:
- `MeasurementPlane(pixel_scale=..., viewing_angle_deg=...)`
`MeasurementPlane(z_tolerance=...)` + a reconstruction-level `CameraModel`/
`CameraModelTolerance`.
- `BeamReconstructor(...)` requires new `camera`/`camera_tolerance` arguments.
- `SyntheticBeamGenerator.generate(...)` requires a `CameraModel` and two pointing
angles instead of scalar `viewing_angle_deg`/`pixel_scale`/`pointing_angle_deg`.
- `ReconstructionResult.pointing_angle_deg``pointing_angle_horizontal_deg` +
`pointing_angle_vertical_deg`.
- `docs/api.md` and `examples/full_pipeline_example.py` need updating to match.
- The existing pitfalls documented in `CLAUDE.md` (Rayleigh-range clipping,
mode-growth overfitting) still apply unchanged; the new degeneracy pitfall above is
additive.
## Deliverables
- Updated `geometry.py` (`CameraModel`, `CameraModelTolerance`, rewritten
`GeometryCalibration`), `data.py`, `fitting.py`, `synthetic.py`, `reconstruct.py`.
- Updated tests covering projection round-trip, keystone behavior, degenerate poses,
tolerance semantics, and end-to-end recovery under offset nominal inputs.
- Updated `docs/api.md` and `examples/full_pipeline_example.py` reflecting the new
public interface.
- Updated `CLAUDE.md` with the new degeneracy pitfall.
+88 -16
View File
@@ -2,10 +2,15 @@
Simulates a gyrotron beam that is mostly the LG_00 fundamental mode with a 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 small admixture of LG_01, viewed by a thermal camera at four distances from
the output window. The camera has an unknown transverse offset/pointing and the output window through a tilted, off-axis pinhole `CameraModel`. The
adds sensor noise; the target also has some thermal-diffusion blur that we camera has an unknown transverse offset/pointing (two independent tilt
correct for. We then reconstruct the mode purity, beam center/pointing, and angles) and adds sensor noise; the target also has some thermal-diffusion
plot the diagnostics. 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: Run with:
@@ -18,7 +23,10 @@ import matplotlib.pyplot as plt
from he11lib import ( from he11lib import (
BeamReconstructor, BeamReconstructor,
CameraModel,
CameraModelTolerance,
DiffusionDeconvolver, DiffusionDeconvolver,
GeometryCalibration,
LGBasis, LGBasis,
SyntheticBeamGenerator, SyntheticBeamGenerator,
plot_center_trace, plot_center_trace,
@@ -33,16 +41,65 @@ WAVELENGTH = 1.76e-3 # radiation wavelength, meters (e.g. a 170 GHz gyrotron)
# --- Ground truth for the synthetic beam (unknown to the reconstructor) --- # --- Ground truth for the synthetic beam (unknown to the reconstructor) ---
TRUE_COEFFICIENTS = {(0, 0): 0.95 + 0j, (0, 1): 0.25 + 0.05j} 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 camera's optical axis TRUE_CENTER = (0.4e-3, -0.3e-3) # beam offset from the optical axis
TRUE_POINTING_DEG = 0.15 # beam pointing (tilt) angle TRUE_POINTING_HORIZONTAL_DEG = 0.15 # beam pointing (horizontal tilt)
CAMERA_VIEWING_ANGLE_DEG = 5.0 # oblique camera viewing angle (known) TRUE_POINTING_VERTICAL_DEG = -0.08 # beam pointing (vertical tilt)
CAMERA_PIXEL_SCALE = 4e-4 # meters/pixel (known calibration)
IMAGE_SHAPE = (81, 81) 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 # Measurement plane distances, meters. Kept within roughly +/-2 Rayleigh
# ranges of z0 so the (widening) beam stays well within the camera frame -- # 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 # planes much farther out would be clipped by the finite frame, which
# degrades the fit. # degrades the fit.
Z_LIST = [0.4, 0.45, 0.55, 0.6] 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) --- # --- Target thermal-diffusion blur (known target material properties) ---
THERMAL_DIFFUSIVITY = 1e-6 # m^2/s THERMAL_DIFFUSIVITY = 1e-6 # m^2/s
@@ -51,26 +108,31 @@ DWELL_TIME = 0.2 # s
def main() -> None: def main() -> None:
basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
generator = SyntheticBeamGenerator( generator = SyntheticBeamGenerator(basis=basis, camera=TRUE_CAMERA)
basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=CAMERA_PIXEL_SCALE
)
planes = generator.generate( planes = generator.generate(
coefficients=TRUE_COEFFICIENTS, coefficients=TRUE_COEFFICIENTS,
z_list=Z_LIST, z_list=Z_LIST,
image_shape=IMAGE_SHAPE,
center=TRUE_CENTER, center=TRUE_CENTER,
pointing_angle_deg=TRUE_POINTING_DEG, pointing_angle_horizontal_deg=TRUE_POINTING_HORIZONTAL_DEG,
viewing_angle_deg=CAMERA_VIEWING_ANGLE_DEG, pointing_angle_vertical_deg=TRUE_POINTING_VERTICAL_DEG,
z_tolerance=Z_TOLERANCE,
nominal_z_offsets=NOMINAL_Z_OFFSETS,
noise_std=2e-4, noise_std=2e-4,
seed=42, seed=42,
) )
# Apply the same thermal-diffusion blur a real target would exhibit. # 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( blur_deconvolver = DiffusionDeconvolver(
thermal_diffusivity=THERMAL_DIFFUSIVITY, dwell_time=DWELL_TIME thermal_diffusivity=THERMAL_DIFFUSIVITY, dwell_time=DWELL_TIME
) )
nominal_calibration = GeometryCalibration(NOMINAL_CAMERA)
for plane in planes: for plane in planes:
plane.flux = blur_deconvolver.blur(plane.flux, plane.pixel_scale) 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 # 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 # of 1 is enough for automatic mode-set growth to find it; growing much
@@ -80,6 +142,8 @@ def main() -> None:
w0=W0, w0=W0,
z0=Z0, z0=Z0,
wavelength=WAVELENGTH, wavelength=WAVELENGTH,
camera=NOMINAL_CAMERA,
camera_tolerance=CAMERA_TOLERANCE,
max_order=1, max_order=1,
deconvolver=blur_deconvolver, deconvolver=blur_deconvolver,
) )
@@ -91,11 +155,19 @@ def main() -> None:
): ):
print(f" LG_{mode[0]},{mode[1]}: {fraction:6.3%} (phase {phase:+.3f} rad)") print(f" LG_{mode[0]},{mode[1]}: {fraction:6.3%} (phase {phase:+.3f} rad)")
print(f"\nFitted pointing angle: {result.pointing_angle_deg:.4f} deg") 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):") print("Fitted beam center per plane (m):")
for plane, (cx, cy) in zip(planes, result.centers): for plane, (cx, cy) in zip(planes, result.centers):
print(f" z={plane.z:.2f} m -> ({cx:.3e}, {cy:.3e})") 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}") print(f"\nUsed phase-retrieval fallback: {result.used_phase_retrieval}")
plot_mode_purity(result) plot_mode_purity(result)
+3 -1
View File
@@ -6,7 +6,7 @@ See docs/ for the full API and design documentation.
from .data import MeasurementPlane, ReconstructionResult from .data import MeasurementPlane, ReconstructionResult
from .deconvolution import DiffusionDeconvolver from .deconvolution import DiffusionDeconvolver
from .fitting import ModalFitter, generate_mode_shells from .fitting import ModalFitter, generate_mode_shells
from .geometry import GeometryCalibration from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration
from .modes import LGBasis from .modes import LGBasis
from .noise import NoiseEstimator from .noise import NoiseEstimator
from .phase_retrieval import PhaseRetrievalResult, PhaseRetriever, propagate_angular_spectrum from .phase_retrieval import PhaseRetrievalResult, PhaseRetriever, propagate_angular_spectrum
@@ -21,6 +21,8 @@ __all__ = [
"DiffusionDeconvolver", "DiffusionDeconvolver",
"ModalFitter", "ModalFitter",
"generate_mode_shells", "generate_mode_shells",
"CameraModel",
"CameraModelTolerance",
"GeometryCalibration", "GeometryCalibration",
"LGBasis", "LGBasis",
"NoiseEstimator", "NoiseEstimator",
+17 -12
View File
@@ -9,24 +9,22 @@ import numpy as np
@dataclass @dataclass
class MeasurementPlane: class MeasurementPlane:
"""A single thermal (flux) image at a known distance from the output window. """A single thermal (flux) image at a nominal distance from the output window.
Parameters Parameters
---------- ----------
flux : 2D array of flux values (already dead-pixel/background/saturation flux : 2D array of flux values (already dead-pixel/background/saturation
corrected upstream). corrected upstream).
z : distance from the output window, in meters. Must be positive. z : nominal distance from the output window, in meters. Must be positive.
pixel_scale : known mm/pixel scale, if calibrated. If None, treated as an z_tolerance : +/- bound, in meters, around the nominal `z` within which
unknown to be estimated jointly during reconstruction. the true distance is refined during fitting. Must be `>= 0`; `0`
viewing_angle_deg : known camera viewing angle relative to the beam axis, means `z` is trusted exactly and held fixed.
in degrees, if calibrated. If None, treated as an unknown.
label : optional human-readable label (e.g. "plane_40cm"). label : optional human-readable label (e.g. "plane_40cm").
""" """
flux: np.ndarray flux: np.ndarray
z: float z: float
pixel_scale: float | None = None z_tolerance: float = 0.0
viewing_angle_deg: float | None = None
label: str | None = None label: str | None = None
def __post_init__(self) -> None: def __post_init__(self) -> None:
@@ -36,6 +34,10 @@ class MeasurementPlane:
) )
if self.z <= 0: if self.z <= 0:
raise ValueError(f"MeasurementPlane.z must be positive, got {self.z}") raise ValueError(f"MeasurementPlane.z must be positive, got {self.z}")
if self.z_tolerance < 0:
raise ValueError(
f"MeasurementPlane.z_tolerance must be >= 0, got {self.z_tolerance}"
)
def validate_planes(planes: list[MeasurementPlane]) -> None: def validate_planes(planes: list[MeasurementPlane]) -> None:
@@ -68,9 +70,11 @@ class ReconstructionResult:
reconstructed_field : reconstructed complex field (at the reference reconstructed_field : reconstructed complex field (at the reference
waist, or as configured). waist, or as configured).
centers : fitted beam transverse center (x, y) in meters, per plane. centers : fitted beam transverse center (x, y) in meters, per plane.
pointing_angle_deg : fitted shared beam pointing angle (tilt), in degrees. pointing_angle_horizontal_deg, pointing_angle_vertical_deg : fitted
geometry : geometry parameters used or fitted (e.g. pixel_scale, shared beam pointing (tilt) angles, in degrees.
viewing_angle_deg), keyed by name. geometry : fitted/held geometry parameters, keyed by name (the 9
`CameraModel` field names from `he11lib.geometry.CAMERA_FIELD_NAMES`,
plus `z_{i}` per plane index `i`).
residuals : per-plane residual maps (measured - modeled flux). residuals : per-plane residual maps (measured - modeled flux).
coefficient_uncertainty : mapping from (p, l) mode index to the coefficient_uncertainty : mapping from (p, l) mode index to the
1-sigma uncertainty on its fitted power fraction. 1-sigma uncertainty on its fitted power fraction.
@@ -81,7 +85,8 @@ class ReconstructionResult:
purity: dict[tuple[int, int], tuple[float, float]] purity: dict[tuple[int, int], tuple[float, float]]
reconstructed_field: np.ndarray reconstructed_field: np.ndarray
centers: list[tuple[float, float]] centers: list[tuple[float, float]]
pointing_angle_deg: float pointing_angle_horizontal_deg: float
pointing_angle_vertical_deg: float
geometry: dict[str, float] = field(default_factory=dict) geometry: dict[str, float] = field(default_factory=dict)
residuals: list[np.ndarray] = field(default_factory=list) residuals: list[np.ndarray] = field(default_factory=list)
coefficient_uncertainty: dict[tuple[int, int], float] = field(default_factory=dict) coefficient_uncertainty: dict[tuple[int, int], float] = field(default_factory=dict)
+147 -64
View File
@@ -8,7 +8,15 @@ import numpy as np
from scipy.optimize import least_squares from scipy.optimize import least_squares
from .data import MeasurementPlane, ReconstructionResult, validate_planes from .data import MeasurementPlane, ReconstructionResult, validate_planes
from .geometry import GeometryCalibration from .geometry import (
CAMERA_FIELD_NAMES,
CameraModel,
CameraModelTolerance,
GeometryCalibration,
camera_from_values,
camera_to_values,
tolerance_to_values,
)
from .modes import LGBasis from .modes import LGBasis
from .noise import NoiseEstimator from .noise import NoiseEstimator
@@ -35,19 +43,31 @@ class ModalFitter:
self, self,
planes: list[MeasurementPlane], planes: list[MeasurementPlane],
modes: list[tuple[int, int]], modes: list[tuple[int, int]],
camera: CameraModel,
camera_tolerance: CameraModelTolerance,
initial_coefficients: dict[tuple[int, int], complex] | None = None, initial_coefficients: dict[tuple[int, int], complex] | None = None,
initial_center: tuple[float, float] = (0.0, 0.0), initial_center: tuple[float, float] = (0.0, 0.0),
initial_tilt_deg: tuple[float, float] = (0.0, 0.0), initial_pointing_deg: tuple[float, float] = (0.0, 0.0),
initial_pixel_scale: float | None = None,
initial_viewing_angle_deg: float = 0.0,
) -> ReconstructionResult: ) -> ReconstructionResult:
"""Jointly fit complex coefficients for `modes` plus center/tilt/geometry.""" """Jointly fit complex coefficients for `modes` plus center/pointing/geometry.
validate_planes(planes)
unknown_scale_idx = [i for i, p in enumerate(planes) if p.pixel_scale is None] Every `CameraModel` field with a nonzero `camera_tolerance` entry,
unknown_angle_idx = [i for i, p in enumerate(planes) if p.viewing_angle_deg is None] and every plane whose `z_tolerance` is nonzero, is refined within
`[nominal - tolerance, nominal + tolerance]`; zero-tolerance fields
are held fixed at their nominal value.
"""
validate_planes(planes)
weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes] weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes]
camera_nominal = camera_to_values(camera)
camera_tol = tolerance_to_values(camera_tolerance)
free_camera_idx = [i for i, t in enumerate(camera_tol) if t > 0]
free_z_idx = [i for i, p in enumerate(planes) if p.z_tolerance > 0]
n_modes = len(modes)
n_always_free = 2 * n_modes + 4 # coefficients + center(2) + pointing(2)
def pack_initial() -> np.ndarray: def pack_initial() -> np.ndarray:
x: list[float] = [] x: list[float] = []
for i, mode in enumerate(modes): for i, mode in enumerate(modes):
@@ -56,98 +76,119 @@ class ModalFitter:
# Nonzero seed for every mode: starting a coefficient at # Nonzero seed for every mode: starting a coefficient at
# exactly 0+0j sits at a flat/degenerate point for the # exactly 0+0j sits at a flat/degenerate point for the
# optimizer and can prevent it from ever leaving zero. # optimizer and can prevent it from ever leaving zero.
c = 1.0 + 0j if i == 0 else 0.1 + 0.05j # A purely-real seed (Im=0) sits exactly on the flat/
# degenerate valley of a single mode's phase (for one
# mode alone, |c|^2 -- not arg(c) -- is all that's
# observable in intensity), giving a zero-gradient
# column that destabilizes trf's trust-region step;
# a small imaginary offset avoids landing on that axis.
c = 1.0 + 0.05j if i == 0 else 0.1 + 0.05j
x += [c.real, c.imag] x += [c.real, c.imag]
x += [initial_center[0], initial_center[1], initial_tilt_deg[0], initial_tilt_deg[1]] x += [initial_center[0], initial_center[1], initial_pointing_deg[0], initial_pointing_deg[1]]
for _ in unknown_scale_idx: for i in free_camera_idx:
x.append(initial_pixel_scale if initial_pixel_scale is not None else 1e-4) x.append(camera_nominal[i])
for _ in unknown_angle_idx: for i in free_z_idx:
x.append(initial_viewing_angle_deg) x.append(planes[i].z)
return np.array(x, dtype=float) return np.array(x, dtype=float)
n_modes = len(modes) def pack_bounds() -> tuple[np.ndarray, np.ndarray]:
lower = [-np.inf] * n_always_free
upper = [np.inf] * n_always_free
for i in free_camera_idx:
lower.append(camera_nominal[i] - camera_tol[i])
upper.append(camera_nominal[i] + camera_tol[i])
for i in free_z_idx:
lower.append(planes[i].z - planes[i].z_tolerance)
upper.append(planes[i].z + planes[i].z_tolerance)
return np.array(lower), np.array(upper)
def unpack(x: np.ndarray): def unpack(x: np.ndarray):
coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)} coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)}
offset = 2 * n_modes offset = 2 * n_modes
x0, y0, tilt_x_deg, tilt_y_deg = x[offset : offset + 4] x0, y0, tilt_h_deg, tilt_v_deg = x[offset : offset + 4]
offset += 4 offset += 4
scales = {}
for idx in unknown_scale_idx:
scales[idx] = x[offset]
offset += 1
angles = {}
for idx in unknown_angle_idx:
angles[idx] = x[offset]
offset += 1
return coeffs, (x0, y0), (tilt_x_deg, tilt_y_deg), scales, angles
def plane_center(x0: float, y0: float, tilt_deg: tuple[float, float], z: float): camera_values = list(camera_nominal)
drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[0])) for i in free_camera_idx:
drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[1])) camera_values[i] = x[offset]
offset += 1
fitted_camera = camera_from_values(camera_values)
z_values = [p.z for p in planes]
for i in free_z_idx:
z_values[i] = x[offset]
offset += 1
return coeffs, (x0, y0), (tilt_h_deg, tilt_v_deg), fitted_camera, z_values
def plane_center(x0: float, y0: float, pointing_deg: tuple[float, float], z: float):
drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[0]))
drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(pointing_deg[1]))
return x0 + drift_x, y0 + drift_y return x0 + drift_x, y0 + drift_y
def model_flux_for_plane(i: int, plane: MeasurementPlane, coeffs, center0, tilt_deg, scales, angles): def model_flux_for_plane(plane, fitted_camera, z, coeffs, center0, pointing_deg):
scale = plane.pixel_scale if plane.pixel_scale is not None else scales[i] calib = GeometryCalibration(fitted_camera)
angle = plane.viewing_angle_deg if plane.viewing_angle_deg is not None else angles[i] x_grid, y_grid = calib.physical_coordinates(plane.flux.shape, z)
calib = GeometryCalibration(plane) cx, cy = plane_center(center0[0], center0[1], pointing_deg, z)
x_grid, y_grid = calib.physical_coordinates(pixel_scale=scale, viewing_angle_deg=angle) field = self.basis.field_superposition(x_grid - cx, y_grid - cy, z, coeffs)
cx, cy = plane_center(center0[0], center0[1], tilt_deg, plane.z)
field = self.basis.field_superposition(x_grid - cx, y_grid - cy, plane.z, coeffs)
return np.abs(field) ** 2 return np.abs(field) ** 2
def residuals(x: np.ndarray) -> np.ndarray: def residuals(x: np.ndarray) -> np.ndarray:
coeffs, center0, tilt_deg, scales, angles = unpack(x) coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(x)
parts = [] parts = []
for i, plane in enumerate(planes): for i, plane in enumerate(planes):
model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) model_flux = model_flux_for_plane(
plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg
)
parts.append(((plane.flux - model_flux) * weights[i]).ravel()) parts.append(((plane.flux - model_flux) * weights[i]).ravel())
return np.concatenate(parts) return np.concatenate(parts)
x0_vec = pack_initial() x0_vec = pack_initial()
# 'trf' + x_scale='jac' handles the very different natural magnitudes lower, upper = pack_bounds()
# of these parameters (coefficients ~O(1), pixel_scale ~O(1e-3), # 'trf' + x_scale='jac' handles the very different natural
# angles ~O(1-90)); plain 'lm' can terminate prematurely on 'xtol' # magnitudes of these parameters (coefficients ~O(1), focal length
# because its unscaled step-size test is dominated by the largest # ~O(1e3-1e4), angles ~O(1-90), z ~O(0.1-1)); plain 'lm' can
# parameters. # terminate prematurely on 'xtol' because its unscaled step-size
# test is dominated by the largest parameters. 'lm' also doesn't
# support bounds, which the tolerance mechanism requires.
opt_result = least_squares( opt_result = least_squares(
residuals, x0_vec, method="trf", x_scale="jac", max_nfev=5000 residuals, x0_vec, method="trf", x_scale="jac", bounds=(lower, upper), max_nfev=5000
) )
coeffs, center0, tilt_deg, scales, angles = unpack(opt_result.x) coeffs, center0, pointing_deg, fitted_camera, z_values = unpack(opt_result.x)
total_power = sum(abs(c) ** 2 for c in coeffs.values()) total_power = sum(abs(c) ** 2 for c in coeffs.values())
if total_power == 0: if total_power == 0:
total_power = 1.0 total_power = 1.0
purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()} purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()}
centers = [plane_center(center0[0], center0[1], tilt_deg, p.z) for p in planes] centers = [
pointing_angle_deg = float(np.hypot(tilt_deg[0], tilt_deg[1])) plane_center(center0[0], center0[1], pointing_deg, z_values[i])
for i in range(len(planes))
]
geometry: dict[str, float] = {} geometry: dict[str, float] = dict(zip(CAMERA_FIELD_NAMES, camera_to_values(fitted_camera)))
for i in range(len(planes)): for i in range(len(planes)):
geometry[f"pixel_scale_{i}"] = ( geometry[f"z_{i}"] = z_values[i]
planes[i].pixel_scale if planes[i].pixel_scale is not None else scales[i]
)
geometry[f"viewing_angle_deg_{i}"] = (
planes[i].viewing_angle_deg if planes[i].viewing_angle_deg is not None else angles[i]
)
residual_maps = [] residual_maps = []
for i, plane in enumerate(planes): for i, plane in enumerate(planes):
model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) model_flux = model_flux_for_plane(
plane, fitted_camera, z_values[i], coeffs, center0, pointing_deg
)
residual_maps.append(plane.flux - model_flux) residual_maps.append(plane.flux - model_flux)
coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power) coefficient_uncertainty = self._estimate_uncertainty(opt_result, modes, coeffs, total_power)
reference_z = min(planes, key=lambda p: abs(p.z - self.basis.z0)).z reference_idx = min(range(len(planes)), key=lambda i: abs(z_values[i] - self.basis.z0))
field_at_reference = self._field_on_default_grid(coeffs, reference_z) field_at_reference = self._field_on_default_grid(coeffs, z_values[reference_idx])
return ReconstructionResult( return ReconstructionResult(
purity=purity, purity=purity,
reconstructed_field=field_at_reference, reconstructed_field=field_at_reference,
centers=centers, centers=centers,
pointing_angle_deg=pointing_angle_deg, pointing_angle_horizontal_deg=pointing_deg[0],
pointing_angle_vertical_deg=pointing_deg[1],
geometry=geometry, geometry=geometry,
residuals=residual_maps, residuals=residual_maps,
coefficient_uncertainty=coefficient_uncertainty, coefficient_uncertainty=coefficient_uncertainty,
@@ -157,23 +198,28 @@ class ModalFitter:
def fit_auto( def fit_auto(
self, self,
planes: list[MeasurementPlane], planes: list[MeasurementPlane],
camera: CameraModel,
camera_tolerance: CameraModelTolerance,
max_order: int = 4, max_order: int = 4,
bic_improvement_threshold: float = 10.0, bic_improvement_threshold: float = 10.0,
) -> ReconstructionResult: ) -> ReconstructionResult:
"""Fit with automatic mode-set growth, capped at `max_order`.""" """Fit with automatic mode-set growth, capped at `max_order`."""
validate_planes(planes) validate_planes(planes)
self._warn_if_degenerate(planes, camera_tolerance)
shells = generate_mode_shells(max_order) shells = generate_mode_shells(max_order)
current_modes = list(shells[0]) current_modes = list(shells[0])
best_result = self.fit(planes, current_modes) best_result = self.fit(planes, current_modes, camera, camera_tolerance)
best_bic = self._bic(planes, best_result, current_modes) best_bic = self._bic(planes, best_result, current_modes, camera_tolerance)
grew_until_cap = True grew_until_cap = True
for shell in shells[1:]: for shell in shells[1:]:
trial_modes = current_modes + shell trial_modes = current_modes + shell
warm_start = self._warm_start_coefficients(best_result, current_modes) warm_start = self._warm_start_coefficients(best_result, current_modes)
trial_result = self.fit(planes, trial_modes, initial_coefficients=warm_start) trial_result = self.fit(
trial_bic = self._bic(planes, trial_result, trial_modes) planes, trial_modes, camera, camera_tolerance, initial_coefficients=warm_start
)
trial_bic = self._bic(planes, trial_result, trial_modes, camera_tolerance)
if trial_bic < best_bic - bic_improvement_threshold: if trial_bic < best_bic - bic_improvement_threshold:
current_modes = trial_modes current_modes = trial_modes
@@ -203,10 +249,47 @@ class ModalFitter:
coeffs[mode] = amplitude * np.exp(1j * phase) coeffs[mode] = amplitude * np.exp(1j * phase)
return coeffs return coeffs
def _bic(self, planes: list[MeasurementPlane], result: ReconstructionResult, modes: list[tuple[int, int]]) -> float: def _warn_if_degenerate(
chi2 = sum(np.sum((r * np.sqrt(self.noise_estimator.weights(p.flux))) ** 2) for r, p in zip(result.residuals, planes)) self, planes: list[MeasurementPlane], camera_tolerance: CameraModelTolerance
) -> None:
"""Warn when free camera+z geometry parameters exceed the plane count.
With only a handful of planes, adding ~7-9 shared camera unknowns
plus one z correction per plane can be practically underdetermined
even though each plane contributes many pixels of data, because
those unknowns are *global* and only weakly constrained by subtle
keystone differences between planes.
"""
free_camera_count = sum(1 for t in tolerance_to_values(camera_tolerance) if t > 0)
free_z_count = sum(1 for p in planes if p.z_tolerance > 0)
free_geometry_count = free_camera_count + free_z_count
if free_geometry_count > len(planes):
warnings.warn(
f"{free_geometry_count} free camera/z geometry parameters "
f"(from nonzero tolerances) but only {len(planes)} measurement "
"planes; the joint fit may be practically underdetermined. "
"Consider tightening CameraModelTolerance / "
"MeasurementPlane.z_tolerance.",
UserWarning,
stacklevel=3,
)
def _bic(
self,
planes: list[MeasurementPlane],
result: ReconstructionResult,
modes: list[tuple[int, int]],
camera_tolerance: CameraModelTolerance,
) -> float:
chi2 = sum(
np.sum((r * np.sqrt(self.noise_estimator.weights(p.flux))) ** 2)
for r, p in zip(result.residuals, planes)
)
n_data = sum(p.flux.size for p in planes) n_data = sum(p.flux.size for p in planes)
n_params = 2 * len(modes) + 4 free_camera_count = sum(1 for t in tolerance_to_values(camera_tolerance) if t > 0)
free_z_count = sum(1 for p in planes if p.z_tolerance > 0)
n_params = 2 * len(modes) + 4 + free_camera_count + free_z_count
return float(chi2 + n_params * np.log(n_data)) return float(chi2 + n_params * np.log(n_data))
def _estimate_uncertainty(self, opt_result, modes, coeffs, total_power): def _estimate_uncertainty(self, opt_result, modes, coeffs, total_power):
+226 -41
View File
@@ -1,64 +1,249 @@
"""Camera geometry correction: pixel-to-physical scale and viewing angle. """Camera geometry: a shared pinhole camera model and pixel<->physical mapping.
Converts a MeasurementPlane's pixel grid into physical (x, y) coordinates in Models the camera as a full pinhole camera (3D position + yaw/pitch/roll
the beam's transverse plane, compensating for an oblique camera viewing orientation + focal length + principal point) shared across all measurement
angle (which compresses the image along the tilt axis by cos(angle)). planes in one reconstruction. Every nominal value on `CameraModel` is paired
Known calibration values (on the MeasurementPlane) are used directly; when a with a `CameraModelTolerance` entry that determines whether `ModalFitter`
value is unknown, an override must be supplied (e.g. by ModalFitter while holds it fixed (tolerance == 0) or refines it within a bound
exploring it as a free parameter). (tolerance > 0) -- `CameraModel` alone is never trusted as exact.
Coordinate conventions
----------------------
World frame: `x` increases along the pixel-column direction, `y` increases
along the pixel-row direction, `z` is distance from the output window along
the beam axis (target planes live at `z = const > 0`).
Camera frame: `X_cam` = right (pixel-column direction), `Y_cam` = down
(pixel-row direction), `Z_cam` = boresight (depth). At
`orientation_deg == (0, 0, 0)`, the camera frame is axis-aligned with the
world frame, so the boresight points along `+z` -- normal to every
`z = const` target plane, with no in-plane rotation.
`orientation_deg = (yaw, pitch, roll)` composes as
`R = R_yaw(about Y) @ R_pitch(about X) @ R_roll(about Z)`, applied to the
camera axes to obtain their world-frame directions.
""" """
from __future__ import annotations from __future__ import annotations
from dataclasses import dataclass, fields
from typing import Sequence
import numpy as np import numpy as np
from .data import MeasurementPlane CAMERA_FIELD_NAMES: tuple[str, ...] = (
"focal_length_px",
"position_x",
"position_y",
"position_z",
"yaw_deg",
"pitch_deg",
"roll_deg",
"principal_point_x",
"principal_point_y",
)
@dataclass
class CameraModel:
"""Nominal pinhole camera parameters, shared across all measurement planes.
Never trusted as exact by itself -- pair with a `CameraModelTolerance`
to express how much each field may be refined during fitting.
Parameters
----------
focal_length_px : focal length, in pixel units.
position : (x, y, z) camera position in the world (beam-axis) frame,
in meters. z=0 is the output window.
orientation_deg : (yaw, pitch, roll), in degrees. All-zero means the
boresight is normal to every z=const target plane, no in-plane
rotation (see module docstring for the full convention).
principal_point : (px, px) offset of the principal point from the frame
center, in pixels.
"""
focal_length_px: float
position: tuple[float, float, float]
orientation_deg: tuple[float, float, float]
principal_point: tuple[float, float] = (0.0, 0.0)
@dataclass
class CameraModelTolerance:
"""+/- bound (same units as `CameraModel`) within which each field is refined.
`0` holds the paired `CameraModel` field fixed at its nominal value;
`> 0` bounds it to `[nominal - tolerance, nominal + tolerance]` during
fitting. All fields must be `>= 0`.
"""
focal_length_px: float
position: tuple[float, float, float]
orientation_deg: tuple[float, float, float]
principal_point: tuple[float, float] = (0.0, 0.0)
def __post_init__(self) -> None:
for f in fields(self):
value = getattr(self, f.name)
components = value if isinstance(value, tuple) else (value,)
for component in components:
if component < 0:
raise ValueError(
f"CameraModelTolerance.{f.name} must be >= 0, got {value}"
)
def camera_to_values(camera: CameraModel) -> list[float]:
"""Flatten a `CameraModel` into the 9 scalars named by `CAMERA_FIELD_NAMES`."""
return [
camera.focal_length_px,
camera.position[0],
camera.position[1],
camera.position[2],
camera.orientation_deg[0],
camera.orientation_deg[1],
camera.orientation_deg[2],
camera.principal_point[0],
camera.principal_point[1],
]
def tolerance_to_values(tolerance: CameraModelTolerance) -> list[float]:
"""Flatten a `CameraModelTolerance` into the 9 scalars named by `CAMERA_FIELD_NAMES`."""
return [
tolerance.focal_length_px,
tolerance.position[0],
tolerance.position[1],
tolerance.position[2],
tolerance.orientation_deg[0],
tolerance.orientation_deg[1],
tolerance.orientation_deg[2],
tolerance.principal_point[0],
tolerance.principal_point[1],
]
def camera_from_values(values: Sequence[float]) -> CameraModel:
"""Inverse of `camera_to_values`: rebuild a `CameraModel` from 9 scalars."""
return CameraModel(
focal_length_px=values[0],
position=(values[1], values[2], values[3]),
orientation_deg=(values[4], values[5], values[6]),
principal_point=(values[7], values[8]),
)
def _rotation_matrix(yaw_deg: float, pitch_deg: float, roll_deg: float) -> np.ndarray:
"""3x3 rotation matrix mapping camera-frame axes to world-frame directions."""
yaw = np.deg2rad(yaw_deg)
pitch = np.deg2rad(pitch_deg)
roll = np.deg2rad(roll_deg)
cy, sy = np.cos(yaw), np.sin(yaw)
cx, sx = np.cos(pitch), np.sin(pitch)
cz, sz = np.cos(roll), np.sin(roll)
r_yaw = np.array([[cy, 0.0, sy], [0.0, 1.0, 0.0], [-sy, 0.0, cy]])
r_pitch = np.array([[1.0, 0.0, 0.0], [0.0, cx, -sx], [0.0, sx, cx]])
r_roll = np.array([[cz, -sz, 0.0], [sz, cz, 0.0], [0.0, 0.0, 1.0]])
return r_yaw @ r_pitch @ r_roll
class GeometryCalibration: class GeometryCalibration:
"""Resolves pixel scale / viewing angle and builds a physical coordinate grid.""" """Resolves the pixel<->physical mapping for a shared pinhole `CameraModel`."""
def __init__(self, plane: MeasurementPlane): def __init__(self, camera: CameraModel):
self.plane = plane self.camera = camera
self._rotation = _rotation_matrix(*camera.orientation_deg)
@property def pixel_coordinates(
def pixel_scale_known(self) -> bool: self, x: np.ndarray, y: np.ndarray, z: float
return self.plane.pixel_scale is not None ) -> tuple[np.ndarray, np.ndarray]:
"""Forward pinhole projection: physical (x, y) at depth z -> centered pixel (row, col)."""
px, py, pz = self.camera.position
dx = x - px
dy = y - py
dz = z - pz
@property r = self._rotation
def viewing_angle_known(self) -> bool: xc = r[0, 0] * dx + r[1, 0] * dy + r[2, 0] * dz
return self.plane.viewing_angle_deg is not None yc = r[0, 1] * dx + r[1, 1] * dy + r[2, 1] * dz
zc = r[0, 2] * dx + r[1, 2] * dy + r[2, 2] * dz
if np.any(zc <= 0):
raise ValueError(
f"One or more target points are behind or edge-on to the "
f"camera at z={z}; cannot project."
)
f = self.camera.focal_length_px
cx, cy = self.camera.principal_point
col = f * xc / zc + cx
row = f * yc / zc + cy
return row, col
def physical_coordinates( def physical_coordinates(
self, self, image_shape: tuple[int, int], z: float
pixel_scale: float | None = None,
viewing_angle_deg: float | None = None,
) -> tuple[np.ndarray, np.ndarray]: ) -> tuple[np.ndarray, np.ndarray]:
"""Physical (x, y) grid matching the plane's flux array shape. """Inverse pinhole projection: pixel grid at depth z -> physical (x, y).
Known values on the MeasurementPlane take precedence; overrides are Casts a ray from the camera through each pixel and intersects it
only used to fill in values that are not known/calibrated. with the world plane z=const. Raises ValueError if the target
plane is edge-on to (parallel to) the view direction or behind the
camera for this pose.
""" """
scale = self.plane.pixel_scale if self.pixel_scale_known else pixel_scale rows, cols = image_shape
angle_deg = (
self.plane.viewing_angle_deg if self.viewing_angle_known else viewing_angle_deg
)
if scale is None:
raise ValueError(
"pixel_scale is not known for this MeasurementPlane and no override was given"
)
if angle_deg is None:
raise ValueError(
"viewing_angle_deg is not known for this MeasurementPlane and no override was given"
)
rows, cols = self.plane.flux.shape
row_idx = np.arange(rows) - rows // 2 row_idx = np.arange(rows) - rows // 2
col_idx = np.arange(cols) - cols // 2 col_idx = np.arange(cols) - cols // 2
col_grid, row_grid = np.meshgrid(col_idx, row_idx) col_grid, row_grid = np.meshgrid(col_idx, row_idx)
cos_angle = np.cos(np.deg2rad(angle_deg)) f = self.camera.focal_length_px
x = col_grid * scale / cos_angle cx, cy = self.camera.principal_point
y = row_grid * scale dir_cam_x = (col_grid - cx) / f
dir_cam_y = (row_grid - cy) / f
dir_cam_z = np.ones_like(dir_cam_x)
r = self._rotation
dir_world_x = r[0, 0] * dir_cam_x + r[0, 1] * dir_cam_y + r[0, 2] * dir_cam_z
dir_world_y = r[1, 0] * dir_cam_x + r[1, 1] * dir_cam_y + r[1, 2] * dir_cam_z
dir_world_z = r[2, 0] * dir_cam_x + r[2, 1] * dir_cam_y + r[2, 2] * dir_cam_z
if np.any(np.abs(dir_world_z) < 1e-12):
raise ValueError(
f"Camera pose is edge-on to the target plane z={z}; no "
"valid ray-plane intersection."
)
px, py, pz = self.camera.position
t = (z - pz) / dir_world_z
if np.any(t <= 0):
raise ValueError(
f"Target plane z={z} is behind the camera for this pose; "
"no valid ray-plane intersection."
)
x = px + t * dir_world_x
y = py + t * dir_world_y
return x, y return x, y
def effective_pixel_scale(self, image_shape: tuple[int, int], z: float) -> float:
"""Isotropic finite-difference approximation of the local pixel scale.
`DiffusionDeconvolver` assumes one isotropic pixel-space blur
kernel; this is only exact for an on-axis, zero-orientation
camera, and an approximation whenever the true projection is
keystoned.
"""
rows, cols = image_shape
if rows < 2 or cols < 2:
raise ValueError(
f"image_shape must be at least 2x2 to compute a finite-difference "
f"pixel scale, got {image_shape}"
)
x, y = self.physical_coordinates(image_shape, z)
mid_row, mid_col = rows // 2, cols // 2
dx = abs(x[mid_row, mid_col + 1] - x[mid_row, mid_col])
dy = abs(y[mid_row + 1, mid_col] - y[mid_row, mid_col])
return float((dx + dy) / 2)
+8 -9
View File
@@ -15,7 +15,7 @@ from dataclasses import dataclass
import numpy as np import numpy as np
from .data import MeasurementPlane, validate_planes from .data import MeasurementPlane, validate_planes
from .geometry import GeometryCalibration from .geometry import CameraModel, GeometryCalibration
def propagate_angular_spectrum( def propagate_angular_spectrum(
@@ -79,22 +79,21 @@ class PhaseRetriever:
def retrieve( def retrieve(
self, self,
planes: list[MeasurementPlane], planes: list[MeasurementPlane],
pixel_scale: float | None = None, camera: CameraModel,
viewing_angle_deg: float | None = None,
max_iterations: int = 200, max_iterations: int = 200,
) -> PhaseRetrievalResult: ) -> PhaseRetrievalResult:
"""Run Gerchberg-Saxton phase retrieval across the given planes. """Run Gerchberg-Saxton phase retrieval across the given planes.
Planes must share the same known (or overridden) pixel_scale and All planes are propagated on one common physical grid, derived
viewing_angle_deg, since all planes are propagated on one common from `camera` at the smallest-z plane's depth (an existing
physical grid. approximation: the shared grid is only exact at that one z, since
other planes may sit at a slightly different true depth under true
perspective projection).
""" """
validate_planes(planes) validate_planes(planes)
ordered = sorted(planes, key=lambda p: p.z) ordered = sorted(planes, key=lambda p: p.z)
x, y = GeometryCalibration(ordered[0]).physical_coordinates( x, y = GeometryCalibration(camera).physical_coordinates(ordered[0].flux.shape, ordered[0].z)
pixel_scale=pixel_scale, viewing_angle_deg=viewing_angle_deg
)
dx = float(x[0, 1] - x[0, 0]) dx = float(x[0, 1] - x[0, 0])
amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered] amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered]
+5 -1
View File
@@ -41,7 +41,11 @@ def plot_center_trace(
ax_y.plot(z_values, y_values, marker="o") ax_y.plot(z_values, y_values, marker="o")
ax_y.set_xlabel("z (m)") ax_y.set_xlabel("z (m)")
ax_y.set_ylabel("center y (m)") ax_y.set_ylabel("center y (m)")
fig.suptitle(f"Beam center (pointing angle {result.pointing_angle_deg:.3g} deg)") fig.suptitle(
"Beam center (pointing angle "
f"h={result.pointing_angle_horizontal_deg:.3g} deg, "
f"v={result.pointing_angle_vertical_deg:.3g} deg)"
)
return fig return fig
+19 -10
View File
@@ -9,6 +9,7 @@ import numpy as np
from .data import MeasurementPlane, ReconstructionResult, validate_planes from .data import MeasurementPlane, ReconstructionResult, validate_planes
from .deconvolution import DiffusionDeconvolver from .deconvolution import DiffusionDeconvolver
from .fitting import ModalFitter, generate_mode_shells from .fitting import ModalFitter, generate_mode_shells
from .geometry import CameraModel, CameraModelTolerance, GeometryCalibration
from .modes import LGBasis from .modes import LGBasis
from .noise import NoiseEstimator from .noise import NoiseEstimator
from .phase_retrieval import PhaseRetriever from .phase_retrieval import PhaseRetriever
@@ -25,12 +26,15 @@ class BeamReconstructor:
Parameters Parameters
---------- ----------
w0, z0, wavelength : known reference beam parameters (see `LGBasis`). w0, z0, wavelength : known reference beam parameters (see `LGBasis`).
camera : nominal shared CameraModel (position/orientation/intrinsics).
camera_tolerance : per-field +/- refinement bound for `camera`; a
zero-tolerance field is held fixed at its nominal value.
max_order : cap on automatic candidate-mode-set growth (see max_order : cap on automatic candidate-mode-set growth (see
`ModalFitter.fit_auto`), and also the mode set used to project the `ModalFitter.fit_auto`), and also the mode set used to project the
phase-retrieval fallback's recovered field onto the LG basis. phase-retrieval fallback's recovered field onto the LG basis.
noise_estimator : shared noise model; defaults to `NoiseEstimator()`. noise_estimator : shared noise model; defaults to `NoiseEstimator()`.
deconvolver : if given, each plane's flux is deblurred (its deconvolver : if given, each plane's flux is deblurred (using
`pixel_scale` must be known) before fitting. `GeometryCalibration(camera).effective_pixel_scale`) before fitting.
force_phase_retrieval : if True, always run the phase-retrieval fallback force_phase_retrieval : if True, always run the phase-retrieval fallback
instead of the modal fit. instead of the modal fit.
phase_retrieval_residual_threshold : if set (and `force_phase_retrieval` phase_retrieval_residual_threshold : if set (and `force_phase_retrieval`
@@ -43,6 +47,8 @@ class BeamReconstructor:
w0: float, w0: float,
z0: float, z0: float,
wavelength: float, wavelength: float,
camera: CameraModel,
camera_tolerance: CameraModelTolerance,
max_order: int = 4, max_order: int = 4,
noise_estimator: NoiseEstimator | None = None, noise_estimator: NoiseEstimator | None = None,
deconvolver: DiffusionDeconvolver | None = None, deconvolver: DiffusionDeconvolver | None = None,
@@ -51,6 +57,8 @@ class BeamReconstructor:
): ):
self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength) self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength)
self.wavelength = wavelength self.wavelength = wavelength
self.camera = camera
self.camera_tolerance = camera_tolerance
self.max_order = max_order self.max_order = max_order
self.noise_estimator = noise_estimator or NoiseEstimator() self.noise_estimator = noise_estimator or NoiseEstimator()
self.deconvolver = deconvolver self.deconvolver = deconvolver
@@ -63,7 +71,9 @@ class BeamReconstructor:
planes = self._deconvolve(planes) planes = self._deconvolve(planes)
fitter = ModalFitter(self.basis, self.noise_estimator) fitter = ModalFitter(self.basis, self.noise_estimator)
result = fitter.fit_auto(planes, max_order=self.max_order) result = fitter.fit_auto(
planes, self.camera, self.camera_tolerance, max_order=self.max_order
)
if self.force_phase_retrieval or self._residual_too_high(result, planes): if self.force_phase_retrieval or self._residual_too_high(result, planes):
result = self._phase_retrieval_fallback(planes) result = self._phase_retrieval_fallback(planes)
@@ -73,13 +83,11 @@ class BeamReconstructor:
def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]:
if self.deconvolver is None: if self.deconvolver is None:
return planes return planes
calib = GeometryCalibration(self.camera)
deblurred = [] deblurred = []
for plane in planes: for plane in planes:
if plane.pixel_scale is None: pixel_scale = calib.effective_pixel_scale(plane.flux.shape, plane.z)
raise ValueError( flux = self.deconvolver.deconvolve(plane.flux, pixel_scale)
"Deconvolution requires a known pixel_scale on every MeasurementPlane."
)
flux = self.deconvolver.deconvolve(plane.flux, plane.pixel_scale)
deblurred.append(replace(plane, flux=flux)) deblurred.append(replace(plane, flux=flux))
return deblurred return deblurred
@@ -101,7 +109,7 @@ class BeamReconstructor:
self, planes: list[MeasurementPlane] self, planes: list[MeasurementPlane]
) -> ReconstructionResult: ) -> ReconstructionResult:
retriever = PhaseRetriever(self.wavelength) retriever = PhaseRetriever(self.wavelength)
pr_result = retriever.retrieve(planes) pr_result = retriever.retrieve(planes, self.camera)
modes = [mode for shell in generate_mode_shells(self.max_order) for mode in shell] modes = [mode for shell in generate_mode_shells(self.max_order) for mode in shell]
dx = float(pr_result.x[0, 1] - pr_result.x[0, 0]) dx = float(pr_result.x[0, 1] - pr_result.x[0, 0])
@@ -116,7 +124,8 @@ class BeamReconstructor:
purity=purity, purity=purity,
reconstructed_field=pr_result.field, reconstructed_field=pr_result.field,
centers=[pr_result.center for _ in planes], centers=[pr_result.center for _ in planes],
pointing_angle_deg=float("nan"), pointing_angle_horizontal_deg=float("nan"),
pointing_angle_vertical_deg=float("nan"),
geometry={}, geometry={},
residuals=[], residuals=[],
coefficient_uncertainty={mode: float("nan") for mode in modes}, coefficient_uncertainty={mode: float("nan") for mode in modes},
+30 -34
View File
@@ -10,6 +10,7 @@ from __future__ import annotations
import numpy as np import numpy as np
from .data import MeasurementPlane from .data import MeasurementPlane
from .geometry import CameraModel, GeometryCalibration
from .modes import LGBasis from .modes import LGBasis
@@ -19,65 +20,60 @@ class SyntheticBeamGenerator:
Parameters Parameters
---------- ----------
basis : LGBasis defining the reference w0, z0, wavelength. basis : LGBasis defining the reference w0, z0, wavelength.
image_shape : (rows, cols) pixel shape of generated images. camera : ground-truth CameraModel (position/orientation/intrinsics) used
pixel_scale : physical size of one pixel, in meters, along the to render each plane via true perspective projection.
non-tilted (y) axis. The tilt/projection axis is assumed to be x.
""" """
def __init__(self, basis: LGBasis, image_shape: tuple[int, int], pixel_scale: float): def __init__(self, basis: LGBasis, camera: CameraModel):
self.basis = basis self.basis = basis
self.image_shape = image_shape self.camera = camera
self.pixel_scale = pixel_scale self.calibration = GeometryCalibration(camera)
def _pixel_grid(self, center: tuple[float, float], viewing_angle_deg: float):
rows, cols = self.image_shape
row_idx = np.arange(rows) - rows // 2
col_idx = np.arange(cols) - cols // 2
col_grid, row_grid = np.meshgrid(col_idx, row_idx)
cos_angle = np.cos(np.deg2rad(viewing_angle_deg))
x = col_grid * self.pixel_scale / cos_angle - center[0]
y = row_grid * self.pixel_scale - center[1]
return x, y
def generate( def generate(
self, self,
coefficients: dict[tuple[int, int], complex], coefficients: dict[tuple[int, int], complex],
z_list: list[float], z_list: list[float],
image_shape: tuple[int, int],
*, *,
center: tuple[float, float] = (0.0, 0.0), center: tuple[float, float] = (0.0, 0.0),
pointing_angle_deg: float = 0.0, pointing_angle_horizontal_deg: float = 0.0,
viewing_angle_deg: float = 0.0, pointing_angle_vertical_deg: float = 0.0,
z_tolerance: float = 0.0,
nominal_z_offsets: dict[float, float] | None = None,
noise_std: float = 0.0, noise_std: float = 0.0,
seed: int | None = None, seed: int | None = None,
) -> list[MeasurementPlane]: ) -> list[MeasurementPlane]:
"""Generate one MeasurementPlane per requested z distance. """Generate one MeasurementPlane per requested (true) z distance.
The beam transverse center drifts linearly with z according to The beam transverse center drifts linearly with z according to the
pointing_angle_deg (tilt of the beam axis along x), starting from two pointing angles, starting from `center` at the basis's
`center` at the basis's reference z0. reference z0. `nominal_z_offsets`, if given, maps a true z (as
given in z_list) to an offset applied to the *nominal* z stored on
the resulting MeasurementPlane -- letting tests verify a fit
recovers the true z despite a deliberately-offset nominal input.
Every resulting plane shares `z_tolerance`.
""" """
rng = np.random.default_rng(seed) rng = np.random.default_rng(seed)
tilt_rad = np.deg2rad(pointing_angle_deg) tilt_h_rad = np.deg2rad(pointing_angle_horizontal_deg)
tilt_v_rad = np.deg2rad(pointing_angle_vertical_deg)
offsets = nominal_z_offsets or {}
planes = [] planes = []
for z in z_list: for z in z_list:
drift_x = (z - self.basis.z0) * np.tan(tilt_rad) drift_x = (z - self.basis.z0) * np.tan(tilt_h_rad)
plane_center = (center[0] + drift_x, center[1]) drift_y = (z - self.basis.z0) * np.tan(tilt_v_rad)
cx = center[0] + drift_x
cy = center[1] + drift_y
x, y = self._pixel_grid(plane_center, viewing_angle_deg) x, y = self.calibration.physical_coordinates(image_shape, z)
field = self.basis.field_superposition(x, y, z, coefficients) field = self.basis.field_superposition(x - cx, y - cy, z, coefficients)
flux = np.abs(field) ** 2 flux = np.abs(field) ** 2
if noise_std > 0: if noise_std > 0:
flux = flux + rng.normal(0.0, noise_std, size=flux.shape) flux = flux + rng.normal(0.0, noise_std, size=flux.shape)
nominal_z = z + offsets.get(z, 0.0)
planes.append( planes.append(
MeasurementPlane( MeasurementPlane(flux=flux, z=nominal_z, z_tolerance=z_tolerance)
flux=flux,
z=z,
pixel_scale=self.pixel_scale,
viewing_angle_deg=viewing_angle_deg,
)
) )
return planes return planes
+13 -9
View File
@@ -10,19 +10,15 @@ def test_measurement_plane_stores_fields():
assert plane.z == 0.3 assert plane.z == 0.3
assert np.array_equal(plane.flux, flux) assert np.array_equal(plane.flux, flux)
assert plane.pixel_scale is None assert plane.z_tolerance == 0.0
assert plane.viewing_angle_deg is None
assert plane.label is None assert plane.label is None
def test_measurement_plane_stores_optional_fields(): def test_measurement_plane_stores_optional_fields():
flux = np.ones((4, 4)) flux = np.ones((4, 4))
plane = MeasurementPlane( plane = MeasurementPlane(flux=flux, z=0.4, z_tolerance=0.01, label="plane_40cm")
flux=flux, z=0.4, pixel_scale=0.1, viewing_angle_deg=5.0, label="plane_40cm"
)
assert plane.pixel_scale == 0.1 assert plane.z_tolerance == 0.01
assert plane.viewing_angle_deg == 5.0
assert plane.label == "plane_40cm" assert plane.label == "plane_40cm"
@@ -39,6 +35,11 @@ def test_measurement_plane_rejects_non_positive_z():
MeasurementPlane(flux=np.ones((4, 4)), z=-0.1) MeasurementPlane(flux=np.ones((4, 4)), z=-0.1)
def test_measurement_plane_rejects_negative_z_tolerance():
with pytest.raises(ValueError, match="z_tolerance"):
MeasurementPlane(flux=np.ones((4, 4)), z=0.3, z_tolerance=-0.01)
def test_validate_planes_rejects_fewer_than_three(): def test_validate_planes_rejects_fewer_than_three():
planes = [ planes = [
MeasurementPlane(flux=np.ones((4, 4)), z=0.3), MeasurementPlane(flux=np.ones((4, 4)), z=0.3),
@@ -82,12 +83,15 @@ def test_reconstruction_result_stores_fields():
purity={(0, 0): (1.0, 0.0)}, purity={(0, 0): (1.0, 0.0)},
reconstructed_field=np.ones((4, 4), dtype=complex), reconstructed_field=np.ones((4, 4), dtype=complex),
centers=[(0.0, 0.0)], centers=[(0.0, 0.0)],
pointing_angle_deg=0.0, pointing_angle_horizontal_deg=0.1,
geometry={"pixel_scale": 0.1, "viewing_angle_deg": 2.0}, pointing_angle_vertical_deg=-0.2,
geometry={"focal_length_px": 2000.0},
residuals=[np.zeros((4, 4))], residuals=[np.zeros((4, 4))],
coefficient_uncertainty={(0, 0): 0.01}, coefficient_uncertainty={(0, 0): 0.01},
used_phase_retrieval=False, used_phase_retrieval=False,
) )
assert result.purity[(0, 0)] == (1.0, 0.0) assert result.purity[(0, 0)] == (1.0, 0.0)
assert result.pointing_angle_horizontal_deg == 0.1
assert result.pointing_angle_vertical_deg == -0.2
assert result.used_phase_retrieval is False assert result.used_phase_retrieval is False
+214 -34
View File
@@ -1,8 +1,11 @@
import warnings
import numpy as np import numpy as np
import pytest import pytest
from he11lib.data import validate_planes from he11lib.data import validate_planes
from he11lib.fitting import ModalFitter, generate_mode_shells from he11lib.fitting import ModalFitter, generate_mode_shells
from he11lib.geometry import CameraModel, CameraModelTolerance
from he11lib.modes import LGBasis from he11lib.modes import LGBasis
from he11lib.synthetic import SyntheticBeamGenerator from he11lib.synthetic import SyntheticBeamGenerator
@@ -10,6 +13,7 @@ W0 = 5e-3
Z0 = 0.5 Z0 = 0.5
WAVELENGTH = 1.76e-3 WAVELENGTH = 1.76e-3
PIXEL_SCALE = 4e-4 PIXEL_SCALE = 4e-4
CAMERA_DISTANCE = 5.0
IMAGE_SHAPE = (61, 61) IMAGE_SHAPE = (61, 61)
Z_LIST = [0.35, 0.5, 0.65, 0.8] Z_LIST = [0.35, 0.5, 0.65, 0.8]
@@ -18,8 +22,21 @@ def make_basis():
return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
def make_generator(basis): def make_camera(pixel_scale=PIXEL_SCALE, position=(0.0, 0.0, -CAMERA_DISTANCE), orientation_deg=(0.0, 0.0, 0.0)):
return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) focal_length_px = (CAMERA_DISTANCE + Z0) / pixel_scale
return CameraModel(
focal_length_px=focal_length_px, position=position, orientation_deg=orientation_deg
)
def zero_tolerance():
return CameraModelTolerance(
focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0)
)
def make_generator(basis, camera):
return SyntheticBeamGenerator(basis=basis, camera=camera)
def test_generate_mode_shells_orders_by_2p_plus_abs_l(): def test_generate_mode_shells_orders_by_2p_plus_abs_l():
@@ -31,13 +48,14 @@ def test_generate_mode_shells_orders_by_2p_plus_abs_l():
def test_fit_recovers_pure_fundamental_mode(): def test_fit_recovers_pure_fundamental_mode():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0 coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0
) )
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result = fitter.fit(planes, modes=[(0, 0)]) result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())
power_fraction, _ = result.purity[(0, 0)] power_fraction, _ = result.purity[(0, 0)]
assert power_fraction == pytest.approx(1.0, abs=1e-6) assert power_fraction == pytest.approx(1.0, abs=1e-6)
@@ -48,12 +66,17 @@ def test_fit_recovers_pure_fundamental_mode():
def test_fit_recovers_two_mode_purity_ratio(): def test_fit_recovers_two_mode_purity_ratio():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j} true_coeffs = {(0, 0): 0.9 + 0j, (1, 0): 0.3 + 0.1j}
planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=1) planes = gen.generate(
coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=1
)
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result = fitter.fit(planes, modes=list(true_coeffs.keys())) result = fitter.fit(
planes, modes=list(true_coeffs.keys()), camera=camera, camera_tolerance=zero_tolerance()
)
true_total = sum(abs(c) ** 2 for c in true_coeffs.values()) true_total = sum(abs(c) ** 2 for c in true_coeffs.values())
for mode, c in true_coeffs.items(): for mode, c in true_coeffs.items():
@@ -64,71 +87,228 @@ def test_fit_recovers_two_mode_purity_ratio():
def test_fit_recovers_center_offset(): def test_fit_recovers_center_offset():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, coefficients={(0, 0): 1.0 + 0j},
z_list=Z_LIST, z_list=Z_LIST,
image_shape=IMAGE_SHAPE,
center=true_center, center=true_center,
noise_std=1e-4, noise_std=1e-4,
seed=2, seed=2,
) )
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result = fitter.fit(planes, modes=[(0, 0)], initial_center=true_center) result = fitter.fit(
planes,
modes=[(0, 0)],
camera=camera,
camera_tolerance=zero_tolerance(),
initial_center=true_center,
)
for cx, cy in result.centers: for cx, cy in result.centers:
assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE) assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE)
assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE) assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE)
def test_fit_recovers_unknown_pixel_scale(): def test_fit_recovers_pointing_angles_independently():
# Use a coarser pixel scale so the (much wider, far-field) beam at the
# outer z distances still fits within the frame -- otherwise pixel scale
# becomes unobservable from clipped images.
basis = make_basis() basis = make_basis()
local_pixel_scale = 1.5e-3 camera = make_camera()
gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale) gen = make_generator(basis, camera)
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3) planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j},
z_list=Z_LIST,
image_shape=IMAGE_SHAPE,
pointing_angle_horizontal_deg=0.3,
pointing_angle_vertical_deg=-0.15,
noise_std=1e-4,
seed=6,
)
# hide the known calibration to force the fitter to solve for it fitter = ModalFitter(basis)
for plane in planes: result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())
plane.pixel_scale = None
plane.viewing_angle_deg = None assert result.pointing_angle_horizontal_deg == pytest.approx(0.3, abs=0.05)
assert result.pointing_angle_vertical_deg == pytest.approx(-0.15, abs=0.05)
def test_fit_holds_zero_tolerance_camera_field_fixed_at_wrong_nominal():
# A tolerance=0 field must stay exactly at its (deliberately wrong)
# nominal value rather than being corrected.
basis = make_basis()
true_camera = make_camera()
gen = make_generator(basis, true_camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=7
)
wrong_focal_length = true_camera.focal_length_px * 1.2
nominal_camera = CameraModel(
focal_length_px=wrong_focal_length,
position=true_camera.position,
orientation_deg=true_camera.orientation_deg,
)
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result = fitter.fit( result = fitter.fit(
planes, planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=zero_tolerance()
modes=[(0, 0)],
initial_pixel_scale=local_pixel_scale * 1.1,
initial_viewing_angle_deg=0.0,
) )
fitted_scales = [result.geometry[f"pixel_scale_{i}"] for i in range(len(planes))] assert result.geometry["focal_length_px"] == wrong_focal_length
for scale in fitted_scales:
assert scale == pytest.approx(local_pixel_scale, rel=0.05)
def test_fit_recovers_offset_camera_field_within_tolerance():
# A tolerance>0 field recovers a ground-truth offset from nominal, but
# within its band.
basis = make_basis()
true_camera = make_camera()
gen = make_generator(basis, true_camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=8
)
offset = true_camera.focal_length_px * 0.02 # 2% off nominal
nominal_camera = CameraModel(
focal_length_px=true_camera.focal_length_px + offset,
position=true_camera.position,
orientation_deg=true_camera.orientation_deg,
)
tolerance = CameraModelTolerance(
focal_length_px=true_camera.focal_length_px * 0.05, # +/-5% band
position=(0.0, 0.0, 0.0),
orientation_deg=(0.0, 0.0, 0.0),
)
fitter = ModalFitter(basis)
result = fitter.fit(planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tolerance)
assert result.geometry["focal_length_px"] == pytest.approx(
true_camera.focal_length_px, rel=0.02
)
def test_fit_clips_out_of_band_ground_truth_to_bound():
# A ground truth placed outside a deliberately too-tight band is
# clipped to the bound rather than escaping it.
basis = make_basis()
true_camera = make_camera()
gen = make_generator(basis, true_camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=9
)
# nominal is 10% off true, but the band only allows +/-1%.
nominal_focal_length = true_camera.focal_length_px * 1.10
nominal_camera = CameraModel(
focal_length_px=nominal_focal_length,
position=true_camera.position,
orientation_deg=true_camera.orientation_deg,
)
tight_tolerance = CameraModelTolerance(
focal_length_px=nominal_focal_length * 0.01,
position=(0.0, 0.0, 0.0),
orientation_deg=(0.0, 0.0, 0.0),
)
fitter = ModalFitter(basis)
result = fitter.fit(
planes, modes=[(0, 0)], camera=nominal_camera, camera_tolerance=tight_tolerance
)
lower_bound = nominal_focal_length - tight_tolerance.focal_length_px
assert result.geometry["focal_length_px"] == pytest.approx(lower_bound, rel=1e-3)
def test_fit_recovers_offset_z_within_tolerance():
basis = make_basis()
camera = make_camera()
gen = make_generator(basis, camera)
true_z_list = [0.35, 0.5, 0.65, 0.8]
offsets = {z: 0.01 for z in true_z_list} # nominal is 1 cm off true
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j},
z_list=true_z_list,
image_shape=IMAGE_SHAPE,
nominal_z_offsets=offsets,
z_tolerance=0.03,
noise_std=1e-4,
seed=10,
)
fitter = ModalFitter(basis)
result = fitter.fit(planes, modes=[(0, 0)], camera=camera, camera_tolerance=zero_tolerance())
for i, true_z in enumerate(true_z_list):
assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005)
def test_fit_auto_does_not_add_modes_for_pure_fundamental(): def test_fit_auto_does_not_add_modes_for_pure_fundamental():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=4 coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=4
) )
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result = fitter.fit_auto(planes, max_order=2) result = fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=2)
assert set(result.purity.keys()) == {(0, 0)} assert set(result.purity.keys()) == {(0, 0)}
def test_fit_auto_grows_to_include_second_mode(): def test_fit_auto_grows_to_include_second_mode():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
true_coeffs = {(0, 0): 0.9 + 0j, (0, 1): 0.4 + 0j} true_coeffs = {(0, 0): 0.9 + 0j, (0, 1): 0.4 + 0j}
planes = gen.generate(coefficients=true_coeffs, z_list=Z_LIST, noise_std=1e-4, seed=5) planes = gen.generate(
coefficients=true_coeffs, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=5
)
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result = fitter.fit_auto(planes, max_order=2) result = fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=2)
assert (0, 1) in result.purity or (0, -1) in result.purity assert (0, 1) in result.purity or (0, -1) in result.purity
def test_fit_auto_warns_when_free_geometry_params_exceed_plane_count():
basis = make_basis()
camera = make_camera()
gen = make_generator(basis, camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j},
z_list=Z_LIST, # 4 planes
image_shape=IMAGE_SHAPE,
z_tolerance=0.05, # +4 free z params
noise_std=1e-4,
seed=11,
)
# +7 free camera params (all but the 2 principal_point components) +
# 4 free z params = 11 free geometry params > 4 planes.
generous_tolerance = CameraModelTolerance(
focal_length_px=camera.focal_length_px * 0.05,
position=(0.01, 0.01, 0.01),
orientation_deg=(2.0, 2.0, 2.0),
principal_point=(0.0, 0.0),
)
fitter = ModalFitter(basis)
with pytest.warns(UserWarning, match="free camera/z geometry parameters"):
fitter.fit_auto(planes, camera=camera, camera_tolerance=generous_tolerance, max_order=1)
def test_fit_auto_does_not_warn_when_geometry_fully_fixed():
basis = make_basis()
camera = make_camera()
gen = make_generator(basis, camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=12
)
fitter = ModalFitter(basis)
with warnings.catch_warnings():
warnings.simplefilter("error", UserWarning)
fitter.fit_auto(planes, camera=camera, camera_tolerance=zero_tolerance(), max_order=1)
+145 -62
View File
@@ -1,77 +1,160 @@
import numpy as np import numpy as np
import pytest import pytest
from he11lib.data import MeasurementPlane from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration
from he11lib.geometry import GeometryCalibration
def test_pixel_scale_known_reflects_plane(): def test_camera_model_tolerance_accepts_zero_and_positive():
plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, pixel_scale=1e-4) CameraModelTolerance(
plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) focal_length_px=0.0,
position=(0.0, 0.0, 0.0),
assert GeometryCalibration(plane_known).pixel_scale_known is True orientation_deg=(1.0, 2.0, 3.0),
assert GeometryCalibration(plane_unknown).pixel_scale_known is False principal_point=(0.5, 0.5),
) # should not raise
def test_viewing_angle_known_reflects_plane(): def test_camera_model_tolerance_rejects_negative_scalar_field():
plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, viewing_angle_deg=10.0) with pytest.raises(ValueError, match="focal_length_px"):
plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) CameraModelTolerance(
focal_length_px=-1.0,
assert GeometryCalibration(plane_known).viewing_angle_known is True position=(0.0, 0.0, 0.0),
assert GeometryCalibration(plane_unknown).viewing_angle_known is False orientation_deg=(0.0, 0.0, 0.0),
def test_physical_coordinates_uses_known_calibration():
plane = MeasurementPlane(
flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0
) )
calib = GeometryCalibration(plane)
x, y = calib.physical_coordinates()
row_idx = np.arange(5) - 2
col_idx = np.arange(5) - 2
expected_x = col_idx * 2e-4
expected_y = row_idx * 2e-4
np.testing.assert_allclose(x[2, :], expected_x)
np.testing.assert_allclose(y[:, 2], expected_y)
def test_physical_coordinates_compresses_x_for_viewing_angle(): def test_camera_model_tolerance_rejects_negative_tuple_component():
plane = MeasurementPlane( with pytest.raises(ValueError, match="position"):
flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=60.0 CameraModelTolerance(
focal_length_px=1.0,
position=(0.0, -0.5, 0.0),
orientation_deg=(0.0, 0.0, 0.0),
) )
calib = GeometryCalibration(plane)
x, y = calib.physical_coordinates()
col_idx = np.arange(5) - 2
expected_x = col_idx * 2e-4 / np.cos(np.deg2rad(60.0))
np.testing.assert_allclose(x[2, :], expected_x)
def test_physical_coordinates_raises_without_calibration_or_override(): def make_on_axis_camera(focal_length_px=2000.0, camera_z=-2.0):
plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) return CameraModel(
calib = GeometryCalibration(plane) focal_length_px=focal_length_px,
position=(0.0, 0.0, camera_z),
with pytest.raises(ValueError, match="pixel_scale"): orientation_deg=(0.0, 0.0, 0.0),
calib.physical_coordinates()
def test_physical_coordinates_accepts_override_for_unknown_values():
plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3)
calib = GeometryCalibration(plane)
x, y = calib.physical_coordinates(pixel_scale=1e-4, viewing_angle_deg=0.0)
col_idx = np.arange(5) - 2
np.testing.assert_allclose(x[2, :], col_idx * 1e-4)
def test_known_calibration_takes_precedence_over_override():
plane = MeasurementPlane(
flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=0.0
) )
calib = GeometryCalibration(plane)
# override should be ignored since plane already specifies calibration
x, _ = calib.physical_coordinates(pixel_scale=999.0, viewing_angle_deg=45.0) def make_tilted_camera():
col_idx = np.arange(5) - 2 return CameraModel(
np.testing.assert_allclose(x[2, :], col_idx * 2e-4) 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)
+28 -8
View File
@@ -1,6 +1,7 @@
import numpy as np import numpy as np
import pytest import pytest
from he11lib.geometry import CameraModel
from he11lib.modes import LGBasis from he11lib.modes import LGBasis
from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum
from he11lib.synthetic import SyntheticBeamGenerator from he11lib.synthetic import SyntheticBeamGenerator
@@ -9,6 +10,7 @@ W0 = 5e-3
Z0 = 0.5 Z0 = 0.5
WAVELENGTH = 1.76e-3 WAVELENGTH = 1.76e-3
PIXEL_SCALE = 3e-4 PIXEL_SCALE = 3e-4
CAMERA_DISTANCE = 5.0
IMAGE_SHAPE = (121, 121) IMAGE_SHAPE = (121, 121)
@@ -16,6 +18,15 @@ def make_basis():
return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
def make_camera():
focal_length_px = (CAMERA_DISTANCE + Z0) / PIXEL_SCALE
return CameraModel(
focal_length_px=focal_length_px,
position=(0.0, 0.0, -CAMERA_DISTANCE),
orientation_deg=(0.0, 0.0, 0.0),
)
def make_grid(): def make_grid():
coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE
x, y = np.meshgrid(coords, coords) x, y = np.meshgrid(coords, coords)
@@ -42,7 +53,6 @@ def test_propagate_matches_lgbasis_analytic_evolution():
propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH) propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH)
analytic = basis.field(x, y, Z0 + dz, p=0, l=0) analytic = basis.field(x, y, Z0 + dz, p=0, l=0)
# compare intensity profiles (phase reference/global constant may differ)
np.testing.assert_allclose( np.testing.assert_allclose(
np.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2) np.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2)
) )
@@ -53,15 +63,19 @@ def test_retrieve_recovers_pure_mode_purity():
# within the frame -- otherwise FFT wraparound/clipping at the edges # within the frame -- otherwise FFT wraparound/clipping at the edges
# degrades angular-spectrum propagation accuracy. # degrades angular-spectrum propagation accuracy.
basis = make_basis() basis = make_basis()
gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) camera = make_camera()
gen = SyntheticBeamGenerator(basis=basis, camera=camera)
z_list = [0.47, 0.5, 0.53] z_list = [0.47, 0.5, 0.53]
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=0) planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=0
)
retriever = PhaseRetriever(wavelength=WAVELENGTH) retriever = PhaseRetriever(wavelength=WAVELENGTH)
result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100) result = retriever.retrieve(planes, camera, max_iterations=100)
dx = float(result.x[0, 1] - result.x[0, 0])
coeffs = basis.project( coeffs = basis.project(
result.field, result.x, result.y, PIXEL_SCALE, result.z, modes=[(0, 0), (1, 0), (0, 1)] result.field, result.x, result.y, dx, result.z, modes=[(0, 0), (1, 0), (0, 1)]
) )
total_power = sum(abs(c) ** 2 for c in coeffs.values()) total_power = sum(abs(c) ** 2 for c in coeffs.values())
purity_00 = abs(coeffs[(0, 0)]) ** 2 / total_power purity_00 = abs(coeffs[(0, 0)]) ** 2 / total_power
@@ -70,15 +84,21 @@ def test_retrieve_recovers_pure_mode_purity():
def test_retrieve_estimates_beam_center(): def test_retrieve_estimates_beam_center():
basis = make_basis() basis = make_basis()
gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) camera = make_camera()
gen = SyntheticBeamGenerator(basis=basis, camera=camera)
z_list = [0.47, 0.5, 0.53] z_list = [0.47, 0.5, 0.53]
true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE) true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE)
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, center=true_center, noise_std=1e-5, seed=1 coefficients={(0, 0): 1.0 + 0j},
z_list=z_list,
image_shape=IMAGE_SHAPE,
center=true_center,
noise_std=1e-5,
seed=1,
) )
retriever = PhaseRetriever(wavelength=WAVELENGTH) retriever = PhaseRetriever(wavelength=WAVELENGTH)
result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100) result = retriever.retrieve(planes, camera, max_iterations=100)
assert result.center[0] == pytest.approx(true_center[0], abs=3 * PIXEL_SCALE) assert result.center[0] == pytest.approx(true_center[0], abs=3 * PIXEL_SCALE)
assert result.center[1] == pytest.approx(true_center[1], abs=3 * PIXEL_SCALE) assert result.center[1] == pytest.approx(true_center[1], abs=3 * PIXEL_SCALE)
+2 -1
View File
@@ -11,7 +11,8 @@ def make_result(**overrides):
purity={(0, 0): (0.9, 0.1), (1, 0): (0.1, -0.2)}, purity={(0, 0): (0.9, 0.1), (1, 0): (0.1, -0.2)},
reconstructed_field=np.zeros((5, 5), dtype=complex), reconstructed_field=np.zeros((5, 5), dtype=complex),
centers=[(0.0, 0.0), (1e-4, -1e-4), (2e-4, -2e-4)], centers=[(0.0, 0.0), (1e-4, -1e-4), (2e-4, -2e-4)],
pointing_angle_deg=0.5, pointing_angle_horizontal_deg=0.3,
pointing_angle_vertical_deg=0.4,
residuals=[np.ones((5, 5)), np.ones((5, 5)) * 2, np.ones((5, 5)) * 3], residuals=[np.ones((5, 5)), np.ones((5, 5)) * 2, np.ones((5, 5)) * 3],
) )
defaults.update(overrides) defaults.update(overrides)
+126 -18
View File
@@ -4,6 +4,7 @@ import pytest
from he11lib.deconvolution import DiffusionDeconvolver from he11lib.deconvolution import DiffusionDeconvolver
from he11lib.fitting import ModalFitter from he11lib.fitting import ModalFitter
from he11lib.geometry import CameraModel, CameraModelTolerance, GeometryCalibration
from he11lib.modes import LGBasis from he11lib.modes import LGBasis
from he11lib.reconstruct import BeamReconstructor from he11lib.reconstruct import BeamReconstructor
from he11lib.synthetic import SyntheticBeamGenerator from he11lib.synthetic import SyntheticBeamGenerator
@@ -12,6 +13,7 @@ W0 = 5e-3
Z0 = 0.5 Z0 = 0.5
WAVELENGTH = 1.76e-3 WAVELENGTH = 1.76e-3
PIXEL_SCALE = 4e-4 PIXEL_SCALE = 4e-4
CAMERA_DISTANCE = 5.0
IMAGE_SHAPE = (61, 61) IMAGE_SHAPE = (61, 61)
Z_LIST = [0.35, 0.5, 0.65, 0.8] Z_LIST = [0.35, 0.5, 0.65, 0.8]
@@ -20,16 +22,36 @@ def make_basis():
return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
def make_generator(basis): def make_camera():
return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) focal_length_px = (CAMERA_DISTANCE + Z0) / PIXEL_SCALE
return CameraModel(
focal_length_px=focal_length_px,
position=(0.0, 0.0, -CAMERA_DISTANCE),
orientation_deg=(0.0, 0.0, 0.0),
)
def zero_tolerance():
return CameraModelTolerance(
focal_length_px=0.0, position=(0.0, 0.0, 0.0), orientation_deg=(0.0, 0.0, 0.0)
)
def make_generator(basis, camera):
return SyntheticBeamGenerator(basis=basis, camera=camera)
def test_reconstruct_recovers_pure_mode_purity(): def test_reconstruct_recovers_pure_mode_purity():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0) gen = make_generator(basis, camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=0
)
reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2) reconstructor = BeamReconstructor(
w0=W0, z0=Z0, wavelength=WAVELENGTH, camera=camera, camera_tolerance=zero_tolerance(), max_order=2
)
result = reconstructor.reconstruct(planes) result = reconstructor.reconstruct(planes)
power_fraction, _ = result.purity[(0, 0)] power_fraction, _ = result.purity[(0, 0)]
@@ -39,13 +61,21 @@ def test_reconstruct_recovers_pure_mode_purity():
def test_reconstruct_recovers_center_offset(): def test_reconstruct_recovers_center_offset():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
gen = make_generator(basis, camera)
true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE)
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, center=true_center, noise_std=1e-4, seed=1 coefficients={(0, 0): 1.0 + 0j},
z_list=Z_LIST,
image_shape=IMAGE_SHAPE,
center=true_center,
noise_std=1e-4,
seed=1,
) )
reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2) reconstructor = BeamReconstructor(
w0=W0, z0=Z0, wavelength=WAVELENGTH, camera=camera, camera_tolerance=zero_tolerance(), max_order=2
)
result = reconstructor.reconstruct(planes) result = reconstructor.reconstruct(planes)
for cx, cy in result.centers: for cx, cy in result.centers:
@@ -55,21 +85,34 @@ def test_reconstruct_recovers_center_offset():
def test_reconstruct_with_deconvolution_corrects_blur(): def test_reconstruct_with_deconvolution_corrects_blur():
basis = make_basis() basis = make_basis()
gen = make_generator(basis) camera = make_camera()
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=2) gen = make_generator(basis, camera)
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, image_shape=IMAGE_SHAPE, noise_std=1e-4, seed=2
)
deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0) deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0)
calib = GeometryCalibration(camera)
blurred_planes = [ blurred_planes = [
replace(p, flux=deconvolver.blur(p.flux, p.pixel_scale)) for p in planes replace(p, flux=deconvolver.blur(p.flux, calib.effective_pixel_scale(p.flux.shape, p.z)))
for p in planes
] ]
# Without deconvolution, blur should measurably hurt purity recovery. # Without deconvolution, blur should measurably hurt purity recovery.
fitter = ModalFitter(basis) fitter = ModalFitter(basis)
result_no_deconv = fitter.fit(blurred_planes, modes=[(0, 0), (1, 0), (0, 1)]) result_no_deconv = fitter.fit(
blurred_planes, modes=[(0, 0), (1, 0), (0, 1)], camera=camera, camera_tolerance=zero_tolerance()
)
purity_no_deconv, _ = result_no_deconv.purity[(0, 0)] purity_no_deconv, _ = result_no_deconv.purity[(0, 0)]
reconstructor = BeamReconstructor( reconstructor = BeamReconstructor(
w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, deconvolver=deconvolver w0=W0,
z0=Z0,
wavelength=WAVELENGTH,
camera=camera,
camera_tolerance=zero_tolerance(),
max_order=2,
deconvolver=deconvolver,
) )
result = reconstructor.reconstruct(blurred_planes) result = reconstructor.reconstruct(blurred_planes)
purity_with_deconv, _ = result.purity[(0, 0)] purity_with_deconv, _ = result.purity[(0, 0)]
@@ -80,12 +123,21 @@ def test_reconstruct_with_deconvolution_corrects_blur():
def test_reconstruct_forces_phase_retrieval_fallback(): def test_reconstruct_forces_phase_retrieval_fallback():
basis = make_basis() basis = make_basis()
gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) camera = make_camera()
gen = SyntheticBeamGenerator(basis=basis, camera=camera)
z_list = [0.47, 0.5, 0.53] z_list = [0.47, 0.5, 0.53]
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=3) planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=3
)
reconstructor = BeamReconstructor( reconstructor = BeamReconstructor(
w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, force_phase_retrieval=True w0=W0,
z0=Z0,
wavelength=WAVELENGTH,
camera=camera,
camera_tolerance=zero_tolerance(),
max_order=2,
force_phase_retrieval=True,
) )
result = reconstructor.reconstruct(planes) result = reconstructor.reconstruct(planes)
@@ -96,17 +148,73 @@ def test_reconstruct_forces_phase_retrieval_fallback():
def test_reconstruct_falls_back_automatically_on_high_residual(): def test_reconstruct_falls_back_automatically_on_high_residual():
basis = make_basis() basis = make_basis()
gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) camera = make_camera()
gen = SyntheticBeamGenerator(basis=basis, camera=camera)
z_list = [0.47, 0.5, 0.53] z_list = [0.47, 0.5, 0.53]
planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, noise_std=1e-5, seed=4) planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE, noise_std=1e-5, seed=4
)
reconstructor = BeamReconstructor( reconstructor = BeamReconstructor(
w0=W0, w0=W0,
z0=Z0, z0=Z0,
wavelength=WAVELENGTH, wavelength=WAVELENGTH,
camera=camera,
camera_tolerance=zero_tolerance(),
max_order=2, max_order=2,
phase_retrieval_residual_threshold=1e-8, phase_retrieval_residual_threshold=1e-8,
) )
result = reconstructor.reconstruct(planes) result = reconstructor.reconstruct(planes)
assert result.used_phase_retrieval is True assert result.used_phase_retrieval is True
def test_reconstruct_recovers_camera_and_z_offset_from_nominal():
# End-to-end: ground truth is offset from the nominal camera/z inputs
# (within their tolerances), simulating realistic calibration error.
basis = make_basis()
true_camera = make_camera()
gen = make_generator(basis, true_camera)
true_z_list = Z_LIST
z_offsets = {z: 0.01 for z in true_z_list}
planes = gen.generate(
coefficients={(0, 0): 1.0 + 0j},
z_list=true_z_list,
image_shape=IMAGE_SHAPE,
nominal_z_offsets=z_offsets,
z_tolerance=0.03,
pointing_angle_horizontal_deg=0.2,
pointing_angle_vertical_deg=-0.1,
noise_std=1e-4,
seed=13,
)
nominal_focal_offset = true_camera.focal_length_px * 0.03
nominal_camera = CameraModel(
focal_length_px=true_camera.focal_length_px + nominal_focal_offset,
position=true_camera.position,
orientation_deg=true_camera.orientation_deg,
)
tolerance = CameraModelTolerance(
focal_length_px=true_camera.focal_length_px * 0.1,
position=(0.0, 0.0, 0.0),
orientation_deg=(0.0, 0.0, 0.0),
)
reconstructor = BeamReconstructor(
w0=W0,
z0=Z0,
wavelength=WAVELENGTH,
camera=nominal_camera,
camera_tolerance=tolerance,
max_order=1,
)
result = reconstructor.reconstruct(planes)
power_fraction, _ = result.purity[(0, 0)]
assert power_fraction > 0.95
assert result.pointing_angle_horizontal_deg == pytest.approx(0.2, abs=0.1)
assert result.pointing_angle_vertical_deg == pytest.approx(-0.1, abs=0.1)
assert result.geometry["focal_length_px"] == pytest.approx(true_camera.focal_length_px, rel=0.03)
for i, true_z in enumerate(true_z_list):
assert result.geometry[f"z_{i}"] == pytest.approx(true_z, abs=0.005)
+63 -40
View File
@@ -1,6 +1,7 @@
import numpy as np import numpy as np
import pytest import pytest
from he11lib.geometry import CameraModel
from he11lib.modes import LGBasis from he11lib.modes import LGBasis
from he11lib.synthetic import SyntheticBeamGenerator from he11lib.synthetic import SyntheticBeamGenerator
@@ -8,21 +9,29 @@ from he11lib.synthetic import SyntheticBeamGenerator
W0 = 5e-3 W0 = 5e-3
Z0 = 0.5 Z0 = 0.5
WAVELENGTH = 1.76e-3 WAVELENGTH = 1.76e-3
PIXEL_SCALE = 2e-4 # 0.2 mm/px PIXEL_SCALE = 2e-4 # 0.2 mm/px, achieved at z=Z0
CAMERA_DISTANCE = 5.0 # camera stands 5 m upstream of the output window
IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel
def make_camera(pixel_scale=PIXEL_SCALE, z0=Z0, camera_distance=CAMERA_DISTANCE):
focal_length_px = (camera_distance + z0) / pixel_scale
return CameraModel(
focal_length_px=focal_length_px,
position=(0.0, 0.0, -camera_distance),
orientation_deg=(0.0, 0.0, 0.0),
)
def make_generator(): def make_generator():
basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH)
return SyntheticBeamGenerator( return SyntheticBeamGenerator(basis=basis, camera=make_camera())
basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE
)
def test_generate_returns_planes_with_requested_z(): def test_generate_returns_planes_with_requested_z():
gen = make_generator() gen = make_generator()
z_list = [0.3, 0.4, 0.5] z_list = [0.3, 0.4, 0.5]
planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list) planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list, image_shape=IMAGE_SHAPE)
assert [p.z for p in planes] == z_list assert [p.z for p in planes] == z_list
assert all(p.flux.shape == IMAGE_SHAPE for p in planes) assert all(p.flux.shape == IMAGE_SHAPE for p in planes)
@@ -30,7 +39,9 @@ def test_generate_returns_planes_with_requested_z():
def test_generate_pure_mode_peak_at_image_center_when_centered(): def test_generate_pure_mode_peak_at_image_center_when_centered():
gen = make_generator() gen = make_generator()
planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(0.0, 0.0)) planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(0.0, 0.0)
)
flux = planes[0].flux flux = planes[0].flux
peak_idx = np.unravel_index(np.argmax(flux), flux.shape) peak_idx = np.unravel_index(np.argmax(flux), flux.shape)
@@ -40,9 +51,9 @@ def test_generate_pure_mode_peak_at_image_center_when_centered():
def test_generate_applies_center_offset(): def test_generate_applies_center_offset():
gen = make_generator() gen = make_generator()
offset_m = 20 * PIXEL_SCALE # 20 pixels offset_m = 20 * PIXEL_SCALE # ~20 pixels at z=Z0
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(offset_m, 0.0) coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, center=(offset_m, 0.0)
) )
flux = planes[0].flux flux = planes[0].flux
@@ -53,35 +64,41 @@ def test_generate_applies_center_offset():
assert peak_idx[1] == pytest.approx(center_col + 20, abs=1) assert peak_idx[1] == pytest.approx(center_col + 20, abs=1)
def test_generate_applies_pointing_angle_as_linear_drift(): def test_generate_applies_pointing_angles_as_2d_linear_drift():
gen = make_generator() gen = make_generator()
pointing_angle_deg = 1.0 # small tilt
z_list = [Z0, Z0 + 0.2] z_list = [Z0, Z0 + 0.2]
planes = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, coefficients={(0, 0): 1 + 0j},
z_list=z_list, z_list=z_list,
image_shape=IMAGE_SHAPE,
center=(0.0, 0.0), center=(0.0, 0.0),
pointing_angle_deg=pointing_angle_deg, pointing_angle_horizontal_deg=1.0,
pointing_angle_vertical_deg=0.5,
) )
peaks_col = [] peaks = []
for plane in planes: for plane in planes:
peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape) peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape)
peaks_col.append(peak_idx[1]) peaks.append(peak_idx)
expected_shift_m = 0.2 * np.tan(np.deg2rad(pointing_angle_deg)) expected_shift_x_m = 0.2 * np.tan(np.deg2rad(1.0))
expected_shift_px = expected_shift_m / PIXEL_SCALE expected_shift_y_m = 0.2 * np.tan(np.deg2rad(0.5))
actual_shift_px = peaks_col[1] - peaks_col[0] expected_shift_col_px = expected_shift_x_m / PIXEL_SCALE
assert actual_shift_px == pytest.approx(expected_shift_px, abs=1) expected_shift_row_px = expected_shift_y_m / PIXEL_SCALE
actual_shift_col_px = peaks[1][1] - peaks[0][1]
actual_shift_row_px = peaks[1][0] - peaks[0][0]
assert actual_shift_col_px == pytest.approx(expected_shift_col_px, abs=1)
assert actual_shift_row_px == pytest.approx(expected_shift_row_px, abs=1)
def test_generate_noise_is_reproducible_with_seed(): def test_generate_noise_is_reproducible_with_seed():
gen = make_generator() gen = make_generator()
planes_a = gen.generate( planes_a = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42 coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42
) )
planes_b = gen.generate( planes_b = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42 coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.01, seed=42
) )
np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux) np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux)
@@ -90,34 +107,40 @@ def test_generate_noise_std_matches_requested_level():
gen = make_generator() gen = make_generator()
noise_std = 0.02 noise_std = 0.02
planes_noisy = gen.generate( planes_noisy = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=noise_std, seed=1 coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=noise_std, seed=1
)
planes_clean = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], image_shape=IMAGE_SHAPE, noise_std=0.0
) )
planes_clean = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.0)
diff = planes_noisy[0].flux - planes_clean[0].flux diff = planes_noisy[0].flux - planes_clean[0].flux
assert np.std(diff) == pytest.approx(noise_std, rel=0.15) assert np.std(diff) == pytest.approx(noise_std, rel=0.15)
def test_generate_viewing_angle_compresses_tilt_axis(): def test_generate_applies_z_tolerance_to_every_plane():
gen = make_generator() gen = make_generator()
planes_straight = gen.generate( planes = gen.generate(
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0 coefficients={(0, 0): 1 + 0j},
z_list=[0.3, 0.4, 0.5],
image_shape=IMAGE_SHAPE,
z_tolerance=0.02,
) )
planes_tilted = gen.generate( assert all(p.z_tolerance == 0.02 for p in planes)
coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0
def test_generate_applies_nominal_z_offset_independent_of_true_z():
gen = make_generator()
true_z_list = [0.3, 0.4, 0.5]
offsets = {0.3: 0.01, 0.4: -0.005, 0.5: 0.0}
planes = gen.generate(
coefficients={(0, 0): 1 + 0j},
z_list=true_z_list,
image_shape=IMAGE_SHAPE,
nominal_z_offsets=offsets,
) )
def width_along_axis(flux, axis): nominal_zs = [p.z for p in planes]
profile = flux[flux.shape[0] // 2, :] if axis == 1 else flux[:, flux.shape[1] // 2] assert nominal_zs == pytest.approx([0.31, 0.395, 0.5])
half_max = profile.max() / 2 # The flux is still rendered at each plane's *true* z (0.3, 0.4, 0.5),
above = np.where(profile >= half_max)[0] # not its offset nominal z -- verified indirectly in Task 7's
return above[-1] - above[0] # end-to-end tolerance-recovery test.
width_straight_x = width_along_axis(planes_straight[0].flux, axis=1)
width_tilted_x = width_along_axis(planes_tilted[0].flux, axis=1)
width_straight_y = width_along_axis(planes_straight[0].flux, axis=0)
width_tilted_y = width_along_axis(planes_tilted[0].flux, axis=0)
# tilt compresses the viewed beam along the tilt (x) axis, y unaffected
assert width_tilted_x < width_straight_x
assert width_tilted_y == pytest.approx(width_straight_y, abs=1)