from pylab import *

## Use simulation to check the coverage probabilities of the usual
## confidence interval for a regression coefficient.

## Number of cases
n = 20

## Number of covariates
p = 3

## Construct a design matrix.
'''X = rnorm(n*p)
X = array(X, c(n,p))
X = X - outer(array(1,n), apply(X,2,mean))  ## Center the columns
X = svd(X)$u                                ## Orthogonalize the columns
r = 0.3                                     ## Introduce a correlation of
X[,3] = r*X[,2] + sqrt(1-r^2)*X[,3]         ## r between columns 2 and 3.
X = cbind(array(1,n), X)
'''

X = randn(n,p)
X = X - outer(ones(n), X.mean(0))
X,_,_ = svd(X,full_matrices=0)
r = 0.3
X[:,2]= r*X[:,1] + sqrt(1-r**2)*X[:,2]
X = c_[ones(n), X]


## Some things we'll need below.
'''
H = svd(X)
H = H$v %*% diag(1/H$d^2) %*% t(H$v) ## equivalent to solve(t(X) %*% X)
dc = c(0,-1,1,0) ## Contrast
V = dc %*% H %*% dc
'''

from collections import namedtuple
SVD = namedtuple('SVD', 'u d v')
H = svd(X, 0)
H = SVD(*H)

H2 = dot( dot(H.v, diag(1/H.d**2)), H.v.T)  #ugly

def mdot(*args):
    pass
    return reduce(dot, args)
    
    
    
H3 = mdot(H.v, diag(1/H.d**2), H.v.T)

# bad
dc = c_[0, -1, 1, 0]   # dc is 2D array!!  ????

V = mdot(dc, H3, dc.T)  #check the type of V

# ok
dc = array([0,-1,1,0])  # dc is now 1D array
V = mdot(dc, H3, dc)


# ok 
dc = r_[0,1,1,0]  #dc is now 1D array
V = mdot(dc, H3, dc)


'''
## Error standard deviation.
sigma = 1

## Number of simulation replications.
nrep = 1000
'''

sigma = 1
nrep = 1000


for b in arange(0,1,0.1):
    # generate the population regression vector
    beta = zeros(p+1)
    beta[1] = b

    # The true value of the contrast. 
    tv = sum(dc * beta) 
    
    # storage for all the estimates of beta_1
    B = empty(nrep)  
    
    # The mean response
    EY = dot(X, beta)
    
    for k in arange(nrep):
        Y = EY + sigma * randn(n)
        
        # use polyfit or lstsq   
        m = lstsq(X,Y)
        
        # m[0] is the coefficients
        beta_hat = m[0]
        mse = 1.0/(n-p) *m[1]  # m[1] sum of square residual ?   
        
        rmse = sqrt(mse) 
        
        pe = sum ( dc * beta_hat)
        
        SE = rmse * sqrt(V)
        
        from scipy import stats
        
        f = stats.t.ppf(0.975, n-p-1)
        LB = pe - f*SE
        UB = pe + f*SE
        B[k]= ((tv > LB) and (tv < UB))  # and versus &
    
    print b, mean(B)
    
'''
## Loop over population values for beta_1.
for (b in seq(0, 1, 0.1)) {

    ## Generate the population regression vector.
    beta = rep(0, p+1)
    beta[2] = b

    ## The true value of the contrast.
    tv = sum(dc * beta)
    
    ## Storage for all the estimates of beta_1.
    B = array(0, nrep)

    ## The mean response.
    EY = X %*% beta

    ## Simulation replications.
    for (k in 1:nrep) {

        ## Generate response data from Y|X.
        Y = EY + sigma*rnorm(n)

        ## Fit the model.
        m = lm(Y ~ 0+X)
	s = summary(m)
        beta_hat = m$coeff
        mse = s$sigma^2 
        rmse = sqrt(mse)

        ## The estimated contrast
        pe = sum(dc * beta_hat)
        
        ## The estimated standard error for the contrast.
        SE = rmse * sqrt(V)

        ## The multiplier for the 95% CI.
        f = qt(0.975, n-p-1)

        ## The lower and upper limit of the CI.
        LB = pe - f*SE
        UB = pe + f*SE

	## Get the indicator of coverage for this sample.
        B[k] = 1*((tv > LB) & (tv < UB))
    }

    ## Print the true beta value and coverage probability.
    print(c(b, mean(B)))
}
'''