*** program for NB2 with multinomial endogeneity           ***
*** mata function for joint mmlogit & nb2                  ***
***                                                        ***
*** author: Partha Deb                                     ***
*** written: January 16, 2006                              ***
*** modified: January 27, 2006                             ***

clear
version 9.1
local mydir "."

mata:
function mtreatnb_lf(
		string scalar lnL, real matrix rmat, string scalar neq
	, string scalar neqall, string scalar nobs, string scalar sim
	, string scalar xb11, string scalar xb12, string scalar xb13
	, string scalar xb14, string scalar xb15, string scalar xb16
	, string scalar xb17, string scalar xb18, string scalar xb19
	, string scalar xb2
	, string scalar alpha
	, string scalar lam1, string scalar lam2, string scalar lam3
	, string scalar lam4, string scalar lam5, string scalar lam6
	, string scalar lam7, string scalar lam8, string scalar lam9
	, string scalar g1, string scalar g2, string scalar g3
	, string scalar g4, string scalar g5, string scalar g6
	, string scalar g7, string scalar g8, string scalar g9
	, string scalar g10, string scalar g11, string scalar g12
	, string scalar g13, string scalar g14, string scalar g15
	, string scalar g16, string scalar g17, string scalar g18
	, string scalar g19, string scalar g20
	, string scalar H)

	{
	neq=st_numscalar(neq)
	neqall=st_numscalar(neqall)
	nobs=st_numscalar(nobs)
	sim=st_numscalar(sim)

	ymml = *(findexternal("_mtreatnb_ymml"))
	y2 = *(findexternal("_mtreatnb_ynb"))
	xmml = *(findexternal("_mtreatnb_xmml"))
	xmml = (xmml, J(nobs,1,1))
	xnb = *(findexternal("_mtreatnb_xnb"))
	xnb = (xnb, J(nobs,1,1))
	st_view(xbmml, ., (xb11,xb12,xb13,xb14,xb15,xb16,xb17,xb18,xb19))
	st_view(xb2, ., xb2)

	st_view(alpha, ., alpha)
	st_view(lam, ., (lam1,lam2,lam3,lam4,lam5,lam6,lam7,lam8,lam9))

	exb = J(nobs,sim*neq,.)
	pmml = J(nobs,sim*neq,.)
	den = J(nobs,sim,1)
	for (j=1; j<=neq; j++) {
		exb[,((j-1)*sim+1)..(j*sim)] = exp(xbmml[,j]:+rmat[,((j-1)*sim+1)..(j*sim)])
		den = den + exb[,((j-1)*sim+1)..(j*sim)]
	}

	mml = (1:-rowsum(ymml)):/den
	for (j=1; j<=neq; j++) {
		pmml[,((j-1)*sim+1)..(j*sim)] = exb[,((j-1)*sim+1)..(j*sim)]:/den
		mml = mml :+ ymml[,j]:*pmml[,((j-1)*sim+1)..(j*sim)]
	}

	psi = 1:/alpha

	xb2mat = J(nobs,sim,0)
	for (j=1; j<=sim; j++) {
		xb2mat[,j] = xb2
	}
	for (j=1; j<=neq; j++) {
		xb2mat = xb2mat + lam[,j]:*rmat[,((j-1)*sim+1)..(j*sim)]
	}
	mu = exp(xb2mat)

	y2mmu = y2:-mu
	muppsi = mu:+psi
	psidmuppsi = psi:/muppsi
	lpalphamu = 1:+alpha:*mu
	nb2 = exp(lngamma(y2:+psi) :- lngamma(psi) :-  lngamma(y2:+1)  ///
							:+ psi:*ln(psidmuppsi) :+ y2:*ln(mu:/(muppsi)) )

	mmlnb2 = mml :* nb2
	L = (rowsum(mmlnb2))/sim
	L=rowmax((L , J(nobs,1,smallestdouble())))
	vlnL = ln(L)


	gmml = J(nobs,sim*neq,.)
	vg = J(nobs,20,.)
	for (j=1; j<=neq; j++) {
		gmml[,((j-1)*sim+1)..(j*sim)] = ymml[,j] :- pmml[,((j-1)*sim+1)..(j*sim)]
		vg[,j] = (1:/L):*rowsum(mmlnb2:*gmml[,((j-1)*sim+1)..(j*sim)])/sim
	}

	gnbb = y2mmu:/lpalphamu
	gnba =  -(digamma(y2:+psi) :- digamma(psi) :+ log(psidmuppsi) ///
					:+ 1 :- (y2:+psi):/muppsi):*psi
	vg[,10] = (1:/L):*rowsum(mmlnb2:*gnbb)/sim
	vg[,11] = (1:/L):*rowsum(mmlnb2:*gnba)/sim
	for (k=12; k<=neq+11; k++) {
		vg[,k] = (1:/L):*rowsum(mmlnb2:*gnbb:*rmat[,((k-12)*sim+1)..((k-11)*sim)])/sim
	}

	nparmsmml = neq*cols(xmml)
	nparmsnb = cols(xnb)+1+neq
	nparms = nparmsmml+nparmsnb
	H = J(nparms,nparms,0)
	c = 1
	for (k=1; k<=neq; k++) {
		r = c
		for (j=k; j<=neq; j++) {
			hmml = -pmml[,((j-1)*sim+1)..(j*sim)] ///
						:*((j==k):-pmml[,((k-1)*sim+1)..(k*sim)])
			h = ((-vg[,j]:*vg[,k] + (1:/L):*rowsum(mmlnb2 ///
			:*gmml[,((j-1)*sim+1)..(j*sim)]:*gmml[,((k-1)*sim+1)..(k*sim)])/sim ///
					:+ (1:/L):*rowsum(mmlnb2:*hmml)/sim) :* xmml)'*xmml
			H[r..(r+rows(h)-1),c..(c+cols(h)-1)] = h
			if (j>k) {
				H[c..(c+cols(h)-1),r..(r+rows(h)-1)] = h'
			}
			r = r + rows(h)
		}
		c = c + cols(h)
	}

	hnbbb = -mu:*(1:+alpha:*y2):/(lpalphamu:^2)
	hnbba = -y2mmu:/(lpalphamu:^2):*mu:*alpha
	hnbaa = gnba :+ (psi:^2) :*(trigamma(y2:+psi) :- trigamma(psi) ///
						:+ 1:/psi :- 2:/muppsi :+ psidmuppsi:/muppsi :+ y2:/(muppsi:^2))

	H[(nparmsmml+1)..(nparmsmml+cols(xnb)),(nparmsmml+1)..(nparmsmml+cols(xnb))] ///
			= ((-vg[,10]:*vg[,10] + (1:/L):*rowsum(mmlnb2:*gnbb:*gnbb)/sim ///
			+ (1:/L):*rowsum(mmlnb2:*hnbbb)/sim) :* xnb)'*xnb

	h = ((-vg[,11]:*vg[,10] + (1:/L):*rowsum(mmlnb2:*gnba:*gnbb)/sim ///
			+ (1:/L):*rowsum(mmlnb2:*hnbba)/sim))'*xnb
	H[nparmsmml+cols(xnb)+1,(nparmsmml+1)..(nparmsmml+cols(xnb))] = h
	H[(nparmsmml+1)..(nparmsmml+cols(xnb)),nparmsmml+cols(xnb)+1] = h'

	for (j=12; j<=neq+11; j++) {
		h = (-vg[,j]:*vg[,10] ///
			+ (1:/L):*rowsum(mmlnb2:*gnbb:*gnbb:*rmat[,((j-12)*sim+1)..((j-11)*sim)])/sim ///
			+ (1:/L):*rowsum(mmlnb2:*hnbbb:*rmat[,((j-12)*sim+1)..((j-11)*sim)])/sim)'*xnb
		H[nparmsmml+cols(xnb)+j-10,(nparmsmml+1)..(nparmsmml+cols(xnb))] = h
		H[(nparmsmml+1)..(nparmsmml+cols(xnb)),nparmsmml+cols(xnb)+j-10] = h'
	}

	H[nparmsmml+cols(xnb)+1,nparmsmml+cols(xnb)+1] ///
			= colsum(-vg[,11]:*vg[,11] + (1:/L):*rowsum(mmlnb2:*gnba:*gnba)/sim ///
			+ (1:/L):*rowsum(mmlnb2:*hnbaa)/sim)

	for (j=12; j<=neq+11; j++) {
		h = colsum(-vg[,j]:*vg[,11] ///
			+ (1:/L):*rowsum(mmlnb2:*gnbb:*gnba:*rmat[,((j-12)*sim+1)..((j-11)*sim)])/sim ///
			+ (1:/L):*rowsum(mmlnb2:*hnbba:*rmat[,((j-12)*sim+1)..((j-11)*sim)])/sim)
		H[nparmsmml+cols(xnb)+j-10,nparmsmml+cols(xnb)+1] = h
		H[nparmsmml+cols(xnb)+1,nparmsmml+cols(xnb)+j-10] = h'
	}

	for (k=12; k<=neq+11; k++) {
		for (j=k; j<=neq+11; j++) {
			h = colsum(-vg[,j]:*vg[,k] ///
				+ (1:/L):*rowsum(mmlnb2:*gnbb:*gnbb:*rmat[,((j-12)*sim+1)..((j-11)*sim)] ///
				:*rmat[,((k-12)*sim+1)..((k-11)*sim)])/sim ///
				:+ (1:/L):*rowsum(mmlnb2:*hnbbb:*rmat[,((j-12)*sim+1)..((j-11)*sim)] ///
				:*rmat[,((k-12)*sim+1)..((k-11)*sim)])/sim)
		H[nparmsmml+cols(xnb)+j-10,nparmsmml+cols(xnb)+k-10] = h
		H[nparmsmml+cols(xnb)+k-10,nparmsmml+cols(xnb)+j-10] = h'
		}
	}

	st_store(., lnL, vlnL)
	st_store(., (g1,g2,g3,g4,g5,g6,g7,g8,g9,g10,g11,g12,g13,g14,g15 ///
			,g16,g17,g18,g19,g20), vg)
	st_matrix("H", H)
	}

	mata mosave mtreatnb_lf(), dir(`mydir') replace

end