diff --git a/crates/fast-mlsirm-py/src/multilevel_bindings.rs b/crates/fast-mlsirm-py/src/multilevel_bindings.rs index f08eb8021..c81470a81 100644 --- a/crates/fast-mlsirm-py/src/multilevel_bindings.rs +++ b/crates/fast-mlsirm-py/src/multilevel_bindings.rs @@ -1,16 +1,20 @@ //! Python bindings for the sparse cross-classified multiple-membership -//! contextual-effects predictor (see `mlsirm_core::multilevel`). +//! contextual-effects predictor and crossed `u_h` MAP estimator +//! (see `mlsirm_core::multilevel`). //! -//! This module performs shape/type marshalling only. The additive sum over -//! membership-weighted context effects, its determinism across worker -//! counts, and its input validation are owned by -//! `mlsirm_core::multilevel::weighted_contextual_effect`. +//! This module performs shape/type marshalling only. Weighted summation, +//! Newton/IRLS estimation of `u_h`, CPU worker partitioning, and the optional +//! wgpu person-score kernel are owned by `mlsirm_core::multilevel`. -use mlsirm_core::multilevel::weighted_contextual_effect as core_weighted_contextual_effect; +use mlsirm_core::multilevel::{ + estimate_crossed_person_effects as core_estimate_crossed_person_effects, + weighted_contextual_effect as core_weighted_contextual_effect, CrossedPersonEffectConfig, +}; +use mlsirm_core::Device; use numpy::{PyArray1, PyReadonlyArray1, ToPyArray}; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use pyo3::types::PyModule; +use pyo3::types::{PyDict, PyModule}; use pyo3::wrap_pyfunction; // Keep the raw extension boundary aligned with the canonical Python design @@ -87,9 +91,119 @@ fn py_weighted_contextual_effect<'py>( Ok(result.to_pyarray(py)) } +/// Estimate crossed / multiple-membership person effects ``u_h``. +/// +/// Parameters +/// ---------- +/// y : numpy.ndarray[float64] +/// Row-major ``n_persons * n_items`` binary responses. Non-finite or +/// negative cells are missing. +/// row_offsets : numpy.ndarray[uint64] +/// CSR row pointer, length ``n_persons + 1``. +/// context_indices : numpy.ndarray[uint64] +/// Flattened context-effect index per membership edge. +/// weights : numpy.ndarray[float64] +/// Membership weight per edge. +/// item_slopes : numpy.ndarray[float64] +/// Known item discriminations, length ``n_items``. +/// item_intercepts : numpy.ndarray[float64] +/// Known item intercepts, length ``n_items``. +/// person_offsets : numpy.ndarray[float64] +/// Optional person-level location offsets, length ``n_persons`` or empty. +/// classification_offsets : numpy.ndarray[uint64] +/// CSR pointer over flattened classifications. +/// n_persons, n_items, n_effects : int +/// Declared problem sizes. +/// prior_precision : float +/// Gaussian prior precision ``1 / sigma_u^2``. +/// max_iter : int +/// Newton iteration budget. +/// tol : float +/// Absolute effect-step tolerance. +/// worker_count : int +/// Deterministic CPU workers. +/// device : str +/// ``cpu``, ``gpu``, or ``auto``. +/// +/// Returns +/// ------- +/// dict +/// ``effects``, ``loglik``, ``n_iter``, ``converged``, ``used_gpu``, +/// and ``termination_reason``. +#[pyfunction(name = "estimate_crossed_person_effects")] +#[allow(clippy::too_many_arguments)] +fn py_estimate_crossed_person_effects<'py>( + py: Python<'py>, + y: PyReadonlyArray1<'_, f64>, + row_offsets: PyReadonlyArray1<'_, u64>, + context_indices: PyReadonlyArray1<'_, u64>, + weights: PyReadonlyArray1<'_, f64>, + item_slopes: PyReadonlyArray1<'_, f64>, + item_intercepts: PyReadonlyArray1<'_, f64>, + person_offsets: PyReadonlyArray1<'_, f64>, + classification_offsets: PyReadonlyArray1<'_, u64>, + n_persons: usize, + n_items: usize, + n_effects: usize, + prior_precision: f64, + max_iter: usize, + tol: f64, + worker_count: usize, + device: &str, +) -> PyResult> { + let row_offsets = row_offsets.as_slice()?; + let context_indices = context_indices.as_slice()?; + if row_offsets.len() > MAX_ROW_OFFSETS { + return Err(PyValueError::new_err(format!( + "row_offsets exceeds maximum supported length of {MAX_ROW_OFFSETS}" + ))); + } + if context_indices.len() > MAX_CONTEXT_MEMBERSHIPS { + return Err(PyValueError::new_err(format!( + "context_indices exceeds maximum supported length of {MAX_CONTEXT_MEMBERSHIPS}" + ))); + } + let device = Device::parse(device) + .ok_or_else(|| PyValueError::new_err("device must be one of 'cpu', 'gpu', or 'auto'"))?; + let row_offsets = checked_usize_values(row_offsets, "row_offsets")?; + let context_indices = checked_usize_values(context_indices, "context_indices")?; + let classification_offsets = + checked_usize_values(classification_offsets.as_slice()?, "classification_offsets")?; + let result = core_estimate_crossed_person_effects( + y.as_slice()?, + &row_offsets, + &context_indices, + weights.as_slice()?, + item_slopes.as_slice()?, + item_intercepts.as_slice()?, + person_offsets.as_slice()?, + &classification_offsets, + n_persons, + n_items, + n_effects, + CrossedPersonEffectConfig { + prior_precision, + max_iter, + tol, + worker_count, + device, + }, + ) + .map_err(PyValueError::new_err)?; + let out = PyDict::new(py); + out.set_item("effects", result.effects.to_pyarray(py))?; + out.set_item("loglik", result.loglik)?; + out.set_item("n_iter", result.n_iter)?; + out.set_item("converged", result.converged)?; + out.set_item("used_gpu", result.used_gpu)?; + out.set_item("termination_reason", result.termination_reason)?; + Ok(out.into()) +} + #[pymodule] #[pyo3(name = "_multilevel_core")] fn fast_mlsirm_multilevel_core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_weighted_contextual_effect, m)?)?; + m.add_function(wrap_pyfunction!(py_estimate_crossed_person_effects, m)?)?; Ok(()) } diff --git a/crates/mlsirm-core/src/gpu_multilevel.rs b/crates/mlsirm-core/src/gpu_multilevel.rs new file mode 100644 index 000000000..805ea97f7 --- /dev/null +++ b/crates/mlsirm-core/src/gpu_multilevel.rs @@ -0,0 +1,342 @@ +//! wgpu person-score reduction for crossed / multiple-membership IRT. +//! +//! The kernel evaluates the `O(n_persons * n_items)` Bernoulli residual and +//! Fisher information that the MAP estimator of `u_h` consumes. Sparse +//! membership accumulation and the dense Newton solve remain on the CPU. +//! Kernels run in f32; the CPU path is the f64 reference. Missing adapters +//! return `None` so the estimator falls back without failing the fit. +//! +//! Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel +//! IRT model. *Psychometrika, 66*, 271-288. +//! + +use std::sync::OnceLock; + +use bytemuck::{Pod, Zeroable}; +use wgpu::util::DeviceExt; + +const WORKGROUP_SIZE: u32 = 64; + +#[repr(C)] +#[derive(Clone, Copy, Pod, Zeroable)] +struct Uniforms { + n_persons: u32, + n_items: u32, + _pad0: u32, + _pad1: u32, +} + +const SHADER: &str = r#" +struct Uniforms { + n_persons: u32, + n_items: u32, + _pad0: u32, + _pad1: u32, +}; + +@group(0) @binding(0) var U: Uniforms; +@group(0) @binding(1) var y: array; +@group(0) @binding(2) var slopes: array; +@group(0) @binding(3) var intercepts: array; +@group(0) @binding(4) var locations: array; +@group(0) @binding(5) var residual: array; +@group(0) @binding(6) var information: array; + +fn logistic(x: f32) -> f32 { + if (x >= 0.0) { + return 1.0 / (1.0 + exp(-x)); + } + let ex = exp(x); + return ex / (1.0 + ex); +} + +@compute @workgroup_size(64) +fn person_score_pass(@builtin(global_invocation_id) gid: vec3) { + let person = gid.x; + if (person >= U.n_persons) { return; } + var score = 0.0; + var weight = 0.0; + let location = locations[person]; + for (var item = 0u; item < U.n_items; item = item + 1u) { + let response = y[person * U.n_items + item]; + if (!(response == response) || response < 0.0) { continue; } + let slope = slopes[item]; + let probability = logistic(slope * location + intercepts[item]); + score = score + slope * (response - probability); + weight = weight + slope * slope * probability * (1.0 - probability); + } + residual[person] = score; + information[person] = weight; +} +"#; + +struct GpuContext { + device: wgpu::Device, + queue: wgpu::Queue, + layout: wgpu::BindGroupLayout, + pipeline: wgpu::ComputePipeline, +} + +static CONTEXT: OnceLock> = OnceLock::new(); + +fn storage_entry(binding: u32, read_only: bool) -> wgpu::BindGroupLayoutEntry { + wgpu::BindGroupLayoutEntry { + binding, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Storage { read_only }, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + } +} + +impl GpuContext { + fn init() -> Option { + let instance = crate::gpu_init::new_instance(); + let adapter = + pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions::default())) + .ok()?; + let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { + label: Some("mlsirm-crossed-person-effects"), + required_limits: adapter.limits(), + ..Default::default() + })) + .ok()?; + let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { + label: Some("mlsirm-crossed-person-effects"), + source: wgpu::ShaderSource::Wgsl(SHADER.into()), + }); + let mut entries = vec![wgpu::BindGroupLayoutEntry { + binding: 0, + visibility: wgpu::ShaderStages::COMPUTE, + ty: wgpu::BindingType::Buffer { + ty: wgpu::BufferBindingType::Uniform, + has_dynamic_offset: false, + min_binding_size: None, + }, + count: None, + }]; + for binding in 1..=6u32 { + entries.push(storage_entry(binding, binding <= 4)); + } + let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor { + label: Some("mlsirm-crossed-person-effects-layout"), + entries: &entries, + }); + let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { + label: Some("mlsirm-crossed-person-effects-pipeline-layout"), + bind_group_layouts: &[Some(&layout)], + immediate_size: 0, + }); + let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { + label: Some("person_score_pass"), + layout: Some(&pipeline_layout), + module: &module, + entry_point: Some("person_score_pass"), + compilation_options: wgpu::PipelineCompilationOptions::default(), + cache: None, + }); + Some(Self { + device, + queue, + layout, + pipeline, + }) + } + + fn get() -> Option<&'static Self> { + CONTEXT.get_or_init(Self::init).as_ref() + } +} + +fn as_f32_finite(values: &[f64]) -> Option> { + values + .iter() + .map(|&value| { + let converted = value as f32; + converted.is_finite().then_some(converted) + }) + .collect() +} + +fn as_f32_responses(values: &[f64]) -> Option> { + values + .iter() + .map(|&value| { + if value.is_infinite() { + return None; + } + Some(value as f32) + }) + .collect() +} + +fn storage(device: &wgpu::Device, data: &[u8], usage: wgpu::BufferUsages) -> wgpu::Buffer { + device.create_buffer_init(&wgpu::util::BufferInitDescriptor { + label: None, + contents: data, + usage, + }) +} + +fn output(device: &wgpu::Device, len: usize) -> wgpu::Buffer { + device.create_buffer(&wgpu::BufferDescriptor { + label: Some("crossed-person-effect-output"), + size: (len.max(1) * std::mem::size_of::()) as u64, + usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, + mapped_at_creation: false, + }) +} + +fn read_f32( + device: &wgpu::Device, + queue: &wgpu::Queue, + source: &wgpu::Buffer, + len: usize, +) -> Option> { + let size = (len.max(1) * std::mem::size_of::()) as u64; + let readback = device.create_buffer(&wgpu::BufferDescriptor { + label: Some("crossed-person-effect-readback"), + size, + usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + let mut encoder = device.create_command_encoder(&Default::default()); + encoder.copy_buffer_to_buffer(source, 0, &readback, 0, size); + queue.submit([encoder.finish()]); + readback.slice(..).map_async(wgpu::MapMode::Read, |_| {}); + device.poll(wgpu::PollType::wait_indefinitely()).ok()?; + let view = readback.slice(..).get_mapped_range().ok()?; + let values: &[f32] = bytemuck::cast_slice(&view); + let result = values.iter().take(len).map(|&value| value as f64).collect(); + drop(view); + readback.unmap(); + Some(result) +} + +/// Reduce per-person Bernoulli scores on a usable GPU. +/// +/// Returns `None` when no adapter is present, a buffer exceeds device limits, +/// or f32 conversion of a finite f64 input overflows. +pub(crate) fn person_irt_scores_gpu( + y: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + locations: &[f64], + n_persons: usize, + n_items: usize, +) -> Option<(Vec, Vec)> { + let context = GpuContext::get()?; + let limits = context.device.limits(); + let y_len = n_persons.checked_mul(n_items)?; + if y.len() != y_len || locations.len() != n_persons { + return None; + } + for &len in &[y_len, n_items, n_persons] { + let bytes = len.checked_mul(std::mem::size_of::())?; + if bytes as u64 > limits.max_buffer_size + || bytes > limits.max_storage_buffer_binding_size as usize + { + return None; + } + } + let y_f32 = as_f32_responses(y)?; + let slopes = as_f32_finite(item_slopes)?; + let intercepts = as_f32_finite(item_intercepts)?; + let locations_f32 = as_f32_finite(locations)?; + let uniforms = Uniforms { + n_persons: u32::try_from(n_persons).ok()?, + n_items: u32::try_from(n_items).ok()?, + _pad0: 0, + _pad1: 0, + }; + let uniform_buf = storage( + &context.device, + bytemuck::bytes_of(&uniforms), + wgpu::BufferUsages::UNIFORM, + ); + let y_buf = storage( + &context.device, + bytemuck::cast_slice(&y_f32), + wgpu::BufferUsages::STORAGE, + ); + let slope_buf = storage( + &context.device, + bytemuck::cast_slice(&slopes), + wgpu::BufferUsages::STORAGE, + ); + let intercept_buf = storage( + &context.device, + bytemuck::cast_slice(&intercepts), + wgpu::BufferUsages::STORAGE, + ); + let location_buf = storage( + &context.device, + bytemuck::cast_slice(&locations_f32), + wgpu::BufferUsages::STORAGE, + ); + let residual_buf = output(&context.device, n_persons); + let information_buf = output(&context.device, n_persons); + let bind_group = context + .device + .create_bind_group(&wgpu::BindGroupDescriptor { + label: Some("crossed-person-effect-bind-group"), + layout: &context.layout, + entries: &[ + wgpu::BindGroupEntry { + binding: 0, + resource: uniform_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 1, + resource: y_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 2, + resource: slope_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 3, + resource: intercept_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 4, + resource: location_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 5, + resource: residual_buf.as_entire_binding(), + }, + wgpu::BindGroupEntry { + binding: 6, + resource: information_buf.as_entire_binding(), + }, + ], + }); + let groups = n_persons.div_ceil(WORKGROUP_SIZE as usize) as u32; + let mut encoder = context + .device + .create_command_encoder(&wgpu::CommandEncoderDescriptor { + label: Some("crossed-person-effect-encoder"), + }); + { + let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor { + label: Some("person_score_pass"), + timestamp_writes: None, + }); + pass.set_pipeline(&context.pipeline); + pass.set_bind_group(0, &bind_group, &[]); + pass.dispatch_workgroups(groups.max(1), 1, 1); + } + context.queue.submit([encoder.finish()]); + let residual = read_f32(&context.device, &context.queue, &residual_buf, n_persons)?; + let information = read_f32(&context.device, &context.queue, &information_buf, n_persons)?; + if residual.iter().any(|value| !value.is_finite()) + || information.iter().any(|value| !value.is_finite()) + { + return None; + } + Some((residual, information)) +} diff --git a/crates/mlsirm-core/src/lib.rs b/crates/mlsirm-core/src/lib.rs index f2719a800..b80f58991 100644 --- a/crates/mlsirm-core/src/lib.rs +++ b/crates/mlsirm-core/src/lib.rs @@ -85,6 +85,8 @@ pub(crate) mod gpu_marginal; pub(crate) mod gpu_plausible; #[cfg(all(feature = "gpu", not(coverage)))] pub(crate) mod gpu_scoring; +#[cfg(all(feature = "gpu", not(coverage)))] +pub(crate) mod gpu_multilevel; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum ModelType { Mirt, diff --git a/crates/mlsirm-core/src/multilevel.rs b/crates/mlsirm-core/src/multilevel.rs index e4ee4a96c..a73973785 100644 --- a/crates/mlsirm-core/src/multilevel.rs +++ b/crates/mlsirm-core/src/multilevel.rs @@ -9,15 +9,31 @@ //! design. Equal-key ordering therefore cannot affect a public result and finite //! input overflow cannot silently become an infinite contextual contribution. //! +//! [`estimate_crossed_person_effects`] is the MAP estimator of the same `u_h` +//! vector for crossed and weighted multiple-membership designs (Fox & Glas, +//! 2001; Browne, Goldstein, & Rasbash, 2001). It does not rewrite the +//! longitudinal OLS / AR utilities reserved for a separate slice. +//! //! Browne, W. J., Goldstein, H., & Rasbash, J. (2001). Multiple membership //! multiple classification (MMMC) models. *Statistical Modelling, 1*(2), //! 103-124. +//! +//! Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel +//! IRT model. *Psychometrika, 66*, 271-288. +//! use std::collections::HashSet; +#[path = "multilevel_estimator.rs"] +mod estimator; #[path = "multilevel_kernel.rs"] mod kernel; +pub use estimator::{ + estimate_crossed_person_effects, CrossedPersonEffectConfig, CrossedPersonEffectEstimate, + MAX_CROSSED_EFFECTS, MAX_CROSSED_ITER, +}; + fn validate_unique_context_indices_per_row( row_offsets: &[usize], context_indices: &[usize], diff --git a/crates/mlsirm-core/src/multilevel_estimator.rs b/crates/mlsirm-core/src/multilevel_estimator.rs new file mode 100644 index 000000000..71e8fa090 --- /dev/null +++ b/crates/mlsirm-core/src/multilevel_estimator.rs @@ -0,0 +1,933 @@ +//! MAP estimator for crossed / multiple-membership person effects `u_h`. +//! +//! This module owns the first buyer-visible random-effect *estimator* for the +//! contextual term already evaluated by [`crate::multilevel::weighted_contextual_effect`]. +//! Persons may belong to several units of one classification at once (weighted +//! multiple membership) and to several classifications at once (crossed / +//! multiple-classification designs). Ordinary one-hot nesting is the singleton +//! special case of the same sparse design. +//! +//! # Linear predictor +//! +//! For binary response `Y_pi` from person `p` and item `i`, +//! +//! ```text +//! eta_pi = a_i * (theta_p + sum_h w_ph * u_h) + b_i +//! ``` +//! +//! with known item slopes `a_i > 0`, known item intercepts `b_i`, known +//! non-negative membership weights that already sum to one within each +//! classification (Browne, Goldstein, & Rasbash, 2001, eq. 1), and an optional +//! person-level offset `theta_p`. The offset is the time-flow compatibility +//! hook: a longitudinal layer may supply already-estimated occasion states. +//! This slice does **not** estimate OLS trends or AR coefficients. +//! +//! # Estimand and identification +//! +//! Fox and Glas (2001) place a level-2 Gaussian prior on the group effects of +//! a multilevel IRT model. This kernel is the matching MAP / ridge point +//! estimator of the flattened effects `u_h`, not their Gibbs sampler: +//! +//! ```text +//! u_h ~ N(0, sigma_u^2) (independent, prior_precision = 1 / sigma_u^2) +//! ``` +//! +//! The reported estimate is re-centered to sum to zero inside each +//! classification so recovered effects are deviations, matching the usual +//! multilevel IRT location constraint. Item intercepts absorb the global +//! location. A classification with fewer than two levels is rejected because +//! centering would leave a non-identified singleton. +//! +//! The estimator is **not** a claim of Fox & Glas MCMC, Jeon & Rabe-Hesketh +//! adaptive quadrature, variance-component ML, or causal contextual effects. +//! +//! # Compute +//! +//! The `O(n_persons * n_items)` Bernoulli score / information reduction is +//! multithreaded on CPU and, when a wgpu adapter is present, offloaded to an +//! f32 GPU kernel with an f64 CPU fallback. The sparse-design Newton system +//! (`n_effects <= 128`) stays on CPU. Results do not depend on `worker_count`. +//! +//! # References +//! +//! Browne, W. J., Goldstein, H., & Rasbash, J. (2001). Multiple membership +//! multiple classification (MMMC) models. *Statistical Modelling, 1*(2), +//! 103-124. +//! +//! Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel +//! IRT model. *Psychometrika, 66*, 271-288. +//! + +use std::thread; + +use crate::mmle::{log_sigmoid, sigmoid_stable}; +use crate::multilevel::weighted_contextual_effect; +use crate::Device; + +/// Hard cap on flattened context effects for the dense Newton system. +pub const MAX_CROSSED_EFFECTS: usize = 128; + +/// Upper bound on Newton iterations accepted from a caller. +pub const MAX_CROSSED_ITER: usize = 10_000; + +type EstimatorResult = Result; + +/// Configuration for the crossed / multiple-membership MAP estimator. +#[derive(Clone, Copy, Debug)] +pub struct CrossedPersonEffectConfig { + /// Gaussian prior precision `1 / sigma_u^2` (Fox & Glas, 2001, level-2). + pub prior_precision: f64, + /// Maximum Newton / IRLS iterations. + pub max_iter: usize, + /// Absolute effect-step convergence tolerance. + pub tol: f64, + /// Deterministic CPU worker count (`>= 1`); does not change the result. + pub worker_count: usize, + /// CPU / GPU / auto device policy for the person-score reduction. + pub device: Device, +} + +impl Default for CrossedPersonEffectConfig { + fn default() -> Self { + Self { + prior_precision: 1.0, + max_iter: 50, + tol: 1e-8, + worker_count: 1, + device: Device::Auto, + } + } +} + +/// MAP estimate of crossed / multiple-membership person effects `u_h`. +#[derive(Clone, Debug, PartialEq)] +pub struct CrossedPersonEffectEstimate { + /// Centered context effects, length `n_effects`. + pub effects: Vec, + /// Observed-data Bernoulli log-likelihood plus Gaussian prior penalty. + pub loglik: f64, + /// Newton iterations actually performed. + pub n_iter: usize, + /// Whether the last step satisfied `tol`. + pub converged: bool, + /// Whether the person-score reduction used the wgpu kernel. + pub used_gpu: bool, + /// Machine-readable termination status. + pub termination_reason: String, +} + +/// Estimate crossed / multiple-membership `u_h` by Gaussian-prior MAP. +/// +/// `y` is row-major `n_persons * n_items`. Finite non-negative observed cells +/// must be exactly `0` or `1`; a cell is treated as missing when it is +/// non-finite or strictly negative (the established `NaN` / `-1` mask +/// contract). `item_slopes` and `item_intercepts` have length `n_items`. +/// `person_offsets` is either empty (treated as zeros) or length `n_persons`. +/// `classification_offsets` is a CSR pointer over the flattened effect table: +/// classification `d` occupies `classification_offsets[d]` +/// .. classification_offsets[d + 1]`. +#[allow(clippy::too_many_arguments)] +pub fn estimate_crossed_person_effects( + y: &[f64], + row_offsets: &[usize], + context_indices: &[usize], + weights: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + person_offsets: &[f64], + classification_offsets: &[usize], + n_persons: usize, + n_items: usize, + n_effects: usize, + config: CrossedPersonEffectConfig, +) -> EstimatorResult { + validate_estimator_inputs( + y, + row_offsets, + context_indices, + weights, + item_slopes, + item_intercepts, + person_offsets, + classification_offsets, + n_persons, + n_items, + n_effects, + &config, + )?; + + let mut effects = vec![0.0_f64; n_effects]; + let mut used_gpu = false; + let mut last_step = f64::INFINITY; + let mut n_iter = 0usize; + + for iteration in 1..=config.max_iter { + n_iter = iteration; + let locations = person_locations( + row_offsets, + context_indices, + weights, + &effects, + person_offsets, + n_persons, + config.worker_count, + )?; + let (residual, information, gpu_hit) = person_scores( + y, + item_slopes, + item_intercepts, + &locations, + n_persons, + n_items, + config.worker_count, + config.device, + )?; + used_gpu = used_gpu || gpu_hit; + let rhs = effect_score( + row_offsets, + context_indices, + weights, + &residual, + &effects, + config.prior_precision, + n_effects, + ); + let mut system = effect_system( + row_offsets, + context_indices, + weights, + &information, + config.prior_precision, + n_effects, + ); + let delta = solve_dense_system(&mut system, &rhs, n_effects)?; + last_step = delta.iter().fold(0.0_f64, |acc, &step| acc.max(step.abs())); + for (effect, step) in effects.iter_mut().zip(delta.iter()) { + *effect += *step; + } + if last_step < config.tol { + break; + } + } + center_classifications(&mut effects, classification_offsets); + + let locations = person_locations( + row_offsets, + context_indices, + weights, + &effects, + person_offsets, + n_persons, + config.worker_count, + )?; + let loglik = bernoulli_map_loglik( + y, + item_slopes, + item_intercepts, + &locations, + &effects, + config.prior_precision, + n_persons, + n_items, + ); + if !loglik.is_finite() { + return Err("crossed person-effect log-likelihood must be finite".to_string()); + } + let converged = last_step < config.tol; + Ok(CrossedPersonEffectEstimate { + effects, + loglik, + n_iter, + converged, + used_gpu, + termination_reason: if converged { + "converged".to_string() + } else { + "max_iter_reached".to_string() + }, + }) +} + +#[allow(clippy::too_many_arguments)] +fn validate_estimator_inputs( + y: &[f64], + row_offsets: &[usize], + context_indices: &[usize], + weights: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + person_offsets: &[f64], + classification_offsets: &[usize], + n_persons: usize, + n_items: usize, + n_effects: usize, + config: &CrossedPersonEffectConfig, +) -> EstimatorResult<()> { + if n_persons < 1 || n_items < 1 || n_effects < 1 { + return Err("n_persons, n_items, and n_effects must be at least one".to_string()); + } + if n_effects > MAX_CROSSED_EFFECTS { + return Err(format!( + "n_effects exceeds the dense Newton cap of {MAX_CROSSED_EFFECTS}" + )); + } + let expected = crate::checked_mul_usize(n_persons, n_items, "response matrix is too large")?; + if y.len() != expected { + return Err("y must have length n_persons * n_items".to_string()); + } + for &response in y { + if response.is_finite() && response >= 0.0 && response != 0.0 && response != 1.0 { + return Err("binary responses must contain only 0 or 1 for observed cells".to_string()); + } + } + if item_slopes.len() != n_items || item_intercepts.len() != n_items { + return Err("item_slopes and item_intercepts must have length n_items".to_string()); + } + if !person_offsets.is_empty() && person_offsets.len() != n_persons { + return Err("person_offsets must be empty or have length n_persons".to_string()); + } + for &slope in item_slopes { + if !slope.is_finite() || slope <= 0.0 { + return Err("item_slopes must be finite and strictly positive".to_string()); + } + } + for &intercept in item_intercepts { + if !intercept.is_finite() { + return Err("item_intercepts must be finite".to_string()); + } + } + for &offset in person_offsets { + if !offset.is_finite() { + return Err("person_offsets must be finite".to_string()); + } + } + if classification_offsets.len() < 2 { + return Err("classification_offsets must contain at least one classification".to_string()); + } + if classification_offsets[0] != 0 + || *classification_offsets.last().expect("non-empty") != n_effects + || classification_offsets + .windows(2) + .any(|window| window[1] <= window[0]) + { + return Err( + "classification_offsets must start at zero, increase strictly, and end at n_effects" + .to_string(), + ); + } + for window in classification_offsets.windows(2) { + if window[1] - window[0] < 2 { + return Err("each classification must contain at least two context levels".to_string()); + } + } + if !config.prior_precision.is_finite() || config.prior_precision <= 0.0 { + return Err("prior_precision must be finite and strictly positive".to_string()); + } + if !(1..=MAX_CROSSED_ITER).contains(&config.max_iter) { + return Err(format!("max_iter must be in 1..={MAX_CROSSED_ITER}")); + } + if !config.tol.is_finite() || config.tol <= 0.0 { + return Err("tol must be finite and strictly positive".to_string()); + } + if config.worker_count == 0 { + return Err("worker_count must be at least one".to_string()); + } + if row_offsets.len() != n_persons + 1 { + return Err("row_offsets must have length n_persons + 1".to_string()); + } + // Touch the public predictor boundary so malformed CSR / non-finite + // weights / duplicate row indices fail with the established messages + // before Newton iteration begins. + let dummy = vec![0.0_f64; n_effects]; + weighted_contextual_effect( + row_offsets, + context_indices, + weights, + &dummy, + config.worker_count, + )?; + Ok(()) +} + +fn person_locations( + row_offsets: &[usize], + context_indices: &[usize], + weights: &[f64], + effects: &[f64], + person_offsets: &[f64], + n_persons: usize, + worker_count: usize, +) -> EstimatorResult> { + let mut locations = + weighted_contextual_effect(row_offsets, context_indices, weights, effects, worker_count)?; + if locations.len() != n_persons { + return Err("contextual predictor length must match n_persons".to_string()); + } + if !person_offsets.is_empty() { + for (location, offset) in locations.iter_mut().zip(person_offsets.iter()) { + *location += *offset; + } + } + Ok(locations) +} + +#[allow(clippy::too_many_arguments)] +fn person_scores( + y: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + locations: &[f64], + n_persons: usize, + n_items: usize, + worker_count: usize, + device: Device, +) -> EstimatorResult<(Vec, Vec, bool)> { + if device != Device::Cpu { + if let Some((residual, information)) = try_person_scores_gpu( + y, + item_slopes, + item_intercepts, + locations, + n_persons, + n_items, + ) { + return Ok((residual, information, true)); + } + if device == Device::Gpu { + eprintln!( + "fast-mlsirm: GPU crossed person-effect scores requested but no usable GPU adapter was found; falling back to the CPU implementation." + ); + } + } + let (residual, information) = person_scores_cpu( + y, + item_slopes, + item_intercepts, + locations, + n_persons, + n_items, + worker_count, + ); + Ok((residual, information, false)) +} + +#[cfg(all(feature = "gpu", not(coverage)))] +fn try_person_scores_gpu( + y: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + locations: &[f64], + n_persons: usize, + n_items: usize, +) -> Option<(Vec, Vec)> { + crate::gpu_multilevel::person_irt_scores_gpu( + y, + item_slopes, + item_intercepts, + locations, + n_persons, + n_items, + ) +} + +#[cfg(any(not(feature = "gpu"), coverage))] +fn try_person_scores_gpu( + _y: &[f64], + _item_slopes: &[f64], + _item_intercepts: &[f64], + _locations: &[f64], + _n_persons: usize, + _n_items: usize, +) -> Option<(Vec, Vec)> { + None +} + +fn person_scores_cpu( + y: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + locations: &[f64], + n_persons: usize, + n_items: usize, + worker_count: usize, +) -> (Vec, Vec) { + let mut residual = vec![0.0_f64; n_persons]; + let mut information = vec![0.0_f64; n_persons]; + if n_persons == 0 { + return (residual, information); + } + let worker_count = worker_count.min(n_persons); + let chunk_size = n_persons.div_ceil(worker_count); + thread::scope(|scope| { + for (chunk_index, (residual_chunk, information_chunk)) in residual + .chunks_mut(chunk_size) + .zip(information.chunks_mut(chunk_size)) + .enumerate() + { + let start = chunk_index * chunk_size; + scope.spawn(move || { + for (offset, (residual_slot, information_slot)) in residual_chunk + .iter_mut() + .zip(information_chunk.iter_mut()) + .enumerate() + { + let person = start + offset; + let (person_residual, person_information) = person_score_one( + &y[person * n_items..(person + 1) * n_items], + item_slopes, + item_intercepts, + locations[person], + ); + *residual_slot = person_residual; + *information_slot = person_information; + } + }); + } + }); + (residual, information) +} + +fn person_score_one( + row: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + location: f64, +) -> (f64, f64) { + let mut residual = 0.0_f64; + let mut information = 0.0_f64; + for (item, &response) in row.iter().enumerate() { + if !response_is_observed(response) { + continue; + } + let slope = item_slopes[item]; + let probability = sigmoid_stable(slope * location + item_intercepts[item]); + residual += slope * (response - probability); + information += slope * slope * probability * (1.0 - probability); + } + (residual, information) +} + +fn response_is_observed(response: f64) -> bool { + response.is_finite() && response >= 0.0 +} + +fn effect_score( + row_offsets: &[usize], + context_indices: &[usize], + weights: &[f64], + residual: &[f64], + effects: &[f64], + prior_precision: f64, + n_effects: usize, +) -> Vec { + let mut score = vec![0.0_f64; n_effects]; + for (person, window) in row_offsets.windows(2).enumerate() { + for edge in window[0]..window[1] { + score[context_indices[edge]] += weights[edge] * residual[person]; + } + } + for (slot, effect) in score.iter_mut().zip(effects.iter()) { + *slot -= prior_precision * *effect; + } + score +} + +fn effect_system( + row_offsets: &[usize], + context_indices: &[usize], + weights: &[f64], + information: &[f64], + prior_precision: f64, + n_effects: usize, +) -> Vec { + let mut system = vec![0.0_f64; n_effects * n_effects]; + for (person, window) in row_offsets.windows(2).enumerate() { + let weight_p = information[person]; + if weight_p == 0.0 { + continue; + } + for left in window[0]..window[1] { + let left_index = context_indices[left]; + let left_weight = weights[left]; + for right in window[0]..window[1] { + system[left_index * n_effects + context_indices[right]] += + left_weight * weights[right] * weight_p; + } + } + } + for effect in 0..n_effects { + system[effect * n_effects + effect] += prior_precision; + } + system +} + +fn solve_dense_system(matrix: &mut [f64], rhs: &[f64], n: usize) -> EstimatorResult> { + if n == 0 { + return Ok(Vec::new()); + } + let mut augmented = vec![0.0_f64; n * (n + 1)]; + for row in 0..n { + let dest = row * (n + 1); + augmented[dest..dest + n].copy_from_slice(&matrix[row * n..(row + 1) * n]); + augmented[dest + n] = rhs[row]; + } + for col in 0..n { + let mut pivot_row = col; + let mut pivot_abs = augmented[col * (n + 1) + col].abs(); + for row in (col + 1)..n { + let candidate = augmented[row * (n + 1) + col].abs(); + if candidate > pivot_abs { + pivot_abs = candidate; + pivot_row = row; + } + } + if pivot_abs < 1e-14 { + return Err("crossed person-effect Newton system is singular".to_string()); + } + if pivot_row != col { + for slot in 0..=n { + augmented.swap(col * (n + 1) + slot, pivot_row * (n + 1) + slot); + } + } + let pivot = augmented[col * (n + 1) + col]; + for row in (col + 1)..n { + let factor = augmented[row * (n + 1) + col] / pivot; + for slot in col..=n { + let value = augmented[col * (n + 1) + slot]; + augmented[row * (n + 1) + slot] -= factor * value; + } + } + } + let mut solution = vec![0.0_f64; n]; + for row in (0..n).rev() { + let mut value = augmented[row * (n + 1) + n]; + for col in (row + 1)..n { + value -= augmented[row * (n + 1) + col] * solution[col]; + } + let pivot = augmented[row * (n + 1) + row]; + if !pivot.is_finite() || pivot.abs() < 1e-14 { + return Err("crossed person-effect Newton system is singular".to_string()); + } + solution[row] = value / pivot; + if !solution[row].is_finite() { + return Err("crossed person-effect Newton step must be finite".to_string()); + } + } + Ok(solution) +} + +fn center_classifications(effects: &mut [f64], classification_offsets: &[usize]) { + for window in classification_offsets.windows(2) { + let slice = &mut effects[window[0]..window[1]]; + let mean = slice.iter().sum::() / slice.len() as f64; + for value in slice { + *value -= mean; + } + } +} + +fn bernoulli_map_loglik( + y: &[f64], + item_slopes: &[f64], + item_intercepts: &[f64], + locations: &[f64], + effects: &[f64], + prior_precision: f64, + n_persons: usize, + n_items: usize, +) -> f64 { + let mut loglik = 0.0_f64; + for person in 0..n_persons { + let location = locations[person]; + for item in 0..n_items { + let response = y[person * n_items + item]; + if !response_is_observed(response) { + continue; + } + let eta = item_slopes[item] * location + item_intercepts[item]; + loglik += if response >= 0.5 { + log_sigmoid(eta) + } else { + log_sigmoid(-eta) + }; + } + } + loglik -= 0.5 * prior_precision * effects.iter().map(|value| value * value).sum::(); + loglik +} + +#[cfg(test)] +mod tests { + use super::*; + + fn crossed_design() -> (Vec, Vec, Vec, Vec) { + // Four persons, two classifications (school, neighborhood), each with + // two levels. Person 1 has weighted multiple school membership. + // Flattened effects: school_0, school_1, neigh_0, neigh_1. + let row_offsets = vec![0, 2, 5, 7, 9]; + let context_indices = vec![0, 2, 0, 1, 3, 1, 2, 1, 3]; + let weights = vec![1.0, 1.0, 0.6, 0.4, 1.0, 1.0, 1.0, 1.0, 1.0]; + let classification_offsets = vec![0, 2, 4]; + ( + row_offsets, + context_indices, + weights, + classification_offsets, + ) + } + + fn simulate_responses( + row_offsets: &[usize], + context_indices: &[usize], + weights: &[f64], + true_effects: &[f64], + intercepts: &[f64], + n_persons: usize, + n_items: usize, + seed: u64, + ) -> Vec { + let locations = + weighted_contextual_effect(row_offsets, context_indices, weights, true_effects, 1) + .unwrap(); + let mut y = vec![0.0_f64; n_persons * n_items]; + let mut state = seed; + for person in 0..n_persons { + for item in 0..n_items { + let probability = sigmoid_stable(locations[person] + intercepts[item]); + state = state.wrapping_mul(6364136223846793005).wrapping_add(1); + let unit = ((state >> 11) as f64) / ((1u64 << 53) as f64); + y[person * n_items + item] = if unit < probability { 1.0 } else { 0.0 }; + } + } + y + } + + #[test] + fn recovers_crossed_membership_effects_below_rmse_gate() { + let (row_offsets, context_indices, weights, classification_offsets) = crossed_design(); + let true_effects = vec![-0.8, 0.8, -0.5, 0.5]; + let n_persons = 4; + let n_items = 80; + let intercepts: Vec = (0..n_items) + .map(|item| -1.2 + 2.4 * (item as f64) / ((n_items - 1) as f64)) + .collect(); + let slopes = vec![1.0; n_items]; + // Replicate the four-person design many times so each context has a + // large effective sample while preserving the crossed weights. + let repeats = 24usize; + let mut row_offsets_rep = vec![0usize]; + let mut context_indices_rep = Vec::new(); + let mut weights_rep = Vec::new(); + for _ in 0..repeats { + for person in 0..n_persons { + let start = row_offsets[person]; + let end = row_offsets[person + 1]; + context_indices_rep.extend_from_slice(&context_indices[start..end]); + weights_rep.extend_from_slice(&weights[start..end]); + row_offsets_rep.push(context_indices_rep.len()); + } + } + let n_rep = n_persons * repeats; + let y = simulate_responses( + &row_offsets_rep, + &context_indices_rep, + &weights_rep, + &true_effects, + &intercepts, + n_rep, + n_items, + 20260818, + ); + let estimate = estimate_crossed_person_effects( + &y, + &row_offsets_rep, + &context_indices_rep, + &weights_rep, + &slopes, + &intercepts, + &[], + &classification_offsets, + n_rep, + n_items, + 4, + CrossedPersonEffectConfig { + prior_precision: 0.05, + max_iter: 40, + tol: 1e-8, + worker_count: 4, + device: Device::Cpu, + }, + ) + .unwrap(); + assert!(estimate.converged, "{}", estimate.termination_reason); + let mse: f64 = estimate + .effects + .iter() + .zip(true_effects.iter()) + .map(|(hat, truth)| (hat - truth) * (hat - truth)) + .sum::() + / 4.0; + let rmse = mse.sqrt(); + assert!( + rmse < 0.20, + "RMSE {rmse} exceeded the crossed u_h recovery gate; hats={:?} truth={:?}", + estimate.effects, + true_effects + ); + } + + #[test] + fn worker_count_does_not_change_the_estimate() { + let (row_offsets, context_indices, weights, classification_offsets) = crossed_design(); + let intercepts = vec![-0.5, 0.0, 0.5, 1.0]; + let slopes = vec![1.0; 4]; + let y = vec![ + 1.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0, + ]; + let config = |worker_count| CrossedPersonEffectConfig { + worker_count, + device: Device::Cpu, + ..CrossedPersonEffectConfig::default() + }; + let one = estimate_crossed_person_effects( + &y, + &row_offsets, + &context_indices, + &weights, + &slopes, + &intercepts, + &[], + &classification_offsets, + 4, + 4, + 4, + config(1), + ) + .unwrap(); + let four = estimate_crossed_person_effects( + &y, + &row_offsets, + &context_indices, + &weights, + &slopes, + &intercepts, + &[], + &classification_offsets, + 4, + 4, + 4, + config(4), + ) + .unwrap(); + assert_eq!(one.effects, four.effects); + assert_eq!(one.loglik, four.loglik); + } + + #[test] + fn rejects_singleton_classification() { + let error = estimate_crossed_person_effects( + &[1.0], + &[0, 1], + &[0], + &[1.0], + &[1.0], + &[0.0], + &[], + &[0, 1], + 1, + 1, + 1, + CrossedPersonEffectConfig::default(), + ) + .unwrap_err(); + assert!(error.contains("at least two context levels")); + } + + #[test] + fn rejects_nonbinary_observed_response() { + let error = estimate_crossed_person_effects( + &[0.5, 0.0], + &[0, 1, 2], + &[0, 1], + &[1.0, 1.0], + &[1.0], + &[0.0], + &[], + &[0, 2], + 2, + 1, + 2, + CrossedPersonEffectConfig::default(), + ) + .unwrap_err(); + assert!(error.contains("binary responses")); + } + + #[test] + fn rejects_zero_worker_count() { + let error = estimate_crossed_person_effects( + &[1.0, 0.0, 1.0, 0.0], + &[0, 1, 2], + &[0, 1], + &[1.0, 1.0], + &[1.0, 1.0], + &[0.0, 0.0], + &[], + &[0, 2], + 2, + 2, + 2, + CrossedPersonEffectConfig { + worker_count: 0, + ..CrossedPersonEffectConfig::default() + }, + ) + .unwrap_err(); + assert!(error.contains("worker_count")); + } + + #[test] + fn person_offsets_shift_the_linear_predictor() { + let (row_offsets, context_indices, weights, classification_offsets) = crossed_design(); + let intercepts = vec![0.0, 0.0]; + let slopes = vec![1.0, 1.0]; + let y = vec![1.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]; + let baseline = estimate_crossed_person_effects( + &y, + &row_offsets, + &context_indices, + &weights, + &slopes, + &intercepts, + &[], + &classification_offsets, + 4, + 2, + 4, + CrossedPersonEffectConfig { + device: Device::Cpu, + ..CrossedPersonEffectConfig::default() + }, + ) + .unwrap(); + let shifted = estimate_crossed_person_effects( + &y, + &row_offsets, + &context_indices, + &weights, + &slopes, + &intercepts, + &[2.0, 2.0, 2.0, 2.0], + &classification_offsets, + 4, + 2, + 4, + CrossedPersonEffectConfig { + device: Device::Cpu, + ..CrossedPersonEffectConfig::default() + }, + ) + .unwrap(); + assert_ne!(baseline.effects, shifted.effects); + } +} diff --git a/docs/GOVERNANCE_INDEX.md b/docs/GOVERNANCE_INDEX.md index db2edfda2..1061cfcd0 100644 --- a/docs/GOVERNANCE_INDEX.md +++ b/docs/GOVERNANCE_INDEX.md @@ -84,7 +84,7 @@ STRIDE focus for this package: | Requirement | Design | Code | Test | | --- | --- | --- | --- | | MLS2PLM point estimate | PRD formula contract | `crates/mlsirm-core`, `fit.py` | recovery RMSE tests | -| Multilevel nesting | MMLE design + Fox & Glas (2001) | `PopulationSpec::Multilevel` | multilevel recovery / contracts | +| Multilevel nesting | MMLE design + Fox & Glas (2001) | `PopulationSpec::Multilevel` + `estimate_crossed_person_effects` | multilevel recovery / contracts | | Temporal occasions | Longitudinal contracts RFC | `TemporalOccasion` | `tests/test_multilevel_*.py` | | Python 3.14 support | ADR-004 | `.github/workflows/ci.yml` | `tests/test_ci_python_314_contract.py` | @@ -98,7 +98,7 @@ flowchart TB end subgraph py [Python package] API[fit config io scoring] - ML[multilevel contracts] + ML[multilevel contracts and crossed u_h] VAL[fail-closed validators] end subgraph rust [Rust crates] diff --git a/docs/adr/0007-multilevel-multiple-membership-temporal.md b/docs/adr/0007-multilevel-multiple-membership-temporal.md index 1b65624d4..e3d451797 100644 --- a/docs/adr/0007-multilevel-multiple-membership-temporal.md +++ b/docs/adr/0007-multilevel-multiple-membership-temporal.md @@ -7,7 +7,7 @@ Date: 2026-08-09 Psychometric and AI-evaluation observations commonly sit inside schools, teams, organizations, prompts, testlets, documents, clients, time periods or other overlapping contexts. Repeated observations also evolve over time. Flattening those structures into independent rows can produce atomistic fallacy, understate uncertainty, confound stable traits with context effects and drift, and misinterpret temporal dependence. -A current open PR contains reusable contract work for nested, cross-classified, multiple-membership and longitudinal designs, but it is not yet protected-integrated. Numerical estimators for the full structures are not accepted production behavior. Therefore this ADR remains Proposed. +Reusable nested, cross-classified, multiple-membership, and longitudinal *contracts* are on protected main. A Rust MAP estimator now recovers crossed / weighted multiple-membership person effects `u_h` with RMSE evidence. OLS/AR longitudinal state estimation and MCMC variance-component engines remain separate slices, so this ADR stays Proposed until those numerical release-rule items are also evidenced. ## Decision diff --git a/docs/changelog.d/565-crossed-multiple-membership-uh.md b/docs/changelog.d/565-crossed-multiple-membership-uh.md new file mode 100644 index 000000000..f54b9977a --- /dev/null +++ b/docs/changelog.d/565-crossed-multiple-membership-uh.md @@ -0,0 +1,8 @@ +# Crossed multiple-membership person effects + +## Added + +- Added a Rust-owned MAP estimator of crossed / weighted multiple-membership person effects `u_h` (Fox & Glas, 2001; Browne, Goldstein, & Rasbash, 2001). Persons may belong to several groups at once; one-hot nesting remains the singleton special case of the same sparse design. +- Added a CPU-multithreaded Bernoulli score/information reduction and an optional wgpu GPU kernel for that hot loop, with f64 CPU fallback when no adapter is present. Sparse Newton accumulation stays on CPU. This slice does not estimate OLS or AR longitudinal states. +- Added `fast_mlsirm.multilevel.estimate_crossed_person_effects` and `CrossedPersonEffectResult` as marshal-only Python access, plus a true-parameter RMSE recovery test against simulated crossed membership weights. +- Enforced the binary-response contract before native discovery and again inside the Rust estimator: finite non-negative observed cells must be exactly `0` or `1`; negative and non-finite cells retain the established missing-data semantics. diff --git a/docs/doctoring/multilevel_crossed_person_effects.md b/docs/doctoring/multilevel_crossed_person_effects.md new file mode 100644 index 000000000..418b3eefc --- /dev/null +++ b/docs/doctoring/multilevel_crossed_person_effects.md @@ -0,0 +1,63 @@ +# Crossed multiple-membership person effects `u_h` + +## Decision + +`fast_mlsirm.multilevel.estimate_crossed_person_effects` estimates the +contextual random effects `u_h` of a binary IRT linear predictor when a +person belongs to more than one group at once. The kernel is Rust-owned. +Python validates, marshals a sealed `ContextMembershipDesign`, and reports +the immutable result. There is no Python numerical fallback and no stub. + +The implemented predictor is + +```text +eta_pi = a_i * (theta_p + sum_h w_ph * u_h) + b_i +``` + +`a_i` and `b_i` are known item parameters. `w_ph` are the Browne, Goldstein, +and Rasbash (2001) membership weights, already normalized to one inside each +classification. `theta_p` is an optional caller-supplied offset so a later or +already-estimated longitudinal state can enter the linear predictor. This +slice does not estimate OLS trends or AR coefficients. + +## Scientific rationale + +Fox and Glas (2001) specify a multilevel IRT model in which person location +depends on a cluster-level random effect with a Gaussian level-2 prior. The +ordinary nested case is one-hot membership: each person belongs to exactly +one unit of one classification. Browne, Goldstein, and Rasbash (2001) extend +the same additive random-effect term to multiple membership (several units of +one classification, weights summing to one) and multiple classification +(several classifications at once, i.e. crossed effects). + +This kernel is the matching MAP / ridge point estimator of the flattened +effects `u_h`, not the Fox and Glas Gibbs sampler and not an MMMC MCMC +variance-component engine. The reported estimate is re-centered to sum to +zero inside each classification so recovered effects are deviations. A +classification with fewer than two levels is rejected because the location +constraint would leave a non-identified singleton. + +## Compute boundary + +The `O(n_persons * n_items)` Bernoulli score and information reduction is +multithreaded on CPU and, when a wgpu adapter is present, executed by an f32 +GPU kernel. Missing adapters fall back to the f64 CPU path. Sparse membership +accumulation and the dense Newton system remain on CPU. `worker_count` does +not change the numerical result. + +## Verification + +Recovery evidence is RMSE against known simulated `u_h` under a crossed +school × neighborhood design that also includes weighted dual-school +membership. Correlation is supplementary only. Interval coverage, +variance-component ML, causal contextual effects, and continuous-time +dynamics are out of scope for this slice. + +## APA 7th references + +Browne, W. J., Goldstein, H., & Rasbash, J. (2001). Multiple membership +multiple classification (MMMC) models. *Statistical Modelling, 1*(2), +103–124. https://doi.org/10.1177/1471082X0100100202 + +Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel IRT +model. *Psychometrika, 66*, 271–288. https://doi.org/10.1007/BF02294839 diff --git a/docs/doctoring/multilevel_longitudinal_measurement.md b/docs/doctoring/multilevel_longitudinal_measurement.md index 7499558d7..09190ab8b 100644 --- a/docs/doctoring/multilevel_longitudinal_measurement.md +++ b/docs/doctoring/multilevel_longitudinal_measurement.md @@ -11,7 +11,8 @@ This first slice performs no statistical estimation. Python owns validation, content identity, bounded collection handling, replay protection, sparse design marshalling, and serialization only. Likelihood, integration, gradients, optimization, uncertainty, CPU multithreading, and any justified GPU batching -remain in Rust. +remain in Rust. Crossed / multiple-membership `u_h` MAP estimation is documented +separately in `docs/doctoring/multilevel_crossed_person_effects.md`. ## Scientific rationale diff --git a/docs/documentation_coverage.md b/docs/documentation_coverage.md index 61be2ca2d..fca329527 100644 --- a/docs/documentation_coverage.md +++ b/docs/documentation_coverage.md @@ -83,7 +83,7 @@ The table below records product truth, not documentation-file presence. “Imple | Latent-space residual interaction | IMPLEMENTED_ON_PROTECTED_MAIN | interpretation remains gated on substantive dimension/testlet/facet diagnosis | | Formal non-nested distinguishability/model comparison | PARTIAL | fail-closed relation-aware comparison exists; additional family-specific evidence and metadata remain incremental | | Adaptive rotation criterion selection | IMPLEMENTED_ON_PROTECTED_MAIN / PARTIAL | Rust-backed criterion registry/multi-start selector/report surfaces are integrated; additional criteria/GPU/recovery remain incremental | -| Multilevel / cross-classified / multiple-membership contracts | IMPLEMENTED_ON_PROTECTED_MAIN / PARTIAL | contextual and longitudinal contracts are integrated; estimator identification/recovery remains separate work | +| Multilevel / cross-classified / multiple-membership contracts | IMPLEMENTED_ON_PROTECTED_MAIN / PARTIAL | contextual and longitudinal contracts are integrated; crossed / multiple-membership `u_h` MAP estimation with RMSE recovery is this kernel; OLS/AR and richer variance-component claims remain separate | | Temporal/longitudinal/drift estimators | PARTIAL | governed contracts/design primitives exist; continuous-time or richer estimator claims require separate recovery evidence | | Automated essay scoring calibration/validation | IMPLEMENTED_ON_PROTECTED_MAIN / PARTIAL | governed essay contracts/validation/reporting exist; generalized rater discrimination/range/drift remains incremental | | Paired automated-vs-reference rating-range evidence | IMPLEMENTED_ON_PROTECTED_MAIN | Rust-owned paired range/compression diagnostic is integrated | diff --git a/docs/traceability/requirements-matrix.md b/docs/traceability/requirements-matrix.md index 3ccc5de87..c7b0e2305 100644 --- a/docs/traceability/requirements-matrix.md +++ b/docs/traceability/requirements-matrix.md @@ -24,7 +24,7 @@ This matrix makes the major product requirements discoverable without reconstruc | Adaptive rotation | PRD-FR-051/052, TRD-ROT | ADR-0009 | protected main contains `crates/mlsirm-core/src/rotation/`, PyO3 bindings, `python/fast_mlsirm/rotation.py`, `rotation_selection.py`, package-root exports, criterion-neutral selection and rotation regression/doctoring evidence | Accepted CPU baseline / planned GPU and broader recovery extensions | | True-parameter recovery | PRD-PRN-003, TRD-TEST-003..006 | ADR-0008 | simulation/recovery reports, Rust/NumPy parity, scheduled statistical studies/recovery contracts | Accepted | | Correlation vs recovery/agreement | PRD-PRN-003, scoring validity requirements | ADR-0008, ADR-0005 | recovery/simulation, agreement/QWK/facets evidence | Accepted: correlation is supplementary association evidence, never sole proof of parameter recovery or interchangeability | -| Multilevel/multiple-membership/temporal | PRD-FR-060..062, TRD-MLT | ADR-0007 | contextual summaries exist; full reusable contract PR remains open and Rust estimator recovery is future work | Proposed/partial / active PR | +| Multilevel/multiple-membership/temporal | PRD-FR-060..062, TRD-MLT | ADR-0007 | protected main has contextual/multilevel contracts; this active PR adds Rust MAP recovery of crossed / multiple-membership `u_h`, while OLS/AR and continuous-time claims remain separate slices | Proposed/partial / active PR | | Accessible standalone reports | PRD-FR-070..072, NFR-004 | ADR-0005 | report renderers, exact-value exports, WCAG-focused regression/doctoring | Accepted/evolving | | Sensitive data / PII utility | privacy/security requirements | ADR-0012 | source-free/digest/opaque-id provenance where implemented; provider error redaction; hosted identity/retention downstream | Accepted reusable policy / Downstream operations | | Continuous execution / documentation governance | TRD-DOC-002 / work-conserving automation | ADR-0013 | single-writer exact branch head; work-conserving when blocked; feasibility-first prioritization | Accepted governance | diff --git a/docs/traceability/research-basis.md b/docs/traceability/research-basis.md index b15ee0540..9e8c25556 100644 --- a/docs/traceability/research-basis.md +++ b/docs/traceability/research-basis.md @@ -127,11 +127,13 @@ Architecture effect: - prevent atomistic flattening; - explicit context dimensions and weighted memberships; - separate repeated occasion ordering from continuous-time dynamics; -- require Rust estimator identification/recovery before production claims. +- recover crossed / multiple-membership `u_h` with Rust MAP + RMSE evidence; +- keep OLS/AR and continuous-time claims on separate estimator slices. Primary basis: - Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel IRT model. *Psychometrika, 66*, 271–288. +- Browne, W. J., Goldstein, H., & Rasbash, J. (2001). Multiple membership multiple classification (MMMC) models. *Statistical Modelling, 1*(2), 103–124. - Uto, M. (2022). A Bayesian many-facet Rasch model with Markov modeling for rater severity drift. *Behavior Research Methods, 55*, 3910–3928. ## 9. Adaptive factor rotation — Proposed diff --git a/python/fast_mlsirm/multilevel/__init__.py b/python/fast_mlsirm/multilevel/__init__.py index 214bb1104..f3d0de806 100644 --- a/python/fast_mlsirm/multilevel/__init__.py +++ b/python/fast_mlsirm/multilevel/__init__.py @@ -1,4 +1,9 @@ -"""Public contextual-membership and longitudinal measurement contracts.""" +"""Public contextual-membership contracts and crossed ``u_h`` estimation. + +Contracts remain the sealed design layer. ``estimate_crossed_person_effects`` +is the Rust-owned MAP estimator of multiple-membership / crossed person +effects (Fox & Glas, 2001; Browne, Goldstein, & Rasbash, 2001). +""" from .contracts import ( ContextMembership, @@ -14,7 +19,11 @@ build_longitudinal_state_spec, build_temporal_occasion, ) -from .estimation import weighted_contextual_effect +from .estimation import ( + CrossedPersonEffectResult, + estimate_crossed_person_effects, + weighted_contextual_effect, +) __all__ = [ "ContextMembership", @@ -29,5 +38,7 @@ "build_longitudinal_design", "build_longitudinal_state_spec", "build_temporal_occasion", + "CrossedPersonEffectResult", + "estimate_crossed_person_effects", "weighted_contextual_effect", ] diff --git a/python/fast_mlsirm/multilevel/estimation.py b/python/fast_mlsirm/multilevel/estimation.py index e0e6b3cba..04b4bf8d8 100644 --- a/python/fast_mlsirm/multilevel/estimation.py +++ b/python/fast_mlsirm/multilevel/estimation.py @@ -1,22 +1,28 @@ -"""Typed Python access to the Rust-native contextual-effects predictor. - -This module performs marshalling only: converting a validated -``ContextMembershipDesign`` (see ``fast_mlsirm.multilevel.contracts``) into -the flat CSR arrays ``mlsirm_core::multilevel::weighted_contextual_effect`` -expects, and converting the caller's per-context random-effect values into -the matching flat vector. The additive sum, its determinism across worker -counts, and its numerical input validation are owned by the Rust core; see -that module's docstring for the full linear-predictor context and the -Browne, Goldstein, and Rasbash (2001) citation. +"""Typed Python access to Rust-native multilevel contextual kernels. + +This module performs marshalling only. The additive predictor +``sum_h w_ph * u_h`` and the MAP estimator of crossed / multiple-membership +person effects ``u_h`` are owned by ``mlsirm_core::multilevel``. See that +crate and the Fox & Glas (2001) / Browne, Goldstein, and Rasbash (2001) +citations on the public estimator. + +Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel +IRT model. *Psychometrika, 66*, 271-288. https://doi.org/10.1007/BF02294839 + +Browne, W. J., Goldstein, H., & Rasbash, J. (2001). Multiple membership +multiple classification (MMMC) models. *Statistical Modelling, 1*(2), +103-124. https://doi.org/10.1177/1471082X0100100202 """ from __future__ import annotations from collections.abc import Mapping +from dataclasses import dataclass import numpy as np from .._multilevel_core_loader import multilevel_core +from ._validation import exact_integer from .contracts import ContextMembershipDesign ContextKey = tuple[str, str] @@ -127,4 +133,278 @@ def weighted_contextual_effect( ) -__all__ = ["ContextKey", "weighted_contextual_effect"] +@dataclass(frozen=True) +class CrossedPersonEffectResult: + """Immutable MAP estimate of crossed / multiple-membership ``u_h``. + + Attributes + ---------- + context_effects: + Mapping from ``(context_dimension_id, context_id)`` to the centered + estimated random effect. Keys follow ``design.context_keys``. + effect_vector: + The same effects as a float64 vector aligned with ``context_keys``. + context_keys: + Deterministic dimension-qualified context identities. + loglik: + Bernoulli log-likelihood plus the Gaussian prior penalty. + n_iter: + Newton / IRLS iterations actually performed. + converged: + Whether the last effect step satisfied the requested tolerance. + used_gpu: + Whether the person-score reduction used the wgpu kernel. + termination_reason: + ``converged`` or ``max_iter_reached``. + """ + + context_effects: dict[ContextKey, float] + effect_vector: np.ndarray + context_keys: tuple[ContextKey, ...] + loglik: float + n_iter: int + converged: bool + used_gpu: bool + termination_reason: str + + +def _csr_from_design( + design: ContextMembershipDesign, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, tuple[ContextKey, ...]]: + """Marshal one sealed design into CSR arrays and classification offsets.""" + context_keys = design.context_keys + key_index = {key: index for index, key in enumerate(context_keys)} + by_observation: dict[str, list] = { + observation_id: [] for observation_id in design.observation_ids + } + for edge in design.memberships: + by_observation[edge.observation_id].append(edge) + + row_offsets: list[int] = [0] + context_indices: list[int] = [] + weights: list[float] = [] + for observation_id in design.observation_ids: + for edge in by_observation[observation_id]: + context_indices.append( + key_index[(edge.context_dimension_id, edge.context_id)] + ) + weights.append(edge.membership_weight) + row_offsets.append(len(context_indices)) + + classification_offsets = [0] + for dimension_id in design.context_dimension_ids: + classification_offsets.append( + classification_offsets[-1] + + sum(key[0] == dimension_id for key in context_keys) + ) + return ( + np.array(row_offsets, dtype=np.uint64), + np.array(context_indices, dtype=np.uint64), + np.array(weights, dtype=np.float64), + np.array(classification_offsets, dtype=np.uint64), + context_keys, + ) + + +def _finite_vector(values: object, name: str, length: int) -> np.ndarray: + """Return one exact 1-D float64 vector of the required length.""" + try: + array = np.asarray(values, dtype=np.float64) + except Exception: + raise ValueError(f"{name} could not be converted safely") from None + if array.ndim != 1 or array.shape[0] != length: + raise ValueError(f"{name} must be a length-{length} vector") + if not np.all(np.isfinite(array)): + raise ValueError(f"{name} must be finite") + return np.ascontiguousarray(array, dtype=np.float64) + + +def _response_matrix(values: object, n_persons: int, n_items: int) -> np.ndarray: + """Return one row-major binary response matrix aligned with the design.""" + try: + array = np.asarray(values, dtype=np.float64) + except Exception: + raise ValueError("responses could not be converted safely") from None + if array.ndim != 2 or array.shape != (n_persons, n_items): + raise ValueError("responses must have shape (n_observations, n_items)") + observed = np.isfinite(array) & (array >= 0.0) + nonbinary = observed & (array != 0.0) & (array != 1.0) + if np.any(nonbinary): + raise ValueError("binary responses must contain only 0 or 1 for observed cells") + return np.ascontiguousarray(array.reshape(-1), dtype=np.float64) + + +def _optional_offsets(values: object | None, n_persons: int) -> np.ndarray: + """Return empty offsets or one finite person-level location vector.""" + if values is None: + return np.zeros(0, dtype=np.float64) + return _finite_vector(values, "person_offsets", n_persons) + + +def _exact_positive_real(value: object, name: str) -> float: + """Return one strictly positive finite real without Boolean coercion.""" + if type(value) not in (int, float) or isinstance(value, bool): + raise ValueError(f"{name} must be a finite real number greater than zero") + number = float(value) + if not np.isfinite(number) or number <= 0.0: + raise ValueError(f"{name} must be a finite real number greater than zero") + return number + + +def _exact_device(value: object) -> str: + """Return one supported compute-device label.""" + if type(value) is not str: + raise ValueError("device must be one of 'cpu', 'gpu', or 'auto'") + device = value.strip().casefold() + if device not in {"cpu", "gpu", "auto"}: + raise ValueError("device must be one of 'cpu', 'gpu', or 'auto'") + return device + + +def estimate_crossed_person_effects( + responses: object, + design: ContextMembershipDesign, + *, + item_intercepts: object, + item_slopes: object | None = None, + person_offsets: object | None = None, + prior_scale: object = 1.0, + max_iter: object = 50, + tol: object = 1e-8, + worker_count: object = 1, + device: object = "auto", +) -> CrossedPersonEffectResult: + """Estimate crossed / multiple-membership person effects ``u_h``. + + The kernel is a Gaussian-prior MAP / Newton estimator of the Fox and Glas + (2001) multilevel IRT group effects, using Browne, Goldstein, and Rasbash + (2001) multiple-membership weights. Persons may belong to several units of + one classification and to several classifications at once. Known item + parameters stay fixed. Optional ``person_offsets`` accept already-estimated + longitudinal locations; this function does not estimate OLS or AR states. + + Parameters + ---------- + responses: + Binary response matrix aligned with ``design.observation_ids`` on axis + 0 and items on axis 1. Finite non-negative observed cells must be + exactly 0 or 1; non-finite or negative cells are treated as missing. + design: + A package-built ``ContextMembershipDesign``. Tampered designs fail + closed before native dispatch. + item_intercepts: + Known item intercepts ``b_i``, length ``n_items``. + item_slopes: + Known item discriminations ``a_i``. When omitted, Rasch slopes of 1 + are used. + person_offsets: + Optional person-level location offsets ``theta_p`` aligned with + ``design.observation_ids``. Use this to consume a longitudinal state + already estimated elsewhere. ``None`` treats every offset as zero. + prior_scale: + Level-2 standard deviation ``sigma_u`` of the Fox and Glas Gaussian + prior. The kernel uses precision ``1 / sigma_u^2``. + max_iter: + Newton / IRLS iteration budget (exact built-in ``int``, ``>= 1``). + tol: + Absolute effect-step convergence tolerance. + worker_count: + Deterministic CPU worker count (exact built-in ``int``, ``>= 1``). + The estimate does not depend on this value. + device: + ``cpu``, ``gpu``, or ``auto``. ``auto`` / ``gpu`` use the wgpu + person-score kernel when an adapter is present and fall back to the + f64 CPU reduction otherwise. + + Returns + ------- + CrossedPersonEffectResult + Centered ``u_h`` estimates, log-likelihood, and termination metadata. + + Raises + ------ + ValueError + If controls, shapes, response values, or the design fail the package + contract. + KeyError + Not used; membership identities come from the sealed design. + + Notes + ----- + Recovered effects are centered to sum to zero inside each classification. + The estimator is a MAP point method, not Fox and Glas Gibbs sampling and + not a variance-component ML claim. + """ + if type(design) is not ContextMembershipDesign: + raise ValueError("design must be an exact ContextMembershipDesign") + _ = design.design_fingerprint + n_persons = len(design.observation_ids) + n_effects = len(design.context_keys) + trusted_max_iter = exact_integer(max_iter, "max_iter", minimum=1, maximum=10_000) + trusted_workers = exact_integer( + worker_count, "worker_count", minimum=1, maximum=10_000 + ) + trusted_tol = _exact_positive_real(tol, "tol") + trusted_scale = _exact_positive_real(prior_scale, "prior_scale") + trusted_device = _exact_device(device) + try: + intercept_count = int(np.asarray(item_intercepts, dtype=np.float64).shape[0]) + except Exception: + raise ValueError("item_intercepts could not be converted safely") from None + intercepts = _finite_vector(item_intercepts, "item_intercepts", intercept_count) + n_items = int(intercepts.shape[0]) + if item_slopes is None: + item_slopes = np.ones(n_items, dtype=np.float64) + slopes = _finite_vector(item_slopes, "item_slopes", n_items) + if np.any(slopes <= 0.0): + raise ValueError("item_slopes must be strictly positive") + y = _response_matrix(responses, n_persons, n_items) + offsets = _optional_offsets(person_offsets, n_persons) + ( + row_offsets, + context_indices, + weights, + classification_offsets, + context_keys, + ) = _csr_from_design(design) + core = multilevel_core() + payload = core.estimate_crossed_person_effects( + y, + row_offsets, + context_indices, + weights, + slopes, + intercepts, + offsets, + classification_offsets, + n_persons, + n_items, + n_effects, + 1.0 / (trusted_scale * trusted_scale), + trusted_max_iter, + trusted_tol, + trusted_workers, + trusted_device, + ) + effect_vector = np.ascontiguousarray(payload["effects"], dtype=np.float64) + context_effects = { + key: float(value) for key, value in zip(context_keys, effect_vector, strict=True) + } + return CrossedPersonEffectResult( + context_effects=context_effects, + effect_vector=effect_vector, + context_keys=context_keys, + loglik=float(payload["loglik"]), + n_iter=int(payload["n_iter"]), + converged=bool(payload["converged"]), + used_gpu=bool(payload["used_gpu"]), + termination_reason=str(payload["termination_reason"]), + ) + + +__all__ = [ + "ContextKey", + "CrossedPersonEffectResult", + "estimate_crossed_person_effects", + "weighted_contextual_effect", +] diff --git a/tests/test_multilevel_crossed_response_contract.py b/tests/test_multilevel_crossed_response_contract.py new file mode 100644 index 000000000..328774f63 --- /dev/null +++ b/tests/test_multilevel_crossed_response_contract.py @@ -0,0 +1,68 @@ +"""Binary-response trust-boundary tests for crossed person-effect estimation.""" + +from __future__ import annotations + +import hashlib + +import numpy as np +import pytest + +import fast_mlsirm.multilevel.estimation as estimation +from fast_mlsirm.multilevel import ( + build_context_membership, + build_context_membership_design, +) + + +def _revision(tag: str) -> str: + """Return one deterministic content fingerprint for the test fixture.""" + return hashlib.sha256(tag.encode("utf-8")).hexdigest() + + +def _design(): + """Build the smallest identified one-classification membership design.""" + return build_context_membership_design( + [ + build_context_membership( + observation_id="person_alpha", + context_dimension_id="school_membership", + context_id="school_east", + membership_weight=1.0, + membership_revision_fingerprint=_revision("alpha-east"), + ), + build_context_membership( + observation_id="person_beta", + context_dimension_id="school_membership", + context_id="school_west", + membership_weight=1.0, + membership_revision_fingerprint=_revision("beta-west"), + ), + ] + ) + + +@pytest.mark.parametrize("invalid_response", [0.5, 1.5, 2.0]) +def test_nonbinary_observed_response_fails_before_native_discovery( + monkeypatch: pytest.MonkeyPatch, + invalid_response: float, +) -> None: + """Finite observed cells outside {0, 1} must never reach the Rust core.""" + core_discoveries = 0 + + def _unexpected_core_discovery(): + nonlocal core_discoveries + core_discoveries += 1 + raise AssertionError("native core must not be discovered for invalid responses") + + monkeypatch.setattr(estimation, "multilevel_core", _unexpected_core_discovery) + responses = np.array([[invalid_response], [0.0]], dtype=np.float64) + + with pytest.raises(ValueError, match="binary responses"): + estimation.estimate_crossed_person_effects( + responses, + _design(), + item_intercepts=np.array([0.0], dtype=np.float64), + device="cpu", + ) + + assert core_discoveries == 0 diff --git a/tests/test_multilevel_crossed_uh_doctoring.py b/tests/test_multilevel_crossed_uh_doctoring.py new file mode 100644 index 000000000..675e1900a --- /dev/null +++ b/tests/test_multilevel_crossed_uh_doctoring.py @@ -0,0 +1,18 @@ +"""Require APA 7th doctoring for the crossed ``u_h`` estimator.""" + +from pathlib import Path + +_ROOT = Path(__file__).resolve().parents[1] +_DOC = _ROOT / "docs" / "doctoring" / "multilevel_crossed_person_effects.md" + + +def test_crossed_uh_doctoring_cites_fox_glas_and_browne_mmmc() -> None: + """The estimator doctoring note must cite both primary papers in APA 7th.""" + note = _DOC.read_text(encoding="utf-8") + assert "Fox, J.-P., & Glas, C. A. W. (2001)" in note + assert "Browne, W. J., Goldstein, H., & Rasbash, J. (2001)" in note + assert "https://doi.org/10.1007/BF02294839" in note + assert "https://doi.org/10.1177/1471082X0100100202" in note + assert "Multiple membership" in note + assert "multiple classification (MMMC) models" in note + assert "does not estimate OLS" in note diff --git a/tests/test_multilevel_crossed_uh_recovery.py b/tests/test_multilevel_crossed_uh_recovery.py new file mode 100644 index 000000000..0f561ab23 --- /dev/null +++ b/tests/test_multilevel_crossed_uh_recovery.py @@ -0,0 +1,215 @@ +"""True-parameter RMSE recovery for crossed multiple-membership ``u_h``. + +This is a real accuracy gate, not a smoke test. A stub, pass-through, or +zero vector fails the RMSE threshold against known simulated context +effects. The design is simultaneously crossed (school × neighborhood) and +weighted multiple-membership (some persons split across two schools). + +Fox, J.-P., & Glas, C. A. W. (2001). Bayesian estimation of a multilevel +IRT model. *Psychometrika, 66*, 271-288. https://doi.org/10.1007/BF02294839 + +Browne, W. J., Goldstein, H., & Rasbash, J. (2001). Multiple membership +multiple classification (MMMC) models. *Statistical Modelling, 1*(2), +103-124. https://doi.org/10.1177/1471082X0100100202 +""" + +from __future__ import annotations + +import hashlib + +import numpy as np +import pytest + +from fast_mlsirm.multilevel import ( + build_context_membership, + build_context_membership_design, + estimate_crossed_person_effects, +) + + +def _revision(tag: str) -> str: + """Return one unique 64-character assignment-revision fingerprint.""" + return hashlib.sha256(tag.encode("utf-8")).hexdigest() + + +def _membership( + observation_id: str, + context_dimension_id: str, + context_id: str, + membership_weight: float, + tag: str, +): + """Build one sealed membership edge for the recovery fixture.""" + return build_context_membership( + observation_id=observation_id, + context_dimension_id=context_dimension_id, + context_id=context_id, + membership_weight=membership_weight, + membership_revision_fingerprint=_revision(tag), + ) + + +def _simulate_crossed_membership_responses( + *, + n_items: int = 28, + seed: int = 20260818, +) -> tuple[object, np.ndarray, dict[tuple[str, str], float], np.ndarray]: + """Simulate a crossed, partially multiple-membership Rasch design. + + Returns the sealed design, responses aligned with + ``design.observation_ids``, the true centered ``u_h`` map, and the known + item intercepts. + """ + true_effects = { + ("school_membership", "school_east"): -1.20, + ("school_membership", "school_west"): -0.40, + ("school_membership", "school_north"): 0.40, + ("school_membership", "school_south"): 1.20, + ("neighborhood_context", "neighborhood_a"): -0.80, + ("neighborhood_context", "neighborhood_b"): 0.00, + ("neighborhood_context", "neighborhood_c"): 0.80, + } + schools = [ + "school_east", + "school_west", + "school_north", + "school_south", + ] + neighborhoods = [ + "neighborhood_a", + "neighborhood_b", + "neighborhood_c", + ] + edges = [] + locations: dict[str, float] = {} + person_index = 0 + for school_index, school_id in enumerate(schools): + partner = schools[(school_index + 1) % len(schools)] + for neighborhood_id in neighborhoods: + for copy in range(8): + person_id = f"person_{person_index:03d}" + person_index += 1 + split = copy % 3 == 0 + if split: + edges.append( + _membership( + person_id, + "school_membership", + school_id, + 0.70, + f"{person_id}-school-a", + ) + ) + edges.append( + _membership( + person_id, + "school_membership", + partner, + 0.30, + f"{person_id}-school-b", + ) + ) + school_effect = ( + 0.70 * true_effects[("school_membership", school_id)] + + 0.30 * true_effects[("school_membership", partner)] + ) + else: + edges.append( + _membership( + person_id, + "school_membership", + school_id, + 1.0, + f"{person_id}-school", + ) + ) + school_effect = true_effects[("school_membership", school_id)] + edges.append( + _membership( + person_id, + "neighborhood_context", + neighborhood_id, + 1.0, + f"{person_id}-neighborhood", + ) + ) + locations[person_id] = ( + school_effect + + true_effects[("neighborhood_context", neighborhood_id)] + ) + design = build_context_membership_design(edges) + intercepts = np.linspace(-1.4, 1.4, n_items, dtype=np.float64) + rng = np.random.default_rng(seed) + responses = np.empty((len(design.observation_ids), n_items), dtype=np.float64) + for row, observation_id in enumerate(design.observation_ids): + eta = locations[observation_id] + intercepts + probability = 1.0 / (1.0 + np.exp(-eta)) + responses[row] = rng.binomial(1, probability) + return design, responses, true_effects, intercepts + + +def _rmse( + estimated: dict[tuple[str, str], float], + truth: dict[tuple[str, str], float], +) -> float: + """Return RMSE of centered context effects against the simulated truth.""" + errors = np.array( + [estimated[key] - truth[key] for key in truth], + dtype=np.float64, + ) + return float(np.sqrt(np.mean(errors**2))) + + +def test_crossed_multiple_membership_uh_recovers_true_effects() -> None: + """Estimated ``u_h`` must recover the simulated crossed membership effects.""" + design, responses, truth, intercepts = _simulate_crossed_membership_responses() + result = estimate_crossed_person_effects( + responses, + design, + item_intercepts=intercepts, + prior_scale=2.0, + max_iter=40, + tol=1e-8, + worker_count=4, + device="cpu", + ) + assert result.converged, result.termination_reason + assert set(result.context_effects) == set(truth) + rmse = _rmse(result.context_effects, truth) + zero_rmse = _rmse({key: 0.0 for key in truth}, truth) + assert zero_rmse > 0.70 + assert rmse < 0.25, ( + f"crossed u_h RMSE {rmse:.4f} exceeded the recovery gate; " + f"hats={result.context_effects} truth={truth}" + ) + + +def test_crossed_estimator_is_not_a_pass_through_of_weights() -> None: + """Membership weights are the design, not the estimand.""" + design, responses, truth, intercepts = _simulate_crossed_membership_responses() + result = estimate_crossed_person_effects( + responses, + design, + item_intercepts=intercepts, + prior_scale=2.0, + device="cpu", + ) + weight_like = { + key: float(sum(key[1] == edge.context_id for edge in design.memberships)) + for key in truth + } + assert _rmse(result.context_effects, truth) < _rmse(weight_like, truth) + + +def test_rejects_non_factory_design_before_native_dispatch() -> None: + """A hand-built design must not reach the Rust estimator.""" + + class FakeDesign: + """Hostile stand-in that is not a sealed ContextMembershipDesign.""" + + with pytest.raises(ValueError, match="ContextMembershipDesign"): + estimate_crossed_person_effects( + np.zeros((1, 2)), + FakeDesign(), # type: ignore[arg-type] + item_intercepts=np.zeros(2), + )