Add caller layout abstraction FactorLabel

Author: leostimpfle

README.md

@@ -217,8 +217,9 @@ let r2 = solver.solve(&another_y, &LsmrOptions::default())?;  // reuses precondi
 | `SchurMode` | `Approximate(ApproxSchurConfig)` \| `Exact` |
 | `Preconditioner` | Opaque built handle — reuse via `Solver::new(.., precond)` (owned or `&`) |
 | `Effect` | `Effect::new(levels: &[u32], intercept: bool, slopes: impl IntoIterator<Item = &[f64]>) -> Result<Self, BuildError>` |
-| `CoefficientAddress` | `{ channel: Channel { term, column }, level: usize }` |
-| `CoefficientLayout` | `index(CoefficientAddress) -> Option<usize>`, `address(usize) -> Option<CoefficientAddress>`, `n_dofs()`, `n_terms()`, `n_levels(term)`, `n_columns(term)` |
+| `FactorLabel` | Opaque caller-visible label; currently constructible from `u32` and readable with `try_as_u32()` |
+| `CoefficientAddress` | `{ channel: Channel { term, column }, level: FactorLabel }` |
+| `CoefficientLayout` | `index(&CoefficientAddress) -> Option<usize>`, `address(usize) -> Option<CoefficientAddress>`, `n_dofs()`, `n_terms()`, `n_levels(term)`, `n_columns(term)` |
 
 ### Varying slopes
 
@@ -240,8 +241,8 @@ let terms = vec![
 let r = solve(terms, &y, None, None, None)?;
 
 // Read firm level 3's x-slope via the layout map (column 0 = intercept, 1 = first slope):
-let at = CoefficientAddress { channel: Channel { term: 0, column: 1 }, level: 3 };
-println!("{}", r.x[r.layout.index(at).unwrap()]);
+let at = CoefficientAddress { channel: Channel { term: 0, column: 1 }, level: 3.into() };
+println!("{}", r.x[r.layout.index(&at).unwrap()]);
 ```
 
 `Solver::new` takes the same `Vec<Effect>`, so a slope design can be reused

crates/within-py/src/results.rs

@@ -71,7 +71,7 @@ pub struct PyUnidentifiedDirection {
     #[pyo3(get)]
     pub term: usize,
     #[pyo3(get)]
-    pub level: usize,
+    pub level: u32,
     #[pyo3(get)]
     pub column: usize,
 }
@@ -114,22 +114,28 @@ impl PyCoefficientLayout {
             .ok_or_else(|| self.term_oob(term))
     }
 
-    fn index(&self, term: usize, level: usize, column: usize) -> PyResult<usize> {
+    fn index(&self, term: usize, level: u32, column: usize) -> PyResult<usize> {
         let at = CoefficientAddress {
             channel: Channel { term, column },
-            level,
+            level: level.into(),
         };
-        self.inner.index(at).ok_or_else(|| {
+        self.inner.index(&at).ok_or_else(|| {
             PyIndexError::new_err(format!(
                 "coefficient address (term={term}, level={level}, column={column}) out of range"
             ))
         })
     }
 
-    fn address(&self, index: usize) -> PyResult<(usize, usize, usize)> {
+    fn address(&self, index: usize) -> PyResult<(usize, u32, usize)> {
         self.inner
             .address(index)
-            .map(|at| (at.channel.term, at.level, at.channel.column))
+            .map(|at| {
+                (
+                    at.channel.term,
+                    at.level.try_as_u32().expect("u32 factor label"),
+                    at.channel.column,
+                )
+            })
             .ok_or_else(|| {
                 PyIndexError::new_err(format!(
                     "x index {index} out of range (n_dofs={})",
@@ -164,7 +170,7 @@ pub(crate) fn into_py_result(py: Python<'_>, result: SolveResult) -> PySolveResu
             .iter()
             .map(|u| PyUnidentifiedDirection {
                 term: u.channel.term,
-                level: u.level,
+                level: u.level.try_as_u32().expect("u32 factor label"),
                 column: u.channel.column,
             })
             .collect(),
@@ -199,7 +205,7 @@ pub(crate) fn into_py_batch_result(
             .iter()
             .map(|u| PyUnidentifiedDirection {
                 term: u.channel.term,
-                level: u.level,
+                level: u.level.try_as_u32().expect("u32 factor label"),
                 column: u.channel.column,
             })
             .collect(),

crates/within/src/domain.rs

@@ -78,11 +78,41 @@ impl<T> Loading<T> {
     }
 }
 
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+enum FactorLabelValue {
+    U32(u32),
+}
+
+/// A caller-visible factor label.
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct FactorLabel(FactorLabelValue);
+
+impl FactorLabel {
+    /// Return the label as `u32`, or `None` when it has another representation.
+    pub fn try_as_u32(&self) -> Option<u32> {
+        match self.0 {
+            FactorLabelValue::U32(value) => Some(value),
+        }
+    }
+}
+
+impl From<u32> for FactorLabel {
+    fn from(value: u32) -> Self {
+        Self(FactorLabelValue::U32(value))
+    }
+}
+
+impl From<&FactorLabel> for FactorLabel {
+    fn from(value: &FactorLabel) -> Self {
+        value.clone()
+    }
+}
+
 /// Mapping between caller-visible factor labels and numerical level positions.
 ///
 /// This is an identity mapping for the existing dense `u32` API. Later
 /// encodings can store an explicit mapping without changing `TermMeta`.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub(crate) struct FactorEncoding {
     n_levels: usize,
 }
@@ -95,6 +125,18 @@ impl FactorEncoding {
     pub(crate) fn n_levels(&self) -> usize {
         self.n_levels
     }
+
+    pub(crate) fn position(&self, label: &FactorLabel) -> Option<usize> {
+        let position = label.try_as_u32()? as usize;
+        (position < self.n_levels).then_some(position)
+    }
+
+    pub(crate) fn label(&self, position: usize) -> Option<FactorLabel> {
+        if position >= self.n_levels {
+            return None;
+        }
+        u32::try_from(position).ok().map(FactorLabel::from)
+    }
 }
 
 /// Per-term metadata; coefficient `c` of `level` lives at `offset + c · n_levels + level`.
@@ -442,6 +484,14 @@ mod tests {
         ));
     }
 
+    #[test]
+    fn u32_factor_label_round_trips() {
+        let label = FactorLabel::from(u32::MAX);
+
+        assert_eq!(label.try_as_u32(), Some(u32::MAX));
+        assert_eq!(FactorLabel::from(&label), label);
+    }
+
     #[test]
     fn validate_weights_checks_count_and_finiteness() {
         let design = Design::from_frame(frame(vec![vec![0, 0, 0, 0, 0]], vec![])).unwrap();

crates/within/src/lib.rs

@@ -42,7 +42,7 @@ pub use config::{
     ApproxCholConfig, ApproxSchurConfig, LocalSolverConfig, LsmrOptions, PreconditionerConfig,
     ReductionStrategy, ScalingConfig, ScalingFailure, SchurMode,
 };
-pub use domain::{Design, Effect};
+pub use domain::{Design, Effect, FactorLabel};
 pub use error::{BuildError, BuildWarning, SolveError, WithinError};
 pub use operator::schwarz::Preconditioner;
 pub use solver::{

crates/within/src/solver.rs

@@ -13,12 +13,12 @@ use crate::channel::Channel;
 use crate::config::{LsmrOptions, PreconditionerConfig};
 use crate::domain::collinearity::detect_collinear_slopes;
 use crate::domain::level_moments::TermMoments;
-use crate::domain::{Design, Effect};
+use crate::domain::{Design, Effect, FactorEncoding};
 use crate::observation::ObservationFrame;
 use crate::operator::design::gather_apply;
 use crate::operator::schwarz::{build_preconditioner, Preconditioner};
 use crate::operator::DesignOperator;
-use crate::{BuildError, BuildWarning, SolveError, WithinError};
+use crate::{BuildError, BuildWarning, FactorLabel, SolveError, WithinError};
 
 mod reparam;
 #[cfg(test)]
@@ -112,27 +112,47 @@ impl From<&Preconditioner> for PreconditionerInput {
 }
 
 /// One coefficient of the design: a [`Channel`] at one level of its term.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CoefficientAddress {
     /// The coefficient column this address sits in.
     pub channel: Channel,
-    /// Level index within the term (`0..n_levels`).
-    pub level: usize,
+    /// Caller-visible factor label.
+    pub level: FactorLabel,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+struct CoefficientPosition {
+    channel: Channel,
+    level: usize,
+}
+
+impl CoefficientPosition {
+    fn to_caller_address(self, design: &Design) -> CoefficientAddress {
+        let level = design.terms[self.channel.term]
+            .encoding
+            .label(self.level)
+            .expect("coefficient position belongs to its term");
+
+        CoefficientAddress {
+            channel: self.channel,
+            level,
+        }
+    }
 }
 
 /// Translates a [`CoefficientAddress`] to its flat index in [`SolveResult::x`]
-/// and back, so callers need not reconstruct the term-major offset formula
-/// (`offset + column * n_levels + level`) by hand.
+/// and back, including the translation between caller-visible labels and
+/// internal compact level positions.
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct CoefficientLayout {
     terms: Vec<TermLayout>,
     n_dofs: usize,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 struct TermLayout {
     offset: usize,
-    n_levels: usize,
+    encoding: FactorEncoding,
     n_columns: usize,
 }
 
@@ -143,7 +163,7 @@ impl CoefficientLayout {
             .iter()
             .map(|t| TermLayout {
                 offset: t.offset,
-                n_levels: t.n_levels(),
+                encoding: t.encoding.clone(),
                 n_columns: t.n_columns(),
             })
             .collect();
@@ -165,7 +185,7 @@ impl CoefficientLayout {
 
     /// Level count of `term`, or `None` if `term` is out of range.
     pub fn n_levels(&self, term: usize) -> Option<usize> {
-        self.terms.get(term).map(|t| t.n_levels)
+        self.terms.get(term).map(|t| t.encoding.n_levels())
     }
 
     /// Coefficient-column count of `term` (`intercept? + slopes`, ordered
@@ -174,12 +194,18 @@ impl CoefficientLayout {
         self.terms.get(term).map(|t| t.n_columns)
     }
 
-    /// Flat [`SolveResult::x`] index of `at`, or `None` if any coordinate is
-    /// out of range.
-    pub fn index(&self, at: CoefficientAddress) -> Option<usize> {
-        let t = self.terms.get(at.channel.term)?;
-        (at.level < t.n_levels && at.channel.column < t.n_columns)
-            .then(|| t.offset + at.channel.column * t.n_levels + at.level)
+    /// Flat [`SolveResult::x`] index of `at`, or `None` if its term,
+    /// column, or caller-visible level label is out of range.
+    pub fn index(&self, at: &CoefficientAddress) -> Option<usize> {
+        let term = self.terms.get(at.channel.term)?;
+
+        if at.channel.column >= term.n_columns {
+            return None;
+        }
+
+        let position = term.encoding.position(&at.level)?;
+
+        Some(term.offset + at.channel.column * term.encoding.n_levels() + position)
     }
 
     /// The address of flat index `i`, or `None` if `i >= n_dofs`.
@@ -191,12 +217,18 @@ impl CoefficientLayout {
         let term = self.terms.partition_point(|t| t.offset <= i) - 1;
         let t = &self.terms[term];
         let within = i - t.offset;
+        let n_levels = t.encoding.n_levels();
+        let level = t
+            .encoding
+            .label(within % n_levels)
+            .expect("coefficient position belongs to the term encoding");
+
         Some(CoefficientAddress {
             channel: Channel {
                 term,
-                column: within / t.n_levels,
+                column: within / n_levels,
             },
-            level: within % t.n_levels,
+            level,
         })
     }
 }
@@ -207,9 +239,10 @@ impl CoefficientLayout {
 pub struct SolveResult {
     /// Fixed-effect coefficients (length = total DOFs across all factors).
     ///
-    /// Term-major: coefficient column `c` of level `level` sits at
-    /// `term_offset + c * n_levels + level`, columns ordered
-    /// `[intercept?, slopes…]`. Slots for unidentified directions hold the
+    /// Term-major by compact level position `p`: coefficient column `c` sits at
+    /// `term_offset + c * n_levels + p`, with columns ordered
+    /// `[intercept?, slopes…]`. Use [`SolveResult::layout`] to translate caller
+    /// labels to these slots. Slots for unidentified directions hold the
     /// minimal-norm value `0`, never NaN; see [`SolveResult::unidentified`].
     pub x: Vec<f64>,
     /// Per-level directions the data cannot identify.
@@ -504,10 +537,15 @@ impl<'a> Solver<'a> {
     /// Per-level directions the data cannot identify, shared across all RHS:
     /// identification depends only on the design and weights, never on `y`.
     fn unidentified(&self) -> Vec<CoefficientAddress> {
-        self.reparam
-            .as_ref()
-            .map(|rp| rp.unidentified.clone())
-            .unwrap_or_default()
+        let Some(reparam) = &self.reparam else {
+            return Vec::new();
+        };
+        reparam
+            .unidentified
+            .iter()
+            .copied()
+            .map(|position| position.to_caller_address(&self.design))
+            .collect()
     }
 
     /// Solve for a single RHS vector with the given LSMR tuning.

crates/within/src/solver/reparam.rs

@@ -3,7 +3,7 @@
 use crate::domain::level_moments::{BasisScratch, TermMoments};
 use crate::domain::Design;
 
-use super::CoefficientAddress;
+use super::CoefficientPosition;
 use crate::channel::Channel;
 use crate::linalg::dot;
 
@@ -16,7 +16,7 @@ mod tests;
 pub(crate) struct SlopeReparam {
     terms: Vec<TermReparam>,
     /// Directions the data cannot identify, ascending in `(term, level, column)`.
-    pub(crate) unidentified: Vec<CoefficientAddress>,
+    pub(super) unidentified: Vec<CoefficientPosition>,
 }
 
 /// One slope-bearing term's whitening state.
@@ -77,7 +77,7 @@ impl TermReparam {
         design: &mut Design<'_>,
         term: usize,
         moments: &TermMoments,
-        unidentified: &mut Vec<CoefficientAddress>,
+        unidentified: &mut Vec<CoefficientPosition>,
     ) -> Self {
         let meta = &design.terms[term];
         let (offset, n_levels) = (meta.offset, meta.n_levels());
@@ -103,14 +103,14 @@ impl TermReparam {
             moments.basis(level, &mut scratch);
             let (w, kept) = (&scratch.basis, &scratch.kept);
             if intercept && moments.w_sum(level) == 0.0 {
-                unidentified.push(CoefficientAddress {
+                unidentified.push(CoefficientPosition {
                     channel: Channel { term, column: 0 },
                     level,
                 });
             }
             for (j, &kept_j) in kept.iter().enumerate() {
                 if !kept_j {
-                    unidentified.push(CoefficientAddress {
+                    unidentified.push(CoefficientPosition {
                         channel: Channel {
                             term,
                             column: slope_columns[j],

crates/within/src/solver/reparam/tests.rs

@@ -77,11 +77,11 @@ fn unidentified_directions_ascend_across_terms() {
     assert_eq!(
         rp.unidentified,
         vec![
-            CoefficientAddress {
+            CoefficientPosition {
                 channel: Channel { term: 0, column: 1 },
                 level: 1,
             },
-            CoefficientAddress {
+            CoefficientPosition {
                 channel: Channel { term: 1, column: 1 },
                 level: 0,
             },