///|
/// 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.
pub(all) struct PSProcessorConfig {
clipping_threshold : Double
extreme_threshold : Double
calibration_method : String
cv_calibration : Bool
} derive(Debug)
///|
pub extend PSProcessorConfig with @moonbitlang/core/debug.Debug::{to_repr}
///|
/// 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).
pub fn PSProcessorConfig::new(
clipping_threshold? : Double = 1.0e-2,
extreme_threshold? : Double = 1.0e-12,
calibration_method? : String = "none",
cv_calibration? : Bool = false,
) -> PSProcessorConfig raise PSConfigError {
// v0.48.0: per-call wrap on each `require` because the body also
// raises PSConfigError::InconsistentCVCalibration; a block-level
// try/catch would partial_match the catch.
ignore(
require(clipping_threshold > 0.0) catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
},
)
ignore(
require(clipping_threshold < 0.5) catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
},
)
ignore(
require(extreme_threshold > 0.0) catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
},
)
ignore(
require(extreme_threshold < 0.5) catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
},
)
ignore(
require(calibration_method == "none" || calibration_method == "isotonic") catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
},
)
if cv_calibration && calibration_method == "none" {
raise PSConfigError::InconsistentCVCalibration
}
{ clipping_threshold, extreme_threshold, calibration_method, cv_calibration, }
}
///|
/// 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`).
pub fn PSProcessorConfig::default() -> PSProcessorConfig {
// The default args are valid (no cv_calibration,
// calibration_method="none"), so the cv_calibration abort
// is unreachable. Construct directly without going
// through the `new()` raise path.
{
clipping_threshold: 1.0e-2,
extreme_threshold: 1.0e-12,
calibration_method: "none",
cv_calibration: false,
}
}
///|
/// 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).
pub struct PSProcessor {
config : PSProcessorConfig
} derive(Debug)
///|
pub extend PSProcessor with @moonbitlang/core/debug.Debug::{to_repr}
///|
pub fn PSProcessor::new(
config? : PSProcessorConfig = PSProcessorConfig::default(),
) -> PSProcessor {
{ config, }
}
///|
/// Convenience constructor mirroring upstream
/// `PSProcessor.from_config`.
pub fn PSProcessor::from_config(config : PSProcessorConfig) -> PSProcessor {
{ config, }
}
///|
pub fn PSProcessor::clipping_threshold(self : PSProcessor) -> Double {
self.config.clipping_threshold
}
///|
pub fn PSProcessor::extreme_threshold(self : PSProcessor) -> Double {
self.config.extreme_threshold
}
///|
pub fn PSProcessor::calibration_method(self : PSProcessor) -> String {
self.config.calibration_method
}
///|
pub fn PSProcessor::cv_calibration(self : PSProcessor) -> Bool {
self.config.cv_calibration
}
///|
/// 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).
pub fn PSProcessor::adjust_ps(
self : PSProcessor,
ps : Array[Double],
treatment : Array[Double],
cv? : Array[(Array[Int], Array[Int])]? = None,
) -> Array[Double] {
try {
validate_treatment(treatment)
require(ps.length() == treatment.length())
let n = ps.length()
require(n > 0)
// Calibration step. v0.37.0: extracted into apply_calibration
// helper that raises InvalidCalibrationError on unknown
// calibration method; catch and re-abort to preserve pre-v0.37.0
// process-death behavior. v0.38.0: apply_calibration now also
// propagates CalibrationFittingError from isotonic_calibrate_cv;
// we re-abort to preserve the pre-v0.38.0 behavior on a malformed
// cv partition.
let calibrated = apply_calibration(self.config, ps, treatment, cv) catch {
InvalidCalibrationError::UnknownMethod(m) =>
abort(
"unknown calibration_method (set in PSProcessorConfig::new): " + m,
)
CalibrationFittingError::IncompleteCVPartition =>
abort("isotonic_calibrate_cv: cv partition does not cover all indices")
_ => abort("apply_calibration: unknown error")
}
// Clip to [eps, 1 - eps].
let lo = self.config.clipping_threshold
let hi = 1.0 - self.config.clipping_threshold
let out : Array[Double] = Array::make(calibrated.length(), 0.0)
for i = 0; i < calibrated.length(); i = i + 1 {
let v = calibrated[i]
if v < lo {
out[i] = lo
} else if v > hi {
out[i] = hi
} else {
out[i] = v
}
}
out
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// 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.
pub fn apply_calibration(
config : PSProcessorConfig,
ps : Array[Double],
treatment : Array[Double],
cv : Array[(Array[Int], Array[Int])]?,
) -> Array[Double] raise Error {
match config.calibration_method {
"none" => ps
"isotonic" =>
if config.cv_calibration {
isotonic_calibrate_cv(ps, treatment, cv)
} else {
let (sorted_x, y_hat) = fit_isotonic(ps, treatment)
predict_isotonic_from_fit(ps, sorted_x, y_hat)
}
_ => raise InvalidCalibrationError::UnknownMethod(config.calibration_method)
}
}
// ---------------------------------------------------------------------------
// Internal: propensity-score + treatment validation
// ---------------------------------------------------------------------------
///|
fn validate_treatment(treatment : Array[Double]) -> Unit {
try {
for t in treatment {
require(t == 0.0 || t == 1.0)
}
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
// ---------------------------------------------------------------------------
// Internal: isotonic regression via PAVA
// ---------------------------------------------------------------------------
///|
/// 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.
pub fn pava(y : Array[Double], weights? : Array[Double] = []) -> Array[Double] {
try {
let n = y.length()
if n == 0 {
return []
}
let w : Array[Double] = if weights.length() == 0 {
Array::make(n, 1.0)
} else {
require(weights.length() == n)
weights
}
// Parallel arrays as a poor man's stack of (sum_y, sum_w, size)
// blocks. MoonBit has no native tuple-array sort, but the
// per-field arrays stay synchronised because every push/pop
// updates all three in lockstep.
let block_sum : Array[Double] = []
let block_w : Array[Double] = []
let block_size : Array[Int] = []
for i = 0; i < n; i = i + 1 {
// v0.104.0: this seeded the block's weighted sum with the
// RAW `y[i]`, so a block's mean came out as
// `(sum of y) / (sum of w)` instead of
// `(sum of w * y) / (sum of w)`. With unit weights the
// two coincide, which is why the unweighted cases still
// passed after the fix; with any non-uniform weight the
// fitted block value is wrong. Caught by
// `validate_pava_with_python.py` against
// `sklearn.isotonic.isotonic_regression`: for
// y = [0.2, 0.5, 0.1, 0.8], w = [1, 1, 0.25, 1] the pooled
// block holds 0.5 and 0.1, whose correct weighted mean is
// (1*0.5 + 0.25*0.1) / 1.25 = 0.42; this computed
// (0.5 + 0.1) / 1.25 = 0.48.
let mut cur_sum = w[i] * y[i]
let mut cur_w = w[i]
let mut cur_size = 1
// Pool with previous blocks as long as the last block's
// (weighted) mean is greater than the new block's mean.
while block_size.length() > 0 {
let last_idx = block_size.length() - 1
let last_mean = block_sum[last_idx] / block_w[last_idx]
let new_mean = cur_sum / cur_w
if last_mean <= new_mean {
break
}
// Pool: pop the last block, fold its content into the new
// block, and continue checking.
cur_size = cur_size + block_size[last_idx]
cur_sum = cur_sum + block_sum[last_idx]
cur_w = cur_w + block_w[last_idx]
block_size.truncate(last_idx)
block_sum.truncate(last_idx)
block_w.truncate(last_idx)
}
block_size.push(cur_size)
block_sum.push(cur_sum)
block_w.push(cur_w)
}
// Flatten the block stack back to a per-element output. Each
// block contributes `block_size[b]` copies of its (weighted)
// mean `block_sum[b] / block_w[b]`.
let out : Array[Double] = Array::make(n, 0.0)
let mut k = 0
for b = 0; b < block_size.length(); b = b + 1 {
let mean = block_sum[b] / block_w[b]
let size = block_size[b]
let mut _i = 0
while _i < size {
out[k] = mean
k = k + 1
_i = _i + 1
}
}
out
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// 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).
pub fn fit_isotonic(
x : Array[Double],
y : Array[Double],
) -> (Array[Double], Array[Double]) {
try {
let n = x.length()
require(y.length() == n)
require(n > 0)
// Sort indices by x.
let order : Array[Int] = Array::makei(n, fn(i) { i })
order.sort_by(fn(a, b) { x[a].compare(x[b]) })
let sorted_x : Array[Double] = Array::make(n, 0.0)
let sorted_y : Array[Double] = Array::make(n, 0.0)
for k = 0; k < n; k = k + 1 {
let i = order[k]
sorted_x[k] = x[i]
sorted_y[k] = y[i]
}
let sorted_y_hat = pava(sorted_y)
(sorted_x, sorted_y_hat)
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// 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.
pub fn predict_isotonic(
fitted_x : Array[Double],
fitted_y_hat : Array[Double],
x_new : Array[Double],
) -> Array[Double] {
try {
let n_new = x_new.length()
let n_fit = fitted_x.length()
require(n_fit > 0)
require(fitted_y_hat.length() == n_fit)
let lo = fitted_y_hat[0]
let hi = fitted_y_hat[n_fit - 1]
let out : Array[Double] = Array::make(n_new, 0.0)
// Pre-compute the right-edge index of the block containing each
// fitted `x`. Since `fitted_x` is sorted in non-decreasing
// order, the right edge of a block is the largest `j` such
// that `fitted_y_hat[j]` is constant. Concretely: walk `j` from
// 0 to `n_fit - 1`; whenever `fitted_y_hat[j]` changes, the
// block ends at `j - 1`. The block containing `j` ends at the
// first `k > j` with `fitted_y_hat[k] != fitted_y_hat[j] - 1`.
// For prediction, the easiest lookup is: given `x_new[i]`,
// find the largest `j` with `fitted_x[j] <= x_new[i]`, then
// return `fitted_y_hat[j]`. We do that with a single linear
// scan that walks `fitted_x` and `x_new` in lockstep.
for i = 0; i < n_new; i = i + 1 {
let x = x_new[i]
if x < fitted_x[0] {
out[i] = lo
} else if x > fitted_x[n_fit - 1] {
out[i] = hi
} else {
// Find the largest j with fitted_x[j] <= x.
let mut j = 0
while j < n_fit - 1 && fitted_x[j + 1] <= x {
j = j + 1
}
let v = fitted_y_hat[j]
// Defensive clip to [0, 1] (PAVA on {0, 1} stays in [0, 1]
// by construction, but the clip makes the contract
// explicit).
if v < 0.0 {
out[i] = 0.0
} else if v > 1.0 {
out[i] = 1.0
} else {
out[i] = v
}
}
}
out
} catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
}
}
///|
/// Predict at the original `x` values. Convenience for the
/// non-CV path where the in-sample prediction IS the calibrated
/// propensity. Since `fitted_y_hat` is the PAVA output on
/// `(sorted_x, sorted_y)`, predicting at the original (unsorted)
/// `x` requires a step-function lookup: for each `x[i]`, find
/// the largest `j` with `sorted_x[j] <= x[i]` and return
/// `fitted_y_hat[j]`.
fn predict_isotonic_from_fit(
x : Array[Double],
sorted_x : Array[Double],
fitted_y_hat : Array[Double],
) -> Array[Double] {
predict_isotonic(sorted_x, fitted_y_hat, x)
}
///|
/// 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 `?`.
pub fn isotonic_calibrate_cv(
x : Array[Double],
y : Array[Double],
cv : Array[(Array[Int], Array[Int])]?,
) -> Array[Double] raise CalibrationFittingError {
let n = x.length()
let folds : Array[(Array[Int], Array[Int])] = match cv {
Some(f) => f
None => default_5fold(n)
}
// v0.48.0: per-call wrap because the body also raises
// CalibrationFittingError::IncompleteCVPartition; a block-level
// try/catch would partial_match the catch.
ignore(
require(folds.length() > 0) catch {
PreconditionError::Violated(loc) =>
abort("precondition failed at " + loc.to_string())
},
)
let out : Array[Double] = Array::make(n, 0.0)
let seen : Array[Bool] = Array::make(n, false)
for fold = 0; fold < folds.length(); fold = fold + 1 {
let (train_idx, test_idx) = folds[fold]
// Pull out the training data.
let x_train : Array[Double] = []
let y_train : Array[Double] = []
let mut x_train_acc = x_train
let mut y_train_acc = y_train
for k = 0; k < train_idx.length(); k = k + 1 {
let i = train_idx[k]
x_train_acc = x_train_acc + [x[i]]
y_train_acc = y_train_acc + [y[i]]
}
// Pull out the test x values.
let x_test : Array[Double] = []
let mut x_test_acc = x_test
for k = 0; k < test_idx.length(); k = k + 1 {
x_test_acc = x_test_acc + [x[test_idx[k]]]
}
// Fit on the training fold, predict on the test fold.
let (sorted_x, sorted_y_hat) = fit_isotonic(x_train_acc, y_train_acc)
let pred = predict_isotonic(sorted_x, sorted_y_hat, x_test_acc)
for k = 0; k < test_idx.length(); k = k + 1 {
let j = test_idx[k]
out[j] = pred[k]
seen[j] = true
}
}
// Each input index must be predicted by exactly one fold.
// Raise if any index is uncovered (a malformed cv partition).
for i = 0; i < n; i = i + 1 {
if !seen[i] {
raise CalibrationFittingError::IncompleteCVPartition
}
}
out
}
///|
/// Default 5-fold partition with `seed=3141`. Matches the
/// upstream `cross_val_predict(cv=5)` default.
fn default_5fold(n : Int) -> Array[(Array[Int], Array[Int])] {
let folds = kfold(n, 5, 3141)
let out : Array[(Array[Int], Array[Int])] = []
let mut out_acc = out
for f = 0; f < folds.length(); f = f + 1 {
out_acc = out_acc + [(folds[f].train_idx, folds[f].test_idx)]
}
out_acc
}
// ---------------------------------------------------------------------------
// Inverse-probability-weight normalisation (v0.119.0)
// ---------------------------------------------------------------------------
///|
/// v0.119.0: normalize inverse probability weights so the mean weight
/// within each treatment group is 1. Ports upstream
/// `doubleml.utils._propensity_score._normalize_ipw`:
///
/// mean_treat1 = mean(treatment / propensity)
/// mean_treat0 = mean((1 - treatment) / (1 - propensity))
/// normalized = treatment * propensity * mean_treat1
/// + (1 - treatment) * (1 - (1 - propensity) * mean_treat0)
///
/// Intuition: the raw weight `treatment / propensity` is unbiased for
/// the treated count but its scale carries the nuisance model's bias,
/// so it inflates the variance of the score. Multiplying by
/// `mean_treat1` removes that constant-of-proportionality factor.
///
/// # The 1e-12 clamp is a DELIBERATE deviation
///
/// Upstream divides by `propensity` with no guard at all. This clamps
/// both denominators at `1e-12`, so a propensity of exactly 0 yields a
/// weight of 0 rather than `inf` or `NaN`.
///
/// The two agree bit-for-bit on any propensity bounded away from zero,
/// which is what `_verify/gen_v119_oracle.py` checks -- the goldens are
/// generated by calling upstream's own function on the fixture. The
/// clamp only matters on a degenerate propensity, where upstream
/// returns `inf` and propagates `NaN` through the whole estimator while
/// this returns a finite, bounded number. That is recorded as a
/// difference rather than hidden: v0.119.0's
/// `v119_normalized_ipw_matches_upstream` pins the agreement, and
/// `v119_normalized_ipw_survives_a_degenerate_propensity` pins the
/// clamp.
///
/// # Why this moved here
///
/// `cvar.mbt` has carried a private copy of exactly this function since
/// its `normalize_ipw` option landed. Two private copies of an
/// upstream-tracked formula is how the two drift apart, so this is the
/// single definition and `cvar.mbt` now calls it. It was promoted from
/// private to `pub` for the same reason: an oracle cross-check cannot
/// import a file-private function.
///
/// `treatment` must be the 0/1 indicator, matching upstream's
/// `treatment_indicator`; `validate_treatment` is the gate.
pub fn normalize_ipw(
propensity : Array[Double],
treatment : Array[Double],
) -> Array[Double] {
let n = propensity.length()
let m_clip = clip_vec(propensity, 1.0e-12, 1.0)
// v0.119.0: the CONTROL denominator is `1 - p`, not `p`. This was
// written as a second `clip_vec(m_clip, ...)` -- which clips the
// propensity to itself and leaves the control weights un-normalised.
// The upstream oracle caught it: on the fixture's row 18 (p ~ 0.09) the
// port returned +0.9019 where upstream returns -0.0136, a sign flip
// that only appears where `1 - p` and `p` disagree.
//
// That it survived the whole v0.119.0 gate suite is the sharper
// point: `cvar.mbt` now DELEGATES here, so the bug was live in CVAR's
// `normalize_ipw` path, and no pre-existing test noticed. One
// external oracle bought more than the entire internal suite.
let om_clip : Array[Double] = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
let om = 1.0 - m_clip[i]
om_clip[i] = if om < 1.0e-12 { 1.0e-12 } else { om }
}
let mut sum_t1 = 0.0
let mut sum_t0 = 0.0
for i = 0; i < n; i = i + 1 {
sum_t1 = sum_t1 + treatment[i] / m_clip[i]
sum_t0 = sum_t0 + (1.0 - treatment[i]) / om_clip[i]
}
let mean_t1 = sum_t1 / n.to_double()
let mean_t0 = sum_t0 / n.to_double()
let out = Array::make(n, 0.0)
for i = 0; i < n; i = i + 1 {
let w1 = treatment[i] * m_clip[i] * mean_t1
let w0 = (1.0 - treatment[i]) * (1.0 - om_clip[i] * mean_t0)
out[i] = w1 + w0
}
out
}
///|
/// v0.119.0: upstream `doubleml.utils._propensity_score
/// ._propensity_score_adjustment` -- the switch that decides whether the
/// propensity enters the score raw or IPW-normalized.
///
/// `normalize_ipw = false` is the IDENTITY, verified against upstream:
/// `gen_v119_oracle.py` reports `adjust_identity_max_diff = 0.0`. That
/// is the whole point of the flag, and it is why the default on
/// `DoubleMLAPO` / `DoubleMLAPOS` is `false` (upstream's
/// `apo.py:97` default) -- turning normalization on must move the
/// answer, and leaving it off must not.
pub fn propensity_score_adjustment(
propensity : Array[Double],
treatment : Array[Double],
do_normalize? : Bool = false,
) -> Array[Double] {
// v0.119.0: the parameter is deliberately NOT named `normalize_ipw`.
// That would shadow the function of the same name, turning the call
// below into an invocation of a Bool. It surfaced as a bare type
// mismatch rather than anything a reader could diagnose.
if do_normalize {
normalize_ipw(propensity, treatment)
} else {
propensity
}
}
///|
/// v0.119.0: upstream `doubleml.utils._propensity_score._trimm`.
///
/// `rule = "truncate"` clips to `[threshold, 1 - threshold]`; anything
/// else (upstream's `None`) is the identity. Kept because
/// `trimming_rule` / `trimming_threshold` are deprecated-but-still-live
/// upstream parameters (see the `DeprecationWarning` at
/// `propensity_score_processing.py:59`) and `DoubleMLDIDCSBinary`
/// still exposes the collapsed form of them.
pub fn trim_predictions(
preds : Array[Double],
rule : String,
threshold : Double,
) -> Array[Double] {
if rule == "truncate" {
clip_vec(preds, threshold, 1.0 - threshold)
} else {
preds
}
}