From 03b63ba03a065ea8fdcd92548a8176aa0a25d320 Mon Sep 17 00:00:00 2001 From: Martino Ferrari Date: Thu, 2 Jul 2026 21:47:40 +0200 Subject: [PATCH] Initial commit: he11lib mode-purity reconstruction library Full implementation of Laguerre-Gauss modal reconstruction for gyrotron beam diagnostics, per the approved design spec, plus tests, docs, and a runnable end-to-end example. Co-Authored-By: Claude Sonnet 5 --- .gitignore | 23 + CLAUDE.md | 104 +++ LICENSE | 674 ++++++++++++++++++ README.md | 42 ++ docs/api.md | 223 ++++++ .../2026-07-02-gyrotron-mode-purity-design.md | 212 ++++++ examples/full_pipeline_example.py | 108 +++ he11lib/__init__.py | 36 + he11lib/data.py | 88 +++ he11lib/deconvolution.py | 58 ++ he11lib/fitting.py | 233 ++++++ he11lib/geometry.py | 64 ++ he11lib/modes.py | 96 +++ he11lib/noise.py | 40 ++ he11lib/phase_retrieval.py | 135 ++++ he11lib/plotting.py | 65 ++ he11lib/reconstruct.py | 124 ++++ he11lib/synthetic.py | 83 +++ pyproject.toml | 27 + tests/__init__.py | 0 tests/conftest.py | 3 + tests/test_data.py | 93 +++ tests/test_deconvolution.py | 67 ++ tests/test_fitting.py | 134 ++++ tests/test_geometry.py | 77 ++ tests/test_modes.py | 103 +++ tests/test_noise.py | 46 ++ tests/test_phase_retrieval.py | 84 +++ tests/test_plotting.py | 64 ++ tests/test_reconstruct.py | 112 +++ tests/test_synthetic.py | 123 ++++ 31 files changed, 3341 insertions(+) create mode 100644 .gitignore create mode 100644 CLAUDE.md create mode 100644 LICENSE create mode 100644 README.md create mode 100644 docs/api.md create mode 100644 docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md create mode 100644 examples/full_pipeline_example.py create mode 100644 he11lib/__init__.py create mode 100644 he11lib/data.py create mode 100644 he11lib/deconvolution.py create mode 100644 he11lib/fitting.py create mode 100644 he11lib/geometry.py create mode 100644 he11lib/modes.py create mode 100644 he11lib/noise.py create mode 100644 he11lib/phase_retrieval.py create mode 100644 he11lib/plotting.py create mode 100644 he11lib/reconstruct.py create mode 100644 he11lib/synthetic.py create mode 100644 pyproject.toml create mode 100644 tests/__init__.py create mode 100644 tests/conftest.py create mode 100644 tests/test_data.py create mode 100644 tests/test_deconvolution.py create mode 100644 tests/test_fitting.py create mode 100644 tests/test_geometry.py create mode 100644 tests/test_modes.py create mode 100644 tests/test_noise.py create mode 100644 tests/test_phase_retrieval.py create mode 100644 tests/test_plotting.py create mode 100644 tests/test_reconstruct.py create mode 100644 tests/test_synthetic.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b77e775 --- /dev/null +++ b/.gitignore @@ -0,0 +1,23 @@ +# Python +__pycache__/ +*.py[cod] +*.egg-info/ +.eggs/ +build/ +dist/ + +# Virtual environments +.venv/ +venv/ + +# Test / tool caches +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ + +# Editors +.vscode/ +.idea/ + +# Claude Code local settings (machine-specific, not project config) +.claude/settings.local.json diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..3a47c72 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,104 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## What this is + +`he11lib` reconstructs the Laguerre-Gauss (LG) modal content ("mode purity") of a +free-space-propagating gyrotron RF beam from a set of thermal (flux) images taken at +different distances from the output window. It accounts for camera geometry (unknown +pixel scale / oblique viewing angle), sensor noise, target thermal-diffusion blur, and +unknown beam center/pointing. + +The full design rationale lives in +`docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md` — read it before +making architectural changes. `docs/api.md` is the API reference for the implemented +public interface; keep it in sync when changing public signatures. + +## Commands + +```bash +pip install -e ".[dev]" # editable install with test dependencies (numpy, scipy, matplotlib, pytest) +pytest # run the full test suite (discovers tests/, per pyproject.toml) +pytest tests/test_modes.py # run one test file +pytest tests/test_modes.py::test_field_at_waist # run one test +python examples/full_pipeline_example.py # runnable end-to-end demo +``` + +There is a project `.venv`; activate it or otherwise ensure the editable install's +dependencies are available before running tests. + +`tests/conftest.py` forces the `Agg` matplotlib backend so plotting tests run headless. + +This project is not (yet) a git repository. + +## Architecture + +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 +`pixel_scale`/`viewing_angle_deg`. Everything downstream is keyed off `LGBasis`, which +defines the mode basis relative to a known waist `w0`/`z0`/`wavelength`. + +Module responsibilities (`he11lib/`): + +- **`data.py`** — `MeasurementPlane`, `ReconstructionResult` (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)`, + 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. +- **`geometry.py`** — `GeometryCalibration`: resolves a plane's pixel-to-physical + coordinate grid, deferring to known `pixel_scale`/`viewing_angle_deg` on the plane + over any override passed in. +- **`noise.py`** — `NoiseEstimator`: automatic per-image noise-std estimation + (Laplacian method) and per-pixel weights for noise-weighted least squares. +- **`deconvolution.py`** — `DiffusionDeconvolver`: optional forward blur / Wiener + 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 + oblique view makes x/y pixel scales differ) — an accepted approximation, not a bug. +- **`synthetic.py`** — `SyntheticBeamGenerator`: forward model that produces + `MeasurementPlane`s from known ground-truth coefficients/center/pointing/geometry. + 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 + core joint nonlinear least-squares fit (complex LG coefficients + beam + center/pointing + unknown geometry) via `scipy.optimize.least_squares`. `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). +- **`phase_retrieval.py`** — `propagate_angular_spectrum` (FFT-based paraxial + free-space propagation) and `PhaseRetriever` (multi-plane Gerchberg-Saxton), the + fallback reconstruction path for when a finite mode basis doesn't fit well. +- **`reconstruct.py`** — `BeamReconstructor`: the orchestrator. Pipeline order: + validate planes → optional deconvolution (requires known `pixel_scale` per plane) → + `ModalFitter.fit_auto` → optional `PhaseRetriever` fallback (forced via + `force_phase_retrieval`, or triggered automatically when the noise-weighted RMS + residual exceeds `phase_retrieval_residual_threshold`). The fallback path projects + the recovered field onto all modes up to `max_order` and produces a + `ReconstructionResult` with `used_phase_retrieval=True`, empty `residuals`, and NaN + `coefficient_uncertainty` (no fit covariance available from phase retrieval). +- **`plotting.py`** — diagnostic figures (`plot_mode_purity`, `plot_center_trace`, + `plot_residuals`); each returns a `Figure` rather than calling `plt.show()`. + +Everything above is re-exported from the top-level `he11lib` package (see +`he11lib/__init__.py`); import from there rather than submodules. + +## Known physics/fitting pitfalls (read before writing new tests or examples) + +These aren't library bugs — they're consequences of realistic optics parameters and of +automatic order-selection being genuinely data-driven — but they've caused most of the +debugging time in this project's history: + +1. **Rayleigh-range / frame clipping.** `w(z)` grows with `|z - z0|` relative to + `zR = pi*w0**2/wavelength`. With typical test parameters (`w0=5e-3`, + `wavelength=1.76e-3`), `zR` is only ~4.46 cm, so z-distances spanning tens of cm put + the beam many Rayleigh ranges out, where `w(z)` can exceed a small test frame — + clipping the beam and corrupting fits, or introducing FFT wraparound artifacts in + `propagate_angular_spectrum`. Keep z-distances within roughly ±1-2 Rayleigh ranges of + `z0`, or enlarge the frame/pixel_scale accordingly. +2. **Automatic mode-set growth can overfit deconvolution artifacts.** Wiener-deconvolved + data always has some residual imperfection; `fit_auto`'s BIC-driven growth will try + to "explain" it with spurious higher-order modes, degrading fitted beam + center/pointing via parameter degeneracy (observed: pointing angle off by 4-6x at + `max_order=3` vs. matching ground truth almost exactly at `max_order=1`, for the same + 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, + unless the test specifically targets growth behavior itself. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..f288702 --- /dev/null +++ b/LICENSE @@ -0,0 +1,674 @@ + GNU GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU General Public License is a free, copyleft license for +software and other kinds of works. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +the GNU General Public License is intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. We, the Free Software Foundation, use the +GNU General Public License for most of our software; it applies also to +any other work released this way by its authors. You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + To protect your rights, we need to prevent others from denying you +these rights or asking you to surrender the rights. Therefore, you have +certain responsibilities if you distribute copies of the software, or if +you modify it: responsibilities to respect the freedom of others. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must pass on to the recipients the same +freedoms that you received. You must make sure that they, too, receive +or can get the source code. And you must show them these terms so they +know their rights. + + Developers that use the GNU GPL protect your rights with two steps: +(1) assert copyright on the software, and (2) offer you this License +giving you legal permission to copy, distribute and/or modify it. + + For the developers' and authors' protection, the GPL clearly explains +that there is no warranty for this free software. For both users' and +authors' sake, the GPL requires that modified versions be marked as +changed, so that their problems will not be attributed erroneously to +authors of previous versions. + + Some devices are designed to deny users access to install or run +modified versions of the software inside them, although the manufacturer +can do so. This is fundamentally incompatible with the aim of +protecting users' freedom to change the software. The systematic +pattern of such abuse occurs in the area of products for individuals to +use, which is precisely where it is most unacceptable. Therefore, we +have designed this version of the GPL to prohibit the practice for those +products. If such problems arise substantially in other domains, we +stand ready to extend this provision to those domains in future versions +of the GPL, as needed to protect the freedom of users. + + Finally, every program is threatened constantly by software patents. +States should not allow patents to restrict development and use of +software on general-purpose computers, but in those that do, we wish to +avoid the special danger that patents applied to a free program could +make it effectively proprietary. To prevent this, the GPL assures that +patents cannot be used to render the program non-free. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Use with the GNU Affero General Public License. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU Affero General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the special requirements of the GNU Affero General Public License, +section 13, concerning interaction through a network will apply to the +combination as such. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + + END OF TERMS AND CONDITIONS + + How to Apply These Terms to Your New Programs + + If you develop a new program, and you want it to be of the greatest +possible use to the public, the best way to achieve this is to make it +free software which everyone can redistribute and change under these terms. + + To do so, attach the following notices to the program. It is safest +to attach them to the start of each source file to most effectively +state the exclusion of warranty; and each file should have at least +the "copyright" line and a pointer to where the full notice is found. + + + Copyright (C) + + This program is free software: you can redistribute it and/or modify + it under the terms of the GNU General Public License as published by + the Free Software Foundation, either version 3 of the License, or + (at your option) any later version. + + This program is distributed in the hope that it will be useful, + but WITHOUT ANY WARRANTY; without even the implied warranty of + MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + GNU General Public License for more details. + + You should have received a copy of the GNU General Public License + along with this program. If not, see . + +Also add information on how to contact you by electronic and paper mail. + + If the program does terminal interaction, make it output a short +notice like this when it starts in an interactive mode: + + Copyright (C) + This program comes with ABSOLUTELY NO WARRANTY; for details type `show w'. + This is free software, and you are welcome to redistribute it + under certain conditions; type `show c' for details. + +The hypothetical commands `show w' and `show c' should show the appropriate +parts of the General Public License. Of course, your program's commands +might be different; for a GUI interface, you would use an "about box". + + You should also get your employer (if you work as a programmer) or school, +if any, to sign a "copyright disclaimer" for the program, if necessary. +For more information on this, and how to apply and follow the GNU GPL, see +. + + The GNU General Public License does not permit incorporating your program +into proprietary programs. If your program is a subroutine library, you +may consider it more useful to permit linking proprietary applications with +the library. If this is what you want to do, use the GNU Lesser General +Public License instead of this License. But first, please read +. diff --git a/README.md b/README.md new file mode 100644 index 0000000..465d46a --- /dev/null +++ b/README.md @@ -0,0 +1,42 @@ +# he11lib + +Mode purity reconstruction of a free-space-propagating gyrotron RF beam from a +set of thermal (flux) images taken at different distances from the output +window. + +Decomposes the beam into Laguerre-Gauss (LG) modes referenced to a known +waist size/location/wavelength, accounting for camera geometry, sensor noise, +target thermal-diffusion blur, and unknown beam centering/pointing. + +See `docs/api.md` for the full API reference and +`examples/full_pipeline_example.py` for a runnable end-to-end +demonstration, and +`docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md` for the +full design. + +## Install (editable, for development) + +```bash +pip install -e ".[dev]" +``` + +## Run tests + +```bash +pytest +``` + +## License + +GPL-3.0-or-later. See `LICENSE`. + +## AI-assisted development disclosure + +The initial implementation of this library (design, code, tests, and +documentation) was produced with AI assistance: + +- **Model:** Claude Sonnet 5 (`claude-sonnet-5`) +- **Agent/harness:** Claude Code CLI, version 2.1.92 + +All output was reviewed and directed by a human collaborator throughout +development. diff --git a/docs/api.md b/docs/api.md new file mode 100644 index 0000000..8aea1f3 --- /dev/null +++ b/docs/api.md @@ -0,0 +1,223 @@ +# he11lib API Reference + +`he11lib` reconstructs the Laguerre-Gauss (LG) modal content ("mode purity") +of a free-space-propagating gyrotron RF beam from a set of thermal (flux) +images taken at different distances from the output window. + +See `examples/full_pipeline_example.py` for a runnable end-to-end +demonstration, and +`docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md` for the +full design rationale. + +Every class/function below is exported from the top-level `he11lib` package +(e.g. `from he11lib import BeamReconstructor`), except where noted. + +## Quick start + +```python +from he11lib import BeamReconstructor, MeasurementPlane + +# planes: a list of >=3 MeasurementPlane objects built from your own +# flux arrays (see MeasurementPlane below). +reconstructor = BeamReconstructor(w0=5e-3, z0=0.5, wavelength=1.76e-3) +result = reconstructor.reconstruct(planes) + +for mode, (power_fraction, phase_rad) in result.purity.items(): + print(mode, power_fraction, phase_rad) +``` + +## `data` — `MeasurementPlane`, `ReconstructionResult` + +### `MeasurementPlane(flux, z, pixel_scale=None, viewing_angle_deg=None, label=None)` + +One measurement: a 2D flux array plus its acquisition metadata. + +- `flux` — 2D `np.ndarray` of flux values. Dead-pixel correction, background + subtraction, and saturation clipping are assumed already handled upstream. +- `z` — nominal distance from the output window, in meters. Must be `> 0`. +- `pixel_scale` — known meters/pixel, or `None` if unknown (then jointly + fit by `ModalFitter`/`BeamReconstructor`). +- `viewing_angle_deg` — known camera viewing angle relative to the beam + axis, in degrees, or `None` if unknown (also jointly fit). +- `label` — optional human-readable identifier. + +### `validate_planes(planes)` + +Raises `ValueError` if there are fewer than 3 planes, planes have +mismatched flux shapes, or `z` values are not all distinct. Called +internally by `ModalFitter.fit`/`fit_auto`, `PhaseRetriever.retrieve`, and +`BeamReconstructor.reconstruct` — you generally don't need to call it +yourself. Not exported from the top-level package; import via +`from he11lib.data import validate_planes` if needed. + +### `ReconstructionResult` + +Output of a full reconstruction (returned by `ModalFitter.fit`/`fit_auto` +and `BeamReconstructor.reconstruct`): + +- `purity: dict[(p, l), (power_fraction, phase_rad)]` +- `reconstructed_field: np.ndarray` — reconstructed complex field. +- `centers: list[(x, y)]` — fitted beam transverse center per plane, meters. +- `pointing_angle_deg: float` — fitted shared beam pointing angle (tilt). +- `geometry: dict[str, float]` — geometry parameters used or fitted (keys + `pixel_scale_{i}`, `viewing_angle_deg_{i}` per plane index `i`). +- `residuals: list[np.ndarray]` — per-plane (measured − modeled) flux maps. + Empty when `used_phase_retrieval` is `True`. +- `coefficient_uncertainty: dict[(p, l), float]` — 1-sigma uncertainty on + each mode's fitted power fraction. `NaN` per mode when + `used_phase_retrieval` is `True`. +- `used_phase_retrieval: bool` — whether the phase-retrieval fallback (not + the modal fit) produced this result. + +## `modes` — `LGBasis` + +`LGBasis(w0, z0, wavelength)` — the LG mode basis referenced to a known +waist radius `w0` (m), waist location `z0` (m), and radiation `wavelength` +(m). + +- `beam_radius(z)` — `w(z)`. +- `inverse_radius_of_curvature(z)` — `1/R(z)` (well-defined, `0`, at the + waist). +- `gouy_phase(z, p, l)` — Gouy phase of mode `(p, l)` at `z`. +- `field(x, y, z, p, l)` — complex `LG_{p,l}` field sampled on the `(x, y)` + grid at distance `z`. +- `field_superposition(x, y, z, coefficients)` — complex field for + `coefficients: dict[(p, l), complex]`. +- `project(complex_field, x, y, dx, z, modes)` — projects `complex_field` + onto each `(p, l)` in `modes`, returning `dict[(p, l), complex]` + coefficients (Riemann-sum inner product; `dx` is the grid spacing). + +## `geometry` — `GeometryCalibration` + +`GeometryCalibration(plane)` wraps a single `MeasurementPlane` and resolves +its pixel-to-physical-coordinate mapping. + +- `pixel_scale_known` / `viewing_angle_known` — `bool` properties. +- `physical_coordinates(pixel_scale=None, viewing_angle_deg=None)` — + returns `(x, y)` physical coordinate grids matching the plane's flux + shape. Values known on the `MeasurementPlane` take precedence over the + `override` arguments; raises `ValueError` if a value is neither known nor + overridden. + +## `noise` — `NoiseEstimator` + +`NoiseEstimator()` — automatic per-image noise estimation (no +user-supplied noise parameter needed). + +- `estimate_std(image)` — fast Laplacian-based (Immerkær 1996) noise + standard-deviation estimate. +- `weights(image)` — per-pixel weights (`1/sigma**2`) for noise-weighted + least squares. + +## `deconvolution` — `DiffusionDeconvolver` + +`DiffusionDeconvolver(thermal_diffusivity, dwell_time)` — optional +correction for lateral thermal-diffusion blur in the absorbing target +(`thermal_diffusivity` in m²/s, `dwell_time` in s). Disabled unless you +pass a `deconvolver` to `BeamReconstructor`. + +- `blur_sigma_m()` — Gaussian blur standard deviation, in meters. +- `blur(image, pixel_scale)` — forward blur (for synthetic testing). +- `deconvolve(image, pixel_scale, noise_to_signal_ratio=1e-3)` — regularized + (Wiener) removal of the blur. + +Note: the blur/deconvolution kernel is isotropic in pixel space. If a +plane has a nonzero `viewing_angle_deg`, its `x` and `y` pixel axes have +different physical scales (see `SyntheticBeamGenerator` below), so +deconvolution is only exact for `viewing_angle_deg == 0`; at oblique +angles it is an approximation. + +## `synthetic` — `SyntheticBeamGenerator` + +`SyntheticBeamGenerator(basis, image_shape, pixel_scale)` — forward model +used to validate the pipeline against known ground truth, and to evaluate +experimental design (e.g. "would these distances separate my modes?"). +`pixel_scale` is the physical pixel size, in meters, along the non-tilted +`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)` + — returns one `MeasurementPlane` per `z` in `z_list`. The beam's + transverse center drifts linearly with `z` according to + `pointing_angle_deg`, starting from `center` at `z0`. + +## `fitting` — `ModalFitter`, `generate_mode_shells` + +### `generate_mode_shells(max_order)` + +Groups candidate `LG_{p,l}` modes into shells of increasing order +`2p + |l|`, up to and including `max_order`. Returns +`list[list[(p, l)]]`, one list of modes per order. + +### `ModalFitter(basis, noise_estimator=None)` + +Core reconstruction path: a joint nonlinear least-squares fit of complex LG +coefficients, beam center/pointing, and (if unknown) 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` + — fits exactly the given candidate `modes`. +- `fit_auto(planes, max_order=4, bic_improvement_threshold=10.0) -> ReconstructionResult` + — starts from `LG_00` and grows the candidate mode set shell-by-shell + (via `generate_mode_shells`), stopping once BIC no longer improves by + more than `bic_improvement_threshold`, capped at `max_order`. Emits a + `UserWarning` (does not raise) if the cap is reached while the fit is + still improving. + +## `phase_retrieval` — `PhaseRetriever`, `propagate_angular_spectrum` + +Fallback reconstruction path for when the modal fit's residual stays high, +or when the mode content isn't well described by a small finite mode set. + +### `propagate_angular_spectrum(field, dx, dz, wavelength)` + +Free-space-propagates a complex `field` (pixel spacing `dx`) by distance +`dz` via the (paraxial) angular-spectrum method — the same propagation +model implicitly assumed by `LGBasis`'s closed-form paraxial modes. + +### `PhaseRetriever(wavelength)` + +- `retrieve(planes, pixel_scale=None, viewing_angle_deg=None, max_iterations=200) -> PhaseRetrievalResult` + — multi-plane Gerchberg-Saxton phase retrieval: propagates a trial + complex field back and forth between planes, enforcing the measured + amplitude (`sqrt(flux)`) at each plane, without assuming a finite mode + basis. + +### `PhaseRetrievalResult` + +`field, x, y, z, center, residual` — the recovered complex field (at the +smallest-`z` plane) on its `(x, y)` grid, the estimated beam center +(intensity centroid), and the final RMS amplitude-mismatch residual. +Project `field` onto `LGBasis` (via `LGBasis.project`) to get a purity +table, as `BeamReconstructor` does internally for its fallback path. + +## `reconstruct` — `BeamReconstructor` + +`BeamReconstructor(w0, z0, wavelength, 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 +diffusion deblurring → `ModalFitter.fit_auto` → optional +`PhaseRetriever` fallback. + +- `reconstruct(planes) -> ReconstructionResult` + 1. Validates `planes` (see `validate_planes`). + 2. If `deconvolver` is set, deblurs each plane (raises `ValueError` if a + plane's `pixel_scale` isn't known). + 3. Runs `ModalFitter(basis, noise_estimator).fit_auto(planes, max_order)`. + 4. Runs the `PhaseRetriever` fallback instead, projecting its recovered + field onto all modes up to `max_order`, if `force_phase_retrieval` is + `True`, or if `phase_retrieval_residual_threshold` is set and the + modal fit's noise-weighted RMS residual exceeds it. In that case + `result.residuals` is empty and `coefficient_uncertainty` is `NaN` + per mode (phase retrieval doesn't produce a fit covariance). + +## `plotting` — diagnostic visualizations + +Each function returns a `matplotlib.figure.Figure` for the caller to +display (`fig.show()`) or save (`fig.savefig(...)`); none of them call +`plt.show()` themselves. + +- `plot_mode_purity(result)` — bar chart of power fraction per mode. +- `plot_center_trace(planes, result)` — fitted beam center `(x, y)` vs. `z`. +- `plot_residuals(planes, result)` — per-plane residual maps. Raises + `ValueError` if `result.residuals` is empty (e.g. after the + phase-retrieval fallback). diff --git a/docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md b/docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md new file mode 100644 index 0000000..9f3e6fe --- /dev/null +++ b/docs/superpowers/specs/2026-07-02-gyrotron-mode-purity-design.md @@ -0,0 +1,212 @@ +# he11lib — Gyrotron Beam Mode Purity Reconstruction + +**Date:** 2026-07-02 +**Status:** Approved for planning + +## Purpose + +A general-purpose, reusable Python library for reconstructing the Laguerre-Gauss +(LG) modal content ("mode purity") of a free-space-propagating gyrotron RF beam, +from a set of thermal (flux) images captured at different distances from the +gyrotron output window. The library must account for real-world measurement +issues: unknown/partial camera geometry, sensor noise, target thermal-diffusion +blur, and unknown beam pointing/centering — not just an idealized intensity fit. + +This is a from-scratch library (empty project directory), intended to be used by +others beyond the initial author. + +## Scope + +- Beam regime: **free-space quasi-optical propagation** (post mode-converter / + output window), not waveguide-confined hybrid modes. +- Mode basis: **Laguerre-Gauss modes** `LG_{p,l}`, referenced to a **known** + waist size `w0` and waist location `z0` supplied by the user (design values + from the gyrotron/mode-converter specs) — these are inputs, not fit + parameters. +- Input data: **pre-processed NumPy flux arrays** (already extracted from raw + NVF radiometric camera files by the user's own pipeline). Dead-pixel + correction, background/ambient subtraction, and saturation-clipping detection + are already handled upstream and out of scope for this library. Because the + data is flux (not temperature), there is no temperature-to-power nonlinearity + to correct. +- Number of measurement planes: **3 to 10**, at distances spanning roughly + 0–100 cm from the output window (e.g. 30/40/50/60 cm), known to a nominal + precision (set by a translation stage or tape measure). +- Camera geometry (pixel-to-length scale, viewing angle relative to beam axis): + may be **known from a prior calibration** (supplied as input) **or unknown** + (estimated jointly with everything else). +- Beam **transverse center and pointing angle** are always unknown and must be + estimated from the images. +- Residual sensor noise (NETD-type, after upstream preprocessing) is + **auto-estimated per image**; no user-supplied noise parameter required for + the default path. +- Target thermal-diffusion blur correction (deconvolution) is **optional**, + parametrized by target thermal diffusivity + exposure/dwell time. +- Deliverables include full API docs and an end-to-end example + script/notebook demonstrating the complete pipeline. + +## Architecture + +Modular/composable design: each concern is an independent, testable component +with a clear interface, wired together by a high-level orchestrator. This +supports both simple end-to-end use and power-user access to individual +stages (e.g., using `LGBasis` standalone, or swapping in a custom noise +model). + +### Components + +- **`data.py` — `MeasurementPlane`, `ReconstructionResult`** + `MeasurementPlane`: container for one measurement — flux array, nominal `z` + distance, optional known pixel scale (mm/px) and viewing angle (deg), + optional metadata (timestamp, label). `ReconstructionResult`: container for + all pipeline outputs (see below). + +- **`geometry.py` — `GeometryCalibration`** + Applies projective/perspective correction to compensate for oblique camera + viewing angle, and converts pixel coordinates to physical length units. If + pixel scale and/or viewing angle are supplied on a `MeasurementPlane`, uses + them directly; if not supplied, exposes them as free parameters to be + solved jointly by the `ModalFitter`. + +- **`noise.py` — `NoiseEstimator`** + Estimates residual per-image noise standard deviation automatically (e.g. + from low-signal/background regions or high-frequency residual content). + Produces per-pixel or per-image weights used by the `ModalFitter`'s + noise-weighted least squares. + +- **`deconvolution.py` — `DiffusionDeconvolver`** + Optional step. Given target material thermal diffusivity and + exposure/dwell time, builds a blur kernel modeling lateral heat spreading + in the absorbing target and deconvolves it from each plane before fitting. + Disabled by default. + +- **`modes.py` — `LGBasis`** + Given reference `w0`, `z0`, generates complex LG mode fields `LG_{p,l}(x, + y, z)` at arbitrary transverse coordinates and axial distance `z`, + including correct Gouy phase and beam-radius evolution. Supports + evaluating a finite candidate set of `(p, l)` indices and projecting an + arbitrary complex field onto the basis (used both by the modal fit and by + the phase-retrieval fallback). + +- **`fitting.py` — `ModalFitter`** + Core reconstruction path. Parameters: complex LG coefficients + (amplitude + phase) for each candidate mode, beam transverse center `(x, + y)` per plane, a shared beam pointing angle (tilt), and any uncalibrated + geometry parameters (pixel scale / viewing angle) not supplied on the + `MeasurementPlane`s. Objective: noise-weighted nonlinear least-squares + residual between modeled `|Σ c_j · LG_j(x, y, z)|²` and measured flux, + summed over all planes. + + **Automatic mode-set growth**: starts from `LG_00`, incrementally adds + candidate modes in shells of increasing order, and stops growing once an + information-criterion (e.g. BIC) / residual-reduction test shows no + meaningful improvement, subject to a configurable maximum order cap (for + tractability and as a safety bound). Warns (does not error) if the cap is + hit while still improving, or if final residuals remain large. + +- **`phase_retrieval.py` — `PhaseRetriever`** + Optional/fallback path. Gerchberg-Saxton-style iterative multi-plane phase + retrieval: propagates a trial complex field back and forth between + measurement planes (via free-space propagation, e.g. angular spectrum), + enforcing the measured amplitude (sqrt of flux) at each plane, without + assuming a finite mode basis. Used when explicitly requested, or + automatically as a fallback when `ModalFitter`'s residual stays high after + mode-set growth completes. The recovered field is then projected onto + `LGBasis` to produce a purity table. + +- **`synthetic.py` — `SyntheticBeamGenerator`** + Forward model: given known mode coefficients, `w0`/`z0`, geometry + (center, tilt, viewing angle, pixel scale), noise level, and optional + diffusion blur, generates synthetic multi-plane flux images. Used for + library validation (recover known ground truth) and for users to test + experimental design (e.g., "would these 4 distances separate my modes?"). + +- **`reconstruct.py` — `BeamReconstructor`** + High-level orchestrator: given a list of `MeasurementPlane`s and + configuration (known `w0`/`z0`, optional deconvolution params, optional + mode-set cap, optional forced phase-retrieval mode), runs geometry + correction → noise estimation → optional deconvolution → modal fit (with + automatic growth) → optional phase-retrieval fallback → produces a + `ReconstructionResult`. Each stage remains independently accessible for + power users. + +- **`plotting.py`** + Diagnostic visualizations: measured vs. reconstructed intensity per plane, + residual maps, mode purity bar chart, beam center/pointing trace across + planes. + +### Data flow + +1. Build a list of `MeasurementPlane` objects (flux array + nominal z + + optional known geometry per plane). +2. `GeometryCalibration` applies projective correction using supplied + pixel-scale/viewing-angle, or flags them as unknowns. +3. `NoiseEstimator` computes per-plane noise weights. +4. `DiffusionDeconvolver` optionally deblurs each plane. +5. `ModalFitter` runs the joint noise-weighted nonlinear least-squares fit + over LG coefficients + center + pointing + any uncalibrated geometry + params, growing the mode set automatically. +6. If requested, or if fit residual remains high, `PhaseRetriever` runs as a + fallback and its result is projected onto `LGBasis`. +7. `BeamReconstructor` assembles a `ReconstructionResult` containing: mode + purity table (power fraction + phase per mode), reconstructed complex + field, fitted beam center/pointing per plane, geometry parameters used or + fitted, per-plane residual maps, and coefficient uncertainties (from the + fit's covariance). + +## Testing strategy + +`SyntheticBeamGenerator` is the backbone of validation: generate synthetic +multi-plane data from known ground-truth mode content, geometry, and noise, +run it through the full pipeline (and through individual components in +isolation), and assert recovered parameters match ground truth within +tolerance. Individual components also get targeted unit tests against known +analytic cases (e.g. `LGBasis` orthogonality and known Gouy phase values, +single-pure-mode recovery, geometry correction on synthetic projective +distortions). + +## Error handling + +Validate only at boundaries: reject malformed `MeasurementPlane` inputs +(mismatched array shapes, non-positive `z`, fewer than 3 planes). Warn +(rather than raise) when automatic mode-set growth hits its configured cap +while still improving, or when final fit residuals remain large — the caller +gets the result plus a diagnostic flag, not a crash. + +## Dependencies + +NumPy, SciPy (`optimize` for the nonlinear least-squares fit, `ndimage` for +projective transforms and deconvolution), Matplotlib for diagnostic +plotting. No GPU requirement. + +## Package layout + +``` +he11lib/ + data.py # MeasurementPlane, ReconstructionResult + geometry.py # GeometryCalibration + noise.py # NoiseEstimator + deconvolution.py # DiffusionDeconvolver + modes.py # LGBasis + fitting.py # ModalFitter (+ automatic mode-set growth) + phase_retrieval.py # PhaseRetriever + synthetic.py # SyntheticBeamGenerator + reconstruct.py # BeamReconstructor (orchestrator) + plotting.py # diagnostic visualizations +docs/ + ... # API docs +examples/ + full_pipeline_example.py # end-to-end demo of the whole pipeline +tests/ + ... +``` + +## Deliverables + +- Full implementation of all components above. +- API documentation for the public interface of every module. +- An end-to-end example (script or notebook) exercising the complete + pipeline on synthetic data, from `SyntheticBeamGenerator` through + `BeamReconstructor` to `plotting.py` diagnostics. +- Test suite covering synthetic ground-truth recovery and per-component unit + tests. diff --git a/examples/full_pipeline_example.py b/examples/full_pipeline_example.py new file mode 100644 index 0000000..d3c3a24 --- /dev/null +++ b/examples/full_pipeline_example.py @@ -0,0 +1,108 @@ +"""End-to-end demonstration of the he11lib reconstruction pipeline. + +Simulates a gyrotron beam that is mostly the LG_00 fundamental mode with a +small admixture of LG_01, viewed by a thermal camera at four distances from +the output window. The camera has an unknown transverse offset/pointing and +adds sensor noise; the target also has some thermal-diffusion blur that we +correct for. We then reconstruct the mode purity, beam center/pointing, and +plot the diagnostics. + +Run with: + + python examples/full_pipeline_example.py +""" + +from __future__ import annotations + +import matplotlib.pyplot as plt + +from he11lib import ( + BeamReconstructor, + DiffusionDeconvolver, + LGBasis, + SyntheticBeamGenerator, + plot_center_trace, + plot_mode_purity, + plot_residuals, +) + +# --- Known reference beam parameters (from the gyrotron/mode-converter design) --- +W0 = 5e-3 # reference waist radius, meters +Z0 = 0.5 # reference waist location, meters from the output window +WAVELENGTH = 1.76e-3 # radiation wavelength, meters (e.g. a 170 GHz gyrotron) + +# --- Ground truth for the synthetic beam (unknown to the reconstructor) --- +TRUE_COEFFICIENTS = {(0, 0): 0.95 + 0j, (0, 1): 0.25 + 0.05j} +TRUE_CENTER = (0.4e-3, -0.3e-3) # beam offset from the camera's optical axis +TRUE_POINTING_DEG = 0.15 # beam pointing (tilt) angle +CAMERA_VIEWING_ANGLE_DEG = 5.0 # oblique camera viewing angle (known) +CAMERA_PIXEL_SCALE = 4e-4 # meters/pixel (known calibration) +IMAGE_SHAPE = (81, 81) +# Measurement plane distances, meters. Kept within roughly +/-2 Rayleigh +# ranges of z0 so the (widening) beam stays well within the camera frame -- +# planes much farther out would be clipped by the finite frame, which +# degrades the fit. +Z_LIST = [0.4, 0.45, 0.55, 0.6] + +# --- Target thermal-diffusion blur (known target material properties) --- +THERMAL_DIFFUSIVITY = 1e-6 # m^2/s +DWELL_TIME = 0.2 # s + + +def main() -> None: + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + generator = SyntheticBeamGenerator( + basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=CAMERA_PIXEL_SCALE + ) + + planes = generator.generate( + coefficients=TRUE_COEFFICIENTS, + z_list=Z_LIST, + center=TRUE_CENTER, + pointing_angle_deg=TRUE_POINTING_DEG, + viewing_angle_deg=CAMERA_VIEWING_ANGLE_DEG, + noise_std=2e-4, + seed=42, + ) + + # Apply the same thermal-diffusion blur a real target would exhibit. + blur_deconvolver = DiffusionDeconvolver( + thermal_diffusivity=THERMAL_DIFFUSIVITY, dwell_time=DWELL_TIME + ) + for plane in planes: + plane.flux = blur_deconvolver.blur(plane.flux, plane.pixel_scale) + + # The ground truth only has order-0 and order-1 content, so a max_order + # of 1 is enough for automatic mode-set growth to find it; growing much + # further would start fitting deconvolution/noise artifacts as spurious + # higher-order modes. + reconstructor = BeamReconstructor( + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + max_order=1, + deconvolver=blur_deconvolver, + ) + result = reconstructor.reconstruct(planes) + + print("Mode purity table (power fraction, phase [rad]):") + for mode, (fraction, phase) in sorted( + result.purity.items(), key=lambda item: -item[1][0] + ): + print(f" LG_{mode[0]},{mode[1]}: {fraction:6.3%} (phase {phase:+.3f} rad)") + + print(f"\nFitted pointing angle: {result.pointing_angle_deg:.4f} deg") + print("Fitted beam center per plane (m):") + for plane, (cx, cy) in zip(planes, result.centers): + print(f" z={plane.z:.2f} m -> ({cx:.3e}, {cy:.3e})") + + print(f"\nUsed phase-retrieval fallback: {result.used_phase_retrieval}") + + plot_mode_purity(result) + plot_center_trace(planes, result) + plot_residuals(planes, result) + plt.show() + + +if __name__ == "__main__": + main() diff --git a/he11lib/__init__.py b/he11lib/__init__.py new file mode 100644 index 0000000..790d428 --- /dev/null +++ b/he11lib/__init__.py @@ -0,0 +1,36 @@ +"""he11lib: mode purity reconstruction of free-space gyrotron beams. + +See docs/ for the full API and design documentation. +""" + +from .data import MeasurementPlane, ReconstructionResult +from .deconvolution import DiffusionDeconvolver +from .fitting import ModalFitter, generate_mode_shells +from .geometry import GeometryCalibration +from .modes import LGBasis +from .noise import NoiseEstimator +from .phase_retrieval import PhaseRetrievalResult, PhaseRetriever, propagate_angular_spectrum +from .plotting import plot_center_trace, plot_mode_purity, plot_residuals +from .reconstruct import BeamReconstructor +from .synthetic import SyntheticBeamGenerator + +__all__ = [ + "MeasurementPlane", + "ReconstructionResult", + "BeamReconstructor", + "DiffusionDeconvolver", + "ModalFitter", + "generate_mode_shells", + "GeometryCalibration", + "LGBasis", + "NoiseEstimator", + "PhaseRetrievalResult", + "PhaseRetriever", + "propagate_angular_spectrum", + "plot_center_trace", + "plot_mode_purity", + "plot_residuals", + "SyntheticBeamGenerator", +] + +__version__ = "0.1.0" diff --git a/he11lib/data.py b/he11lib/data.py new file mode 100644 index 0000000..b88b889 --- /dev/null +++ b/he11lib/data.py @@ -0,0 +1,88 @@ +"""Data containers for he11lib: measurement inputs and reconstruction outputs.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import numpy as np + + +@dataclass +class MeasurementPlane: + """A single thermal (flux) image at a known distance from the output window. + + Parameters + ---------- + flux : 2D array of flux values (already dead-pixel/background/saturation + corrected upstream). + z : distance from the output window, in meters. Must be positive. + pixel_scale : known mm/pixel scale, if calibrated. If None, treated as an + unknown to be estimated jointly during reconstruction. + viewing_angle_deg : known camera viewing angle relative to the beam axis, + in degrees, if calibrated. If None, treated as an unknown. + label : optional human-readable label (e.g. "plane_40cm"). + """ + + flux: np.ndarray + z: float + pixel_scale: float | None = None + viewing_angle_deg: float | None = None + label: str | None = None + + def __post_init__(self) -> None: + if self.flux.ndim != 2: + raise ValueError( + f"MeasurementPlane.flux must be a 2D array, got shape {self.flux.shape}" + ) + if self.z <= 0: + raise ValueError(f"MeasurementPlane.z must be positive, got {self.z}") + + +def validate_planes(planes: list[MeasurementPlane]) -> None: + """Validate a list of MeasurementPlanes for use in reconstruction. + + Raises ValueError if there are fewer than 3 planes, shapes mismatch + across planes, or z distances are not distinct. + """ + if len(planes) < 3: + raise ValueError( + f"At least 3 measurement planes are required, got {len(planes)}" + ) + + shapes = {p.flux.shape for p in planes} + if len(shapes) > 1: + raise ValueError(f"All MeasurementPlanes must have the same shape, got {shapes}") + + z_values = [p.z for p in planes] + if len(set(z_values)) != len(z_values): + raise ValueError(f"MeasurementPlane z distances must be distinct, got {z_values}") + + +@dataclass +class ReconstructionResult: + """Output of a full mode-purity reconstruction. + + Parameters + ---------- + purity : mapping from (p, l) mode index to (power_fraction, phase_rad). + reconstructed_field : reconstructed complex field (at the reference + waist, or as configured). + centers : fitted beam transverse center (x, y) in meters, per plane. + pointing_angle_deg : fitted shared beam pointing angle (tilt), in degrees. + geometry : geometry parameters used or fitted (e.g. pixel_scale, + viewing_angle_deg), keyed by name. + residuals : per-plane residual maps (measured - modeled flux). + coefficient_uncertainty : mapping from (p, l) mode index to the + 1-sigma uncertainty on its fitted power fraction. + used_phase_retrieval : whether the phase-retrieval fallback was used + instead of (or to seed) the modal fit. + """ + + purity: dict[tuple[int, int], tuple[float, float]] + reconstructed_field: np.ndarray + centers: list[tuple[float, float]] + pointing_angle_deg: float + geometry: dict[str, float] = field(default_factory=dict) + residuals: list[np.ndarray] = field(default_factory=list) + coefficient_uncertainty: dict[tuple[int, int], float] = field(default_factory=dict) + used_phase_retrieval: bool = False diff --git a/he11lib/deconvolution.py b/he11lib/deconvolution.py new file mode 100644 index 0000000..a76b6eb --- /dev/null +++ b/he11lib/deconvolution.py @@ -0,0 +1,58 @@ +"""Optional correction for lateral thermal-diffusion blur in the target. + +The absorbing target spreads heat laterally during the camera's exposure / +dwell time, blurring the true beam profile by a Gaussian kernel whose width +follows from the target's thermal diffusivity. This module provides both +the forward blur (for synthetic testing) and a regularized (Wiener) +deconvolution to remove it. +""" + +from __future__ import annotations + +import numpy as np +from scipy.ndimage import gaussian_filter + + +class DiffusionDeconvolver: + """Models and removes lateral thermal-diffusion blur in a target. + + Parameters + ---------- + thermal_diffusivity : target material thermal diffusivity, in m^2/s. + dwell_time : camera exposure/integration time, in s. + """ + + def __init__(self, thermal_diffusivity: float, dwell_time: float): + self.thermal_diffusivity = thermal_diffusivity + self.dwell_time = dwell_time + + def blur_sigma_m(self) -> float: + """Gaussian blur standard deviation, in meters, from 2D heat diffusion.""" + return np.sqrt(2 * self.thermal_diffusivity * self.dwell_time) + + def blur(self, image: np.ndarray, pixel_scale: float) -> np.ndarray: + """Apply the forward diffusion blur to `image` (for synthetic testing).""" + sigma_px = self.blur_sigma_m() / pixel_scale + return gaussian_filter(image, sigma=sigma_px) + + def deconvolve( + self, image: np.ndarray, pixel_scale: float, noise_to_signal_ratio: float = 1e-3 + ) -> np.ndarray: + """Remove the diffusion blur via regularized (Wiener) deconvolution.""" + sigma_px = self.blur_sigma_m() / pixel_scale + psf = self._gaussian_psf(image.shape, sigma_px) + + otf = np.fft.fft2(np.fft.ifftshift(psf)) + image_ft = np.fft.fft2(image) + wiener_filter = np.conj(otf) / (np.abs(otf) ** 2 + noise_to_signal_ratio) + deconvolved = np.fft.ifft2(image_ft * wiener_filter).real + return deconvolved + + @staticmethod + def _gaussian_psf(shape: tuple[int, int], sigma_px: float) -> np.ndarray: + rows, cols = 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) + psf = np.exp(-(row_grid**2 + col_grid**2) / (2 * sigma_px**2)) + return psf / psf.sum() diff --git a/he11lib/fitting.py b/he11lib/fitting.py new file mode 100644 index 0000000..4c3184f --- /dev/null +++ b/he11lib/fitting.py @@ -0,0 +1,233 @@ +"""Joint nonlinear least-squares modal fit with automatic mode-set growth.""" + +from __future__ import annotations + +import warnings + +import numpy as np +from scipy.optimize import least_squares + +from .data import MeasurementPlane, ReconstructionResult, validate_planes +from .geometry import GeometryCalibration +from .modes import LGBasis +from .noise import NoiseEstimator + + +def generate_mode_shells(max_order: int) -> list[list[tuple[int, int]]]: + """Group candidate LG_{p,l} modes into shells of increasing order 2p+|l|.""" + shells: list[list[tuple[int, int]]] = [[] for _ in range(max_order + 1)] + for p in range(0, max_order + 1): + for l in range(-max_order, max_order + 1): + order = 2 * p + abs(l) + if order <= max_order: + shells[order].append((p, l)) + return shells + + +class ModalFitter: + """Fits LG mode coefficients, beam center/pointing, and geometry to measured planes.""" + + def __init__(self, basis: LGBasis, noise_estimator: NoiseEstimator | None = None): + self.basis = basis + self.noise_estimator = noise_estimator or NoiseEstimator() + + def fit( + self, + planes: list[MeasurementPlane], + modes: list[tuple[int, int]], + initial_coefficients: dict[tuple[int, int], complex] | None = None, + initial_center: tuple[float, float] = (0.0, 0.0), + initial_tilt_deg: tuple[float, float] = (0.0, 0.0), + initial_pixel_scale: float | None = None, + initial_viewing_angle_deg: float = 0.0, + ) -> ReconstructionResult: + """Jointly fit complex coefficients for `modes` plus center/tilt/geometry.""" + validate_planes(planes) + + unknown_scale_idx = [i for i, p in enumerate(planes) if p.pixel_scale is None] + unknown_angle_idx = [i for i, p in enumerate(planes) if p.viewing_angle_deg is None] + weights = [np.sqrt(self.noise_estimator.weights(p.flux)) for p in planes] + + def pack_initial() -> np.ndarray: + x: list[float] = [] + for i, mode in enumerate(modes): + c = (initial_coefficients or {}).get(mode) + if c is None: + # Nonzero seed for every mode: starting a coefficient at + # exactly 0+0j sits at a flat/degenerate point for the + # optimizer and can prevent it from ever leaving zero. + c = 1.0 + 0j if i == 0 else 0.1 + 0.05j + x += [c.real, c.imag] + x += [initial_center[0], initial_center[1], initial_tilt_deg[0], initial_tilt_deg[1]] + for _ in unknown_scale_idx: + x.append(initial_pixel_scale if initial_pixel_scale is not None else 1e-4) + for _ in unknown_angle_idx: + x.append(initial_viewing_angle_deg) + return np.array(x, dtype=float) + + n_modes = len(modes) + + def unpack(x: np.ndarray): + coeffs = {mode: complex(x[2 * i], x[2 * i + 1]) for i, mode in enumerate(modes)} + offset = 2 * n_modes + x0, y0, tilt_x_deg, tilt_y_deg = x[offset : 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): + drift_x = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[0])) + drift_y = (z - self.basis.z0) * np.tan(np.deg2rad(tilt_deg[1])) + return x0 + drift_x, y0 + drift_y + + def model_flux_for_plane(i: int, plane: MeasurementPlane, coeffs, center0, tilt_deg, scales, angles): + scale = plane.pixel_scale if plane.pixel_scale is not None else scales[i] + angle = plane.viewing_angle_deg if plane.viewing_angle_deg is not None else angles[i] + calib = GeometryCalibration(plane) + x_grid, y_grid = calib.physical_coordinates(pixel_scale=scale, viewing_angle_deg=angle) + 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 + + def residuals(x: np.ndarray) -> np.ndarray: + coeffs, center0, tilt_deg, scales, angles = unpack(x) + parts = [] + for i, plane in enumerate(planes): + model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) + parts.append(((plane.flux - model_flux) * weights[i]).ravel()) + return np.concatenate(parts) + + x0_vec = pack_initial() + # 'trf' + x_scale='jac' handles the very different natural magnitudes + # of these parameters (coefficients ~O(1), pixel_scale ~O(1e-3), + # angles ~O(1-90)); plain 'lm' can terminate prematurely on 'xtol' + # because its unscaled step-size test is dominated by the largest + # parameters. + opt_result = least_squares( + residuals, x0_vec, method="trf", x_scale="jac", max_nfev=5000 + ) + + coeffs, center0, tilt_deg, scales, angles = unpack(opt_result.x) + + total_power = sum(abs(c) ** 2 for c in coeffs.values()) + if total_power == 0: + total_power = 1.0 + 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] + pointing_angle_deg = float(np.hypot(tilt_deg[0], tilt_deg[1])) + + geometry: dict[str, float] = {} + for i in range(len(planes)): + geometry[f"pixel_scale_{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 = [] + for i, plane in enumerate(planes): + model_flux = model_flux_for_plane(i, plane, coeffs, center0, tilt_deg, scales, angles) + residual_maps.append(plane.flux - model_flux) + + 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 + field_at_reference = self._field_on_default_grid(coeffs, reference_z) + + return ReconstructionResult( + purity=purity, + reconstructed_field=field_at_reference, + centers=centers, + pointing_angle_deg=pointing_angle_deg, + geometry=geometry, + residuals=residual_maps, + coefficient_uncertainty=coefficient_uncertainty, + used_phase_retrieval=False, + ) + + def fit_auto( + self, + planes: list[MeasurementPlane], + max_order: int = 4, + bic_improvement_threshold: float = 10.0, + ) -> ReconstructionResult: + """Fit with automatic mode-set growth, capped at `max_order`.""" + validate_planes(planes) + shells = generate_mode_shells(max_order) + + current_modes = list(shells[0]) + best_result = self.fit(planes, current_modes) + best_bic = self._bic(planes, best_result, current_modes) + + grew_until_cap = True + for shell in shells[1:]: + trial_modes = current_modes + shell + warm_start = self._warm_start_coefficients(best_result, current_modes) + trial_result = self.fit(planes, trial_modes, initial_coefficients=warm_start) + trial_bic = self._bic(planes, trial_result, trial_modes) + + if trial_bic < best_bic - bic_improvement_threshold: + current_modes = trial_modes + best_result = trial_result + best_bic = trial_bic + else: + grew_until_cap = False + break + + if grew_until_cap and len(shells) > 1: + warnings.warn( + "Automatic mode-set growth hit the configured max_order cap " + f"({max_order}) while still improving the fit; consider raising max_order.", + stacklevel=2, + ) + + return best_result + + def _warm_start_coefficients( + self, previous_result: ReconstructionResult, previous_modes: list[tuple[int, int]] + ) -> dict[tuple[int, int], complex]: + """Reconstruct approximate complex coefficients from a previous fit's purity.""" + coeffs = {} + for mode in previous_modes: + fraction, phase = previous_result.purity[mode] + amplitude = np.sqrt(max(fraction, 0.0)) + coeffs[mode] = amplitude * np.exp(1j * phase) + return coeffs + + def _bic(self, planes: list[MeasurementPlane], result: ReconstructionResult, modes: list[tuple[int, int]]) -> 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_params = 2 * len(modes) + 4 + return float(chi2 + n_params * np.log(n_data)) + + def _estimate_uncertainty(self, opt_result, modes, coeffs, total_power): + try: + jac = opt_result.jac + cov = np.linalg.pinv(jac.T @ jac) + except np.linalg.LinAlgError: + return {mode: float("nan") for mode in modes} + + uncertainty = {} + for i, mode in enumerate(modes): + var_re = cov[2 * i, 2 * i] + var_im = cov[2 * i + 1, 2 * i + 1] + c = coeffs[mode] + sigma_c = np.sqrt(max(var_re, 0) + max(var_im, 0)) + uncertainty[mode] = float(2 * abs(c) * sigma_c / total_power) + return uncertainty + + def _field_on_default_grid(self, coeffs, z: float, n: int = 128, half_width_in_w: float = 6.0): + w_z = self.basis.beam_radius(z) + extent = half_width_in_w * w_z + coords = np.linspace(-extent, extent, n) + x, y = np.meshgrid(coords, coords) + return self.basis.field_superposition(x, y, z, coeffs) diff --git a/he11lib/geometry.py b/he11lib/geometry.py new file mode 100644 index 0000000..0b423e5 --- /dev/null +++ b/he11lib/geometry.py @@ -0,0 +1,64 @@ +"""Camera geometry correction: pixel-to-physical scale and viewing angle. + +Converts a MeasurementPlane's pixel grid into physical (x, y) coordinates in +the beam's transverse plane, compensating for an oblique camera viewing +angle (which compresses the image along the tilt axis by cos(angle)). +Known calibration values (on the MeasurementPlane) are used directly; when a +value is unknown, an override must be supplied (e.g. by ModalFitter while +exploring it as a free parameter). +""" + +from __future__ import annotations + +import numpy as np + +from .data import MeasurementPlane + + +class GeometryCalibration: + """Resolves pixel scale / viewing angle and builds a physical coordinate grid.""" + + def __init__(self, plane: MeasurementPlane): + self.plane = plane + + @property + def pixel_scale_known(self) -> bool: + return self.plane.pixel_scale is not None + + @property + def viewing_angle_known(self) -> bool: + return self.plane.viewing_angle_deg is not None + + def physical_coordinates( + self, + pixel_scale: float | None = None, + viewing_angle_deg: float | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Physical (x, y) grid matching the plane's flux array shape. + + Known values on the MeasurementPlane take precedence; overrides are + only used to fill in values that are not known/calibrated. + """ + scale = self.plane.pixel_scale if self.pixel_scale_known else pixel_scale + 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 + col_idx = np.arange(cols) - cols // 2 + col_grid, row_grid = np.meshgrid(col_idx, row_idx) + + cos_angle = np.cos(np.deg2rad(angle_deg)) + x = col_grid * scale / cos_angle + y = row_grid * scale + return x, y diff --git a/he11lib/modes.py b/he11lib/modes.py new file mode 100644 index 0000000..112adfc --- /dev/null +++ b/he11lib/modes.py @@ -0,0 +1,96 @@ +"""Laguerre-Gauss mode basis for free-space gyrotron beam propagation.""" + +from __future__ import annotations + +import math + +import numpy as np +from scipy.special import eval_genlaguerre + + +class LGBasis: + """Laguerre-Gauss mode basis referenced to a known waist size/location. + + Parameters + ---------- + w0 : reference waist radius, in meters. + z0 : reference waist location, in meters. + wavelength : radiation wavelength, in meters (e.g. ~1.76 mm for a + 170 GHz gyrotron). + """ + + def __init__(self, w0: float, z0: float, wavelength: float): + self.w0 = w0 + self.z0 = z0 + self.wavelength = wavelength + self.k = 2 * np.pi / wavelength + self.zR = np.pi * w0**2 / wavelength + + def beam_radius(self, z: float) -> float: + """Beam radius w(z).""" + return self.w0 * math.sqrt(1.0 + ((z - self.z0) / self.zR) ** 2) + + def inverse_radius_of_curvature(self, z: float) -> float: + """1/R(z), well-defined (=0) at the waist.""" + dz = z - self.z0 + return dz / (dz**2 + self.zR**2) + + def gouy_phase(self, z: float, p: int, l: int) -> float: + """Gouy phase for mode (p, l) at distance z.""" + order = 2 * p + abs(l) + 1 + return order * math.atan((z - self.z0) / self.zR) + + def field(self, x: np.ndarray, y: np.ndarray, z: float, p: int, l: int) -> np.ndarray: + """Complex LG_{p,l} field sampled on the given (x, y) grid at distance z.""" + w_z = self.beam_radius(z) + r2 = x**2 + y**2 + r = np.sqrt(r2) + phi = np.arctan2(y, x) + abs_l = abs(l) + + c_pl = math.sqrt(2 * math.factorial(p) / (math.pi * math.factorial(p + abs_l))) + radial = ( + c_pl + / w_z + * (r * math.sqrt(2) / w_z) ** abs_l + * eval_genlaguerre(p, abs_l, 2 * r2 / w_z**2) + * np.exp(-r2 / w_z**2) + ) + curvature_phase = self.k * r2 * self.inverse_radius_of_curvature(z) / 2 + gouy = self.gouy_phase(z, p, l) + phase = np.exp(1j * (curvature_phase + gouy - l * phi)) + return radial * phase + + def field_superposition( + self, + x: np.ndarray, + y: np.ndarray, + z: float, + coefficients: dict[tuple[int, int], complex], + ) -> np.ndarray: + """Complex field for a superposition of modes with given complex coefficients.""" + total = np.zeros_like(x, dtype=complex) + for (p, l), coeff in coefficients.items(): + total = total + coeff * self.field(x, y, z, p, l) + return total + + def project( + self, + complex_field: np.ndarray, + x: np.ndarray, + y: np.ndarray, + dx: float, + z: float, + modes: list[tuple[int, int]], + ) -> dict[tuple[int, int], complex]: + """Project a complex field onto the given candidate modes at distance z. + + Returns the complex coefficient of each mode via the inner product + , using a Riemann-sum approximation of the integral + over the (x, y) grid (spacing dx, assumed square/uniform). + """ + coefficients: dict[tuple[int, int], complex] = {} + for p, l in modes: + mode_field = self.field(x, y, z, p, l) + coefficients[(p, l)] = np.sum(np.conj(mode_field) * complex_field) * dx**2 + return coefficients diff --git a/he11lib/noise.py b/he11lib/noise.py new file mode 100644 index 0000000..ea80183 --- /dev/null +++ b/he11lib/noise.py @@ -0,0 +1,40 @@ +"""Automatic residual sensor noise estimation. + +Uses the Immerkaer (1996) fast noise-variance estimator: convolving the +image with a Laplacian-of-Gaussian-like kernel suppresses smooth signal +content (the beam profile) while passing high-frequency noise, giving a +closed-form estimate of the noise standard deviation without needing a +dedicated background region. +""" + +from __future__ import annotations + +import numpy as np +from scipy.signal import convolve2d + +_LAPLACIAN_KERNEL = np.array( + [ + [1, -2, 1], + [-2, 4, -2], + [1, -2, 1], + ], + dtype=float, +) + + +class NoiseEstimator: + """Estimates residual per-image noise and produces least-squares weights.""" + + def estimate_std(self, image: np.ndarray) -> float: + """Estimate the additive Gaussian noise standard deviation in `image`.""" + rows, cols = image.shape + convolved = convolve2d(image, _LAPLACIAN_KERNEL, mode="valid") + sigma = np.sqrt(np.pi / 2) * np.sum(np.abs(convolved)) / ( + 6 * (cols - 2) * (rows - 2) + ) + return float(sigma) + + def weights(self, image: np.ndarray) -> np.ndarray: + """Per-pixel weights (1/sigma^2) for noise-weighted least squares.""" + sigma = self.estimate_std(image) + return np.full(image.shape, 1.0 / sigma**2) diff --git a/he11lib/phase_retrieval.py b/he11lib/phase_retrieval.py new file mode 100644 index 0000000..daa9c87 --- /dev/null +++ b/he11lib/phase_retrieval.py @@ -0,0 +1,135 @@ +"""Gerchberg-Saxton-style multi-plane phase retrieval (fallback reconstruction path). + +Used when the finite-mode-basis fit (ModalFitter) leaves a high residual, or +when explicitly requested. Iteratively propagates a trial complex field +between measurement planes via free-space (angular-spectrum) propagation, +enforcing the measured amplitude at each plane without assuming a finite +mode basis. The recovered field can then be projected onto an LGBasis for a +purity table. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np + +from .data import MeasurementPlane, validate_planes +from .geometry import GeometryCalibration + + +def propagate_angular_spectrum( + field: np.ndarray, dx: float, dz: float, wavelength: float +) -> np.ndarray: + """Free-space-propagate a complex field by distance dz via the angular spectrum method. + + `dx` is the (square) pixel spacing of `field`. Evanescent components + (where the transverse spatial frequency exceeds 1/wavelength) are + dropped, which is an excellent approximation for paraxial beams. + """ + ny, nx = field.shape + fx = np.fft.fftfreq(nx, d=dx) + fy = np.fft.fftfreq(ny, d=dx) + fx_grid, fy_grid = np.meshgrid(fx, fy) + + k = 2 * np.pi / wavelength + kz_squared = k**2 - (2 * np.pi * fx_grid) ** 2 - (2 * np.pi * fy_grid) ** 2 + kz = np.sqrt(np.maximum(kz_squared, 0.0)) + propagating = kz_squared >= 0 + + transfer_function = np.where(propagating, np.exp(1j * kz * dz), 0.0) + + field_ft = np.fft.fft2(field) + return np.fft.ifft2(field_ft * transfer_function) + + +@dataclass +class PhaseRetrievalResult: + """Result of multi-plane phase retrieval. + + field : recovered complex field at distance `z` (the smallest-z plane), + on the (x, y) physical grid. + x, y : physical coordinate grids matching `field`'s shape. + z : distance at which `field` is defined. + center : estimated beam transverse center (intensity centroid), in + meters, on the same grid. + residual : final RMS mismatch between enforced and propagated amplitude + across all planes (diagnostic of convergence quality). + """ + + field: np.ndarray + x: np.ndarray + y: np.ndarray + z: float + center: tuple[float, float] + residual: float + + +class PhaseRetriever: + """Multi-plane Gerchberg-Saxton phase retrieval. + + Parameters + ---------- + wavelength : radiation wavelength, in meters. + """ + + def __init__(self, wavelength: float): + self.wavelength = wavelength + + def retrieve( + self, + planes: list[MeasurementPlane], + pixel_scale: float | None = None, + viewing_angle_deg: float | None = None, + max_iterations: int = 200, + ) -> PhaseRetrievalResult: + """Run Gerchberg-Saxton phase retrieval across the given planes. + + Planes must share the same known (or overridden) pixel_scale and + viewing_angle_deg, since all planes are propagated on one common + physical grid. + """ + validate_planes(planes) + ordered = sorted(planes, key=lambda p: p.z) + + x, y = GeometryCalibration(ordered[0]).physical_coordinates( + pixel_scale=pixel_scale, viewing_angle_deg=viewing_angle_deg + ) + dx = float(x[0, 1] - x[0, 0]) + + amplitudes = [np.sqrt(np.clip(p.flux, 0, None)) for p in ordered] + z_values = [p.z for p in ordered] + + field = amplitudes[0].astype(complex).copy() + residual = float("inf") + + for _ in range(max_iterations): + residual = 0.0 + # forward sweep + for i in range(len(ordered) - 1): + dz = z_values[i + 1] - z_values[i] + field = propagate_angular_spectrum(field, dx, dz, self.wavelength) + mismatch = np.abs(field) - amplitudes[i + 1] + residual += float(np.sum(mismatch**2)) + phase = np.angle(field) + field = amplitudes[i + 1] * np.exp(1j * phase) + # backward sweep + for i in range(len(ordered) - 1, 0, -1): + dz = z_values[i - 1] - z_values[i] + field = propagate_angular_spectrum(field, dx, dz, self.wavelength) + mismatch = np.abs(field) - amplitudes[i - 1] + residual += float(np.sum(mismatch**2)) + phase = np.angle(field) + field = amplitudes[i - 1] * np.exp(1j * phase) + + n_pixels = sum(a.size for a in amplitudes) + rms_residual = float(np.sqrt(residual / n_pixels)) if n_pixels else 0.0 + + intensity = np.abs(field) ** 2 + total_intensity = np.sum(intensity) + cx = float(np.sum(x * intensity) / total_intensity) + cy = float(np.sum(y * intensity) / total_intensity) + + return PhaseRetrievalResult( + field=field, x=x, y=y, z=z_values[0], center=(cx, cy), residual=rms_residual + ) diff --git a/he11lib/plotting.py b/he11lib/plotting.py new file mode 100644 index 0000000..ccb3138 --- /dev/null +++ b/he11lib/plotting.py @@ -0,0 +1,65 @@ +"""Diagnostic visualizations for reconstruction results. + +Each function returns a `matplotlib.figure.Figure` for the caller to display +(`fig.show()` / inline in a notebook) or save (`fig.savefig(...)`). +""" + +from __future__ import annotations + +import matplotlib.figure +import matplotlib.pyplot as plt + +from .data import MeasurementPlane, ReconstructionResult + + +def plot_mode_purity(result: ReconstructionResult) -> matplotlib.figure.Figure: + """Bar chart of power fraction per candidate LG mode.""" + modes = list(result.purity.keys()) + fractions = [result.purity[mode][0] for mode in modes] + labels = [f"LG_{p},{l}" for p, l in modes] + + fig, ax = plt.subplots() + ax.bar(labels, fractions) + ax.set_ylabel("Power fraction") + ax.set_xlabel("Mode") + ax.set_title("Mode purity") + return fig + + +def plot_center_trace( + planes: list[MeasurementPlane], result: ReconstructionResult +) -> matplotlib.figure.Figure: + """Fitted beam transverse center (x, y) as a function of z across planes.""" + z_values = [p.z for p in planes] + x_values = [c[0] for c in result.centers] + y_values = [c[1] for c in result.centers] + + fig, (ax_x, ax_y) = plt.subplots(1, 2, sharex=True) + ax_x.plot(z_values, x_values, marker="o") + ax_x.set_xlabel("z (m)") + ax_x.set_ylabel("center x (m)") + ax_y.plot(z_values, y_values, marker="o") + ax_y.set_xlabel("z (m)") + ax_y.set_ylabel("center y (m)") + fig.suptitle(f"Beam center (pointing angle {result.pointing_angle_deg:.3g} deg)") + return fig + + +def plot_residuals( + planes: list[MeasurementPlane], result: ReconstructionResult +) -> matplotlib.figure.Figure: + """Per-plane residual maps (measured - modeled flux).""" + if not result.residuals: + raise ValueError( + "No per-plane residuals are available on this ReconstructionResult " + "(e.g. the phase-retrieval fallback was used, which does not produce them)." + ) + + n = len(planes) + fig, axes = plt.subplots(1, n, squeeze=False) + axes = axes[0] + for ax, plane, residual in zip(axes, planes, result.residuals): + im = ax.imshow(residual) + ax.set_title(f"z={plane.z:g} m") + fig.colorbar(im, ax=ax) + return fig diff --git a/he11lib/reconstruct.py b/he11lib/reconstruct.py new file mode 100644 index 0000000..0d22b4e --- /dev/null +++ b/he11lib/reconstruct.py @@ -0,0 +1,124 @@ +"""High-level orchestrator wiring together the full reconstruction pipeline.""" + +from __future__ import annotations + +from dataclasses import replace + +import numpy as np + +from .data import MeasurementPlane, ReconstructionResult, validate_planes +from .deconvolution import DiffusionDeconvolver +from .fitting import ModalFitter, generate_mode_shells +from .modes import LGBasis +from .noise import NoiseEstimator +from .phase_retrieval import PhaseRetriever + + +class BeamReconstructor: + """Runs the complete mode-purity reconstruction pipeline. + + Given a list of `MeasurementPlane`s, this orchestrates (optional) + diffusion deblurring, the joint modal least-squares fit with automatic + mode-set growth, and an optional Gerchberg-Saxton phase-retrieval + fallback -- producing a single `ReconstructionResult`. + + Parameters + ---------- + w0, z0, wavelength : known reference beam parameters (see `LGBasis`). + max_order : cap on automatic candidate-mode-set growth (see + `ModalFitter.fit_auto`), and also the mode set used to project the + phase-retrieval fallback's recovered field onto the LG basis. + noise_estimator : shared noise model; defaults to `NoiseEstimator()`. + deconvolver : if given, each plane's flux is deblurred (its + `pixel_scale` must be known) before fitting. + force_phase_retrieval : if True, always run the phase-retrieval fallback + instead of the modal fit. + phase_retrieval_residual_threshold : if set (and `force_phase_retrieval` + is False), the phase-retrieval fallback runs automatically whenever + the modal fit's noise-weighted RMS residual exceeds this value. + """ + + def __init__( + self, + w0: float, + z0: float, + wavelength: float, + max_order: int = 4, + noise_estimator: NoiseEstimator | None = None, + deconvolver: DiffusionDeconvolver | None = None, + force_phase_retrieval: bool = False, + phase_retrieval_residual_threshold: float | None = None, + ): + self.basis = LGBasis(w0=w0, z0=z0, wavelength=wavelength) + self.wavelength = wavelength + self.max_order = max_order + self.noise_estimator = noise_estimator or NoiseEstimator() + self.deconvolver = deconvolver + self.force_phase_retrieval = force_phase_retrieval + self.phase_retrieval_residual_threshold = phase_retrieval_residual_threshold + + def reconstruct(self, planes: list[MeasurementPlane]) -> ReconstructionResult: + """Run the full pipeline and return a `ReconstructionResult`.""" + validate_planes(planes) + planes = self._deconvolve(planes) + + fitter = ModalFitter(self.basis, self.noise_estimator) + result = fitter.fit_auto(planes, max_order=self.max_order) + + if self.force_phase_retrieval or self._residual_too_high(result, planes): + result = self._phase_retrieval_fallback(planes) + + return result + + def _deconvolve(self, planes: list[MeasurementPlane]) -> list[MeasurementPlane]: + if self.deconvolver is None: + return planes + deblurred = [] + for plane in planes: + if plane.pixel_scale is None: + raise ValueError( + "Deconvolution requires a known pixel_scale on every MeasurementPlane." + ) + flux = self.deconvolver.deconvolve(plane.flux, plane.pixel_scale) + deblurred.append(replace(plane, flux=flux)) + return deblurred + + def _residual_too_high( + self, result: ReconstructionResult, planes: list[MeasurementPlane] + ) -> bool: + if self.phase_retrieval_residual_threshold is None: + return False + total = 0.0 + n_pixels = 0 + for residual_map, plane in zip(result.residuals, planes): + weights = self.noise_estimator.weights(plane.flux) + total += float(np.sum((residual_map**2) * weights)) + n_pixels += residual_map.size + rms = np.sqrt(total / n_pixels) if n_pixels else 0.0 + return rms > self.phase_retrieval_residual_threshold + + def _phase_retrieval_fallback( + self, planes: list[MeasurementPlane] + ) -> ReconstructionResult: + retriever = PhaseRetriever(self.wavelength) + pr_result = retriever.retrieve(planes) + + 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]) + coeffs = self.basis.project(pr_result.field, pr_result.x, pr_result.y, dx, pr_result.z, modes) + + total_power = sum(abs(c) ** 2 for c in coeffs.values()) + if total_power == 0: + total_power = 1.0 + purity = {mode: (abs(c) ** 2 / total_power, float(np.angle(c))) for mode, c in coeffs.items()} + + return ReconstructionResult( + purity=purity, + reconstructed_field=pr_result.field, + centers=[pr_result.center for _ in planes], + pointing_angle_deg=float("nan"), + geometry={}, + residuals=[], + coefficient_uncertainty={mode: float("nan") for mode in modes}, + used_phase_retrieval=True, + ) diff --git a/he11lib/synthetic.py b/he11lib/synthetic.py new file mode 100644 index 0000000..215b34d --- /dev/null +++ b/he11lib/synthetic.py @@ -0,0 +1,83 @@ +"""Forward model: synthetic thermal (flux) images from known ground truth. + +Used to validate the reconstruction pipeline (recover known mode content) +and to help users evaluate experimental design (e.g. whether a given set of +measurement distances will separate candidate modes). +""" + +from __future__ import annotations + +import numpy as np + +from .data import MeasurementPlane +from .modes import LGBasis + + +class SyntheticBeamGenerator: + """Generates synthetic multi-plane flux images for a known ground-truth beam. + + Parameters + ---------- + basis : LGBasis defining the reference w0, z0, wavelength. + image_shape : (rows, cols) pixel shape of generated images. + pixel_scale : physical size of one pixel, in meters, along the + 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): + self.basis = basis + self.image_shape = image_shape + self.pixel_scale = pixel_scale + + 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( + self, + coefficients: dict[tuple[int, int], complex], + z_list: list[float], + *, + center: tuple[float, float] = (0.0, 0.0), + pointing_angle_deg: float = 0.0, + viewing_angle_deg: float = 0.0, + noise_std: float = 0.0, + seed: int | None = None, + ) -> list[MeasurementPlane]: + """Generate one MeasurementPlane per requested z distance. + + The beam transverse center drifts linearly with z according to + pointing_angle_deg (tilt of the beam axis along x), starting from + `center` at the basis's reference z0. + """ + rng = np.random.default_rng(seed) + tilt_rad = np.deg2rad(pointing_angle_deg) + + planes = [] + for z in z_list: + drift_x = (z - self.basis.z0) * np.tan(tilt_rad) + plane_center = (center[0] + drift_x, center[1]) + + x, y = self._pixel_grid(plane_center, viewing_angle_deg) + field = self.basis.field_superposition(x, y, z, coefficients) + flux = np.abs(field) ** 2 + + if noise_std > 0: + flux = flux + rng.normal(0.0, noise_std, size=flux.shape) + + planes.append( + MeasurementPlane( + flux=flux, + z=z, + pixel_scale=self.pixel_scale, + viewing_angle_deg=viewing_angle_deg, + ) + ) + return planes diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..e6f09a6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "he11lib" +version = "0.1.0" +description = "Mode purity reconstruction of free-space gyrotron beams from multi-plane thermal (flux) images." +readme = "README.md" +requires-python = ">=3.10" +license = { text = "GPL-3.0-or-later" } +dependencies = [ + "numpy>=1.24", + "scipy>=1.10", + "matplotlib>=3.7", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.4", +] + +[tool.setuptools.packages.find] +include = ["he11lib*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..290cc21 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,3 @@ +import matplotlib + +matplotlib.use("Agg") diff --git a/tests/test_data.py b/tests/test_data.py new file mode 100644 index 0000000..b4d2a52 --- /dev/null +++ b/tests/test_data.py @@ -0,0 +1,93 @@ +import numpy as np +import pytest + +from he11lib.data import MeasurementPlane, ReconstructionResult, validate_planes + + +def test_measurement_plane_stores_fields(): + flux = np.ones((4, 4)) + plane = MeasurementPlane(flux=flux, z=0.3) + + assert plane.z == 0.3 + assert np.array_equal(plane.flux, flux) + assert plane.pixel_scale is None + assert plane.viewing_angle_deg is None + assert plane.label is None + + +def test_measurement_plane_stores_optional_fields(): + flux = np.ones((4, 4)) + plane = MeasurementPlane( + 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.viewing_angle_deg == 5.0 + assert plane.label == "plane_40cm" + + +def test_measurement_plane_rejects_non_2d_flux(): + with pytest.raises(ValueError, match="2D"): + MeasurementPlane(flux=np.ones((4, 4, 3)), z=0.3) + + +def test_measurement_plane_rejects_non_positive_z(): + with pytest.raises(ValueError, match="positive"): + MeasurementPlane(flux=np.ones((4, 4)), z=0.0) + + with pytest.raises(ValueError, match="positive"): + MeasurementPlane(flux=np.ones((4, 4)), z=-0.1) + + +def test_validate_planes_rejects_fewer_than_three(): + planes = [ + MeasurementPlane(flux=np.ones((4, 4)), z=0.3), + MeasurementPlane(flux=np.ones((4, 4)), z=0.4), + ] + with pytest.raises(ValueError, match="[Aa]t least 3"): + validate_planes(planes) + + +def test_validate_planes_rejects_mismatched_shapes(): + planes = [ + MeasurementPlane(flux=np.ones((4, 4)), z=0.3), + MeasurementPlane(flux=np.ones((5, 5)), z=0.4), + MeasurementPlane(flux=np.ones((4, 4)), z=0.5), + ] + with pytest.raises(ValueError, match="same shape"): + validate_planes(planes) + + +def test_validate_planes_rejects_duplicate_z(): + planes = [ + MeasurementPlane(flux=np.ones((4, 4)), z=0.3), + MeasurementPlane(flux=np.ones((4, 4)), z=0.3), + MeasurementPlane(flux=np.ones((4, 4)), z=0.5), + ] + with pytest.raises(ValueError, match="distinct"): + validate_planes(planes) + + +def test_validate_planes_accepts_valid_list(): + planes = [ + MeasurementPlane(flux=np.ones((4, 4)), z=0.3), + MeasurementPlane(flux=np.ones((4, 4)), z=0.4), + MeasurementPlane(flux=np.ones((4, 4)), z=0.5), + ] + validate_planes(planes) # should not raise + + +def test_reconstruction_result_stores_fields(): + result = ReconstructionResult( + purity={(0, 0): (1.0, 0.0)}, + reconstructed_field=np.ones((4, 4), dtype=complex), + centers=[(0.0, 0.0)], + pointing_angle_deg=0.0, + geometry={"pixel_scale": 0.1, "viewing_angle_deg": 2.0}, + residuals=[np.zeros((4, 4))], + coefficient_uncertainty={(0, 0): 0.01}, + used_phase_retrieval=False, + ) + + assert result.purity[(0, 0)] == (1.0, 0.0) + assert result.used_phase_retrieval is False diff --git a/tests/test_deconvolution.py b/tests/test_deconvolution.py new file mode 100644 index 0000000..d9a654e --- /dev/null +++ b/tests/test_deconvolution.py @@ -0,0 +1,67 @@ +import numpy as np +import pytest + +from he11lib.deconvolution import DiffusionDeconvolver + + +def gaussian_bump(n, sigma_px): + coords = np.arange(n) - n // 2 + xx, yy = np.meshgrid(coords, coords) + return np.exp(-(xx**2 + yy**2) / (2 * sigma_px**2)) + + +def profile_std(image): + n = image.shape[0] + coords = np.arange(n) - n // 2 + weights = image[n // 2, :] + weights = np.clip(weights, 0, None) + mean = np.sum(coords * weights) / np.sum(weights) + var = np.sum(weights * (coords - mean) ** 2) / np.sum(weights) + return np.sqrt(var) + + +def test_blur_sigma_m_from_diffusivity_and_dwell_time(): + diffusivity = 1e-6 # m^2/s + dwell_time = 0.5 # s + deconvolver = DiffusionDeconvolver(thermal_diffusivity=diffusivity, dwell_time=dwell_time) + + expected_sigma = np.sqrt(2 * diffusivity * dwell_time) + assert deconvolver.blur_sigma_m() == pytest.approx(expected_sigma) + + +def test_blur_widens_a_sharp_peak(): + deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.5) + pixel_scale = 2e-4 + sharp = gaussian_bump(101, sigma_px=3) + + blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale) + + assert profile_std(blurred) > profile_std(sharp) + + +def test_deconvolve_reduces_error_relative_to_blurred(): + deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.3) + pixel_scale = 2e-4 + sharp = gaussian_bump(101, sigma_px=4) + + blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale) + deconvolved = deconvolver.deconvolve(blurred, pixel_scale=pixel_scale) + + error_blurred = np.sum((blurred - sharp) ** 2) + error_deconvolved = np.sum((deconvolved - sharp) ** 2) + assert error_deconvolved < error_blurred + + +def test_deconvolve_narrows_width_back_toward_original(): + deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=0.3) + pixel_scale = 2e-4 + sharp = gaussian_bump(101, sigma_px=4) + + blurred = deconvolver.blur(sharp, pixel_scale=pixel_scale) + deconvolved = deconvolver.deconvolve(blurred, pixel_scale=pixel_scale) + + std_sharp = profile_std(sharp) + std_blurred = profile_std(blurred) + std_deconvolved = profile_std(deconvolved) + + assert std_sharp < std_deconvolved < std_blurred diff --git a/tests/test_fitting.py b/tests/test_fitting.py new file mode 100644 index 0000000..ff0c170 --- /dev/null +++ b/tests/test_fitting.py @@ -0,0 +1,134 @@ +import numpy as np +import pytest + +from he11lib.data import validate_planes +from he11lib.fitting import ModalFitter, generate_mode_shells +from he11lib.modes import LGBasis +from he11lib.synthetic import SyntheticBeamGenerator + +W0 = 5e-3 +Z0 = 0.5 +WAVELENGTH = 1.76e-3 +PIXEL_SCALE = 4e-4 +IMAGE_SHAPE = (61, 61) +Z_LIST = [0.35, 0.5, 0.65, 0.8] + + +def make_basis(): + return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + + +def make_generator(basis): + return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + + +def test_generate_mode_shells_orders_by_2p_plus_abs_l(): + shells = generate_mode_shells(max_order=2) + assert shells[0] == [(0, 0)] + assert set(shells[1]) == {(0, 1), (0, -1)} + assert set(shells[2]) == {(0, 2), (0, -2), (1, 0)} + + +def test_fit_recovers_pure_fundamental_mode(): + basis = make_basis() + gen = make_generator(basis) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0 + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)]) + + power_fraction, _ = result.purity[(0, 0)] + assert power_fraction == pytest.approx(1.0, abs=1e-6) + for cx, cy in result.centers: + assert cx == pytest.approx(0.0, abs=2 * PIXEL_SCALE) + assert cy == pytest.approx(0.0, abs=2 * PIXEL_SCALE) + + +def test_fit_recovers_two_mode_purity_ratio(): + basis = make_basis() + gen = make_generator(basis) + 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) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=list(true_coeffs.keys())) + + true_total = sum(abs(c) ** 2 for c in true_coeffs.values()) + for mode, c in true_coeffs.items(): + expected_fraction = abs(c) ** 2 / true_total + recovered_fraction, _ = result.purity[mode] + assert recovered_fraction == pytest.approx(expected_fraction, abs=0.03) + + +def test_fit_recovers_center_offset(): + basis = make_basis() + gen = make_generator(basis) + true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, + z_list=Z_LIST, + center=true_center, + noise_std=1e-4, + seed=2, + ) + + fitter = ModalFitter(basis) + result = fitter.fit(planes, modes=[(0, 0)], initial_center=true_center) + + for cx, cy in result.centers: + assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE) + assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE) + + +def test_fit_recovers_unknown_pixel_scale(): + # 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() + local_pixel_scale = 1.5e-3 + gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=local_pixel_scale) + planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=3) + + # hide the known calibration to force the fitter to solve for it + for plane in planes: + plane.pixel_scale = None + plane.viewing_angle_deg = None + + fitter = ModalFitter(basis) + result = fitter.fit( + planes, + 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))] + for scale in fitted_scales: + assert scale == pytest.approx(local_pixel_scale, rel=0.05) + + +def test_fit_auto_does_not_add_modes_for_pure_fundamental(): + basis = make_basis() + gen = make_generator(basis) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=4 + ) + + fitter = ModalFitter(basis) + result = fitter.fit_auto(planes, max_order=2) + + assert set(result.purity.keys()) == {(0, 0)} + + +def test_fit_auto_grows_to_include_second_mode(): + basis = make_basis() + gen = make_generator(basis) + 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) + + fitter = ModalFitter(basis) + result = fitter.fit_auto(planes, max_order=2) + + assert (0, 1) in result.purity or (0, -1) in result.purity diff --git a/tests/test_geometry.py b/tests/test_geometry.py new file mode 100644 index 0000000..6fdd399 --- /dev/null +++ b/tests/test_geometry.py @@ -0,0 +1,77 @@ +import numpy as np +import pytest + +from he11lib.data import MeasurementPlane +from he11lib.geometry import GeometryCalibration + + +def test_pixel_scale_known_reflects_plane(): + plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, pixel_scale=1e-4) + plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) + + assert GeometryCalibration(plane_known).pixel_scale_known is True + assert GeometryCalibration(plane_unknown).pixel_scale_known is False + + +def test_viewing_angle_known_reflects_plane(): + plane_known = MeasurementPlane(flux=np.ones((5, 5)), z=0.3, viewing_angle_deg=10.0) + plane_unknown = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) + + assert GeometryCalibration(plane_known).viewing_angle_known is True + assert GeometryCalibration(plane_unknown).viewing_angle_known is False + + +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(): + plane = MeasurementPlane( + flux=np.ones((5, 5)), z=0.3, pixel_scale=2e-4, viewing_angle_deg=60.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(): + plane = MeasurementPlane(flux=np.ones((5, 5)), z=0.3) + calib = GeometryCalibration(plane) + + with pytest.raises(ValueError, match="pixel_scale"): + 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) + col_idx = np.arange(5) - 2 + np.testing.assert_allclose(x[2, :], col_idx * 2e-4) diff --git a/tests/test_modes.py b/tests/test_modes.py new file mode 100644 index 0000000..305bce9 --- /dev/null +++ b/tests/test_modes.py @@ -0,0 +1,103 @@ +import numpy as np +import pytest + +from he11lib.modes import LGBasis + + +W0 = 5e-3 # 5 mm waist +Z0 = 0.5 # waist at 0.5 m +WAVELENGTH = 1.76e-3 # ~170 GHz gyrotron + + +def make_grid(w, half_widths_in_w=6.0, n=300): + extent = half_widths_in_w * w + coords = np.linspace(-extent, extent, n) + dx = coords[1] - coords[0] + x, y = np.meshgrid(coords, coords) + return x, y, dx + + +def test_beam_radius_at_waist_equals_w0(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + assert basis.beam_radius(Z0) == pytest.approx(W0) + + +def test_beam_radius_at_one_rayleigh_range(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + zR = np.pi * W0**2 / WAVELENGTH + assert basis.beam_radius(Z0 + zR) == pytest.approx(W0 * np.sqrt(2)) + + +def test_gouy_phase_zero_at_waist(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + assert basis.gouy_phase(Z0, p=0, l=0) == pytest.approx(0.0) + assert basis.gouy_phase(Z0, p=1, l=2) == pytest.approx(0.0) + + +def test_gouy_phase_at_one_rayleigh_range_for_fundamental(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + zR = np.pi * W0**2 / WAVELENGTH + # order (2p+|l|+1) = 1 for p=0,l=0; atan(1) = pi/4 + assert basis.gouy_phase(Z0 + zR, p=0, l=0) == pytest.approx(np.pi / 4) + + +def test_fundamental_mode_matches_analytic_gaussian_at_waist(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + x, y, _ = make_grid(W0, n=50) + r2 = x**2 + y**2 + + field = basis.field(x, y, Z0, p=0, l=0) + expected_intensity_shape = np.exp(-2 * r2 / W0**2) + intensity = np.abs(field) ** 2 + + # shapes should match up to a constant normalization factor + ratio = intensity / expected_intensity_shape + ratio_center = ratio[len(ratio) // 2, len(ratio) // 2] + np.testing.assert_allclose(ratio / ratio_center, 1.0, atol=1e-6) + + +def test_mode_is_normalized_at_waist(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + x, y, dx = make_grid(W0, n=400) + + field = basis.field(x, y, Z0, p=0, l=0) + total_power = np.sum(np.abs(field) ** 2) * dx**2 + assert total_power == pytest.approx(1.0, rel=2e-3) + + +def test_mode_is_normalized_away_from_waist(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + z = Z0 + 0.2 + w_z = basis.beam_radius(z) + x, y, dx = make_grid(w_z, n=400) + + field = basis.field(x, y, z, p=1, l=2) + total_power = np.sum(np.abs(field) ** 2) * dx**2 + assert total_power == pytest.approx(1.0, rel=2e-3) + + +def test_modes_are_orthogonal(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + z = Z0 + 0.1 + w_z = basis.beam_radius(z) + x, y, dx = make_grid(w_z, n=400) + + field_a = basis.field(x, y, z, p=0, l=0) + field_b = basis.field(x, y, z, p=1, l=0) + inner_product = np.sum(field_a * np.conj(field_b)) * dx**2 + assert abs(inner_product) < 1e-3 + + +def test_project_recovers_known_coefficients(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + z = Z0 + 0.15 + w_z = basis.beam_radius(z) + x, y, dx = make_grid(w_z, n=400) + + true_coeffs = {(0, 0): 0.8 + 0.1j, (1, 0): 0.2 - 0.3j} + field = basis.field_superposition(x, y, z, true_coeffs) + + recovered = basis.project(field, x, y, dx, z, modes=list(true_coeffs.keys())) + + for mode, coeff in true_coeffs.items(): + assert recovered[mode] == pytest.approx(coeff, abs=5e-3) diff --git a/tests/test_noise.py b/tests/test_noise.py new file mode 100644 index 0000000..9ba2f1c --- /dev/null +++ b/tests/test_noise.py @@ -0,0 +1,46 @@ +import numpy as np +import pytest + +from he11lib.noise import NoiseEstimator + + +def test_estimate_std_recovers_known_noise_on_flat_image(): + rng = np.random.default_rng(0) + true_std = 0.05 + image = np.ones((200, 200)) * 10.0 + rng.normal(0, true_std, size=(200, 200)) + + estimated = NoiseEstimator().estimate_std(image) + assert estimated == pytest.approx(true_std, rel=0.15) + + +def test_estimate_std_recovers_known_noise_on_smooth_bump(): + rng = np.random.default_rng(1) + x = np.linspace(-3, 3, 200) + xx, yy = np.meshgrid(x, x) + smooth = np.exp(-(xx**2 + yy**2)) + true_std = 0.01 + image = smooth + rng.normal(0, true_std, size=smooth.shape) + + estimated = NoiseEstimator().estimate_std(image) + assert estimated == pytest.approx(true_std, rel=0.35) + + +def test_estimate_std_near_zero_for_noise_free_image(): + x = np.linspace(-3, 3, 100) + xx, yy = np.meshgrid(x, x) + smooth = np.exp(-(xx**2 + yy**2)) + + estimated = NoiseEstimator().estimate_std(smooth) + assert estimated < 1e-6 + + +def test_weights_are_uniform_and_match_shape(): + rng = np.random.default_rng(2) + image = np.ones((50, 50)) * 5.0 + rng.normal(0, 0.1, size=(50, 50)) + + weights = NoiseEstimator().weights(image) + assert weights.shape == image.shape + assert np.allclose(weights, weights.flat[0]) + + expected_std = NoiseEstimator().estimate_std(image) + assert weights.flat[0] == pytest.approx(1.0 / expected_std**2, rel=1e-6) diff --git a/tests/test_phase_retrieval.py b/tests/test_phase_retrieval.py new file mode 100644 index 0000000..ad66a5d --- /dev/null +++ b/tests/test_phase_retrieval.py @@ -0,0 +1,84 @@ +import numpy as np +import pytest + +from he11lib.modes import LGBasis +from he11lib.phase_retrieval import PhaseRetriever, propagate_angular_spectrum +from he11lib.synthetic import SyntheticBeamGenerator + +W0 = 5e-3 +Z0 = 0.5 +WAVELENGTH = 1.76e-3 +PIXEL_SCALE = 3e-4 +IMAGE_SHAPE = (121, 121) + + +def make_basis(): + return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + + +def make_grid(): + coords = (np.arange(IMAGE_SHAPE[0]) - IMAGE_SHAPE[0] // 2) * PIXEL_SCALE + x, y = np.meshgrid(coords, coords) + return x, y + + +def test_propagate_round_trip_recovers_original_field(): + basis = make_basis() + x, y = make_grid() + field = basis.field(x, y, Z0, p=0, l=0) + + forward = propagate_angular_spectrum(field, PIXEL_SCALE, dz=0.05, wavelength=WAVELENGTH) + back = propagate_angular_spectrum(forward, PIXEL_SCALE, dz=-0.05, wavelength=WAVELENGTH) + + np.testing.assert_allclose(back, field, atol=1e-3 * np.max(np.abs(field))) + + +def test_propagate_matches_lgbasis_analytic_evolution(): + basis = make_basis() + x, y = make_grid() + field_at_waist = basis.field(x, y, Z0, p=0, l=0) + + dz = 0.05 + propagated = propagate_angular_spectrum(field_at_waist, PIXEL_SCALE, dz=dz, wavelength=WAVELENGTH) + 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.abs(propagated) ** 2, np.abs(analytic) ** 2, atol=1e-2 * np.max(np.abs(analytic) ** 2) + ) + + +def test_retrieve_recovers_pure_mode_purity(): + # Keep z distances close to the waist so the (widening) beam stays well + # within the frame -- otherwise FFT wraparound/clipping at the edges + # degrades angular-spectrum propagation accuracy. + basis = make_basis() + gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + 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) + + retriever = PhaseRetriever(wavelength=WAVELENGTH) + result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100) + + coeffs = basis.project( + result.field, result.x, result.y, PIXEL_SCALE, result.z, modes=[(0, 0), (1, 0), (0, 1)] + ) + total_power = sum(abs(c) ** 2 for c in coeffs.values()) + purity_00 = abs(coeffs[(0, 0)]) ** 2 / total_power + assert purity_00 > 0.9 + + +def test_retrieve_estimates_beam_center(): + basis = make_basis() + gen = SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + z_list = [0.47, 0.5, 0.53] + true_center = (15 * PIXEL_SCALE, -8 * PIXEL_SCALE) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=z_list, center=true_center, noise_std=1e-5, seed=1 + ) + + retriever = PhaseRetriever(wavelength=WAVELENGTH) + result = retriever.retrieve(planes, viewing_angle_deg=0.0, max_iterations=100) + + 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) diff --git a/tests/test_plotting.py b/tests/test_plotting.py new file mode 100644 index 0000000..a376bb9 --- /dev/null +++ b/tests/test_plotting.py @@ -0,0 +1,64 @@ +import matplotlib.figure +import numpy as np +import pytest + +from he11lib.data import MeasurementPlane, ReconstructionResult +from he11lib.plotting import plot_center_trace, plot_mode_purity, plot_residuals + + +def make_result(**overrides): + defaults = dict( + purity={(0, 0): (0.9, 0.1), (1, 0): (0.1, -0.2)}, + reconstructed_field=np.zeros((5, 5), dtype=complex), + centers=[(0.0, 0.0), (1e-4, -1e-4), (2e-4, -2e-4)], + pointing_angle_deg=0.5, + residuals=[np.ones((5, 5)), np.ones((5, 5)) * 2, np.ones((5, 5)) * 3], + ) + defaults.update(overrides) + return ReconstructionResult(**defaults) + + +def make_planes(): + return [ + MeasurementPlane(flux=np.zeros((5, 5)), z=0.3), + MeasurementPlane(flux=np.zeros((5, 5)), z=0.5), + MeasurementPlane(flux=np.zeros((5, 5)), z=0.7), + ] + + +def test_plot_mode_purity_draws_one_bar_per_mode(): + result = make_result() + fig = plot_mode_purity(result) + + assert isinstance(fig, matplotlib.figure.Figure) + ax = fig.axes[0] + assert len(ax.patches) == len(result.purity) + + +def test_plot_center_trace_plots_one_point_per_plane(): + planes = make_planes() + result = make_result() + fig = plot_center_trace(planes, result) + + assert isinstance(fig, matplotlib.figure.Figure) + ax = fig.axes[0] + line = ax.lines[0] + assert len(line.get_xdata()) == len(planes) + + +def test_plot_residuals_draws_one_axes_per_plane(): + planes = make_planes() + result = make_result() + fig = plot_residuals(planes, result) + + assert isinstance(fig, matplotlib.figure.Figure) + image_axes = [ax for ax in fig.axes if ax.images] + assert len(image_axes) == len(planes) + + +def test_plot_residuals_raises_when_phase_retrieval_used_without_residuals(): + planes = make_planes() + result = make_result(residuals=[], used_phase_retrieval=True) + + with pytest.raises(ValueError, match="residuals"): + plot_residuals(planes, result) diff --git a/tests/test_reconstruct.py b/tests/test_reconstruct.py new file mode 100644 index 0000000..2c20c91 --- /dev/null +++ b/tests/test_reconstruct.py @@ -0,0 +1,112 @@ +from dataclasses import replace + +import pytest + +from he11lib.deconvolution import DiffusionDeconvolver +from he11lib.fitting import ModalFitter +from he11lib.modes import LGBasis +from he11lib.reconstruct import BeamReconstructor +from he11lib.synthetic import SyntheticBeamGenerator + +W0 = 5e-3 +Z0 = 0.5 +WAVELENGTH = 1.76e-3 +PIXEL_SCALE = 4e-4 +IMAGE_SHAPE = (61, 61) +Z_LIST = [0.35, 0.5, 0.65, 0.8] + + +def make_basis(): + return LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + + +def make_generator(basis): + return SyntheticBeamGenerator(basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE) + + +def test_reconstruct_recovers_pure_mode_purity(): + basis = make_basis() + gen = make_generator(basis) + planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=0) + + reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2) + result = reconstructor.reconstruct(planes) + + power_fraction, _ = result.purity[(0, 0)] + assert power_fraction == pytest.approx(1.0, abs=1e-3) + assert result.used_phase_retrieval is False + + +def test_reconstruct_recovers_center_offset(): + basis = make_basis() + gen = make_generator(basis) + true_center = (10 * PIXEL_SCALE, -5 * PIXEL_SCALE) + planes = gen.generate( + coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, center=true_center, noise_std=1e-4, seed=1 + ) + + reconstructor = BeamReconstructor(w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2) + result = reconstructor.reconstruct(planes) + + for cx, cy in result.centers: + assert cx == pytest.approx(true_center[0], abs=2 * PIXEL_SCALE) + assert cy == pytest.approx(true_center[1], abs=2 * PIXEL_SCALE) + + +def test_reconstruct_with_deconvolution_corrects_blur(): + basis = make_basis() + gen = make_generator(basis) + planes = gen.generate(coefficients={(0, 0): 1.0 + 0j}, z_list=Z_LIST, noise_std=1e-4, seed=2) + + deconvolver = DiffusionDeconvolver(thermal_diffusivity=1e-6, dwell_time=30.0) + blurred_planes = [ + replace(p, flux=deconvolver.blur(p.flux, p.pixel_scale)) for p in planes + ] + + # Without deconvolution, blur should measurably hurt purity recovery. + fitter = ModalFitter(basis) + result_no_deconv = fitter.fit(blurred_planes, modes=[(0, 0), (1, 0), (0, 1)]) + purity_no_deconv, _ = result_no_deconv.purity[(0, 0)] + + reconstructor = BeamReconstructor( + w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, deconvolver=deconvolver + ) + result = reconstructor.reconstruct(blurred_planes) + purity_with_deconv, _ = result.purity[(0, 0)] + + assert purity_with_deconv > purity_no_deconv + assert purity_with_deconv > 0.9 + + +def test_reconstruct_forces_phase_retrieval_fallback(): + basis = make_basis() + gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) + 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) + + reconstructor = BeamReconstructor( + w0=W0, z0=Z0, wavelength=WAVELENGTH, max_order=2, force_phase_retrieval=True + ) + result = reconstructor.reconstruct(planes) + + assert result.used_phase_retrieval is True + power_fraction, _ = result.purity[(0, 0)] + assert power_fraction > 0.9 + + +def test_reconstruct_falls_back_automatically_on_high_residual(): + basis = make_basis() + gen = SyntheticBeamGenerator(basis=basis, image_shape=(121, 121), pixel_scale=3e-4) + 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) + + reconstructor = BeamReconstructor( + w0=W0, + z0=Z0, + wavelength=WAVELENGTH, + max_order=2, + phase_retrieval_residual_threshold=1e-8, + ) + result = reconstructor.reconstruct(planes) + + assert result.used_phase_retrieval is True diff --git a/tests/test_synthetic.py b/tests/test_synthetic.py new file mode 100644 index 0000000..90e0e9d --- /dev/null +++ b/tests/test_synthetic.py @@ -0,0 +1,123 @@ +import numpy as np +import pytest + +from he11lib.modes import LGBasis +from he11lib.synthetic import SyntheticBeamGenerator + + +W0 = 5e-3 +Z0 = 0.5 +WAVELENGTH = 1.76e-3 +PIXEL_SCALE = 2e-4 # 0.2 mm/px +IMAGE_SHAPE = (161, 161) # odd so there's a well-defined center pixel + + +def make_generator(): + basis = LGBasis(w0=W0, z0=Z0, wavelength=WAVELENGTH) + return SyntheticBeamGenerator( + basis=basis, image_shape=IMAGE_SHAPE, pixel_scale=PIXEL_SCALE + ) + + +def test_generate_returns_planes_with_requested_z(): + gen = make_generator() + z_list = [0.3, 0.4, 0.5] + planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=z_list) + + assert [p.z for p in planes] == z_list + assert all(p.flux.shape == IMAGE_SHAPE for p in planes) + + +def test_generate_pure_mode_peak_at_image_center_when_centered(): + gen = make_generator() + planes = gen.generate(coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(0.0, 0.0)) + flux = planes[0].flux + + peak_idx = np.unravel_index(np.argmax(flux), flux.shape) + center_idx = (IMAGE_SHAPE[0] // 2, IMAGE_SHAPE[1] // 2) + assert peak_idx == center_idx + + +def test_generate_applies_center_offset(): + gen = make_generator() + offset_m = 20 * PIXEL_SCALE # 20 pixels + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], center=(offset_m, 0.0) + ) + flux = planes[0].flux + + peak_idx = np.unravel_index(np.argmax(flux), flux.shape) + center_row = IMAGE_SHAPE[0] // 2 + center_col = IMAGE_SHAPE[1] // 2 + assert peak_idx[0] == center_row + assert peak_idx[1] == pytest.approx(center_col + 20, abs=1) + + +def test_generate_applies_pointing_angle_as_linear_drift(): + gen = make_generator() + pointing_angle_deg = 1.0 # small tilt + z_list = [Z0, Z0 + 0.2] + planes = gen.generate( + coefficients={(0, 0): 1 + 0j}, + z_list=z_list, + center=(0.0, 0.0), + pointing_angle_deg=pointing_angle_deg, + ) + + peaks_col = [] + for plane in planes: + peak_idx = np.unravel_index(np.argmax(plane.flux), plane.flux.shape) + peaks_col.append(peak_idx[1]) + + expected_shift_m = 0.2 * np.tan(np.deg2rad(pointing_angle_deg)) + expected_shift_px = expected_shift_m / PIXEL_SCALE + actual_shift_px = peaks_col[1] - peaks_col[0] + assert actual_shift_px == pytest.approx(expected_shift_px, abs=1) + + +def test_generate_noise_is_reproducible_with_seed(): + gen = make_generator() + planes_a = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42 + ) + planes_b = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=0.01, seed=42 + ) + np.testing.assert_array_equal(planes_a[0].flux, planes_b[0].flux) + + +def test_generate_noise_std_matches_requested_level(): + gen = make_generator() + noise_std = 0.02 + planes_noisy = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], noise_std=noise_std, seed=1 + ) + 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 + assert np.std(diff) == pytest.approx(noise_std, rel=0.15) + + +def test_generate_viewing_angle_compresses_tilt_axis(): + gen = make_generator() + planes_straight = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=0.0 + ) + planes_tilted = gen.generate( + coefficients={(0, 0): 1 + 0j}, z_list=[Z0], viewing_angle_deg=60.0 + ) + + def width_along_axis(flux, axis): + profile = flux[flux.shape[0] // 2, :] if axis == 1 else flux[:, flux.shape[1] // 2] + half_max = profile.max() / 2 + above = np.where(profile >= half_max)[0] + return above[-1] - above[0] + + 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)