version 11.0
capture mata mata drop m_ipt_ate()
capture mata mata drop m_ipt_logit()
capture mata mata drop m_ipt_logit_crit()
capture mata mata drop m_ipt_logit_phi()
mata:

void m_ipt_ate(string scalar varlist, string scalar touse, string scalar clusterid, string scalar weight, string scalar optroutine)

	{

	real matrix YDh_X, h_X, t_X, D1_ipt_logit_results, D0_ipt_logit_results
	real matrix m_delta1, m_delta0, m_hat, M11_hat, M22_hat, M31_hat, M32_hat, M33_hat, M_hat, OMEGA_hat, V_hat
	real colvector Y, D, wgts, delta1_ipt, delta0_ipt, theta_ipt, v1_ipt, v0_ipt, pi1_ipt, pi0_ipt, m_gamma, selvec
	real rowvector m_g
	real scalar N, M, gamma_ipt, NG
	string colvector t_X_names
	string matrix pscore_D1_names, pscore_D0_names, coef_names

	YDh_X = Y = D = h_X = .
	st_view(YDh_X, ., tokens(varlist), touse)
	st_subview(Y, YDh_X, ., 1)
	st_subview(D, YDh_X, ., 2)
	st_subview(h_X, YDh_X, ., (3\.))

	N = rows(h_X)
	M = cols(h_X)

	if (weight != "") {
			st_view(wgts, ., tokens(weight), touse)
			wgts = wgts :/ (colsum(wgts) :/ N)
	}
	else {
			wgts = J(N,1,1)
	}

	t_X = (h_X,J(N,1,1))

	display("")
	display("")
	display("Computing treated subsample tilt")
	D1_ipt_logit_results 	= m_ipt_logit(D, h_X, wgts, optroutine)
	delta1_ipt 				= D1_ipt_logit_results[1,.]'
	v1_ipt 					= t_X*delta1_ipt
	pi1_ipt 				= (N^-1) :* (D :/ invlogit(v1_ipt))
	m_delta1 				= D1_ipt_logit_results[1+2*(1+M)+1..1+2*(1+M)+N,1..1+M]

	display("")
	display("")
	display("Computing control subsample tilt")
	D0_ipt_logit_results 	= m_ipt_logit((1 :- D), h_X, wgts, optroutine)
	delta0_ipt 				= -D0_ipt_logit_results[1,.]'
	v0_ipt 					= t_X*delta0_ipt
	pi0_ipt 				= (N^-1) :* ((1 :- D) :/ (1 :- invlogit(v0_ipt)))
	m_delta0 				= D0_ipt_logit_results[1+2*(1+M)+1..1+2*(1+M)+N,1..1+M]

	gamma_ipt 				= sum(pi1_ipt :* Y - pi0_ipt :* Y)
	m_gamma 				= wgts :* ((N :* pi1_ipt :* Y) - (N :* pi0_ipt :* (Y :+ gamma_ipt)))

	m_hat 					= (m_delta1, m_delta0, m_gamma)

	M11_hat = D1_ipt_logit_results[(1+1+M+1)..(1+2*(1+M)),.]

	M22_hat = D0_ipt_logit_results[(1+1+M+1)..(1+2*(1+M)),.]

	M31_hat = colsum(wgts :* ((D        :* -exp(-v1_ipt)) :*  Y              ) :* t_X)
	M32_hat = colsum(wgts :* (((1 :- D) :*  exp(v0_ipt))  :* (Y :+ gamma_ipt)) :* t_X)
	M33_hat = -N

	M_hat   = (M11_hat, J(1+M,1+M,0), J(1+M,1,0) \ J(1+M,1+M,0), M22_hat, J(1+M,1,0) \ M31_hat,M32_hat,M33_hat)

	OMEGA_hat = J(2*(1+M)+1,2*(1+M)+1,0)

	if (clusterid != "") {
		st_view(G, ., tokens(clusterid), touse)
		NG = max(G)
		st_numscalar("nclust",NG)
		for (g=1; g<=NG; g++) {
			selvec = (G:==g)
			m_g    = colsum(m_hat :* selvec)
			OMEGA_hat = OMEGA_hat + m_g'*m_g
		}
		V_hat		= luinv(M_hat)*OMEGA_hat*luinv(M_hat)'
	}
	else {
		OMEGA_hat 	= m_hat'*m_hat
		V_hat		= luinv(M_hat)*OMEGA_hat*luinv(M_hat)'
	}

	t_X_names = tokens(varlist)
	t_X_names = t_X_names[.,3..1+1+M]
	t_X_names = (t_X_names, "_cons")'

	pscore_D1_names 	= (("delta1" :* J(1+M,1,1)), t_X_names)
	pscore_D0_names 	= (("delta0" :* J(1+M,1,1)), t_X_names)
	coef_names 			= (pscore_D1_names \ pscore_D0_names \ ("ate", "gamma"))

	st_matrix("V_hat",V_hat)
	st_matrixrowstripe("V_hat", coef_names)
	st_matrixcolstripe("V_hat", coef_names)

	theta_hat = (delta1_ipt' ,delta0_ipt', gamma_ipt)
	st_matrix("theta_hat",theta_hat)
	st_matrixcolstripe("theta_hat", coef_names)

	st_numscalar("nobs",N)

	}

real matrix m_ipt_logit(real colvector D, real matrix h_X, real colvector wgts, string scalar optroutine)

	{

	real matrix vcov_ipt, hess_ipt, m2_ipt, A
	real rowvector delta_ipt

	t_Xt_X = cross(h_X,1,h_X,1)
	if (rank(t_Xt_X) < cols(h_X) + 1) {
		errprintf("singular or near singular matrix")
		exit(499)
	}
	t_XD = cross(h_X,1,D,0)
	delta_sv = cholsolve(t_Xt_X,t_XD)'

	M = moptimize_init()
	moptimize_init_evaluator(M, &m_ipt_logit_crit())
	moptimize_init_which(M, "max")

	if (optroutine=="e2") {
		moptimize_init_evaluatortype(M, "e2")
	}
	else {
		moptimize_init_evaluatortype(M, "e1")
		moptimize_init_technique(M, "bfgs")
	}

	moptimize_init_depvar(M, 1, D)
	moptimize_init_eq_indepvars(M, 1, h_X)
	moptimize_init_eq_coefs(M, 1, delta_sv)
	moptimize_init_vcetype(M,"robust")
	moptimize_init_userinfo(M,1,wgts)

	moptimize(M)

	delta_ipt 	= moptimize_result_eq_coefs(M)
	vcov_ipt    = moptimize_result_V(M,1)
	hess_ipt	= moptimize_result_Hessian(M)
	m2_ipt      = (moptimize_result_scores(M) :* h_X, moptimize_result_scores(M))
	A 			= (delta_ipt \ vcov_ipt  \ hess_ipt \ m2_ipt)

	return(A)

	}

function m_ipt_logit_crit(transmorphic M, real scalar todo,
			  		      real rowvector delta, crit, foc, soc)

	{

	real scalar N
	real colvector t_Xdelta
	real colvector D, phi, phi1, phi2, soc_i
	real matrix phi012

	t_Xdelta 	= moptimize_util_xb(M, delta, 1)
	D 			= moptimize_util_depvar(M, 1)
	N           = length(D)
	phi012 		= m_ipt_logit_phi(t_Xdelta, N)
	phi         = phi012[.,1]

	wgts        = moptimize_util_userinfo(M,1)

	crit  		= moptimize_util_sum(M, wgts :* (D :* phi :- t_Xdelta))

	if (todo>=1) {
				    phi1    = phi012[.,2]
					foc 	= wgts :* (D :* phi1 :- 1)
					if (todo==2) {
									phi2    = phi012[.,3]
									soc_i 	= wgts :* (D :* phi2)
									soc     = moptimize_util_matsum(M, 1,1, soc_i, crit)
								 }
                }
	}

real matrix m_ipt_logit_phi(real colvector v, real scalar N)
	{

    real scalar a, b, c, v_star
	real colvector phi, phi1, phi2
	real matrix A

	a = -(N - 1)*(1 + log(1/(N - 1)) + 0.5*(log(1/(N - 1)))^2)
	b = N + (N - 1)*log(1/(N - 1))
	c = -(N-1)
	v_star = log(1/(N - 1))

	phi        =  (v:>v_star) :* (v - exp(-v))   + (v:<=v_star) :* (a :+ (b:*v) + (0.5*c):*(v:^2))
	phi1       =  (v:>v_star) :* (1 :+ exp(-v))  + (v:<=v_star) :* (b :+ (c:*v))
	phi2       =  (v:>v_star) :* (-exp(-v))      + (v:<=v_star) :* c
	A          =  phi,phi1,phi2

	return(A)

    }

end

program define iptATE, eclass
	version 11
    syntax varlist(numeric) [if] [in] [pweight] [, cluster(string) optroutine(string) BALance]
   	marksample touse

	tokenize `varlist'
	local outcomevar 	`1'
	local treatdum 		`2'
	macro shift
	macro shift
	local controlvars 	`*'

	if "`balance'" != "" {
		tempvar controldum
		quietly gen double `controldum' = 1 - `treatdum' if `touse'

		display ""
		display ""
		display "Mean covariate values by treatment status"
		display "----------------------------------------------------------------------------------------"
		display "-                Variable     T = 0          T = 1          Diff            p-value    -"
		display "----------------------------------------------------------------------------------------"
		foreach controlvar in `controlvars' {
			quietly {
					reg `controlvar' `controldum' `treatdum' [`pweight'] if `touse', nocons cluster(`cluster') robust
					lincom `controldum'
					local cm 	= r(estimate)
					local cm_se = r(se)
					lincom `treatdum'
					local tm 	= r(estimate)
					local tm_se = r(se)
					lincom `treatdum' - `controldum'
					local dm 	= r(estimate)
					local dm_se = r(se)
					test `treatdum' = `controldum'
					local dm_pv = r(p)
					}
		display as result %25s abbrev("`controlvar'",25) "     " %9.3f `cm'    "      " %9.3f `tm'    "      " %9.3f `dm'
		display as result %25s ""                        "    (" %9.3f `cm_se' ")    (" %9.3f `tm_se' ")    (" %9.3f `dm_se' ")   " %9.3f `dm_pv'
		display ""
		}
		display "----------------------------------------------------------------------------------------"
		display ""
	}

	if "`weight'" != "" {
		tempvar wgts
		quietly gen double `wgts' `exp' if `touse'
	}
	else {
		tempvar wgts
		quietly gen double `wgts' = 1 if `touse'
	}

	if "`cluster'" != "" {
	    tempvar clusterid
	   	egen `clusterid' = group(`cluster') if `touse'
		mata: m_ipt_ate("`varlist'", "`touse'", "`clusterid'","`wgts'","`optroutine'")
	}
	else {
	    mata: m_ipt_ate("`varlist'", "`touse'", "`cluster'","`wgts'","`optroutine'")
	}

	display ""
	display ""
	display "Inverse probability tilting propensity score & ATE estimates"
	display "Outcome variable    : `outcomevar'"
	display "Treatment indicator : `treatdum'"
	display "Control variables   : `controlvars'"
	ereturn post theta_hat V_hat [`weight'], esample(`touse')
	ereturn scalar N = nobs
	ereturn display

	if "`weight'" != "" {
		display "Probability weights option specified: `weight' `exp'"
	}

	if "`cluster'" != "" {
		display "Total number of primary sampling units: " nclust
	}
	display "Total number of observations: " nobs

    display ""
	display "Test of equality of two tilting coefficient vectors"
	quietly: test [delta1]_cons = [delta0]_cons
	test [delta1=delta0], common accum
end