1use std::sync::Arc;
6
7use antecedent_core::{
8 AverageEffectQuery, ExecutionContext, KernelPolicy, TemporalEffectQuery, VariableId,
9};
10use antecedent_data::TemporalIndexer;
11use antecedent_data::{
12 DiscoveryEstimationSplit, PanelData, PanelUnit, TableView, TabularData, TimeIndex,
13 TimeSeriesData, ValidityBitmap,
14};
15use antecedent_estimate::{
16 EffectEstimate, EstimationWorkspace, IpwTarget, LinearAdjustmentAte, OverlapPolicy,
17 OverlapReport, TemporalLinearAdjustment,
18};
19use antecedent_identify::IdentifiedEstimand;
20use antecedent_kernels::erfc;
21use antecedent_stats::{
22 FaerBackend, GlmOptions, PropensityFit, PropensityWorkspace, fit_propensity_diagnostic,
23};
24
25use crate::error::ValidationError;
26
27#[derive(Clone, Copy, Debug)]
29pub struct TemporalRefitContext<'a> {
30 pub indexer: &'a TemporalIndexer,
32 pub temporal_query: &'a TemporalEffectQuery,
34 pub split: Option<&'a DiscoveryEstimationSplit>,
36 pub kernel_policy: &'a KernelPolicy,
38 pub time_index: Option<&'a TimeIndex>,
40 pub panel: Option<&'a PanelData>,
42}
43
44impl TemporalRefitContext<'_> {
45 #[must_use]
47 pub fn is_panel(&self) -> bool {
48 self.panel.is_some()
49 }
50}
51
52#[derive(Clone, Debug)]
54#[non_exhaustive]
55pub struct RefutationReport {
56 pub refuter: Arc<str>,
58 pub original_ate: f64,
60 pub refuted_ate: f64,
62 pub comparison: f64,
66 pub informative: bool,
68 pub passed: bool,
70 pub failure_condition: Option<Arc<str>>,
72 pub replicates: u32,
74}
75
76impl RefutationReport {
77 #[must_use]
79 #[allow(clippy::too_many_arguments)]
80 pub fn new(
81 refuter: impl Into<Arc<str>>,
82 original_ate: f64,
83 refuted_ate: f64,
84 comparison: f64,
85 informative: bool,
86 passed: bool,
87 failure_condition: Option<Arc<str>>,
88 replicates: u32,
89 ) -> Self {
90 Self {
91 refuter: refuter.into(),
92 original_ate,
93 refuted_ate,
94 comparison,
95 informative,
96 passed,
97 failure_condition,
98 replicates,
99 }
100 }
101}
102
103#[derive(Clone, Copy, Debug)]
105pub struct RefutationProblem<'a> {
106 pub data: &'a TabularData,
108 pub estimand: &'a IdentifiedEstimand,
110 pub query: &'a AverageEffectQuery,
112 pub original: &'a EffectEstimate,
114 pub estimator: Option<&'a str>,
116 pub temporal: Option<TemporalRefitContext<'a>>,
118}
119
120impl RefutationProblem<'_> {
121 #[must_use]
123 pub fn treatment(&self) -> VariableId {
124 self.query.treatment
125 }
126
127 #[must_use]
129 pub fn outcome(&self) -> VariableId {
130 self.query.outcome
131 }
132}
133
134pub(crate) fn with_replaced_float(
136 data: &TabularData,
137 id: VariableId,
138 values: Arc<[f64]>,
139) -> Result<TabularData, ValidationError> {
140 data.with_replaced_float(id, values).map_err(ValidationError::from)
141}
142
143pub(crate) fn with_extra_float(
145 data: &TabularData,
146 name: &str,
147 values: Arc<[f64]>,
148) -> Result<(TabularData, VariableId), ValidationError> {
149 data.with_appended_float(name, values).map_err(ValidationError::from)
150}
151
152pub(crate) fn fit_once(
154 estimator: &LinearAdjustmentAte,
155 data: &TabularData,
156 estimand: &IdentifiedEstimand,
157 query: &AverageEffectQuery,
158 workspace: &mut EstimationWorkspace,
159 ctx: &ExecutionContext,
160) -> Result<EffectEstimate, ValidationError> {
161 let prep = estimator.prepare(data, estimand, query).map_err(ValidationError::from)?;
162 estimator
163 .fit(&prep, workspace, ctx, antecedent_core::AssumptionSet::new())
164 .map_err(ValidationError::from)
165}
166
167pub(crate) fn refit_effect(
169 problem: &RefutationProblem<'_>,
170 data: &TabularData,
171 estimand: &IdentifiedEstimand,
172 extra_contemporaneous: &[VariableId],
173 workspace: &mut EstimationWorkspace,
174 ctx: &ExecutionContext,
175) -> Result<EffectEstimate, ValidationError> {
176 let Some(temporal) = problem.temporal else {
177 let est = linear_estimator_no_bootstrap();
178 return fit_once(&est, data, estimand, problem.query, workspace, ctx);
179 };
180 let mut estimator = TemporalLinearAdjustment::new();
181 estimator.inner.bootstrap_replicates = 0;
182 estimator.inner.overlap = OverlapPolicy::ExplicitOverride;
183 if let Some(panel) = temporal.panel {
184 let rebuilt = panel_from_stacked(panel, data)?;
185 let prep = if extra_contemporaneous.is_empty() {
186 let (prep, _cluster_ids, _panel_times) = estimator
187 .prepare_panel(
188 &rebuilt,
189 estimand,
190 temporal.temporal_query,
191 temporal.indexer,
192 temporal.split,
193 temporal.kernel_policy,
194 )
195 .map_err(ValidationError::from)?;
196 prep
197 } else {
198 panel_prepare_with_extras(
199 &estimator,
200 &rebuilt,
201 estimand,
202 &temporal,
203 extra_contemporaneous,
204 )?
205 };
206 return estimator
207 .fit(&prep, workspace, ctx, antecedent_core::AssumptionSet::new())
208 .map_err(ValidationError::from);
209 }
210 let time_index = temporal.time_index.ok_or(ValidationError::NotApplicable {
211 message: "temporal series refit requires time_index",
212 })?;
213 let series = TimeSeriesData::try_new(data.storage().clone(), time_index.clone())
214 .map_err(ValidationError::from)?;
215 let prep = estimator
216 .prepare_with_extras(
217 &series,
218 estimand,
219 temporal.temporal_query,
220 temporal.indexer,
221 temporal.split,
222 temporal.kernel_policy,
223 extra_contemporaneous,
224 )
225 .map_err(ValidationError::from)?;
226 estimator
227 .fit(&prep, workspace, ctx, antecedent_core::AssumptionSet::new())
228 .map_err(ValidationError::from)
229}
230
231fn panel_prepare_with_extras(
232 estimator: &TemporalLinearAdjustment,
233 panel: &PanelData,
234 estimand: &IdentifiedEstimand,
235 temporal: &TemporalRefitContext<'_>,
236 extra: &[VariableId],
237) -> Result<antecedent_estimate::PreparedEstimationProblem, ValidationError> {
238 let mut all_t = Vec::new();
240 let mut all_y = Vec::new();
241 let mut all_covs: Vec<(VariableId, Vec<f64>)> = Vec::new();
242 let mut adj_keys: Vec<VariableId> = Vec::new();
243 let mut active = 0.0;
244 let mut control = 0.0;
245 let mut treatment_delta = 0.0;
246 let mut first = true;
247 for unit in panel.units() {
248 let prep = estimator
249 .prepare_with_extras(
250 &unit.series,
251 estimand,
252 temporal.temporal_query,
253 temporal.indexer,
254 temporal.split,
255 temporal.kernel_policy,
256 extra,
257 )
258 .map_err(ValidationError::from)?;
259 if first {
260 active = prep.active;
261 control = prep.control;
262 treatment_delta = prep.treatment_delta;
263 adj_keys = prep.adjustment_set.to_vec();
264 all_covs = adj_keys.iter().map(|&id| (id, Vec::new())).collect();
265 first = false;
266 }
267 all_t.extend_from_slice(&prep.treatment);
268 all_y.extend_from_slice(&prep.design.outcome);
269 let nrows = prep.design.nrows;
270 for (i, (_id, dest)) in all_covs.iter_mut().enumerate() {
271 let base = (2 + i) * nrows;
272 dest.extend_from_slice(&prep.design.matrix[base..base + nrows]);
273 }
274 }
275 let cov_refs: Vec<(VariableId, &[f64])> =
276 all_covs.iter().map(|(id, v)| (*id, v.as_slice())).collect();
277 let selected: Vec<usize> = (0..all_t.len()).collect();
278 let design =
279 antecedent_stats::CompiledDesign::linear_adjustment(&all_t, &cov_refs, &all_y, &selected)
280 .map_err(ValidationError::from)?;
281 Ok(antecedent_estimate::PreparedEstimationProblem {
282 design,
283 method: Arc::from("temporal.linear.adjustment.panel"),
284 adjustment_set: Arc::from(adj_keys),
285 overlap: OverlapPolicy::ExplicitOverride,
286 treatment_delta,
287 target_population: antecedent_core::TargetPopulation::AllObserved,
288 treatment: Arc::from(all_t),
289 active,
290 control,
291 })
292}
293
294pub(crate) fn panel_from_stacked(
296 original: &PanelData,
297 stacked: &TabularData,
298) -> Result<PanelData, ValidationError> {
299 let expected = original.total_rows();
300 if stacked.row_count() != expected {
301 return Err(ValidationError::data_msg(format!(
302 "stacked panel refute rows {} != panel total_rows {expected}",
303 stacked.row_count()
304 )));
305 }
306 let mut offset = 0usize;
307 let mut units = Vec::with_capacity(original.unit_count());
308 for u in original.units() {
309 let n = u.series.row_count();
310 let slice = slice_tabular(stacked, offset, n)?;
311 let series =
312 TimeSeriesData::try_new(slice.storage().clone(), u.series.time_index().clone())
313 .map_err(ValidationError::from)?;
314 units.push(PanelUnit { unit_id: u.unit_id, series });
315 offset += n;
316 }
317 PanelData::try_new(Arc::from(units)).map_err(ValidationError::from)
318}
319
320fn slice_tabular(
321 data: &TabularData,
322 start: usize,
323 len: usize,
324) -> Result<TabularData, ValidationError> {
325 use antecedent_data::{Float64Column, OwnedColumn, OwnedColumnarStorage};
326 let storage = data.storage();
327 let end = start + len;
328 if end > data.row_count() {
329 return Err(ValidationError::NotApplicable { message: "panel slice out of range" });
330 }
331 let mut cols = Vec::with_capacity(storage.columns().len());
332 for col in storage.columns() {
333 match col {
334 OwnedColumn::Float64(c) => {
335 let values: Arc<[f64]> = Arc::from(c.values[start..end].to_vec());
336 let validity = ValidityBitmap::all_valid(len);
337 cols.push(OwnedColumn::Float64(
338 Float64Column::new(c.id, values, validity).map_err(ValidationError::from)?,
339 ));
340 }
341 _ => {
342 return Err(ValidationError::NotApplicable {
343 message: "panel refute slice requires float64 columns",
344 });
345 }
346 }
347 }
348 let mask = storage.analysis_mask().map(|m| {
349 let mut bytes = vec![0u8; len.div_ceil(8)];
350 for i in 0..len {
351 if m.is_valid(start + i) {
352 bytes[i / 8] |= 1 << (i % 8);
353 }
354 }
355 ValidityBitmap::from_bytes(bytes, len)
356 });
357 let mask = mask.transpose().map_err(ValidationError::from)?;
358 let weights = storage.weights().map(|w| Arc::<[f64]>::from(w[start..end].to_vec()));
359 let new_storage = OwnedColumnarStorage::try_new(storage.schema().clone(), cols, mask, weights)
360 .map_err(ValidationError::from)?;
361 Ok(TabularData::new(new_storage))
362}
363
364pub fn stack_panel_tabular(panel: &PanelData) -> Result<TabularData, ValidationError> {
366 use antecedent_data::{Float64Column, OwnedColumn, OwnedColumnarStorage};
367 let total = panel.total_rows();
368 let schema = panel.schema().clone();
369 let n_cols = schema.len();
370 let mut col_vals: Vec<Vec<f64>> = (0..n_cols).map(|_| Vec::with_capacity(total)).collect();
371 for u in panel.units() {
372 for (j, dest) in col_vals.iter_mut().enumerate() {
373 let id =
374 VariableId::from_raw(u32::try_from(j).map_err(|_| {
375 ValidationError::data_msg("panel column index exceeds u32::MAX")
376 })?);
377 let vals = u.series.float64_values(id).map_err(ValidationError::from)?;
378 dest.extend_from_slice(&vals);
379 }
380 }
381 let mut cols = Vec::with_capacity(n_cols);
382 for (j, values) in col_vals.into_iter().enumerate() {
383 let id = VariableId::from_raw(
384 u32::try_from(j)
385 .map_err(|_| ValidationError::data_msg("panel column index exceeds u32::MAX"))?,
386 );
387 cols.push(OwnedColumn::Float64(
388 Float64Column::new(id, Arc::from(values), ValidityBitmap::all_valid(total))
389 .map_err(ValidationError::from)?,
390 ));
391 }
392 let storage =
393 OwnedColumnarStorage::try_new(schema, cols, None, None).map_err(ValidationError::from)?;
394 Ok(TabularData::new(storage))
395}
396
397pub(crate) struct DiagnosticPropensityColumns {
399 pub scores: Vec<f64>,
401 pub treatment: Vec<f64>,
403 pub outcome: Option<Vec<f64>>,
405}
406
407pub(crate) fn fit_diagnostic_propensity(
411 problem: &RefutationProblem<'_>,
412 glm_options: &GlmOptions,
413 include_outcome_in_mask: bool,
414 propensity: &mut PropensityWorkspace,
415) -> Result<DiagnosticPropensityColumns, ValidationError> {
416 let mut ids = vec![problem.treatment()];
417 if include_outcome_in_mask {
418 ids.push(problem.outcome());
419 }
420 ids.extend_from_slice(&problem.estimand.adjustment_set);
421 let row_mask = problem.data.complete_case_mask(&ids).map_err(ValidationError::from)?;
422 let treatment = problem
423 .data
424 .float64_masked(problem.treatment(), &row_mask)
425 .map_err(ValidationError::from)?;
426 let outcome = if include_outcome_in_mask {
427 Some(
428 problem
429 .data
430 .float64_masked(problem.outcome(), &row_mask)
431 .map_err(ValidationError::from)?,
432 )
433 } else {
434 None
435 };
436 let nrows = treatment.len();
437 let ncols = 1 + problem.estimand.adjustment_set.len();
438 let mut design = vec![0.0; nrows * ncols];
439 for r in design.iter_mut().take(nrows) {
440 *r = 1.0;
441 }
442 for (i, &z) in problem.estimand.adjustment_set.iter().enumerate() {
443 let col = problem.data.float64_masked(z, &row_mask).map_err(ValidationError::from)?;
444 let base = (1 + i) * nrows;
445 design[base..base + nrows].copy_from_slice(&col);
446 }
447 let backend = FaerBackend;
448 let fit: PropensityFit = fit_propensity_diagnostic(
449 &design,
450 nrows,
451 ncols,
452 &treatment,
453 &backend,
454 propensity,
455 glm_options,
456 )
457 .map_err(ValidationError::from)?;
458 Ok(DiagnosticPropensityColumns { scores: fit.scores, treatment, outcome })
459}
460
461pub(crate) fn diagnostic_overlap_report(
463 problem: &RefutationProblem<'_>,
464 glm_options: &GlmOptions,
465 policy: OverlapPolicy,
466) -> Result<OverlapReport, ValidationError> {
467 let mut ws = PropensityWorkspace::default();
468 diagnostic_overlap_report_with(problem, glm_options, policy, &mut ws)
469}
470
471pub(crate) fn diagnostic_overlap_report_with(
473 problem: &RefutationProblem<'_>,
474 glm_options: &GlmOptions,
475 policy: OverlapPolicy,
476 propensity: &mut PropensityWorkspace,
477) -> Result<OverlapReport, ValidationError> {
478 let cols = fit_diagnostic_propensity(problem, glm_options, false, propensity)?;
479 let weights: Vec<f64> = cols
481 .treatment
482 .iter()
483 .zip(cols.scores.iter())
484 .map(|(&t, &p)| {
485 let p = p.clamp(1e-9, 1.0 - 1e-9);
486 if t > 0.5 { 1.0 / p } else { 1.0 / (1.0 - p) }
487 })
488 .collect();
489 Ok(OverlapReport::from_propensities(
490 &cols.scores,
491 Some(&weights),
492 policy,
493 Some(&cols.treatment),
494 Some(IpwTarget::Ate),
495 None,
496 ))
497}
498
499#[must_use]
501pub(crate) fn linear_estimator_no_bootstrap() -> LinearAdjustmentAte {
502 let mut estimator = LinearAdjustmentAte::new();
503 estimator.bootstrap_replicates = 0;
504 estimator
505}
506
507#[derive(Clone, Copy, Debug, Eq, PartialEq)]
509pub(crate) enum NoiseReplaceTarget {
510 Treatment,
512 Outcome,
514}
515
516#[allow(clippy::too_many_arguments)]
521pub(crate) fn noise_replace_refute(
522 problem: &RefutationProblem<'_>,
523 workspace: &mut EstimationWorkspace,
524 ctx: &ExecutionContext,
525 _estimator: &LinearAdjustmentAte,
526 replicates: u32,
527 alpha: f64,
528 target: NoiseReplaceTarget,
529 stream_base: u64,
530 refuter_id: &'static str,
531 failure_label: &'static str,
532) -> Result<RefutationReport, ValidationError> {
533 if replicates < 2 {
534 return Err(ValidationError::NotApplicable {
535 message: "noise-replace refuter requires replicates >= 2",
536 });
537 }
538 let replace_id = match target {
539 NoiseReplaceTarget::Treatment => problem.treatment(),
540 NoiseReplaceTarget::Outcome => problem.outcome(),
541 };
542 let n = problem.data.row_count();
543 let mut ates = Vec::with_capacity(replicates as usize);
544 for r in 0..replicates {
545 let mut noise = vec![0.0; n];
546 fill_gaussian(&mut noise, ctx, stream_base.wrapping_add(u64::from(r)));
547 let data = with_replaced_float(problem.data, replace_id, Arc::from(noise))?;
548 let est = refit_effect(problem, &data, problem.estimand, &[], workspace, ctx)?;
549 ates.push(est.ate);
550 }
551 let mean_ate = ates.iter().sum::<f64>() / f64::from(replicates);
552 let p_value = replicate_p_value(&ates, 0.0);
553 let passed = p_value >= alpha;
554 Ok(RefutationReport {
555 refuter: Arc::from(refuter_id),
556 original_ate: problem.original.ate,
557 refuted_ate: mean_ate,
558 comparison: p_value,
559 informative: true,
560 passed,
561 failure_condition: (!passed).then(|| {
562 Arc::from(format!(
563 "{failure_label} ATE distribution (mean {mean_ate}) is inconsistent with zero \
564 (p={p_value} < alpha={alpha})"
565 ))
566 }),
567 replicates,
568 })
569}
570
571pub(crate) fn replicate_p_value(samples: &[f64], hypothesized: f64) -> f64 {
574 if samples.len() < 2 {
575 return 1.0;
576 }
577 #[allow(clippy::cast_precision_loss)]
578 let mean = samples.iter().sum::<f64>() / samples.len() as f64;
579 let sd = sample_sd(samples);
580 let scale = mean.abs().max(hypothesized.abs()).max(1.0);
581 if !sd.is_finite() {
582 return 1.0;
583 }
584 if sd <= 1e-12 * scale {
585 return if (hypothesized - mean).abs() <= 1e-9 * scale { 1.0 } else { 0.0 };
586 }
587 let z = (hypothesized - mean) / sd;
588 erfc(z.abs() / std::f64::consts::SQRT_2)
589}
590
591pub(crate) fn float64_full(
593 data: &TabularData,
594 id: VariableId,
595) -> Result<Vec<f64>, ValidationError> {
596 data.float64_values(id).map_err(ValidationError::from)
597}
598
599pub(crate) fn with_row_subset(
602 data: &TabularData,
603 keep_fraction: f64,
604 ctx: &ExecutionContext,
605 stream_id: u64,
606) -> Result<TabularData, ValidationError> {
607 let n = data.row_count();
608 let mut rng = ctx.rng.stream(stream_id);
609 let mut bytes = vec![0u8; n.div_ceil(8)];
610 for i in 0..n {
611 if rng.next_f64() < keep_fraction {
612 bytes[i / 8] |= 1 << (i % 8);
613 }
614 }
615 let mask = ValidityBitmap::from_bytes(bytes, n).map_err(ValidationError::from)?;
616 data.with_analysis_mask(mask).map_err(ValidationError::from)
617}
618
619pub(crate) fn with_resampled_rows(
625 data: &TabularData,
626 resample_ids: &[VariableId],
627 row_idx: &[usize],
628 keep: &[bool],
629) -> Result<TabularData, ValidationError> {
630 let mut out = data.clone();
631 for &id in resample_ids {
632 let full = float64_full(&out, id)?;
633 let resampled: Vec<f64> = row_idx.iter().map(|&i| full[i]).collect();
634 out = with_replaced_float(&out, id, Arc::from(resampled))?;
635 }
636 if keep.iter().all(|&k| k) {
637 return Ok(out);
638 }
639 let n = keep.len();
640 let mut bytes = vec![0u8; n.div_ceil(8)];
641 for (i, &k) in keep.iter().enumerate() {
642 if k {
643 bytes[i / 8] |= 1 << (i % 8);
644 }
645 }
646 let mask = ValidityBitmap::from_bytes(bytes, n).map_err(ValidationError::from)?;
647 out.with_analysis_mask(mask).map_err(ValidationError::from)
648}
649
650pub(crate) fn complete_case_rows(
652 data: &TabularData,
653 ids: &[VariableId],
654) -> Result<(Vec<bool>, Vec<usize>), ValidationError> {
655 let mask = data.complete_case_mask(ids).map_err(ValidationError::from)?;
656 let valid: Vec<usize> = mask.iter().enumerate().filter_map(|(i, &k)| k.then_some(i)).collect();
657 Ok((mask, valid))
658}
659
660pub(crate) fn masked_sample_sd(
662 data: &TabularData,
663 id: VariableId,
664 mask: &[bool],
665) -> Result<f64, ValidationError> {
666 let vals = data.float64_masked(id, mask).map_err(ValidationError::from)?;
667 Ok(sample_sd(&vals))
668}
669
670pub(crate) fn sample_sd(values: &[f64]) -> f64 {
672 antecedent_stats::sample_std(values)
673}
674
675pub(crate) fn fill_gaussian(out: &mut [f64], ctx: &ExecutionContext, stream_id: u64) {
680 let mut rng = ctx.rng.stream(stream_id);
681 antecedent_kernels::fill_standard_normal(&mut rng, out);
682}