externalize observation weights from Store layer (closes #28)

CHANGELOG.md

@@ -32,6 +32,26 @@ Modified LSMR is now the sole iterative solver, replacing CG and GMRES.
   type parameter; the `LocalSolveInvoker` trait, `DefaultLocalSolveInvoker`,
   and `with_strategy_and_invoker` constructor are removed.
 - LSMR vector kernels parallelized via Rayon.
+- **BREAKING:** `ObservationStore` trait renamed to `Store`;
+  `WeightedDesign` to `Design`; `WeightedDesignOperator` to
+  `DesignOperator` (closes #28).
+- **BREAKING:** Observation weights externalized from the store layer.
+  `FactorMajorStore::new` and `ArrayStore::new` drop their weights
+  argument. `Solver::from_design` / `from_design_with_preconditioner` and
+  `build_preconditioner` gain `Option<Vec<f64>>` / `Option<&[f64]>`
+  weights parameters. Python `solve` / `Solver` signatures unchanged.
+- **BREAKING:** `Design` is now pure data + layout. The `matvec_d`,
+  `rmatvec_dt`, `rmatvec_wdt`, `gramian_diagonal`, and `uid_weight`
+  methods are removed — use `DesignOperator::new(&design, weights).apply`
+  / `apply_adjoint` instead.
+- `DesignOperator::new` validates that `weights.len() == design.n_rows`
+  and panics on mismatch; the scratch `Mutex<Vec<f64>>` field is gone and
+  weighted `apply` / `apply_adjoint` no longer allocate. The weighted
+  `apply` fuses the `W^{1/2}` multiply into the last gather pass, so
+  there is no trailing scale loop.
+- `build_preconditioner` now returns `WithinError::WeightCountMismatch`
+  for wrong-length weights instead of panicking on out-of-bounds access
+  inside `CrossTab` assembly.
 
 ### Removed
 
@@ -41,9 +61,11 @@ Modified LSMR is now the sole iterative solver, replacing CG and GMRES.
   `SolverParams.max_refinements`, Python `CG`/`GMRES`/
   `MultiplicativeSchwarz`, `MultiplicativeSchwarzPreconditioner`,
   `ResidualUpdater`, `OperatorResidualUpdater`, `IdentityOperator`).
-- **BREAKING:** `Gramian`, `GramianOperator`, `DesignOperator`,
-  `build_schwarz`, `FeSchwarz`, and `WithinError::Overflow` removed from
-  the `within` public surface. LSMR uses `WeightedDesignOperator` directly.
+- **BREAKING:** `Gramian`, `GramianOperator`, the previous bare
+  `DesignOperator`, `build_schwarz`, `FeSchwarz`, and
+  `WithinError::Overflow` removed from the `within` public surface. LSMR
+  uses the rectangular weighted design operator directly (see also #28,
+  which later renamed it to `DesignOperator`).
 - **BREAKING:** `schwarz_precond::solve::{cg, gmres}` and the `solve`
   module removed; use crate-root `lsmr`/`mlsmr`.
   `schwarz_precond::schwarz::{additive, multiplicative}` flattened into
@@ -60,6 +82,8 @@ Modified LSMR is now the sole iterative solver, replacing CG and GMRES.
   variants moved onto `SolveError`. `PyFePreconditioner.apply` raises
   `RuntimeError` on local-solver failure instead of returning NaNs
   (closes #29).
+- **BREAKING:** `ObservationWeights` enum removed; `Store::weight` and
+  `Store::is_unweighted` removed from the trait (closes #28).
 
 ## [0.1.0] - 2026-03-12
 

crates/within-py/src/lib.rs

@@ -43,8 +43,8 @@ use within::config::{
     ApproxCholConfig, ApproxSchurConfig, LocalSolverConfig, Preconditioner, ReductionStrategy,
     SolverParams, DEFAULT_DENSE_SCHUR_THRESHOLD,
 };
-use within::domain::WeightedDesign;
-use within::observation::{FactorMajorStore, ObservationWeights};
+use within::domain::Design;
+use within::observation::FactorMajorStore;
 use within::{
     solve as solve_native, solve_batch as solve_batch_native, FePreconditioner, Operator,
     SolveResult, Solver,
@@ -700,27 +700,27 @@ impl PySolver {
         let factor_levels: Vec<Vec<u32>> = (0..n_factors)
             .map(|f| cats.column(f).iter().copied().collect())
             .collect();
-        let w = match &weights {
-            Some(w) => ObservationWeights::Dense(w.as_array().iter().copied().collect()),
-            None => ObservationWeights::Unit,
-        };
-        let store = FactorMajorStore::new(factor_levels, w, n_obs).map_err(value_err)?;
-        let design = WeightedDesign::from_store(store).map_err(value_err)?;
+        let store = FactorMajorStore::new(factor_levels, n_obs).map_err(value_err)?;
+        let design = Design::from_store(store).map_err(value_err)?;
+        let weights_vec: Option<Vec<f64>> = weights
+            .as_ref()
+            .map(|w| w.as_array().iter().copied().collect());
 
         // Handle pre-built FePreconditioner separately (uses a different constructor);
         // all other variants go through extract_preconditioner_config.
-        let solver =
-            if let Some(Ok(fe)) = preconditioner.map(|o| o.downcast::<PyFePreconditioner>()) {
-                let fe_precond = fe.get().inner.clone();
-                py.allow_threads(|| {
-                    Solver::from_design_with_preconditioner(design, &params, fe_precond)
-                })
+        let solver = if let Some(Ok(fe)) =
+            preconditioner.map(|o| o.downcast::<PyFePreconditioner>())
+        {
+            let fe_precond = fe.get().inner.clone();
+            py.allow_threads(|| {
+                Solver::from_design_with_preconditioner(design, weights_vec, &params, fe_precond)
+            })
+            .map_err(value_err)?
+        } else {
+            let precond = extract_preconditioner_config(py, preconditioner)?;
+            py.allow_threads(|| Solver::from_design(design, weights_vec, &params, precond.as_ref()))
                 .map_err(value_err)?
-            } else {
-                let precond = extract_preconditioner_config(py, preconditioner)?;
-                py.allow_threads(|| Solver::from_design(design, &params, precond.as_ref()))
-                    .map_err(value_err)?
-            };
+        };
 
         Ok(Self { solver })
     }

crates/within/README.md

@@ -51,14 +51,15 @@ assert!(result.converged);
 
 The crate is organized in four layers:
 
-1. **`observation`** — Per-observation factor levels and weights via
-   `FactorMajorStore` and the `ObservationStore` trait.
+1. **`observation`** — Per-observation factor levels via `FactorMajorStore`
+   and the `Store` trait. Observation weights are not owned here — they
+   flow as `Option<&[f64]>` to the operator layer.
 
-2. **`domain`** — Domain decomposition. `WeightedDesign` wraps a store with
+2. **`domain`** — Domain decomposition. `Design` wraps a store with
    factor metadata; `build_local_domains` constructs factor-pair subdomains
    with partition-of-unity weights for the Schwarz preconditioner.
 
-3. **`operator`** — Linear algebra primitives. `WeightedDesignOperator`
+3. **`operator`** — Linear algebra primitives. `DesignOperator`
    (rectangular `sqrt(W) D` for LSMR) and Schwarz preconditioner builders
    that wire approximate Cholesky local solvers into the generic
    `schwarz-precond` framework.

crates/within/benches/fixest.rs

@@ -10,9 +10,9 @@ use schwarz_precond::Operator;
 use within::config::{
     ApproxCholConfig, LocalSolverConfig, Preconditioner, ReductionStrategy, SolverParams,
 };
-use within::domain::WeightedDesign;
-use within::observation::{FactorMajorStore, ObservationWeights};
-use within::operator::WeightedDesignOperator;
+use within::domain::Design;
+use within::observation::FactorMajorStore;
+use within::operator::DesignOperator;
 use within::Solver;
 
 // ===========================================================================
@@ -45,10 +45,7 @@ impl Case {
     }
 }
 
-fn generate_fixest_like_case(
-    case: Case,
-    seed: u64,
-) -> (WeightedDesign<FactorMajorStore>, Vec<f64>) {
+fn generate_fixest_like_case(case: Case, seed: u64) -> (Design<FactorMajorStore>, Vec<f64>) {
     let mut rng = SmallRng::seed_from_u64(seed);
     let n_years = 10usize;
     let n_indiv_per_firm = 23usize;
@@ -76,17 +73,18 @@ fn generate_fixest_like_case(
         vec![indiv_id, year, firm_id]
     };
 
-    let store = FactorMajorStore::new(factor_levels, ObservationWeights::Unit, case.n_obs)
-        .expect("valid factor-major store");
-    let design = WeightedDesign::from_store(store).expect("valid design");
+    let store = FactorMajorStore::new(factor_levels, case.n_obs).expect("valid factor-major store");
+    let design = Design::from_store(store).expect("valid design");
 
     let mut x_true = vec![0.0; design.n_dofs];
     for x in &mut x_true {
         *x = rng.random_range(-1.0..1.0);
     }
 
     let mut y = vec![0.0; case.n_obs];
-    design.matvec_d(&x_true, &mut y);
+    DesignOperator::new(&design, None)
+        .apply(&x_true, &mut y)
+        .expect("apply succeeds");
     for yi in &mut y {
         *yi += 0.1 * rng.random_range(-1.0..1.0);
     }
@@ -123,15 +121,15 @@ fn configure_group<'a>(
 fn run_smoke(
     group: &mut BenchmarkGroup<'_, WallTime>,
     label: &str,
-    design: &WeightedDesign<FactorMajorStore>,
+    design: &Design<FactorMajorStore>,
     y: &[f64],
 ) {
     group.bench_function(BenchmarkId::new(label, ""), |b| {
         b.iter(|| run_lsmr_one_level(design, y, false))
     });
 }
 
-fn run_lsmr_one_level(design: &WeightedDesign<FactorMajorStore>, y: &[f64], ac2: bool) {
+fn run_lsmr_one_level(design: &Design<FactorMajorStore>, y: &[f64], ac2: bool) {
     let params = SolverParams {
         tol: TOL,
         maxiter: MAXITER,
@@ -140,7 +138,7 @@ fn run_lsmr_one_level(design: &WeightedDesign<FactorMajorStore>, y: &[f64], ac2:
     let cfg = one_level_local_solver(ac2);
     let precond = Preconditioner::Additive(cfg, ReductionStrategy::Auto);
     let solver =
-        Solver::from_design(design.clone(), &params, Some(&precond)).expect("solver build");
+        Solver::from_design(design.clone(), None, &params, Some(&precond)).expect("solver build");
     let _ = solver.solve(y).expect("solve");
 }
 
@@ -275,13 +273,13 @@ fn matvec_cases() -> [Case; 4] {
 }
 
 fn bench_matvec(c: &mut Criterion) {
-    let mut group = configure_group(c, "matvec_weighted_design", 50, 200);
+    let mut group = configure_group(c, "design_operator_apply", 50, 200);
     for case in matvec_cases() {
         let label = case.label();
         let (design, _y) = generate_fixest_like_case(case, 42);
         let n_dofs = design.n_dofs;
         let n_obs = design.n_rows;
-        let op = WeightedDesignOperator::new(&design);
+        let op = DesignOperator::new(&design, None);
         let x: Vec<f64> = (0..n_dofs).map(|i| (i as f64).sin()).collect();
         let mut y = vec![0.0; n_obs];
         group.bench_function(BenchmarkId::new("apply", &label), |b| {

crates/within/benches/store_backend.rs

@@ -8,8 +8,8 @@ use rand::rngs::SmallRng;
 use rand::{Rng, SeedableRng};
 
 use within::config::{Preconditioner, SolverParams};
-use within::domain::WeightedDesign;
-use within::observation::{ArrayStore, FactorMajorStore, ObservationWeights};
+use within::domain::Design;
+use within::observation::{ArrayStore, FactorMajorStore};
 use within::Solver;
 
 const TOL: f64 = 1e-6;
@@ -100,10 +100,9 @@ fn bench_store_backends(c: &mut Criterion) {
                 let factor_levels: Vec<Vec<u32>> = (0..p.categories_c.ncols())
                     .map(|q| p.categories_c.column(q).to_vec())
                     .collect();
-                let store =
-                    FactorMajorStore::new(factor_levels, ObservationWeights::Unit, *n_obs).unwrap();
-                let design = WeightedDesign::from_store(store).unwrap();
-                let solver = Solver::from_design(design, &p.params, precond_ref).unwrap();
+                let store = FactorMajorStore::new(factor_levels, *n_obs).unwrap();
+                let design = Design::from_store(store).unwrap();
+                let solver = Solver::from_design(design, None, &p.params, precond_ref).unwrap();
                 let r = solver.solve(&p.y).unwrap();
                 assert!(r.converged);
             });
@@ -112,10 +111,9 @@ fn bench_store_backends(c: &mut Criterion) {
         // ArrayStore C-order: zero-copy, strided columns.
         group.bench_function(BenchmarkId::new("Array(C)", &p.label), |b| {
             b.iter(|| {
-                let store =
-                    ArrayStore::new(p.categories_c.view(), ObservationWeights::Unit).unwrap();
-                let design = WeightedDesign::from_store(store).unwrap();
-                let solver = Solver::from_design(design, &p.params, precond_ref).unwrap();
+                let store = ArrayStore::new(p.categories_c.view()).unwrap();
+                let design = Design::from_store(store).unwrap();
+                let solver = Solver::from_design(design, None, &p.params, precond_ref).unwrap();
                 let r = solver.solve(&p.y).unwrap();
                 assert!(r.converged);
             });
@@ -124,10 +122,9 @@ fn bench_store_backends(c: &mut Criterion) {
         // ArrayStore F-order: zero-copy, contiguous columns.
         group.bench_function(BenchmarkId::new("Array(F)", &p.label), |b| {
             b.iter(|| {
-                let store =
-                    ArrayStore::new(p.categories_f.view(), ObservationWeights::Unit).unwrap();
-                let design = WeightedDesign::from_store(store).unwrap();
-                let solver = Solver::from_design(design, &p.params, precond_ref).unwrap();
+                let store = ArrayStore::new(p.categories_f.view()).unwrap();
+                let design = Design::from_store(store).unwrap();
+                let solver = Solver::from_design(design, None, &p.params, precond_ref).unwrap();
                 let r = solver.solve(&p.y).unwrap();
                 assert!(r.converged);
             });

crates/within/examples/solve_demo.rs

@@ -19,19 +19,22 @@ fn main() {
     }
 
     // Build design to compute D * x_true.
-    use within::observation::{FactorMajorStore, ObservationWeights};
-    use within::WeightedDesign;
+    use schwarz_precond::Operator as _;
+    use within::observation::FactorMajorStore;
+    use within::operator::DesignOperator;
+    use within::Design;
 
     let factor_levels = vec![categories.column(0).to_vec(), categories.column(1).to_vec()];
-    let store = FactorMajorStore::new(factor_levels, ObservationWeights::Unit, n_obs)
-        .expect("valid factor-major store");
-    let design = WeightedDesign::from_store(store).expect("valid design");
+    let store = FactorMajorStore::new(factor_levels, n_obs).expect("valid factor-major store");
+    let design = Design::from_store(store).expect("valid design");
 
     // True coefficient vector: x_true[j] = (j mod 7) - 3.
     let total_dofs = design.n_dofs;
     let x_true: Vec<f64> = (0..total_dofs).map(|j| (j % 7) as f64 - 3.0).collect();
     let mut y = vec![0.0; n_obs];
-    design.matvec_d(&x_true, &mut y);
+    DesignOperator::new(&design, None)
+        .apply(&x_true, &mut y)
+        .expect("apply succeeds");
     // Add small deterministic perturbation so the system is not trivially exact.
     for (i, yi) in y.iter_mut().enumerate() {
         *yi += 0.01 * ((i * 7 + 3) % 13) as f64 - 0.06;