externalize observation weights from Store layer (closes #28)#40
Merged
Conversation
Remove weight(obs) and is_unweighted() from the ObservationStore trait, delete the ObservationWeights enum, and drop the weights field from FactorMajorStore and ArrayStore. Weights now flow as Option<&[f64]> through the operator/preconditioner builders and live owned at the Solver layer as Option<Vec<f64>>. - WeightedDesignOperator::new(design, weights) takes the borrowed slice and computes sqrt_weights once; apply/apply_adjoint match on &self.sqrt_weights so the unweighted hot path stays branch-free. - WeightedDesign::rmatvec_wdt(weights, r, x) and the static WeightedDesign::uid_weight(weights, uid) helper replace the prior store-driven branches. gramian_diagonal and its five tests are deleted as dead code (no production callers). - build_preconditioner, build_local_domains, CrossTab::build_for_pair*, and the accumulate_cross_block family all thread Option<&[f64]>. - Solver gains a weights: Option<Vec<f64>> field; from_design and from_design_with_preconditioner take it. orchestrate::solve / solve_batch and the PyO3 bridge build the owned vector from the optional numpy weights input. Python signatures unchanged.
…perator→DesignOperator Mechanical identifier rename across the workspace plus README and module docstrings. No semantic change; every call site, trait bound, and lib.rs re-export is updated to the new names.
Design is now pure data + layout: store, factors, n_rows, n_dofs,
from_store, n_factors. The matvec helpers (gather_add, scatter_add,
scatter_add_single_factor, level_from_column_or_store, ScatterStrategy,
PAR_THRESHOLD, SCATTER_LOCAL_THRESHOLD, factor_columns) move into
operator.rs as module-private free functions taking &Design. The
DesignOperator::apply / apply_adjoint methods now call them directly.
- Removed Design::matvec_d, rmatvec_dt, rmatvec_wdt, uid_weight (and the
internal gather_add/scatter_add methods). The uid_weight read in
cross_tab.rs inlines to weights.map_or(1.0, |w| w[uid]).
- Solver::solve diagnostic computes D^T W v via apply_adjoint(W^{1/2} v)
instead of the dedicated rmatvec_wdt method; results are bit-identical
because D^T W = D^T W^{1/2} W^{1/2} algebraically.
- Tests/examples/benches that called design.matvec_d(...) etc. now use
DesignOperator::new(&design, None).apply(...) (and likewise for the
adjoint). The internal tests in src/domain.rs and src/operator/tests.rs
are pruned to remove redundant matvec coverage that DesignOperator
already exercises.
…gather
Adopt the cleaner apply/apply_adjoint pattern explored on the closed
refactor/store-design-operator branch:
- Replace gather_add + scatter_add with gather_apply + scatter_apply,
both taking a closure parameter. gather_apply folds the optional
W^{1/2} multiply into the LAST factor's pass, so apply does exactly Q
sweeps over dst regardless of weights (no trailing scale loop).
scatter_apply computes sw[i] * x[i] inline through a closure, so
apply_adjoint allocates no per-call scratch and DesignOperator no
longer holds a Mutex<Vec<f64>>.
- Add length assertion in DesignOperator::new: panics on mismatched
weights length. Solver entry points keep their fallible
validate_weights check, so callers going through Solver / solve()
never trigger the panic.
- ScatterStrategy::pick replaces the ad-hoc match in scatter_add.
- Rename a few stale test/bench labels that still said matvec_d /
rmatvec_dt: test_matvec_d -> test_apply_unweighted_values;
test_rmatvec_dt -> test_apply_adjoint_unweighted_values;
test_large_design_adjoint_property_matvec_d_rmatvec_dt -> drop the
method-name suffix; matvec_weighted_design bench group ->
design_operator_apply.
build_preconditioner is a public entry point that takes Option<&[f64]> weights without going through Solver::from_design. Direct callers could pass mismatched-length weights and trigger an index-out-of-bounds panic deep inside accumulate_cross_block. Add validate_weights at the top so the failure surfaces as the usual WithinError::WeightCountMismatch. Drop residual matvec_d / rmatvec_dt / rmatvec_wdt references in tests/domain.rs doc comments and test names.
…er validation Two behavioral additions from the codex-review follow-ups (4f1410f, 75573ab) weren't captured in the #28 changelog entries: - DesignOperator::new now asserts weight length matches design.n_rows, the Mutex<Vec<f64>> scratch field is removed, and weighted apply fuses the sqrt(W) multiply into the last gather pass. - build_preconditioner now returns WithinError::WeightCountMismatch instead of panicking on out-of-bounds access in CrossTab assembly when weights are mis-sized.
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
Resolves #28. Weights were fused into the storage layer alongside factor levels, which forced every layer above (Design, Operator, preconditioner, solver) to thread an
is_unweighted()predicate through hot paths. This PR moves weights out ofStoreand into theSolver, then separatesDesign(data + layout) fromDesignOperator(linear map).Three bisectable commits, each compiling and passing the full test suite independently:
Store— removeweight()/is_unweighted()from the trait, deleteObservationWeights, drop the weights field fromFactorMajorStore/ArrayStore. Weights flow asOption<&[f64]>through builders and live owned at theSolverlayer asOption<Vec<f64>>.DesignOperator::new(design, weights)buildssqrt_weightsonce;apply/apply_adjointmatch on&self.sqrt_weightsso the unweighted hot path stays branch-free. The deadgramian_diagonalmethod and its five tests are deleted (no production callers).ObservationStore→Store,WeightedDesign→Design,WeightedDesignOperator→DesignOperator. No semantic change.Design—matvec_d,rmatvec_dt,rmatvec_wdt,uid_weightand the internal scatter/gather helpers move fromdomain.rstooperator.rsas module-private functions.Designis now pure data + layout:store,factors,n_rows,n_dofs,from_store,n_factors.Solver::solve's residual diagnostic computesD^T W vasapply_adjoint(weighted_rhs(v))— bit-identical via theD^T W = D^T W^{1/2} W^{1/2}algebraic identity.Python API surface unchanged —
within.solve,within.solve_batch,within.Solverall preserve their existing signatures. Breaking changes are confined to the Rust crate (callers who constructStore/Design/ operators directly, or implement theStoretrait). CHANGELOG updated.Net diff: 27 files, +671 / −997.
Test plan
cargo build --workspace— cleancargo test --workspace— 181 passed, 0 failedcargo test -p within --release -- --ignored --nocapture— 1 passedpixi run develop && pixi run test— 97 passed (Python signatures preserved)cargo bench -p within --bench fixest -- --quick— no regression on the unweighted hot pathtest_uniform_weights_matches_unweighted,test_high_level_solve_weighted,test_solve_batch_weighted,prop_weighted_adjoint_propertyall passcargo fmt --all+cargo clippy --workspace --all-targets— clean