///|
pub struct DesignMatrix {
values : Array[Array[Double]]
column_names : Array[String]
}
///|
pub fn make_design_matrix(
covariates : Array[Array[Double]],
include_intercept : Bool,
) -> DesignMatrix {
let width = if covariates.length() == 0 { 0 } else { covariates[0].length() }
let intercept_columns = if include_intercept { 1 } else { 0 }
let names : Array[String] = Array::new(capacity=width + intercept_columns)
if include_intercept {
names.push("intercept")
}
for j in 0.. Array[Array[Double]] {
let result : Array[Array[Double]] = Array::new(capacity=covariates.length())
for i in 0.. Array[Array[Double]] {
let result : Array[Array[Double]] = Array::new(capacity=matrix.length())
let n = if matrix.length() < column.length() {
matrix.length()
} else {
column.length()
}
for i in 0.. Array[Double] {
let result = Array::new(capacity=design.length())
for row in design {
result.push(dot(row, coefficients))
}
result
}
///|
pub fn marginal_effects(
model : ModelFit,
covariates : Array[Array[Double]],
) -> Array[Double] {
let result = Array::new(capacity=model.coefficients.length())
for coefficient in model.coefficients {
result.push(coefficient)
}
let _ = covariates
result
}
///|
pub fn interaction_average_effect(
model : ModelFit,
covariates : Array[Array[Double]],
treatment_column : Int,
) -> Double {
if covariates.length() == 0 {
return 0.0
}
let treated = covariates.copy()
let control = covariates.copy()
for row in treated {
if treatment_column >= 0 && treatment_column < row.length() {
row[treatment_column] = 1.0
}
}
for row in control {
if treatment_column >= 0 && treatment_column < row.length() {
row[treatment_column] = 0.0
}
}
mean(
subtract_vectors(
predict_outcomes(model, treated),
predict_outcomes(model, control),
),
)
}
///|
pub fn polynomial_design(
covariates : Array[Array[Double]],
degree : Int,
) -> DesignMatrix {
let values = polynomial_features(covariates, degree)
let names : Array[String] = Array::new()
if values.length() > 0 {
for j in 0..