///|
pub struct SubgroupEffect {
name : String
sample_size : Int
treated_count : Int
control_count : Int
estimate : Double
standard_error : Double
}
///|
fn subgroup_effect(
outcomes : Array[Double],
treatment : Array[Bool],
indices : Array[Int],
name : String,
) -> SubgroupEffect {
let subgroup_outcomes = Array::new(capacity=indices.length())
let subgroup_treatment = Array::new(capacity=indices.length())
for index in indices {
subgroup_outcomes.push(outcomes[index])
subgroup_treatment.push(treatment[index])
}
let treated = group_values(subgroup_outcomes, subgroup_treatment, true)
let control = group_values(subgroup_outcomes, subgroup_treatment, false)
let estimate = mean(treated) - mean(control)
let standard_error = if treated.length() == 0 || control.length() == 0 {
0.0
} else {
(variance(treated) / treated.length().to_double() +
variance(control) / control.length().to_double()).sqrt()
}
{
name,
sample_size: indices.length(),
treated_count: treated.length(),
control_count: control.length(),
estimate,
standard_error,
}
}
///|
pub fn threshold_subgroup_effects(
covariate : Array[Double],
outcomes : Array[Double],
treatment : Array[Bool],
threshold : Double,
) -> Array[SubgroupEffect] {
let below = Array::new()
let above = Array::new()
for i in 0.. Array[SubgroupEffect] {
let actual_groups = if groups < 1 { 1 } else { groups }
let result : Array[SubgroupEffect] = Array::new(capacity=actual_groups)
for group in 0..= lower && (covariate[i] < upper || is_last) {
indices.push(i)
}
}
result.push(
subgroup_effect(outcomes, treatment, indices, "quantile-{group + 1}"),
)
}
result
}
///|
pub fn conditional_average_treatment_effect(
predicted_treated : Array[Double],
predicted_control : Array[Double],
) -> Array[Double] {
let n = if predicted_treated.length() < predicted_control.length() {
predicted_treated.length()
} else {
predicted_control.length()
}
let result = Array::new(capacity=n)
for i in 0.. Double {
let n = if outcomes.length() < individual_effects.length() {
outcomes.length()
} else {
individual_effects.length()
}
if n == 0 {
return 0.0
}
let mut total = 0.0
for i in 0..= 0.0 && treatment[i] {
total += outcomes[i]
}
if individual_effects[i] < 0.0 && !treatment[i] {
total += outcomes[i]
}
}
total / n.to_double()
}
///|
pub fn uplift_curve(
outcomes : Array[Double],
treatment : Array[Bool],
scores : Array[Double],
bins : Int,
) -> Array[Double] {
let order = Array::new(capacity=scores.length())
for i in 0.. 0 && scores[order[j - 1]] < scores[key] {
order[j] = order[j - 1]
j -= 1
}
order[j] = key
}
let result = Array::new()
let actual_bins = if bins < 1 { 1 } else { bins }
for bin in 1..<=actual_bins {
let count = scores.length() * bin / actual_bins
let selected_y = Array::new(capacity=count)
let selected_t = Array::new(capacity=count)
for position in 0..