# set some ipython notebook properties
%matplotlib inline

# set degree of verbosity (adapt to INFO for more verbose output)
import logging
logging.basicConfig(level=logging.WARNING)

# set figure sizes
import pylab
pylab.rcParams['figure.figsize'] = (10.0, 8.0)

# set display width for pandas data frames
import pandas as pd
pd.set_option('display.width', 1000)

"""
Explanation: FaST-LMM Manual
Factored Spectrally Transformed Linear Mixed Models
Version 0.2.15
Microsoft Research, July 15, 2015
Introduction
FaST-LMM, which stands for Factored Spectrally Transformed Linear Mixed Models, is a program for performing both single-SNP and SNP-set genome-wide association studies (GWAS) on extremely large data sets.  This release contains the improvements described in Widmer et al., Scientific Reports 2014, and tests for epistasis.
For help with the software, please contact fastlmm@microsoft.com.
Citing FaST-LMM
If you use FaST-LMM in any published work, please cite the relevant manuscript(s) describing it.
Data preparation
This version of FaST-LMM is designed for use with randomly ascertained data with Gaussian residuals.  If you have case-control data with substantial ascertainment bias, you should first transform your phenotype(s) using LEAP (Weissbrod et al., arXiv 2014), which is available at https://github.com/omerwe/LEAP.  If you are analyzing continuous phenotypes with non-Gaussian residuals, you should first transform your phenotype(s) using Warped-LMM (Fusi et al., Nature Commun 2014), available at https://github.com/MicrosoftGenomics/WarpedLMM.
FaST-LMM uses four input files containing (1) the SNP data to be tested, (2) the SNP data used to determine the genetic similarity matrix (GSM) between individuals (which can be different from 1), (3) the phenotype data, and (4, optionally) a set of covariates.
SNP files can be in PLINK format (ped/map, tped/tfam, bed/bim/fam, or fam/dat/map).  For the most speed, use the binary format in SNP major order.  See the PLINK manual http://pngu.mgh.harvard.edu/~purcell/plink/ (Purcell et al., Am J Hum Genet 2007) for further details.  FaST-LMM also supports Hdf5 file format http://www.hdfgroup.org/HDF5/whatishdf5.html.  See https://github.com/MicrosoftGenomics/PySnpTools for more details.  Note that each SNP is standardized to have mean zero and standard deviation one across all individuals before processing.  Missing values are mean imputed.
The required file containing the phenotype uses the PLINK alternate phenotype format with no header.  The covariate file also uses this format (with additional columns for multiple covariates).
Notebook preparation and general use
To prepare this notebook to run analyses, please run the following script.
End of explanation
"""


# import the algorithm
from fastlmm.association import single_snp_leave_out_one_chrom

# set up data
##############################
bed_fn = "../../tests/datasets/synth/all"
pheno_fn = "../../tests/datasets/synth/pheno_10_causals.txt"
cov_fn = "../../tests/datasets/synth/cov.txt"

# run gwas
###################################################################
results_df = single_snp_leave_out_one_chrom(bed_fn,  pheno_fn, covar=cov_fn)

# print head of results data frame
import pandas as pd
pd.set_option('display.width', 1000)
print results_df.head(n=10)

# manhattan plot
import pylab
import fastlmm.util.util as flutil
flutil.manhattan_plot(results_df.as_matrix(["Chr", "ChrPos", "PValue"]),pvalue_line=1e-5,xaxis_unit_bp=False)
pylab.show()

# qq plot
from fastlmm.util.stats import plotp
plotp.qqplot(results_df["PValue"].values, xlim=[0,5], ylim=[0,5])

"""
Explanation: If you would like to run any of the code below from the command line, first copy it into a file (e.g., test.py), and then run it by typing python text.py at the command line.
If you would like to see all the options for a function just type ? &lt;function_name&gt; to an ipython prompt.
Single-SNP association testing
Traditional analysis: LMM(all)
First, let's run a standard LMM analysis in which the GSM uses (almost) all available SNPs.  The model for this analysis is called LMM(all) in  Widmer et al., Scientific Reports 2014.  We'll apply this model to the synthetic data in tests\datasets\synth.  The data has 500 samples with 5000 SNPs, and was generated from a Balding-Nichols model with FST=0.05.
When using a linear mixed model for association analysis, it is important to avoid proximal contamination (Lippert et al., Nat Meth 2011). To understand proximal contamination, first note that a LMM with no fixed effects, using a realized relationship matrix for the GSM (as FaST-LMM does), is mathematically equivalent to linear regression of the SNPs on the phenotype, with weights integrated over independent Normal distributions having the same variance (e.g., Hayes et al., Genet Res 2009).  That is, a LMM using a given set of SNPs for the GSM is equivalent to a form of linear regression using those SNPs as covariates to correct for confounding.  This equivalence implies that, when testing a given SNP, that SNP (and SNPs physically close to it) should be excluded from the computation of the GSM.  If not, when testing a particular SNP, we would also be using that same SNP as a covariate, making the log likelihood of the null model higher than it should be, thus leading to deflation of the test statistic and loss of power.
Excluding the SNP you are testing and those SNPs in close proximity to it from the GSM in a naïve way is extremely computationally expensive.  A computationally efficient approach for performing the exclusion is to use a GSM computed from all but chromosome $i$ when testing SNPs on chromosome $i$ (Lippert et al., Nat Meth 2011).  We call this approach leave out one chromosome (LOOC).  The analysis here does this.
End of explanation
"""


# example script for two kernel feature selection
# this takes a couple of minutes to run
import numpy as np
import pysnptools.util
from pysnptools.snpreader import Bed

from fastlmm.feature_selection.feature_selection_two_kernel import FeatureSelectionInSample
from pysnptools.util.pheno import loadOnePhen,loadPhen
from fastlmm.association import single_snp

# define file names
bed_fn = "../../tests/datasets/synth/all"
pheno_fn = "../../tests/datasets/synth/pheno_10_causals.txt"
cov_fn = "../../tests/datasets/synth/cov.txt"

# load data
###################################################################
snp_reader = Bed(bed_fn)
pheno = loadOnePhen(pheno_fn)
cov = loadPhen(cov_fn)

# intersect sample ids
snp_reader, pheno, cov = pysnptools.util.intersect_apply([snp_reader, pheno, cov])

# read in snps

# partition snps on chr5 vs rest
test_chr = 5
G0 = snp_reader[:,snp_reader.pos[:,0] != test_chr].read(order='C').standardize()
test_snps = snp_reader[:,snp_reader.pos[:,0] == test_chr].read(order='C').standardize()


y = pheno['vals'][:,0]
y -= y.mean()
y /= y.std()

# load covariates
X_cov = cov['vals']
X_cov.flags.writeable = False

# invoke feature selection to learn which SNPs to use to build G1
print "running feature selection conditioned on background kernel"
# partition data into the first 50 SNPs on chr1 and all but chr1

select = FeatureSelectionInSample(max_log_k=7, n_folds=7, order_by_lmm=True, measure="ll", random_state=42)
best_k, feat_idx, best_mix, best_delta = select.run_select(G0.val, G0.val, y, cov=X_cov)    

# plot out of sample error
select.plot_results(measure="ll")
# select.plot_results(measure="mse")

# print results
print "best_k:", best_k
print "best_mix:", best_mix
print "best_delta:", best_delta


###############################
# use selected SNPs to build G1
print feat_idx
G1 = G0[:,feat_idx]

results_df = single_snp(test_snps, pheno, G0=G0, G1=G1, mixing=best_mix, h2=None)

print 
print "results:"
print "#"*40
print results_df.head()

"""
Explanation: Improving power: LMM(all+selecct)
In Widmer et al., Scientific Reports 2014, we have shown that power can be increased while still maintaining control of type I error by using two GSMs: one based on all SNPs (G0) and one based on selected SNPs that are highly correlated with the phenotype (G1).  The model is called LMM(select + all).  This approach has greater computational demands (we recommend using a cluster computer when analyzing large data sets). Here is an example of how to apply this model to the synthetic data.
End of explanation
"""


# example script for two kernel without feature selection
import numpy as np
import pysnptools.util
from pysnptools.snpreader import Bed

from pysnptools.util.pheno import loadOnePhen,loadPhen
from fastlmm.association import single_snp

# define file names
bed_fn = "../../tests/datasets/synth/all"
pheno_fn = "../../tests/datasets/synth/pheno_10_causals.txt"
cov_fn = "../../tests/datasets/synth/cov.txt"

# load data
###################################################################
snp_reader = Bed(bed_fn)
pheno = loadOnePhen(pheno_fn)
cov = loadPhen(cov_fn)

# partition snps on chr5 vs chr6 vs rest
test_chr = 5
G1_chr = 6

G0 = snp_reader[:,(snp_reader.pos[:,0] != test_chr) & (snp_reader.pos[:,0] != G1_chr)]
G1 = snp_reader[:,snp_reader.pos[:,0] == G1_chr]
test_snps = snp_reader[:,snp_reader.pos[:,0] == test_chr]

results_df = single_snp(test_snps, pheno, G0=G0, G1=G1, mixing=None, h2=None)

print 
print "results:"
print "#"*40
print results_df.head()

"""
Explanation: Aside: In some applications, you may want to use two kernels constructed from two pre-specified sets of SNPs (i.e., with no feature selection). Here we show how to do that and how to simultaneously find h2 and the mixing weight between the kernels.
End of explanation
"""


from pysnptools.snpreader import Bed
from fastlmm.util import compute_auto_pcs

# define file names
bed_fn = "../../tests/datasets/synth/all"
cov_fn = "pcs_cov.txt"

# run PCgeno
#TODO: rename to auto_pcs
result = compute_auto_pcs(bed_fn, output_file_name=cov_fn)

print "selected number of PCs:", result["vals"].shape[1]

"""
Explanation: Improving speed and memory use when there is little family strucutre: LMM(select)+PCs
In the same publication, we have shown that a simpler, more computationally efficient model can be used when the data is  confounded only by population structure and not by family structure or cryptic relatedness.  Under these circumstances, we have found that a model with a single GSM based on selected SNPs in combination with principle components as fixed-effect covariates yields good control of type I error and power.  
This model, called LMM(select)+PCs, should be used with caution.  Even if you explicitly remove closely related individuals from your data, cryptic relatedness may remain.
To use this model, first identify the principle components to be used with the PCgeno algorithm as described in Widmer et al., Scientific Reports 2014.
End of explanation
"""


# import algorithms
from fastlmm.util.run_fastlmmc import runFASTLMM, runLMMSELECT

# set some file paths for fastlmmc
phen_fn = "../../tests/datasets/synth/pheno_10_causals.txt"
out_dir = "."
out_fn = "out.csv"
fastlmm_path = "../../external/fastlmmc"

# consists of two fastlmmc calls, one that does feature selection and one that runs GWAS
output = runLMMSELECT(bed_fn, phen_fn, out_dir, out_fn, bfileSim=bed_fn, fastlmm_path=fastlmm_path)

"""
Explanation: Now you can call FaST-LMM with feature selection and with these PCs as covariates. Below, we show how to do this with the C++ version of FaST-LMM. Soon, this functionality will be available in the python version.
End of explanation
"""


# import the algorithm and reader
from fastlmm.association import epistasis
from pysnptools.snpreader import Bed

# define file names
bed_reader = Bed("../../tests/datasets/synth/all")
pheno_fn = "../../tests/datasets/synth/pheno_10_causals.txt"
cov_fn = "../../tests/datasets/synth/cov.txt"

# partition data into the first 50 SNPs on chr1 and all but chr1
G0 = bed_reader[:,bed_reader.pos[:,0] != 1]
test_snps = bed_reader[:,bed_reader.pos[:,0] == 1][:,0:50]

# run epistasis analysis
results_df = epistasis(test_snps, pheno_fn, G0=G0, covar=cov_fn)

# print head of results data frame
import pandas as pd
pd.set_option('display.width', 1000)
print results_df.head(n=10)

# qq plot
from fastlmm.util.stats import plotp
plotp.qqplot(results_df["PValue"].values, xlim=[0,5], ylim=[0,5])

"""
Explanation: Epistasis
You can test for epistatic interactions between pairs of SNPs as well.  Here is an example analysis applied to the same synthetic data.  Note that this version of the code uses a likelihood ratio test based on maximum-likelihood estimates.  A REML-based version is in the works.
End of explanation
"""


# this will take a few minutes to run

# import the algorithm and reader
from fastlmm.association import snp_set

# define file names
test_snps_fn = "../../tests/datasets/synth/chr1"
pheno_fn = "../../tests/datasets/synth/pheno_10_causals.txt"
cov_fn = "../../tests/datasets/synth/cov.txt"
set_list_fn = "../../tests/datasets/synth/chr1.sets.txt"

# run SNP-set analysis
results_df = snp_set(test_snps_fn, set_list_fn, pheno_fn, covar=cov_fn, test="lrt")

# print head of results data frame
import pandas as pd
pd.set_option('display.width', 1000)
print results_df.head(n=10)

# qq plot
from fastlmm.util.stats import plotp
plotp.qqplot(results_df["P-value_adjusted"].values, xlim=[0,5], ylim=[0,5])


"""
Explanation: SNP-set association testing
SNP-set association testing is performed similarly to single-SNP testing.  The output columns are described in the Appendix.  Both LRT and score tests are supported, and the analysis may be run with or without a background GSM to correct for confounding.  Here is an example that uses the LRT and no background GSM--note the inflation in the results.  If you include a background GSM (using G0) type I error will be better controlled.  Again, proximal contamination is avoided using LOOC.
End of explanation
"""


