# Find project root and include ensure_packages.jl
project_root = let
current = pwd()
while !isfile(joinpath(current, "Project.toml")) && !isfile(joinpath(current, "_quarto.yml"))
parent = dirname(current)
parent == current && break
current = parent
end
current
end
include(joinpath(project_root, "scripts", "ensure_packages.jl"))
@auto_using Random Distributions CairoMakie
Random.seed!(42)
# Example: Time-varying treatment and confounder
# L_t: Disease severity (confounder)
# A_t: Treatment (depends on L_t)
# Y_t: Outcome (depends on L_t and A_t)
# Step 1: Fit models from observed data (simulated here for demonstration)
# In practice, these would be fitted from real data
# Confounder model: P(L_t | A_{t-1}, L_{t-1})
function fit_confounder_model(data)
"""
Fit model for confounder evolution.
NOTE: This is a simplified example. In practice, use flexible models
(GLM with splines, random forests, Super Learner) to avoid misspecification bias.
"""
# Simplified: L_t = α + β*A_{t-1} + γ*L_{t-1} + noise
# In production, use flexible models (GLM with splines, machine learning)
return (α=0.1, β=-0.2, γ=0.8, σ=0.1)
end
# Outcome model: E[Y_t | A_t, L_t, Y_{t-1}]
function fit_outcome_model(data)
"""
Fit model for outcome expectation.
NOTE: This is a simplified example. In practice, use flexible models
to avoid misspecification bias in G-computation.
"""
# Simplified: Y_t = α + β*A_t + γ*L_t + δ*Y_{t-1} + noise
return (α=0.0, β=0.5, γ=-0.3, δ=0.7, σ=0.1)
end
# Simulate observed data (for demonstration)
n_subjects = 100
T = 5
data = []
for i in 1:n_subjects
L = [rand(Normal(0.5, 0.2))] # Initial severity
A = Int[]
Y = [rand(Normal(0.0, 0.1))] # Initial outcome
for t in 2:T
# Treatment assignment (depends on L_{t-1})
p_treat = 1 / (1 + exp(-2*(L[end] - 0.5))) # Higher severity → more likely to treat
push!(A, rand(Bernoulli(p_treat)))
# Confounder evolution (depends on A_{t-1} and L_{t-1})
L_new = 0.1 - 0.2*A[end] + 0.8*L[end] + rand(Normal(0, 0.1))
push!(L, max(0, min(1, L_new))) # Bound to [0,1]
# Outcome (depends on A_t and L_t)
Y_new = 0.5*A[end] - 0.3*L[end] + 0.7*Y[end] + rand(Normal(0, 0.1))
push!(Y, Y_new)
end
push!(data, (L=L, A=A, Y=Y))
end
# Fit models
confounder_model = fit_confounder_model(data)
outcome_model = fit_outcome_model(data)
# Step 2: G-computation - simulate under treatment strategy
# Strategy: Always treat (A_t = 1 for all t)
function g_compute(strategy, confounder_model, outcome_model, n_sim::Int=1000)
"""
Perform G-computation: simulate outcomes under a treatment strategy.
Args:
strategy: Function (t, L, Y) -> A_t (treatment at time t)
confounder_model: Fitted model for confounder evolution
outcome_model: Fitted model for outcome expectation
n_sim: Number of Monte Carlo simulations
Returns:
Vector of final outcomes under the strategy
"""
# Pre-allocate for performance
outcomes = Vector{Float64}(undef, n_sim)
for sim in 1:n_sim
L = [rand(Normal(0.5, 0.2))] # Initial severity
Y = [rand(Normal(0.0, 0.1))] # Initial outcome
for t in 2:T
# Treatment according to strategy
# Strategy function takes (t, L, Y) where L and Y are current state vectors
A_t = strategy(t, L, Y)
# Simulate confounder under strategy
L_new = confounder_model.α + confounder_model.β*A_t +
confounder_model.γ*L[end] + rand(Normal(0, confounder_model.σ))
push!(L, max(0, min(1, L_new)))
# Simulate outcome under strategy
Y_new = outcome_model.α + outcome_model.β*A_t +
outcome_model.γ*L[end] + outcome_model.δ*Y[end] +
rand(Normal(0, outcome_model.σ))
push!(Y, Y_new)
end
outcomes[sim] = Y[end]
end
return outcomes
end
# Strategy 1: Always treat
strategy_always_treat = (t, L, Y) -> 1.0
# Strategy 2: Never treat
strategy_never_treat = (t, L, Y) -> 0.0
# Compute expected outcomes
outcomes_treat = g_compute(strategy_always_treat, confounder_model, outcome_model)
outcomes_no_treat = g_compute(strategy_never_treat, confounder_model, outcome_model)
# Visualise
let
fig = Figure(size = (800, 400))
ax = Axis(fig[1, 1], title = "G-Computation: Outcome Distributions Under Strategies",
xlabel = "Final Outcome", ylabel = "Density")
hist!(ax, outcomes_treat, bins=30, normalization=:pdf,
label="Always treat", color=(:blue, 0.5))
hist!(ax, outcomes_no_treat, bins=30, normalization=:pdf,
label="Never treat", color=(:red, 0.5))
axislegend(ax)
fig # Only this gets displayed
end
println("G-computation results:")
println(" E[Y | always treat] = ", round(mean(outcomes_treat), digits=3))
println(" E[Y | never treat] = ", round(mean(outcomes_no_treat), digits=3))
println(" Treatment effect = ", round(mean(outcomes_treat) - mean(outcomes_no_treat), digits=3))