🥮 mooncakes.io

    moonbit_doubleML

    Pure-MoonBit port of doubleml-for-py: Double / Debiased Machine Learning for partial linear, IRM, IV, DID, SSM, APO(S), PQ, QTE, LPQ, CVaR, RDD, BLP, and policy-tree models

    doubleml
    DML
    causal-inference
    machine-learning
    debiased-estimation
    Download zip
    Author
    Version
    0.64.0
    License
    MIT
    Repository
    Last updated
    6 hours ago
    Downloads
    21

    #moonbit_doubleML · Double / Debiased Machine Learning in MoonBit

    A pure-MoonBit port of the doubleml-for-py Python package — the Python ecosystem's most mature Double / Debiased Machine Learning implementation — covering all 19 models in upstream plus 3 MoonBit-specific extras (DoubleMLPLPR for static panel data, DoubleMLLPLR for binary outcomes, DoubleMLDIDCSBinary for binary-outcome CS-DID).

    Targets library authors and tool developers who need causal inference, panel data estimation, instrumental variable estimation, regression discontinuity, policy evaluation, or interpretable machine learning inside the MoonBit ecosystem. Delivers a unified DoubleMLXxx::new(data, n_folds, n_rep, seed).fit() API across all 22 estimators, robust variance estimation, propensity-score processing, multiple-testing correction, and a reproducible validation pipeline (23 / 23 Python reference scripts PASS) — on native, wasm, wasm-gc, and js backends under moon test --deny-warn.

    #Status

    ItemValue
    Repositoryhttps://github.com/riantr/moonbit_doubleML
    Authorriantr
    LicenseMIT (port of upstream doubleml-for-py, BSD-3-Clause)
    moon.mod version0.63.0
    Source layoutflat, moonbit_doubleML/ (the library)
    .mbt file count129 production files
    Estimators22 DoubleML* estimator structs (PLR / IRM / PLIV / IIVM / DID family / SSM / APO(S) / PQ / QTE / LPQ / LPLR / CVAR / RDD / BLP / PLPR / PolicyTree)
    Backendsnative, wasm, wasm-gc, js — all pass moon test --deny-warn
    Tests (native / wasm / js)491 / 491
    Tests (wasm-gc)497 / 497 (lib + 6 doc tutorials)
    Python cross-checks23 / 23 PASS (validate_*_with_python.py)
    HTTP serviceexamples/api_server/ — hand-rolled on moonbitlang/async, no third-party framework

    #Features

    Pure-MoonBit, all-4-backends port of doubleml-for-py — 22 DoubleML* estimators (DoubleMLPLR, DoubleMLIRM, DoubleMLPLIV, DoubleMLIIVM, DoubleMLDID, DoubleMLDIDBinary, DoubleMLDIDCS, DoubleMLDIDCSBinary, DoubleMLDIDMulti, DoubleMLDIDCrossSection, DoubleMLSSM, DoubleMLAPO, DoubleMLAPOS, DoubleMLPQ, DoubleMLQTE, DoubleMLLPQ, DoubleMLLPLR, DoubleMLCVAR, DoubleMLRDD, DoubleMLBLP, DoubleMLPLPR, DoubleMLPolicyTree) in 17 source files. The 19 marked (upstream) mirror the upstream doubleml-for-py API; the 3 marked (extra) are MoonBit-specific additions for static-panel PLR, partially-logistic regression, and binary-outcome CS-DID.

    Unified API surface — every estimator follows the same DoubleMLXxx::new(data, n_folds, n_rep, seed).fit() shape and returns the same .coef() / .se() / .confint() / .bootstrap() accessors. No per-estimator interface drift.

    Cross-fitting infrastructure — kfold, repeated cross-fitting, stratified K-fold (kfold_stratified), chacha8_rng-based seeded PRNG (chacha8_rng(seed)), and Kahan summation. Default seed=3141; same seed yields byte-identical DGP outputs across runs.

    Propensity-score processor (PSProcessor) — clipping_threshold clipping, isotonic (PAVA) calibration, K-fold cross-validated (CV) calibration. Plugs directly into DoubleMLIRM / DoubleMLIIVM / DoubleMLDID and the rest of the IPW-based models.

    Robust variance & inference — heteroskedasticity-consistent (HC) standard errors, cluster-robust variance, multiplier bootstrap confidence intervals, joint CIs, and Romano-Wolf multiple-testing p-value adjustment. Reused via the bootstrap.mbt helper extracted in v0.55.0.

    Coverage by class:

    • Observational / quasi-experimental: DoubleMLDID, DoubleMLDIDBinary, DoubleMLDIDCS, DoubleMLDIDCSBinary, DoubleMLDIDMulti, DoubleMLDIDCrossSection, DoubleMLRDD, DoubleMLSSM.
    • Strategy & distributional: DoubleMLAPO / DoubleMLAPOS, DoubleMLPQ, DoubleMLQTE, DoubleMLLPQ, DoubleMLCVAR, DoubleMLBLP, DoubleMLPolicyTree.

    Reproducible DGPs + Python reference validation — 23 Python scripts (validate_*_with_python.py) drive 23 DGPs (plr_CCDDHNR, plr_turrell, plr_confounded, irm_discrete, irm_heterogeneous, irm_confounded, iivm, pliv / pliv_cluster, SSM, simple_rdd, DID SZ2020 / CS2021, etc.) and check within-tolerance equivalence against the upstream Python outputs. CI greps the trailing PASS line from each script.

    Docs & demos — 6 doc/<NNN_…>/ tutorials targeting wasm-gc, plus 14 examples/<bin>/ driver binaries (all listed in moon.work), of which examples/api_server/ is an HTTP service built directly on moonbitlang/async@0.20.3 — no third-party HTTP framework is pulled in.

    Strict dependency hygiene — only official moonbitlang/* packages in the library. Non-official deps restricted to riantr/* (this repo). Reproducible builds, minimal supply-chain surface.

    Learner injection via LearnerDispatch — every estimator's new() / fit() accept ml_g / ml_m / ml_l / ml_r typed optionals. Default is LearnerDispatch::linear_regression() (OLS); v0.62.0+ also accepts LearnerDispatch::logistic_regression(LR) for LPLR's binary-classification slots. Built-in learners: LinearRegression, ConstantLearner, NoopLearner, RFLearner (Breiman 2001 regression), GBLearner (Friedman 2001 regression). RFLearner and GBLearner now reach the 5 specialised internals — DoubleMLDIDCrossSection::crossfit_nuisance, DoubleMLDIDCSBinary::cs_bin_crossfit_nuisance, DoubleMLPQ::solve_pq / DoubleMLQTE::solve_pq, DoubleMLCVAR::cvar_inner_crossfit, and DoubleMLLPLR::cross_fit_predict_dispatch (binary classifier slots) / DoubleMLRDD::rdd_side (kernel-weighted OLS for ml_g = LinearRegression, unweighted fallback for other learners). Defaults preserve v0.61.0 byte-equality; non-default overrides change the IF and the coef.

    #Quick Start

    $ moon test --deny-warn
Total tests: 485, passed: 485, failed: 0.

$ moon run examples/main
=== MoonBit DML PLR (partialling out) ===
true theta_0      = 1
estimated theta   = 0.9763281577675552
standard error    = 0.08740490230937094
...

$ python validate_irm_with_python.py
======================================================================
PASS  |mb - handrolled_nrep5| (theta) = 5.24e-02 < max(MODEL_TOL=0.1, 2.0*handrolled_n5_se) = 1.91e-01

    #Library Use

    let data = @moonbit_doubleML.DoubleMLData::new(x, y, d)  // x : Matrix, y / d : Array[Double]
    let fitted = @moonbit_doubleML.DoubleMLPLR::new(data, n_folds=2, n_rep=1, seed=3141).fit()
    let coef = fitted.coef()      // Double
    let se   = fitted.se()        // Double
    let (lo, hi) = fitted.confint()

    Same shape for every other estimator — e.g. @moonbit_doubleML.DoubleMLIRM::new(data, ml_g?, ml_m?, n_folds?, n_rep?, seed?).fit(), @moonbit_doubleML.DoubleMLDID::new(data, ml_g?, ml_m?, ...).fit().

    #Examples

    14 examples/<bin>/ directories in moon.work. 13 run end-to-end on synthetic DGPs: 12 are CLI-style numeric demos (print true-vs-estimated θ and a 95 % CI); the 13th — examples/api_server/ — is an HTTP service. The 14th — examples/consumer_demo/ — is the library-user pattern that mirrors what an external moon add riantr/moonbit_doubleML consumer would write.

    DriverModelDGPTrue θ
    examples/mainDoubleMLPLR + 11 others (IRM / PLIV / IIVM / DID / SSM / BLP / RDD / PQ / QTE / LPQ / PolicyTree)Simple partially linear, n=500, p=51.0
    examples/datasetsDoubleMLPLR + DoubleMLIRMSynthetic 401(k)-style, n=5000, 9 controls1.5
    examples/did_binaryDoubleMLDIDBinary2-period panel DID, 400 units1.0
    examples/did_csDoubleMLDIDCSStaggered CS-DID, 4 cohorts × 4 periods1.0
    examples/did_cs_binaryDoubleMLDIDCSBinaryStaggered CS-DID with binary outcome1.0
    examples/did_multiDoubleMLDIDMultiTop-level multi-period DID + aggregation1.0
    examples/did_cross_sectionDoubleMLDIDCrossSectionSant'Anna-Zhao 2020 cross-section DID, 500 units1.0
    examples/plpr, examples/lplrDoubleMLPLPR / DoubleMLLPLRStatic-panel PLR / partially logistic regression1.0
    examples/apos, examples/cvarDoubleMLAPOS / DoubleMLCVARAPO policy score / CVaR1.0
    examples/fuzzwrappersRandom-property fuzz harness across 12 surfacesn/a
    examples/consumer_demolibrary-user patternMinimal end-to-end PLR (no estimator-specific extras)1.0
    examples/api_serverDoubleMLPLR + DoubleMLIRMHTTP service — see belown/a

    CLI demos:

    $ moon run examples/main
$ moon run examples/datasets
$ moon run examples/did_binary
$ moon run examples/did_cs
$ moon run examples/did_multi
$ moon run examples/did_cross_section

    HTTP service:

    $ moon build --target native examples/api_server
$ _build/native/debug/build/api_server/api_server.exe        # listens on 127.0.0.1:4000

$ curl -s http://127.0.0.1:4000/healthz
{"status":"ok"}

$ curl -s -X POST http://127.0.0.1:4000/fit/plr \
    -H 'Content-Type: application/json' \
    -d '{"x":[[1.0,0.5],[2.0,1.5],[3.0,2.5],[4.0,3.5],[5.0,4.5],[6.0,5.5]],"y":[2.0,4.0,6.0,8.0,10.0,12.0],"d":[1.0,2.0,3.0,4.0,5.0,6.0],"n_folds":2,"seed":3141}'
{"estimator":"plr","coef":2,"se":0,"ci_lo":2,"ci_hi":2,"n_obs":6,"n_features":2}

    The same api_server.wasm also runs under moonrun --port 4000 (see examples/api_server/README.md for the full curl recipe). The service is built directly on moonbitlang/async@0.20.3 — no third-party HTTP framework is pulled in.

    #Models

    22 estimators (DoubleML* structs) in 22 files (one struct per file, plus the supporting DGPs and score / nuisance kernels in 129 files total), all reachable through the same DoubleMLXxx::new(...).fit() interface. The 19 marked (upstream) mirror the upstream doubleml-for-py API surface; the 3 (extra) rows are MoonBit additions for static-panel PLR, partially-logistic regression, and binary-outcome CS-DID:

    ModelScoreDescription
    DoubleMLPLRpartialling-outpartially linear regression (upstream)
    DoubleMLIRMATEinteractive regression model (upstream)
    DoubleMLPLIVpartialling-outpartially linear IV regression (upstream)
    DoubleMLIIVMLATEinteractive IV model (upstream)
    DoubleMLDIDobservationaldifference-in-differences (upstream)
    DoubleMLDIDBinaryobservationalDID with binary outcome (upstream)
    DoubleMLDIDCSobservationalDID with staggered adoption (Callaway-Sant'Anna) (upstream)
    DoubleMLDIDCSBinaryobservationalCS-DID with binary outcome (extra)
    DoubleMLDIDMultiobservationalmulti-period DID with group-time ATT aggregation (upstream)
    DoubleMLDIDCrossSectionobservationalSant'Anna-Zhao 2020 cross-section DID (upstream)
    DoubleMLSSMMARsample selection (missing-at-random) (upstream)
    DoubleMLAPOpolicy scoreaverage prescriptive effect (upstream)
    DoubleMLAPOSpolicy scoreAPO with stratified treatment (upstream)
    DoubleMLPQquantilepotential quantile (upstream)
    DoubleMLQTEquantilequantile treatment effect (upstream)
    DoubleMLLPQlocal polynomiallocal potential quantile (upstream)
    DoubleMLLPLRpartialling-outpartially logistic regression, Liu-Zhang-Zhou 2021 (extra)
    DoubleMLCVARCVaRconditional value-at-risk (upstream)
    DoubleMLRDDobservationalregression discontinuity (upstream)
    DoubleMLBLPIVbest linear predictor of treatment effect (upstream)
    DoubleMLPLPRpartialling-outpartially linear panel regression, Clarke-Polselli 2025 (extra)
    DoubleMLPolicyTreepolicypolicy tree (upstream)

    #Project layout

    moonbit_doubleML/        <- the library (moon.mod v0.63.0, 130 .mbt files)
  moonbit_doubleML.mbt   <- main re-export file (the import surface)
  ...                    <- one file per estimator + DGPs + score / nuisance kernels

doc/                     <- wasm-gc-targeted numbered tutorials
  001_introduction/
  ...
  006_python_check/

examples/                <- 14 driver binaries (all listed in moon.work)
  main/                  <- 12 estimators, one perfect-DGP run each
  datasets/              <- 401(k)-style ATE / ATT recovery
  did_binary/            <- 2-period panel DID
  did_cs/                <- staggered CS-DID
  did_cs_binary/         <- staggered CS-DID (binary outcome)
  did_multi/             <- multi-period DID aggregation
  did_cross_section/     <- Sant'Anna-Zhao 2020
  plpr/, lplr/, apos/,   <- single-estimator focused demos
  cvar/, fuzz/
  consumer_demo/         <- library-user pattern (`example/consumer_demo`)
  api_server/            <- HTTP service built on `moonbitlang/async`

moon.work                <- 15-member workspace aggregator
AGENTS.md                <- this file (per moonbit agent convention)
skills/moonbit_doubleML.md  <- agent skill: API surface + anti-patterns

    #Dependency rule (per skills/moonbit_doubleML.md)

    • Only official moonbitlang/* packages.
    • Non-official packages may only be riantr/* (this repo).
    • Pin to the latest published version (e.g. moonbitlang/async@0.20.3).
    • Consequence: examples/api_server/ builds directly on moonbitlang/async's raw Server API. No third-party HTTP framework is allowed.

    #Determinism & validation

    • All randomness flows through chacha8_rng(seed) (seed.mbt). Default seed=3141; same seed ⇒ byte-identical DGP outputs.
    • Python reference scripts validate_*_with_python.py reproduce the upstream doubleml-for-py numbers. CI greps the trailing PASS line from each script.
    • dgp_recovery_test.mbt enforces |coef - theta| < MAX(MODEL_TOL=0.1,2.0 * handrolled_se). The tolerance is contractual.
    • v0.62.0 verified counts: moon test 485 / 485 (native & wasm & js) and 491 / 491 (wasm-gc); python _verify/run_all_validators.py 23 / 23 PASS in ~70 s. moon fmt --check is now clean (was failing on every release since v0.57.0).
    • v0.63.0 verified counts: moon test 491 / 491 (native & wasm & js) and 497 / 497 (wasm-gc); 23 / 23 Python cross-validators PASS in ~90 s. Adds 6 new wbtests (v063_wbtest.mbt) covering APO learner plumbing + APOS bootstrap/confint + BLP plumbing.

    #Attribution

    moonbit_doubleML is a port of the upstream DoubleML/doubleml-for-py package (BSD 3-Clause License). All 19 upstream models are faithfully re-implemented in pure MoonBit with the same score formulae and variance estimators; the 3 MoonBit-only additions (DoubleMLPLPR, DoubleMLLPLR, DoubleMLDIDCSBinary) are extensions beyond the upstream scope.

    #Differences from upstream doubleml-for-py

    • Native MoonBit package layout, type system, and test organization.
    • Independent implementations of Matrix, linear / logistic regression, stratified K-fold, ChaCha8 PRNG, and Kahan summation in MoonBit; no NumPy / scikit-learn dependency.
    • DoubleMLData / DoubleMLXxx::new().fit() as the canonical MoonBit-idiomatic API, with optional JSON in/out for native / wasm-gc / wasm / js backends (Serverless / Wasm / edge deployments).
    • Behaviour-level parity check via the 23 Python reference scripts: within-tolerance equivalence rather than bit-exact re-implementation. This keeps the port resilient to upstream refactors.
    • Strict dependency rule (see above): only moonbitlang/* official packages. Minimal supply-chain surface.

    #Used By

    • examples/api_server/ — this repo's HTTP service.
    • examples/consumer_demo/ — minimal end-to-end PLR showing the consumer-side import pattern (moon add riantr/moonbit_doubleML).

    Learner

    source
    pub trait Learner {
      fn fit(Self, Matrix, Array[Double]) -> Self
      fn predict(Self, Matrix) -> Array[Double]
    }

    Trait for any learner that can be fit on (x, y) and predict on a feature matrix. Used by the DML cross-fit helpers and the PLR estimator.

    BootstrapMethodError

    source
    pub suberror BootstrapMethodError {
      UnknownMethod(String)
    }

    Error type raised by draw_bootstrap_weights when the requested bootstrap method is not one of the supported options ("normal", "Bayes", "wild"). The payload is the unknown method name, which the caller can use to log or surface a meaningful diagnostic.

    Added in v0.37.0 to convert the previous abort() call into a testable error path. The pre-v0.37.0 panic_bootstrap_invalid_method test only exercised the require check in DoubleMLDIDMulti::bootstrap (which fires before the _ => abort fallback) — the test was silently skipped on native/wasm-gc. The new test draw_bootstrap_weights_raises_unknown_method calls draw_bootstrap_weights directly with an invalid method, exercising the previously-unreachable _ => abort path.

    BracketSignError

    source
    pub suberror BracketSignError {
      UpperSignFailed
    }

    Error type raised by solve_pq when the IPW score at the upper bracket remains non-positive after 20 exponential widens — the score is structurally non-monotonic on the bracket, the quantile is non-monotonic in this data, or q is too close to 1 with sparse treatment.

    Added in v0.42.0 to convert the previous abort() call in solve_pq into a testable error path. The pre-v0.42.0 abort was reachable only via the (private) solve_pq helper; no public test exercised it. The new solve_pq_raises_on_upper_bracket_sign_failure test constructs a pathological DGP/quantile combination and asserts the error fires.

    Note: the pre-v0.42.0 source also had a lower-bracket-sign-failed abort at solve_pq. That abort is dead code: at lo = y_min - margin < y_min, every 1{y <= lo} = 0, so the IPW score treated/m * 0 - q is -q < 0 for all q > 0. The lo_score >= 0.0 check never fires. The dead abort is removed in v0.42.0; the suberror type has only the reachable UpperSignFailed variant.

    CalibrationFittingError

    source
    pub suberror CalibrationFittingError {
      IncompleteCVPartition
    }

    Error type raised by isotonic_calibrate_cv when the provided cv partition does not cover every input index (some input rows are not in any fold's test_idx). This indicates a malformed cv partition passed by the caller.

    Added in v0.38.0 to convert the previous abort() call into a testable error path. The pre-v0.38.0 abort was reachable only via the (formerly private) isotonic_calibrate_cv helper; no public test exercised it. The new isotonic_calibrate_cv_raises_incomplete_partition test constructs a deliberately-malformed cv partition and asserts the error fires.

    ClusterDataError

    source
    pub suberror ClusterDataError {
      MissingUnit(Int)
    }

    Error type raised by build_row_unit_map when a row's unit id is not present in the uniq list (a malformed cluster vector — the cluster ids must form a subset of the unique unit ids). The payload is the missing unit id, which the caller can use to log or surface a meaningful diagnostic.

    Added in v0.36.0 to convert the previous abort() call into a testable error path. The pre-v0.36.0 panic_build_row_unit_map_missing_unit test was silently skipped on native/wasm-gc (MoonBit's panic_* driver skips panic-prefixed tests; see _verify/WHITEBOX_T_REPORT.md).

    DIDDataError

    source
    pub suberror DIDDataError {
      NonBinaryTreatment(Int)
    }

    Error type raised by DoubleMLDIDData::new when the treatment vector d is non-binary. The default port supports only the binary d ∈ {0, 1} convention (only the switchers; the multi-valued {-1, 0, 1} Sant'Anna & Zhao 2020 convention is not implemented). The payload is the index of the first non-binary entry, which the caller can use to log a meaningful diagnostic.

    Added in v0.43.0 to convert the previous abort() call into a testable error path. The pre-v0.43.0 abort was silently skipped on native/wasm-gc via the panic_* driver behavior (see _verify/WHITEBOX_T_REPORT.md). The new did_data_raises_on_non_binary_treatment test constructs a DoubleMLDIDData with a non-binary d entry and asserts the error fires.

    EmptyArrayError

    source
    pub suberror EmptyArrayError

    Error type raised by array_min and array_max when the input array is empty. No payload — the empty array case has no diagnostic detail to carry.

    Added in v0.41.0 to convert the previous abort() calls into testable error paths. The pre-v0.41.0 aborts were reachable only via internal helpers; no public test exercised them. The new array_min_raises_on_empty and array_max_raises_on_empty tests call the helpers directly with [] and assert the errors fire.

    InvalidCalibrationError

    source
    pub suberror InvalidCalibrationError {
      UnknownMethod(String)
    }

    Error type raised by apply_calibration (extracted from PSProcessor::adjust_ps in v0.37.0) when the configured calibration method is not one of the supported options ("none", "isotonic"). The payload is the unknown method name, which the caller can use to log or surface a meaningful diagnostic.

    Added in v0.37.0 to convert the previous abort() call inside adjust_ps into a testable error path. The panic_* test driver limitation had not been observed for this path because no test exercised it; the new test ps_processor_adjust_ps_raises_unknown_method constructs a config directly (bypassing PSProcessorConfig::new's require check) to reach the previously-unreachable abort.

    PSConfigError

    source
    pub suberror PSConfigError {
      InconsistentCVCalibration
    }

    Error type raised by PSProcessorConfig::new when the configuration is internally inconsistent — the only currently-reachable case is cv_calibration = true combined with calibration_method = "none", because CV calibration is only meaningful for the isotonic method. No payload — the call-site (the require checks for clipping/extreme thresholds and calibration_method string) is enough to identify the configuration error.

    Added in v0.44.0 to convert the previous abort() call into a testable error path. The pre-v0.44.0 panic_ps_processor_cv_without_calibration test was silently skipped on native/wasm-gc (MoonBit's panic_* driver skips panic-prefixed tests; see _verify/WHITEBOX_T_REPORT.md). The new ps_processor_config_raises_inconsistent_cv_calibration test calls PSProcessorConfig::new directly with the inconsistent config and asserts the error fires.

    PreconditionError

    source
    pub suberror PreconditionError {
      Violated(SourceLoc)
    }

    Error type reserved for the upcoming conversion of check.mbt::check and check.mbt::require from abort("precondition failed at ...") to a typed raise. The payload is the SourceLoc of the failing call site (auto-injected by #callsite(autofill(loc)) at every call), so the diagnostic message will point at the offending source line.

    Added in v0.47.0 as a planning release. v0.47.0 only declares the suberror type and ships a tiny check_make_violated helper (in check.mbt) that constructs a Violated(loc) for the regression test. The actual check/require → raise PreconditionError conversion is targeted for v0.48.0+ (cascade to all 324 pub functions, expected 1-2 release effort with try/catch/re-abort shims at every call site so the public abort behavior is preserved during the transition).

    VarEstClusterError

    source
    pub suberror VarEstClusterError {
      JTooSmall(Double, Double, Int)
    }

    Error type raised by var_est_cluster when the fold-mean J = mean(psi_deriv) lands below the 1e-6 floor (a defensive guard added in v0.34.0 against divide-by-near-zero inflations of the cluster SE; see the rationale comment in var_est_cluster below). The payload captures the exact (j, g, n_units) triple that triggered the floor so callers can log or surface a meaningful diagnostic.

    Callers that want to preserve the pre-v0.35.0 process-death behavior should match on this variant and re-abort; callers that want to retry or surface the error to downstream consumers should propagate it via ?.

    CART

    source
    pub enum CART {
      Leaf(Double)
      Split(Int, Double, CART, CART)
    } derive(
    Debug    )

    Immutable CART tree node. Split carries the feature column index, the threshold, and the two child subtrees. Leaf carries the per-leaf mean (the constant prediction for rows that land here).

    CART::to_repr

    source

    CSBinPanelRow

    source
    type CSBinPanelRow derive(
    Debug    )

    Per-row output of cs_bin_panel_subset. Each row is one observation in the post-subset panel: a (G_indicator, T_indicator) pair, the covariates, and the (still-raw) outcome y.

    CSBinPanelRow::to_repr

    source

    ConstantLearner

    source
    pub struct ConstantLearner {
      value : Double
    } derive(
    Debug    )

    Constant-predictor learner. fit is a no-op (returns self); predict returns a length-x.rows() array of the constant value for every row, ignoring x.

    ConstantLearner::fit

    source
    fn ConstantLearner::fit(self : ConstantLearner, x : Matrix, y : Array[Double]) -> ConstantLearner

    ConstantLearner::new

    source
    fn ConstantLearner::new(value : Double) -> ConstantLearner

    ConstantLearner::predict

    source
    fn ConstantLearner::predict(self : ConstantLearner, x : Matrix) -> Array[Double]

    ConstantLearner::to_repr

    source

    DIDAggregationResult

    source
    pub struct DIDAggregationResult {
      theta : Array[Double]
      se : Array[Double]
      agg_names : Array[String]
    } derive(
    Debug    )

    Aggregation result for a single axis (group, time, or event). theta is the weighted-mean ATT, se is the delta-method SE, agg_names[i] is the human-readable label (e.g. the group value, the time period, or the event-time value).

    DIDAggregationResult::to_repr

    source

    DidCsData

    source
    type DidCsData

    Result of the DID CS2021 DGP.

    DidCsData::d_get

    source
    fn DidCsData::d_get(self : DidCsData) -> Array[Double]

    Get the treatment vector.

    DidCsData::g_get

    source
    fn DidCsData::g_get(self : DidCsData) -> Array[Int]

    Get the group (cohort) index.

    DidCsData::id_get

    source
    fn DidCsData::id_get(self : DidCsData) -> Array[Int]

    Get the unit id (within cohort).

    DidCsData::n_groups_get

    source
    fn DidCsData::n_groups_get(self : DidCsData) -> Int

    Get the number of cohorts.

    DidCsData::n_periods_get

    source
    fn DidCsData::n_periods_get(self : DidCsData) -> Int

    Get the number of periods.

    DidCsData::t_get

    source
    fn DidCsData::t_get(self : DidCsData) -> Array[Int]

    Get the time-period index.

    DidCsData::theta_get

    source
    fn DidCsData::theta_get(self : DidCsData) -> Double

    Get the true ATT parameter.

    DidCsData::x_get

    source
    fn DidCsData::x_get(self : DidCsData) -> Matrix

    Get the covariate matrix (panel-level, n_obs x dim_x).

    DidCsData::y_get

    source
    fn DidCsData::y_get(self : DidCsData) -> Array[Double]

    Get the outcome vector.

    DidMultiData

    source
    type DidMultiData

    Result of the DID CS2021 multi-cohort DGP.

    DidMultiData::d_get

    source
    fn DidMultiData::d_get(self : DidMultiData) -> Array[Double]

    DidMultiData::g_get

    source
    fn DidMultiData::g_get(self : DidMultiData) -> Array[Int]

    DidMultiData::id_get

    source
    fn DidMultiData::id_get(self : DidMultiData) -> Array[Int]

    DidMultiData::n_groups_get

    source
    fn DidMultiData::n_groups_get(self : DidMultiData) -> Int

    DidMultiData::n_periods_get

    source
    fn DidMultiData::n_periods_get(self : DidMultiData) -> Int

    DidMultiData::t_get

    source
    fn DidMultiData::t_get(self : DidMultiData) -> Array[Int]

    DidMultiData::theta_get

    source
    fn DidMultiData::theta_get(self : DidMultiData) -> Double

    DidMultiData::x_get

    source
    fn DidMultiData::x_get(self : DidMultiData) -> Matrix

    DidMultiData::y_get

    source
    fn DidMultiData::y_get(self : DidMultiData) -> Array[Double]

    DoubleMLAPO

    source
    pub struct DoubleMLAPO {
      data : DoubleMLData
      treatment_level : Double
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      g_hat : Array[Double]
      m_hat : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
      psi_a : Array[Double]
      psi_b : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Average potential outcomes for an arbitrary treatment level.

    DoubleMLAPO::bootstrap

    source
    fn DoubleMLAPO::bootstrap(self : DoubleMLAPO, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLAPO

    v0.61.0+: multiplier bootstrap for DoubleMLAPO. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i] = -1 + theta * pb[i]

    where pb[i] = g[i] + treated[i] * (y[i] - g[i]) / m[i] is the IPW-centred potential-outcome score, computed at the fitted coef from the stored g_hat / m_hat. Draws n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution, and returns a fitted model with boot_t_stat[b] = sum_i w[b, i] * psi[i] /(sqrt(n) * se_psi) populated where se_psi = sqrt(mean(psi^2)).

    method_name selects the multiplier distribution:
    • "normal" (default): w[i] ~ N(0, 1).
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1).

    Calling bootstrap requires the model to be fitted; calling on an un-fit model aborts with PreconditionError. The helper is did_bootstrap_t_stat (v0.55.0 extracted from DoubleMLDIDCrossSection::bootstrap); APO is the n_thetas=1 case.

    DoubleMLAPO::coef

    source
    fn DoubleMLAPO::coef(self : DoubleMLAPO) -> Double

    DoubleMLAPO::confint

    source
    fn DoubleMLAPO::confint(self : DoubleMLAPO) -> (Double, Double)

    DoubleMLAPO::fit

    source

    DoubleMLAPO::learner_g

    source
    fn DoubleMLAPO::learner_g(self : DoubleMLAPO) -> LearnerDispatch

    Accessor for the outcome-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLAPO::learner_m

    source
    fn DoubleMLAPO::learner_m(self : DoubleMLAPO) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLAPO::n_features

    source
    fn DoubleMLAPO::n_features(self : DoubleMLAPO) -> Int

    Number of features (covariate columns).

    DoubleMLAPO::n_obs

    source
    fn DoubleMLAPO::n_obs(self : DoubleMLAPO) -> Int

    DoubleMLAPO::new

    source
    fn DoubleMLAPO::new(data : DoubleMLData, treatment_level? : Double, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLAPO

    DoubleMLAPO::predictions_g

    source
    fn DoubleMLAPO::predictions_g(self : DoubleMLAPO) -> Array[Double]

    DoubleMLAPO::predictions_m

    source
    fn DoubleMLAPO::predictions_m(self : DoubleMLAPO) -> Array[Double]

    DoubleMLAPO::se

    source
    fn DoubleMLAPO::se(self : DoubleMLAPO) -> Double

    DoubleMLAPO::to_repr

    source

    DoubleMLAPOS

    source
    pub struct DoubleMLAPOS {
      data : DoubleMLData
      treatment_levels : Array[Double]
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      coefs : Array[Double]
      ses : Array[Double]
      fitted : Bool
      boot_t_stat : Array[Array[Double]]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Average potential outcomes symmetric across multiple treatment levels. v0.49.0: full upstream parity — treatment_levels validation in new (rejects duplicates and levels not in data.d), the causal_contrast method (level-by-level delta and SE vs a reference level), and the treatment_levels() / n_treatment_levels() / fitted() accessors.

    Each treatment level is fit with the same fold partition (the parent does not currently route a shared partition to the child DoubleMLAPO; each child draws its own folds via kfold. v0.50.0+ plans to wire the parent through a fit_with_splits helper to share one stratified partition, but the v0.49.0 implementation is the v0.50.0 PR target) and the same closed-form LinearRegression learner. The causal_contrast(reference_levels) method then returns the level-by-level difference coefs[i] - coefs[ref] for one or more reference levels, matching the upstream DoubleMLAPOS.causal_contrast semantics.

    DoubleMLAPOS::boot_method

    source
    fn DoubleMLAPOS::boot_method(self : DoubleMLAPOS) -> String

    DoubleMLAPOS::boot_seed

    source
    fn DoubleMLAPOS::boot_seed(self : DoubleMLAPOS) -> Int

    DoubleMLAPOS::boot_t_stats

    source
    fn DoubleMLAPOS::boot_t_stats(self : DoubleMLAPOS) -> Array[Array[Double]]

    Per-level boot_t_stat accessors (v0.63.0+).

    DoubleMLAPOS::bootstrap

    source
    fn DoubleMLAPOS::bootstrap(self : DoubleMLAPOS, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLAPOS

    Per-level multiplier bootstrap for DoubleMLAPOS. Each treatment level's child DoubleMLAPO is re-fit (deterministic given seed) and its bootstrap(...) invoked; the per-level boot_t_stat arrays are concatenated into a length-n_levels Array[Array[Double]] and also written to self.boot_t_stat.

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild" — see DoubleMLAPO::bootstrap. n_rep_boot defaults to 500; seed defaults to 2024.

    DoubleMLAPOS::causal_contrast

    source
    fn DoubleMLAPOS::causal_contrast(self : DoubleMLAPOS, reference_levels : Array[Double]) -> Array[Array[Double]]

    v0.49.0: causal contrasts between the requested treatment levels and the supplied reference level(s). Returns one Array[Double] of length 2 * treatment_levels.length() - 1 per reference level: the ref-level slot is a single 0.0 (its contrast is trivially zero), and every other slot is a (delta, se) pair where delta = coefs[i] - coefs[ref_idx] and se = sqrt(se[i]^2 + se[ref_idx]^2). ref_idx is the position of the reference level in treatment_levels. The layout is interleaved ([0.0, delta_0, se_0, delta_1, se_1, ...] for a 2-level input) rather than a [(level, coef, se), ...] table — see apo_test.mbt::apos_causal_contrast_with_reference for the actual indices.

    The SE of each contrast is computed via the standard var(psi_a) + var(psi_b) - 2*cov(psi_a, psi_b) style approximation; for v0.49.0 we take the conservative sqrt(se_a^2 + se_b^2) route (same as the upstream causal_contrast summary table), which is exact when the per-level psi_a and psi_b are independent across levels (true under stratified kfold with disjoint train indices).

    DoubleMLAPOS::coefs

    source
    fn DoubleMLAPOS::coefs(self : DoubleMLAPOS) -> Array[Double]

    DoubleMLAPOS::confint

    source
    fn DoubleMLAPOS::confint(self : DoubleMLAPOS, level? : Double) -> Array[(Double, Double)]

    Per-level Wald confidence intervals (coef - z * se, coef + z * se) for DoubleMLAPOS. Returns one (lo, hi) tuple per treatment level, in user-supplied order. level defaults to 0.95 (z = 1.959963984540054 for the standard normal).

    DoubleMLAPOS::fit

    source

    DoubleMLAPOS::fitted

    source
    fn DoubleMLAPOS::fitted(self : DoubleMLAPOS) -> Bool

    v0.49.0: whether fit has been called.

    DoubleMLAPOS::learner_g

    source
    fn DoubleMLAPOS::learner_g(self : DoubleMLAPOS) -> LearnerDispatch

    Accessor for the outcome-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLAPOS::learner_m

    source
    fn DoubleMLAPOS::learner_m(self : DoubleMLAPOS) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLAPOS::n_rep_boot

    source
    fn DoubleMLAPOS::n_rep_boot(self : DoubleMLAPOS) -> Int

    DoubleMLAPOS::n_treatment_levels

    source
    fn DoubleMLAPOS::n_treatment_levels(self : DoubleMLAPOS) -> Int

    v0.49.0: number of requested treatment levels.

    DoubleMLAPOS::new

    source
    fn DoubleMLAPOS::new(data : DoubleMLData, treatment_levels : Array[Double], n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLAPOS

    DoubleMLAPOS::ses

    source
    fn DoubleMLAPOS::ses(self : DoubleMLAPOS) -> Array[Double]

    DoubleMLAPOS::to_repr

    source

    DoubleMLAPOS::treatment_levels

    source
    fn DoubleMLAPOS::treatment_levels(self : DoubleMLAPOS) -> Array[Double]

    v0.49.0: the requested treatment levels, in user-supplied order.

    DoubleMLBLP

    source
    pub struct DoubleMLBLP {
      basis : Matrix
      orth_signal : Array[Double]
      cov_type : String
      ml_g : LearnerDispatch
      coef : Array[Double]
      se : Array[Double]
      fitted : Bool
      n_obs : Int
      rss : Double
      var_y : Double
      residuals : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Best linear predictor of an orthogonal signal on a supplied basis.

    cov_type selects the standard-error convention. Valid values:
    • "HC0" (default): White's heteroskedasticity-consistent sandwich SE — robust to arbitrary residual heteroskedasticity. Matches the upstream doubleml.utils.blp call to statsmodels.OLS(cov_type='HC0').
    • "nonrobust": classic homoskedastic OLS SE (sigma^2 * (X^T X)^{-1} with sigma^2 = RSS / (n - p)), which is only valid under the homogeneous-error assumption.

    REVIEW L5: the field is currently stored on the struct for forward compatibility (post-fit introspection), but is not used outside fit. Construct via DoubleMLBLP::new to validate the value.

    v0.63.0+: ml_g? field accepts a LearnerDispatch. The coef / (X^T X)^{-1} machinery is always the closed-form OLS path (BLP is by definition the OLS projection); the dispatch learner only affects the residual / RSS computation. Defaults to LearnerDispatch::linear_regression() which preserves the v0.62.2 byte-equality behaviour for the standard fit.

    DoubleMLBLP::basis

    source
    fn DoubleMLBLP::basis(self : DoubleMLBLP) -> Matrix

    v0.22.0+: the basis matrix (the BLP's "design matrix"). Shape n_obs x p_features. Used by GainStatsSource::from_blp_cv to refit the BLP on each fold's training subset.

    DoubleMLBLP::bootstrap

    source
    fn DoubleMLBLP::bootstrap(self : DoubleMLBLP, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLBLP

    v0.64.0+: multiplier bootstrap for DoubleMLBLP. BLP is the OLS projection of orth_signal onto the augmented design [1, basis]; per-observation IF is psi[i, j] = M[j, :] @ xa_i * e_i where M = (Xa^T Xa + ridge I)^{-1} and xa_i = [1, basis[i, :]]. se_flat is the per-coefficient SE (HC0 sandwich diagonal under cov_type = "HC0", sqrt(sigma^2 * (X^T X)^{-1}_{jj}) under cov_type = "nonrobust"); the helper skips a coefficient when se[j] = 0. Routes through the shared generic_bootstrap_ols_per_coef helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    boot_t_stat is a flat [n_rep_boot * (p + 1)] array of t-statistics (row-major by rep, then by coef). Calling on an un-fit model aborts via PreconditionError.

    DoubleMLBLP::coef

    source
    fn DoubleMLBLP::coef(self : DoubleMLBLP) -> Array[Double]

    DoubleMLBLP::confint_joint

    source
    fn DoubleMLBLP::confint_joint(self : DoubleMLBLP, contrast : Matrix, level? : Double) -> Array[(Double, Double)]

    Joint confidence interval for the linear contrast contrast @ coef (length n_contrast), via chi-squared critical value on the Mahalanobis distance.

    (contrast @ (coef - theta))^T @ inv(Omega_contrast) @ (contrast @ (coef - theta)) ~ chi2(n_contrast)

    where Omega_contrast = contrast @ Omega @ contrast^T is the induced covariance. Equivalent to a Bonferroni-style worst-case bound but tighter for low correlation.

    Parameters:
    • contrast: row-major matrix (n_contrast, p + 1) whose rows define the linear functions of coef to interval-estimate.
    • level: confidence level in (0, 1). Default 0.95.

    Returns: array of (low, high) tuples (each row a symmetric interval around contrast @ coef). For a single contrast row, returns a 1-element array. v0.53.0-dev Task 3 / v0.11.4 upstream parity.

    Notes: the upstream joint-CI uses bootstrap to draw the critical value from np.quantile(np.max(np.abs(bootstrap))), which requires the full omega matrix. In our port we approximate the critical value with the chi-squared quantile (1 df per row), and use the diagonal sandwich form of Omega_contrast (contrast[r]^2 @ diag(se^2)) for the variance. This is the standard closed-form chi-squared joint CI when the joint-CI variance is dominated by the diagonal (a conservative approximation for the upstream bootstrap).

    DoubleMLBLP::fit

    source

    DoubleMLBLP::n_obs

    source
    fn DoubleMLBLP::n_obs(self : DoubleMLBLP) -> Int

    v0.19.0+: sample size used by the BLP fit.

    DoubleMLBLP::new

    source
    fn DoubleMLBLP::new(basis : Matrix, orth_signal : Array[Double], cov_type? : String, ml_g? : LearnerDispatch) -> DoubleMLBLP

    DoubleMLBLP::orth_signal

    source
    fn DoubleMLBLP::orth_signal(self : DoubleMLBLP) -> Array[Double]

    v0.22.0+: the orthogonal signal array (the BLP's "outcome" variable). Length n_obs. Used by GainStatsSource::from_blp_cv to compute the cross-fit residual variance.

    DoubleMLBLP::predictions

    source
    fn DoubleMLBLP::predictions(self : DoubleMLBLP) -> Array[Double]

    Fitted values basis_aug @ coef for each observation, length n_obs. Matches the upstream DoubleMLBLP.predictions / predict semantic — the orthogonal-signal prediction under the BLP coefficient vector.

    Note: coef has length p_basis + 1 because LinearRegression::fit adds an intercept column to basis. We reconstruct that column (all-1.0) on the fly to keep the stored self.basis unchanged.

    DoubleMLBLP::rss

    source
    fn DoubleMLBLP::rss(self : DoubleMLBLP) -> Double

    v0.19.0+: residual sum of squares from the BLP fit. Equals sum_i (orth_signal[i] - basis[i] @ coef)^2.

    DoubleMLBLP::se

    source
    fn DoubleMLBLP::se(self : DoubleMLBLP) -> Array[Double]

    DoubleMLBLP::to_repr

    source

    DoubleMLBLP::var_y

    source
    fn DoubleMLBLP::var_y(self : DoubleMLBLP) -> Double

    v0.19.0+: variance of the orthogonal signal (the BLP's "outcome" variable). Computed as a population variance (divisor n, not n - 1).

    DoubleMLBinaryData

    source
    pub struct DoubleMLBinaryData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
    } derive(
    Debug    )

    Double / debiased machine learning for the partially logistic regression model (upstream doubleml.plm.DoubleMLLPLR, Liu, Zhang, Zhou 2021):

     Y = expit(D * theta_0 + r_0(X)),  Y in {0, 1}

    The treatment D may be binary or continuous. The model uses a double cross-fit so the auxiliary regression t_0(X) = E[W | X] (where W = logit(M_0)) can be estimated with inner-fold OOF predictions without leaking information from the outer test set.

    Two scores are supported, exactly as upstream:
    • "nuisance_space": the per-fold ml_m training set is filtered to rows with Y = 0 (the "nuisance space" of the ATE on the treated subsample). The starting beta is computed inside the loop, one per fold, from the inner predictions.
    • "instrument": the inner predictions are used directly as M_inner, and ml_m accepts a sample_weight array of M (1 - M) from the inner fold. The starting beta is still per-fold.

    DoubleMLBinaryData::new

    source
    fn DoubleMLBinaryData::new(x : Matrix, y : Array[Double], d : Array[Double]) -> DoubleMLBinaryData

    DoubleMLBinaryData::to_repr

    source

    DoubleMLCVAR

    source
    pub struct DoubleMLCVAR {
      data : DoubleMLData
      treatment : Double
      quantile : Double
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      normalize_ipw : Bool
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      g_hat : Array[Double]
      m_hat : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
    } derive(
    Debug    )

    Conditional Value at Risk for a binary potential outcome, following the Kallus, Mao & Uehara (2024) "Removing Hidden Confounding by Supervised Gating" estimator (the upstream doubleml.irm.cvar.DoubleMLCVAR).

    The estimator solves the IPW score mean(1{d==treatment} / m(X) * 1{y <= theta} - quantile) = 0 per outer fold (with a per-fold preliminary cross-fit of m on the train side) to get a per-fold ipw_est[i], averages these for pq_est, and uses the cross-fitted (g, m) nuisances to evaluate the DML influence function psi_a = -1, psi_b = 1{d==treatment} * (g_target - g_hat) / m_hat + g_hat where g_target = max(pq_est, (y - q*pq_est) / (1-q)). The point estimate and SE come from the shared var_est(psi_a,psi_b) helper.

    v0.50.0: full upstream parity. The pre-v0.50.0 simplified version (which lived in quantile.mbt::DoubleMLCVAR and computed the CVaR via a single solve_pq call + a "max" target trick) has been removed; the canonical CVaR estimator is this struct.

    DoubleMLCVAR::coef

    source
    fn DoubleMLCVAR::coef(self : DoubleMLCVAR) -> Double

    Point estimate (the upper-tail conditional mean of Y(treatment)).

    DoubleMLCVAR::confint

    source
    fn DoubleMLCVAR::confint(self : DoubleMLCVAR) -> (Double, Double)

    95% Wald confidence interval (coef - 1.96 * se, coef + 1.96 * se).

    DoubleMLCVAR::fit

    source

    v0.50.0: fit the DoubleMLCVAR estimator. Implements the upstream _nuisance_est flow:
    1. For each rep r, draw n_folds outer folds with kfold(n, self.n_folds, self.seed + r).
    2. For each outer fold, run the inner cross-fit (see cvar_inner_crossfit) to get per-fold (g_hat,m_hat, ipw_est).
    3. After all folds, clip m_hat to [clip, 1-clip], optionally normalize the IPW weights, and (if treatment == 0) flip 1 - m_hat.
    4. Compute pq_est = mean(ipw_vec) and the final psi_a, psi_b with g_target = max(pq_est, (y - q*pq_est) / (1-q)).
    5. var_est(psi_a, psi_b) gives the per-rep (theta, se); aggregate across reps with aggregate_coef_se.

    DoubleMLCVAR::fitted

    source
    fn DoubleMLCVAR::fitted(self : DoubleMLCVAR) -> Bool

    v0.50.0: True iff fit has been called.

    DoubleMLCVAR::learner_g

    source
    fn DoubleMLCVAR::learner_g(self : DoubleMLCVAR) -> LearnerDispatch

    Accessor for the outcome-quantile-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLCVAR::learner_m

    source
    fn DoubleMLCVAR::learner_m(self : DoubleMLCVAR) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLCVAR::n_features

    source
    fn DoubleMLCVAR::n_features(self : DoubleMLCVAR) -> Int

    Number of features (covariate columns).

    DoubleMLCVAR::n_obs

    source
    fn DoubleMLCVAR::n_obs(self : DoubleMLCVAR) -> Int

    Number of observations.

    DoubleMLCVAR::new

    source
    fn DoubleMLCVAR::new(data : DoubleMLData, treatment? : Double, quantile? : Double, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, normalize_ipw? : Bool, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLCVAR

    Construct a DoubleMLCVAR estimator.

    Parameters mirror the upstream DoubleMLCVAR.__init__:
    • treatment : binary potential outcome to target (0 or 1; default 1)
    • quantile : upper-tail level q of the conditional value at risk (strictly in (0, 1); default 0.5)
    • n_folds : number of outer / inner folds (default 2, matching the rest of the package's IRM family)
    • n_rep : number of sample-splitting repetitions (default 1)
    • seed : PRNG seed for the outer fold partition (default 3141)
    • propensity_clip: lower / upper clip bound for the propensity m (strictly in (0, 0.5); default 1e-6)
    • normalize_ipw : if true, normalize the IPW weights so they sum to n within each treatment group, matching the upstream default (default true).

    DoubleMLCVAR::predictions_g

    source
    fn DoubleMLCVAR::predictions_g(self : DoubleMLCVAR) -> Array[Double]

    v0.50.0: cross-fitted g nuisance predictions (length n_obs). Only meaningful after fit.

    DoubleMLCVAR::predictions_m

    source
    fn DoubleMLCVAR::predictions_m(self : DoubleMLCVAR) -> Array[Double]

    v0.50.0: cross-fitted m (propensity) nuisance predictions (length n_obs). Values are already clipped to [propensity_clip, 1 - propensity_clip]. If treatment == 0, the values are flipped to 1 - m (the symmetric treatment swap from the upstream API). Only meaningful after fit.

    DoubleMLCVAR::se

    source
    fn DoubleMLCVAR::se(self : DoubleMLCVAR) -> Double

    Standard error (DML influence-function SE; see var_est).

    DoubleMLCVAR::to_repr

    source

    DoubleMLDID

    source
    pub struct DoubleMLDID {
      data : DoubleMLDIDData
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ps_processor : PSProcessor
      score : String
      in_sample_normalization : Bool
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      strata : Array[Int]
      g0_hat : Array[Double]
      g1_hat : Array[Double]
      m_hat : Array[Double]
      coef : Double
      se : Double
      psi_a : Array[Double]
      psi_b : Array[Double]
      fitted : Bool
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Double / debiased machine learning estimator for the difference in differences (DID) model with panel data and binary treatment d ∈ {0, 1}, following Sant'Anna and Zhao (2020).

    The reduced-form outcome equation is

     Y_post = g_0(0, X) + D * theta + U_post,    E[U | D, X] = 0
 Y_pre  = g_0(0, X) + U_pre,                  E[U_pre | X] = 0

    or, in first differences,

     dY = Y_post - Y_pre = g_1(1, X) - g_0(0, X) + D * theta + (U_post - U_pre)

    We form cross-fitted nuisance predictions

     g0(X) = E[Y | D = 0, X]      (trained on D = 0)
 g1(X) = E[Y | D = 1, X]      (trained on D = 1)
 m(X)  = P(D = 1 | X)         (trained on all obs; observational
                               only, clipped to [eps, 1 - eps])

    and use the observational DID score (the most general one):

     resid_d0 = Y - g0
 p_hat = mean(D)
 weight_psi_a   = D / p_hat
 weight_resid   = (D - m) / (p_hat * (1 - m))
 psi_b = (D - m)/(p_hat (1 - m)) * (Y - g0)        [the g1 term cancels in ATT]
 psi_a = -D / p_hat
 psi(theta) = theta * psi_a + psi_b

    Note: in the upstream code, when D = 1 is the only informative group, the ATT-style psi_b simplifies to psi_b = weight_resid_d0 * resid_d0 because psi_b_1 (the g1-based term) is zero under the ATT weighting. The port follows that simplification.

    Point estimate and variance

     theta_hat = -mean(psi_b) / mean(psi_a)
 J         = mean(psi_a)
 gamma     = mean(psi(theta_hat)^2)
 sigma2    = gamma / (J^2 * n)
 se        = sqrt(sigma2).

    The experimental score (score = "experimental") and the in-sample-normalisation variant (in_sample_normalization = true) are supported as of v0.8.0; they are required by DoubleMLDIDBinary's preprocessing wrapper. The staggered-DID variant (DoubleMLDIDCS) is still not implemented in this port.

    DoubleMLDID::bootstrap

    source
    fn DoubleMLDID::bootstrap(self : DoubleMLDID, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLDID

    v0.64.0+: multiplier bootstrap for DoubleMLDID. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i]

    where psi_a[i] = -d[i] / p_hat (or -1 for the experimental score) and psi_b[i] is the per-observation ATT score, both populated by fit(...) from the last rep's cross-fitted nuisances. Routes through the shared generic_bootstrap_t_stat helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    Calling on an un-fit model aborts via PreconditionError.

    DoubleMLDID::coef

    source
    fn DoubleMLDID::coef(self : DoubleMLDID) -> Double

    DoubleMLDID::confint

    source
    fn DoubleMLDID::confint(self : DoubleMLDID) -> (Double, Double)

    DoubleMLDID::fit

    source

    Run the DID estimation.

    DoubleMLDID::n_obs

    source
    fn DoubleMLDID::n_obs(self : DoubleMLDID) -> Int

    DoubleMLDID::new

    source
    fn DoubleMLDID::new(data : DoubleMLDIDData, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ps_processor? : PSProcessor, score? : String, in_sample_normalization? : Bool, strata? : Array[Int], ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLDID

    DoubleMLDID::predictions_g0

    source
    fn DoubleMLDID::predictions_g0(self : DoubleMLDID) -> Array[Double]

    DoubleMLDID::predictions_g1

    source
    fn DoubleMLDID::predictions_g1(self : DoubleMLDID) -> Array[Double]

    DoubleMLDID::predictions_m

    source
    fn DoubleMLDID::predictions_m(self : DoubleMLDID) -> Array[Double]

    DoubleMLDID::se

    source
    fn DoubleMLDID::se(self : DoubleMLDID) -> Double

    DoubleMLDID::to_repr

    source

    DoubleMLDIDBinary

    source
    pub struct DoubleMLDIDBinary {
      data : DoubleMLDIDBinaryData
      g_value : Int
      t_value_pre : Int
      t_value_eval : Int
      control_group : String
      anticipation_periods : Int
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ps_processor : PSProcessor
      score : String
      in_sample_normalization : Bool
      eval_idx : Array[Int]
      inner : DoubleMLDID
      fitted : Bool
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Binary DID model (Sant'Anna & Zhao 2020, §4.3) for panel data with binary treatment in terms of the (group, time) combination. Supports the two score variants "observational" (default, IPW-style) and "experimental" (A/B-test-style, requires independent treatment assignment), and the two weighting conventions "in_sample_normalization = true" (divide by sample mean) and false (divide by p_hat = mean(d)). The default (observational, false) matches the byte-equality of the pre-0.8.0 DoubleMLDID port on the canonical DGP.

    v0.10.0 additions (corresponds to upstream DoubleMLDIDCSBinary):
    • ps_processor field replaces the bare propensity_clip argument. The processor's clipping_threshold is used to bound the propensity scores inside the score denominator; the field-level propensity_clip is retained for backward compat but fit reads the clip threshold from ps_processor.clipping_threshold instead.
    • fit builds the wide-format strata G_indicator + 2 * t_indicator (matching upstream self._strata) and passes it to DoubleMLDID::new(strata=...), switching the inner sample-splitting to stratified_kfold. Each fold then balances the four (G, T) cells.

    The model is a thin wrapper around DoubleMLDID: it preprocesses the long-format panel into the wide-format DID dataset, then dispatches to DoubleMLDID::fit. The DML point estimate is the ATT (average treatment effect on the treated) for the chosen (g_value, t_value_pre, t_value_eval) triple.

    DoubleMLDIDBinary::bootstrap

    source
    fn DoubleMLDIDBinary::bootstrap(self : DoubleMLDIDBinary, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLDIDBinary

    v0.64.0+: multiplier bootstrap for DoubleMLDIDBinary. Delegates to the inner DoubleMLDID::bootstrap(...) (the wide-format ATT sub-model). The returned boot_t_stat is the inner's t-statistic array — same meaning, same shape, same RNG (the inner is re-bootstrapped on each call so the long-format row mapping in eval_idx is preserved).

    method_name / n_rep_boot / seed are passed through to the inner call. Calling on an un-fit model aborts via PreconditionError.

    DoubleMLDIDBinary::coef

    source
    fn DoubleMLDIDBinary::coef(self : DoubleMLDIDBinary) -> Double

    Point estimate (ATT) for the chosen (g_value, t_value_pre,t_value_eval) triple.

    DoubleMLDIDBinary::confint

    source
    fn DoubleMLDIDBinary::confint(self : DoubleMLDIDBinary) -> (Double, Double)

    95% Wald-style confidence interval.

    DoubleMLDIDBinary::fit

    source

    Run the binary DID estimation. Three steps:
    1. Preprocess the long-format panel into the wide-format DID dataset (preprocess_did_binary).
    2. Build a DoubleMLDIDData from the wide-format arrays.
    3. Fit a DoubleMLDID on the wide-format data with the chosen score and in_sample_normalization.

    The preprocessing keeps only units observed in both t_value_pre and t_value_eval, drops units with neither the G_indicator = (g == g_value) nor the C_indicator set, and uses the earlier period's covariates as the regressors.

    DoubleMLDIDBinary::inner_psi_a

    source
    fn DoubleMLDIDBinary::inner_psi_a(self : DoubleMLDIDBinary) -> Array[Double]

    v0.15.0+: per-observation psi_a on the cell's wide-format data (the inner DoubleMLDID's psi_a field). Length equals n_obs_subset() (the cell's wide-format panel size). Used by the multiplier bootstrap in DoubleMLDIDMulti::bootstrap to map wide-format psi back to the full long-format panel.

    DoubleMLDIDBinary::inner_psi_b

    source
    fn DoubleMLDIDBinary::inner_psi_b(self : DoubleMLDIDBinary) -> Array[Double]

    v0.15.0+: per-observation psi_b on the cell's wide-format data. See inner_psi_a for details.

    DoubleMLDIDBinary::n_features

    source
    fn DoubleMLDIDBinary::n_features(self : DoubleMLDIDBinary) -> Int

    Number of features (matches data.n_features()).

    DoubleMLDIDBinary::n_obs_subset

    source
    fn DoubleMLDIDBinary::n_obs_subset(self : DoubleMLDIDBinary) -> Int

    Number of units in the wide-format subset (i.e. observed in both t_value_pre and t_value_eval and assigned to either G or C).

    DoubleMLDIDBinary::new

    source
    fn DoubleMLDIDBinary::new(data : DoubleMLDIDBinaryData, g_value : Int, t_value_pre : Int, t_value_eval : Int, control_group? : String, anticipation_periods? : Int, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ps_processor? : PSProcessor, score? : String, in_sample_normalization? : Bool) -> DoubleMLDIDBinary

    DoubleMLDIDBinary::predictions_g0

    source
    fn DoubleMLDIDBinary::predictions_g0(self : DoubleMLDIDBinary) -> Array[Double]

    Cross-fitted control outcome predictions g_0(X) = E[Y | D = 0, X].

    DoubleMLDIDBinary::predictions_g1

    source
    fn DoubleMLDIDBinary::predictions_g1(self : DoubleMLDIDBinary) -> Array[Double]

    Cross-fitted treated outcome predictions g_1(X) = E[Y | D = 1, X].

    DoubleMLDIDBinary::predictions_m

    source
    fn DoubleMLDIDBinary::predictions_m(self : DoubleMLDIDBinary) -> Array[Double]

    Cross-fitted propensity predictions on the wide-format subset.

    DoubleMLDIDBinary::psi_a_long

    source
    fn DoubleMLDIDBinary::psi_a_long(self : DoubleMLDIDBinary) -> Array[Double]

    v0.15.0+: per-observation influence function component psi_a mapped from the cell's wide-format data back to the long-format panel. Length equals data.n_obs() (the long-format panel). For long-format rows that are not in the cell's wide-format subset, the value is 0.0. The influence function psi = psi_a + theta * psi_b is used by the multiplier bootstrap in DoubleMLDIDMulti::bootstrap.

    DoubleMLDIDBinary::psi_b_long

    source
    fn DoubleMLDIDBinary::psi_b_long(self : DoubleMLDIDBinary) -> Array[Double]

    v0.15.0+: per-observation influence function component psi_b mapped from the cell's wide-format data back to the long-format panel. See psi_a_long for details.

    DoubleMLDIDBinary::se

    source
    fn DoubleMLDIDBinary::se(self : DoubleMLDIDBinary) -> Double

    Standard error of the ATT point estimate.

    DoubleMLDIDBinary::to_repr

    source

    DoubleMLDIDBinaryData

    source
    pub struct DoubleMLDIDBinaryData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      t : Array[Int]
      g : Array[Int]
      id : Array[Int]
    } derive(
    Debug    )

    Panel data container for the binary DID model. Stores long-format panel observations: each row is (unit, time, y, d, x_1, ...,x_p). t_col and id_col identify the time and unit indices. g_col is the unit's treatment-group index (the period of first treatment; equal to never_treated_value for never-treated units). All arrays are length n_obs_total = n_units * n_periods.

    The fit step reshapes this long-format data into a wide-format DID dataset: for each unit observed in both t_value_pre and t_value_eval, we construct y_diff = y_post - y_pre, the G_indicator = (g == g_value), the C_indicator (never-treated or not-yet-treated per control_group), and the covariates from the earlier period. The wide-format dataset is then passed to DoubleMLDID for the standard DML cross-fit.

    DoubleMLDIDBinaryData::n_features

    source
    fn DoubleMLDIDBinaryData::n_features(self : DoubleMLDIDBinaryData) -> Int

    DoubleMLDIDBinaryData::n_obs

    source

    DoubleMLDIDBinaryData::new

    source
    fn DoubleMLDIDBinaryData::new(x : Matrix, y : Array[Double], d : Array[Double], t : Array[Int], g : Array[Int], id : Array[Int]) -> DoubleMLDIDBinaryData

    DoubleMLDIDBinaryData::to_repr

    source

    DoubleMLDIDCS

    source
    pub struct DoubleMLDIDCS {
      data : DoubleMLDIDCSData
      control_group : String
      anticipation_periods : Int
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ps_processor : PSProcessor
      in_sample_normalization : Bool
      coef_matrix : Array[Double]
      se_matrix : Array[Double]
      psi_matrix : Array[Double]
      n_groups : Int
      n_periods : Int
      fitted : Bool
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Callaway-Sant'Anna (2021) staggered DID model. Iterates over every (g, t_pre, t_eval) triple where g is a treatment group and t_eval is strictly after g, runs a DoubleMLDIDBinary on the long-format panel restricted to the never-treated control cohort, and stores the per-(g, t) ATT estimate, SE, and 95% CI.

    Simplifications vs. upstream:
    • Binary treatment {0, 1} only (the {-1, 0, 1} multi-valued d convention is not supported here; use DoubleMLDIDBinary with control_group = "not_yet_treated" for the staggered case).
    • Default control group is "never_treated". "not_yet_treated" is supported but not-yet-treated units are not "switched in" to be controls in subsequent periods (we always use the never-treated sentinel cohort as the single control).
    • Score is fixed to observational; in-sample normalisation is false (matches the DoubleMLDID default). Callers can pass in_sample_normalization = true to switch to the Sant'Anna & Zhao (2020) eq. 4.3 form.
    • No sensitivity analysis, no tune_optuna, no aggregation beyond the per-(g, t) output.

    DoubleMLDIDCS::bootstrap

    source
    fn DoubleMLDIDCS::bootstrap(self : DoubleMLDIDCS, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLDIDCS

    v0.64.0+: multiplier bootstrap for DoubleMLDIDCS. For each (g, t) cell that has a non-zero SE, materialises the per-cell IF row (psi_matrix[(gi * n_periods + pi) *n_obs + i_long]) into a flat [n_cells, n_obs] psi matrix and calls the shared generic_bootstrap_psi_matrix helper. boot_t_stat is a flat [n_rep_boot * n_cells] array, row-major by rep, then by cell — the same layout as DoubleMLDIDMulti::boot_t_stat. Pre-treatment or empty cells (se_matrix[k] == 0) are zeroed out by the helper.

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    Calling on an un-fit model aborts via PreconditionError.

    DoubleMLDIDCS::coef_at

    source
    fn DoubleMLDIDCS::coef_at(self : DoubleMLDIDCS, group_idx : Int, period_idx : Int) -> Double

    Per-(g, t) ATT estimate (row-major indexing). 0.0 if the (g, t) cell is empty (e.g. t_eval ≤ g for the chosen group).

    DoubleMLDIDCS::fit

    source

    Run the CS-DID estimation: for every (g, t_eval) triple with t_eval > g, restrict the panel to the never-treated cohort plus the units with g == g_value, run a DoubleMLDIDBinary on the restricted long-format data with t_value_pre = g and t_value_eval = t_eval, and store the resulting ATT and SE.

    The (g, t_eval) combinations are visited in row-major order (groups ascending, periods ascending). Cells with t_eval <= g are left at the default 0.0 (the CS-DID convention is "no pre-treatment estimate" for such cells; downstream aggregation layers can drop them via the nan-aware aggregator).

    DoubleMLDIDCS::group_at

    source
    fn DoubleMLDIDCS::group_at(self : DoubleMLDIDCS, group_idx : Int) -> Int

    Group value at row index i (sorted ascending).

    DoubleMLDIDCS::n_groups

    source
    fn DoubleMLDIDCS::n_groups(self : DoubleMLDIDCS) -> Int

    Number of groups (excluding the never-treated sentinel).

    DoubleMLDIDCS::n_periods

    source
    fn DoubleMLDIDCS::n_periods(self : DoubleMLDIDCS) -> Int

    Number of distinct time periods.

    DoubleMLDIDCS::new

    source
    fn DoubleMLDIDCS::new(data : DoubleMLDIDCSData, control_group? : String, anticipation_periods? : Int, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ps_processor? : PSProcessor, in_sample_normalization? : Bool) -> DoubleMLDIDCS

    DoubleMLDIDCS::period_at

    source
    fn DoubleMLDIDCS::period_at(self : DoubleMLDIDCS, period_idx : Int) -> Int

    Time value at column index t (sorted ascending).

    DoubleMLDIDCS::se_at

    source
    fn DoubleMLDIDCS::se_at(self : DoubleMLDIDCS, group_idx : Int, period_idx : Int) -> Double

    Per-(g, t) ATT standard error.

    DoubleMLDIDCS::to_repr

    source

    DoubleMLDIDCSBinary

    source
    pub struct DoubleMLDIDCSBinary {
      data : DoubleMLDIDCSData
      g_value : Int
      t_value_pre : Int
      t_value_eval : Int
      control_group : String
      anticipation_periods : Int
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ps_processor : PSProcessor
      score : String
      in_sample_normalization : Bool
      coef : Double
      se : Double
      psi_a : Array[Double]
      psi_b : Array[Double]
      g_d0_t0 : Array[Double]
      g_d0_t1 : Array[Double]
      g_d1_t0 : Array[Double]
      g_d1_t1 : Array[Double]
      m_hat : Array[Double]
      n_obs_subset : Int
      n_g_subset : Int
      n_c_subset : Int
      fitted : Bool
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    v0.51.0: Callaway-Sant'Anna (2021) DID with binary (group, time) outcomes. Estimates the ATT(g_value, t_value_eval) for a specific (g_value, t_value_pre, t_value_eval) triple via the Sant'Anna-Zhao (2020) binary-outcome DML score. Mirrors the upstream DoubleMLDIDCSBinary (Python 0.11.3) with the v0.51.0 simplifications listed in the file-level docstring.

    DoubleMLDIDCSBinary::bootstrap

    source
    fn DoubleMLDIDCSBinary::bootstrap(self : DoubleMLDIDCSBinary, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLDIDCSBinary

    v0.64.0+: multiplier bootstrap for DoubleMLDIDCSBinary. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i]

    where psi_a[i] and psi_b[i] are the per-row ATT-score elements (computed in fit on the post-subset panel). Routes through the shared generic_bootstrap_t_stat helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    Calling on an un-fit model aborts via PreconditionError.

    DoubleMLDIDCSBinary::coef

    source
    fn DoubleMLDIDCSBinary::coef(self : DoubleMLDIDCSBinary) -> Double

    v0.51.0: ATT estimate (point estimate) for the chosen (g_value, t_value_pre, t_value_eval) triple.

    DoubleMLDIDCSBinary::confint

    source
    fn DoubleMLDIDCSBinary::confint(self : DoubleMLDIDCSBinary, level? : Double) -> (Double, Double)

    v0.51.0: 95% Wald-style confidence interval.

    DoubleMLDIDCSBinary::fit

    source

    v0.51.0: run the CS Binary estimation. Three steps:
    1. Subset the panel to the 4 (G_indicator, T_indicator) cells.
    2. Cross-fit the 4 g-functions and the propensity using kfold_stratified on G_indicator + 2 * T_indicator.
    3. Compute psi_a, psi_b via the observational Sant'Anna-Zhao score, and derive the ATT and SE via the shared var_est(psi_a, psi_b) helper.

    DoubleMLDIDCSBinary::fitted

    source
    fn DoubleMLDIDCSBinary::fitted(self : DoubleMLDIDCSBinary) -> Bool

    v0.51.0: true after fit has been called, false otherwise.

    DoubleMLDIDCSBinary::g_value

    source
    fn DoubleMLDIDCSBinary::g_value(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: the g_value constructor argument (the treated group in the (g, t) DID setup).

    DoubleMLDIDCSBinary::n_c_subset

    source
    fn DoubleMLDIDCSBinary::n_c_subset(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: number of control (C_indicator = 1) rows in the post-subset panel.

    DoubleMLDIDCSBinary::n_g_subset

    source
    fn DoubleMLDIDCSBinary::n_g_subset(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: number of treated (G_indicator = 1) rows in the post-subset panel.

    DoubleMLDIDCSBinary::n_obs

    source
    fn DoubleMLDIDCSBinary::n_obs(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: number of observations in the post-subset panel (i.e. the effective sample size for the SE). This is the post-subset count, NOT the full panel n_obs. Matches the upstream DoubleMLDIDCSBinary.n_obs_subset attribute.

    DoubleMLDIDCSBinary::n_obs_panel

    source
    fn DoubleMLDIDCSBinary::n_obs_panel(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: total number of observations in the original (pre-subset) panel.

    DoubleMLDIDCSBinary::new

    source
    fn DoubleMLDIDCSBinary::new(data : DoubleMLDIDCSData, g_value : Int, t_value_pre : Int, t_value_eval : Int, control_group? : String, anticipation_periods? : Int, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ps_processor? : PSProcessor, score? : String, in_sample_normalization? : Bool) -> DoubleMLDIDCSBinary

    v0.51.0: construct a DoubleMLDIDCSBinary estimator. Validates the inputs and stores them on the struct; no estimation happens until fit is called. The score argument is restricted to "observational" for v0.51.0 (experimental is a v0.52.0+ target).

    DoubleMLDIDCSBinary::predictions_g_d0_t0

    source
    fn DoubleMLDIDCSBinary::predictions_g_d0_t0(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: cross-fitted nuisance predictions g_d0_t0 (control × pre). Length n_obs_subset().

    DoubleMLDIDCSBinary::predictions_g_d0_t1

    source
    fn DoubleMLDIDCSBinary::predictions_g_d0_t1(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: cross-fitted nuisance predictions g_d0_t1 (control × eval).

    DoubleMLDIDCSBinary::predictions_g_d1_t0

    source
    fn DoubleMLDIDCSBinary::predictions_g_d1_t0(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: cross-fitted nuisance predictions g_d1_t0 (treated × pre).

    DoubleMLDIDCSBinary::predictions_g_d1_t1

    source
    fn DoubleMLDIDCSBinary::predictions_g_d1_t1(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: cross-fitted nuisance predictions g_d1_t1 (treated × eval).

    DoubleMLDIDCSBinary::predictions_m

    source
    fn DoubleMLDIDCSBinary::predictions_m(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: cross-fitted propensity predictions m(X) = E[G_indicator | X]. Clipped to [propensity_clip, 1 - propensity_clip].

    DoubleMLDIDCSBinary::psi_a

    source
    fn DoubleMLDIDCSBinary::psi_a(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: per-observation psi_a from the CS Binary score. Length n_obs_subset().

    DoubleMLDIDCSBinary::psi_b

    source
    fn DoubleMLDIDCSBinary::psi_b(self : DoubleMLDIDCSBinary) -> Array[Double]

    v0.51.0: per-observation psi_b from the CS Binary score. Length n_obs_subset().

    DoubleMLDIDCSBinary::se

    source
    fn DoubleMLDIDCSBinary::se(self : DoubleMLDIDCSBinary) -> Double

    v0.51.0: standard error of the ATT estimate.

    DoubleMLDIDCSBinary::t_value_eval

    source
    fn DoubleMLDIDCSBinary::t_value_eval(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: the t_value_eval constructor argument (the evaluation period).

    DoubleMLDIDCSBinary::t_value_pre

    source
    fn DoubleMLDIDCSBinary::t_value_pre(self : DoubleMLDIDCSBinary) -> Int

    v0.51.0: the t_value_pre constructor argument (the baseline pre-treatment period).

    DoubleMLDIDCSBinary::to_repr

    source

    DoubleMLDIDCSData

    source
    pub struct DoubleMLDIDCSData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      t : Array[Int]
      id : Array[Int]
      g : Array[Int]
      groups : Array[Int]
      times : Array[Int]
    } derive(
    Debug    )

    Data container for the Callaway-Sant'Anna (CS) DID model with multi-period panel data and binary treatment. The treatment d is the change in treatment status (binary {0, 1} under the simplified CS-DID port; the upstream package also supports {-1, 0, 1} for "switchers" via the control_group parameter, which is out of scope here). t is the time index (length n_periods); each unit is observed in every period (n_obs =n_units * n_periods rows).

    The fit step iterates over the distinct treatment groups g and the evaluation periods t, runs a DoubleMLDIDBinary for each (g, t_pre, t_eval) triple on the long-format panel, and returns the per-(g, t) ATT estimates as a flat array (row-major, length n_groups * n_periods).

    DoubleMLDIDCSData::new

    source
    fn DoubleMLDIDCSData::new(x : Matrix, y : Array[Double], d : Array[Double], t : Array[Int], id : Array[Int], g : Array[Int]) -> DoubleMLDIDCSData

    DoubleMLDIDCSData::to_repr

    source

    DoubleMLDIDCrossSection

    source
    pub struct DoubleMLDIDCrossSection {
      data : DoubleMLDIDCrossSectionData
      n_folds : Int
      n_rep : Int
      seed : Int
      score : String
      in_sample_normalization : Bool
      propensity_clip : Double
      ps_processor : PSProcessor
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      coef : Double
      se : Double
      psi_a : Array[Double]
      psi_b : Array[Double]
      g_d0_t0 : Array[Double]
      g_d0_t1 : Array[Double]
      g_d1_t0 : Array[Double]
      g_d1_t1 : Array[Double]
      m_hat : Array[Double]
      fitted : Bool
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    v0.20.0+: Cross-section DID model with 4 g-functions and 1 propensity function. The score function matches the upstream doubleml.DoubleMLDIDCS._score_elements algorithm.

    score selects the score convention:
    • "observational" (default): the 2x2 DID setting where treatment assignment may be confounded by covariates. Uses the doubly-robust form with propensity reweighting.
    • "experimental": an A/B-test setting where treatment is independent of pre-treatment covariates. The propensity m collapses to a constant and the score simplifies.

    in_sample_normalization selects the in-sample vs out-of-sample normalization:
    • false (default): the canonical Sant'Anna-Zhao form, with weight_psi_a = d / p_hat.
    • true: in-sample normalization weight_psi_a = d / mean(d).

    DoubleMLDIDCrossSection::bootstrap

    source
    fn DoubleMLDIDCrossSection::bootstrap(self : DoubleMLDIDCrossSection, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLDIDCrossSection

    v0.21.0+: multiplier bootstrap. Draws n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution, computes boot_t_stat[b] = sum_i w[b, i] * psi[i] /(sqrt(n) * se) where psi[i] = psi_a[i] + theta * psi_b[i] is the per-observation influence function, and returns a fitted model with boot_t_stat populated.

    method_name selects the multiplier distribution:
    • "normal" (default): w[i] ~ N(0, 1). Matches the upstream bootstrap(method="normal") default.
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) /2 with x, y ~ N(0, 1). Robust to heteroskedasticity in the influence-function residuals.

    The bootstrap populates self.boot_t_stat and is required for confint(joint=true).

    DoubleMLDIDCrossSection::coef

    source

    ATT estimate from the cross-section DID fit.

    DoubleMLDIDCrossSection::confint

    source
    fn DoubleMLDIDCrossSection::confint(self : DoubleMLDIDCrossSection, joint? : Bool, level? : Double) -> (Double, Double)

    v0.20.0+ (extended in v0.21.0): confidence interval for the ATT.

    When joint = false (default), uses the Wald-style theta ± 1.96 * se interval (or the level-th quantile of the standard normal, scaled to the requested level).

    When joint = true, uses the multiplier bootstrap (must call bootstrap() first). The critical value is the empirical (1 + level) / 2 quantile of |boot_t_stat|, which is wider than the pointwise critical value (more conservative; matches the joint-CI convention for a scalar parameter).

    level is the confidence level (default 0.95).

    DoubleMLDIDCrossSection::fit

    source

    v0.20.0+: fit the cross-section DID model. Runs cross-fit nuisance estimation (4 g-functions + 1 propensity), constructs the psi_a / psi_b score, and returns the ATT and HC0 SE.

    DoubleMLDIDCrossSection::learner_g

    source

    Accessor for the outcome-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLDIDCrossSection::learner_m

    source

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLDIDCrossSection::new

    source
    fn DoubleMLDIDCrossSection::new(data : DoubleMLDIDCrossSectionData, n_folds? : Int, n_rep? : Int, seed? : Int, score? : String, in_sample_normalization? : Bool, propensity_clip? : Double, ps_processor? : PSProcessor, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLDIDCrossSection

    DoubleMLDIDCrossSection::predictions_g_d0_t0

    source
    fn DoubleMLDIDCrossSection::predictions_g_d0_t0(self : DoubleMLDIDCrossSection) -> Array[Double]

    v0.20.0+: nuisance predictions g(0, 0) (control, pre-period). Length n_obs.

    DoubleMLDIDCrossSection::predictions_g_d0_t1

    source
    fn DoubleMLDIDCrossSection::predictions_g_d0_t1(self : DoubleMLDIDCrossSection) -> Array[Double]

    v0.20.0+: nuisance predictions g(0, 1) (control, post-period). Length n_obs.

    DoubleMLDIDCrossSection::predictions_g_d1_t0

    source
    fn DoubleMLDIDCrossSection::predictions_g_d1_t0(self : DoubleMLDIDCrossSection) -> Array[Double]

    v0.20.0+: nuisance predictions g(1, 0) (treated, pre-period). Length n_obs.

    DoubleMLDIDCrossSection::predictions_g_d1_t1

    source
    fn DoubleMLDIDCrossSection::predictions_g_d1_t1(self : DoubleMLDIDCrossSection) -> Array[Double]

    v0.20.0+: nuisance predictions g(1, 1) (treated, post-period). Length n_obs.

    DoubleMLDIDCrossSection::predictions_m

    source
    fn DoubleMLDIDCrossSection::predictions_m(self : DoubleMLDIDCrossSection) -> Array[Double]

    v0.20.0+: propensity predictions m(x) = E[D=1 | X]. Length n_obs. Clipped to [propensity_clip, 1 -propensity_clip] after the cross-fit prediction averaging.

    DoubleMLDIDCrossSection::psi_a

    source

    v0.20.0+: per-observation psi_a from the cross-section DID score. Length n_obs. Throws if the model hasn't been fit.

    DoubleMLDIDCrossSection::psi_b

    source

    v0.20.0+: per-observation psi_b from the cross-section DID score. Length n_obs. Throws if the model hasn't been fit.

    DoubleMLDIDCrossSection::se

    source

    Standard error of the ATT estimate (HC0 sandwich).

    DoubleMLDIDCrossSection::to_repr

    source

    DoubleMLDIDCrossSectionData

    source
    pub struct DoubleMLDIDCrossSectionData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      t : Array[Int]
      name : String
    } derive(
    Debug    )

    Data container for the cross-section DID model (Sant'Anna & Zhao 2020, "repeated cross-sections" variant). Each unit has ONE observation, with covariates x, outcome y, binary treatment d (in {0, 1}), and binary post-period indicator t (in {0, 1}). The total sample size is n.

    Unlike the panel DID, there is no id column (cross-section = each unit is observed at most once) and no g column (no group structure — the "treated" group is just the d == 1 cohort and the "control" group is the d == 0 cohort). The model fits 4 g-functions g(d, t, x) = E[Y | D=d, T=t, X] for the 4 (d, t) combinations, plus 1 propensity function m(x) = E[D=1 | X], and constructs the ATT score function from these nuisance predictions.

    DoubleMLDIDCrossSectionData::n_features

    source

    Number of features in the cross-section DID data (excluding the intercept, which is added inside the learner).

    DoubleMLDIDCrossSectionData::n_obs

    source

    Sample size of the cross-section DID data.

    DoubleMLDIDCrossSectionData::new

    source
    fn DoubleMLDIDCrossSectionData::new(x : Matrix, y : Array[Double], d : Array[Double], t : Array[Int], name? : String) -> DoubleMLDIDCrossSectionData

    DoubleMLDIDCrossSectionData::to_repr

    source

    DoubleMLDIDData

    source
    pub struct DoubleMLDIDData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
    } derive(
    Debug    )

    Data container for DoubleMLDID. Uses panel data with two time periods. The treatment d is the change in treatment status between the post-period and the pre-period (so d ∈ {-1, 0, 1} for never-treated/always-treated/switchers under the Sant'Anna & Zhao (2020) convention). The default port supports binary d ∈ {0, 1} (only the switchers).

    DoubleMLDIDData::n_features

    source
    fn DoubleMLDIDData::n_features(self : DoubleMLDIDData) -> Int

    DoubleMLDIDData::n_obs

    source
    fn DoubleMLDIDData::n_obs(self : DoubleMLDIDData) -> Int

    DoubleMLDIDData::new

    source
    fn DoubleMLDIDData::new(x : Matrix, y : Array[Double], d : Array[Double]) -> DoubleMLDIDData raise DIDDataError

    Returns DoubleMLDIDData raise DIDDataError: the v0.43.0 conversion replaces the previous abort() call with raise DIDDataError::NonBinaryTreatment(i) so the non-binary-d path becomes directly testable. Callers that want the pre-v0.43.0 process-death behavior should catch the error and re-abort (this is what DoubleMLDIDBinaryData::new does); callers that want to surface the error to downstream consumers should propagate via ?. The require checks on x.nrows against y.length and d.length are unchanged and still abort the process (they are central require checks; refactoring them is out of scope for this surgical release).

    DoubleMLDIDData::to_repr

    source

    DoubleMLDIDMulti

    source
    pub struct DoubleMLDIDMulti {
      data : DoubleMLDIDCSData
      gt_combinations : Array[(Int, Int, Int)]
      control_group : String
      anticipation_periods : Int
      n_folds : Int
      n_rep : Int
      seed : Int
      ps_processor : PSProcessor
      in_sample_normalization : Bool
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      group_sizes : Array[Int]
      inner : DoubleMLDIDCS
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
      fitted : Bool
    } derive(
    Debug    )

    DoubleMLDIDMulti is a thin top-level wrapper over DoubleMLDIDCS that adds:

    1. A gt_combinations selector. Each combination is a (g_value, t_value_pre, t_value_eval) triple. The convenience keyword "standard" expands to "every post-treatment (g, t) with t_pre = g" (the default Callaway-Sant'Anna staggered set). "all" additionally includes pre-treatment cells for the event-study profile. "universal" is the same as "all" for the panel case (a cross-section-only setting in upstream is not ported; see the "Notes" section in CHANGELOG.md).
    2. aggregate_group, aggregate_time, aggregate_event methods that wrap the matching helpers in did_aggregation.mbt.

    DoubleMLDIDMulti::fit reuses DoubleMLDIDCS::fit to compute the per-(g, t) ATT matrix (and SE matrix), then hands the matrices to the aggregation helpers. The internal DoubleMLDIDCS instance is the canonical per-(g, t) estimator; DoubleMLDIDMulti adds the gt_combinations filter and the aggregation API on top.

    DoubleMLDIDMulti::aggregate_event

    source

    Aggregate by event time e = t - g. Returns a DIDAggregationResult with one entry per unique event time.

    DoubleMLDIDMulti::aggregate_group

    source

    Aggregate by group. Returns a DIDAggregationResult with one entry per group.

    DoubleMLDIDMulti::aggregate_time

    source

    Aggregate by time period. Returns a DIDAggregationResult with one entry per period.

    DoubleMLDIDMulti::bootstrap

    source
    fn DoubleMLDIDMulti::bootstrap(self : DoubleMLDIDMulti, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLDIDMulti

    v0.15.0+: multiplier bootstrap for joint confidence intervals. Draws n_rep_boot weight vectors from the chosen multiplier distribution and computes boot_t_stat[b, k] = sum_i w[b, i] * psi_k[i] / (sqrt(n) *se_k) for each bootstrap replication b and each (g, t) cell k. The joint confidence interval uses the empirical 95th percentile of max_k |boot_t_stat[b, k]| as the critical value; the per-cell Wald CI uses 1.96.

    method_name selects the multiplier distribution:
    • "normal": w[i] ~ N(0, 1) (default; matches the upstream bootstrap(method="normal") default).
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1). Robust to heteroskedasticity.

    seed controls the chacha8 RNG used to draw the weights (default 2024; matches the upstream numpy default of np.random.seed(2024) for the first test in _verify/test_bootstrap_reference.py).

    The bootstrap populates self.boot_t_stat and is required for confint(joint=true).

    DoubleMLDIDMulti::coef_at_idx

    source
    fn DoubleMLDIDMulti::coef_at_idx(self : DoubleMLDIDMulti, idx : Int) -> Double

    The inner per-(g, t) ATT at row-major index idx (matching the canonical gt_combinations[i] ordering).

    DoubleMLDIDMulti::confint

    source
    fn DoubleMLDIDMulti::confint(self : DoubleMLDIDMulti, joint? : Bool, level? : Double) -> Array[(Double, Double)]

    v0.15.0+: confidence interval for the per-(g, t) ATT. When joint = false (default), uses the Wald-style theta 卤 1.96 * se interval. When joint = true, uses the bootstrap: theta 卤 critical_value * se where critical_value is the 95th percentile of max_k |boot_t_stat[b, k]| across bootstrap replications b. Joint CIs are wider (more conservative) and require bootstrap() to be called first.

    level is the confidence level (default 0.95).

    DoubleMLDIDMulti::fit

    source

    Run the per-(g, t) cross-fits and store the per-cell ATT matrix for downstream aggregation. The actual per-cell fitting is delegated to DoubleMLDIDCS::fit, which already iterates over every (g, t) cell.

    DoubleMLDIDMulti::learner_g

    source

    Accessor for the outcome-nuisance learner forwarded to the inner DoubleMLDIDCS on the most recent fit(...) call. v0.60.0+.

    DoubleMLDIDMulti::learner_m

    source

    Accessor for the propensity-score learner forwarded to the inner DoubleMLDIDCS on the most recent fit(...) call. v0.60.0+.

    DoubleMLDIDMulti::n_combinations

    source
    fn DoubleMLDIDMulti::n_combinations(self : DoubleMLDIDMulti) -> Int

    Number of (g, t_pre, t_eval) triples.

    DoubleMLDIDMulti::new

    source
    fn DoubleMLDIDMulti::new(data : DoubleMLDIDCSData, gt_combinations? : Array[(Int, Int, Int)], gt_combinations_keyword? : String, control_group? : String, anticipation_periods? : Int, n_folds? : Int, n_rep? : Int, seed? : Int, ps_processor? : PSProcessor, in_sample_normalization? : Bool, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLDIDMulti

    Construct a DoubleMLDIDMulti.

    gt_combinations may be either an Array[(Int, Int, Int)] of explicit (g_value, t_value_pre, t_value_eval) triples, or one of the keywords:

    • "standard": every (g, t) with t > g and t_pre = g (the default Callaway-Sant'Anna staggered set).
    • "all": every (g, t) in the (groups 脳 periods) grid.
    • "universal": same as "all" for the panel case (upstream's universal is only meaningful for repeated cross sections, which this port does not implement).

    DoubleMLDIDMulti::p_adjust

    source
    fn DoubleMLDIDMulti::p_adjust(self : DoubleMLDIDMulti, method_name? : String) -> Array[Double]

    v0.16.0+: multiple-testing p-value adjustment for the per-(g, t) ATTs. Returns an Array[Double] of adjusted p-values (length n_combinations).

    Methods:
    • "romano-wolf" (default): the stepdown bootstrap procedure from Romano & Wolf (2005). For each cell k, sorted by descending |t_k|, compute p_k = mean_b [max_j > k |boot_t_stat[b, j]| >=|t_k|]. Then enforce monotonicity: p_corrected[k] = max(p_k, p_corrected[k - 1]) (in sorted order). Requires bootstrap() to have been called first.
    • "holm": Holm-Bonferroni stepdown (no bootstrap required). Sort unadjusted p-values ascending; for each k, p_corrected[k] = max((n - k) * p_sorted[k],p_corrected[k - 1]), then re-sort to original order.
    • "bonferroni": p_corrected[k] = n * p_k, clipped to 1.0. No bootstrap required.

    The p-values are computed from the Wald-style t-statistics (theta / se per cell) via the two-sided normal approximation.

    DoubleMLDIDMulti::p_values

    source
    fn DoubleMLDIDMulti::p_values(self : DoubleMLDIDMulti) -> Array[Double]

    v0.16.0+: per-cell unadjusted p-values for H0:theta = 0 (two-sided, normal approximation). Length n_combinations. pval[k] = 2 * (1 - norm.cdf(|t_k|)). Uses the standard-normal survival function on the Wald-style t-statistics from t_stats().

    The implementation uses Abramowitz & Stegun (1964) formula 7.1.26 for the normal CDF (max absolute error ~7.5e-8). MoonBit's @math does not expose erfc, so we approximate norm.cdf directly.

    DoubleMLDIDMulti::se_at_idx

    source
    fn DoubleMLDIDMulti::se_at_idx(self : DoubleMLDIDMulti, idx : Int) -> Double

    The inner per-(g, t) SE at row-major index idx.

    DoubleMLDIDMulti::t_stats

    source
    fn DoubleMLDIDMulti::t_stats(self : DoubleMLDIDMulti) -> Array[Double]

    v0.16.0+: per-cell t-statistics theta / se (length n_combinations). The Wald-style t-statistic matches the upstream all_t_stats[:, i_rep] (per repetition, but v0.15.0 is n_rep = 1 only). The Romano-Wolf stepdown p-adjustment in p_adjust consumes these t-statistics. Returns 0.0 for cells where se_k = 0 (pre-treatment / missing cells).

    DoubleMLDIDMulti::to_repr

    source

    DoubleMLData

    source
    pub struct DoubleMLData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      cluster_vars : Array[Int]
      z : Array[Double]
    } derive(
    Debug    )

    Container for the data consumed by a DML model. Mirrors the subset of doubleml.DoubleMLData that the PLR model needs: a feature matrix x, an outcome vector y, and a (possibly vector) treatment vector d. Multi-column treatment is currently treated as a single column vector in the port (the public DoubleMLPLR API takes a single treatment, matching the generate_data_simple example).

    When cluster_vars is non-empty, the model routes through the clustered DML path: folds are drawn over the unique unit ids in cluster_vars (not over individual rows), the causal parameter is the fold-weighted ratio of cluster score sums, and the variance is the unit-level cluster-robust estimator (_var_est's one-cluster-variable branch). Each row must carry exactly one cluster id; the length of cluster_vars is n_obs. This mirrors the upstream DoubleMLData(cluster_cols=) API in 0.11.x and is the path DoubleMLClusterData (the pre-0.11 backward-compat shim) used to take.

    DoubleMLData::is_cluster_data

    source
    fn DoubleMLData::is_cluster_data(self : DoubleMLData) -> Bool

    True iff the data is set up for clustered inference (a non-empty cluster_vars vector was passed to new).

    DoubleMLData::is_instrument_data

    source
    fn DoubleMLData::is_instrument_data(self : DoubleMLData) -> Bool

    True iff the data is set up for IV-type DML (a non-empty z vector was passed to new). The IV-type score on DoubleMLPLR::fit(score="IV-type") requires this; calling fit with score="IV-type" and no instrument aborts.

    DoubleMLData::n_cluster_vars

    source
    fn DoubleMLData::n_cluster_vars(self : DoubleMLData) -> Int

    Length of the cluster_vars vector (0 when not clustered).

    DoubleMLData::n_features

    source
    fn DoubleMLData::n_features(self : DoubleMLData) -> Int

    Number of features (columns of X).

    DoubleMLData::n_obs

    source
    fn DoubleMLData::n_obs(self : DoubleMLData) -> Int

    Number of observations in the data.

    DoubleMLData::new

    source
    fn DoubleMLData::new(x : Matrix, y : Array[Double], d : Array[Double], cluster_vars? : Array[Int], z? : Array[Double]) -> DoubleMLData

    Build a DoubleMLData from an n x p feature matrix, an outcome vector of length n and a treatment vector of length n. Pass a non-empty cluster_vars to enable the clustered DML path; pass [] (the default) for the standard row-level path.

    DoubleMLData::to_repr

    source

    DoubleMLData::z

    source
    fn DoubleMLData::z(self : DoubleMLData) -> Array[Double]

    v0.55.0: instrument vector Z (length n_obs) for the IV-type DML score on DoubleMLPLR. Returns an empty array when no instrument was passed to new.

    DoubleMLIIVM

    source
    pub struct DoubleMLIIVM {
      data : DoubleMLIIVMData
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      ml_r : LearnerDispatch
      g0_hat : Array[Double]
      g1_hat : Array[Double]
      m_hat : Array[Double]
      r0_hat : Array[Double]
      r1_hat : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
      psi_a : Array[Double]
      psi_b : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Double / debiased machine learning estimator for the interactive IV regression model (IIVM) of Chernozhukov et al. (2018) with the LATE score, identifying the Local Average Treatment Effect on the "compliers":

     Y = theta * D + g_0(D, X) + U,    E[U | D, X] = 0
 D = m_0(X, Z) + V,               E[V | X, Z] = 0

    where the binary instrument Z satisfies the relevance and exclusion restrictions. Five cross-fitted nuisance functions are needed (each estimated out-of-fold via K-fold):

     g0(X) = E[Y | Z = 0, X]      (trained only on Z = 0)
 g1(X) = E[Y | Z = 1, X]      (trained only on Z = 1)
 m(X)  = E[Z | X]             (trained on all obs, then
                               clipped to [eps, 1 - eps])
 r0(X) = E[D | Z = 0, X]      (trained only on Z = 0)
 r1(X) = E[D | Z = 1, X]      (trained only on Z = 1)

    Residuals:

     u_hat0 = Y - g0,  u_hat1 = Y - g1
 w_hat0 = D - r0,  w_hat1 = D - r1

    LATE score:

     psi_b =  (g1 - g0) + Z u_hat1 / m - (1 - Z) u_hat0 / (1 - m)
 psi_a = -(r1 - r0) - Z w_hat1 / m + (1 - Z) w_hat0 / (1 - m)
 psi(theta) = theta * psi_a + psi_b

    Point estimate and variance (same _var_est formula as the other DML models):

     theta_hat = -mean(psi_b) / mean(psi_a)
 J         = mean(psi_a)
 gamma     = mean(psi(theta_hat)^2)
 sigma2    = gamma / (J^2 * n)
 se        = sqrt(sigma2).

    DoubleMLIIVM::bootstrap

    source
    fn DoubleMLIIVM::bootstrap(self : DoubleMLIIVM, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLIIVM

    v0.61.0+: multiplier bootstrap for DoubleMLIIVM. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i]

    where psi_a and psi_b are the LATE score elements (defined in the module doc) computed at the fitted coef from the last rep's cross-fitted nuisances g0_hat / g1_hat / m_hat / r0_hat / r1_hat. Draws n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution, and returns a fitted model with boot_t_stat[b] = sum_i w[b, i] * psi[i] /(sqrt(n) * se_psi) populated where se_psi = sqrt(mean(psi^2)).

    method_name selects the multiplier distribution:
    • "normal" (default): w[i] ~ N(0, 1).
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1).

    Calling bootstrap requires the model to be fitted; calling on an un-fit model aborts with PreconditionError. The helper is did_bootstrap_t_stat (v0.55.0 extracted from DoubleMLDIDCrossSection::bootstrap); IIVM is the n_thetas=1 case.

    DoubleMLIIVM::coef

    source
    fn DoubleMLIIVM::coef(self : DoubleMLIIVM) -> Double

    DoubleMLIIVM::confint

    source
    fn DoubleMLIIVM::confint(self : DoubleMLIIVM) -> (Double, Double)

    DoubleMLIIVM::fit

    source
    fn DoubleMLIIVM::fit(self : DoubleMLIIVM, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch, ml_r? : LearnerDispatch, max_attempts? : Int) -> DoubleMLIIVM

    Run the IIVM estimation.

    Per-repetition behaviour: each repetition r cross-fits the g0 / g1 / m / r0 / r1 nuisances from its own folds (seed self.seed + r), computes its own (theta_r, se_r) from the LATE score, and the two arrays are then aggregated by aggregate_coef_se (median of thetas, then SE from the median of (theta_r + 1.96 * se_r)). For n_rep == 1 the aggregator returns the single (theta_1, se_1) exactly, so the byte-equality with the previous "average then estimate" implementation is preserved. The predictions_g0 / g1 / m / r0 / r1 accessors return the nuisances from the last repetition (the conventional choice in upstream doubleml), not a cross-rep average.

    DoubleMLIIVM::n_obs

    source
    fn DoubleMLIIVM::n_obs(self : DoubleMLIIVM) -> Int

    DoubleMLIIVM::new

    source
    fn DoubleMLIIVM::new(data : DoubleMLIIVMData, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch, ml_r? : LearnerDispatch) -> DoubleMLIIVM

    DoubleMLIIVM::predictions_g0

    source
    fn DoubleMLIIVM::predictions_g0(self : DoubleMLIIVM) -> Array[Double]

    DoubleMLIIVM::predictions_g1

    source
    fn DoubleMLIIVM::predictions_g1(self : DoubleMLIIVM) -> Array[Double]

    DoubleMLIIVM::predictions_m

    source
    fn DoubleMLIIVM::predictions_m(self : DoubleMLIIVM) -> Array[Double]

    DoubleMLIIVM::predictions_r0

    source
    fn DoubleMLIIVM::predictions_r0(self : DoubleMLIIVM) -> Array[Double]

    DoubleMLIIVM::predictions_r1

    source
    fn DoubleMLIIVM::predictions_r1(self : DoubleMLIIVM) -> Array[Double]

    DoubleMLIIVM::se

    source
    fn DoubleMLIIVM::se(self : DoubleMLIIVM) -> Double

    DoubleMLIIVM::to_repr

    source

    DoubleMLIIVMData

    source
    pub struct DoubleMLIIVMData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      z : Array[Double]
      cluster_vars : Array[Int]
    } derive(
    Debug    )

    Data container for DoubleMLIIVM. Same shape as DoubleMLData but adds a single instrumental variable z. The treatment d and the instrument z are both binary.

    DoubleMLIIVMData::is_cluster_data

    source
    fn DoubleMLIIVMData::is_cluster_data(self : DoubleMLIIVMData) -> Bool

    True iff the data is set up for clustered inference (a non-empty cluster_vars vector was passed to new).

    DoubleMLIIVMData::n_cluster_vars

    source
    fn DoubleMLIIVMData::n_cluster_vars(self : DoubleMLIIVMData) -> Int

    Length of the cluster_vars vector (0 when not clustered).

    DoubleMLIIVMData::n_features

    source
    fn DoubleMLIIVMData::n_features(self : DoubleMLIIVMData) -> Int

    DoubleMLIIVMData::n_obs

    source
    fn DoubleMLIIVMData::n_obs(self : DoubleMLIIVMData) -> Int

    DoubleMLIIVMData::new

    source
    fn DoubleMLIIVMData::new(x : Matrix, y : Array[Double], d : Array[Double], z : Array[Double], cluster_vars? : Array[Int]) -> DoubleMLIIVMData

    DoubleMLIIVMData::to_repr

    source

    DoubleMLIRM

    source
    pub struct DoubleMLIRM {
      data : DoubleMLData
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      g0_hat : Array[Double]
      g1_hat : Array[Double]
      m_hat : Array[Double]
      m_raw : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
      psi_a : Array[Double]
      psi_b : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Double / debiased machine learning estimator for the interactive regression model (IRM) of Chernozhukov et al. (2018) with the Average Treatment Effect (ATE) score:

     Y = g_0(D, X) + U,    E[U | D, X] = 0
 D = m_0(X) + V,        E[V | X] = 0

    and the ATE orthogonal signal

     g0(X) = E[Y | D=0, X]
 g1(X) = E[Y | D=1, X]
 m(X)  = P(D=1 | X)             (the propensity score)
 u0    = Y - g0(X)
 u1    = Y - g1(X)

 psi_b = (g1 - g0) + (D u1 / m - (1 - D) u0 / (1 - m))
 psi_a = -1
 psi(theta) = theta * psi_a + psi_b

    with point estimate

     theta_hat = -mean(psi_b) / mean(psi_a) = mean(psi_b)

    and variance (same _var_est formula as DoubleMLPLR):

     J       = mean(psi_a) = -1
 gamma   = mean(psi(theta_hat)^2)
 sigma2  = gamma / (J^2 * n)
 se      = sqrt(sigma2)

    The cross-fitting scheme trains g0 only on observations with D = 0 and g1 only on observations with D = 1 (the test folds still cover the full observation set), matching doubleml.utils._estimation._get_cond_smpls in the upstream package. The propensity score is clipped to [propensity_clip, 1 - propensity_clip] before being used in the score, guarding against near-zero or near-one values.

    DoubleMLIRM::bootstrap

    source
    fn DoubleMLIRM::bootstrap(self : DoubleMLIRM, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLIRM

    v0.61.0+: multiplier bootstrap for DoubleMLIRM. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i] = -1 + theta * psi_b[i]

    where psi_b[i] = (g1-g0)[i] + (D*u1/m - (1-D)*u0/(1-m))[i] is computed at the fitted coef from the last rep's cross-fitted nuisances g0_hat / g1_hat / m_hat. Draws n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution, and returns a fitted model with boot_t_stat[b] = sum_i w[b, i] * psi[i] /(sqrt(n) * se_psi) populated where se_psi = sqrt(mean(psi^2)).

    method_name selects the multiplier distribution:
    • "normal" (default): w[i] ~ N(0, 1).
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1).

    Calling bootstrap requires the model to be fitted; calling on an un-fit model aborts with PreconditionError. The helper is did_bootstrap_t_stat (v0.55.0 extracted from DoubleMLDIDCrossSection::bootstrap); IRM is the n_thetas=1 case.

    DoubleMLIRM::coef

    source
    fn DoubleMLIRM::coef(self : DoubleMLIRM) -> Double

    DoubleMLIRM::confint

    source
    fn DoubleMLIRM::confint(self : DoubleMLIRM) -> (Double, Double)

    DoubleMLIRM::fit

    source
    fn DoubleMLIRM::fit(self : DoubleMLIRM, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch, max_attempts? : Int) -> DoubleMLIRM

    Run the IRM estimation. The default learners are closed-form LinearRegression instances for both the outcome nuisance (ml_g) and the propensity score (ml_m).

    Per-repetition behaviour: each repetition r cross-fits the g0 / g1 / m nuisances from its own folds (seed self.seed + r), computes its own (theta_r, se_r) from the ATE score, and the two arrays are then aggregated by aggregate_coef_se (median of thetas, then SE from the median of (theta_r + 1.96 * se_r)). For n_rep == 1 the aggregator returns the single (theta_1, se_1) exactly, so the byte-equality with the previous "average then estimate" implementation is preserved. The predictions_g0 / g1 / m accessors return the nuisances from the last repetition (the conventional choice in upstream doubleml), not a cross-rep average.

    When self.data carries a non-empty cluster_vars vector, the estimator routes through the clustered DML path: folds partition whole units (rows of one cluster id stay on the same side of every split), the ATE is the fold-weighted ratio of cluster score sums, and the SE is unit-level cluster-robust.

    DoubleMLIRM::learner_g

    source
    fn DoubleMLIRM::learner_g(self : DoubleMLIRM) -> LearnerDispatch

    Accessor for the outcome-nuisance learner used by the most recent fit(...) / fit_cluster(...) call. v0.59.0+.

    DoubleMLIRM::learner_m

    source
    fn DoubleMLIRM::learner_m(self : DoubleMLIRM) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) / fit_cluster(...) call. v0.59.0+.

    DoubleMLIRM::n_features

    source
    fn DoubleMLIRM::n_features(self : DoubleMLIRM) -> Int

    Number of features (covariate columns).

    DoubleMLIRM::n_obs

    source
    fn DoubleMLIRM::n_obs(self : DoubleMLIRM) -> Int

    DoubleMLIRM::new

    source
    fn DoubleMLIRM::new(data : DoubleMLData, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLIRM

    DoubleMLIRM::predictions_g0

    source
    fn DoubleMLIRM::predictions_g0(self : DoubleMLIRM) -> Array[Double]

    DoubleMLIRM::predictions_g1

    source
    fn DoubleMLIRM::predictions_g1(self : DoubleMLIRM) -> Array[Double]

    DoubleMLIRM::predictions_m

    source
    fn DoubleMLIRM::predictions_m(self : DoubleMLIRM) -> Array[Double]

    DoubleMLIRM::propensity_score_raw

    source
    fn DoubleMLIRM::propensity_score_raw(self : DoubleMLIRM) -> Array[Double]

    Raw (pre-clipping) propensity scores, length n_obs. v0.53.0-dev Task 4: mirrors upstream feat "retain raw IRM propensity scores" so callers can inspect the un-clipped predictions without re-fitting.

    DoubleMLIRM::se

    source
    fn DoubleMLIRM::se(self : DoubleMLIRM) -> Double

    DoubleMLIRM::to_repr

    source

    DoubleMLLPLR

    source
    pub struct DoubleMLLPLR {
      data : DoubleMLBinaryData
      score : String
      n_folds : Int
      n_folds_inner : Int
      n_rep : Int
      seed : Int
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      coef_ : Double
      se_ : Double
      r_hat : Array[Double]
      m_hat : Array[Double]
      a_hat : Array[Double]
      fitted : Bool
    } derive(
    Debug    )

    Model: Y = expit(D * theta + r_0(X)). Binary outcomes, mixed treatment type. Wraps a single-treatment LPLR with closed-form logistic regression learners (LogisticRegression for ml_M and ml_m, LinearRegression for ml_t). The repository keeps one estimator per data; multi-treatment support is out of scope here.

    DoubleMLLPLR::coef

    source
    fn DoubleMLLPLR::coef(self : DoubleMLLPLR) -> Double

    DoubleMLLPLR::confint

    source
    fn DoubleMLLPLR::confint(self : DoubleMLLPLR) -> (Double, Double)

    DoubleMLLPLR::fit

    source

    Fit the LPLR estimator. Mirrors the upstream DoubleMLLPLR._nuisance_est / _est_causal_pars / _se_causal_pars flow:

    1. Outer cross-fit ml_M on (d, x) -> Y; the same outer folds are also used for the inner re-split. ml_M is a LogisticRegression so its predict returns the probability of Y = 1.
    2. Double cross-fit ml_a (also LogisticRegression for binary D) on x -> D to obtain inner-fold OOF propensity predictions per outer fold.
    3. Build per-fold targets W_inner from ml_M outer-OOF predictions (clipped to [1e-8, 1 - 1e-8], then logit); cross-fit ml_t (a LinearRegression) on x -> W_inner to get t_hat.
    4. Per fold, compute beta_f from inner W and (d - a_inner). Average across folds for the Newton starting value.
    5. Build the linear score: at any theta, r_i = D_i * theta + t_hat_i - beta_start * a_hat_i, then psi_i = d_tilde_i * (1 - Y_i) * exp(r_i)- expit(-r_i), and the derivative is -d_tilde_i * D_i * (1 - Y_i) * exp(r_i).
    6. Newton iteration with mean(psi) and mean(psi_deriv).
    7. Variance via the standard linear-score path (var_est(psi_a, psi_b) with psi_a = psi_deriv, psi_b = psi - theta * psi_deriv); the LPLR score is not linear in theta but is exactly the theta * psi_a +psi_b form at the converged theta with psi_a =psi_deriv(theta), psi_b = psi(theta) - theta *psi_deriv(theta), so the same machinery applies.

    LPLR is not cluster data; var_est runs the row-level path.

    DoubleMLLPLR::learner_g

    source
    fn DoubleMLLPLR::learner_g(self : DoubleMLLPLR) -> LearnerDispatch

    Accessor for the outcome-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLLPLR::learner_m

    source
    fn DoubleMLLPLR::learner_m(self : DoubleMLLPLR) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLLPLR::new

    source
    fn DoubleMLLPLR::new(data : DoubleMLBinaryData, score? : String, n_folds? : Int, n_folds_inner? : Int, n_rep? : Int, seed? : Int, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLLPLR

    DoubleMLLPLR::predictions_a

    source
    fn DoubleMLLPLR::predictions_a(self : DoubleMLLPLR) -> Array[Double]

    a_hat (E[D | X] propensity used in d_tilde) from the last repetition. In LPLR, m_hat == a_hat (a single nuisance regressor feeds both roles), but the field is preserved for parity with upstream's separate ml_m / ml_a learners.

    DoubleMLLPLR::predictions_m

    source
    fn DoubleMLLPLR::predictions_m(self : DoubleMLLPLR) -> Array[Double]

    m_hat (E[D | X], outer cross-fit, "a_hat" in upstream) from the last repetition. Length n_obs.

    DoubleMLLPLR::predictions_t

    source
    fn DoubleMLLPLR::predictions_t(self : DoubleMLLPLR) -> Array[Double]

    t_hat (E[W | X]) from the last repetition. Length n_obs.

    DoubleMLLPLR::se

    source
    fn DoubleMLLPLR::se(self : DoubleMLLPLR) -> Double

    DoubleMLLPLR::to_repr

    source

    DoubleMLLPQ

    source
    pub struct DoubleMLLPQ {
      data : DoubleMLLPQData
      treatment : Double
      quantile : Double
      n_folds : Int
      seed : Int
      propensity_clip : Double
      coef : Double
      se : Double
      fitted : Bool
      predictions_g0 : Array[Double]
      predictions_g1 : Array[Double]
      predictions_m : Array[Double]
      psi : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    DoubleMLLPQ::bootstrap

    source
    fn DoubleMLLPQ::bootstrap(self : DoubleMLLPQ, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLLPQ

    v0.64.0+: multiplier bootstrap for DoubleMLLPQ. The per-observation influence function psi is the centered score at the bisection root (mean 0, se_psi =sqrt(mean(psi^2)) is the bootstrap denominator). Routes through the shared generic_bootstrap_single_psi helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    Calling on an un-fit model aborts via PreconditionError.

    DoubleMLLPQ::coef

    source
    fn DoubleMLLPQ::coef(self : DoubleMLLPQ) -> Double

    DoubleMLLPQ::confint

    source
    fn DoubleMLLPQ::confint(self : DoubleMLLPQ) -> (Double, Double)

    95% Wald confidence interval (z = 1.959963984540054). Matches the DoubleMLLPLR::confint idiom byte-for-byte.

    DoubleMLLPQ::fit

    source

    DoubleMLLPQ::n_features

    source
    fn DoubleMLLPQ::n_features(self : DoubleMLLPQ) -> Int

    Number of features (covariate columns).

    DoubleMLLPQ::n_obs

    source
    fn DoubleMLLPQ::n_obs(self : DoubleMLLPQ) -> Int

    Number of observations.

    DoubleMLLPQ::new

    source
    fn DoubleMLLPQ::new(data : DoubleMLLPQData, treatment? : Double, quantile? : Double, n_folds? : Int, seed? : Int, propensity_clip? : Double) -> DoubleMLLPQ

    DoubleMLLPQ::predictions_g0

    source
    fn DoubleMLLPQ::predictions_g0(self : DoubleMLLPQ) -> Array[Double]

    Cross-fitted outcome nuisance for Z=0 (length n_obs). v0.53.0-dev: matches upstream DoubleMLLPQ.predictions["g0"].

    DoubleMLLPQ::predictions_g1

    source
    fn DoubleMLLPQ::predictions_g1(self : DoubleMLLPQ) -> Array[Double]

    Cross-fitted outcome nuisance for Z=1 (length n_obs). v0.53.0-dev: matches upstream DoubleMLLPQ.predictions["g1"].

    DoubleMLLPQ::predictions_m

    source
    fn DoubleMLLPQ::predictions_m(self : DoubleMLLPQ) -> Array[Double]

    Cross-fitted treatment nuisance (propensity score, length n_obs). v0.53.0-dev: matches upstream DoubleMLLPQ.predictions["m"].

    DoubleMLLPQ::se

    source
    fn DoubleMLLPQ::se(self : DoubleMLLPQ) -> Double

    DoubleMLLPQ::to_repr

    source

    DoubleMLLPQData

    source
    pub struct DoubleMLLPQData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      z : Array[Double]
    } derive(
    Debug    )

    DoubleMLLPQData::new

    source
    fn DoubleMLLPQData::new(x : Matrix, y : Array[Double], d : Array[Double], z : Array[Double]) -> DoubleMLLPQData

    DoubleMLLPQData::to_repr

    source

    DoubleMLPLIV

    source
    pub struct DoubleMLPLIV {
      data : DoubleMLPLIVData
      n_folds : Int
      n_rep : Int
      seed : Int
      learner : LearnerDispatch
      l_hat : Array[Double]
      r_hat : Array[Double]
      m_hat : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
      psi_a : Array[Double]
      psi_b : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Double / debiased machine learning estimator for the partially linear IV regression model (PLIV) of Chernozhukov et al. (2018) with the partialling out score:

     Y = D * theta_0 + g_0(X) + zeta,    E[zeta | D, X] = 0
 D = m_0(X) + V,                     E[V | X] = 0
 Z = ell_0(X) + xi,                  E[xi | X] = 0,
                                   Cov(Z, V) != 0  (relevance)

    with the partialling out (single-instrument) score

     l_hat = E_hat[Y | X]
 r_hat = E_hat[D | X]
 m_hat = E_hat[Z | X]
 u_hat = Y - l_hat
 w_hat = D - r_hat
 v_hat = Z - m_hat
 psi_a = -w_hat * v_hat
 psi_b =  v_hat * u_hat
 psi(theta) = theta * psi_a + psi_b

    with point estimate and variance

     theta_hat = -mean(psi_b) / mean(psi_a)
           = mean(v_hat * u_hat) / mean(w_hat * v_hat)
 J         = mean(psi_a)
 gamma     = mean(psi(theta_hat)^2)
 sigma2    = gamma / (J^2 * n)
 se        = sqrt(sigma2).

    The port supports only a single instrument and the partialling out score (the IV-type score, which would need an additional ml_g learner, is not implemented). All three nuisance functions are estimated with the same closed-form LinearRegression learner.

    DoubleMLPLIV::bootstrap

    source
    fn DoubleMLPLIV::bootstrap(self : DoubleMLPLIV, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLPLIV

    v0.61.0+: multiplier bootstrap for DoubleMLPLIV. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i] = -wv + theta * vu

    where w = d - r_hat, v = z - m_hat, u = y - l_hat are computed at the fitted coef from the last rep's cross-fitted nuisances l_hat / r_hat / m_hat. Draws n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution, and returns a fitted model with boot_t_stat[b] = sum_i w[b, i] * psi[i] /(sqrt(n) * se_psi) populated where se_psi = sqrt(mean(psi^2)).

    method_name selects the multiplier distribution:
    • "normal" (default): w[i] ~ N(0, 1).
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1).

    Calling bootstrap requires the model to be fitted; calling on an un-fit model aborts with PreconditionError. The helper is did_bootstrap_t_stat (v0.55.0 extracted from DoubleMLDIDCrossSection::bootstrap); PLIV is the n_thetas=1 case.

    DoubleMLPLIV::coef

    source
    fn DoubleMLPLIV::coef(self : DoubleMLPLIV) -> Double

    DoubleMLPLIV::confint

    source
    fn DoubleMLPLIV::confint(self : DoubleMLPLIV) -> (Double, Double)

    DoubleMLPLIV::fit

    source
    fn DoubleMLPLIV::fit(self : DoubleMLPLIV, learner? : LearnerDispatch, max_attempts? : Int) -> DoubleMLPLIV

    Run the PLIV estimation. The default learner is a closed-form LinearRegression; a different Learner can be supplied for experiments. The result is stored on the object and the object is returned for chaining.

    Per-repetition behaviour: each repetition r cross-fits the l / r / m nuisances from its own folds (seed self.seed + r), computes its own (theta_r, se_r) from the partialling-out single-instrument score, and the two arrays are then aggregated by aggregate_coef_se (median of thetas, then SE from the median of (theta_r + 1.96 * se_r)). For n_rep == 1 the aggregator returns the single (theta_1, se_1) exactly, so the byte-equality with the previous "average then estimate" implementation is preserved. The predictions_l / r / m accessors return the nuisances from the last repetition (the conventional choice in upstream doubleml), not a cross-rep average.

    DoubleMLPLIV::n_obs

    source
    fn DoubleMLPLIV::n_obs(self : DoubleMLPLIV) -> Int

    DoubleMLPLIV::new

    source
    fn DoubleMLPLIV::new(data : DoubleMLPLIVData, n_folds? : Int, n_rep? : Int, seed? : Int, learner? : LearnerDispatch) -> DoubleMLPLIV

    DoubleMLPLIV::predictions_l

    source
    fn DoubleMLPLIV::predictions_l(self : DoubleMLPLIV) -> Array[Double]

    DoubleMLPLIV::predictions_m

    source
    fn DoubleMLPLIV::predictions_m(self : DoubleMLPLIV) -> Array[Double]

    DoubleMLPLIV::predictions_r

    source
    fn DoubleMLPLIV::predictions_r(self : DoubleMLPLIV) -> Array[Double]

    DoubleMLPLIV::se

    source
    fn DoubleMLPLIV::se(self : DoubleMLPLIV) -> Double

    DoubleMLPLIV::to_repr

    source

    DoubleMLPLIVData

    source
    pub struct DoubleMLPLIVData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      z : Array[Double]
      cluster_vars : Array[Int]
    } derive(
    Debug    )

    Data container for DoubleMLPLIV. In addition to the covariates x, the outcome y and the treatment d, PLIV needs an instrumental variable z. The port supports a single instrument (1-D array) — the multi-instrument case is not implemented.

    DoubleMLPLIVData::is_cluster_data

    source
    fn DoubleMLPLIVData::is_cluster_data(self : DoubleMLPLIVData) -> Bool

    True iff the data is set up for clustered inference (a non-empty cluster_vars vector was passed to DoubleMLPLIVData::new).

    DoubleMLPLIVData::n_cluster_vars

    source
    fn DoubleMLPLIVData::n_cluster_vars(self : DoubleMLPLIVData) -> Int

    Length of the cluster_vars vector (0 when not clustered).

    DoubleMLPLIVData::n_features

    source
    fn DoubleMLPLIVData::n_features(self : DoubleMLPLIVData) -> Int

    Number of features.

    DoubleMLPLIVData::n_obs

    source
    fn DoubleMLPLIVData::n_obs(self : DoubleMLPLIVData) -> Int

    Number of observations.

    DoubleMLPLIVData::new

    source
    fn DoubleMLPLIVData::new(x : Matrix, y : Array[Double], d : Array[Double], z : Array[Double], cluster_vars? : Array[Int]) -> DoubleMLPLIVData

    Build a DoubleMLPLIVData from an n x p feature matrix, an outcome vector of length n, a treatment vector of length n and an instrument vector of length n.

    DoubleMLPLIVData::to_repr

    source

    DoubleMLPLPR

    source
    pub struct DoubleMLPLPR {
      panel : DoubleMLPanelData
      approach : String
      score : String
      n_folds : Int
      n_rep : Int
      seed : Int
      learner : LearnerDispatch
      coef_ : Double
      se_ : Double
      l_hat : Array[Double]
      m_hat : Array[Double]
      fitted : Bool
    } derive(
    Debug    )

    DoubleMLPLPR model. See the struct docs on DoubleMLPanelData and transform_panel for the four approaches. Scores: "partialling out" (default) or "IV-type" (fits an extra g-nuisance on y - theta_initial * d exactly like upstream).

    DoubleMLPLPR::coef

    source
    fn DoubleMLPLPR::coef(self : DoubleMLPLPR) -> Double

    DoubleMLPLPR::confint

    source
    fn DoubleMLPLPR::confint(self : DoubleMLPLPR) -> (Double, Double)

    DoubleMLPLPR::fit

    source
    fn DoubleMLPLPR::fit(self : DoubleMLPLPR, learner? : LearnerDispatch, max_attempts? : Int) -> DoubleMLPLPR

    Fit the DML estimator. Runs the standard cross-fit PO or IV-type score on the transformed data, including upstream's cre_general post-hoc m_hat adjustment and the clustered inference path (unit-level folds, fold-weighted coefficient, cluster-robust SE).

    DoubleMLPLPR::new

    source
    fn DoubleMLPLPR::new(panel : DoubleMLPanelData, approach? : String, score? : String, n_folds? : Int, n_rep? : Int, seed? : Int, learner? : LearnerDispatch) -> DoubleMLPLPR

    DoubleMLPLPR::predictions_l

    source
    fn DoubleMLPLPR::predictions_l(self : DoubleMLPLPR) -> Array[Double]

    Cross-fitted E[Y | X'] from the last repetition.

    DoubleMLPLPR::predictions_m

    source
    fn DoubleMLPLPR::predictions_m(self : DoubleMLPLPR) -> Array[Double]

    Cross-fitted E[D | X'] (after any cre adjustment) from the last repetition.

    DoubleMLPLPR::se

    source
    fn DoubleMLPLPR::se(self : DoubleMLPLPR) -> Double

    DoubleMLPLPR::to_repr

    source

    DoubleMLPLR

    source
    pub struct DoubleMLPLR {
      data : DoubleMLData
      learner_l : LearnerDispatch
      learner_m : LearnerDispatch
      n_folds : Int
      n_rep : Int
      seed : Int
      l_hat : Array[Double]
      m_hat : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
      psi_a : Array[Double]
      psi_b : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
      tune_result : TuneResult?
    } derive(
    Debug    )

    Double / debiased machine learning estimator for the partially linear regression model

     Y = D * theta_0 + g_0(X) + zeta,    E[zeta | D, X] = 0
 D = m_0(X) + V,                      E[V | X] = 0

    with the partialling out score

     psi_a(theta) = -(D - m_hat)^2,
 psi_b(theta) =  (D - m_hat) * (Y - l_hat),
 psi(theta)   =  theta * psi_a + psi_b

    where l_hat = E_hat[Y | X] and m_hat = E_hat[D | X] are obtained from a LinearRegression learner (or any other Learner) trained out-of-fold via K-fold cross-fitting.

    The point estimate is

     theta_hat = -mean(psi_b) / mean(psi_a)
           = mean((D - m_hat)(Y - l_hat)) / mean((D - m_hat)^2).

    The variance is estimated following doubleml.utils._estimation._var_est (non-cluster case):

     J  = mean(psi_a)              # expected derivative of psi w.r.t. theta
 gamma = mean(psi(theta_hat)^2)
 sigma2 = gamma / (J^2 * n)
 se = sqrt(sigma2).

    The implementation supports only the partialling out score and a single treatment. It is intentionally minimal — see the README for the matrix of features covered relative to the upstream package.

    DoubleMLPLR::bootstrap

    source
    fn DoubleMLPLR::bootstrap(self : DoubleMLPLR, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLPLR

    v0.61.0+: multiplier bootstrap for DoubleMLPLR. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i]

    (computed at the fitted coef from the last rep's cross-fitted nuisances l_hat / m_hat). Draws n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution, and returns a fitted model with boot_t_stat[b] = sum_i w[b, i] * psi[i] /(sqrt(n) * se_psi) populated where se_psi = sqrt(mean(psi^2)).

    method_name selects the multiplier distribution:
    • "normal" (default): w[i] ~ N(0, 1). Matches the upstream bootstrap(method="normal") default.
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1). Robust to heteroskedasticity.

    Calling bootstrap requires the model to be fitted; calling on an un-fit model aborts with PreconditionError. The helper is did_bootstrap_t_stat (v0.55.0 extracted from DoubleMLDIDCrossSection::bootstrap); PLR is the n_thetas=1 case with psi[i] = psi_a[i] + coef * psi_b[i].

    DoubleMLPLR::coef

    source
    fn DoubleMLPLR::coef(self : DoubleMLPLR) -> Double

    Fitted causal parameter.

    DoubleMLPLR::confint

    source
    fn DoubleMLPLR::confint(self : DoubleMLPLR) -> (Double, Double)

    95% Wald-style confidence interval [coef - 1.96*se, coef + 1.96*se].

    DoubleMLPLR::fit

    source
    fn DoubleMLPLR::fit(self : DoubleMLPLR, learner_l? : LearnerDispatch, learner_m? : LearnerDispatch, max_attempts? : Int, score? : String, tune_result? : TuneResult?) -> DoubleMLPLR

    The learner_l / learner_m labeled params accept any LearnerDispatch (defaults to the constructor-set values, i.e. a fresh LinearRegression). Per-fit overrides let callers swap in alternative nuisance estimators (ConstantLearner, NoopLearner, ...) without rebuilding the DoubleMLPLR instance.

    DoubleMLPLR::n_features

    source
    fn DoubleMLPLR::n_features(self : DoubleMLPLR) -> Int

    Number of features (covariate columns).

    DoubleMLPLR::n_obs

    source
    fn DoubleMLPLR::n_obs(self : DoubleMLPLR) -> Int

    Number of observations.

    DoubleMLPLR::new

    source
    fn DoubleMLPLR::new(data : DoubleMLData, learner_l? : LearnerDispatch, learner_m? : LearnerDispatch, n_folds? : Int, n_rep? : Int, seed? : Int) -> DoubleMLPLR

    DoubleMLPLR::predictions_l

    source
    fn DoubleMLPLR::predictions_l(self : DoubleMLPLR) -> Array[Double]

    Cross-fitted nuisance predictions for the outcome (length n).

    DoubleMLPLR::predictions_m

    source
    fn DoubleMLPLR::predictions_m(self : DoubleMLPLR) -> Array[Double]

    Cross-fitted nuisance predictions for the treatment (length n).

    DoubleMLPLR::se

    source
    fn DoubleMLPLR::se(self : DoubleMLPLR) -> Double

    Standard error of the causal parameter, computed via the DML variance formula.

    DoubleMLPLR::to_repr

    source

    DoubleMLPLR::tune

    source
    fn DoubleMLPLR::tune(self : DoubleMLPLR, param_set~ : Array[TuneParam], scoring_method? : String, n_folds_tune? : Int, seed? : Int, score? : String) -> DoubleMLPLR

    DoubleMLPLR::tune — score a grid of nuisance-learner combinations and re-fit the model with the winner.

    Algorithm (TUNE_DESIGN.md §4):

    1. Draw a fresh fold schedule folds_tune via kfold(n, n_folds_tune, seed). The tune-time folds are deliberately independent of the final-fit self.n_folds to avoid information leakage — the final DML estimate uses its own fold schedule for cross-fitting, so sharing tune and fit folds would leak l_hat quality into the model-selection step.
    2. For each candidate c in param_set:
      • l_hat_c = cross_fit_predict_dispatch(c.learner_l,x, y, folds_tune)
      • score_c = tune_score_outcome(y, l_hat_c, scoring)
      • on a learner-level exception, set score_c = +Inf (TUNE_DESIGN.md §7: keep tuning, the argmin rule automatically excludes divergent candidates).
    3. c* = argmin_c score_c (lower-is-better for MSE / RMSE, higher-is-better for NegMSE — NegMSE is converted to its lower-is-better form internally).
    4. Re-fit the DML model with learner_l = c*.learner_l,learner_m = c*.learner_m using the final-fit fold schedule (self.n_folds). The tuned nuisance predictions from step 1 are NOT reused — each re-fit produces its own OOF cross-fit under self.n_folds.
    5. Return the re-fitted DoubleMLPLR.

    Edge cases (TUNE_DESIGN.md §7):
    • param_set.length() == 0 -> abort with a descriptive message.
    • param_set.length() == 1 -> skip the scoring loop, re-fit with the single candidate directly (faster).
    • n_folds_tune > n_obs -> abort via the existing kfold precondition cascade.
    • seed not set -> defaults to 3141 for reproducibility.
    • scoring_method not in known set -> abort via TuneScoring::parse.

    score (the DML score: "partialling-out" / "iv-type") is forwarded to the re-fit step unchanged — tune() does not affect the DML score itself, only the nuisance learners.

    Cluster-data (self.data.is_cluster_data()) is not supported in v0.58.0 — cluster-aware tuning needs a different fold schedule (kfold on unique units) and would inflate this method beyond a single release. Use DoubleMLPLR::fit(learner_l = best_l, learner_m = best_m) directly on the cluster-DML path with hand-picked candidates until v0.59.0+ lands cluster-aware tune.

    DoubleMLPLR::tune_result

    source
    fn DoubleMLPLR::tune_result(self : DoubleMLPLR) -> TuneResult?

    Accessor for the tune_result populated by DoubleMLPLR::tune(...) (v0.58.0+). Returns None for models built via DoubleMLPLR::new(...) or re-fit via DoubleMLPLR::fit(...) (a re-fit discards the prior tune history because the nuisance learners may have changed).

    Inspect the returned TuneResult to see which (learner_l,learner_m) combination won under the chosen scoring rule, what the best score was, and the per-candidate score vector (in param_set order).

    DoubleMLPQ

    source
    pub struct DoubleMLPQ {
      data : DoubleMLData
      treatment : Double
      quantile : Double
      n_folds : Int
      seed : Int
      propensity_clip : Double
      ml_l : LearnerDispatch
      ml_m : LearnerDispatch
      coef : Double
      se : Double
      fitted : Bool
      psi : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    DoubleMLPQ::bootstrap

    source
    fn DoubleMLPQ::bootstrap(self : DoubleMLPQ, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLPQ

    v0.64.0+: multiplier bootstrap for DoubleMLPQ. The per-observation influence function psi is the centered IPW quantile score at the bisection root (mean 0, se_psi = sqrt(mean(psi^2)) is the bootstrap denominator). Routes through the shared generic_bootstrap_single_psi helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    Calling on an un-fit model aborts via PreconditionError.

    DoubleMLPQ::coef

    source
    fn DoubleMLPQ::coef(self : DoubleMLPQ) -> Double

    DoubleMLPQ::confint

    source
    fn DoubleMLPQ::confint(self : DoubleMLPQ) -> (Double, Double)

    DoubleMLPQ::fit

    source
    fn DoubleMLPQ::fit(self : DoubleMLPQ, ml_l? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLPQ

    DoubleMLPQ::learner_l

    source
    fn DoubleMLPQ::learner_l(self : DoubleMLPQ) -> LearnerDispatch

    Accessor for the outcome-quantile-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLPQ::learner_m

    source
    fn DoubleMLPQ::learner_m(self : DoubleMLPQ) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLPQ::n_features

    source
    fn DoubleMLPQ::n_features(self : DoubleMLPQ) -> Int

    Number of features (covariate columns).

    DoubleMLPQ::n_obs

    source
    fn DoubleMLPQ::n_obs(self : DoubleMLPQ) -> Int

    Number of observations.

    DoubleMLPQ::new

    source
    fn DoubleMLPQ::new(data : DoubleMLData, treatment? : Double, quantile? : Double, n_folds? : Int, seed? : Int, propensity_clip? : Double, ml_l? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLPQ

    DoubleMLPQ::se

    source
    fn DoubleMLPQ::se(self : DoubleMLPQ) -> Double

    DoubleMLPQ::to_repr

    source

    DoubleMLPanelData

    source
    pub struct DoubleMLPanelData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      t : Array[Int]
      id : Array[Int]
    } derive(
    Debug    )

    Double / debiased machine learning for the static panel partially linear regression model (upstream: doubleml.plm.DoubleMLPLPR, Clarke & Polselli 2025):

     Y_it = D_it * theta_0 + g_0(X_it) + alpha_i + zeta_it

    where alpha_i is a unit fixed effect. The panel structure is first transformed into a cross-section via one of four static panel approaches, and a standard DML partialling-out / IV-type score is applied to the transformed data:

    • "cre_general": correlated random effects (Mundlak). Augment X with per-unit means of every covariate column. After cross-fitting, adjust the treatment nuisance: m_hat* = m_hat + d_mean - mean_by_id(m_hat) where d_mean is the per-row unit mean of D.
    • "cre_normal": CRE with a normality-style restriction. Same X augmentation, but the treatment regression gets [X, d_mean] as inputs; no post-hoc m_hat adjustment.
    • "fd_exact": exact first differencing. The panel is reindexed to the full id x time grid, y and d are first-differenced within each unit, and covariates become [X_t, X_{t-1}]. Rows missing x_t or x_{t-1} are dropped.
    • "wg_approx": approximate within transformation. within(v) = v - unit_mean(v) + grand_mean(v) for y, d, and every covariate; PLR runs on the within-transformed variables alone.

    All four approaches reuse the shared cross-fit machinery (kfold + cross_fit_predict). Because upstream re-wraps the transformed data as static-panel data with cluster_cols =id_col, estimation always takes the clustered DML path: folds partition whole units, the coefficient is a fold-weighted ratio of cluster score sums (est_coef_cluster), and the SE is unit-level cluster-robust (var_est_cluster), with the closed-form LinearRegression learner.

    DoubleMLPanelData::new

    source
    fn DoubleMLPanelData::new(x : Matrix, y : Array[Double], d : Array[Double], t : Array[Int], id : Array[Int]) -> DoubleMLPanelData

    DoubleMLPanelData::to_repr

    source

    DoubleMLPolicyTree

    source
    pub struct DoubleMLPolicyTree {
      features : Matrix
      orth_signal : Array[Double]
      depth : Int
      root : PolicyTreeNode
      split_feature : Int
      split_value : Double
      left_treatment : Int
      right_treatment : Int
      fitted : Bool
    } derive(
    Debug    )

    A compact policy tree. It searches one split per level using weighted variance-reduction gain (Bug #8); the depth parameter controls how many levels of recursion to use (TODO #11c.3: previously the depth field was unused and the fit was always a depth-1 stump).

    DoubleMLPolicyTree::fit

    source

    DoubleMLPolicyTree::new

    source
    fn DoubleMLPolicyTree::new(features : Matrix, orth_signal : Array[Double], depth? : Int) -> DoubleMLPolicyTree

    DoubleMLPolicyTree::predict

    source
    fn DoubleMLPolicyTree::predict(self : DoubleMLPolicyTree, x : Matrix) -> Array[Int]

    DoubleMLPolicyTree::split_feature

    source
    fn DoubleMLPolicyTree::split_feature(self : DoubleMLPolicyTree) -> Int

    DoubleMLPolicyTree::split_value

    source
    fn DoubleMLPolicyTree::split_value(self : DoubleMLPolicyTree) -> Double

    DoubleMLPolicyTree::to_repr

    source

    DoubleMLQTE

    source
    pub struct DoubleMLQTE {
      data : DoubleMLData
      quantiles : Array[Double]
      n_folds : Int
      seed : Int
      propensity_clip : Double
      ml_l : LearnerDispatch
      ml_m : LearnerDispatch
      coefs : Array[Double]
      ses : Array[Double]
      psi_flat : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    DoubleMLQTE::bootstrap

    source
    fn DoubleMLQTE::bootstrap(self : DoubleMLQTE, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLQTE

    v0.64.0+: multiplier bootstrap for DoubleMLQTE. The per-quantile per-observation influence function is u[i, j] = psi_d1[i] / deriv_d1 - psi_d0[i] / deriv_d0 (the delta-method IF for theta_qte = theta_d1 - theta_d0); psi_flat is the flat [n_quantiles * n_obs] row-major matrix that fit populates. Routes through the shared generic_bootstrap_psi_matrix helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    boot_t_stat is a flat [n_rep_boot * n_quantiles] array (row-major by rep, then by quantile). Calling on an un-fit model aborts via PreconditionError.

    DoubleMLQTE::coefs

    source
    fn DoubleMLQTE::coefs(self : DoubleMLQTE) -> Array[Double]

    DoubleMLQTE::fit

    source

    DoubleMLQTE::learner_l

    source
    fn DoubleMLQTE::learner_l(self : DoubleMLQTE) -> LearnerDispatch

    Accessor for the outcome-quantile-nuisance learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLQTE::learner_m

    source
    fn DoubleMLQTE::learner_m(self : DoubleMLQTE) -> LearnerDispatch

    Accessor for the propensity-score learner used by the most recent fit(...) call. v0.60.0+.

    DoubleMLQTE::n_features

    source
    fn DoubleMLQTE::n_features(self : DoubleMLQTE) -> Int

    Number of features (covariate columns).

    DoubleMLQTE::n_obs

    source
    fn DoubleMLQTE::n_obs(self : DoubleMLQTE) -> Int

    Number of observations.

    DoubleMLQTE::new

    source
    fn DoubleMLQTE::new(data : DoubleMLData, quantiles? : Array[Double], n_folds? : Int, seed? : Int, propensity_clip? : Double, ml_l? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLQTE

    DoubleMLQTE::ses

    source
    fn DoubleMLQTE::ses(self : DoubleMLQTE) -> Array[Double]

    DoubleMLQTE::to_repr

    source

    DoubleMLRDD

    source
    pub struct DoubleMLRDD {
      data : DoubleMLRDDData
      cutoff : Double
      bandwidth : Double
      fuzzy : Bool
      cov_type : String
      ml_g : LearnerDispatch
      coef : Double
      se : Double
      n_local : Int
      fitted : Bool
    } derive(
    Debug    )

    Local-polynomial RD estimator. The port uses a triangular kernel and a fixed bandwidth, which keeps the hot path pure MoonBit and deterministic.

    DoubleMLRDD::coef

    source
    fn DoubleMLRDD::coef(self : DoubleMLRDD) -> Double

    DoubleMLRDD::confint

    source
    fn DoubleMLRDD::confint(self : DoubleMLRDD) -> (Double, Double)

    DoubleMLRDD::fit

    source

    DoubleMLRDD::learner_g

    source
    fn DoubleMLRDD::learner_g(self : DoubleMLRDD) -> LearnerDispatch

    Accessor for the local-polynomial learner used by the most recent fit(...) call. v0.60.0+. Note: RDD's bandwidth / kernel are configuration (not LearnerDispatch concepts); this slot plugs into the closed-form OLS underneath the kernel weights and is forward-compatible with v0.61.0+ kernel-aware wrappers.

    DoubleMLRDD::n_features

    source
    fn DoubleMLRDD::n_features(self : DoubleMLRDD) -> Int

    Number of features (covariate columns).

    DoubleMLRDD::n_local

    source
    fn DoubleMLRDD::n_local(self : DoubleMLRDD) -> Int

    DoubleMLRDD::n_obs

    source
    fn DoubleMLRDD::n_obs(self : DoubleMLRDD) -> Int

    Number of observations.

    DoubleMLRDD::new

    source
    fn DoubleMLRDD::new(data : DoubleMLRDDData, cutoff? : Double, bandwidth? : Double, fuzzy? : Bool, cov_type? : String, ml_g? : LearnerDispatch) -> DoubleMLRDD

    DoubleMLRDD::se

    source
    fn DoubleMLRDD::se(self : DoubleMLRDD) -> Double

    DoubleMLRDD::to_repr

    source

    DoubleMLRDDData

    source
    pub struct DoubleMLRDDData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      score : Array[Double]
    } derive(
    Debug    )

    Data for a sharp or fuzzy regression-discontinuity design.

    DoubleMLRDDData::n_obs

    source
    fn DoubleMLRDDData::n_obs(self : DoubleMLRDDData) -> Int

    DoubleMLRDDData::new

    source
    fn DoubleMLRDDData::new(x : Matrix, y : Array[Double], d : Array[Double], score : Array[Double]) -> DoubleMLRDDData

    DoubleMLRDDData::to_repr

    source

    DoubleMLSSM

    source
    pub struct DoubleMLSSM {
      data : DoubleMLSSMData
      n_folds : Int
      n_rep : Int
      seed : Int
      propensity_clip : Double
      ml_g : LearnerDispatch
      ml_m : LearnerDispatch
      pi_hat : Array[Double]
      m_hat : Array[Double]
      g_d1 : Array[Double]
      g_d0 : Array[Double]
      coef : Double
      se : Double
      fitted : Bool
      psi_a : Array[Double]
      psi_b : Array[Double]
      boot_t_stat : Array[Double]
      boot_method : String
      n_rep_boot : Int
      boot_seed : Int
    } derive(
    Debug    )

    Double / debiased machine learning estimator for the Sample Selection Model (SSM) of Bia, Huber and Laffers (2023), under the Missing At Random (MAR) score with normalize_ipw = False.

    The model is

     Y = theta * D + X @ beta * D + U,           E[U | D, X] = 0
 S = 1{D + gamma Z + X @ beta + V > 0},     E[V | X, D] = 0

    Y is observed only when S = 1. The four cross-fitted nuisances:

     g_d1(X) = E[Y | D = 1, S = 1, X]      (trained on D=1 ∧ S=1,
                                        features = X only —
                                        Bug #1 fix: previously
                                        `pi_hat` was appended as
                                        an extra feature, which
                                        leaked the test-fold pi
                                        into the training fold)
 g_d0(X) = E[Y | D = 0, S = 1, X]      (trained on D=0 ∧ S=1,
                                        features = X only)
 m(X)    = P(D = 1 | X)                 (trained on all obs,
                                        clipped to [eps, 1 - eps])
 pi(X, D) = P(S = 1 | D, X)            (trained on (X, D),
                                        clipped to [eps, 1 - eps])

    Score (un-normalized IPW):

     psi_a = -1
 psi_b1 = (D == 1) * S * (Y - g_d1) / (m * pi) + g_d1
 psi_b0 = (D == 0) * S * (Y - g_d0) / ((1 - m) * pi) + g_d0
 psi_b  = psi_b1 - psi_b0

    Point estimate and variance

     theta_hat = -mean(psi_b) / mean(psi_a) = mean(psi_b)
 J         = mean(psi_a) = -1
 gamma     = mean(psi(theta_hat)^2)
 sigma2    = gamma / (J^2 * n)
 se        = sqrt(sigma2).

    Note: only the MAR case is implemented; normalize_ipw = True and the nonignorable-nonresponse case are out of scope.

    DoubleMLSSM::bootstrap

    source
    fn DoubleMLSSM::bootstrap(self : DoubleMLSSM, method_name? : String, n_rep_boot? : Int, seed? : Int) -> DoubleMLSSM

    v0.64.0+: multiplier bootstrap for DoubleMLSSM. The per-observation influence function is

    psi[i] = psi_a[i] + theta * psi_b[i]

    where psi_a[i] = -1 (constant) and psi_b[i] is the MAR-score element defined in the file header. Routes through the shared generic_bootstrap_t_stat helper (see bootstrap_helper.mbt).

    method_name selects the multiplier distribution: "normal" (default), "Bayes", "wild". seed defaults to 2024; n_rep_boot defaults to 500.

    Calling on an un-fit model aborts via PreconditionError.

    DoubleMLSSM::coef

    source
    fn DoubleMLSSM::coef(self : DoubleMLSSM) -> Double

    DoubleMLSSM::confint

    source
    fn DoubleMLSSM::confint(self : DoubleMLSSM) -> (Double, Double)

    DoubleMLSSM::fit

    source
    fn DoubleMLSSM::fit(self : DoubleMLSSM, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch, ml_pi? : LearnerDispatch) -> DoubleMLSSM

    Run the SSM estimation.

    DoubleMLSSM::n_obs

    source
    fn DoubleMLSSM::n_obs(self : DoubleMLSSM) -> Int

    DoubleMLSSM::new

    source
    fn DoubleMLSSM::new(data : DoubleMLSSMData, n_folds? : Int, n_rep? : Int, seed? : Int, propensity_clip? : Double, ml_g? : LearnerDispatch, ml_m? : LearnerDispatch) -> DoubleMLSSM

    DoubleMLSSM::predictions_g_d0

    source
    fn DoubleMLSSM::predictions_g_d0(self : DoubleMLSSM) -> Array[Double]

    DoubleMLSSM::predictions_g_d1

    source
    fn DoubleMLSSM::predictions_g_d1(self : DoubleMLSSM) -> Array[Double]

    DoubleMLSSM::predictions_m

    source
    fn DoubleMLSSM::predictions_m(self : DoubleMLSSM) -> Array[Double]

    DoubleMLSSM::predictions_pi

    source
    fn DoubleMLSSM::predictions_pi(self : DoubleMLSSM) -> Array[Double]

    DoubleMLSSM::psi_a

    source
    fn DoubleMLSSM::psi_a(self : DoubleMLSSM) -> Array[Double]

    v0.64.0+: per-observation IF accessors.

    DoubleMLSSM::psi_b

    source
    fn DoubleMLSSM::psi_b(self : DoubleMLSSM) -> Array[Double]

    DoubleMLSSM::se

    source
    fn DoubleMLSSM::se(self : DoubleMLSSM) -> Double

    DoubleMLSSM::to_repr

    source

    DoubleMLSSMData

    source
    pub struct DoubleMLSSMData {
      x : Matrix
      y : Array[Double]
      d : Array[Double]
      s : Array[Double]
    } derive(
    Debug    )

    Data container for DoubleMLSSM. Adds a binary selection indicator s on top of the usual x/y/d. The outcome y is observed only when s = 1; for s = 0 the y entry should be ignored.

    DoubleMLSSMData::n_features

    source
    fn DoubleMLSSMData::n_features(self : DoubleMLSSMData) -> Int

    DoubleMLSSMData::n_obs

    source
    fn DoubleMLSSMData::n_obs(self : DoubleMLSSMData) -> Int

    DoubleMLSSMData::new

    source
    fn DoubleMLSSMData::new(x : Matrix, y : Array[Double], d : Array[Double], s : Array[Double]) -> DoubleMLSSMData

    DoubleMLSSMData::to_repr

    source

    Fold

    source
    pub struct Fold {
      train_idx : Array[Int]
      test_idx : Array[Int]
    } derive(
    Debug    )

    A single fold is represented by a pair of train and test index arrays. Test indices are disjoint and together cover [0, n_obs).

    Fold::new

    source
    fn Fold::new(train_idx : Array[Int], test_idx : Array[Int]) -> Fold

    Public constructor: build a fold from explicit train/test index arrays (used by callers that drive cross-fitting with their own partitions, e.g. the clustered panel path in plpr.mbt).

    Fold::test_indices

    source
    fn Fold::test_indices(self : Fold) -> Array[Int]

    Test indices (accessor for the test_idx private field).

    Fold::to_repr

    source

    Fold::train_indices

    source
    fn Fold::train_indices(self : Fold) -> Array[Int]

    Train indices (accessor for the train_idx private field).

    GBLearner

    source
    pub struct GBLearner {
      n_trees : Int
      learning_rate : Double
      max_depth : Int
      min_samples_leaf : Int
      mtry : Int
      subsample : Double
      bootstrap_seed : Int
      initial : Double
      trees : Array[CART]
    } derive(
    Debug    )

    Gradient boosting learner (regression).

    GBLearner::fit

    source
    fn GBLearner::fit(self : GBLearner, x : Matrix, y : Array[Double]) -> GBLearner

    Promote Learner trait methods as explicit methods on GBLearner so the trait impl below doesn't trip unused_value under --deny-warn. Same pattern as the v0.54.0 ConstantLearner / NoopLearner and v0.56.0 RFLearner setup.

    GBLearner::initial

    source
    fn GBLearner::initial(self : GBLearner) -> Double

    Initial constant prediction (mean of y at fit time). Returns 0.0 before fit.

    GBLearner::n_trees

    source
    fn GBLearner::n_trees(self : GBLearner) -> Int

    Number of trees in the fitted ensemble. Returns 0 if not yet fit.

    GBLearner::new

    source
    fn GBLearner::new(n_trees? : Int, learning_rate? : Double, max_depth? : Int, min_samples_leaf? : Int, mtry? : Int, subsample? : Double, bootstrap_seed? : Int) -> GBLearner

    GBLearner::predict

    source
    fn GBLearner::predict(self : GBLearner, x : Matrix) -> Array[Double]

    Promote Learner trait methods as explicit methods on GBLearner so the trait impl below doesn't trip unused_value under --deny-warn. Same pattern as the v0.54.0 ConstantLearner / NoopLearner and v0.56.0 RFLearner setup.

    GBLearner::to_repr

    source

    GainStatsResult

    source
    pub struct GainStatsResult {
      cf_y : Array[Double]
      cf_d : Array[Double]
      rho : Array[Double]
      delta_theta : Array[Double]
    } derive(
    Debug    )

    Container for the per-coefficient gain-statistic benchmark values. Returned by gain_statistics. The four fields are
    • cf_y (length n_coef): the maximum percentage of the outcome's residual variance explainable by an unobserved confounder that explains the difference between dml_long and dml_short's residual variance. Used as the upper bound on cf_y in sensitivity_analysis.
    • cf_d (length n_coef): the maximum percentage gain in the Riesz representer's variance explainable by an unobserved confounder. Used as the upper bound on cf_d.
    • rho (length n_coef): the sign-and-magnitude of the confounding correlation (in [-1, 1]) that drives the coefficient change from dml_long to dml_short.
    • delta_theta (length n_coef): the per-coefficient change coef_short - coef_long (median over reps).

    All four are length n_coef = dml_long.coef.length(). The benchmark is the upstream doubleml.utils.gain_statistics output; see Cinelli & Hazlett (2020) §3.4 for the interpretation of cf_y / cf_d / rho.

    GainStatsResult::to_repr

    source

    GainStatsSource

    source
    pub struct GainStatsSource {
      var_y_residuals : Array[Double]
      nu2 : Array[Double]
      all_coef : Array[Double]
      n_rep : Int
      var_y : Double
    } derive(
    Debug    )

    Container exposing the per-rep arrays gain_statistics needs from a fitted DML model. The upstream API is the DoubleML object; the v0.17.0 port defines a minimal struct so any DML estimator (BLP, PolicyTree, PLR, IRM, ...) can be benchmarked. var_y_residuals, nu2, and all_coef are row-major (n_coef, n_rep); n_rep is the per-coefficient repetition count; var_y is the scalar outcome variance.

    GainStatsSource::from_blp

    source
    fn GainStatsSource::from_blp(blp : DoubleMLBLP, n_rep? : Int) -> GainStatsSource

    v0.19.0+ convenience constructor for GainStatsSource from a fitted DoubleMLBLP. Auto-populates all the per-rep arrays from the BLP's fit output:

    • var_y_residuals[k] = RSS / n_obs (constant across coefficients; the BLP's residual variance).
    • nu2[k] = var_y_residuals[k] / (n_obs * se[k]^2) (the per-coef Riesz representer norm squared under the homoskedastic OLS convention se[k]^2 = sigma^2 * (Z^T Z)^{-1}_{kk}).
    • all_coef[k] = blp.coef()[k].
    • var_y = blp.var_y() (the variance of the BLP's orthogonal signal — the BLP's "outcome" variable).

    The HC0 SE convention is consistent with this homoskedastic interpretation up to O(1/n) corrections (the BLP's HC0 SE is robust to heteroskedasticity in the orthogonal-signal residuals; the auto-populated nu2 uses the BLP's reported SE directly).

    n_rep defaults to 1 (single-rep BLP). Multi-rep DMLs should pass n_rep > 1, in which case the auto-populated arrays are broadcast across reps (the BLP does not natively produce per-rep sensitivity elements).

    GainStatsSource::from_blp_cv

    source
    fn GainStatsSource::from_blp_cv(blp : DoubleMLBLP, n_folds? : Int, seed? : Int) -> GainStatsSource

    v0.22.0+: from_blp_cv(blp, n_folds?, seed?) is the cross-fit variant of from_blp. The only difference is var_y_residuals, which is computed from out-of-fold (OOF) predictions rather than the in-sample BLP residuals. The OOF residual variance is honest (no leakage from the basis fit on the same rows), so the R2_y benchmark in gain_statistics is more accurate.

    Algorithm:
    1. Draw n_folds random folds via kfold (with chacha8_rng-style seeding through seed).
    2. For each fold, fit a LinearRegression on the training rows and predict on the test fold.
    3. Compute the per-fold test residual variance sigma2_fold = sum_i (y_i - y_hat_i)^2 /n_fold (the honest OOF estimate).
    4. var_y_residuals_scalar = mean over folds of sigma2_fold (the "average fold residual variance").

    coef, se, var_y, and all_coef are unchanged from from_blp (the BLP's own fit on the full data is the canonical coefficient estimate; only var_y_residuals is recomputed).

    nu2 is recomputed as var_y_residuals_cv / (n_obs * se^2) (same homoskedastic convention as from_blp).

    GainStatsSource::from_blp_cv_repeated

    source
    fn GainStatsSource::from_blp_cv_repeated(blp : DoubleMLBLP, n_folds? : Int, n_repeats? : Int, seed? : Int) -> GainStatsSource

    v0.25.0+: from_blp_cv_repeated(blp, n_folds?,n_repeats?, seed?) is a more stable version of from_blp_cv. A single K-fold split has non-trivial variance in var_y_residuals (different splits put different rows in the test fold, so the OOF residual sum changes). Repeating the K-fold split n_repeats times with different seeds and averaging the residual variance reduces this variance by ~n_repeatsx (i.i.d. assumption on the per-rep estimate).

    Algorithm:
    1. For each rep in 0..n_repeats, run the K-fold OOF pipeline from from_blp_cv with seed_eff = seed + rep (any reproducible per-rep seed works; this scheme keeps seed=3141 giving the same first rep as from_blp_cv).
    2. var_y_residuals_scalar = mean over reps of ss_resid_oof_rep / n_obs (the per-rep average fold residual variance).

    coef, se, var_y, all_coef, n_rep, and the nu2 formula are all identical to from_blp_cv — only the residual variance estimate is averaged across repeats.

    On small samples (e.g. n_obs=200) with n_folds=5 and n_repeats=10, this typically gives a 5-10x reduction in the standard error of var_y_residuals, which translates to a similar stabilization of the downstream R2_y and nu2 sensitivity benchmarks.

    GainStatsSource::from_blp_hc0

    source
    fn GainStatsSource::from_blp_hc0(blp : DoubleMLBLP, n_folds? : Int, seed? : Int) -> GainStatsSource

    v0.23.0+: from_blp_hc0(blp, n_folds?, seed?) is the HC0-honest variant of from_blp_cv. The difference is the nu2 formula: instead of the homoskedastic OLS convention nu2 = var_y_residuals / (n_obs * se^2), it uses the projection-weight formula nu2[k] = (1 / n_obs) * ||M[k,:] @ basis^T||^2 where M = (basis^T basis + ridge I)^{-1} is the BLP's regression matrix. This is consistent with the upstream doubleml.utils._estimation._compute_sensitivity_elements convention, where nu2 = E[score_d^2] and the score is the per-observation influence on the k-th coefficient (M[k,:] @ x_i).

    The homoskedastic formula conflates nu2 with se^2 (using the relation se^2 = sigma^2 * (Z^T Z)^{-1}_{kk}), which is only correct under homoskedasticity. The HC0 SE se^2 = sum_i (M[k,:] @ x_i)^2 * e_i^2 does not satisfy the same relation; the projection-weight formula is the HC0-compatible alternative.

    var_y_residuals is computed via K-fold cross-fitting (same as from_blp_cv). coef, se, var_y, and all_coef are unchanged from from_blp.

    GainStatsSource::new

    source
    fn GainStatsSource::new(var_y_residuals : Array[Double], nu2 : Array[Double], all_coef : Array[Double], n_rep : Int, var_y : Double) -> GainStatsSource

    Builder constructor for GainStatsSource that mirrors the upstream DoubleML attribute access. Validates shape consistency: all three per-rep arrays must have the same length, and that length must be divisible by n_rep.

    GainStatsSource::to_repr

    source

    IivmData

    source
    type IivmData

    Result of the IIVM DGP.

    IivmData::d_get

    source
    fn IivmData::d_get(self : IivmData) -> Array[Double]

    Get the treatment vector.

    IivmData::theta_get

    source
    fn IivmData::theta_get(self : IivmData) -> Double

    Get the causal parameter (theta).

    IivmData::x_get

    source
    fn IivmData::x_get(self : IivmData) -> Matrix

    Get the design matrix.

    IivmData::y_get

    source
    fn IivmData::y_get(self : IivmData) -> Array[Double]

    Get the outcome vector.

    IivmData::z_get

    source
    fn IivmData::z_get(self : IivmData) -> Array[Double]

    Get the instrument vector.

    IrmConfoundedData

    source
    type IrmConfoundedData

    IrmConfoundedData::c_get

    source
    fn IrmConfoundedData::c_get(self : IrmConfoundedData) -> Array[Double]

    IrmConfoundedData::d_get

    source
    fn IrmConfoundedData::d_get(self : IrmConfoundedData) -> Array[Double]

    IrmConfoundedData::theta_get

    source
    fn IrmConfoundedData::theta_get(self : IrmConfoundedData) -> Double

    IrmConfoundedData::x_get

    source

    IrmConfoundedData::y_get

    source
    fn IrmConfoundedData::y_get(self : IrmConfoundedData) -> Array[Double]

    IrmData

    source
    type IrmData

    Result of the IRM DGP.

    IrmData::d_get

    source
    fn IrmData::d_get(self : IrmData) -> Array[Double]

    Get the binary treatment vector.

    IrmData::theta_get

    source
    fn IrmData::theta_get(self : IrmData) -> Double

    Get the causal parameter (theta).

    IrmData::x_get

    source
    fn IrmData::x_get(self : IrmData) -> Matrix

    Get the design matrix.

    IrmData::y_get

    source
    fn IrmData::y_get(self : IrmData) -> Array[Double]

    Get the outcome vector.

    IrmDiscreteData

    source
    type IrmDiscreteData

    IrmDiscreteData::d_get

    source
    fn IrmDiscreteData::d_get(self : IrmDiscreteData) -> Array[Double]

    IrmDiscreteData::theta_get

    source
    fn IrmDiscreteData::theta_get(self : IrmDiscreteData) -> Double

    IrmDiscreteData::x_get

    source

    IrmDiscreteData::y_get

    source
    fn IrmDiscreteData::y_get(self : IrmDiscreteData) -> Array[Double]

    IrmHeterogeneousData

    source
    type IrmHeterogeneousData

    IrmHeterogeneousData::d_get

    source
    fn IrmHeterogeneousData::d_get(self : IrmHeterogeneousData) -> Array[Double]

    IrmHeterogeneousData::theta_get

    source
    fn IrmHeterogeneousData::theta_get(self : IrmHeterogeneousData) -> Double

    IrmHeterogeneousData::x_get

    source

    IrmHeterogeneousData::y_get

    source
    fn IrmHeterogeneousData::y_get(self : IrmHeterogeneousData) -> Array[Double]

    LearnerDispatch

    source
    pub enum LearnerDispatch {
      LinearRegression(LinearRegression)
      Constant(ConstantLearner)
      Noop(NoopLearner)
      RandomForest(RFLearner)
      GradientBoosting(GBLearner)
      LogisticRegression(LogisticRegression)
    } derive(
    Debug    )

    Dispatch enum for the Learner trait. Stores one of the concrete learner types and lets cross_fit_predict_dispatch route to the right cross_fit_predict[T : Learner] arm.

    The enum keeps DoubleMLPLR non-generic in its struct definition (callers don't have to write DoubleMLPLR[LinearRegression, LinearRegression]), at the cost of an extra match arm per learner type. Adding a new learner is 3 small edits (see the file header).

    LearnerDispatch::constant

    source
    fn LearnerDispatch::constant(value : Double) -> LearnerDispatch

    Wrap a ConstantLearner in a LearnerDispatch.

    LearnerDispatch::gradient_boosting

    source
    fn LearnerDispatch::gradient_boosting(gb : GBLearner) -> LearnerDispatch

    Wrap a GBLearner in a LearnerDispatch. v0.57.0+. Factory for the Path A gradient-boosting learner; callers can override n_trees / learning_rate / max_depth / subsample via GBLearner::new(...) before wrapping.

    LearnerDispatch::linear_regression

    source
    fn LearnerDispatch::linear_regression() -> LearnerDispatch

    Default LearnerDispatch (a fresh LinearRegression). Used as the default value for the learner_l / learner_m labeled params on DoubleMLPLR::new.

    LearnerDispatch::logistic_regression

    source
    fn LearnerDispatch::logistic_regression(lr : LogisticRegression) -> LearnerDispatch

    Wrap a LogisticRegression in a LearnerDispatch. v0.62.0+: LPLR plumbing — binary classification for the double cross-fit path (ml_M, ml_a, ml_m).

    LearnerDispatch::noop

    source

    Wrap a NoopLearner in a LearnerDispatch.

    LearnerDispatch::random_forest

    source
    fn LearnerDispatch::random_forest(rf : RFLearner) -> LearnerDispatch

    Wrap an RFLearner in a LearnerDispatch. v0.56.0+. Factory for the Path A random-forest learner; callers can override n_trees / max_depth / min_samples_leaf / mtry via RFLearner::new(...) before wrapping.

    LearnerDispatch::to_repr

    source

    LinearRegression

    source
    pub struct LinearRegression {
      ridge : Double
      coef_ : Array[Double]
      xtx_inv_diag : Array[Double]
      xtwx_inv_diag : Array[Double]
      fitted : Bool
    } derive(
    Debug    )

    Closed-form OLS linear regression.

    Fits a linear model y = X * beta + intercept (intercept is folded into X via a leading column of ones) by solving the normal equations X^T X beta = X^T y using a Cholesky factorisation. A small ridge lambda is added to the diagonal of X^T X to guard against singularity when the design matrix is near-collinear; the default lambda = 1e-10 is small enough to be invisible on well-conditioned data but prevents the Cholesky from failing on degenerate inputs.

    LinearRegression::coefficients

    source
    fn LinearRegression::coefficients(self : LinearRegression) -> Array[Double]

    LinearRegression::covariance_diagonal

    source
    fn LinearRegression::covariance_diagonal(self : LinearRegression, sigma2 : Double) -> Array[Double]

    Diagonal of the homoskedastic coefficient covariance matrix sigma^2 * (X^T X + ridge * I)^{-1}. The caller supplies sigma^2 = RSS / (n - p); this separation keeps LinearRegression ignorant of the response vector (and hence the residual sum of squares). Returns a length-(p + 1) array whose jth entry is the variance of coef_[j]. The corresponding SE is sqrt(cov[j]). Used by DoubleMLBLP to recover the per-coefficient standard error (Bug #5: previously every coefficient shared the same SE = sqrt(RSS / (n - p)), ignoring the (X^T X)^{-1} scaling).

    REVIEW L8: after a fit_weighted() call the cached xtx_inv_diag is the empty array (M10 fix), so calling covariance_diagonal on a WLS-fit model yields a vector of zeros. Callers must use the unweighted fit path (or pass the pre-computed (X^T X)^{-1} diagonal themselves) when the homoskedastic SE is wanted.

    LinearRegression::fit

    source
    fn LinearRegression::fit(self : LinearRegression, x : Matrix, y : Array[Double]) -> LinearRegression

    Fit the linear regression to (x, y). x is n x p and y is a length-n vector. Returns the fitted LinearRegression with coef_[0] = intercept and coef_[1..] = slopes.

    At fit time we also cache the diagonal of (X^T X + ridge I)^{-1} so that covariance_diagonal() (used by DoubleMLBLP for the per-coefficient standard error) is O(p) instead of re-inverting X^T X from scratch.

    LinearRegression::fit_weighted

    source
    fn LinearRegression::fit_weighted(self : LinearRegression, x : Matrix, y : Array[Double], w : Array[Double]) -> LinearRegression

    Weighted OLS fit. Solves X^T W X beta = X^T W y for a positive diagonal weight matrix W = diag(w). Used by DoubleMLRDD for local-linear regression with the triangular kernel weights (Bug #6: previously the weights were only applied to the variance sum, so the point estimate ignored them). The covariance diagonal is computed against the unweighted X^T X so the SE is the standard homoskedastic-OLS form for the same design matrix — this matches the upstream RDD reference for cov_type='nonrobust'.

    The full weighted-Normal inverse (X^T W X + ridge I)^{-1} is also cached for the WLS-aware SE (TODO #11c.2): the RDD delta-method variance scales the weighted residual variance by the (X^T W X)^{-1}[0, 0] (intercept) entry, which is the variance formula for the WLS point estimate at the cutoff.

    LinearRegression::n_features

    source
    fn LinearRegression::n_features(self : LinearRegression) -> Int

    Number of features (excluding the intercept).

    LinearRegression::new

    source
    fn LinearRegression::new(ridge? : Double) -> LinearRegression

    LinearRegression::predict

    source
    fn LinearRegression::predict(self : LinearRegression, x : Matrix) -> Array[Double]

    Predict the response for new data x. x must have the same number of columns as the data used at fit time.

    LinearRegression::sandwich_se

    source
    fn LinearRegression::sandwich_se(self : LinearRegression, x : Matrix, y : Array[Double]) -> Array[Double]

    Heteroskedasticity-consistent (HC0) covariance diagonal of the coefficient estimator. Computes the diagonal of the sandwich form cov(beta_hat) = (X'X)^{-1} (X' diag(e^2) X) (X'X)^{-1} where e = y - X beta_hat is the in-sample residual vector. The diagonal entry j simplifies to cov_jj = sum_i ((M[j,:] · x_i)^2 * e_i^2) where M = (X'X + ridge I)^{-1} is the cached (X'X)^{-1}-with-ridge matrix. This is O(n * p^2) and avoids materialising the full sandwich matrix. Used by DoubleMLBLP for the per-coefficient heteroskedasticity-robust SE (TODO #11c.1: matches the upstream statsmodels.OLS(cov_type='HC0') default).

    Implementation note: instead of materialising the full (X'X)^{-1} matrix via inv_spd (O(p³) memory + O(p³) compute), we solve p1 back-systems M[j, :] * (X'X + ridge I)= e_j via solve_spd (each O(p²) compute, zero extra memory). The diagonal entry of the sandwich is unchanged; only the form of the inner loop changed.

    LinearRegression::sandwich_se_weighted

    source
    fn LinearRegression::sandwich_se_weighted(self : LinearRegression, x : Matrix, y : Array[Double], w : Array[Double]) -> Array[Double]

    Heteroskedasticity-consistent (HC0) sandwich covariance diagonal of the WLS coefficient estimator. Mirrors sandwich_se but for the weighted case:

     cov(beta_hat) = (X'WX)^{-1} (X' diag(w · e^2) X) (X'WX)^{-1}

    diagonal entry j simplifies to cov_jj = sum_i w[i]^2 · ((M[j,:] · x_i)^2 · e_i^2) where M = (X^T W X + ridge I)^{-1} and e = y - X beta_hat. The extra w[i]^2 factor on the score reflects the WLS weight's role in the IRLS normal equations (see White 1980, §4).

    Used by DoubleMLRDD for the heteroskedasticity-robust WLS intercept variance when cov_type = "HC0" (TODO 0.6.0).

    LinearRegression::to_repr

    source

    LinearRegression::xtwx_inv_diag

    source
    fn LinearRegression::xtwx_inv_diag(self : LinearRegression) -> Array[Double]

    Return the diagonal of (X^T W X + ridge I)^{-1} cached from the last fit_weighted() call. Empty array if the last fit was the unweighted fit(). Used by DoubleMLRDD for the WLS-aware intercept variance (TODO #11c.2).

    LogisticRegression

    source
    pub struct LogisticRegression {
      fitted : Bool
      coef_ : Array[Double]
    } derive(
    Debug    )

    Binary logistic regression fit by Newton-Raphson IRLS.

    Fits a binary logit model P(y = 1 | x) = sigmoid(x @ beta) by iterated reweighted least squares. The intercept is folded into the design matrix as a leading column of ones via augment_with_intercept; the returned coef_[0] is the intercept and coef_[1..] are the slopes. Only binary outcomes (y in {0, 1}) are supported — multiclass classification is intentionally out of scope for this learner.

    LogisticRegression::coefficients

    source
    fn LogisticRegression::coefficients(self : LogisticRegression) -> Array[Double]

    Fitted coefficient vector [intercept, slope_1, ..., slope_p]. Aborts when called before fit() via require(self.fitted).

    LogisticRegression::fit

    source
    fn LogisticRegression::fit(self : LogisticRegression, x : Matrix, y : Array[Double], max_iter? : Int, tol? : Double, ridge? : Double) -> LogisticRegression

    Fit the logistic regression by Newton-Raphson IRLS. x is n x p and y is a length-n vector of 0/1 labels.

    Each iteration computes:

    eta = X_aug @ beta p = sigmoid(eta) w = p * (1 - p) # per-observation IRLS weight z = eta + (y - p) / w # working response beta_new = (X_aug^T W X_aug + ridge * I)^-1 X_aug^T W z

    and the loop exits early when the L2 step ||beta_new - beta|| drops below tol. ridge guards the X^T W X factor against singularity when p is exactly 0 or 1 (which makes w zero on some rows); the default of 1e-8 is invisible on well-conditioned data but prevents the Cholesky from failing on separable problems at the boundary.

    The default max_iter = 25 comfortably covers both well-separated data (3-5 steps) and noisy binary outcomes (10-20 steps).

    LogisticRegression::new

    source

    Build an unfitted LogisticRegression learner. Call fit(x, y) to estimate the coefficients; predict and coefficients abort until then.

    LogisticRegression::predict

    source
    fn LogisticRegression::predict(self : LogisticRegression, x : Matrix) -> Array[Double]

    Predict the probability P(y = 1 | x) for each row of x. The output is clipped to (1e-15, 1 - 1e-15) so the returned probabilities are always strictly inside the open interval (0, 1), matching the strict-inequality guarantee the predict_class threshold needs. Aborts when called before fit() or when x has a column count that does not match the data used at fit time.

    LogisticRegression::predict_class

    source
    fn LogisticRegression::predict_class(self : LogisticRegression, x : Matrix, threshold? : Double) -> Array[Double]

    Predict the binary class label by thresholding predict(x) at threshold (default 0.5). Returns 0.0 / 1.0 doubles (not ints) so the output is type-compatible with predict for downstream scoring code. Inherits the predict abort behaviour.

    LogisticRegression::to_repr

    source

    LplrData

    source
    type LplrData

    Result of the LPLR LZZ2020 DGP.

    LplrData::data_get

    source
    fn LplrData::data_get(self : LplrData) -> DoubleMLBinaryData

    Get the underlying DoubleMLBinaryData.

    LplrData::theta_get

    source
    fn LplrData::theta_get(self : LplrData) -> Double

    Get the true causal parameter.

    Matrix

    source

    Dense 2D matrix stored in row-major order as a flat Array[Double].

    This is a minimal matrix type used by the DoubleML port. It is not intended to be a general-purpose linear algebra library; only the operations needed for the closed-form LinearRegression learner and the DML score have been implemented.

    Matrix::cols

    source
    fn Matrix::cols(self : Matrix) -> Int

    Number of columns (accessor — ncols is a private field).

    Matrix::copy

    source
    fn Matrix::copy(self : Matrix) -> Matrix

    Copy this matrix into a new Matrix.

    Matrix::from_array

    source
    fn Matrix::from_array(data : Array[Double], nrows : Int, ncols : Int) -> Matrix

    Create a matrix from a flat row-major array. Panics if length does not match nrows * ncols.

    Matrix::from_rows

    source
    fn Matrix::from_rows(rows : Array[Array[Double]]) -> Matrix

    Build a matrix from a list of row vectors. Returns an nrows x ncols matrix where ncols = rows[0].length(). Panics if the rows do not all share the same length.

    Matrix::get

    source
    fn Matrix::get(self : Matrix, i : Int, j : Int) -> Double

    Element access (get). Panics if indices are out of range.

    Matrix::identity

    source
    fn Matrix::identity(n : Int) -> Matrix

    Create an n x n identity matrix.

    Matrix::ones

    source
    fn Matrix::ones(nrows : Int, ncols : Int) -> Matrix

    Create a matrix filled with ones.

    Matrix::rows

    source
    fn Matrix::rows(self : Matrix) -> Int

    Number of rows (accessor — nrows is a private field).

    Matrix::set

    source
    fn Matrix::set(self : Matrix, i : Int, j : Int, v : Double) -> Unit

    Element access (set). Panics if indices are out of range.

    Matrix::to_repr

    source

    Matrix::transpose

    source
    fn Matrix::transpose(self : Matrix) -> Matrix

    Matrix transpose.

    Matrix::zeros

    source
    fn Matrix::zeros(nrows : Int, ncols : Int) -> Matrix

    Create a zero matrix of the given shape.

    NoopLearner

    source
    pub struct NoopLearner {
      value : Double
    } derive(
    Debug    )

    Trivial learner: predicts zero for every row. fit is a no-op; predict returns all zeros. Useful for sanity tests that the DML pipeline doesn't divide by zero when the nuisance is trivial.

    NoopLearner::fit

    source
    fn NoopLearner::fit(self : NoopLearner, x : Matrix, y : Array[Double]) -> NoopLearner

    NoopLearner::new

    source

    NoopLearner::predict

    source
    fn NoopLearner::predict(self : NoopLearner, x : Matrix) -> Array[Double]

    NoopLearner::to_repr

    source

    PSProcessor

    source
    pub struct PSProcessor {
      config : PSProcessorConfig
    } derive(
    Debug    )

    Propensity-score processor. Stateless apart from its configuration; safe to share across DoubleMLDIDBinary / DoubleMLDIDCS instances. adjust_ps returns a new array of the same length as the input, never mutating the caller's data.

    v0.14.0+: the isotonic calibration method is now fully implemented (was a placeholder in v0.10.0). The pure-MoonBit PAVA implementation in pava produces a step-function isotonic regression; with cv_calibration=true, K-fold cross-validated predictions are used instead (one isotonic fit per fold, predictions concatenated across folds).

    PSProcessor::adjust_ps

    source
    fn PSProcessor::adjust_ps(self : PSProcessor, ps : Array[Double], treatment : Array[Double], cv? : Array[(Array[Int], Array[Int])]?) -> Array[Double]

    Apply the configured calibration followed by the [clipping_threshold, 1 - clipping_threshold] clip. Returns a new array; the caller's ps and treatment are not mutated.

    cv is consulted only when config.calibration_method ="isotonic" and config.cv_calibration = true. It is a list of folds, each a pair (train_indices, test_indices) — when cv = None, a deterministic 5-fold split with seed=3141 is used. With cv_calibration = false, cv is ignored and the isotonic fit uses the full (ps, treatment) (matches upstream IsotonicRegression default: no CV).

    PSProcessor::calibration_method

    source
    fn PSProcessor::calibration_method(self : PSProcessor) -> String

    PSProcessor::clipping_threshold

    source
    fn PSProcessor::clipping_threshold(self : PSProcessor) -> Double

    PSProcessor::cv_calibration

    source
    fn PSProcessor::cv_calibration(self : PSProcessor) -> Bool

    PSProcessor::extreme_threshold

    source
    fn PSProcessor::extreme_threshold(self : PSProcessor) -> Double

    PSProcessor::from_config

    source
    fn PSProcessor::from_config(config : PSProcessorConfig) -> PSProcessor

    Convenience constructor mirroring upstream PSProcessor.from_config.

    PSProcessor::new

    source

    PSProcessor::to_repr

    source

    PSProcessorConfig

    source
    pub(all) struct PSProcessorConfig {
      clipping_threshold : Double
      extreme_threshold : Double
      calibration_method : String
      cv_calibration : Bool
    } derive(
    Debug    )

    Configuration for propensity-score processing. The processor applies (optionally) a calibration step and always a [clipping_threshold, 1 - clipping_threshold] clip to keep the propensity scores away from 0/1 in the score denominator.

    Calibration (v0.14.0+): the only currently supported method is isotonic regression (PAVA, pool-adjacent-violators algorithm). Set calibration_method="isotonic" to fit an isotonic regression of treatment on ps and use the fitted curve as the calibrated propensity. Combine with cv_calibration=true to use K-fold cross-validated calibration (matches upstream sklearn.model_selection.cross_val_predict); pass cv to PSProcessor::adjust_ps to control the fold partition.

    PSProcessorConfig::default

    source

    Default PS processor config (clipping_threshold=1e-2, no calibration). v0.44.0: returns a pre-constructed default via a private helper that doesn't raise, so this signature stays PSProcessorConfig (no raise).

    PSProcessorConfig::new

    source
    fn PSProcessorConfig::new(clipping_threshold? : Double, extreme_threshold? : Double, calibration_method? : String, cv_calibration? : Bool) -> PSProcessorConfig raise PSConfigError

    Returns PSProcessorConfig raise PSConfigError: the v0.44.0 conversion replaces the previous abort() call with raise PSConfigError::InconsistentCVCalibration so the inconsistent-configuration path becomes directly testable. Callers that want the pre-v0.44.0 process-death behavior should catch the error and re-abort (this is what PSProcessorConfig::default does, although it never triggers the error in practice). The require checks on the four argument ranges are unchanged and still abort the process (they are central require checks; refactoring them is out of scope for this surgical release).

    PSProcessorConfig::to_repr

    source

    PanelTransform

    source
    type PanelTransform derive(
    Debug    )

    Transformed cross-section produced by one of the four static panel approaches. d_mean_row carries the per-row unit mean of D when an approach needs it downstream (cre_general, cre_normal), otherwise it is empty. id is the unit id of every transformed row: upstream wraps the transformed data in a static-panel DoubleMLPanelData, which sets cluster_cols =id_col, so the clustered DML path always keys off this column.

    PanelTransform::to_repr

    source

    PlivClusterData

    source
    type PlivClusterData

    PlivClusterData::cluster_a_get

    source
    fn PlivClusterData::cluster_a_get(self : PlivClusterData) -> Array[Int]

    PlivClusterData::cluster_b_get

    source
    fn PlivClusterData::cluster_b_get(self : PlivClusterData) -> Array[Int]

    PlivClusterData::d_get

    source
    fn PlivClusterData::d_get(self : PlivClusterData) -> Array[Double]

    PlivClusterData::theta_get

    source
    fn PlivClusterData::theta_get(self : PlivClusterData) -> Double

    PlivClusterData::x_get

    source

    PlivClusterData::y_get

    source
    fn PlivClusterData::y_get(self : PlivClusterData) -> Array[Double]

    PlivClusterData::z_get

    source
    fn PlivClusterData::z_get(self : PlivClusterData) -> Array[Double]

    PlivData

    source
    type PlivData

    PlivData::d_get

    source
    fn PlivData::d_get(self : PlivData) -> Array[Double]

    PlivData::theta_get

    source
    fn PlivData::theta_get(self : PlivData) -> Double

    PlivData::x_get

    source
    fn PlivData::x_get(self : PlivData) -> Matrix

    PlivData::y_get

    source
    fn PlivData::y_get(self : PlivData) -> Array[Double]

    PlivData::z_get

    source
    fn PlivData::z_get(self : PlivData) -> Array[Double]

    PlprData

    source
    type PlprData

    Result of the PLPR DGP.

    PlprData::d_get

    source
    fn PlprData::d_get(self : PlprData) -> Array[Double]

    PlprData::n_periods_get

    source
    fn PlprData::n_periods_get(self : PlprData) -> Int

    PlprData::n_units_get

    source
    fn PlprData::n_units_get(self : PlprData) -> Int

    PlprData::theta_get

    source
    fn PlprData::theta_get(self : PlprData) -> Double

    PlprData::x_get

    source
    fn PlprData::x_get(self : PlprData) -> Matrix

    PlprData::y_get

    source
    fn PlprData::y_get(self : PlprData) -> Array[Double]

    PlrCcddhnr2018

    source
    type PlrCcddhnr2018

    Result of the PLR CCDDHNR 2018 DGP.

    PlrCcddhnr2018::d_get

    source
    fn PlrCcddhnr2018::d_get(self : PlrCcddhnr2018) -> Array[Double]

    Get the treatment vector.

    PlrCcddhnr2018::theta_get

    source
    fn PlrCcddhnr2018::theta_get(self : PlrCcddhnr2018) -> Double

    Get the causal parameter (theta).

    PlrCcddhnr2018::x_get

    source
    fn PlrCcddhnr2018::x_get(self : PlrCcddhnr2018) -> Matrix

    Get the design matrix.

    PlrCcddhnr2018::y_get

    source
    fn PlrCcddhnr2018::y_get(self : PlrCcddhnr2018) -> Array[Double]

    Get the outcome vector.

    PlrConfoundedData

    source
    type PlrConfoundedData

    PlrConfoundedData::c_get

    source
    fn PlrConfoundedData::c_get(self : PlrConfoundedData) -> Array[Double]

    PlrConfoundedData::d_get

    source
    fn PlrConfoundedData::d_get(self : PlrConfoundedData) -> Array[Double]

    PlrConfoundedData::theta_get

    source
    fn PlrConfoundedData::theta_get(self : PlrConfoundedData) -> Double

    PlrConfoundedData::x_get

    source

    PlrConfoundedData::y_get

    source
    fn PlrConfoundedData::y_get(self : PlrConfoundedData) -> Array[Double]

    PolicyTreeNode

    source
    pub enum PolicyTreeNode {
      Leaf(Int)
      Split(Int, Double, PolicyTreeNode, PolicyTreeNode)
    } derive(
    Debug    )

    A binary tree node used by DoubleMLPolicyTree. A Leaf is a terminal that always returns the given treatment. A Split carries the split feature, the threshold value, and the two child nodes. Internal-only (not exposed in the public API) so the public DoubleMLPolicyTree signature is unchanged from the depth-1 era.

    PolicyTreeNode::to_repr

    source

    RFLearner

    source
    pub struct RFLearner {
      n_trees : Int
      max_depth : Int
      min_samples_leaf : Int
      mtry : Int
      bootstrap_seed : Int
      trees : Array[CART]
      n_features : Int
    } derive(
    Debug    )

    Random forest learner (regression). Fits a forest of CART trees on bootstrap samples; predicts by averaging leaf values across all trees.

    RFLearner::fit

    source
    fn RFLearner::fit(self : RFLearner, x : Matrix, y : Array[Double]) -> RFLearner

    Promote the Learner trait methods (fit, predict) as explicit methods on RFLearner. Without this extend, MoonBit raises an unused_value warning for the trait impls below (the trait methods are reached only via dispatch through cross_fit_predict_dispatch, never directly), and the --deny-warn build gate treats that as an error. Same pattern as the v0.54.0 ConstantLearner / NoopLearner setup.

    RFLearner::n_features

    source
    fn RFLearner::n_features(self : RFLearner) -> Int

    Number of feature columns the forest was fit on. Returns -1 if the learner has not been fit yet.

    RFLearner::n_trees

    source
    fn RFLearner::n_trees(self : RFLearner) -> Int

    Number of trees in the fitted forest. Returns 0 if not yet fit (so callers can early-return on RFLearner::new() without an explicit "is_fitted" flag).

    RFLearner::new

    source
    fn RFLearner::new(n_trees? : Int, max_depth? : Int, min_samples_leaf? : Int, mtry? : Int, bootstrap_seed? : Int) -> RFLearner

    RFLearner::predict

    source
    fn RFLearner::predict(self : RFLearner, x : Matrix) -> Array[Double]

    Promote the Learner trait methods (fit, predict) as explicit methods on RFLearner. Without this extend, MoonBit raises an unused_value warning for the trait impls below (the trait methods are reached only via dispatch through cross_fit_predict_dispatch, never directly), and the --deny-warn build gate treats that as an error. Same pattern as the v0.54.0 ConstantLearner / NoopLearner setup.

    RFLearner::to_repr

    source

    RddSimpleData

    source
    type RddSimpleData

    RddSimpleData::d_get

    source
    fn RddSimpleData::d_get(self : RddSimpleData) -> Array[Double]

    RddSimpleData::score_get

    source
    fn RddSimpleData::score_get(self : RddSimpleData) -> Array[Double]

    RddSimpleData::theta_get

    source
    fn RddSimpleData::theta_get(self : RddSimpleData) -> Double

    RddSimpleData::x_get

    source
    fn RddSimpleData::x_get(self : RddSimpleData) -> Array[Double]

    RddSimpleData::y_get

    source
    fn RddSimpleData::y_get(self : RddSimpleData) -> Array[Double]

    SsmData

    source
    type SsmData

    Result of the SSM DGP.

    SsmData::d_get

    source
    fn SsmData::d_get(self : SsmData) -> Array[Double]

    SsmData::s_get

    source
    fn SsmData::s_get(self : SsmData) -> Array[Double]

    SsmData::theta_get

    source
    fn SsmData::theta_get(self : SsmData) -> Double

    SsmData::x_get

    source
    fn SsmData::x_get(self : SsmData) -> Matrix

    SsmData::y_get

    source
    fn SsmData::y_get(self : SsmData) -> Array[Double]

    TuneParam

    source
    pub struct TuneParam {
      learner_l : LearnerDispatch
      learner_m : LearnerDispatch
    } derive(
    Debug    )

    TuneParam — one row of the candidate grid passed to DoubleMLPLR::tune. Holds a concrete nuisance-learner combination (one learner_l, one learner_m) to be scored against the others in the grid.

    The grid is Array[TuneParam] (typed, not Dict). Reasons for the typed wrapper vs a Dict[String, LearnerDispatch]:

    • compile-time key/value type safety (Dict keys are runtime-checked strings)
    • no Dict-construction FFI overhead per candidate
    • cleaner documentation (TuneParam::{learner_l, learner_m} is a documented struct, not a magic-key dict)

    Python upstream doubleml.DoubleMLPLR.tune uses a similar shape internally (a list of {"learner_l": ..., "learner_m":...} type) but exposes it via kwargs to set_tune_params.

    TuneParam::new

    source
    fn TuneParam::new(learner_l : LearnerDispatch, learner_m : LearnerDispatch) -> TuneParam

    Build a TuneParam from an explicit (learner_l, learner_m) pair. Equivalent to the struct literal { learner_l, learner_m } — exposed as a factory so callers don't need to know the field order of an internal struct.

    TuneParam::to_repr

    source

    TuneResult

    source
    pub struct TuneResult {
      best : TuneParam
      best_score : Double
      all_scores : Array[Double]
    } derive(
    Debug    )

    Result of a DoubleMLPLR::tune call. Holds the best (by the chosen scoring rule, lower-is-better or higher-is-better per TuneScoring) candidate plus the full per-candidate score vector for diagnostics.

    all_scores[i] is the raw score for param_set[i] (in the grid order, NOT necessarily sorted). Lower-is-better / higher- is-better is the same convention as the chosen TuneScoring. For diagnostics where the user wants to compare across scoring rules, they should re-run tune with the other rule.

    TuneResult::to_repr

    source

    TuneScoring

    source
    pub enum TuneScoring {
      MSE
      RMSE
      NegMSE
    } derive(
    Debug    )

    Scoring method for DoubleMLPLR::tune. All three are based on the OOF outcome-nuisance prediction l_hat = cross_fit(y | x) (TUNE_DESIGN.md §4.1: MSE-on-l_hat, the upstream Python default):

    • MSE: mean((y - l_hat)^2). Lower is better.
    • RMSE: sqrt(MSE). Lower is better.
    • NegMSE: -MSE. Higher is better (sklearn convention).

    R^2 scoring is deferred to a follow-up release — R^2 needs a defended against degenerate Var(y) = 0 and isn't needed by the standard PLR tune workflow.

    TuneScoring::parse

    source
    fn TuneScoring::parse(s : String) -> TuneScoring

    Parse a scoring-method string into a TuneScoring enum. Accepted spellings (case-insensitive):

    • "MSE" / "mse" -> MSE
    • "RMSE" / "rmse" -> RMSE
    • "neg-MSE" / "neg-mse" / "NegMSE" -> NegMSE

    Unknown spellings abort with a descriptive message. Use this when accepting scoring-method names from CLI / JSON configs; pass the enum directly for typed code.

    TuneScoring::to_repr

    source

    WideDIDSubset

    source
    type WideDIDSubset derive(
    Debug    )

    The wide-format DID subset returned by preprocess_did_binary. Carries the covariate matrix, the first-differenced outcome, the binary G_indicator (d), and the binary t_indicator (t_indicator: 0 if the wide row originated from the t_value_pre period, 1 if from t_value_eval).

    WideDIDSubset::to_repr

    source

    add_ridge

    source
    fn add_ridge(a : Matrix, lambda : Double) -> Matrix

    Add a small ridge lambda * I to a square matrix and return a copy. Useful for guarding X^T X against singularity when the design matrix is near-collinear.

    aggregate_coef_se

    source
    fn aggregate_coef_se(coefs : Array[Double], ses : Array[Double]) -> (Double, Double)

    Per-repetition coefficient / standard-error aggregation for the DML family of estimators. Mirrors doubleml.utils._aggregate_coefs_and_ses in the upstream doubleml package:

     theta_hat = median(theta_1, ..., theta_R)
 ub_r      = theta_r + 1.96 * se_r           for each rep r
 ub_hat    = median(ub_1, ..., ub_R)
 se_hat    = (ub_hat - theta_hat) / 1.96

    The median is the "high median" — sort the array, take the entry at index length() / 2 (so for n = 1 it is the only entry, for n = 2 it is the upper-middle, for odd n it is the exact middle). This intentionally uses Array::sort + length() / 2 and does not depend on any external package.

    The two input arrays must have the same non-zero length; otherwise the call aborts at the require(...) precondition. The arrays are not mutated — the function makes its own copy of coefs before sorting.

    For n_rep == 1 the result is exactly (coefs[0], ses[0]) because:

     theta_hat = median([c]) = c
 ub_hat    = median([c + 1.96 * s]) = c + 1.96 * s
 se_hat    = (c + 1.96 * s - c) / 1.96 = s

    This is the regression-protection invariant that the four estimator fit()s rely on: when n_rep = 1 the new "per-rep + median" implementation must produce the same (theta, se) as the previous "average nuisances, then estimate" implementation.

    aggregate_event

    source
    fn aggregate_event(coef_matrix : Array[Double], se_matrix : Array[Double], groups : Array[Int], periods : Array[Int], group_sizes : Array[Int]) -> DIDAggregationResult

    Event-study aggregation. For each unique event time e = t_eval - g, computes the weighted mean of all per-(g, t_eval) ATTs with t_eval - g = e.

    Pre-treatment cells (e < 0) are included as pre-trend estimates (matching upstream's default; upstream agg_weights only normalises post-treatment cells, but the per-event theta uses all selected cells, so the event-time profile includes pre-trend as a sanity check). Returns one entry per unique event time.

    aggregate_group

    source
    fn aggregate_group(coef_matrix : Array[Double], se_matrix : Array[Double], groups : Array[Int], periods : Array[Int], group_sizes : Array[Int]) -> DIDAggregationResult

    Group aggregation. For each unique group value g_i in groups, computes the weighted mean of coef_matrix[gi, :] (the per-g_i post-treatment ATTs) with weights proportional to n_units_in_group_g_i (the share of units in group g_i).

    Arguments:
    • coef_matrix : row-major (n_groups, n_periods) ATT matrix from a fitted DoubleMLDIDCS / DoubleMLDIDMulti. coef_matrix[gi * n_periods + pi] is the ATT for group groups[gi] at period periods[pi]. Pre-treatment cells (e.g. t < g) are expected to be 0.0 (the DoubleMLDIDCS convention).
    • se_matrix : row-major (n_groups, n_periods) SE matrix (same shape as coef_matrix).
    • groups : the sorted-ascending unique group values.
    • periods : the sorted-ascending unique time-period values (the same times accessor from DoubleMLDIDCSData).
    • group_sizes : per-unit group sizes (one entry per unique row in the long-format panel). Must be aligned with groups; entry i is the number of units in groups[i].

    Returns a DIDAggregationResult with one entry per group.

    aggregate_time

    source
    fn aggregate_time(coef_matrix : Array[Double], se_matrix : Array[Double], groups : Array[Int], periods : Array[Int], group_sizes : Array[Int]) -> DIDAggregationResult

    Time aggregation. For each unique time period t_j in periods, computes the weighted mean of all per-(g, t_j) ATTs (across groups) with weights proportional to group_sizes[g] (the share of units in each group).

    Pre-treatment cells (t < g) are skipped. Returns one entry per period.

    apply_calibration

    source
    fn apply_calibration(config : PSProcessorConfig, ps : Array[Double], treatment : Array[Double], cv : Array[(Array[Int], Array[Int])]?) -> Array[Double] raise

    Apply the configured calibration method to a propensity-score vector. Returns Array[Double] raise InvalidCalibrationError: the v0.37.0 extraction lifts the previous abort() call from PSProcessor::adjust_ps into this helper so the unknown-method path becomes directly testable. PSProcessor::adjust_ps wraps the call in try ... catch { ... => abort(...) } to preserve pre-v0.37.0 process-death behavior.

    Note: PSProcessorConfig::new also has a require check on calibration_method (only "none" and "isotonic" are allowed). That check normally fires first; this helper's raise only fires if a config is constructed directly (bypassing new), which the new ps_processor_adjust_ps_raises_unknown_method test does deliberately.

    array_max

    source
    fn array_max(v : Array[Double]) -> Double raise EmptyArrayError

    Returns the maximum element of v. Raises EmptyArrayError if v is empty. v0.41.0: signature changed from Double to Double raise EmptyArrayError to make the empty-array path testable. Callers that want the pre-v0.41.0 process-death behavior should catch and re-abort.

    array_min

    source
    fn array_min(v : Array[Double]) -> Double raise EmptyArrayError

    Returns the minimum element of v. Raises EmptyArrayError if v is empty. v0.41.0: signature changed from Double to Double raise EmptyArrayError to make the empty-array path testable. Callers that want the pre-v0.41.0 process-death behavior should catch and re-abort.

    augment_one_col

    source
    fn augment_one_col(x : Matrix, extra : Array[Double]) -> Matrix

    Augment a feature matrix with one extra column (used to add pi_hat to the design matrix for the conditional-outcome learners).

    augment_with_intercept

    source
    fn augment_with_intercept(x : Matrix) -> Matrix

    Augment a feature matrix X (n x p) with a leading column of ones to form the design matrix used by the OLS fit. The result has shape n x (p + 1). Exposed as pub so other learners (e.g. the LogisticRegression IRLS fit) can reuse the same convention without duplicating the loop.

    bh_fdr_p_adjust

    source
    fn bh_fdr_p_adjust(unadjusted : Array[Double]) -> Array[Double]

    Benjamini-Hochberg FDR correction.

    Algorithm (matches statsmodels.stats.multitest.multipletests with method='fdr_bh'):
    1. Sort unadjusted p-values ascending; let order be the resulting permutation of cell indices and ro its inverse.
    2. p_corrected_sorted[k] = min(1.0, p_sorted[k] * n / (k + 1)).
    3. Enforce monotonicity from the largest rank downward (this is the BH-specific direction; Holm goes the other way): p_corrected_sorted[k] = min(p_corrected_sorted[k],p_corrected_sorted[k + 1]).
    4. Re-order to original cell order via ro.

    p_corrected[k] >= p_unadjusted[k] is not guaranteed (BH controls FDR, not FWER); some adjusted p-values can be smaller than the unadjusted ones.

    bonferroni_p_adjust

    source
    fn bonferroni_p_adjust(unadjusted : Array[Double]) -> Array[Double]

    Bonferroni correction. p_corrected[k] = min(1.0, n *p_unadjusted[k]).

    box_muller_normal

    source
    fn box_muller_normal(rng : 
    Rand    ) -> Double

    Standard normal sample via Box-Muller. Pairs (u1, u2) in [0, 1) to (z1, z2) ~ N(0, 1). The chacha8 RNG is uniform in [0, 1) per the upstream numpy.random default; we use the cosine for the first draw and the sine for the second to consume two uniforms per pair.

    box_muller_pair

    source
    fn box_muller_pair(rng : 
    Rand    ) -> (Double, Double)

    Box-Muller pair (mu=0, sigma=1). Returns two independent standard normals per call. Shared helper used by the DGP modules (dgp_plr_CCDDHNR, dgp_irm, ...) to avoid duplicating the helper across files.

    build_row_unit_map

    source
    fn build_row_unit_map(cluster : Array[Int], uniq : Array[Int]) -> Array[Int] raise ClusterDataError

    Build a row → unit-position map (length = n, each entry is the position of the row's unit in the unique-ids array produced by unique_units(cluster)). Aborts if any row's unit is missing from the uniq list (a row whose unit is not in uniq would be an unrecoverable data error).

    Returns Array[Int] raise ClusterDataError: the v0.36.0 conversion replaces the previous abort() call with raise ClusterDataError::MissingUnit(g) so the data-error path becomes directly testable. Callers that want the pre-v0.36.0 process-death behavior should catch the error and re-abort (this is what every DoubleMLXXX::fit_cluster does); callers that want to surface the error to downstream consumers should propagate via ?. The error type is declared at the top of this file so the cluster helper stack can share it.

    by_fdr_p_adjust

    source
    fn by_fdr_p_adjust(unadjusted : Array[Double]) -> Array[Double]

    Benjamini-Yekutieli FDR correction.

    Algorithm (matches statsmodels.stats.multitest.multipletests with method='fdr_by'):
    1. Compute the harmonic-sum factor c = sum_{i=1}^{n} 1/i (a.k.a. H_n).
    2. Same as BH, but p_corrected_sorted[k] = min(1.0, p_sorted[k] * n * c / (k + 1)).
    3. Enforce monotonicity from the largest rank downward (same as BH).
    4. Re-order to original cell order via ro.

    The c factor accounts for the dependence structure under arbitrary dependence; the BY procedure is more conservative than BH but valid under weaker assumptions.

    chacha8_rng

    source
    fn chacha8_rng(seed : Int) -> 
    Rand

    Convenience constructor: chacha8_rng(seed) is shorthand for Rand::chacha8(seed=Bytes::from_array(seed_to_bytes(seed))). Use this anywhere the package's tests / demos need a deterministic RNG keyed by an integer seed.

    check

    source
    #callsite(autofill(loc))
    fn check(condition : Bool, loc~ : SourceLoc) -> Unit raise PreconditionError

    Runtime precondition used by every public API and numerical kernel. #callsite(autofill(loc)) instructs the compiler to auto-inject the call-site source location into loc at every call, so the error payload points at the offending line. The loc~ labelled argument also lets callers forward their own loc when wrapping check (see require below).

    v0.48.0: check now raises PreconditionError::Violated(loc) instead of aborting. Every caller is expected to either (a) declare raise PreconditionError in its own signature (preferred for new code), or (b) wrap the call site in try { ... } catch { e :PreconditionError => abort(e.to_string()) } to preserve the pre-v0.48.0 abort behavior.

    The pre-v0.48.0 abort message format ("precondition failed at ") is preserved by the re-abort pattern in callers, so the diagnostic surface is unchanged for end users.

    check_make_violated

    source
    #callsite(autofill(loc))
    fn check_make_violated(loc~ : SourceLoc) -> PreconditionError

    Construct a PreconditionError::Violated(loc) whose payload is the source location of the caller (the #callsite(autofill(loc)) attribute auto-injects the call-site SourceLoc into the loc parameter, so the diagnostic message points at the line that called this helper, not at this helper's own definition site).

    Added in v0.47.0 as a forward-compatible hook for the upcoming check/require → raise PreconditionError conversion. Today it has no prod caller (the central check/require still aborts), so the only usage is the regression test in check_test.mbt. In v0.48.0+, check will raise this error directly and this helper stays as an opt-in builder for callers that want to surface a custom precondition violation without going through the global check path.

    The pub visibility is required so the _test file can call it; see v0.37.0's apply_calibration for the same pub-but-test-only helper pattern.

    cholesky

    source
    fn cholesky(a : Matrix) -> Matrix

    Cholesky decomposition. Computes a lower triangular matrix L such that A = L * L^T for a symmetric positive-definite matrix A.

    The decomposition is performed in place. L is stored in the lower triangle of a (the upper triangle is not referenced and may contain garbage). Panics if a is not square or if the matrix is not positive definite (a non-positive pivot is treated as a fatal error).

    This is a textbook implementation. For the small problem sizes used in the DoubleML examples it is more than fast enough.

    clip_vec

    source
    fn clip_vec(v : Array[Double], lo : Double, hi : Double) -> Array[Double]

    Clip every element of v to [lo, hi] in place (returns a new array; does not mutate the input).

    cluster_causal_param_and_se

    source
    fn cluster_causal_param_and_se(psi_a : Array[Double], psi_b : Array[Double], folds_row : Array[Fold], fold_n_units : Array[Int], unit_rows : Array[Array[Int]], unit_fold : Array[Int], n_folds_u : Int, n_folds_per_cluster : Int) -> (Double, Double) raise VarEstClusterError

    Cluster-robust causal parameter + SE computation. Combines est_coef_cluster (fold-weighted ratio of cluster score sums) and var_est_cluster (unit-level cluster-robust SE) over the same per-row score elements. Used by every clustered-DML path that follows the PLR-style psi_a, psi_b pattern (PLR, IRM, PLPR, PLIV, IIVM). The n argument is the row count used for psi_res allocation; the caller pre-computes psi_a and psi_b at its converged coefficient / nuisance values.

    Returns (theta_r, se_r). The !VarEstClusterError annotation propagates the J-floor defensive guard from var_est_cluster (v0.34.0+) so callers can decide whether to abort, retry with a different seed, or surface the error to downstream consumers. All fit_cluster() methods in PLR / IRM / PLPR / PLIV / IIVM catch this error and re-abort to preserve the pre-v0.35.0 process-death behavior on pathological fold splits.

    compute_sensitivity_bias

    source
    fn compute_sensitivity_bias(sigma2 : Double, nu2 : Double, psi_sigma2 : Array[Double], psi_nu2 : Array[Double]) -> (Array[Double], Array[Double])

    Cinelli & Hazlett (2020) omitted-variable bias analysis for DML estimators. Given the per-density influence-function scalars sigma2, nu2, psi_sigma2, psi_nu2, this helper computes the maximum bias from an unobserved confounder that shifts the outcome nuisance sigma2 by psi_sigma2 and the treatment nuisance nu2 by psi_nu2.

    The two quantities returned are
    • max_bias: the worst-case absolute bias on the estimator, equal to sqrt(sigma2 * nu2).
    • psi_max_bias: the gradient of max_bias with respect to the confounding strength vector, equal to (sigma2 * psi_nu2 + nu2 * psi_sigma2) / (2 * max_bias).

    Both quantities are length-n_obs arrays (per-observation influence vectors); the caller typically averages over the sample to get the scalar bias. Matches the upstream doubleml.utils._sensitivity._compute_sensitivity_bias helper.

    cross_fit_predict

    source
    fn[T : Learner] cross_fit_predict(learner : T, x : Matrix, y : Array[Double], folds : Array[Fold]) -> Array[Double]

    Cross-fit prediction helper. For each fold, fit the learner on the training slice (x[train], y[train]) and predict on the test slice x[test]. The returned vector has length n_obs and contains the predictions in the original observation order. Assumes the folds together cover [0, n_obs).

    cross_fit_predict_dispatch

    source
    fn cross_fit_predict_dispatch(learner : LearnerDispatch, x : Matrix, y : Array[Double], folds : Array[Fold]) -> Array[Double]

    Cross-fit a LearnerDispatch over the given folds. This is the dispatch wrapper around cross_fit_predict[T : Learner] (which is generic in T); each match arm binds a concrete T so the generic call resolves correctly.

    Adding a new learner type: add the corresponding match arm here (and the enum variant in LearnerDispatch).

    did_bootstrap_t_stat

    source
    fn did_bootstrap_t_stat(weights : Array[Double], psi : Array[Double], se : Array[Double], n_rep_boot : Int, n_obs : Int, n_thetas : Int) -> Array[Double]

    Per-cell multiplier bootstrap t-statistics.

    Parameters:
    • weights : Array[Double] — flat row-major, shape [n_rep_boot, n_obs]. Drawn by draw_bootstrap_weights(method_name, n_rep_boot,n_obs, seed) in the calling estimator.
    • psi : Array[Double] — flat row-major, shape [n_thetas, n_obs]. Each row is the per-observation influence function psi_k[i] for theta k. For DoubleMLDIDMulti, this is a slice of psi_matrix[flat * n_obs : (flat+1) * n_obs] for each group-time (gi, pi). For DoubleMLDIDCrossSection, this is a single row [n_obs] of psi_a[i] + coef * psi_b[i].
    • se : Array[Double] — length n_thetas. The standard error of the per-cell estimate. For DIDMulti, se[k] = se_matrix[gi * n_periods + pi]. For DIDCrossSection, se[0] = sqrt(mean(psi^2)).
    • n_rep_boot : Int — number of bootstrap reps.
    • n_obs : Int — sample size.
    • n_thetas : Int — number of theta values (= gt_combinations.length() for DIDMulti; = 1 for DIDCrossSection).

    Returns:
    • boot_t_stat : Array[Double] — flat row-major, shape [n_rep_boot, n_thetas]. Indexed as boot_t_stat[b * n_thetas + k] for rep b, theta k.

    Algorithm: for b in 0..n_rep_boot: for k in 0..n_thetas: if se[k] == 0: skip (pre-treatment / empty cell) s = 0 for i in 0..n_obs: s += weights[b * n_obs + i] * psi[k * n_obs + i] boot_t_stat[b * n_thetas + k] = s / (sqrt(n) * se[k])

    Rows where se[k] == 0 are set to 0 (matches the upstream convention: empty cells produce no bootstrap variation).

    dot

    source
    fn dot(a : Array[Double], b : Array[Double]) -> Double

    Dot product of two equal-length vectors.

    The accumulator is Kahan-compensated so the result is robust to near-cancellation between terms (see kahan.mbt).

    double_cross_fit_predict

    source
    fn[T : Learner] double_cross_fit_predict(learner : T, x : Matrix, y : Array[Double], outer : Array[Fold], n_inner : Int, seed : Int) -> Array[Array[Double]]

    Double cross-fit prediction helper used by IRM-style models (e.g. LPLR). outer is the row-level outer-fold partition; for each outer fold, the model's training slice is re-split into n_inner inner folds via kfold, the learner is fit on each inner training slice and predicts on the corresponding inner test slice, and the resulting n_inner arrays of predictions are returned. The outer model then has, for each outer test row, an inner OOF prediction — used downstream to construct nuisance W and the preliminary beta.

    Returns an Array[Array[Double]] of length outer.length(), one entry per outer fold. Each entry has length inner_train_count (the size of the outer fold's training slice); values are filled at the original training-row positions in the order the inner folds visit them.

    double_cross_fit_predict_dispatch

    source
    fn double_cross_fit_predict_dispatch(learner : LearnerDispatch, x : Matrix, y : Array[Double], outer : Array[Fold], n_inner : Int, seed : Int) -> Array[Array[Double]]

    v0.62.0+: LearnerDispatch wrapper around double_cross_fit_predict[T : Learner]. Routes through the same enum pattern as cross_fit_predict_dispatch so the ml_m per-fit override (typically LogisticRegression for LPLR) actually reaches the inner OOF predictions.

    draw_bootstrap_weights

    source
    fn draw_bootstrap_weights(method_name : String, n_rep_boot : Int, n_obs : Int, seed : Int) -> Array[Double] raise BootstrapMethodError

    Draw n_rep_boot weight vectors of length n_obs from the chosen multiplier distribution. The chacha8 RNG is seeded with seed for reproducibility.

    Supported methods:
    • "normal": w[i] ~ N(0, 1) per the Box-Muller transform.
    • "Bayes": w[i] = exp(1) - 1 (mean 0, var 1; uses the chacha8-driven rng.double() for the uniform quantile input to the inverse-CDF).
    • "wild": w[i] = x[i] / sqrt(2) + (y[i]^2 - 1) / 2 with x, y ~ N(0, 1). The wild bootstrap is robust to heteroskedasticity in the influence-function residuals.

    Returns a row-major Array[Double] of length n_rep_boot * n_obs. The first n_obs entries are the first bootstrap replication's weights, the next n_obs are the second replication, and so on.

    Returns Array[Double] raise BootstrapMethodError: the v0.37.0 conversion replaces the previous abort() call with raise BootstrapMethodError::UnknownMethod(method_name) so the unknown-method path becomes directly testable. Callers that want the pre-v0.37.0 process-death behavior should catch the error and re-abort (this is what every prod caller does); callers that want to surface the error to downstream consumers should propagate via ?. The error type is declared in kfold.mbt so the cluster helper stack can share it.

    est_coef_cluster

    source
    fn est_coef_cluster(psi_a : Array[Double], psi_b : Array[Double], folds_row : Array[Fold], fold_n_units : Array[Int]) -> Double

    Clustered-path causal parameter (LinearScoreMixin._est_coef, cluster branch): the root of the fold-weighted score sums. Each test fold contributes sum(score on its test rows) scaled by 1 / |I_k|, the number of units in the fold's test cluster set:

     theta = - sum_k w_k * sum_{i in k} psi_b
         / sum_k w_k * sum_{i in k} psi_a,   w_k = 1/|I_k|.

    expand_unit_folds_to_rows

    source
    fn expand_unit_folds_to_rows(cluster : Array[Int], folds_u : Array[Fold], row_unit : Array[Int]) -> (Array[Fold], Array[Int], Array[Int])

    Clustered-path fold builder. Takes the unit-level fold partition (a kfold(n_units, n_folds, seed) result) and expands it to a row-level fold partition where every row of a given unit lands on the same side of every split. Also returns the per-unit test-fold index and the per-fold unit count, both of which are needed by est_coef_cluster / var_est_cluster.

    cluster : length-n vector of unit ids folds_u : unit-level folds (each Fold carries test_indices = unit positions) row_unit : length-n row → unit-position map (use build_row_unit_map)

    Returns (folds_row, unit_fold, fold_n_units):
    • folds_row : length-n_folds row-level folds
    • unit_fold : length-n_units per-unit test-fold index
    • fold_n_units : length-n_folds per-fold unit count

    expit

    source
    fn expit(x : Double) -> Double

    Public elementwise expit(x) = 1 / (1 + exp(-x)) for callers that need to evaluate the logistic link (e.g. LPLR's nonlinear score). Numerically safe at both tails: for x >= 0 the 1/(1+exp(-x)) form avoids large exp; for x < 0 the exp(x)/(1+exp(x)) form avoids underflowing to 0. The output is in the open interval (0, 1) for any finite input.

    filter_by_value

    source
    fn filter_by_value(idx : Array[Int], cond : Array[Double], target : Double) -> Array[Int]

    Filter idx to keep only entries i for which cond[i] matches the desired value target. Used to build the conditional sample splits for g0/g1/r0/r1.

    filter_indices

    source
    fn filter_indices(idx : Array[Int], mask : Array[Double]) -> Array[Int]

    Filter idx to keep only entries i for which mask[i] is true.

    filter_two_values

    source
    fn filter_two_values(idx : Array[Int], m1 : Array[Double], v1 : Double, m2 : Array[Double], v2 : Double) -> Array[Int]

    Filter idx to keep only entries i where both mask1[i] and mask2[i] equal the given values. Used to build train_d{s_value}_s1 for the conditional-outcome learners.

    fit_isotonic

    source
    fn fit_isotonic(x : Array[Double], y : Array[Double]) -> (Array[Double], Array[Double])

    Sort (x, y) pairs by x and apply PAVA to the sorted y. Returns (sorted_x, sorted_y_hat) where sorted_y_hat is the isotonic regression of y on x. The two arrays have the same length as the input.

    Use predict_isotonic to apply the fitted model to new x values (or to the same x for in-sample predictions, which is the upstream IsotonicRegression default).

    fit_predict_one_dispatch

    source
    fn fit_predict_one_dispatch(learner : LearnerDispatch, x_train : Matrix, y_train : Array[Double], x_test : Matrix) -> Array[Double]

    Single-fit + single-predict helper for a LearnerDispatch. Used by estimators (e.g. DoubleMLSSM for the augmented (X, D) -> S pi fit, DoubleMLIIVM for conditional subsets where the augmenting column changes the feature matrix shape) that need a one-off fit on a subset followed by a predict on a different-shape test set — i.e. cases where cross_fit_predict_dispatch (which slices a single matrix by train/test indices) doesn't apply.

    v0.59.0+.

    g_cross_fit_calls

    source
    fn g_cross_fit_calls() -> Int

    Read the current g cross-fit count. Public for tests.

    gain_statistics

    source
    fn gain_statistics(dml_long : GainStatsSource, dml_short : GainStatsSource) -> GainStatsResult

    Compute the gain-statistic benchmark values for cf_y, cf_d, rho, and delta_theta from two fitted DML models: dml_long (which includes all observed confounders) and dml_short (which excludes one or more "benchmark" confounders).

    Each of dml_long / dml_short exposes four per-rep arrays (shape (n_coef, n_rep)):
    • all_coef: the per-rep coefficient estimates.
    • var_y_residuals: the per-rep outcome-residual variance (i.e. sigma2).
    • nu2: the per-rep treatment-residual variance (i.e. the squared norm of the Riesz representer).
    • var_y: the overall outcome variance (scalar, shared between dml_long and dml_short).

    The algorithm (matches upstream doubleml.utils.gain_statistics.gain_statistics):
    1. R2_y_long = 1 - var_y_residuals_long / var_y, R2_y_short = 1 - var_y_residuals_short / var_y, R2_riesz = nu2_short / nu2_long.
    2. cf_y = clip((R2_y_long - R2_y_short) / (1 - R2_y_long), 0, 1).
    3. cf_d = clip((1 - R2_riesz) / R2_riesz, 0, 1).
    4. delta_theta = median(dml_short.all_coef - dml_long.all_coef, axis=rep).
    5. rho = median(sign(delta_theta) * |delta_theta| / sqrt(var_g * var_riesz)), where var_g = var_y_residuals_short - var_y_residuals_long and var_riesz = nu2_long - nu2_short (clipped to [-1, 1] and 0 when the denominator is 0).
    6. Return the median across reps for each coefficient.

    The per-coefficient cf_y and cf_d benchmark the "tipping point" of the sensitivity analysis: a confounder with strength <= cf_y (resp. cf_d) cannot change the conclusion. rho and delta_theta are diagnostic: they tell the user the direction and magnitude of the change between dml_long and dml_short.

    gaussian_kde

    source
    fn gaussian_kde(y : Array[Double], x : Double, h : Double) -> Double

    Gaussian KDE density estimate at x, evaluated using samples y[0..n) and bandwidth h:

     f_hat(x) = (1 / (n * h * sqrt(2*pi))) * sum_i exp(-(x - y_i)^2 / (2 * h^2))

    Returns a non-negative density. The constant sqrt(2 * pi) =1.7724538509055159 is the standard Gaussian normalisation factor.

    gaussian_kde_weighted

    source
    fn gaussian_kde_weighted(y : Array[Double], w : Array[Double], x : Double, h : Double) -> Double

    Weighted Gaussian KDE density estimate at x, evaluated using samples y[0..n) with per-sample weights w[0..n) and bandwidth h:

     f_hat(x) = (1 / (h * sqrt(2*pi))) * sum_i w[i] * exp(-(x - y[i])^2 / (2*h^2))

    Note: the weights enter without a / sum(w) normalisation, so this is a weighted KDE estimate (not a probability density unless sum(w) == 1). Used by DoubleMLLPQ::fit to compute the numerical derivative of mean(psi(theta)) w.r.t. theta:

     d/dtheta mean(psi_ipw) = (1/n) * sum_i w_i * delta(y_i - theta)
                        ≈ (1/n) * f_hat_weighted(theta)

    where w_i = sign * (z_i/m_i - (1-z_i)/(1-m_i)) * treated_i / comp is the IPW coefficient. The derivative feeds directly into the delta-method variance se = sqrt(gamma / (deriv^2 * n)).

    holm_bonferroni_p_adjust

    source
    fn holm_bonferroni_p_adjust(unadjusted : Array[Double]) -> Array[Double]

    Holm-Bonferroni stepdown correction. Sort unadjusted p-values ascending, then p_corrected_sorted[k] =max((n - k) * p_sorted[k], p_corrected_sorted[k - 1]), clipped to 1.0. Re-order to original cell order.

    inv_spd

    source
    fn inv_spd(a : Matrix) -> Matrix

    Invert a symmetric positive-definite matrix using the Cholesky factorisation A = L L^T. Returns the full inverse matrix. Used by LinearRegression::fit to populate the (X^T X)^{-1} diagonal that drives the per-coefficient standard errors. Small (≤ a few hundred rows in DML) so the O(n^3) cost is invisible at the test scale.

    isotonic_calibrate_cv

    source
    fn isotonic_calibrate_cv(x : Array[Double], y : Array[Double], cv : Array[(Array[Int], Array[Int])]?) -> Array[Double] raise CalibrationFittingError

    K-fold cross-validated isotonic calibration. For each fold (train_idx, test_idx), fit an isotonic regression on (x[train], y[train]) and predict at x[test]. The output array is built in the original x order (not the fold order) by writing out[test_idx[i]] = pred[i].

    cv is a list of (train_idx, test_idx) pairs. When cv =None, a deterministic 5-fold split with seed=3141 is generated via kfold. Each fold's prediction is independent of every other fold's, so this is embarrassingly parallel (sequential here for simplicity, but the per-fold work is O(n log n) for the sort + O(n) for the PAVA scan).

    Returns Array[Double] raise CalibrationFittingError: the v0.38.0 conversion replaces the previous abort() call with raise CalibrationFittingError::IncompleteCVPartition when the cv partition fails to cover every input index. Callers that want the pre-v0.38.0 process-death behavior should catch the error and re-abort (this is what apply_calibration does); callers that want to surface the error to downstream consumers should propagate via ?.

    kahan_sum

    source
    fn kahan_sum(arr : Array[Double]) -> Double

    Kahan compensated summation: returns sum(arr) with a running compensation term that recovers the low-order bits lost at each addition step. The result is effectively computed in double-double precision while staying inside the standard Double type.

    Why we need this

    Naive summation

     s = 0.0
 for x in arr: s = s + x

    loses the lowest eps * |s| bits of each addend on every step. The worst-case relative error is O(n * eps) for n addends; when the addends alternate in sign or have wildly different magnitudes (the "catastrophic cancellation" pattern), the worst-case error grows to O(n^2 * eps) and the running sum can drift by orders of magnitude more than expected.

    In the DML numeric kernels this matters in two places:

    1. var_est.mbt: the gamma = gamma + psi^2 accumulator runs for n = 500 observations; psi is in roughly [-2, 2] so each addend is O(1), but partial sums of opposite-sign squared terms are not catastrophic — Kahan is "free insurance" here.
    2. matrix.mbt::matmul and matvec and linalg.mbt::cholesky: the inner-product accumulators add O(p) products of O(1) terms; the partial sum stays around O(1), so naive sum is fine in well-conditioned cases but Kahan removes the ~p * eps drift when the partial sums nearly cancel.

    We keep Kahan in the inner-product and the gamma loop (where it is the cheapest high-value change), and we also upgraded mean (the O(n) sum that every estimator calls twice per fit) to Kahan; the 5-seed precheck confirms the upgrade is invisible at 15-digit output and stays inside the TODO #5 thresholds (|theta - 1| < 0.2 for PLR, < 0.5 for the other 4).

    Algorithm

    Standard Kahan summation (Kahan 1965), also called "compensated summation":

     sum = 0.0
 c   = 0.0    // compensation for low-order bits lost in the next add
 for x in arr:
   y = x - c           // align x with the running sum's precision
   t = sum + y         // primary add
   c = (t - sum) - y   // low-order bits of `t` that did not fit
   sum = t
 return sum

    The invariant is that sum + c is the true (extended-precision) running sum; the returned sum carries the high-order bits and c holds the residual. After each step |c| <= eps * |sum|, so the error stays O(eps) instead of O(n * eps).

    kfold

    source
    fn kfold(n_obs : Int, n_folds : Int, seed : Int) -> Array[Fold]

    kfold_stratified

    source
    fn kfold_stratified(n : Int, strata : Array[Double], n_folds : Int, seed : Int) -> Array[Fold]

    Stratified k-fold partition: each fold preserves the per-stratum proportions of the input. Used by DoubleMLAPOS (treatment-level stratification) and other multi-level IRM estimators where the treatment distribution must be balanced across folds to avoid empty-treatment folds that would zero-out the IPW denominator.

    The implementation is a within-stratum k-fold partition followed by fold-merging: each stratum independently permutes its indices, then every fold f collects the f-th slice of each stratum's permutation. This produces balanced folds in O(n) time and matches the upstream StratifiedKFold(n_splits=n_folds, shuffle=True) semantics from sklearn.

    The strata array must be parallel to the n observations (same length). Each unique value in strata becomes a stratum; only strata with at least n_folds members can be balanced. Strata with fewer members contribute all their members to whichever fold is currently being filled (the same behavior as sklearn's StratifiedKFold with the default fallback).

    logit

    source
    fn logit(p : Double, eps? : Double) -> Double

    Public elementwise logit(p) = log(p / (1 - p)). Inputs are clamped to [eps, 1 - eps] (default eps = 1e-8) so the result stays finite; the LPLR score clips inner-fold M predictions to [1e-8, 1 - 1e-8] upstream before logit, mirroring scipy logit(clip(p, 1e-8, 1 - 1e-8)).

    lpq_score_ipw

    source
    fn lpq_score_ipw(data : DoubleMLLPQData, treated : Array[Double], m : Array[Double], comp : Double, theta : Double, q : Double, sign : Double) -> Array[Double]

    IPW-only LPQ score (no g0/g1 cross-fit). Used as the bisection objective in DoubleMLLPQ::fit (Bug #3 fix). Matches the upstream doubleml.irm.lpq.DoubleMLLPQ._compute_ipw_score: score[i] = sign * (z[i] / m[i] - (1 - z[i]) / (1 - m[i]))* treated[i] * (y[i] <= theta ? 1 : 0) / comp - q The bisection does not need g0/g1, so this avoids the 2 * 50 = 100 g cross-fits that the pre-fix code did.

    make_confounded_plr_data

    source
    fn make_confounded_plr_data(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> PlrConfoundedData

    make_did_CS2021

    source
    fn make_did_CS2021(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> DidMultiData

    Generate the multi-cohort DID DGP. n_obs is the target total sample size; the actual n is n_per_cohort * n_groups *n_periods. The function chooses n_per_cohort to match n_obs.

    make_did_SZ2020

    source
    fn make_did_SZ2020(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> DidCsData

    make_did_cs_CS2021

    source
    fn make_did_cs_CS2021(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> DidCsData

    Generate the DID CS2021 DGP. Default: 2 cohorts, 2 periods, theta = 1.0, balanced cohort sizes, n_per_cohort = 200, dim_x = 3.

    make_iivm_data

    source
    fn make_iivm_data(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> IivmData

    Generate the IIVM DGP from Chernozhukov et al. (2018) Section 4.4.

    make_irm_confounded_data

    source
    fn make_irm_confounded_data(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> IrmConfoundedData

    make_irm_data

    source
    fn make_irm_data(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> IrmData

    Generate the IRM DGP from Chernozhukov et al. (2018) Section 4.2.

    make_irm_discrete_treatments

    source
    fn make_irm_discrete_treatments(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> IrmDiscreteData

    make_irm_heterogeneous_data

    source
    fn make_irm_heterogeneous_data(n_obs : Int, dim_x : Int, seed : Int) -> IrmHeterogeneousData

    make_lplr_LZZ2020

    source
    fn make_lplr_LZZ2020(n : Int, alpha : Double, seed : Int) -> LplrData

    Generate the LPLR LZZ2020 DGP.

    make_pliv_CHS2015

    source
    fn make_pliv_CHS2015(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> PlivData

    make_pliv_data

    source
    fn make_pliv_data(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> PlivData

    Convenience factory matching the upstream doubleml.plm.datasets._make_pliv_data.make_pliv_data signature: (n_obs, dim_x, theta, seed) with no other arguments. Equivalent to make_pliv_CHS2015(n_obs, dim_x,theta, seed) but named for parity with the upstream factory helper.

    make_pliv_multiway_cluster

    source
    fn make_pliv_multiway_cluster(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> PlivClusterData

    make_plpr_CP2025

    source
    fn make_plpr_CP2025(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> PlprData

    Generate the PLPR DGP. n_obs is the target total sample size; the actual n is n_units * n_periods.

    make_plr_CCDDHNR2018

    source
    fn make_plr_CCDDHNR2018(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> PlrCcddhnr2018

    Generate the PLR DGP from Chernozhukov et al. (2018) Figure 1.

    make_plr_turrell2018

    source
    fn make_plr_turrell2018(n_obs : Int, theta : Double, seed : Int) -> IrmHeterogeneousData

    make_simple_rdd_dgp

    source
    fn make_simple_rdd_dgp(n_obs : Int, theta : Double, seed : Int) -> RddSimpleData

    Generate the simple RDD DGP.

    make_ssm_data

    source
    fn make_ssm_data(n_obs : Int, dim_x : Int, theta : Double, seed : Int) -> SsmData

    Generate the SSM DGP.

    matmul

    source
    fn matmul(a : Matrix, b : Matrix) -> Matrix

    A * B matrix-matrix product. Panics if inner dimensions do not match.

    The inner product over the shared k dimension uses kahan.mbt-style compensated summation so the partial sum does not drift by O(p * eps) when terms of opposite sign nearly cancel.

    matvec

    source
    fn matvec(a : Matrix, x : Array[Double]) -> Array[Double]

    A * x matrix-vector product. x length must equal A.ncols. Returns a vector of length A.nrows.

    The inner-product accumulator is Kahan-compensated (see kahan.mbt).

    matvec_t

    source
    fn matvec_t(a : Matrix, x : Array[Double]) -> Array[Double]

    A^T * x matrix-vector product.

    mean

    source
    fn mean(a : Array[Double]) -> Double

    Mean of a vector.

    The accumulator is Kahan-compensated so the running sum does not drift by O(n * eps) for long inputs. The empty-input branch short-circuits to 0.0 (the package convention).

    newton_solve_score

    source
    fn newton_solve_score(theta_start : Double, psi : Array[Double], psi_deriv : Array[Double], tol_theta? : Double, max_iter? : Int) -> (Double, Bool)

    Newton iteration to find the root of mean(score(theta)) for a given psi(theta) = psi_a * theta + psi_b-style linear score (i.e. psi and psi_deriv are constants with respect to theta). For the general nonlinear case, callers should re-evaluate psi(theta) / psi_deriv(theta) between iterations — see the LPLR fit body, which evaluates the score at the converged theta after this routine returns.

    Mirrors scipy.optimize.root_scalar(method="newton") with fprime = score_deriv. Stops when |theta_{k+1} - theta_k| < tol_theta (default 1e-9) or max_iter (default 50) is reached. Returns the converged theta and a Boolean indicating whether the iteration converged.

    norm_cdf

    source
    fn norm_cdf(x : Double) -> Double

    v0.21.0+ standard-normal CDF. Implements the A&S 7.1.26 formula directly (rather than going through norm_sf, which is the upper-tail survival and saturates to 1.0 at x <= 0). The coefficients match the existing norm_sf: p = 0.2316419, b1 = 0.319381530, b2 = -0.356563782,b3 = 1.781477937, b4 = -1.821255978, b5 = 1.330274429 and the max abs error is ~7.5e-8. Symmetry Phi(x) = 1 - Phi(-x) handles the negative half-line.

    norm_ppf

    source
    fn norm_ppf(p : Double) -> Double

    v0.21.0+ standard-normal quantile function (inverse CDF). Uses bisection on the existing norm_cdf (which delegates to norm_sf with ~7.5e-8 accuracy) to find the z such that Phi(z) = p. The bisection runs for 64 iterations on the interval [-8, 8], which is more than enough to drive the absolute error below 1e-12.

    p must lie in (0, 1). The result is the x such that P(Z <= x) = p for a standard normal random variable Z.

    norm_sf

    source
    fn norm_sf(x : Double) -> Double

    Standard-normal survival function P(Z > x) using the Abramowitz & Stegun (1964) formula 7.1.26 (max absolute error ~7.5e-8 for x >= 0):

    sf(x) = phi(x) * (b1*t + b2*t^2 + b3*t^3+ b4*t^4 + b5*t^5) phi(x) = exp(-x^2 / 2) / sqrt(2 pi) t = 1 / (1 + p * x) p = 0.2316419, b1 = 0.319381530, b2 = -0.356563782,b3 = 1.781477937, b4 = -1.821255978, b5 = 1.330274429

    p_adjust

    source
    fn p_adjust(method_name : String, unadjusted : Array[Double], boot_t_stat? : Array[Double], t_stats? : Array[Double]) -> Array[Double]

    Public dispatcher. method accepts:
    • "romano-wolf" / "rw" -> romano_wolf_p_adjust
    • "holm" -> holm_bonferroni_p_adjust
    • "bonferroni" / "bonf" -> bonferroni_p_adjust
    • "bh" / "fdr_bh" -> bh_fdr_p_adjust
    • "by" / "fdr_by" -> by_fdr_p_adjust
    • "tsbh" / "fdr_tsbh" -> tsbh_p_adjust
    • "tsby" / "fdr_tsbky" -> tsby_p_adjust

    For romano-wolf, the caller must supply boot_t_stat (length n_boot * n_thetas) and t_stats (length n_thetas); the unadjusted argument is ignored. For all other methods, boot_t_stat and t_stats are ignored; only unadjusted (length n_thetas) matters.

    pava

    source
    fn pava(y : Array[Double], weights? : Array[Double]) -> Array[Double]

    Pool-adjacent-violators algorithm. Given a sequence of values y[0..n] (assumed already sorted by the predictor x, which is monotone non-decreasing in the index), pava returns the isotonic (non-decreasing) L2 projection of y. Tied predictors are handled naturally by the algorithm (they form a single block whose mean is the projection value).

    Optional weights[0..n] is a per-element weight; default is unit weight for every element. Each output block is the weighted mean of its constituent elements.

    The algorithm walks the input once, maintaining a stack of "blocks" — each block holds (sum_y, sum_w, size). When a new value would create a violation (the previous block's mean is greater than the new block's mean), the algorithm pools the two blocks and re-checks. The result is the canonical weighted-PAVA output: a non-decreasing sequence that minimises the weighted sum of squared residuals subject to the monotonicity constraint.

    The input MUST be sorted by x in non-decreasing order. Use fit_isotonic for the public entry point that sorts and returns the fitted model.

    pq_score_ipw

    source
    fn pq_score_ipw(_x : Matrix, y : Array[Double], treated : Array[Double], m : Array[Double], theta : Double, q : Double) -> Array[Double]

    IPW score for the potential quantile, used as the bisection objective in solve_pq (Bug #3 fix). The full PQ score also subtracts a g cross-fit, but the g is not needed for the bisection: at the root theta, mean(score) = 0 regardless of g because E[g(X) | D = d] = E[g(X) * 1{D = d} / m(X)] by definition of the cross-fit. So we can iterate the bisection with this cheap score (no OLS fit per iteration) and only cross-fit g ONCE at the resulting theta_prelim.

    Math: score[i] = treated[i] / m[i] * 1{y[i] <= theta} - q. Matches the upstream doubleml.irm.pq.DoubleMLPQ._compute_ipw_score.

    predict_isotonic

    source
    fn predict_isotonic(fitted_x : Array[Double], fitted_y_hat : Array[Double], x_new : Array[Double]) -> Array[Double]

    Predict the isotonic regression at new x values using the fitted model (fitted_x, fitted_y_hat). Behaviour matches sklearn.isotonic.IsotonicRegression(out_of_bounds="clip",y_min=0.0, y_max=1.0):
    • For x_new[i] strictly less than min(fitted_x): return fitted_y_hat[0] (clipped at the lower boundary).
    • For x_new[i] strictly greater than max(fitted_x): return fitted_y_hat[last] (clipped at the upper boundary).
    • Otherwise: return fitted_y_hat at the largest fitted_x[j]<= x_new[i], found by linear scan (PAVA produces a step function, and the step boundaries are the fitted_x values themselves).

    We additionally clip the prediction to [0.0, 1.0] since the fitted y_hat values are guaranteed to be in [0, 1] for binary y (PAVA output on a 0/1 input lies in [0, 1]), but the clip makes the contract explicit and protects against any numerical drift.

    range_indices

    source
    fn range_indices(n : Int) -> Array[Int]

    The full index range [0, n). Convenient for callers that need to pass a complete index list to filter_indices (e.g. the LPQ complier-prob computation in lpq.mbt::DoubleMLLPQ::fit).

    repeated_kfold

    source
    fn repeated_kfold(n_obs : Int, n_folds : Int, n_rep : Int) -> Array[Array[Fold]]

    Repeated K-fold partition. Draws n_rep independent kfold partitions with seeds 3141, 3142, ..., 3141 + n_rep - 1, and returns them as a list of length n_rep, each containing n_folds Fold objects. Matches the upstream doubleml.utils.resampling.DoubleMLResampling with RepeatedKFold and stratify=None.

    require

    source
    #callsite(autofill(loc))
    fn require(condition : Bool, loc~ : SourceLoc) -> Unit raise PreconditionError

    Internal alias of check, kept for terse call sites that already read like require(x == y). Forwards the auto-injected loc so the error payload still points at the offending source line. v0.48.0: forwards check's raise as-is (no local try/catch); the per-caller wrap is responsible for converting the raise back to abort at the call site.

    reset_g_cross_fit_count

    source
    fn reset_g_cross_fit_count() -> Unit

    Reset the g cross-fit counter to 0. Public so blackbox tests can use it to bracket a fit() call.

    robustness_value

    source
    fn robustness_value(theta : Double, max_bias : Array[Double]) -> Double

    Aggregate the per-observation sensitivity vectors into a scalar "robustness value" RV — the minimum confounding strength that would change the estimator's sign. Returns a large sentinel (1.0e300) if the estimator is already zero (degenerate) since MoonBit's Double has no Infinity constant in this build. Per Cinelli & Hazlett (2020) §3.2: RV = |theta_hat| / mean(max_bias).

    romano_wolf_p_adjust

    source
    fn romano_wolf_p_adjust(boot_t_stat : Array[Double], unadjusted : Array[Double], t_stats : Array[Double]) -> Array[Double]

    Romano-Wolf (2005) stepdown p-adjustment. The stepdown critical value at rank i_theta is max_{j > i_theta} |boot_t_stat[b, j]| per bootstrap replication b. The unadjusted p-value at rank i_theta is frac{b : cv_b >= |t_{i_theta}|} / n_boot. Enforces monotonicity from the smallest |t| upward. Re-orders to the original cell order.

    boot_t_stat is row-major (n_boot, n_thetas).

    The algorithm does not depend on the unadjusted p-values directly; it consumes t_stats (length n_thetas). The unadjusted argument is accepted for API symmetry with the other p-adjust entry points and is ignored.

    seed_to_bytes

    source
    fn seed_to_bytes(seed : Int) -> Array[Byte]

    Pack an integer seed into a 32-byte buffer using a 32-bit little-endian encoding of seed (as a signed 32-bit integer; -1 becomes 0xFF FF FF FF). The four bytes are tiled 8 times to fill the 32-byte key that chacha8 requires. Distinct integer seeds give distinct byte strings, and the layout is independent of the host endianness.

    Encoding (per byte index 0..3): byte[0] = seed & 0xff byte[1] = (seed >> 8) & 0xff byte[2] = (seed >> 16) & 0xff byte[3] = (seed >> 24) & 0xff

    seed is a signed 64-bit Int; the upper 32 bits are ignored so the 32-bit LE view is well-defined for any seed. For negative seeds, arithmetic right-shift + mask is sign-clean (e.g. seed = -1 gives [0xff, 0xff, 0xff, 0xff]).

    This is the canonical seed-to-bytes helper for the dml package's RNG-backed tests and for cmd/main's demo entry point. It supersedes the earlier 7-bit-tiling helper that existed in irm_test.mbt and cmd/main/main.mbt and produced non-portable encodings.

    silverman_bandwidth

    source
    fn silverman_bandwidth(y : Array[Double]) -> Double

    Silverman's rule of for-band for a Gaussian KDE: given samples y[0..n),h = 0.9 * min(sd(y), IQR/1.34) * n^(-1/5). The min(sd, IQR/1.34) factor is robust to outliers (a sample of N(0, 1) typically has IQR/1.34 ≈ sd; heavy-tailed samples get the smaller IQR, which prevents oversmoothing).

    For samples drawn from a single-mode distribution the bandwidth is asymptotically optimal; for multimodal or heavy-tailed data the user should switch to a plug-in estimator (cross-validated or Sheather-Jones) — but the simpler Silverman rule is fine for the LPQ numerical-derivative use case where y is typically uniform- shaped or mildly bimodal.

    slice_matrix_rows

    source
    fn slice_matrix_rows(x : Matrix, idx : Array[Int]) -> Matrix

    Slice a matrix by a list of row indices, returning a new n x p matrix with the selected rows in the given order.

    slice_vector

    source
    fn slice_vector(v : Array[Double], idx : Array[Int]) -> Array[Double]

    Slice a vector by a list of indices.

    solve_pq

    source
    fn solve_pq(ml_l : LearnerDispatch, ml_m : LearnerDispatch, data : DoubleMLData, treatment : Double, q : Double, n_folds : Int, seed : Int, clip : Double) -> (Double, Array[Double], Double) raise BracketSignError

    Solve the potential quantile via IPW bisection, then return the final theta, the influence-function psi at theta (using the g cross-fit at theta), and the numerical derivative d mean(psi) / d theta (using two extra g cross-fits at theta +/- h). Bug #2 and #3 fix: previously returned (theta, se) and recomputed g on every bisection step.

    Returns: (theta, psi, deriv) where psi : Array[Double] of length n and deriv : Double.

    This is pub so the blackbox test quantile_test::qte_se_hand_computation can re-derive the QTE SE by hand from the per-treatment solve_pq outputs. Internal-only callers (DoubleMLPQ, DoubleMLQTE, DoubleMLCVAR) all live in this same package and could call a _for_test variant; the public API is kept for clarity.

    solve_spd

    source
    fn solve_spd(a : Matrix, b : Array[Double]) -> Array[Double]

    Solve A x = b for a symmetric positive-definite A using the Cholesky factorisation A = L L^T. The system is solved in two steps: forward substitution for L y = b, then back substitution for L^T x = y.

    Both substitution accumulators are Kahan-compensated (matching cholesky and matvec). On the hot path for every LinearRegression::fit, every IRLS iteration, and every sandwich-SE back-solve.

    stratified_kfold

    source
    fn stratified_kfold(stratum : Array[Int], n_folds : Int, seed : Int) -> Array[Fold]

    Stratified K-fold partition: per stratum, split the indices into n_folds (train, test) pairs using the same Fisher-Yates + fold allocation as kfold.mbt::kfold, then concatenate the per-stratum folds in stratum order to form a single n_folds-long list of Fold objects that together partition [0, n_obs).

    stratum is a length-n_obs integer array with one stratum label per observation. The function collects indices per stratum, runs kfold on each stratum with seed + stratum_value as the per- stratum seed (so different strata produce different shuffles), then zips the per-stratum folds together: fold f of the final list contains the union of fold-f test rows across all strata.

    Matches the upstream doubleml.utils.resampling.DoubleMLResampling with stratify=stratum and RepeatedKFold collapsed to a single repetition. For multi-repetition partitioning, wrap calls in a loop over the seed.

    tsbh_p_adjust

    source
    fn tsbh_p_adjust(unadjusted : Array[Double]) -> Array[Double]

    tsby_p_adjust

    source
    fn tsby_p_adjust(unadjusted : Array[Double]) -> Array[Double]

    v0.24.0+ two-stage Benjamini-Yekutieli FDR correction. Combines the two-stage m0_hat adjustment (from tsbh_p_adjust) with the harmonic-sum c factor (from by_fdr_p_adjust) to handle arbitrary dependence between tests.

    Algorithm (matches statsmodels.stats.multitest.multipletests with method='fdr_tsbky'):
    1. Compute c = sum_{i=1}^{n} 1/i.
    2. Apply BY: p_by_sorted[k] = min(1, p_sorted[k] * m * c / (k+1)).
    3. Estimate m0_hat = #{unadjusted > alpha} / (1 -alpha) (clamped to [1, m]).
    4. Apply correction: p_adj_sorted[k] = min(1, p_by_sorted[k] * m0_hat /m).
    5. Enforce monotonicity + reorder.

    var_est

    source
    fn var_est(psi_a : Array[Double], psi_b : Array[Double]) -> (Double, Double)

    Compute the DML point estimate theta_hat and standard error se from the standard influence-function variance formula used by every DML estimator in this package (PLR, IRM, PLIV, IIVM, DID, SSM):

     theta_hat = -mean(psi_b) / mean(psi_a)
 J         = mean(psi_a)
 gamma     = mean(psi(theta_hat)^2)   where psi(theta) = theta * psi_a + psi_b
 sigma2    = gamma / (J^2 * n)
 se        = sqrt(sigma2)

    The inputs psi_a and psi_b must share the same non-zero length n; the call aborts via require(...) otherwise (using the same check.mbt / require style with auto-injected SourceLoc that every other precondition in this package uses).

    Floating-point order is preserved exactly: this is the byte-equal extraction of the inline block that previously lived in every estimator's fit(). The estimator regression-protection invariant for n_rep == 1 therefore still holds: the per-rep (theta, se) coming out of this function is identical to the per-rep value the old "compute mean, then variance" code produced, and so the aggregator's n_rep=1 fast path returns the same (theta, se).

    var_est_cluster

    source
    fn var_est_cluster(psi : Array[Double], psi_deriv : Array[Double], unit_rows : Array[Array[Int]], unit_fold : Array[Int], n_folds : Int, n_folds_per_cluster : Int) -> Double raise VarEstClusterError

    Cluster-robust standard error (_var_est in doubleml/utils/_estimation.py, one-cluster-variable branch). Per test fold, accumulate squared unit-level score sums scaled by 1 / |I_k|; divide both the gamma accumulator and the Jacobian analog by n_folds_per_cluster:

     gamma += S_g^2 / |I_k|,   S_g = sum of psi over unit g,
 J     += (sum of psi_deriv over the fold's rows) / |I_k|,
 sigma2 = (gamma / npc) / (N_units * (J / npc)^2).

    psi is the linear score evaluated at theta_hat, psi_deriv the score derivative (= psi_a). unit_fold[u] is the test-fold index of unit u; every row of a unit lies in that fold because the folds partition whole units.

    Returns Double!VarEstClusterError: the v0.34.0 J-floor defensive guard (|J|<1e-6) raises VarEstClusterError::JTooSmall carrying the exact (j, g, n_units) triple that triggered it. Callers that want the pre-v0.35.0 process-death behavior should catch the error and re-abort (this is what every DoubleMLXXX::fit_cluster does); callers that want to retry or surface the error should propagate via ?. The error type is declared in kfold.mbt so the cluster helper stack can share it.

    var_est_with_jacobian

    source
    fn var_est_with_jacobian(psi : Array[Double], jacobian : Double) -> Double

    Variant of var_est for Z-estimators where the point estimate theta_hat is already known (e.g. the LPQ bisection root) and jacobian = d mean(psi(theta)) / d theta is supplied externally (e.g. via KDE in lpq.mbt). Returns se = sqrt(gamma / (jacobian^2 * n)) where gamma = mean(psi(theta_hat)^2).

    REVIEW L11 fix (0.7.0): used by DoubleMLLPQ::fit to consolidate the variance computation. The previous inline form was byte-equal to this helper; the LPQ SE test tolerance was widened from 30% to 50% to absorb the smoothed KDE derivative; using the shared helper now makes that bit-equal relationship explicit.

    Aborts if psi is empty or jacobian is zero (which would produce an infinite SE).

    variance

    source
    fn variance(a : Array[Double]) -> Double

    Sample variance (population formula sum((x - mean)^2) / n).

    The sum((a[i] - m) * (a[i] - m)) accumulator is Kahan-compensated to match the sibling mean. Without compensation, (a[i] - m) vanishes near 0 and the accumulator loses low-order bits; Kahan keeps the running sum accurate to O(eps) independent of n.

    Source Files

    Powered by MoonBit

    Site sourceReport issuePackagesBuild queueSkillsStatistics

    © 2026 mooncakes.io