feat: dispersive compiler pipeline (v0.1.5) - #3
Merged
Conversation
Delete examples/results/ (5.1 GB, untracked) and examples/fitting.ipynb (fitting notebooks belong in refl-analysis, not refloxide) from the working tree, and gitignore the remaining generated example-output directories so they can't be accidentally staged again.
…cision First commit of this file. CompiledReflectivityModel/compile_model predate this session as uncommitted work-in-progress from the energy-dispersive refactor and are not yet fully consolidated per tmp/PLAN.md's audit (naming, Model/Objective generations, etc. still need to land) — this commit does not represent finished, reviewed work, just the current state being captured. Adds a module docstring recording tmp/PLAN.md Task 1's benchmark result: Python-side per-energy tensor materialization is 5-8% of total wall time across 8-48 layer structures and shrinks as structure size grows, so it stays a Python responsibility rather than moving into a fused Rust batch entry point. See the docstring for the full methodology and numbers.
orientation.py's AdaptiveOrientationScatterer/bookended_orientation_angles/ attach_to_structure had zero call sites anywhere in src/tests/examples (confirmed by grep across .py, .ipynb, .md) — dead code from an earlier, simpler depth-orientation design superseded by BookendedOrientationProfile's exponential gradient. Deleted the file and its __init__.py re-exports. Note: this __init__.py diff also carries pre-existing uncommitted export additions (compile_model, DeferredScatterer, DispersiveMaterialSLD, OocUniTensorScatterer, Probe, SlabEnergyPlan, TabulatedUniTensorSLD, etc.) that predate this session's work and are not yet fully consolidated per tmp/PLAN.md's audit — captured here alongside the orientation.py cleanup rather than held back further, per the same unfinished-work convention as 618fe6e.
Recovers the coverage lost when this file was deleted uncommitted on this branch. Keeps only the core numeric parity assertion (fused Rust path vs assembled Structure.slabs()/tensor() through the plain uniaxial kernel) — per Testing Policy exception #2, this is the one thing that can't be verified by reading the code, only by running both and comparing. Drops the old benchmark-style timing assertions (test_fused_benchmark_100_slabs, test_legacy_triple_tensor_slabs_slower_than_fused) and the tensor()-call-count check (test_slabs_calls_tensor_once) — verified that one manually instead (tensor() is called exactly once per slabs() call, confirmed via a scratch monkeypatch, no regression found) since it's an inspection-level concern, not a testing-harness one.
Per tmp/USAGE.md's target flat module layout: refloxide.tmm holds deliverables 1 and 2 (uniaxial_reflectivity(_batch), the fused book-ended kernel), re-exported from the compiled extension with real docstrings and types. refloxide.model will compose these with Structure/ReflectModel; this module is also the direct escape hatch for power users who want to bypass Structure entirely.
…rialSLD Per tmp/USAGE.md's design principles 2-4: Scatterer is a subclassable ABC (tensor_at/parameters required, slab_row_at/__call__ provided) so a user can add a material with no built-in class by calling refloxide.optics primitives directly, matching real refnx/pyref idiom exactly -- the `|` operator and `SLD(...)(thick, rough)` call convention are unchanged. Structure/Slab hold exactly one Parameter per geometry field regardless of how many energies they're later evaluated at (no per-energy binding at construction). MaterialSLD is the first unified scatterer: always energy-deferred (tensor_at(energy_ev) resolves periodictable.xsf.index_of_refraction fresh per call, no fixed construction-time energy), verified numerically against the existing DispersiveMaterialSLD/Probe path -- Testing Policy exception #2, since floating-point agreement between the two independent implementations can't be verified by reading the code. Not yet done: UniTensorSLD unification, ReflectModel, Objective, BookendedOrientationProfile's 3-slab fix -- Structure.slab_rows_at(energy) is the seam ReflectModel will call once it exists.
Per tmp/USAGE.md's target flat module layout -- built earlier in this branch's work but not committed at the time (caught while wiring UniTensorSLD in model.py on top of it). refloxide.optics: thin re-export of the Rust lab-tensor/OOC primitives (interp_ooc_linear, uniaxial_lab_tensor, molecular_index_at_ooc, tensor_to_slab_row, etc.) with real docstrings and types. refloxide.data.OpticalConstants: the cached, polars-backed OOC table every dispersive Scatterer should share instead of loading its own copy. from_file normalizes paths via Path.resolve() so path spelling can't defeat the cache; from_dataframe/from_source accept polars or pandas DataFrames (identity-keyed) or an existing instance. cache_size() gives public introspection for the sharing guarantee, pinned by tests/test_optical_constants_cache.py (Testing Policy exception #3 -- a memory/performance contract a future edit could silently break).
Same pattern as MaterialSLD: energy always resolved at tensor_at(energy_ev) call time, no fixed construction-time energy. Backed by refloxide.data.OpticalConstants (accepts a polars/pandas DataFrame, CSV path, or an existing OpticalConstants instance) and refloxide.optics.uniaxial_lab_tensor -- so three UniTensorSLD instances referencing the same OOC source automatically share one cached, interpolatable table (verified in test_unitensorsld_shares_one_cached_opticalconstants_across_instances). Verified numerically against the existing OocUniTensorScatterer/Probe path (Testing Policy exception #2 -- two independent implementations, can't be checked by reading the code alone).
ReflectModel(structure) takes no energy/pol at construction; __call__(q, energy) always returns a Reflectivity NamedTuple (r.s, r.p -- both channels, always available, per tmp/USAGE.md: polarization is inferred from data at Objective time, not chosen on the model). Scalar energy gives shape (len(q),) per channel; array energy gives (len(q), len(energy)), matching the existing CompiledReflectivityModel's (q, E) matrix convention. Reflectivity being a NamedTuple means `r.s`/`r.p` attribute access and `r_s, r_p = model(q, e)` tuple-unpacking both work. Structure gained tensor_rows_at(energy) alongside slab_rows_at(energy) -- the kernel needs the full (N, 3, 3) anisotropic tensor separately from the packed isotropic-average [thick, delta, beta, rough] rows. Verified against the legacy DispersiveStructure's raw materialized layers/tensor fed through the same refloxide.tmm.uniaxial_reflectivity kernel (not against DispersiveReflectModel.model() directly, since that applies default instrument correction stages ReflectModel deliberately doesn't own -- comparing against it would be smeared vs unsmeared curves, not a real parity check).
Per tmp/USAGE.md: refloxide.objective.Objective subclasses refnx.analysis.Objective to reuse its parameter/logp/setp/ varying_parameters bookkeeping (same strategy the existing pxr.objective.ReflectivityObjective already uses successfully) but is its own class with its own constructor and log-likelihood -- not a re-export, and not built around refnx's Data1D/GlobalObjective concepts. refloxide.data.ReflectDataset carries q/energy/pol/r/r_err as flat, independent rows (ragged multi-energy data is fine, no rectangular grid assumed); Objective groups by (energy, pol) so N rows at one energy trigger exactly one ReflectModel call, and reads each row's prediction off whichever channel (.s or .p) its pol selects -- the model itself never picks a channel, matching ReflectModel's design. Verified: logl() matches a hand-computed Gaussian log-likelihood for a single group and for multiple energy/pol groups summed (Testing Policy exception #2), and plugs into an unmodified refnx.analysis.CurveFitter end to end (constructs, exposes varying_parameters(), completes a short differential_evolution fit).
…r ooc/energy Real signature bug: the function took a fourth vacuum_slab argument that was immediately discarded (del vacuum_slab) -- never used for anything. A book-ended profile has exactly two book-ends (surface-adjacent, substrate-adjacent) and one bulk value it interpolates between, so three slabs was already correct in spirit; now it's correct in the signature too. No internal callers existed (confirmed by grep), so this is a clean break, not a deprecation. BookendedOrientationProfile's ooc/energy are now optional, keyword-only, deferred by default (bind_ooc(ooc, energy=...) attaches them later), matching ReflectModel's late-binding design in tmp/USAGE.md. Accessing `.anchor`/`.probe_at()`/`.tensor()` before binding raises a clear RuntimeError naming exactly what to call, rather than silently using a wrong energy or crashing on None. Verified: bind_ooc() produces numerically identical tensors to eager construction with the same ooc/energy (Testing Policy exception #2), and the unbound-access error path is pinned so it can't silently regress into using a stale/None value instead of raising. Note: bookended.py also carries pre-existing uncommitted changes unrelated to this fix, same as previous commits in this session (618fe6e, 13dad08).
Thin wrapper over bookended_from_three_slabs so construction reads as BookendedOrientationProfile.from_slabs(...), matching this codebase's other from_* classmethod constructors (OocAnchor.from_dataframe, OpticalConstants.from_file/from_dataframe). Delegates entirely to the existing, tested module-level function -- no logic duplicated. Not yet wired into refloxide.model (tracked separately): this class still implements the older pxr.plugin.structure.Component protocol (.tensor()/.slabs(), one component contributing many rows), not refloxide.model.Component's protocol (tensor_at(energy_ev)/ slab_row_at(energy_ev), one row per component) -- composes with the legacy Structure today, not yet ReflectModel.
Adds scale_s/scale_p/bkg/dq/q_offset/theta_offset_s/theta_offset_p to refloxide.model.ReflectModel, porting the correction-stage machinery from the original, load-bearing pxr.plugin.model.ReflectModel -- the actual "first class citizen" ask. This is also a real performance fix, not just a rename: the old ReflectModel called refloxide.pxr.tjf4x4's pure-Python kernel port (never the Rust extension, unless patch_pyref was applied); this one always calls refloxide.tmm's Rust kernel, including through resolution smearing (same log-grid + Gaussian-convolve + spline technique, now backed by Rust instead of pure Python). Added .anisotropy(q, energy) matching the old ReflectModel.anisotropy(), reusing the same per-channel theta-offset machinery instead of toggling a `self.pol` attribute and calling model() twice. Verified manually against the old ReflectModel for every correction stage individually (scale+bkg, q_offset, differing per-channel theta offsets, dq smearing, anisotropy) -- all match to float64 precision. Note: the old model's pol='s'/'p' labels are a documented, intentional inversion (reflectivity_for_pol reads pol='s' from [:,1,1]) kept only for historical pyref-dataset compatibility; this port uses the native, non-inverted kernel labeling (Reflectivity.s = [:,0,0]), so parity holds at the raw-channel level, not the swapped pol-label level -- documented on Reflectivity itself.
Objective._predicted previously called ReflectModel once per (energy, pol) group -- N energies sharing an identical q grid (a common multi-energy experimental setup) meant N separate Rust calls, each independently re-materializing the structure's layers/tensor arrays. _build_kernel_batches groups (energy, pol) entries by identical q content once at construction time and routes multi-energy batches through ReflectModel's array-energy path (one uniaxial_reflectivity_batch call), mirroring what refloxide.pxr.objective.ReflectivityObjective's _build_reflectivity_eval_batches already did for the old Objective generation. Pinned with a real test asserting the batch count itself collapses to 1 for 3 energies sharing a q grid -- a pure correctness check wouldn't catch a regression here, since the unbatched fallback is still correct, just slower and more memory-hungry.
… test Objective gains an optional anisotropy_weight (default 0.0, unchanged standard unnormalized logl). When nonzero, blends the base Gaussian log-likelihood with an extra term comparing model.anisotropy(q, energy) against the data's own (r_p - r_s)/(r_p + r_s) at every energy with matching s/p rows on the same q grid, then normalizes by len(data) -- matching refloxide.pxr.plugin.fitters.AnisotropyObjective's exact formula (including its per-point normalization, which only applies in this weighted mode). _build_anisotropy_pairs finds the matching-q-grid s/p groups; excludes energies without a pair rather than raising, since a dataset can legitimately mix anisotropy-comparable and incomparable energies. Per this port's own design (Testing Policy exception #2), added ONE lightweight parity test file rather than restoring the old suites: tests/test_legacy_parity.py checks every correction stage individually plus the full anisotropy-weighted logl against the old, load-bearing ReflectModel/AnisotropyObjective. Correction stages match to 1e-8; anisotropy-weighted logl matches within 0.05 (the old XrayReflectDataset computes anisotropy via tolerance-based q interpolation, this implementation requires an exact q match -- verified during development that both use the identical formula, the residual gap is that interpolation, not a bug). This completes the "port Model and Objective as first-class citizens" work: refloxide.model.ReflectModel and refloxide.objective.Objective now have full correction-stage/anisotropy feature parity with the original pxr.plugin stack, always backed by the Rust kernel (the old stack never was, unless patch_pyref was applied), plus a real memory/speed improvement (batching same-q-grid energies into one kernel call) the old stack didn't have either.
Cell-marked (# %%) interactive comparison, runnable top-to-bottom or stepped through in a Python Interactive window. Three things, same vacuum/SiO2/Si structure at 700 eV: 1. Numeric parity vs stock, unpatched pyref.fitting.ReflectModel (pyref's own pure-Python kernel -- confirmed exact, ~1e-16). 2. Speed: unpatched pyref vs refloxide (4x), and pyref patched with refloxide's Rust kernel vs refloxide (~1.1x) -- isolates how much of the win is "just Rust" (most of it) vs the new API/batching on top. 3. Fitting: refloxide.objective.Objective vs pyref.fitting.AnisotropyObjective on the same synthetic noisy s+p dataset -- matching logl, matching recovered thickness, comparable fit time. Verified by actually running it (MPLBACKEND=Agg for the non-interactive check) -- caught and fixed one real difference along the way: pyref's Structure has constrained parameters that can't have .vary set directly, which the disable-everything-then-reenable-one-parameter pattern needs to skip.
Expose OOC lookup, lab-frame rotation, and slab-row packing so energy-deferred scatterers stay thin parameter wrappers.
Ship deferred OOC scatterers, compile_structure/SlabEnergyPlan, and a batched ReflectivityObjective as the native dispersive fitting path.
Add MixedUniTensorSLD, BookendedComponent, and depth-profile coverage so graded films share geometry with the dispersive compiler path.
Retain BatchedGlobalObjective-era glue for pickle migration while new fits prefer the compile/ReflectivityObjective pipeline.
Add patch_pyref_if_needed and pyref_patch_report so kernel swaps stay opt-in and stock pyref comparisons remain reliable.
Capture bookended/ZnPc workflows and document the compile pipeline as the preferred energy-dispersive fitting path.
Align Rust photon-energy params, refresh package docs, and delete unused stacks/plotting/unittensor surfaces after example rewrites.
Align package metadata for the dispersive-compiler release tag and allow the OOC PyO3 binding's argument count under -D warnings.
Drop runtime ipykernel, add pyarrow, upgrade vulnerable transitive deps, exclude legacy pxr/plugin from ty, and fix ruff nits.
Expand the Release workflow to publish manylinux, macOS, and Windows wheels; add a PR wheels smoke job; drop unused lapack for portable links.
Release Linux builds multiple ABIs with --find-interpreter; installing dist/*.whl failed when pip tried an incompatible cp312 wheel.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
compile_structure/compile_model/ReflectivityObjective, plusrefloxide.model/objective/data/tmm/optics.molecular_index_at_ooc, lab tensors, slab-row packing) and consistentenergy_evnaming on the public kernel surface.pxr.stacks/plotting/plugin.unittensor; keeps load-bearingpxr.pluginhelpers for refl-analysis until migration.pyproject.toml,Cargo.toml,src/refloxide/__init__.py, locks) for the post-mergev0.1.5release tag.-D warningsfor CI (too_many_argumentson the OOC PyO3 binding).Test plan
lint,type-check,test,rust-check, scans)git tag v0.1.5 && git push origin v0.1.5to trigger Release workflow (PyPI + GitHub release)make develop && uv run pytestandmake release-smoke