%calculateGWASpowerR2 Computes power and PGS R-squared in a meta-GWAS.
%
%  Author: R. de Vlaming
%  Date: May 2, 2016.
%
%  [dP,dR2] = calculateGWASpowerR2(iC,vN,vH2,mCGR,iM,iS) returns meta-GWAS
%  power per trait-associated SNP and expected PGS R-squared in a hold-out
%  sample with SNP-weightes based on the meta-analysis results.
%
%  Meta-analysis of GWAS results from iC studies with sample size per
%  study specified in vector vN (iC-dimensional) and SNP-heritability per
%  study (including hold-out sample) in vector vH2 ((iC+1)-dimensional). 
%  Cross-study genetic correlations are specified in (iC+1)-by-(iC+1) 
%  matrix mCGR. The effective number of causal SNPs is given by iM and the
%  total effective number of SNPs by iS.
%
%  Example: Compute (i) power of a meta-analysis of two studies with 40k
%  and 60k individuals and SNP-heritability of 60% and 40% respectively,
%  for a trait with 1k independent causal SNPs, with a genetic correlation
%  of 0.9 between the meta-analysis studies, assuming 100k independent
%  SNPs in total in the genome, and (ii) expected PGS R-square in the
%  hold-out sample, with heritability 50% and a genetic correlation of 0.8
%  with study 1 and 0.7 with study 2, using SNP-weights from the
%  meta-analysis to construct the PGS.
%
%     mCGR = [1.0 0.9 0.8 ;...
%             0.9 1.0 0.7 ;...
%             0.8 0.7 1.0];
%     [dP,dR2] = calculateGWASpowerR2(2,[40000 60000],[0.6 0.4 0.5],...
%                                      mCGR,1000,100000);
%
function [dP,dR2] = calculateGWASpowerR2(iC,vN,vH2,mCGR,iM,iS)
% determine threshold for genome-wide significant Z-score
dPval = 5*(10^(-8));
dZhit = norminv(dPval/2,0,1);

% determine sample-size of the meta-analysis
iN = sum(vN(1:iC));

% compute lowest eigenvalue of the CGR matrix
dMinEigCGR = min(eig(mCGR));

if dMinEigCGR > 0 % if CGR matrix is positive definite: continue

    % retrieve Cholesky decomposition of CGR matrix
    mGamma = chol(mCGR)';
    
    % set Var(Z-stat) of meta and numerator corr(PGS,Y) in hold-out
    dVarZmeta = 0;
    dNumeratorR = 0;
    
    % for each genetic factor in the meta-analysis
    for iI = 1:iC
        
        % set the standard-error contributed to the Z-statistic
        dThisSE = 0;
        
        % for each study affected by this genetic factor
        for iJ = iI:iC
        
            % update the standard-error contributed to the Z-statistic
            dThisSE = dThisSE + ...
                      (mGamma(iJ,iI)*vN(iJ)*sqrt(vH2(iJ)/(iM-vH2(iJ))));
            
            % update numerator of expected R-squared according to relation
            % hold-out study and current meta study with this factor
            dNumeratorR = dNumeratorR + ...
                          (mGamma(iC+1,iI)*mGamma(iJ,iI)*...
                           vN(iJ)*sqrt(vH2(iJ)/(iM-vH2(iJ))));
        end
        
        % add contribution to the Var(Z-stat) of this factor
        dThisVar = dThisSE^2;
        dVarZmeta = dVarZmeta + dThisVar;
        
    end
    
    % comute final R-squared numerator and denominator, and Var(Z-stat)
    dNumeratorR2   = dNumeratorR^2;
    dDenominatorR2 = dVarZmeta + ((iN*iS)/iM);
    dVarZmeta = 1 + (dVarZmeta/iN);
    
    % compute power resulting from Var(Z-stat)
    dP = 1 - erf(abs(dZhit)/sqrt(2*(dVarZmeta)));
    
    % compute PGS R-squared for the hold-out sample
    dR2 = vH2(iC+1)*dNumeratorR2/dDenominatorR2;

else % if CGR matrix is not positive definite: terminate
    
    disp(['ERROR: CGR matrix is not positive definite.']);
    
    % return NaNs
    dP  = NaN;
    dR2 = NaN;
    
end
end