code: package xyz.lejon.bayes.models.regression;
  
  import java.io.IOException;
  import java.util.Arrays;
  
  import org.jblas.DoubleMatrix;
  import org.jblas.Solve;
  
  import xyz.lejon.bayes.models.probit.HorseshoeDOProbit;
  import xyz.lejon.configuration.OLSConfiguration;
  import xyz.lejon.sampling.BasicRDists;
  import xyz.lejon.sampling.FastMVNSamplerJBLAS;
  import xyz.lejon.sampling.InverseGammaDistribution;
  import xyz.lejon.utils.BlasOps;
  import xyz.lejon.utils.MatrixOps;
  
  /**
   * Multiple Bayesian Linear Regression with a horseshoe prior on the betas 
   * 
   * @author Leif Jonsson
   *
   */
  public class LinearRegressionJBlasHSPrior extends LinearRegressionHSPrior {
  	protected DoubleMatrix Xd;
  	protected DoubleMatrix XtX;
  	protected DoubleMatrix ys;
  	protected DoubleMatrix Xdt;
  	protected DoubleMatrix yty;
  	protected DoubleMatrix myPrecision;
  	
  	public LinearRegressionJBlasHSPrior(OLSConfiguration config, double[][] xs, double[] ys) throws IOException {
  		super(config, xs, ys);
  		Lambda = new double[noCovariates];
  		Arrays.fill(Lambda, 1.0); 
  		for (int j = 0; j < noCovariates; j++) {
  			betas[j] = BasicRDists.rnorm(0.0, 1.0);
  		}
  		
  		this.ys = new DoubleMatrix(ys);
  		
  		Xd = new DoubleMatrix(xs);
  		Xdt = (new DoubleMatrix(xs)).transpose();
  		XtX = Xdt.mmul(Xd);
  		yty = this.ys.transpose().mmul(this.ys);
  
  		// OLS estimate
  		//DoubleMatrix betahat = BlasOps.blasInvert(XtX).mmul(Xdt).mmul(this.ys);
  		//betas = betahat.toArray();
  		//System.out.println("BetaHat = " + MatrixOps.arrToStr(betas));
  				
  		myPrecision = new DoubleMatrix(XtX.rows,XtX.columns);
  		myPrecision.copy(XtX);
  		
  		System.out.println("################# Running LinearRegressionJBlasHSPrior #################");
  	}
  	
  	public static double sampleSigmaSquared(DoubleMatrix mu, DoubleMatrix precision, DoubleMatrix yty, int noObservations) {
  		double vn = v0  + noObservations / (double) 2;
  
  		DoubleMatrix tmp = yty.sub(mu.transpose().mmul(precision.mmul(mu)));
  		
  		if(tmp.rows!=1 && tmp.columns!=1) {
  			throw new IllegalStateException("intermediary has wrong dims!");
  		}
  		
  		double taun = tau0 + tmp.get(0) / 2;
  		
  		InverseGammaDistribution ig = new InverseGammaDistribution(vn, taun);
  
  		return ig.nextInverseGamma();
  	}
  	
  	@Override
  	public double sampleSigmaSquared(Object muIn, Object precisionIn) {
  		DoubleMatrix mu = (DoubleMatrix) muIn;
  		DoubleMatrix precision = (DoubleMatrix) precisionIn;
  
  		return sampleSigmaSquared(mu,precision,yty,noRows);
  	}
  	
  	@Override
  	public void sampleBeta() {
  		Tau    = HorseshoeDOProbit.sampleTau(Tau, betas, Lambda, Sigma, useIntercept);
  		Lambda = HorseshoeDOProbit.sampleLambda(Tau, betas, Lambda, Sigma);
  				
  		sampledTaus.add(Tau);
  		sampledLambdas.add(Lambda.clone());
  
  		// Rebuild updated precision matrix
  		for (int i = 0; i < myPrecision.rows; i++) {
  			double update;
  			// Don't regularize the intercept
  			if(i==0 && useIntercept) {
  				update = c;
  			} else {
  				update = (Math.pow(Tau, 2) * Math.pow(Lambda[i],2));
  			}
  			myPrecision.put(i, i, XtX.get(i, i) + 1 / update);
  		}
  		
  		DoubleMatrix localTilde = Xdt.mmul(ys);
  		DoubleMatrix localMu = new DoubleMatrix(Xd.columns, 1);
  
  		localMu = Solve.solve(myPrecision, localTilde);
  		DoubleMatrix pInv = BlasOps.blasInvert(myPrecision);
  		
  		for (int i = 0; i < pInv.rows; i++) {
  			for (int j = 0; j < pInv.columns; j++) {				
  				if(Double.isNaN(pInv.get(i,i))) {
  					System.err.println(
  							"Covariance matrix is broken"
  									+ " betas=" + MatrixOps.arrToStr(betas) + "\n"
  									+ " Tau=" + Tau + "\n"
  									+ " Lambda=" + MatrixOps.arrToStr(Lambda) + "\n"
  							);
  					System.exit(-1);
  				}
  			}
  		}
  
  		betas = sampleBeta(localMu, pInv, Sigma);
  		
  		Sigma  = Math.sqrt(sampleSigmaSquared(localMu, myPrecision));
  	}
  
  	public static double [] sampleBeta(DoubleMatrix localMu, DoubleMatrix pInv, double sigma) {
  		FastMVNSamplerJBLAS mvn = new FastMVNSamplerJBLAS(pInv.muli(Math.pow(sigma,2)));
  		double [] betas = mvn.sample(localMu);
  		for (int i = 0; i < betas.length; i++) {
  			if(Double.isNaN(betas[i])) {
  				throw new IllegalArgumentException(
  						"Sampled invalid beta:"
  						+ " betas=" + MatrixOps.arrToStr(betas) + "\n"
  						);
  			}
  		}
  		return betas;
  	}
  }
  