optmed

R and Python implementations of MaxIE and MaxCor — closed-form solvers for learning the optimal composite mediator — from:

Zihuai He (2026). Learning the Optimal Composite Mediator: Closed-Form Solution and Inference. (under review). Preprint available at [link].

---

Overview

Given a treatment A, an outcome Y, and a high-dimensional panel of candidate mediators X (n × p), this package finds the linear composite M = Xw* that maximises mediation through the panel.

Two objectives are supported:

Solver	Objective	Use when
solve_maxie	Indirect effect h(w) = α̂(w) β̂(w)	Primary estimand; interpretable on the original scale
solve_maxcor	Mediation index f*(w) = cor(Xw, A) · cor(Zw, Y)	Scale-free; robust in the suppression regime

Both are solved in closed form via a primal–dual argument that reduces the non-convex optimisation to a one-dimensional bisection on a two-dimensional subspace. Cost is dominated by a single Cholesky factorisation — orders of magnitude faster than numerical optimisation.

A companion cosine test provides an asymptotic t-distribution test for the global null of no composite mediation, with no additional computational cost. The test automatically selects the primal (n > p, T ~ t(p−1)) or dual (p ≥ n, T ~ t(n−2)) formulation based on problem dimensions.

---

Files

File	Description
optmed.R	R implementation (base R only, no external dependencies)
optmed.py	Python implementation (requires NumPy, SciPy, pandas)

Both files implement identical algorithms and reproduce all tables and figures in the paper.

---

Quick start

R

source("optmed.R")

# Simulate data
dat <- generate_data(n = 500, p = 20, seed = 1, scenario = "partial")

# Global test: does any composite mediate?
test <- cosine_test(dat$X, dat$A, dat$Y)
cat("cos(phi) =", round(test$cos_phi, 3),
    "  T =", round(test$T, 3),
    "  p =", signif(test$p_value, 3), "\n")

# Learn the optimal composite mediator (MaxIE)
fit <- solve_maxie(dat$X, dat$A, dat$Y)
w   <- fit$w_consistent

# Mediation statistics for the learned composite
M    <- dat$X %*% w
stats <- mediation_stats(M, dat$A, dat$Y)
print(stats)

Python

from optmed import cosine_test, solve_maxie, mediation_stats, generate_data

# Simulate data
dat = generate_data(n=500, p=20, seed=1, scenario="partial")

# Global test
test = cosine_test(dat["X"], dat["A"], dat["Y"])
print(f"cos(phi) = {test['cos_phi']:.3f}  T = {test['T']:.3f}  p = {test['p_value']:.2e}")

# Learn the optimal composite mediator (MaxIE)
fit = solve_maxie(dat["X"], dat["A"], dat["Y"])
w   = fit["w_consistent"]

# Mediation statistics
M     = dat["X"] @ w
stats = mediation_stats(M, dat["A"], dat["Y"])
print(stats)

---

Main functions

cosine_test(X, A, Y)

Global test for the existence of any composite mediator. Returns cos_phi, T, p_value (two-sided), p_value_consistent, p_value_suppression, df, and method ("primal" or "dual").

solve_maxie(X, A, Y)

Closed-form maximiser of the indirect effect h(w) = α̂(w)β̂(w). Returns w_consistent (consistent mediation) and w_suppression (suppression mediation), plus T, df, and method.

solve_maxcor(X, A, Y)

Semi-closed-form maximiser of the mediation index f*(w) = cor(Xw, A) · cor(Zw, Y). Returns w_consistent and w_suppression.

mediation_stats(M, A, Y)

Conventional product-of-coefficients statistics for a composite score M. Returns α̂, β̂, indirect/direct/total effects, proportion mediated, mediation index f*, and Sobel test.

generate_data(n, p, seed, scenario)

Simulation utility used in the paper. Scenarios:

Code	Description	cos(φ)
"independent" (S1)	α ⊥ β (disjoint supports)	≈ 0
"quarter" (S2)	25% support overlap	≈ 0.23
"partial" (S3)	50% support overlap	≈ 0.47
"same" (S4)	α = β	1

---

Reproducing paper results

# R — reproduces all tables and figures
source("optmed.R")
run_all()

# Python — reproduces all tables and figures
python optmed.py

Output files (tables printed to console; figures saved as PDF):

Output	Function
Table 1: mean indirect effect h	table_obj_h()
Table 2: closed-form vs numerical	table_time()
Table 3: mean mediation index f*	table_obj_f()
Figure 1: QQ plots under H₀	figure_null_dist()
Figure 2: power (primal)	figure_power_props()
Figure 3: power (dual)	figure_power_dual()
Table 5: cosine test vs IUT	table_inference()

---

Dependencies

R: base R only (≥ 4.0). No external packages required.

Python: numpy, scipy, pandas. Install with:

pip install numpy scipy pandas

---

Citation

@article{he2026optmed,
  title  = {Learning the Optimal Composite Mediator:
            Closed-Form Solution and Inference},
  author = {He, Zihuai},
  year   = {2026},
  note   = {Under review}
}

---

License

MIT License. Copyright © 2026 Zihuai He, Stanford University.