def mapper_gather_lots(i_fold_and_pair):
i_fold, (train_idx, test_idx) = i_fold_and_pair
logging.info("Working on GWAS_1K and k search, chrom={0}, i_fold={1}".format(test_chr, i_fold))
G_train = G_for_chrom[train_idx,:]
#Precompute whole x whole standardized on train
from fastlmm.association.single_snp import _internal_determine_block_size, _block_size_from_GB_goal
min_count = _internal_determine_block_size(G_for_chrom, None, None, force_full_rank, force_low_rank)
block_size = _block_size_from_GB_goal(GB_goal, G_for_chrom.iid_count, min_count)
K_whole_unittrain = _SnpWholeWithTrain(whole=G_for_chrom,train_idx=train_idx, standardizer=Unit(), block_size=block_size).read()
assert np.array_equal(K_whole_unittrain.iid,G_for_chrom.iid),"real assert"
K_train = K_whole_unittrain[train_idx]
single_snp_result = single_snp(test_snps=G_train, K0=K_train, pheno=pheno, #iid intersection means when can give the whole covariate and pheno
covar=covar, leave_out_one_chrom=False,
GB_goal=GB_goal, force_full_rank=force_full_rank, force_low_rank=force_low_rank, mixing=mixing, h2=h2)
is_all = (i_fold == n_folds) if n_folds > 1 else True
k_list_in = [0] + [int(k) for k in k_list if 0 < k and k < len(single_snp_result)]
if is_all:
top_snps = list(single_snp_result.SNP[:max_k])
else:
top_snps = None
if i_fold == n_folds:
k_index_to_nLL = None
else:
k_index_to_nLL = []
for k in k_list_in:
top_k = G_for_chrom[:,G_for_chrom.sid_to_index(single_snp_result.SNP[:k])]
logging.info("Working on chr={0}, i_fold={1}, and K_{2}".format(test_chr,i_fold,k))
top_k_train = top_k[train_idx,:] if k > 0 else None
fastlmm = FastLMM(force_full_rank=force_full_rank, force_low_rank=force_low_rank,GB_goal=GB_goal)
fastlmm.fit(K0_train=K_train, K1_train=top_k_train, X=covar, y=pheno,mixing=mixing,h2=h2) #iid intersection means when can give the whole covariate and pheno
top_k_test = top_k[test_idx,:] if k > 0 else None
K0_whole_test = K_whole_unittrain[:,test_idx]
nLL = fastlmm.score(K0_whole_test=K0_whole_test,K1_whole_test=top_k_test,X=covar,y=pheno) #iid intersection means when can give the whole covariate and pheno
k_index_to_nLL.append(nLL)
if i_fold > 0:
k_list_in = None
return k_list_in, top_snps, k_index_to_nLL
def mapper_gather_lots(i_fold_and_pair):
i_fold, (train_idx, test_idx) = i_fold_and_pair
logging.info("Working on GWAS_1K and k search, chrom={0}, i_fold={1}".format(test_chr, i_fold))
G_train = G_for_chrom[train_idx,:]
#Precompute whole x whole standardized on train
from fastlmm.association.single_snp import _internal_determine_block_size, _block_size_from_GB_goal
min_count = _internal_determine_block_size(G_for_chrom, None, None, force_full_rank, force_low_rank)
block_size = _block_size_from_GB_goal(GB_goal, G_for_chrom.iid_count, min_count)
K_whole_unittrain = _SnpWholeWithTrain(whole=G_for_chrom,train_idx=train_idx, standardizer=Unit(), block_size=block_size).read()
assert np.array_equal(K_whole_unittrain.iid,G_for_chrom.iid),"real assert"
K_train = K_whole_unittrain[train_idx]
single_snp_result = single_snp(test_snps=G_train, K0=K_train, pheno=pheno, #iid intersection means when can give the whole covariate and pheno
covar=covar, leave_out_one_chrom=False,
GB_goal=GB_goal, force_full_rank=force_full_rank, force_low_rank=force_low_rank, mixing=mixing, h2=h2)
is_all = (i_fold == n_folds) if n_folds > 1 else True
k_list_in = [0] + [int(k) for k in k_list if 0 < k and k < len(single_snp_result)]
if is_all:
top_snps = list(single_snp_result.SNP[:max_k])
else:
top_snps = None
if i_fold == n_folds:
k_index_to_nLL = None
else:
k_index_to_nLL = []
for k in k_list_in:
top_k = G_for_chrom[:,G_for_chrom.sid_to_index(single_snp_result.SNP[:k])]
logging.info("Working on chr={0}, i_fold={1}, and K_{2}".format(test_chr,i_fold,k))
top_k_train = top_k[train_idx,:] if k > 0 else None
fastlmm = FastLMM(force_full_rank=force_full_rank, force_low_rank=force_low_rank,GB_goal=GB_goal)
fastlmm.fit(K0_train=K_train, K1_train=top_k_train, X=covar, y=pheno,mixing=mixing,h2raw=h2) #iid intersection means when can give the whole covariate and pheno
top_k_test = top_k[test_idx,:] if k > 0 else None
K0_whole_test = K_whole_unittrain[:,test_idx]
nLL = fastlmm.score(K0_whole_test=K0_whole_test,K1_whole_test=top_k_test,X=covar,y=pheno) #iid intersection means when can give the whole covariate and pheno
k_index_to_nLL.append(nLL)
if i_fold > 0:
k_list_in = None
return k_list_in, top_snps, k_index_to_nLL
def k_index_to_nLL_mapper(k):
_, G_in, pheno_in, covar_in = _fixup(test_snps,
G,
pheno,
covar,
count_A1=count_A1)
nll_sum = 0
mse_sum = 0
n_folds_in = 0
for fold_index, (train_idx,
test_idx) in _kfold(G.iid_count,
n_folds,
seed,
end_with_all=False,
iid_to_index=G.iid_to_index):
assert set(train_idx).isdisjoint(set(test_idx)), "real assert"
top_snps_in_fold = fold_index_to_top_snps[fold_index][:k]
sid_idx_in_fold = G_in.sid_to_index(top_snps_in_fold)
G_train = G_in[train_idx, sid_idx_in_fold] if k > 0 else None
fastlmm = FastLMM(force_full_rank=force_full_rank,
force_low_rank=force_low_rank,
GB_goal=GB_goal)
fastlmm.fit(
K0_train=G_train,
X=covar_in[train_idx, :],
y=pheno_in[train_idx, :],
h2raw=h2
) #iid intersection means when can give the whole covariate and pheno
G_test = G_in[
test_idx, sid_idx_in_fold] if k > 0 else KernelIdentity(
G_in.iid, G_in.iid[test_idx]
) #!!! instead of this, which blows up when # of iids is large, should switch to linear regression model with k is 0
nll, mse = fastlmm.score(
K0_whole_test=G_test,
X=covar_in[test_idx, :],
y=pheno_in[test_idx, :],
return_mse_too=True
) #iid intersection means when can give the whole covariate and pheno
nll_sum += nll
mse_sum += mse
n_folds_in += 1
logging.info("k={0},nLL={1},average mse={2}".format(
k, nll_sum, mse_sum / n_folds_in))
return nll_sum
def k_index_to_nLL_mapper(k):
_, G_in, pheno_in, covar_in = _fixup(test_snps, G, pheno, covar,count_A1=count_A1)
nll_sum=0
mse_sum = 0
n_folds_in = 0
for fold_index, (train_idx, test_idx) in _kfold(G.iid_count, n_folds, seed, end_with_all=False,iid_to_index=G.iid_to_index):
assert set(train_idx).isdisjoint(set(test_idx)), "real assert"
top_snps_in_fold = fold_index_to_top_snps[fold_index][:k]
sid_idx_in_fold = G_in.sid_to_index(top_snps_in_fold)
G_train = G_in[train_idx,sid_idx_in_fold] if k > 0 else None
fastlmm = FastLMM(force_full_rank=force_full_rank, force_low_rank=force_low_rank,GB_goal=GB_goal)
fastlmm.fit(K0_train=G_train, X=covar_in[train_idx,:], y=pheno_in[train_idx,:], h2raw=h2) #iid intersection means when can give the whole covariate and pheno
G_test = G_in[test_idx,sid_idx_in_fold] if k > 0 else KernelIdentity(G_in.iid,G_in.iid[test_idx]) #!!! instead of this, which blows up when # of iids is large, should switch to linear regression model with k is 0
nll,mse = fastlmm.score(K0_whole_test=G_test,X=covar_in[test_idx,:],y=pheno_in[test_idx,:],return_mse_too=True) #iid intersection means when can give the whole covariate and pheno
nll_sum += nll
mse_sum += mse
n_folds_in += 1
logging.info("k={0},nLL={1},average mse={2}".format(k,nll_sum,mse_sum / n_folds_in))
return nll_sum