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 <noreply@anthropic.com>
This commit is contained in:
Martino Ferrari
+23
@@ -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
|
||||
@@ -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.
|
||||
@@ -0,0 +1,674 @@
|
||||
GNU GENERAL PUBLIC LICENSE
|
||||
Version 3, 29 June 2007
|
||||
|
||||
Copyright (C) 2007 Free Software Foundation, Inc. <https://fsf.org/>
|
||||
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.
|
||||
|
||||
<one line to give the program's name and a brief idea of what it does.>
|
||||
Copyright (C) <year> <name of author>
|
||||
|
||||
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 <https://www.gnu.org/licenses/>.
|
||||
|
||||
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:
|
||||
|
||||
<program> Copyright (C) <year> <name of author>
|
||||
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
|
||||
<https://www.gnu.org/licenses/>.
|
||||
|
||||
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
|
||||
<https://www.gnu.org/licenses/why-not-lgpl.html>.
|
||||
@@ -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.
|
||||
+223
@@ -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).
|
||||
@@ -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.
|
||||
@@ -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()
|
||||
@@ -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"
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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
|
||||
<LG_mode | field>, 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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
)
|
||||
@@ -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
|
||||
@@ -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,
|
||||
)
|
||||
@@ -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
|
||||
@@ -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"]
|
||||
@@ -0,0 +1,3 @@
|
||||
import matplotlib
|
||||
|
||||
matplotlib.use("Agg")
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user