*! version 1.1.0  2012-11-30
*! Distribution free estimation of heteroskedastic binary response models
*! Jason R. Blevins and Shakeeb Khan

program define dfbr, eclass
   version 11
   capt findfile lmoremata.mlib
   if _rc {
     di as error "-moremata- is required; type {stata ssc install moremata}"
     error 499
   }
   syntax varlist(fv) [if] [in] [, {SIEVE|LOCAL} ///
                       NORMAL BANDwidth(real 0.0) K1(real 0.0) K2(real 0.0) ///
                       BASIS(varlist fv) ///
                       NMITER(integer -1) NMDELTA(real 1.0) ///
                       BREP(integer 50) SEED(integer 0) NOCONSTant ///
                       Level(integer $S_level)]
   marksample touse

   /*
    * Sanity checks
    */
   quietly count if `touse'
   if `r(N)' == 0 {
     error 2000
   }
   if ("`sieve'" != "") + ("`local'" != "") > 1 {
     di as err "Only one of sieve or local may be used"
     exit 198
   }
   if ("`local'" == "") + ("`normal'" != "") > 1 {
     di as err "normal option invalid for sieve nonlinear least squares"
     exit 198
   }
   if `bandwidth' < 0 {
     di as error "bandwidth must be positive"
     exit 198
   }
   if (`brep' != 0) + (`brep' < 2) > 1 {
     di as error "brep must either be zero or an integer greater than 1"
     exit 198
   }
   if ("`sieve'" != "") + (`bandwidth' != 0) > 1 {
     di as error "bandwidth may only be specified for local nonlinear least squares"
     exit 198
   }
   if (`k1' < 0) + (`k2' < 0) > 0 {
     di as error "Bandwidth constants K1 and K2 must be positive"
     exit 198
   }
   if (`k1' == 0) + (`k2' > 0) > 1 {
     di as error "Both bandwidth constants K1 and K2 must be specified together"
     exit 198
   }
   if (`k2' == 0) + (`k1' > 0) > 1 {
     di as error "Both bandwidth constants K1 and K2 must be specified together"
     exit 198
   }
   if `nmdelta' <= 0 {
     di as error "Nelder-Mead delta must be positive"
     exit 198
   }

   /* Seed the RNG (if provided) */
   if (`seed' != 0) {
     set seed `seed'
   }

   /*
    * Parse dependent and independent variables, handling factor
    * variables and checking for multicollinearity.  Require at least
    * one independent variable to be present.
    */
   gettoken y xvars : varlist
   dfbr_rmcoll `xvars'
   local xvars "`r(varlist)'"
   _fv_check_depvar `y'
   if "`xvars'" == "" {
     di as err "No independent variables specified"
     exit 102
   }
   local nvar: word count `xvars'

   /* Require at least two independent variables when noconstant is given */
   if ("`noconstant'" != "") + (`nvar' < 2) > 1 {
     di as err "constant term required with only one independent variable"
     exit 102
   }

   /* Store coefficient names, handling normalization of the last one */
   local lastvar: word `nvar' of `xvars'
   if ("`noconstant'" == "") {
     local coeffs "_cons `xvars'"
   }
   else {
     local coeffs "`xvars'"
   }
   local coeffs: subinstr local coeffs "`lastvar'" "", word

   /*
    * Store starting values
    */

   /* OLS */
   tempvar ydiff
   local olsvars: subinstr local xvars "`lastvar'" "", word
   quietly gen `ydiff' = `y' - `lastvar'
   quietly regress `ydiff' `olsvars', `noconstant'
   matrix define b_ols = e(b)

   /* LAD */
   if ("`noconstant'" == "") {
     quietly qreg `ydiff' `olsvars'
     matrix define b_lad = e(b)
   }
   else {
     matrix define b_lad = J(1,`nvar',-1)
   }

   /* Probit */
   quietly probit `y' `xvars', `noconstant'
   matrix define b_probit = e(b)

   /* Logit */
   quietly logit `y' `xvars', `noconstant'
   matrix define b_logit = e(b)

   /*
    * Perform local or sieve nonlinear least squares estimation as requested
    */
   if "`local'" != "" {
     local title "Local Nonlinear Least Squares (LNLLS)"

     if "`normal'" != "" {
       local method "jackknife"
       mata: jklnlls("`y'", "`xvars'", "`noconstant'", "`touse'", `k1', `k2', `brep')
     }
     else {
       local method "local"
       mata: lnllsf("`y'", "`xvars'", "`noconstant'", "`touse'", `bandwidth', `brep')
     }
   }
   else {
      local title "Sieve Nonlinear Least Squares (SNLLS)"
      local method "sieve"

      /* Use independent variables as default basis */
      if "`basis'" == "" {
        local basis `xvars'
      }
      dfbr_rmcoll `basis'
      local basis `r(varlist)'
      mata: snlls("`y'", "`xvars'", "`noconstant'", "`touse'", "`basis'", `brep')
   }

   /* Collect postestimation results */
   matrix b = r(b)
   matrix start = r(start)
   local N = r(N)
   matrix coleq b = ""
   matrix colnames b = `coeffs'

   if `brep' > 1 {
     matrix V = r(V)
     matrix BS = r(BS)
     matrix coleq V = ""
     matrix roweq V = ""
     matrix colnames V = `coeffs'
     matrix rownames V = `coeffs'
     ereturn post b V, depname("`y'") obs(`N') esample(`touse')
     ereturn scalar brep = `brep'
     if (`seed' != 0) {
       ereturn scalar seed = `seed'
     }
     ereturn local vce "bootstrap"
     ereturn matrix BS = BS
   }
   else {
     ereturn post b, depname("`y'") obs(`N') esample(`touse')
   }
   ereturn scalar N = r(N)
   ereturn scalar K = r(K)
   ereturn scalar h = r(h)
   ereturn scalar level = `level'
   ereturn matrix start = start
   ereturn local cmdname "dfbr"
   ereturn local title "`title'"
   ereturn local method "`method'"

   if "`method'" == "local" {
     ereturn scalar h = r(h)
   }
   else if "`method'" == "jackknife" {
     ereturn scalar k1 = r(k1)
     ereturn scalar k2 = r(k2)
     ereturn scalar w1 = r(w1)
     ereturn scalar w2 = r(w2)
   }
   else if "`method'" == "sieve" {
     matrix alpha = r(alpha)
     if `brep' > 1 {
       matrix V_alpha = r(V_alpha)
       matrix BS_alpha = r(BS_alpha)
     }
     local basis "_cons `basis'"
     dfbr_basis_coeffs `basis'
     local basis_coeffs `r(varlist)'
     matrix coleq alpha = ""
     matrix colnames alpha = `coeffs' `basis_coeffs'
     ereturn matrix alpha = alpha
     if `brep' > 1 {
       matrix coleq V_alpha = ""
       matrix roweq V_alpha = ""
       matrix colnames V_alpha = `coeffs' `basis_coeffs'
       matrix rownames V_alpha = `coeffs' `basis_coeffs'
       ereturn matrix V_alpha = V_alpha
       ereturn matrix BS_alpha = BS_alpha
     }
     ereturn local basis "`basis'"
   }

   /* Format and display results */
   di
   di as text "`title'" _continue
   di _col(49) as text "Number of obs   = " as result %12.0f e(N)
   if ("`method'" == "local") {
     di _col(49) as text "Bandwidth       = " as result %12.5e `r(h)'
   }
   ereturn display, level(`level')
   di as txt "Coefficient on `lastvar' normalized to 1."
   if ("`method'" == "sieve") {
     di as txt "Sieve basis: `basis'"
   }
end

program dfbr_rmcoll, rclass
   version 11
   syntax varlist(fv) [if]
   marksample touse

   fvexpand `varlist' if `touse'
   local xvars "`r(varlist)'"
   _rmcoll `xvars' if `touse', expand `noconstant'
   local xvars ""
   foreach var in `r(varlist)' {
     if strpos("`var'", "o.") == 0 {
       local xvars "`xvars' `var'"
     }
   }
   return local varlist "`xvars'"
end

program dfbr_basis_coeffs, rclass
   version 11
   local coeffs ""
   while "`1'" != "" {
     local coeffs "`coeffs' g_`1'"
     macro shift
   }
   local coeffs: subinstr local coeffs "c." "_", all
   return local varlist "`coeffs'"
end

/* Constant macros */
local SQRT2 = sqrt(2)
local SQRT3 = sqrt(3)
local ALPHA_F = -0.5 * (1 - `SQRT2' + `SQRT3')
local BETA_F = 1.0
local N_START = 6

version 11
mata:
mata set matastrict on

/*----------------------------------------------------------------------------*
 * Local and sieve nonlinear least squares Mata routines                      *
 *----------------------------------------------------------------------------*/

/*
 * Prepare data for local NLLS estimation.
 */
void lnlls_setup(string scalar yvar,
                 string scalar xvars,
                 string scalar noconstant,
                 string scalar touse)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real matrix X0, Y0, X, data
  real colvector w
  real rowvector temp
  real scalar r

  /* Obtain dependent variable and regressors */
  st_view(Y0, ., tokens(yvar), touse)
  st_view(X0, ., tokens(xvars), touse)

  /* Number of observations and regressors */
  N = rows(Y0)
  K = cols(X0)

  /* Store dependent variable */
  Y = Y0

  /* Append a constant term unless noconstant is given */
  if (noconstant == "") {
    X = J(N,1,1), X0
    constant = 1
  }
  else {
    X = X0
    K = K - 1
    constant = 0
  }

  /* Apply standardization, save mean and s.d. */
  X = apply_standardization(X, constant, 1)

  /* Handle regressor with normalized coefficient */
  X1 = X[.,(1..K)]
  X2 = X[.,K+1]

  /*
   * Prepare starting values
   */

  /* Zeros */
  all_start = J(`N_START', K, 0)

  /* OLS */
  temp = st_matrix("b_ols")
  all_start[2,.] = temp[1,(K,1..K-1)]

  /* LAD */
  temp = st_matrix("b_lad")
  all_start[3,.] = temp[1,(K,1..K-1)]


  /* Probit */
  temp = st_matrix("b_probit")
  if (constant > 0) {
    /* Constant last, normalized coefficient next to last */
    all_start[4,.] = temp[1,(K+1,1..K-1)] / temp[1,K]
  }
  else {
    /* Normalized coefficient last */
    all_start[4,.] = temp[1,(K,1..K-1)] / temp[1,K+1]
  }

  /* Logit */
  temp = st_matrix("b_logit")
  if (constant > 0) {
    /* Constant last, normalized coefficient next to last */
    all_start[5,.] = temp[1,(K+1,1..K-1)] / temp[1,K]
  }
  else {
    /* Normalized coefficient last */
    all_start[5,.] = temp[1,(K,1..K-1)] / temp[1,K+1]
  }

  /* Ones */
  all_start[6,.] = J(1,K,1)

  /* Scale starting values */
  for (r = 2; r < `N_START'; r++) {
    all_start[r,.] = apply_scaling(all_start[r,.])
  }
}

/* Standardize explanatory variables */
real matrix apply_standardization(real matrix mat, real scalar hasconst,
                                  real scalar store)
{
  external real scalar K, N, h, opt_normal, verbose
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real colvector m_tmp, sd_tmp
  real scalar nvar

  nvar = cols(mat)
  sd_tmp = sqrt(diagonal(variance(mat)))
  if (hasconst > 0) {
    m_tmp = mean(mat)'
    m_tmp[1] = 0
    sd_tmp[1] = 1
  }
  else {
    m_tmp = J(nvar,1,0)
  }

  /* Standardize variables */
  for (i = 1; i <= nvar; i++) {
    mat[.,i] = (mat[.,i] :- m_tmp[i]) / sd_tmp[i]
  }

  /* Store mean and standard deviation */
  if (store > 0) {
    m = m_tmp
    sd = sd_tmp
  }

  return(mat)
}

/*
 * Apply scaling of beta due to standardization.
 */
real matrix apply_scaling(real rowvector beta)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real rowvector beta_scaled

  beta_scaled = J(1,K,0)

  if (constant > 0) {
    beta_scaled[1] = beta[1] + m[K+1]
    for (i = 2; i <= K; i++) {
      beta_scaled[1] = beta_scaled[1] + (sd[K+1] / sd[i]) * beta_scaled[i] * m[i]
    }
  }

  for (i = 1; i <= K; i++) {
    beta_scaled[i] = (sd[i] / sd[K+1]) * beta[i]
  }

  return(beta_scaled)
}

/*
 * Reverse scaling of beta due to standardization.
 */
real matrix reverse_scaling(real rowvector beta_scaled)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real rowvector beta

  beta = J(1,K,0)

  for (i = 1; i <= K; i++) {
    beta[i] = (sd[K+1] / sd[i]) * beta_scaled[i]
  }

  if (constant > 0) {
    beta[1] = beta[1] - m[K+1]
    for (i = 2; i <= K; i++) {
      beta[1] = beta[1] - (sd[K+1] / sd[i]) * beta_scaled[i] * m[i]
    }
  }

  return(beta)
}

/*
 * Main interface function for LNLLSF estimation in Stata.
 */
void lnllsf(string scalar yvar,
            string scalar xvars,
            string scalar noconstant,
            string scalar touse,
            real scalar bandwidth,
            real scalar brep)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real rowvector beta_hat
  real colvector w
  real matrix data, V, BS
  struct mm_bsstats scalar B

  /* Construct a single data matrix for estimation */
  lnlls_setup(yvar, xvars, noconstant, touse)
  data = Y, X1, X2
  w = J(N,1,1)

  /* Use the alternative regression function */
  opt_normal = 0

  /* Verbosity during optimization */
  verbose = 0

  /* Feasible optimal bandwidth estimation */
  if (bandwidth <= 0) {
    /* Estimate once using default bandwidth */
    h = N^(-1/5)
    beta_hat = lnlls_estimate(data, w)

    /* Calculate optimal bandwidth and store in global h */
    h = lnllsf_bandwidth(h, 0.1, beta_hat)
  }
  else {
    /* Store user provide bandwidth in global h */
    h = bandwidth
  }

  /* Estimate once and possibly bootstrap */
  if (brep > 1) {
    B = mm_bs(&lnlls_estimate(), data, w, brep)
    BS = B.rstat
    beta_hat = mm_bs_report(B, ("b"))
    V = mm_bs_report(B, ("v"))
    st_matrix("r(V)", V)
    st_matrix("r(BS)", BS)
  } else {
    beta_hat = lnlls_estimate(data, w)
  }

  /* Additional results */
  st_matrix("r(b)", beta_hat)
  st_matrix("r(start)", start)
  st_numscalar("r(N)", N)
  st_numscalar("r(K)", K)
  st_numscalar("r(h)", h)
}

/*
 * Main interface function for jackknife local nonlinear least squares
 * estimation in Stata.
 */
void jklnlls(string scalar yvar,
             string scalar xvars,
             string scalar noconstant,
             string scalar touse,
             real scalar k1in,
             real scalar k2in,
             real scalar brep)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real scalar w1, w2, k1, k2
  real rowvector beta_hat
  real colvector w
  real matrix data, V, BS
  struct mm_bsstats scalar B

  /* Construct a single data matrix for estimation */
  lnlls_setup(yvar, xvars, noconstant, touse)
  data = Y, X1, X2
  w = J(N,1,1)

  /* Use the normal CDF as regression function */
  opt_normal = 1

  /* Verbosity during optimization */
  verbose = 0

  /* Feasible optimal jackknife weights */
  if (k1in == 0 | k2in == 0) {
    k1 = 0.56334
    k2 = 1.0180
    w1 = 2.2389
    w2 = -1.2389
  }
  else {
    k1 = k1in
    k2 = k2in
    w1 = k2 / (k2 - k1)
    w2 = 1 - w1
  }

  /* Estimate once and possibly bootstrap */
  if (brep > 1) {
    B = mm_bs(&jklnlls_estimate(), data, w, brep)
    BS = B.rstat
    beta_hat = mm_bs_report(B, ("b"))
    V = mm_bs_report(B, ("v"))
    st_matrix("r(V)", V)
    st_matrix("r(BS)", BS)
  } else {
    beta_hat = jklnlls_estimate(data, w)
  }

  /* Additional results */
  st_matrix("r(b)", beta_hat)
  st_matrix("r(start)", start)
  st_numscalar("r(N)", N)
  st_numscalar("r(K)", K)
  st_numscalar("r(w1)", w2)
  st_numscalar("r(w2)", w1)
  st_numscalar("r(k1)", k2)
  st_numscalar("r(k2)", k1)
}

/*
 * Obtain local NLLS estimates using the selected regression function.
 */
real rowvector lnlls_estimate(real matrix data,
                              real colvector w)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real scalar objfn
  real matrix estimates, beta_hat, dups
  real colvector values, indices

  /* Data matrix form data = [ Y, X1, X2 ] */
  N = rows(data)
  Y = data[.,1]
  X1 = data[.,(2..K+1)]
  X2 = data[.,K+2]

  /* Initialize storage for estimates and functional values */
  estimates = J(`N_START', K, 0)
  values = J(`N_START', 1, 0)

  /* For each starting value... */
  for (r = 1; r <= `N_START'; r++) {
    start = all_start[r,.]
    estimates[r,.] = lnlls_estimate_once(data)
    values[r] = objfn
  }

  /* Return best estimate */
  minindex(values, 1, indices, dups)
  if (rows(indices) > 0) {
    r = indices[1]
    beta_hat = estimates[r,.]
    return(beta_hat)
  }
  else {
    return(J(1,K,.))
  }
}

/*
 * Obtain LNLLS estimates using the current starting values and
 * selected regression function.
 */
real rowvector lnlls_estimate_once(real matrix data)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real scalar objfn
  real matrix beta_hat
  real scalar err, nmiter, nmdelta

  /* Set up optimization problem */
  S = optimize_init()
  if (opt_normal == 1) {
    optimize_init_evaluator(S, &lnlls_gf2())
  }
  else {
    optimize_init_evaluator(S, &lnllsf_gf2())
  }
  optimize_init_params(S, start)
  optimize_init_which(S, "min")
  optimize_init_evaluatortype(S, "gf2")
  if (verbose <= 0) {
    optimize_init_tracelevel(S, "none")
  }

  /* Perform at most 10*K iterations of Nelder-Mead */
  nmiter = strtoreal(st_local("nmiter"))
  nmdelta = strtoreal(st_local("nmdelta"))
  if (nmiter < 0) {
    nmiter = 10 * K
  }
  if (nmiter != 0) {
    if (verbose > 0) {
      printf("(Using Nelder-Mead)\n")
    }
    optimize_init_technique(S, "nm")
    optimize_init_conv_maxiter(S, nmiter)
    optimize_init_conv_warning(S, "off")
    optimize_init_verbose(S, verbose)
    optimize_init_nmsimplexdeltas(S, J(1,K,nmdelta))
    err = _optimize(S)
    beta_hat = optimize_result_params(S)
  }
  else {
    beta_hat = start
  }

  /* Finish optimization using BFGS */
  if (optimize_result_converged(S) != 1) {
    if (verbose > 0) {
      printf("(Switching to BFGS)\n")
      optimize_init_conv_warning(S, "on")
    }
    optimize_init_params(S, beta_hat)
    optimize_init_conv_maxiter(S, st_numscalar("c(maxiter)"))
    optimize_init_technique(S, "bfgs")
    err = _optimize(S)
  }

  /* Return values */
  if (err == 0) {
    beta_hat = optimize_result_params(S)
    beta_hat = reverse_scaling(beta_hat)
    objfn = optimize_result_value(S)
    return(beta_hat)
  }
  else {
    objfn = .
    return(J(1,K,.))
  }
}

/*
 * Obtain jackknife local NLLS estimates.
 */
real rowvector jklnlls_estimate(real matrix data,
                                real colvector w)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real scalar w1, w2, k1, k2
  real matrix beta_hat_2, beta_hat_1

  h = k1 * N^(-1/5)
  beta_hat_1 = lnlls_estimate(data, w)
  h = k2 * N^(-1/5)
  beta_hat_2 = lnlls_estimate(data, w)
  return(w1 * beta_hat_1 + w2 * beta_hat_2)
}

/*
 * Local nonlinear least squares (probit) objective function,
 * gradient, and Hessian.
 */
void lnlls_gf2(todo, param, sq_resid, grad, H)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real colvector beta, idxh, f0, f1, f2, s
  real scalar i

  /* Stata stores parameter vectors as rows */
  beta = param'

  /* Precalculate linear index and regression function values */
  idxh = (X1 * beta + X2) / h
  f0 = normal(idxh)
  s = Y - f0

  /* Function */
  sq_resid = s:^2

  if (todo >= 1) {
    /* Gradient */
    f1 = normalden(idxh)
    grad = -(2 / h) * s :* f1 :* X1

    if (todo == 2) {
      /* Hessian */
      f2 = -idxh :* normalden(idxh)
      H = J(K,K,0)
      for (i = 1; i <= N; i++) {
        H = H + (s[i] * f2[i] - f1[i]^2) * X1[i,.]' * X1[i,.]
      }
      H = -2 / h^2 * H
    }
  }
}

/*
 * Local nonlinear least squares (F) objective function, gradient, and Hessian.
 */
void lnllsf_gf2(todo, param, sq_resid, grad, H)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real colvector beta, idxh, f0, f1, f2, s
  real scalar i

  /* Stata stores parameter vectors as rows */
  beta = param'

  /* Precalculate linear index and regression function values */
  idxh = (X1 * beta + X2) / h
  f0 = J(N, 1, 0.5 - `ALPHA_F' - `BETA_F') ///
     + 2 * `ALPHA_F' * normal(idxh) ///
     + 2 * `BETA_F' * normal(`SQRT2' * idxh)
  s = Y - f0

  /* Function */
  sq_resid = s:^2

  if (todo >= 1) {
    /* Gradient */
    f1 = 2 * `ALPHA_F' * normalden(idxh) ///
       + 2 * `SQRT2' * `BETA_F' * normalden(`SQRT2' * idxh)
    grad = -(2 / h) * s :* f1 :* X1

    if (todo == 2) {
      /* Hessian */
      f2 = -2 * `ALPHA_F' * idxh :* normalden(idxh) ///
         - 4 * `SQRT2' * `BETA_F' * idxh :* normalden(`SQRT2' * idxh)
      H = J(K,K,0)
      for (i = 1; i <= N; i++) {
        H = H + (s[i] * f2[i] - f1[i]^2) * X1[i,.]' * X1[i,.]
      }
      H = -2 / h^2 * H
    }
  }
}

/*
 * Feasible optimal bandwidth calculation for LNLLSF estimator.
 */
real scalar lnllsf_bandwidth(real scalar h,
                             real scalar delta,
                             real rowvector beta)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  real matrix Q, Qinv, D
  real colvector A, T, T_delta, idx, idxh, SY
  real colvector regf0_idxh, regf1_idxh, regf2_idxh
  real scalar i, h_delta, h_opt, acorr, lambda, top, bot

  /* Precalculate some useful quantities */
  idx = X1 * beta' + X2
  idxh = idx / h
  regf0_idxh = J(N, 1, 0.5 - `ALPHA_F' - `BETA_F') ///
             + 2 * `ALPHA_F' * normal(idxh) ///
             + 2 * `BETA_F' * normal(`SQRT2' * idxh)
  regf1_idxh = 2 * `ALPHA_F' * normalden(idxh) ///
             + 2 * `SQRT2' * `BETA_F' * normalden(`SQRT2' * idxh)
  regf2_idxh = -2 * `ALPHA_F' * idxh :* normalden(idxh) ///
             - 4 * `SQRT2' * `BETA_F' * idxh :* normalden(`SQRT2' * idxh)
  SY = Y - regf0_idxh

  /* Estimate Q */
  Q = J(K, K, 0)
  for (i = 1; i <= N; i++) {
    Q = Q + (SY[i] * regf2_idxh[i] - regf1_idxh[i]^2) * X1[i,.]' * X1[i,.]
  }
  Q = -2 * Q / (N * h^2)
  Qinv = qrinv(Q)

  /* Estimate D */
  D = J(K, K, 0)
  T = J(K,1,0)
  for (i = 1; i <= N; i++) {
    T = -2 * SY[i] * regf1_idxh[i] * X1[i,.]' / h
    D = D + T * T'
  }
  D = (h / N) * D

  /* Estimate A (use bandwidth h^delta) */
  h_delta = h^delta
  T_delta = J(K, 1, 0)
  regf1_idxh = 2 * `ALPHA_F' * normalden(idx / h_delta) ///
             + 2 * `SQRT2' * `BETA_F' * normalden(`SQRT2' / h_delta * idx)
  for (i = 1; i <= N; i++) {
    T = -2 * SY[i] * regf1_idxh[i] * X1[i,.]' / h_delta
    T_delta = T_delta + T
  }
  T_delta = T_delta / N
  A = T_delta / h_delta^2

  /* Apply small sample correction to A */
  acorr = 1 - (h / h_delta)^2
  A = A / acorr

  /* Calculate optimal bandwidth */
  top = trace(Qinv * D * Qinv)
  bot = 4 * A' * Qinv * Qinv' * A
  lambda = top / bot
  h_opt = (lambda / N)^(1/5)

  /* Return value */
  return(h_opt)
}

/*
 * Main interface function for sieve NLLS estimation in Stata.
 */
void snlls(string scalar yvar,
           string scalar xvars,
           string scalar noconstant,
           string scalar touse,
           string scalar basis,
           real scalar brep)
{
  external real scalar K, N, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real matrix Z
  external real scalar L
  real rowvector alpha_hat, beta_hat
  real colvector w
  real matrix data, V, BS
  struct mm_bsstats scalar B

  /* Construct a single data matrix for estimation */
  snlls_setup(yvar, xvars, basis, noconstant, touse)
  data = Y, X1, X2, Z
  w = J(N,1,1)

  /* Verbosity during optimization */
  verbose = 0

  /* Estimate once and possibly bootstrap */
  if (brep > 1) {
    B = mm_bs(&snlls_estimate(), data, w, brep)
    alpha_hat = mm_bs_report(B, ("b"))
    BS = B.rstat
    V = mm_bs_report(B, ("v"))
    st_matrix("r(V_alpha)", V)
    st_matrix("r(V)", V[1..K,1..K])
    st_matrix("r(BS_alpha)", BS)
    st_matrix("r(BS)", BS[.,1..K])
  } else {
    alpha_hat = snlls_estimate(data, w)
  }
  beta_hat = alpha_hat[.,1..K]

  /* Additional results */
  st_matrix("r(b)", beta_hat)
  st_matrix("r(alpha)", alpha_hat)
  st_matrix("r(start)", all_start)
  st_numscalar("r(N)", N)
  st_numscalar("r(K)", K)
  st_numscalar("r(h)", h)
}

/*
 * Prepare data for sieve NLLS estimation.
 */
void snlls_setup(string scalar yvar,
                 string scalar xvars,
                 string scalar basis,
                 string scalar noconstant,
                 string scalar touse)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real matrix Z
  external real scalar L
  real matrix Z0

  lnlls_setup(yvar, xvars, noconstant, touse)
  st_view(Z0, ., tokens(basis), touse)
  Z = J(N,1,1), Z0
  L = cols(Z)

  /* Standardize basis variables */
  Z = apply_standardization(Z, 1, 0)

  /* Add columns to starting values for sieve coefficients */
  all_start = all_start, J(`N_START', L, 0)
}

/*
 * Perform sieve NLLS estimation and return parameter values.
 */
real rowvector snlls_estimate(real matrix data,
                              real colvector w)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real matrix Z
  external real scalar L, objfn
  real matrix alpha_hat, estimates, dups
  real colvector values, indices

  /* Data matrix X has form data = [ Y, X1, X2, Z ] */
  N = rows(data)
  Y = data[.,1]
  X1 = data[.,(2..K+1)]
  X2 = data[.,K+2]
  Z = data[.,(K+2+1..K+2+L)]

  /* Initialize storage for estimates and functional values */
  estimates = J(`N_START', K+L, 0)
  values = J(`N_START', 1, 0)

  /* For each starting value... */
  for (r = 1; r <= `N_START'; r++) {
    start = all_start[r,.]
    estimates[r,.] = snlls_estimate_once(data)
    values[r] = objfn
  }

  /* Return best estimate */
  minindex(values, 1, indices, dups)
  if (rows(indices) > 0) {
    r = indices[1]
    alpha_hat = estimates[r,.]
    return(alpha_hat)
  }
  else {
    return(J(1,K+L,.))
  }
}

/*
 * Obtain sieve NLLS estimates using the current starting values.
 */
real rowvector snlls_estimate_once(real matrix X)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real matrix Z
  external real scalar L
  external real scalar objfn
  real matrix alpha_hat
  real scalar err, nmiter, nmdelta

  /* Set up optimization problem */
  S = optimize_init()
  optimize_init_evaluator(S, &snlls_gf2())
  optimize_init_params(S, start)
  optimize_init_which(S, "min")
  optimize_init_evaluatortype(S, "gf2")
  if (verbose <= 0) {
    optimize_init_tracelevel(S, "none")
  }

  /* Perform at most 10*K iterations of Nelder-Mead */
  nmiter = strtoreal(st_local("nmiter"))
  nmdelta = strtoreal(st_local("nmdelta"))
  if (nmiter < 0) {
    nmiter = 10 * K
  }
  if (nmiter != 0) {
    if (verbose > 0) {
      printf("(Using Nelder-Mead)\n")
    }
    optimize_init_technique(S, "nm")
    optimize_init_conv_maxiter(S, nmiter)
    optimize_init_conv_warning(S, "off")
    optimize_init_verbose(S, verbose)
    optimize_init_nmsimplexdeltas(S, J(1,K+L,nmdelta))
    err = _optimize(S)
    alpha_hat = optimize_result_params(S)
  }
  else {
    alpha_hat = start
  }

  /* Finish optimization using BFGS */
  if (optimize_result_converged(S) != 1) {
    if (verbose > 0) {
      printf("(Switching to BFGS)\n")
      optimize_init_conv_warning(S, "on")
    }
    optimize_init_params(S, alpha_hat)
    optimize_init_conv_maxiter(S, st_numscalar("c(maxiter)"))
    optimize_init_technique(S, "bfgs")
    err = _optimize(S)
  }

  /* Return values */
  if (err == 0) {
    alpha_hat = optimize_result_params(S)
    alpha_hat[1..K] = reverse_scaling(alpha_hat[1..K])
    objfn = optimize_result_value(S)
    return(alpha_hat)
  }
  else {
    objfn = .
    return(J(1,K+L,.))
  }
}

/*
 * Sieve nonlinear least squares objective function, gradient, and Hessian.
 */
void snlls_gf2(todo, param, sq_resid, grad, H)
{
  external real scalar K, N, h, opt_normal, verbose, constant
  external real rowvector start
  external real colvector Y, X2, m, sd
  external real matrix X1, all_start
  external real matrix Z
  external real scalar L
  real colvector alpha, beta, gamma, idx, idz, idx_g, g, idxg, phi0, phi1, phi2
  real matrix H11, H12, H22
  real scalar i

  /* Stata stores parameter vectors as rows */
  alpha = param'
  beta = alpha[1..K,.]
  gamma = alpha[K+1..K+L,.]

  /* Precalculate linear index and regression function values */
  idx = (X1 * beta + X2)
  idz = Z * gamma
  g = exp(idz)
  idx_g = idx :* g
  phi0 = normal(idx_g)
  s = Y - phi0

  /* Function */
  sq_resid = s:^2

  if (todo >= 1) {
    /* Gradient */
    phi1 = normalden(idx_g)
    grad = J(N,K+L,0)
    grad[.,1..K] = -2 * g :* s :* phi1 :* X1
    grad[.,K+1..K+L] = -2 * g :* s :* phi1 :* idx :* Z

    if (todo == 2) {
      /* Hessian */
      H11 = J(K,K,0)
      H22 = J(L,L,0)
      H12 = J(K,L,0)

      /* Derivative of normal pdf is phi'(u) = -x * phi(u) */
      phi2 = -idx :* phi1

      for (i = 1; i <= N; i++) {
        H11 = H11 + g[i]^2 * (s[i] * phi2[i] - phi1[i]^2) * X1[i,.]' * X1[i,.]

        H22 = H22 + g[i] * ( (s[i] * phi2[i] - phi1[i]^2) * idx[i]^2 * g[i] ///
                           + s[i] * phi1[i] * idx[i] ) * Z[i,.]' * Z[i,.]

        H12 = H12 + g[i] * ( s[i] * ( phi1[i] + idx[i] * phi2[i] * g[i] ) ///
                           - phi1[i]^2 * g[i] * idx[i] ) * X1[i,.]' * Z[i,.]
      }
      H = -2 * ( H11, H12 \ H12', H22 )
    }
  }
}

end