// A pooled regression that can accommodate thousands of firm-specific slope coefficients.
// The methodology is from Appendix A of Byzalov and Basu (2023),
// "The misuse of regression-based x-Scores as dependent variables".
// This code uses a third-party command rangestat. Install it using "ssc install rangestat"

program bigreg, eclass
	version 15
	syntax varlist [if], firmspecX(varlist) firmID(varname) [cluster(varname)]

	// The arguments are:
	//   varlist = Y pooledX
	//     where pooledX are the X variables with pooled slope coefficients.
	//   firmspecX(varlist) = X variables with firm-specific slope coefficients.
	//     Do not use factor notation (#, ##, etc.) or time-series operators (L., D., etc.) in varlist or firmspecX.
	//     Firm fixed effect (i.e., firm-specific intercepts) are included automatically (i.e., you should not
	//     include a constant term in firmspecX for them).
	//   firmID(varname) = index for firm-specific coefficients,
	// 	   where "firm" is shorthand for any relevant group of observations (e.g., household, industry, state).
	//   cluster(varname) = cluster variable for s.e. If it is omitted, then the s.e. are clustered by firmID.
	//     The cluster variable may differ from firmID (e.g., firmID = gvkey and cluster = industry).
	//     Only one-way clustering is supported (but you can easily replace reg in stage 2 with 2-way-clustered reg).

	qui{
		marksample touse
		markout `touse' `firmspecX' // remove observation from the estimation sample if any of `firmspecX' are missing
  		                            // (missing values for `varlist' are handled by marksample)

		tokenize `varlist'
		local Y "`1'"
		macro shift 1
		local pooledX "`*'"

		if "`cluster'"=="" local cluster="`firmID'"

		/***************** Stage 1 *************/
		// Run firmID-specific regressions of `Y' and each `pooledX' on `firmspecX' and compute residuals.
		// We use rangestat b/c it is MUCH faster than a by(firmID) loop for reg & predict

		tempname r
		local r_list = "" // the list of residuals for stage 2
		foreach var of varlist `Y' `pooledX'{
			rangestat (reg) `var' `firmspecX' if `touse', by(`firmID') interval(`Y' . .)  // firmID-specific regression
															// interval(`Y' . .) selects all available obs for the firm

			// manually compute residuals (rangestat does not have an equivalent of "predict, resid")
			gen `r'_`var'=`var' - b_cons if `touse'
			foreach x of varlist `firmspecX'{
				replace `r'_`var' = `r'_`var'- `x'*b_`x' if `touse'
			}
			drop reg_nobs-se_cons  // clean up the variables created by rangestat
			local r_list = "`r_list' `r'_`var'" // add variable name with the `r'_ prefix to the list
		}

		/************* Stage 2 – regress the residual for Y on the residuals for the pooled X ****************/
		reg `r_list' if `touse', cluster(`cluster') nocons   // the reported s.e. here are good and do NOT need
				// a correction for stage 1 estimation noise or degrees of freedom – see footnote 39 in appendix A.
				// The basic intuition is that b/c stage 2 estimates are identical to those in the full single-stage
				// model (A.1) with firm-specific slopes, these estimates are invariant to stage 1 estimation noise.
				// Therefore, the required s.e. correction for this noise is precisely zero.

		// piggyback on e() stuff created by the stage 2 regression and just tweak a few things
		tempname b
		matrix `b' = e(b)
		matrix colnames `b'=`pooledX'  // re-label the coefficients with the original variable names (without the `r'_ prefix)
		ereturn repost b=`b', rename

		ereturn local depvar "`Y'"
		ereturn local cmdline `"bigreg `0'"'
		ereturn local title "Linear regression with firm-specific slope coefficients"

		tempname TSS RSS r2
		sum `Y' if `touse'
		scalar `TSS' = r(Var)*(r(N)-1)  // r(Var) uses df=N-1
		scalar `RSS' = e(rss)
		sca `r2' = 1 - `RSS'/`TSS'

		ereturn scalar r2 = `r2'	// explanatory power of the full regression including all the firms-specific coeffs
		ereturn local cmd "bigreg"
	}

	dis ""
	dis "Linear regression with firmID-specific slope coefficients"
	dis _column(3) as text "FirmID = " as result "`firmID'"
	dis _column(3) as text "N = " as result e(N) as text "  (total obs. for firms with sufficient data to estimate firm-specific coefficients)"
	dis _column(3) as text "R2 = " as result %4.3f `r2' as text "  (total explanatory power of the full model with firm-specific coefficients)"
	dis _column(3) as text "Only the pooled coefficients are tabulated below."
	ereturn display
end