<repo_name>GoreLab/sorghum-multi-trait<file_sep>/codes/bayesian_network_analysis.py

#------------------------------------------------Modules-----------------------------------------------------#

# Load libraries:
import pandas as pd
import numpy as np
import os
import pickle
import re
import pystan as ps
import argparse
parser = argparse.ArgumentParser()


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Get flags:
parser.add_argument("-y", "--y", dest = "y", default = "error", help="Name of the file with the phenotypes")
parser.add_argument("-x", "--x", dest = "x", default = 'error', help="Name of the file with the features")
parser.add_argument("-model", "--model", dest = "model", default = "BN", help="Name of the model that can be: 'BN' or 'PBN', or 'DBN'")
parser.add_argument("-rpath", "--rpath", dest = "rpath", help="The path of the repository")
parser.add_argument("-opath", "--opath", dest = "opath", help="The path of the folder with general outputs")
parser.add_argument("-cvpath", "--cvpath", dest = "cvpath", help="The path of the folder to receive outputs of the step of cross-validation analysis")

# Parse the paths:
args = parser.parse_args()

# Subset arguments:
y = args.y
x = args.x
model = args.model
REPO_PATH = args.rpath
OUT_PATH = args.opath
CV_OUT_PATH = args.cvpath


#---------------------------------------------Loading data---------------------------------------------------#

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Get each trait file names:
if model=='PBN':
	y_0 = y.split("&")[0]
	y_1 = y.split("&")[1]
	x_0 = x.split("&")[0]
	x_1 = x.split("&")[1]

# Load data:
if (model=='BN') or bool(re.search('DBN', model)):
	# Reading adjusted means:
	y = pd.read_csv(args.y, index_col=0)
	# Reading feature matrix:
	X = pd.read_csv(args.x, index_col=0)

# Load data:
if model=='PBN':
	# Reading adjusted means for both traits:
	y_0 = pd.read_csv(y_0, index_col=0)
	y_1 = pd.read_csv(y_1, index_col=0)
	# Reading feature matrix for both traits:
	X_0 = pd.read_csv(x_0, index_col=0)
	X_1 = pd.read_csv(x_1, index_col=0)


#------------------------------------------Data input for stan-----------------------------------------------#

if model == 'BN':
	# Create an empty list to receive the dictionaries:
	dict_stan = []
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Drop the first column used just for mapping:
	X = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	for t in group:
		dict_stan.append(dict(p_z = X[index==t].shape[1],
							  n = X[index==t].shape[0],
						 	  Z = X[index==t],
						 	  y = y[index==t].values.flatten(),
						 	  phi = y[index==t].max().values[0]*10)) 

if model == 'PBN':
	# Create an empty list to receive the dictionaries:
	dict_stan = [] 
	# Subset time indexes and groups:
	index1 = X_1.iloc[:,0].values
	index0 = X_0.iloc[:,0].values
	group1 = X_1.iloc[:,0].unique()
	group0 = X_0.iloc[:,0].unique()
	# Drop the first column used just for mapping:
	X_0 = X_0.drop(X_0.columns[0], axis=1)
	X_1 = X_1.drop(X_1.columns[0], axis=1)
	# Case where the first input file have measures over time:
	if len(group1)==1:
		# Build dictionaries:
		for t in group0:
			dict_stan.append(dict(p_z = X_0.shape[1],
								  n_0 = X_0[index0==t].shape[0],
							 	  Z_0 = X_0[index0==t],
							 	  y_0 = y_0[index0==t].values.flatten(),
							 	  n_1 = X_1.shape[0],
							 	  Z_1 = X_1,
							 	  y_1 = y_1.values.flatten(),
							 	  phi = pd.concat([y_0[index0==t], y_1], axis=0).max().values[0]*10))
	# Case where the second input file have measures over time:
	if len(group0)==1:
		# Build dictionaries:
		for t in group1:
			dict_stan.append(dict(p_z = X_0.shape[1],
								  n_0 = X_0.shape[0],
							 	  Z_0 = X_0,
							 	  y_0 = y_0.values.flatten(),
							 	  n_1 = X_1[index1==t].shape[0],
							 	  Z_1 = X_1[index1==t],
							 	  y_1 = y_1[index1==t].values.flatten(),
							 	  phi = pd.concat([y_0, y_1[index1==t]], axis=0).max().values[0]*10))
	# Case where all input files have measures over time:
	if (len(group0)!=1) & (len(group1)!=1):
		# Build dictionaries:
		for t0, t1 in zip(group0, group1):
			dict_stan.append(dict(p_z = X_0.shape[1],
								  n_0 = X_0[index0==t0].shape[0],
							 	  Z_0 = X_0[index0==t0],
							 	  y_0 = y_0[index0==t0].values.flatten(),
							 	  n_1 = X_1[index1==t1].shape[0],
							 	  Z_1 = X_1[index1==t1],
							 	  y_1 = y_1[index1==t1].values.flatten(),
							 	  phi = pd.concat([y_0[index0==t0], y_1[index1==t1]], axis=0).max().values[0]*10))

if model == 'DBN-0~1':
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Subset the time covariate and the features:
	Z = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	dict_stan = dict(p_z = Z[index==group[0]].shape[1],
					 p_res = len(group),
					 n_0 = Z[index==group[0]].shape[0],
				 	 Z_0 = Z[index==group[0]],
				 	 y_0 = y[index==group[0]].values.flatten(),
					 n_1 = Z[index==group[1]].shape[0],
				 	 Z_1 = Z[index==group[1]],
				 	 y_1 = y[index==group[1]].values.flatten(),
				 	 phi = y.max().values[0]*10) 

if model == 'DBN-0~2':
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Subset the time covariate and the features:
	Z = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	dict_stan = dict(p_z = Z[index==group[0]].shape[1],
					 p_res = len(group),
					 n_0 = Z[index==group[0]].shape[0],
				 	 Z_0 = Z[index==group[0]],
				 	 y_0 = y[index==group[0]].values.flatten(),
					 n_1 = Z[index==group[1]].shape[0],
				 	 Z_1 = Z[index==group[1]],
				 	 y_1 = y[index==group[1]].values.flatten(),
					 n_2 = Z[index==group[2]].shape[0],
				 	 Z_2 = Z[index==group[2]],
				 	 y_2 = y[index==group[2]].values.flatten(),			 	  
				 	 phi = y.max().values[0]*10) 

if model == 'DBN-0~3':
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Subset the time covariate and the features:
	Z = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	dict_stan = dict(p_z = Z[index==group[0]].shape[1],
					 p_res = len(group),
					 n_0 = Z[index==group[0]].shape[0],
				 	 Z_0 = Z[index==group[0]],
				 	 y_0 = y[index==group[0]].values.flatten(),
					 n_1 = Z[index==group[1]].shape[0],
				 	 Z_1 = Z[index==group[1]],
				 	 y_1 = y[index==group[1]].values.flatten(),
					 n_2 = Z[index==group[2]].shape[0],
				 	 Z_2 = Z[index==group[2]],
				 	 y_2 = y[index==group[2]].values.flatten(),			 	  
					 n_3 = Z[index==group[3]].shape[0],
				 	 Z_3 = Z[index==group[3]],
				 	 y_3 = y[index==group[3]].values.flatten(),
				 	 phi = y.max().values[0]*10) 

if model == 'DBN-0~4':
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Subset the time covariate and the features:
	Z = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	dict_stan = dict(p_z = Z[index==group[0]].shape[1],
					 p_res = len(group),
					 n_0 = Z[index==group[0]].shape[0],
				 	 Z_0 = Z[index==group[0]],
				 	 y_0 = y[index==group[0]].values.flatten(),
					 n_1 = Z[index==group[1]].shape[0],
				 	 Z_1 = Z[index==group[1]],
				 	 y_1 = y[index==group[1]].values.flatten(),
					 n_2 = Z[index==group[2]].shape[0],
				 	 Z_2 = Z[index==group[2]],
				 	 y_2 = y[index==group[2]].values.flatten(),			 	  
					 n_3 = Z[index==group[3]].shape[0],
				 	 Z_3 = Z[index==group[3]],
				 	 y_3 = y[index==group[3]].values.flatten(),
					 n_4 = Z[index==group[4]].shape[0],
				 	 Z_4 = Z[index==group[4]],
				 	 y_4 = y[index==group[4]].values.flatten(),
				 	 phi = y.max().values[0]*10) 

if model == 'DBN-0~5':
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Subset the time covariate and the features:
	Z = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	dict_stan = dict(p_z = Z[index==group[0]].shape[1],
					 p_res = len(group),
					 n_0 = Z[index==group[0]].shape[0],
				 	 Z_0 = Z[index==group[0]],
				 	 y_0 = y[index==group[0]].values.flatten(),
					 n_1 = Z[index==group[1]].shape[0],
				 	 Z_1 = Z[index==group[1]],
				 	 y_1 = y[index==group[1]].values.flatten(),
					 n_2 = Z[index==group[2]].shape[0],
				 	 Z_2 = Z[index==group[2]],
				 	 y_2 = y[index==group[2]].values.flatten(),			 	  
					 n_3 = Z[index==group[3]].shape[0],
				 	 Z_3 = Z[index==group[3]],
				 	 y_3 = y[index==group[3]].values.flatten(),
					 n_4 = Z[index==group[4]].shape[0],
				 	 Z_4 = Z[index==group[4]],
				 	 y_4 = y[index==group[4]].values.flatten(),
					 n_5 = Z[index==group[5]].shape[0],
				 	 Z_5 = Z[index==group[5]],
				 	 y_5 = y[index==group[5]].values.flatten(),
				 	 phi = y.max().values[0]*10) 

if model == 'DBN-0~6':
	# Subset time indexes and groups:
	index = X.iloc[:,0].values
	group = X.iloc[:,0].unique()
	# Subset the time covariate and the features:
	Z = X.drop(X.columns[0], axis=1)
	# Build dictionaries:
	dict_stan = dict(p_z = Z[index==group[0]].shape[1],
					 p_res = len(group),
					 n_0 = Z[index==group[0]].shape[0],
				 	 Z_0 = Z[index==group[0]],
				 	 y_0 = y[index==group[0]].values.flatten(),
					 n_1 = Z[index==group[1]].shape[0],
				 	 Z_1 = Z[index==group[1]],
				 	 y_1 = y[index==group[1]].values.flatten(),
					 n_2 = Z[index==group[2]].shape[0],
				 	 Z_2 = Z[index==group[2]],
				 	 y_2 = y[index==group[2]].values.flatten(),			 	  
					 n_3 = Z[index==group[3]].shape[0],
				 	 Z_3 = Z[index==group[3]],
				 	 y_3 = y[index==group[3]].values.flatten(),
					 n_4 = Z[index==group[4]].shape[0],
				 	 Z_4 = Z[index==group[4]],
				 	 y_4 = y[index==group[4]].values.flatten(),
					 n_5 = Z[index==group[5]].shape[0],
				 	 Z_5 = Z[index==group[5]],
				 	 y_5 = y[index==group[5]].values.flatten(),
					 n_6 = Z[index==group[6]].shape[0],
				 	 Z_6 = Z[index==group[6]],
				 	 y_6 = y[index==group[6]].values.flatten(),
				 	 phi = y.max().values[0]*10) 


#--------------------------------------Running the Bayesian Network------------------------------------------#

# Set directory:
os.chdir(REPO_PATH + "/codes")

# For running the Bayesian Network model:s
if model == 'BN':
	# Compile the Bayesian Network:
	model_stan = ps.StanModel(file='bayesian_network.stan')
	# Create an empty list:
	fit = []
	# Fit the model:
	for t in range(len(group)):
		fit.append(model_stan.sampling(data=dict_stan[t], chains=4, iter=400))

# Run the Pleiotropic Bayesian Network model:
if model == 'PBN':
	# Compile the Pleiotropic Bayesian Network:
	model_stan = ps.StanModel(file='pleiotropic_bayesian_network.stan')
	# Create an empty list:
	fit = []
	# Fit the model:
	if len(group1)==1:
		for t in range(len(group0)):
			fit.append(model_stan.sampling(data=dict_stan[t], chains=4, iter=400))
	if len(group0)==1:
		for t in range(len(group1)):
			fit.append(model_stan.sampling(data=dict_stan[t], chains=4, iter=400))
	if (len(group0)!=1) & (len(group1)!=1):
		for t in range(len(group0)):
			fit.append(model_stan.sampling(data=dict_stan[t], chains=4, iter=400))

# Compile the DBN model:
if model == 'DBN-0~1':
	model_stan = ps.StanModel(file='dynamic_bayesian_network_0_1.stan')
	# Fit the model:
	fit = model_stan.sampling(data=dict_stan, chains=4, iter=400)

# Compile the DBN model:
if model == 'DBN-0~2':
	model_stan = ps.StanModel(file='dynamic_bayesian_network_0_2.stan')
	# Fit the model:
	fit = model_stan.sampling(data=dict_stan, chains=4, iter=400)

# Compile the DBN model:
if model == 'DBN-0~3':
	model_stan = ps.StanModel(file='dynamic_bayesian_network_0_3.stan')
	# Fit the model:
	fit = model_stan.sampling(data=dict_stan, chains=4, iter=400)

# Compile the DBN model:
if model == 'DBN-0~4':
	model_stan = ps.StanModel(file='dynamic_bayesian_network_0_4.stan')
	# Fit the model:
	fit = model_stan.sampling(data=dict_stan, chains=4, iter=400)

# Compile the DBN model:
if model == 'DBN-0~5':
	model_stan = ps.StanModel(file='dynamic_bayesian_network_0_5.stan')
	# Fit the model:
	fit = model_stan.sampling(data=dict_stan, chains=4, iter=400)

# Compile the DBN model:
if model == 'DBN-0~6':
	model_stan = ps.StanModel(file='dynamic_bayesian_network_0_6.stan')
	# Fit the model:
	fit = model_stan.sampling(data=dict_stan, chains=4, iter=400)


#---------------------------------Saving outputs from the Bayesian Network-----------------------------------#

# Set directory:
os.chdir(CV_OUT_PATH)

# Save stan fit object and model for the BN:
if model == 'BN':
	for t in range(len(group)):
		with open('output_' + model.lower() + '_fit_' + str(t) + '.pkl', 'wb') as f:
		    pickle.dump({'model' : model_stan, 'fit' : fit[t]}, f, protocol=-1)

# Save stan fit object and model for the PBN:
if model == 'PBN':
	if len(group1)==1:
		for t in range(len(group0)):
			with open('output_' + model.lower() + '_fit_' + str(t) + '.pkl', 'wb') as f:
			    pickle.dump({'model' : model_stan, 'fit' : fit[t]}, f, protocol=-1)
	if len(group0)==1:
		for t in range(len(group1)):
			with open('output_' + model.lower() + '_fit_' + str(t) + '.pkl', 'wb') as f:
			    pickle.dump({'model' : model_stan, 'fit' : fit[t]}, f, protocol=-1)
	if (len(group0)!=1) & (len(group1)!=1):
		for t in range(len(group0)):
			with open('output_' + model.lower() + '_fit_' + str(t) + '.pkl', 'wb') as f:
			    pickle.dump({'model' : model_stan, 'fit' : fit[t]}, f, protocol=-1)

# Save stan fit object and model for the DBN:
if bool(re.search('DBN', model)):
	with open('output_' + model.lower() + '.pkl', 'wb') as f:
	    pickle.dump({'model' : model_stan, 'fit' : fit}, f, protocol=-1)

<file_sep>/codes/genomic_prediction_results.py

#------------------------------------------------Modules-----------------------------------------------------#

# Load libraries:
import pandas as pd
import numpy as np
import os
import pickle
import argparse
from scipy.stats.stats import pearsonr
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
parser = argparse.ArgumentParser()

# Turn off interactive mode:
plt.ioff()


#-------------------------------------------Add flags to the code--------------------------------------------#

# Get flags:
parser.add_argument("-rpath", "--rpath", dest = "rpath", help="The path of the repository")
parser.add_argument("-opath", "--opath", dest = "opath", help="The path of the folder with general outputs")

# Parse the paths:
args = parser.parse_args()

# Subset arguments:
REPO_PATH = args.rpath
OUT_PATH = args.opath
# REPO_PATH = '/workdir/jp2476/sorghum-multi-trait'
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'


#-----------------------------------------Read train and test data-------------------------------------------#

# Type of forward-chaining cross-validation (fcv) schemes:
fcv_type = ['fcv-30~45', 'fcv-30~60', 'fcv-30~75', 'fcv-30~90', 'fcv-30~105',
      'fcv-30~only', 'fcv-45~only', 'fcv-60~only', 'fcv-75~only', 'fcv-90~only', 'fcv-105~only']

# Create a list of the 5-fold cross-validation (cv5f) folds:
cv5f_fold = ['k0', 'k1', 'k2', 'k3', 'k4']

# Create a list with the traits set:
trait_set = ['drymass', 'height']

# Initialize list to receive the outputs:
y = dict()
X = dict()

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Read cv5f files:
for s in range(len(trait_set)):
  for j in range(len(cv5f_fold)):
    # Create the suffix of the file name for the cv5f data case:
    index = 'cv5f_' + trait_set[s] + '_' + cv5f_fold[j] + '_tst'
    # Read data:
    y[index] = pd.read_csv('y_' + index + '.csv', header = 0, index_col=0)
    X[index] = pd.read_csv('x_' + index + '.csv', header = 0, index_col=0)

# Read fcv files for height:
for t in range(len(fcv_type)):
  for j in range(len(cv5f_fold)):
    # Create the suffix of the file name for the fcv data case:
    index = fcv_type[t] + '_height_tst'
    # Read data:
    y[index] = pd.read_csv('y_' + index + '.csv', header = 0, index_col=0)
    X[index] = pd.read_csv('x_' + index + '.csv', header = 0, index_col=0)


#----------------------------Organize data into a new format for easy subsetting-----------------------------#

# Read adjusted means:
df = pd.read_csv("adjusted_means.csv", index_col=0)
df.dap = df.dap.fillna(0).astype(int)

# Different DAP measures:
dap_group = ['30', '45', '60', '75', '90', '105', '120']  

# Initialize dictionary to receive observations for cv5f analysis:
y_obs_cv5f = dict()

# Store observations for height stratified by modelling scenario for cv5f analysis:
for d in dap_group:
  y_obs_cv5f['cv5f_height_dap:' + d] = df.y_hat[(df.trait=='height') & (df.dap==int(d))] 

# Store observations for drymass stratified by modelling scenario for cv5f analysis:
y_obs_cv5f['cv5f_drymass'] = df.y_hat[df.trait=='drymass'] 

# Initialize dictionary to receive observations for fcv analysis:
y_obs_fcv = dict()

# Store observations for height stratified by modelling scenario for fcv analysis:
dap_group = ['30', '45', '60', '75', '90', '105']  
for d in dap_group:
  y_obs_fcv['fcv_height_for!trained!on:' + d] = y['fcv-' + d + '~only_height_tst']

# Store observations for height stratified by modelling scenario for fcv analysis:
dap_group = ['30~45', '30~60', '30~75', '30~90', '30~105']
for d in dap_group:
  y_obs_fcv['fcv_height_for!trained!on:' + d] = y['fcv-' + d + '_height_tst']


#--------------------------------Compute predictions for the cv5f scheme-------------------------------------#

# Initialize list to receive the predictions:
y_pred_cv5f = dict()

# Different DAP measures:
dap_group = ['30', '45', '60', '75', '90', '105', '120']  

# Compute predictions for the BN model:
for s in trait_set:
  # Compute predictions for drymass:
  if s=='drymass':
    for j in cv5f_fold:
      # Set the directory:
      os.chdir(OUT_PATH + '/cv/BN/cv5f/drymass/' + j)
      # Load stan fit object and model:
      with open("output_bn_fit_0.pkl", "rb") as f:
          data_dict = pickle.load(f)
      # Index the fit object and model
      out = data_dict['fit'].extract()
      # Index and subsetting the feature matrix:
      index1 = 'cv5f_drymass_' + j + '_tst'
      index2 = 'bn_cv5f_drymass'
      X_tmp = X[index1].drop(X[index1].columns[0], axis=1)
      # Create the indexes for creating a data frame to receive predictions:
      index = ["s_" + str(i) for i in range(out['mu'].size)]
      col_name = X_tmp.index
      # Initialize a matrix to receive the posterior predictions:
      tmp = pd.DataFrame(index=index, columns=col_name)
      # Compute predictions:
      for sim in range(out['mu'].size):
        # Subset parameters:
        mu = out['mu'][sim]
        alpha = out['alpha'][sim,:]
        # Prediction:
        tmp.iloc[sim] = (mu + X_tmp.dot(alpha)).values
      # Store prediction:
      if j=='k0':
        y_pred_cv5f[index2] = tmp
      if j!='k0':
        y_pred_cv5f[index2] = pd.concat([y_pred_cv5f[index2], tmp], axis=1)
  # Compute predictions for height:
  if s=='height':
    for d in range(len(dap_group)):
      for j in cv5f_fold:
        # Set the directory:
        os.chdir(OUT_PATH + '/cv/BN/cv5f/height/' + j)
        # Load stan fit object and model:
        with open("output_bn_fit_" + str(d) + ".pkl", "rb") as f:
            data_dict = pickle.load(f)
        # Index the fit object and model
        out = data_dict['fit'].extract()
        # Index and subsetting the feature matrix:
        index1 = 'cv5f_height_' + j + '_tst'
        index2 = 'bn_cv5f_height_trained!on!dap:' + dap_group[d]
        X_tmp = X[index1][X[index1].iloc[:,0] == int(dap_group[d])]
        X_tmp = X_tmp.drop(X_tmp.columns[0], axis=1)
        # Create the indexes for creating a data frame to receive predictions:
        index = ["s_" + str(i) for i in range(out['mu'].size)]
        col_name = X_tmp.index
        # Initialize a matrix to receive the posterior predictions:
        tmp = pd.DataFrame(index=index, columns=col_name)
        # Compute predictions:
        for sim in range(out['mu'].size):
          # Subset parameters:
          mu = out['mu'][sim]
          alpha = out['alpha'][sim,:]
          # Prediction:
          tmp.iloc[sim] = (mu + X_tmp.dot(alpha)).values
        # Store prediction:
        if j=='k0':
          y_pred_cv5f[index2] = tmp
        if j!='k0':
          y_pred_cv5f[index2] = pd.concat([y_pred_cv5f[index2], tmp], axis=1)

# Compute predictions for the PBN model:
for d in range(len(dap_group)):
  for j in cv5f_fold:
    # Set the directory:
    os.chdir(OUT_PATH + '/cv/PBN/cv5f/drymass-height/' + j)
    # Load stan fit object and model:
    with open("output_pbn_fit_" + str(d) + ".pkl", "rb") as f:
        data_dict = pickle.load(f)
    # Index the fit object and model
    out = data_dict['fit'].extract()
    # Index and subset the feature matrix:
    index1_0 = 'cv5f_drymass_' + j + '_tst'
    index1_1 = 'cv5f_height_' + j + '_tst'
    index2_0 = 'pbn_cv5f_drymass_trained!on!dap:' + dap_group[d]
    index2_1 = 'pbn_cv5f_height_trained!on!dap:' + dap_group[d]
    X_tmp_0 = X[index1_0]
    X_tmp_0 = X_tmp_0.drop(X_tmp_0.columns[0], axis=1)
    X_tmp_1 = X[index1_1][X[index1_1].iloc[:,0] == int(dap_group[d])]
    X_tmp_1 = X_tmp_1.drop(X_tmp_1.columns[0], axis=1)
    # Create the indexes for creating a data frame to receive predictions:
    index = ["s_" + str(i) for i in range(out['mu_0'].size)]
    col_name = X_tmp_0.index
    # Initialize a matrix to receive the posterior predictions:
    tmp_0 = pd.DataFrame(index=index, columns=col_name)
    # Create the indexes for creating a data frame to receive predictions:
    index = ["s_" + str(i) for i in range(out['mu_0'].size)]
    col_name = X_tmp_1.index
    # Initialize a matrix to receive the posterior predictions:
    tmp_1 = pd.DataFrame(index=index, columns=col_name)
    # Compute predictions:
    for sim in range(out['mu_0'].size):
      # Subset parameters:
      mu_0 = out['mu_0'][sim]
      mu_1 = out['mu_1'][sim]
      alpha_0 = out['alpha_0'][sim,:]
      alpha_1 = out['alpha_1'][sim,:]
      eta_0 = out['eta_0'][sim,:]
      eta_1 = out['eta_1'][sim,:]
      # Prediction:
      tmp_0.iloc[sim] = (mu_0 + X_tmp_0.dot(alpha_0 + eta_0)).values
      tmp_1.iloc[sim] = (mu_1 + X_tmp_1.dot(alpha_1 + eta_1)).values
    # Store prediction:
    if j=='k0':
      y_pred_cv5f[index2_0] = tmp_0
      y_pred_cv5f[index2_1] = tmp_1
    if j!='k0':
      y_pred_cv5f[index2_0] = pd.concat([y_pred_cv5f[index2_0], tmp_0], axis=1)
      y_pred_cv5f[index2_1] = pd.concat([y_pred_cv5f[index2_1], tmp_1], axis=1)

# List drymass predictions for ensambling:
tmp = []
for i in range(len(dap_group)):
  tmp.append(y_pred_cv5f['pbn_cv5f_drymass_trained!on!dap:' + dap_group[i]])

# Ensamble predictions for drymass:
y_pred_cv5f['pbn_cv5f_drymass_ensambled'] = (tmp[0] + tmp[1] + tmp[2] + tmp[3] + tmp[4] + tmp[5] + tmp[6]) / 7

# Remove predictions that will not be used anymore:
for d in range(len(dap_group)):
  y_pred_cv5f.pop('pbn_cv5f_drymass_trained!on!dap:' + dap_group[d])

# Compute predictions for the DBN model:
for j in cv5f_fold:
  # Set the directory:
  os.chdir(OUT_PATH + '/cv/DBN/cv5f/height/' + j)
  # Load stan fit object and model:
  with open("output_dbn-0~6.pkl", "rb") as f:
      data_dict = pickle.load(f)
  # Index the fit object and model
  out = data_dict['fit'].extract()
  for d in range(len(dap_group)):
    # Index and subset the feature matrix:
    index1 = 'cv5f_height_' + j + '_tst'
    index2 = 'dbn_cv5f_height_trained!on!dap:' + dap_group[d]
    X_tmp = X[index1][X[index1].iloc[:,0] == int(dap_group[d])]
    Z_tmp = X_tmp.drop(X_tmp.columns[0], axis=1)
    # Create the indexes for creating a data frame to receive predictions:
    index = ["s_" + str(i) for i in range(out['mu_' + str(d)].size)]
    col_name = X_tmp.index
    # Initialize a matrix to receive the posterior predictions:
    tmp = pd.DataFrame(index=index, columns=col_name)
    # Compute predictions:
    for sim in range(out['mu_0'].size):
      # Subset parameters:
      mu = out['mu_' + str(d)][sim]
      alpha = out['alpha_' + str(d)][sim,:]
      # Prediction:
      tmp.iloc[sim] = (mu + Z_tmp.dot(alpha)).values
    # Store prediction:
    if j=='k0':
      y_pred_cv5f[index2] = tmp
    if j!='k0':
      y_pred_cv5f[index2] = pd.concat([y_pred_cv5f[index2], tmp], axis=1)


#--------------------------------Computing predictions for the fcv scheme------------------------------------#

# Initialize list to receive the predictions:
y_pred_fcv = dict()

# Initialize list to receive the expectaions:
expect_fcv = dict()

# Different DAP measures:
dap_group = ['30', '45', '60', '75', '90', '105']  

# Compute predictions for the BN model:
for d in range(len(dap_group)):
  # Set the directory:
  os.chdir(OUT_PATH + '/cv/BN/fcv-' + dap_group[d] + '~only/height/')
  # Load stan fit object and model:
  with open("output_bn_fit_0.pkl", "rb") as f:
    data_dict = pickle.load(f)
  # Index the fit object and model
  out = data_dict['fit'].extract()
  # Index and subset the feature matrix:
  index1 = 'fcv-' + dap_group[d] + '~only_height_tst'
  index2 = 'bn_fcv_height_trained!on!dap:' + dap_group[d]
  X_tmp = X[index1].drop(X[index1].columns[0], axis=1)
  # Create the indexes for creating a data frame to receive predictions:
  index = ["s_" + str(i) for i in range(out['mu'].size)]
  col_name = X_tmp.index
  # Initialize a matrix to receive the posterior predictions:
  tmp = pd.DataFrame(index=index, columns=col_name)
  # Compute predictions:
  for sim in range(out['mu'].size):
    # Subset parameters:
    mu = out['mu'][sim]
    alpha = out['alpha'][sim,:]
    # Prediction:
    tmp.iloc[sim] = (mu + X_tmp.dot(alpha)).values
  # Store predictions:
  y_pred_fcv[index2] = tmp
  # Store expectations:
  expect_fcv[index2] = pd.DataFrame(out['expectation'], index=index, columns=df.index[df.dap==int(dap_group[d])])

# Compute predictions for the PBN model:
for d in range(len(dap_group)):
  # Set the directory:
  os.chdir(OUT_PATH + '/cv/PBN/fcv-' + dap_group[d] + '~only/drymass-height/')
  # Load stan fit object and model:
  with open("output_pbn_fit_0.pkl", "rb") as f:
    data_dict = pickle.load(f)
  # Index the fitted object and model
  out = data_dict['fit'].extract()
  # Index and subset the feature matrix:
  index1 = 'fcv-' + dap_group[d] + '~only_height_tst'
  index2 = 'pbn_fcv_height_trained!on!dap:' + dap_group[d]
  X_tmp = X[index1].drop(X[index1].columns[0], axis=1)
  # Create the indexes for creating a data frame to receive predictions:
  index = ["s_" + str(i) for i in range(out['mu_1'].size)]
  col_name = X_tmp.index
  # Initialize a matrix to receive the posterior predictions:
  tmp = pd.DataFrame(index=index, columns=col_name)
  # Compute predictions:
  for sim in range(out['mu_1'].size):
    # Subset parameters:
    mu = out['mu_1'][sim]
    alpha = out['alpha_1'][sim,:]
    eta = out['eta_1'][sim,:]
    # Prediction:
    tmp.iloc[sim] = (mu + X_tmp.dot(alpha + eta)).values
  # Store predictions:
  y_pred_fcv[index2] = tmp
  # Store expectations:
  expect_fcv[index2] = pd.DataFrame(out['expectation_1'], index=index, columns=df.index[df.dap==int(dap_group[d])])

# Different DAP measures:
fcv_type = ['fcv-30~45', 'fcv-30~60', 'fcv-30~75', 'fcv-30~90', 'fcv-30~105']
dap_index = ['0~1', '0~2', '0~3', '0~4', '0~5']

# Compute predictions for the DBN model:
for c in range(len(fcv_type)):
  # Set the directory:
  os.chdir(OUT_PATH + '/cv/DBN/' + fcv_type[c] + '/height/')
  # Load stan fit object and model:
  with open("output_dbn-" + dap_index[c] + ".pkl", "rb") as f:
    data_dict = pickle.load(f)
  # Index the fit object and model
  out = data_dict['fit'].extract()
  # Get the last time point used for training:
  upper = dap_index[c].split('~')[1]
  # Index and subset the feature matrix:
  index1 = fcv_type[c] +'_height_tst'
  index2 = 'dbn_fcv_height_trained!on!dap:' + fcv_type[c].split('-')[1]
  Z_tmp = X[index1].drop(X[index1].columns[0], axis=1)
  # Create the indexes for creating a data frame to receive predictions:
  index = ["s_" + str(i) for i in range(out['mu_0'].size)]
  col_name = Z_tmp.index
  # Initialize a matrix to receive the posterior predictions:
  tmp = pd.DataFrame(index=index, columns=col_name)
  # Compute predictions:
  for sim in range(out['mu_0'].size):
    # Subset parameters:
    mu = out['mu_' + upper][sim]
    alpha = out['alpha_' + upper][sim,:]
    # Prediction:
    tmp.iloc[sim] = (mu + Z_tmp.dot(alpha)).values
  # Prediction:
  y_pred_fcv[index2] = tmp
  # Store expectations:
  expect_fcv[index2] = pd.DataFrame(out['expectation_'  + upper], index=index, columns=df.index[df.dap==int(dap_group[int(upper)])])


#--------------------Compute prediction accuracies for the cv5f scheme (table and figures)--------------------#

# Different models:
models = ['BN', 'PBN', 'DBN', 'MTi-GBLUP', 'MTr-GBLUP']

# Different DAP measures:
dap_group = ['30', '45', '60', '75', '90', '105', '120']  

# Number of Monte Carlo simulations:
n_sim = 800

# Initialize table to store prediction accuracy based on the posterior mean of the predictions from the Bayesian models:
cv5f_table_mean = pd.DataFrame(index = ['DB'] + ['PH_' + str(i) for i in dap_group], columns=models)

# Models for computing predictive accuracy:
models = ['BN', 'PBN', 'DBN']

# Accuracies based on the posterior means of the predictions for cv5f related to the height prediction:
for m in models:
  for d in dap_group:
    index = y_pred_cv5f[m.lower() + '_cv5f_height_trained!on!dap:' + d].columns
    cor_tmp = np.round(pearsonr(y_pred_cv5f[m.lower() + '_cv5f_height_trained!on!dap:' + d].mean(axis=0), y_obs_cv5f['cv5f_height_dap:' + d][index])[0],2)
    cv5f_table_mean.loc['PH_' + str(d), m] = cor_tmp

# Prediction accuracies based on the posterior means of the predictions for cv5f related to dry biomass prediction:
index = y_pred_cv5f['bn_cv5f_drymass'].columns
cv5f_table_mean.loc['DB', 'BN'] = np.round(pearsonr(y_pred_cv5f['bn_cv5f_drymass'].mean(axis=0), y_obs_cv5f['cv5f_drymass'][index])[0],2)

index = y_pred_cv5f['pbn_cv5f_drymass_ensambled'].columns
cv5f_table_mean.loc['DB', 'PBN'] = np.round(pearsonr(y_pred_cv5f['pbn_cv5f_drymass_ensambled'].mean(axis=0), y_obs_cv5f['cv5f_drymass'][index])[0],2)

# Initialize table to store the standard deviation of the prediction accuracies based on the samples of the posterior distribution of the predictions obtained by the Bayesian models:
cv5f_table_std = pd.DataFrame(index = ['DB'] + ['PH_' + str(i) for i in dap_group], columns=models)

# Standard deviation of the prediction accuracies based on the samples of the posterior distribution of the predictions obtained by the Bayesian models for height:
for m in models:
  # Vector to store the predictive accuracy of each posterior draw:
  cor_tmp = pd.DataFrame(index=["s_" + str(i) for i in range(n_sim)], columns=['acc'])
  # To compute the standard deviation of the accuracies in the posterior level:
  for d in dap_group:
    index1 = y_pred_cv5f[m.lower() + '_cv5f_height_trained!on!dap:' + d].columns
    index2 = y_pred_cv5f[m.lower() + '_cv5f_height_trained!on!dap:' + d].index
    for s in index2:
      cor_tmp.loc[s] = pearsonr(y_pred_cv5f[m.lower() + '_cv5f_height_trained!on!dap:' + d].loc[s], y_obs_cv5f['cv5f_height_dap:' + d][index1])[0]
      print('Model: {}, DAP: {}, Simulation: {}'.format(m.upper(), d, s))
    # Printing the standard deviation of the predictive accuracies:
    cv5f_table_std.loc['PH_' + str(d), m] = float(np.round(cor_tmp.std(axis=0),3))

# Different models:
models = ['bn_biomass', 'pbn_biomass']

# Standard deviation of the prediction accuracies based on the samples of the posterior distribution of the predictions obtained by the Bayesian models for dry biomass:
for m in models:
  # Vector to store the predictive accuracy of each posterior draw:
  cor_tmp = pd.DataFrame(index=["s_" + str(i) for i in range(n_sim)], columns=['acc'])
  # To compute the standard deviation of the accuracies in the posterior level:
  for d in dap_group:
    if m=='bn_biomass':
      index1 = y_pred_cv5f['bn_cv5f_drymass'].columns
      index2 = y_pred_cv5f['bn_cv5f_drymass'].index
    if m=='pbn_biomass':
      index1 = y_pred_cv5f['pbn_cv5f_drymass_ensambled'].columns
      index2 = y_pred_cv5f['pbn_cv5f_drymass_ensambled'].index
    for s in index2:
      if m=='bn_biomass':
        cor_tmp.loc[s] = pearsonr(y_pred_cv5f['bn_cv5f_drymass'].loc[s], y_obs_cv5f['cv5f_drymass'][index1])[0]
      if m=='pbn_biomass':
        cor_tmp.loc[s] = pearsonr(y_pred_cv5f['pbn_cv5f_drymass_ensambled'].loc[s], y_obs_cv5f['cv5f_drymass'][index1])[0]
      print('Model: {}, Simulation: {}'.format(m.upper(), s))
    if m=='bn_biomass':
      # Printing the standard deviation of the predictive accuracies:
      cv5f_table_std.loc['DB', 'BN'] = float(np.round(cor_tmp.std(axis=0),3))
    if m=='pbn_biomass':
      # Printing the standard deviation of the predictive accuracies:
      cv5f_table_std.loc['DB', 'PBN'] = float(np.round(cor_tmp.std(axis=0),3))


#-----------------------------Compute prediction accuracies for the fcv scheme--------------------------------#

# Store into a list different DAP values:
dap_group = ['30', '45', '60', '75', '90', '105', '120']

# List of models to use:
model_set = ['bn', 'pbn']
 
# Dictionary to receive the accuracy matrices:
cor_dict = dict()

# Compute correlation for the Baysian Network and Pleiotropic Bayesian Network model under fcv scheme:
for k in model_set:
  # Create an empty correlation matrix:
  cor_tmp = np.empty([len(dap_group)]*2)
  cor_tmp[:] = np.nan
  for i in range(len(dap_group[:-1])):
    # Subset predictions for correlation computation:
    y_pred_tmp = y_pred_fcv[k + '_fcv_height_trained!on!dap:' + dap_group[i]].mean(axis=0)
    y_obs_tmp = y_obs_fcv['fcv_height_for!trained!on:' + dap_group[i]].y_hat
    for j in range(len(dap_group)):
      # Conditional to compute correlation just forward in time:
      if (j>i):
        # Subset indexes for subsetting the data:
        subset = df[df.dap==int(dap_group[j])].index
        # Build correlation matrix for the Bayesian Network model:
        cor_tmp[i, j] = np.round(pearsonr(y_pred_tmp[subset], y_obs_tmp[subset])[0],4)
  # Store the computed correlation matrix for the Bayesian network model under fcv scheme:
  cor_dict['fcv_' + k] = cor_tmp

# Store into a list different DAP values intervals:
dap_group1 = ['30~45', '30~60', '30~75', '30~90', '30~105']
dap_group2 = ['30', '45', '60', '75', '90', '105', '120']

# Create an empty correlation matrix:
cor_tmp = np.empty([len(dap_group)]*2)
cor_tmp[:] = np.nan

# Compute correlation for the Dynamic Bayesian Network model under fcv scheme:
for i in range(len(dap_group1)):
  # Subset predictions for correlation computation:
  y_pred_tmp = y_pred_fcv['dbn_fcv_height_trained!on!dap:' + dap_group1[i]].mean(axis=0)
  y_obs_tmp = y_obs_fcv['fcv_height_for!trained!on:' + dap_group1[i]].y_hat
  for j in range(len(dap_group2)):    
    # Get the upper bound of the interval:
    upper = int(dap_group1[i].split('~')[1])
    # Conditional to compute correlation just forward in time:
    if (int(dap_group2[j])>upper):
      # Subset indexes for subsetting the data:
      subset = df[df.dap==int(dap_group2[j])].index
      # Build correlation matrix for the Bayesian Network model:
      cor_tmp[i, j] = np.round(pearsonr(y_pred_tmp[subset], y_obs_tmp[subset])[0],4)

# Store the computed correlation matrix for the Bayesian network model under fcv scheme:
cor_dict['fcv_dbn'] = cor_tmp

# Builda mask just to filter accuracies common across models:
mask = np.isnan(cor_dict['fcv_dbn'])

# Eliminate values not shared across all models:
cor_dict['fcv_bn'][mask] = np.nan
cor_dict['fcv_pbn'][mask] = np.nan

# Eliminate rows and columns without any correlation value:
cor_dict['fcv_bn'] = cor_dict['fcv_bn'][0:5,2:7]
cor_dict['fcv_pbn'] = cor_dict['fcv_pbn'][0:5,2:7]
cor_dict['fcv_dbn'] = cor_dict['fcv_dbn'][0:5,2:7]

# Print predictive accuracies
print(cor_dict['fcv_bn'])
print(cor_dict['fcv_pbn'])
print(cor_dict['fcv_dbn'])

# Set directory:
os.chdir(REPO_PATH + "/figures")

# List of models to use:
model_set = ['bn', 'pbn', 'dbn']

# Generate accuracy heatmaps:
for i in model_set:
  # Labels for plotting the heatmap for the Pleiotropic Bayesian Network or Bayesian Network:
  if (i=='bn') | (i=='pbn'):
    labels_axis0 = ['DAP 45*', 'DAP 60*', 'DAP 75*', 'DAP 90*', 'DAP 105*']
  # Labels for plotting the heatmap for the Dynamic Bayesian model:
  if i=='dbn':
    labels_axis0 = ['DAP 30:45*', 'DAP 30:60*', 'DAP 30:75*', 'DAP 30:90*', 'DAP 30:105*']
  # Labels for plotting the heatmap:
  labels_axis1 = ['DAP 60', 'DAP 75', 'DAP 90', 'DAP 105', 'DAP 120']
  # Heat map of the adjusted means across traits:
  heat = sns.heatmap(np.flip(np.flip(cor_dict['fcv_' + i],axis=1), axis=0),
             linewidths=0.25,
             cmap='YlOrBr',
             vmin=0.3,
             vmax=1,
             annot=True,
             annot_kws={"size": 18},
             xticklabels=labels_axis0 ,
             yticklabels=labels_axis1)
  heat.set_ylabel('')    
  heat.set_xlabel('')
  plt.xticks(rotation=25)
  plt.yticks(rotation=45)
  plt.savefig("heatplot_fcv_" + i + "_accuracy.pdf", dpi=350)
  plt.savefig("heatplot_fcv_" + i + "_accuracy.png", dpi=350)
  plt.clf()


#------------------------------Compute coincidence index based on lines (CIL)---------------------------------#

# Dictionary to receive CIL:
cil_dict = dict()

# Store into a list different DAP values intervals:
dap_group1 = ['30~45', '30~60', '30~75', '30~90', '30~105']
dap_group2 = ['30', '45', '60', '75', '90', '105', '120']

# Compute CIL for the Dynamic Bayesian network model (best Bayesian model on the fcv):
for i in range(len(dap_group1)):
	# Subset predictions for correlation computation:
	y_pred_tmp = y_pred_fcv['dbn_fcv_height_trained!on!dap:' + dap_group1[i]]
	y_obs_tmp = y_obs_fcv['fcv_height_for!trained!on:' + dap_group1[i]].y_hat
	for j in range(len(dap_group2)):
		# Get the upper bound of the interval:
		upper = int(dap_group1[i].split('~')[1])
		# Conditional to compute correlation just forward in time:
		if (int(dap_group2[j])>upper):
			# Subset indexes for subsetting the data:
			subset = df[df.dap==int(dap_group2[j])].index
			# Get the number of selected individuals for 20% selection intensity:
			n_selected = int(y_obs_tmp[subset].size * 0.2)
			# Build the indexes for computing the CIL:
			top_rank_obs = np.argsort(y_obs_tmp[subset])[::-1][0:n_selected]
			# Vector for storing the indicators:
			ind_vec = pd.DataFrame(index=y_pred_tmp[subset].index, columns=y_pred_tmp[subset].columns)
			# Build CIL matrix for the Bayesian Network model:
			for sim in range(y_pred_tmp.shape[0]):
				# Top predicted lines:
				top_rank_pred = np.argsort(y_pred_tmp[subset].iloc[sim])[::-1][0:n_selected]
				# Top indicator:
				ind_tmp = top_rank_pred.isin(top_rank_obs)
				ind_vec.loc['s_' + str(sim), ind_tmp.index] = ind_tmp
			# Index to store CIL into dictionary:
			index = 'dbn_' + dap_group1[i] + '_' + dap_group2[j]
			# Compute CIL:
			cil_dict[index]=ind_vec.mean(axis=0)
			print('Model: dbn, DAP_i: {}, DAP_j: {}'.format(dap_group1[i], dap_group2[j]))

# Set directory:
os.chdir(REPO_PATH + "/figures")

# DAP groups used for plotting
dap_group = ['45', '60', '75', '90', '105']

# Generating panel plot:
for i in range(len(dap_group)):
	# Subset CIL for plotting:
	cil=cil_dict['dbn_30~' + dap_group[i] + '_120']
	# Remove entries from other time points with nan:
	cil = cil[~np.isnan(cil)]
	# Get the order of the first slice used to train, to plot the order of the other slices:
	if i==0:
		# Get the order decreasing of the CIL:
		order_index = np.argsort(cil)[::-1]
		# Individuais displaying CIL higher then 80%:
		mask = cil.iloc[order_index] > 0.5
	# Subset CIL for plotting:
	p1 = cil.iloc[order_index][mask].plot.barh(color='red', figsize=(10,12))
	p1.set(yticklabels=df.id_gbs[mask.index])
	p1.tick_params(axis='y', labelsize=5)
	p1.tick_params(axis='x', labelsize=12)
	plt.xlabel('Top 20% posterior coincidence index based on lines', fontsize=20)
	plt.ylabel('Sorghum lines', fontsize=20)
	plt.xlim(0.5, 1)
	plt.savefig('cilplot_fcv_' + 'dbn_30~' + dap_group[i] + '_120.pdf', dpi=350)
	plt.savefig('cilplot_fcv_' + 'dbn_30~' + dap_group[i] + '_120.png', dpi=350)
	plt.clf()


#----Compute coincidence index for dry biomass indirect selection using height adjusted means time series-----#

# Store into a list different DAP values:
dap_group = ['30', '45', '60', '75', '90', '105', '120']

# List of models to use:
model_set = ['bn', 'pbn']

# Compute coincidence index for the Bayesian network and Pleiotropic Bayesian Network model:
for k in model_set:
  for i in range(len(dap_group[:-1])):
  	# Subset expectations and observations for coincidence index computation:
  	expect_tmp = expect_fcv[k + '_fcv_height_trained!on!dap:' + dap_group[i]].copy(deep=True)
  	# Get the name of the lines from the entries of the expectation:
  	expect_tmp.columns = df.loc[expect_tmp.columns].id_gbs
  	# Subset data frame:
  	df_tmp = df[df.trait=='drymass']
  	df_tmp.index = df_tmp.id_gbs
  	# Subset observations:
  	y_obs_tmp = df_tmp.y_hat[expect_tmp.columns]
  	# Get the number of selected individuals for 20% selection intensity:
  	n_selected = int(y_obs_tmp.size * 0.2)
  	# Build the indexes for computing the coincidence index:
  	top_rank_obs = np.argsort(y_obs_tmp)[::-1][0:n_selected]
  	# Vector for coincidence indexes:
  	ci_post_tmp = pd.DataFrame(index=expect_tmp.index, columns=['ci'])
  	# Compute the coincidence index:
  	for sim in range(ci_post_tmp.shape[0]):
  		# Build the indexes for computing the coincidence index:
  		top_rank_pred = np.argsort(expect_tmp.loc['s_' + str(sim)])[::-1][0:n_selected]
  		# Compute coincidence index in the top 20% or not: 
  		ci_post_tmp.loc['s_' + str(sim)] = top_rank_pred.isin(top_rank_obs).mean(axis=0)
  	if (k==model_set[0]) & (i==0):
  	  # Store the coincidence index:
  	  ci_post = pd.DataFrame(columns=['post', 'model', 'dap'])
  	  ci_post['model'] = np.repeat(k.upper(), ci_post_tmp.shape[0])
  	  ci_post['dap'] = np.repeat((dap_group[i] + '*'), ci_post_tmp.shape[0])
  	  ci_post['post'] = ci_post_tmp.ci.values
  	else:
  	  # Store the coincidence index:
  	  tmp1 = np.repeat(k.upper(), ci_post_tmp.shape[0])  
  	  tmp2 = np.repeat((dap_group[i] + '*'), ci_post_tmp.shape[0])
  	  tmp = pd.DataFrame({'post': ci_post_tmp.ci.values, 'model': tmp1, 'dap': tmp2})   
  	  ci_post = pd.concat([ci_post, tmp], axis=0)
  	print('Model: {}, DAP_i: {}'.format(k, dap_group[i]))

# Store into a list different DAP values interv
dap_group = ['30~45', '30~60', '30~75', '30~90', '30~105']

# Compute coincidence index for the Dynamic Bayesian network:
for i in range(len(dap_group)):
	# Subset expectations:
	expect_tmp = expect_fcv['dbn_fcv_height_trained!on!dap:' + dap_group[i]].copy(deep=True)
	# Get the name of the lines from the entries of the expectation:
	expect_tmp.columns = df.loc[expect_tmp.columns].id_gbs
	# Sub set data frame:
	df_tmp = df[df.trait=='drymass']
	df_tmp.index = df_tmp.id_gbs
	# Subset observations:
	y_obs_tmp = df_tmp.y_hat[expect_tmp.columns]
	# Get the upper bound of the interval:
	upper = int(dap_group[i].split('~')[1])
	# Get the number of selected individuals for 20% selection intensity:
	n_selected = int(y_obs_tmp.size * 0.2)
	# Build the indexes for computing the coincidence index:
	top_rank_obs = np.argsort(y_obs_tmp)[::-1][0:n_selected]
	# Vector for coincidence indexes:
	ci_post_tmp = pd.DataFrame(index=expect_tmp.index, columns=['ci'])
	# Compute the coincidence index:
	for sim in range(ci_post_tmp.shape[0]):
		# Build the indexes for computing the coincidence index:
		top_rank_pred = np.argsort(expect_tmp.loc['s_' + str(sim)])[::-1][0:n_selected]
		# Compute coincidence index in the top 20% or not: 
		ci_post_tmp.iloc[sim] = top_rank_pred.isin(top_rank_obs).mean(axis=0)
	# Store the coincidence index:
	tmp1 = np.repeat('DBN', ci_post_tmp.shape[0])  
	tmp2 = np.repeat((str(upper) + '*'), ci_post_tmp.shape[0])
	tmp = pd.DataFrame({'post': ci_post_tmp.ci.values, 'model': tmp1, 'dap': tmp2})   
	ci_post = pd.concat([ci_post, tmp], axis=0)
	print('Model: dbn, DAP_i: {}'.format(dap_group[i]))

# Change data type:
ci_post['post'] = ci_post['post'].astype(float)

# Change labels:
ci_post.columns = ['Coincidence index posterior values', 'Model', 'Days after planting']

# Set directory:
os.chdir(REPO_PATH + "/figures")

# Plot coincidence indexes:
plt.figure(figsize=(20,15))
with sns.plotting_context(font_scale=1):
	ax = sns.violinplot(x='Days after planting',
	                    y='Coincidence index posterior values',
	                    data=ci_post,
	                    hue='Model')
	plt.ylim(0.12, 0.4)
	ax.tick_params(labelsize=30)
	plt.xlabel('Days after planting', fontsize=40)
	plt.ylabel(	'Coincidence index posterior values', fontsize=40)
	plt.legend(fontsize='xx-large', title_fontsize=40)
	plt.savefig("ci_plot.pdf", dpi=350)
	plt.savefig("ci_plot.png", dpi=350)
	plt.clf()


#--------------------------Accuracy heatmap for the multivariate linear mixed model--------------------------#

# Set directory:
os.chdir(REPO_PATH + "/tables")

# Initialize a dictionary to receive the accuracies:
gblup_dict = dict()

# Load predictive accuracies of the GBLUP models and store on the final table:
gblup_dict['MTi-GBLUP_cv5f'] = pd.read_csv('acc_MTi-GBLUP_cv5f.csv', header = 0, index_col=0)
gblup_dict['MTr-GBLUP_cv5f'] = pd.read_csv('acc_MTr-GBLUP_cv5f.csv', header = 0, index_col=0)
cv5f_table_mean['MTi-GBLUP'][gblup_dict['MTi-GBLUP_cv5f'].index] = np.round(gblup_dict['MTi-GBLUP_cv5f'].values.flatten(),2)
cv5f_table_mean['MTr-GBLUP'][gblup_dict['MTr-GBLUP_cv5f'].index] = np.round(gblup_dict['MTr-GBLUP_cv5f'].values.flatten(),2)

# Load correlation matrices:
gblup_dict['MTi-GBLUP_fcv'] = pd.read_csv('acc_MTi-GBLUP_fcv.csv', header = 0, index_col=0)
gblup_dict['MTr-GBLUP_fcv'] = pd.read_csv('acc_MTr-GBLUP_fcv.csv', header = 0, index_col=0)

# List of models to use:
model_set = ['MTi-GBLUP', 'MTr-GBLUP']

# Set directory:
os.chdir(REPO_PATH + "/figures")

# Generate accuracy heatmaps:
for i in range(len(model_set)):
  # Labels for plotting the heatmap for the Dynamic Bayesian model:
  labels_axis0 = ['DAP 30:45*', 'DAP 30:60*', 'DAP 30:75*', 'DAP 30:90*', 'DAP 30:105*']
  # Labels for plotting the heatmap:
  labels_axis1 = ['DAP 60', 'DAP 75', 'DAP 90', 'DAP 105', 'DAP 120']
  # Heat map of the adjusted means across traits:
  heat = sns.heatmap(np.round(np.array(gblup_dict[model_set[i] + '_fcv']),4),
             linewidths=0.25,
             cmap='YlOrBr',
             vmin=0.3,
             vmax=1,
             annot=True,
             annot_kws={"size": 18},
             xticklabels=labels_axis0,
             yticklabels=labels_axis1)
  heat.set_ylabel('')    
  heat.set_xlabel('')
  plt.xticks(rotation=25)
  plt.yticks(rotation=45)
  plt.savefig("heatplot_fcv_" + model_set[i] + "_accuracy.pdf", dpi=350)
  plt.savefig("heatplot_fcv_" + model_set[i] + "_accuracy.png", dpi=350)
  plt.clf()

# Set directory:
os.chdir(REPO_PATH + "/tables")

# Save tables with the results from the cv5f:
cv5f_table_mean.to_csv("cv5f_accuracy_table_mean_all_genomic_prediction_models.csv")
cv5f_table_std.to_csv("cv5f_accuracy_table_std_bayesian_genomic_prediction_models.csv")


###############################################################

# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'

# # Set the directory to store processed data:
# os.chdir(OUT_PATH + "/processed_data")

# data = [y, X, y_pred_cv5f, y_pred_fcv, y_obs_cv5f, y_obs_fcv, cil_dict, expect_fcv, df]
# np.savez('mtrait_results.npz', data)

# # Loading data:
# os.chdir(OUT_PATH + "/processed_data")
# container = np.load('mtrait_results.npz', allow_pickle=True)
# data = [container[key] for key in container]
# y = data[0][0]
# X = data[0][1]
# y_pred_cv5f = data[0][2]
# y_pred_fcv = data[0][3]
# y_obs_cv5f = data[0][4]
# y_obs_fcv = data[0][5]
# cil_dict = data[0][6]
# expect_fcv = data[0][7]
# df = data[0][8]


<file_sep>/codes/create_list_files_names_cross-validation.sh
#!/bin/bash

# Reading the flag:
while getopts ":o:" opt; do
  case $opt in
    o)
      OUT_PATH=$OPTARG;
      ;;
    \?)
      echo "Invalid option: -$OPTARG" >&2
      exit 1
      ;;
    :)
      echo "Option -$OPTARG requires an argument." >&2
      exit 1
      ;;
  esac
done

# Set directory where the data is:
cd ${OUT_PATH}/processed_data

# List training data files names related to the 5-fold cross-validation scheme (cv5f) scheme and storing it for latter usage:
ls y*cv5f*trn* > y_cv5f_bn_trn_files.txt
ls x*cv5f*trn* > x_cv5f_bn_trn_files.txt
ls y*cv5f*height*trn* > y_cv5f_dbn_trn_files.txt
ls x*cv5f*height*trn* > x_cv5f_dbn_trn_files.txt

# List training files names related to the cv5f scheme and storing it for latter usage on PBN model
grep "drymass" y_cv5f_bn_trn_files.txt > tmp1.txt 
grep "height" y_cv5f_bn_trn_files.txt > tmp2.txt 
paste -d'&' tmp1.txt tmp2.txt > y_cv5f_pbn_trn_files.txt
grep "drymass" x_cv5f_bn_trn_files.txt > tmp1.txt 
grep "height" x_cv5f_bn_trn_files.txt > tmp2.txt 
paste -d'&' tmp1.txt tmp2.txt > x_cv5f_pbn_trn_files.txt
rm tmp1.txt tmp2.txt

# List training phenotypic data files names related to the forward-chaining cross-validation (fcv) scheme and storing it for latter usage (BN):
ls y*fcv*only*trn* > y_fcv_bn_trn_files.txt
ls x*fcv*only*trn* > x_fcv_bn_trn_files.txt

# List training phenotypic data files names related to the fcv scheme and storing it for latter usage (PBN):
echo 'y_fcv_drymass_trn.csv' > tmp1.txt
echo 'x_fcv_drymass_trn.csv' > tmp2.txt
for i in $(seq 1 5); do 

  echo 'y_fcv_drymass_trn.csv' >> tmp1.txt
  echo 'x_fcv_drymass_trn.csv' >> tmp2.txt

done;
paste -d'&' tmp1.txt y_fcv_bn_trn_files.txt > y_fcv_pbn_trn_files.txt
paste -d'&' tmp2.txt x_fcv_bn_trn_files.txt > x_fcv_pbn_trn_files.txt
rm tmp1.txt tmp2.txt

# List training data files names related to the fcv scheme and storing it for latter usage (DBN):
ls y*fcv*height*trn* > y_fcv_dbn_trn_files.txt
ls x*fcv*height*trn* > x_fcv_dbn_trn_files.txt
sed -i '/only/d' y_fcv_dbn_trn_files.txt
sed -i '/only/d' x_fcv_dbn_trn_files.txt

# Create a text file to store the different types of Dynamic Bayesian network models for latter usage (DBN);
echo "DBN-0~5" > dbn_models_fcv_list.txt
echo "DBN-0~1" >> dbn_models_fcv_list.txt
echo "DBN-0~2" >> dbn_models_fcv_list.txt
echo "DBN-0~3" >> dbn_models_fcv_list.txt
echo "DBN-0~4" >> dbn_models_fcv_list.txt

# List of files to get train data index into R to perform Multivariate GBLUP models analysis in R:
cat y_cv5f_dbn_trn_files.txt > cv5f_mti-gblup_files.txt
cat y_cv5f_pbn_trn_files.txt > cv5f_mtr-gblup_files.txt
cat y_fcv_dbn_trn_files.txt > fcv_mti-gblup_files.txt
echo 'y_fcv_drymass_trn.csv' > tmp1.txt
for i in $(seq 1 4); do 

  echo 'y_fcv_drymass_trn.csv' >> tmp1.txt

done;
paste -d'&' tmp1.txt y_fcv_dbn_trn_files.txt > fcv_mtr-gblup_files.txt
rm tmp1.txt

# Create a text file to store the different types of Multi Time GBLUP (MTi-GBLUP) models for latter usage;
echo "MTi-GBLUP-0~5" > mti-gblup_models_fcv_list.txt
echo "MTi-GBLUP-0~1" >> mti-gblup_models_fcv_list.txt
echo "MTi-GBLUP-0~2" >> mti-gblup_models_fcv_list.txt
echo "MTi-GBLUP-0~3" >> mti-gblup_models_fcv_list.txt
echo "MTi-GBLUP-0~4" >> mti-gblup_models_fcv_list.txt

# Create a text file to store the different types of Multi Trait GBLUP (MTr-GBLUP) models for latter usage;
echo "MTr-GBLUP-0~5" > mtr-gblup_models_fcv_list.txt
echo "MTr-GBLUP-0~1" >> mtr-gblup_models_fcv_list.txt
echo "MTr-GBLUP-0~2" >> mtr-gblup_models_fcv_list.txt
echo "MTr-GBLUP-0~3" >> mtr-gblup_models_fcv_list.txt
echo "MTr-GBLUP-0~4" >> mtr-gblup_models_fcv_list.txt

<file_sep>/codes/process_emmreml_output.R

#-------------------------------------------Loading libraries------------------------------------------------#

# Libraries to manipulate data:
library(optparse)
library(data.table)
library(stringr)
library(magrittr)
library(EMMREML)
library(tidyr)


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Set the flags:
option_list = list(
  make_option(c("-r", "--rpath"), type="character", default=NULL, 
              help="The path of the repository", metavar="character"),
  make_option(c("-o", "--opath"), type="character", default=NULL, 
              help="The path of the folder to receive outputs", metavar="character")
) 

# Parse the arguments:
opt_parser = OptionParser(option_list=option_list)
args = parse_args(opt_parser)

# Subset arguments:
REPO_PATH = args$rpath
OUT_PATH = args$opath
# REPO_PATH = '/workdir/jp2476/sorghum-multi-trait'
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'


#--------------Load outputs, predict and compute accuracies from 5-fold cross-validation (cv5f)-------------------#

# Types of cross-validation:
cv_type = paste0('k', 0:4)

# Initialize list:
y_pred_lst = list()

# Create a variable to store part of the directory:
PARTIAL_PATH = '/cv/MTi-GBLUP/cv5f/height'

# Read cv5f outputs from the Multiple Time Linear Mixed model:
for (i in 1:length(cv_type)) {

  # Set the directory:
  setwd(paste0(OUT_PATH, PARTIAL_PATH , '/', cv_type[i]))

  # Load the entire workspace:
  load("output_mixed_models.RData")

  # Naming the columns of the genomic estimated breeding values output:
  colnames(out$Gpred) = as.character(df_melt$id_gbs)
  rownames(out$Gpred) = rownames(Y)

  # Broadcasting the mean:
  mu = matrix(rep(out$Bhat,each=ncol(out$Gpred[,id_tst])), ncol=ncol(out$Gpred[,id_tst]), byrow=TRUE)

  # Subset predictions:
  y_pred_tmp = t(mu) + t(out$Gpred[,id_tst])

  if (i==1) {

    # Subset predictions:
    y_pred_lst[['MTi-GBLUP']][['cv5f']] = y_pred_tmp

  }
  else {

    # Subset predictions and update the prediction matrix:
    y_pred_lst[['MTi-GBLUP']][['cv5f']] = rbind(y_pred_lst[['MTi-GBLUP']][['cv5f']], y_pred_tmp)

  }

  print(paste0("We are at the k: ", i))

}

# Initialize a list:
acc_lst = list()

# Saving predictive accuracy of the cv5f for the MTi-GBLUP:
index = df_melt$id_gbs %>% droplevels %>% as.character
acc_lst[['MTi-GBLUP']][['cv5f']] = diag(cor(y_pred_lst[['MTi-GBLUP']][['cv5f']][index,], t(Y)[index,]))

# Create a variable to store part of the directory:
PARTIAL_PATH = '/cv/MTr-GBLUP/cv5f/drymass-height/'

# Read cv5f outputs from the Multiple Time Linear Mixed model:
for (i in 1:length(cv_type)) {

  # Set the directory:
  setwd(paste0(OUT_PATH, PARTIAL_PATH , '/', cv_type[i]))

  # Load the entire workspace:
  load("output_mixed_models.RData")

  # Naming the columns of the genomic estimated breeding values output:
  colnames(out$Gpred) = as.character(df_melt$id_gbs)
  rownames(out$Gpred) = rownames(Y)

  # Broadcasting the mean:
  mu = matrix(rep(out$Bhat,each=ncol(out$Gpred[,id_tst])), ncol=ncol(out$Gpred[,id_tst]), byrow=TRUE)

  # Subset predictions:
  y_pred_tmp = t(mu) + t(out$Gpred[,id_tst])

  if (i==1) {

    # Subset predictions:
    y_pred_lst[['MTr-GBLUP']][['cv5f']] = y_pred_tmp

  }
  else {

    # Subset predictions and update the prediction matrix:
    y_pred_lst[['MTr-GBLUP']][['cv5f']] = rbind(y_pred_lst[['MTr-GBLUP']][['cv5f']], y_pred_tmp)

  }

  print(paste0("We are at the k: ", i))

}

# Saving predictive accuracy of the cv5f for the Mr-GBLUP:
index = df_melt$id_gbs %>% droplevels %>% as.character
acc_lst[['MTr-GBLUP']][['cv5f']] = diag(cor(y_pred_lst[['MTr-GBLUP']][['cv5f']][index,], t(Y)[index,]))


#----------Load outputs, predict and compute accuracies from forward chaining cross-validation (fcv)--------------#

# Create a variable to store part of the directory:
PARTIAL_PATH = '/cv/MTi-GBLUP'

# Types of cross-validation:
cv_type = paste0('fcv-30~', seq(45,105,15))

# Time points to correlate:
dap = seq(30,120,15)

# Read CV1 outputs from the Multiple Time Linear Mixed model:
for (i in 1:length(cv_type)) {

  # Set the directory:
  setwd(paste0(OUT_PATH, PARTIAL_PATH, '/', cv_type[i], '/height/'))

  # Load the entire workspace:
  load("output_mixed_models.RData")

  # Naming the columns of the genomic estimated breeding values output:
  colnames(out$Gpred) = as.character(df_melt$id_gbs)
  rownames(out$Gpred) = rownames(Y)

  # Get the upper bound of the time inverval used for training:
  upper = str_split(cv_type[i], '~', simplify = TRUE)[1,2] %>% as.character

  # Get a mask of the time points only used for training:
  mask = str_detect(rownames(out$Gpred), upper)

  # Broadcast the mean and subset just the one related to the last time point:
  mu = matrix(rep(out$Bhat[mask,],each=ncol(out$Gpred)), ncol=ncol(out$Gpred), byrow=TRUE) %>% data.matrix

  # Compute predictions:
  y_pred_lst[['MTi-GBLUP']][[cv_type[i]]] = mu + t(out$Gpred[mask,])

  # Initialize a matrix to receive the accuracies:
  if (i == 1) {

    acc_matrix = matrix(NA, length(dap), length(dap))
    rownames(acc_matrix) = dap
    colnames(acc_matrix) = dap

  }

  # Intervals to be subset to compute accuracies
  begin = which(dap==upper)+1
  end = length(dap)

  # Compute accuracies:
  for (j in dap[begin:end]){
      
      # Subset observde data:
      y_tmp = df[(df$trait=='height') & (df$dap==j),]$y_hat
      names(y_tmp) = df[(df$trait=='height') & (df$dap==j),]$id

      # Get inbred line names to correlate in the right pair:
      index = colnames(y_pred_lst[['MTi-GBLUP']][[cv_type[i]]])

      # Compute predictive accuraces:
      acc_matrix[as.character(j), upper] = cor(c(y_pred_lst[['MTi-GBLUP']][[cv_type[i]]]), c(y_tmp[index]))

  }

  print(paste0("We are at the cross-validation: ", cv_type[i]))

}

# Saving predictive accuracy of the fcv for the MTi-GBLUP:
acc_lst[['MTi-GBLUP']][['fcv']] = acc_matrix[3:nrow(acc_matrix), 2:(nrow(acc_matrix)-1)]

# Create a variable to store part of the directory:
PARTIAL_PATH = '/cv/MTr-GBLUP'

# Types of cross-validation:
cv_type = paste0('fcv-30~', seq(45,105,15))

# Time points to correlate:
dap = seq(30,120,15)

# Read CV1 outputs from the Multiple Time Linear Mixed model:
for (i in 1:length(cv_type)) {

  # Set the directory:
  setwd(paste0(OUT_PATH, PARTIAL_PATH, '/', cv_type[i], '/drymass-height/'))

  # Load the entire workspace:
  load("output_mixed_models.RData")

  # Naming the columns of the genomic estimated breeding values output:
  colnames(out$Gpred) = as.character(df_melt$id_gbs)
  rownames(out$Gpred) = rownames(Y)

  # Get the upper bound of the time inverval used for training:
  upper = str_split(cv_type[i], '~', simplify = TRUE)[1,2] %>% as.character

  # Get a mask of the time points only used for training:
  mask = str_detect(rownames(out$Gpred), upper)

  # Broadcast the mean and subset just the one related to the last time point:
  mu = matrix(rep(out$Bhat[mask,],each=ncol(out$Gpred)), ncol=ncol(out$Gpred), byrow=TRUE) %>% data.matrix

  # Compute predictions:
  y_pred_lst[['MTr-GBLUP']][[cv_type[i]]] = mu + t(out$Gpred[mask,])

  # Initialize a matrix to receive the accuracies:
  if (i == 1) {

    acc_matrix = matrix(NA, length(dap), length(dap))
    rownames(acc_matrix) = dap
    colnames(acc_matrix) = dap

  }

  # Intervals to be subset to compute accuracies
  begin = which(dap==upper)+1
  end = length(dap)

  # Compute accuracies:
  for (j in dap[begin:end]){

      # Subset observde data:
      y_tmp = df[(df$trait=='height') & (df$dap==j),]$y_hat
      names(y_tmp) = df[(df$trait=='height') & (df$dap==j),]$id

      # Subset of observed inbred lines for both drymass and height:
      subset = names(y_tmp) %in% colnames(y_pred_lst[['MTr-GBLUP']][[cv_type[i]]])
      y_tmp = y_tmp[subset]

      # Get inbred line names to correlate in the right pair:
      index = colnames(y_pred_lst[['MTr-GBLUP']][[cv_type[i]]])

      # Compute predictive accuraces:
      acc_matrix[as.character(j), upper] = cor(c(y_pred_lst[['MTr-GBLUP']][[cv_type[i]]]), c(y_tmp[index]))

  }

  print(paste0("We are at the cross-validation: ", cv_type[i]))

}

# Saving predictive accuracy of the fcv for the MTi-GBLUP:
acc_lst[['MTr-GBLUP']][['fcv']] = acc_matrix[3:nrow(acc_matrix), 2:(nrow(acc_matrix)-1)]

# Set of models to compute the coincidence index:
model_set = c('MTi-GBLUP', 'MTr-GBLUP')

# Set of upper time points:
upper_set = seq(45,105,15)

# Create a vector to store the coincidence indexes:
ci_tmp = rep(NA, length(upper_set))
names(ci_tmp) = paste0('fcv-30~', upper_set)

# Initialize a list to store results:
ci_lst = list()

for (m in model_set) {
  for (i in upper_set) {

    # Subset the expectation of plant height for the i^th DAP:
    yhat_height = y_pred_lst[[m]][[paste0('fcv-30~', i)]]
    tmp = yhat_height %>% colnames
    yhat_height = c(yhat_height)
    names(yhat_height) = tmp

    # Subset adjusted mean for drymass related:
    yhat_drymass = df[df$trait=='drymass',]$y_hat
    names(yhat_drymass) = df[df$trait=='drymass',]$id %>% as.character

    # Number of individuals selected to order genotypes:
    n_selected = as.integer(length(yhat_drymass) * 0.2)

    # Get the index of the top observed selected inbred lines for dry mass:
    top_lines_obs = yhat_drymass[order(yhat_drymass, decreasing=TRUE)][1:n_selected] %>% names

    # Get the index of the predicted to be best selected inbred lines for height:
    top_lines_pred = yhat_height[order(yhat_height, decreasing=TRUE)][1:n_selected] %>% names

    # Compute the coincidence index:
    ci_tmp[paste0('fcv-30~', i)] = mean(top_lines_pred %in% top_lines_obs)

  }

  # Store the coincidence indexes in a list:
  ci_lst[[m]] = ci_tmp

}

# Rename entries of the cv5f accuracy:
names(acc_lst[['MTi-GBLUP']][['cv5f']]) = paste0('PH_', seq(30,120,15))
names(acc_lst[['MTr-GBLUP']][['cv5f']]) = c('DB', paste0('PH_', seq(30,120,15)))

# Set directory:
setwd(paste0(REPO_PATH, "/tables"))

# Write accuracy files:
for (i in names(acc_lst[["MTi-GBLUP"]])) write.csv(acc_lst[["MTi-GBLUP"]][[i]], file=paste0("acc_MTi-GBLUP_", i, ".csv"))
for (i in names(acc_lst[["MTr-GBLUP"]])) write.csv(acc_lst[["MTr-GBLUP"]][[i]], file=paste0("acc_MTr-GBLUP_", i, ".csv"))

# Write coincidence indexes:
for (i in names(ci_lst)) write.csv(ci_lst[[i]], file=paste0('coincidence_index_', i, '.csv'))
<file_sep>/README.md
#  **Novel Bayesian Networks for Genomic Prediction of Developmental Traits in Biomass Sorghum**
![language: R](https://img.shields.io/badge/language-R-blue.svg)
![language: Python](https://img.shields.io/badge/language-Python-green.svg)
![status: WIP](https://img.shields.io/badge/status-WorkInProgress-red.svg)

<p align="center"><a href="https://github.com/GoreLab/sorghum-multi-trait/blob/master/edited_figures/figure_3_forward_cross_validation_figures/inkscape/heatplot_fcv_accuracy.png"><img src="https://github.com/GoreLab/sorghum-multi-trait/blob/master/edited_figures/figure_3_forward_cross_validation_figures/inkscape/heatplot_fcv_accuracy.png" width="555" height="328"/></a>

# **Abstract**

The ability to connect genetic information between traits over time allow Bayesian networks to offer a powerful probabilistic framework to construct genomic prediction models. In this study, we phenotyped a diversity panel of 869 biomass sorghum ($Sorghum$ $bicolor$ (L.) Moench] lines, which had been genotyped with 100,435 SNP markers, for plant height (PH) with biweekly measurements from 30 to 120 days after planting (DAP) and for end-of-season dry biomass yield (DBY) in four environments. We evaluated five genomic prediction models: Bayesian network (BN), Pleiotropic Bayesian network (PBN), Dynamic Bayesian network (DBN), multi-trait GBLUP (MTr-GBLUP), and multi-time GBLUP (MTi-GBLUP) models. In 5-fold cross-validation, prediction accuracies ranged from 0.48 (PBN) to 0.51 (MTr-GBLUP) for DBY and from 0.47 (DBN, DAP120) to 0.74 (MTi-GBLUP, DAP60) for PH. Forward-chaining cross-validation further improved prediction accuracies (36.4-52.4\%) of the DBN, MTi-GBLUP and MTr-GBLUP models for PH (training slice: 30-45 DAP) relative to the BN and PBN models. Coincidence indices (target: biomass, secondary: PH) and a coincidence index based on lines (PH time series) showed that the ranking of lines by PH changed minimally after 45 DAP. These results suggest a two-level indirect selection method for PH at harvest (first-level target trait) and DBY (second-level target trait) could be conducted earlier in the season based on ranking of lines by PH at 45 DAP (secondary trait). With the advance of high-throughput phenotyping technologies, our proposed two-level indirect selection framework could be valuable for enhancing genetic gain per unit of time when selecting on developmental traits.

# **Guidelines**

Instructions to reproduce the project, codes description, and data availability on the file `README.txt`.
<file_sep>/codes/create_directories.sh
#!/bin/bash

# Reading the flag:
while getopts ":p:" opt; do
  case $opt in
    p)
      ROOT_PATH=$OPTARG;
      ;;
    \?)
      echo "Invalid option: -$OPTARG" >&2
      exit 1
      ;;
    :)
      echo "Option -$OPTARG requires an argument." >&2
      exit 1
      ;;
  esac
done

# Create the folder for receiving the outputs of the codes:
mkdir ${ROOT_PATH}/output_sorghum-multi-trait;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/processed_data;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/heritabilities;

# Create CV directories on Bayesian network (BN) output folder:
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/height/k0;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/height/k1;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/height/k2;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/height/k3;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/height/k4;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/drymass;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/drymass/k0;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/drymass/k1;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/drymass/k2;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/drymass/k3;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/cv5f/drymass/k4;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-30~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-30~only/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-45~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-45~only/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-60~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-60~only/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-75~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-75~only/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-90~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-90~only/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-105~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/BN/fcv-105~only/height;

# Create CV directories on Pleiotropic Bayesian network (PBN) output folder:
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f/drymass-height/k0;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f/drymass-height/k1;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f/drymass-height/k2;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f/drymass-height/k3;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/cv5f/drymass-height/k4;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-30~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-30~only/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-45~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-45~only/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-60~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-60~only/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-75~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-75~only/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-90~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-90~only/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-105~only;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/PBN/fcv-105~only/drymass-height;

# Create CV directories on Dynamic Bayesian network (DBN) output folder:
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f/height/k0;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f/height/k1;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f/height/k2;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f/height/k3;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/cv5f/height/k4;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~45;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~45/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~60;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~60/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~75;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~75/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~90;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~90/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~105;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/DBN/fcv-30~105/height;

# Create CV directories for the Multiple Time Linear Mixed model (MTi-GBLUP) output folder:
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f/height/k0;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f/height/k1;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f/height/k2;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f/height/k3;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/cv5f/height/k4;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~45;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~45/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~60;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~60/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~75;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~75/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~90;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~90/height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~105;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~105/height;

# Create CV directories for the Multiple Trait Linear Mixed model (MTr-GBLUP) output folder:
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f/drymass-height/k0;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f/drymass-height/k1;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f/drymass-height/k2;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f/drymass-height/k3;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/cv5f/drymass-height/k4;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~45;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~45/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~60;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~60/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~75;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~75/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~90;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~90/drymass-height;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~105;
mkdir ${ROOT_PATH}/output_sorghum-multi-trait/cv/MTr-GBLUP/fcv-30~105/drymass-height;

<file_sep>/codes/create_figures_phenotypic_data_analysis.py

#------------------------------------------------Modules-----------------------------------------------------#

# Import python modules
import os
import numpy as np
import pandas as pd
import matplotlib
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
parser = argparse.ArgumentParser()

# Turn off interactive mode:
plt.ioff()

#-----------------------------------------Adding flags to the code-------------------------------------------#

# Get flags:
parser.add_argument("-rpath", "--rpath", dest = "rpath", help="The path of the repository")
parser.add_argument("-opath", "--opath", dest = "opath", help="The path of the folder to receive outputs")

# Parse the paths:
args = parser.parse_args()

# Subset arguments:
REPO_PATH = args.rpath
OUT_PATH = args.opath
# REPO_PATH = '/workdir/jp2476/sorghum-multi-trait'
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'


#-----------------------------------------------Load data----------------------------------------------------#

# Set directory:
os.chdir(REPO_PATH + "/clean_repository/codes")

# Import functions:
from functions import * 

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Read adjusted means:
df = pd.read_csv("adjusted_means.csv", index_col=0)

# Change class of the dap:
df.dap = df.dap.fillna(0).astype(int)


#--------------------------------------------Generate Figures------------------------------------------------#

# Set directory:
os.chdir(REPO_PATH + "/clean_repository/figures")

# Add a new column in the data frame for plotting:
tmp = df.pivot(index='id_gbs', columns='dap', values='y_hat')

# Labels for plotting the heatmap:
labels = ["Dry biomass",
  		  "Height DAP 30",
		  "Height DAP 45",
		  "Height DAP 60",
		  "Height DAP 75",
		  "Height DAP 90",
		  "Height DAP 105",
		  "Height DAP 120"]

# Heat map of the adjusted means across traits:
heat = sns.heatmap(tmp.corr(),
         linewidths=0.25,
         cmap='YlOrBr',
         vmin=0.09,
         vmax=1,
         annot=True,
         annot_kws={"size": 12},
         xticklabels=labels ,
         yticklabels=labels)
heat.set_ylabel('')    
heat.set_xlabel('')
heat.tick_params(labelsize=7.6)
plt.xticks(rotation=25)
plt.yticks(rotation=45)
plt.savefig("heatplot_traits_adjusted_means.pdf", dpi=350)
plt.savefig("heatplot_traits_adjusted_means.png", dpi=350)
plt.clf()

# Density plot of the adjusted means from dry mass:
den_dm = sns.kdeplot(df.y_hat[df.trait=="drymass"], bw=1, shade=True, legend=False)
den_dm.set_ylabel('Density')    
den_dm.set_xlabel('Biomass (t/ha)')
den_dm.get_lines()[0].set_color('#006d2c')
x = den_dm.get_lines()[0].get_data()[0]
y = den_dm.get_lines()[0].get_data()[1]
plt.fill_between(x,y, color='#006d2c').set_alpha(.25)
plt.savefig("denplot_drymass_adjusted_means.pdf", dpi=350)
plt.savefig("denplot_drymass_adjusted_means.png", dpi=350)
plt.clf()

# Box plot of the adjusted means from height measures:
box_ph = sns.boxplot(x='dap', y='y_hat',
					 data=df[df.trait=="height"])
box_ph.set_ylabel('Height (cm)')    
box_ph.set_xlabel('Days after Planting')
colors = ['#fee391', '#fec44f', '#fe9929', '#ec7014', '#cc4c02', '#993404', '#662506']

for i in range(1,len(colors)): 
	box_ph.artists[i].set_facecolor(colors[i])

plt.savefig("boxplot_height_adjusted_means.pdf", dpi=350)
plt.savefig("boxplot_height_adjusted_means.png", dpi=350)
plt.clf()

# Set directory:
os.chdir(OUT_PATH + "/heritabilities")

# Read heritability values:
h2_table = pd.read_csv('heritabilities.csv', index_col=0)

# Add new labels to the h2_table for plotting:
h2_table['trait'] = labels
h2_table

# Add colors to be ploted:
h2_table['colors'] = ['#006d2c', '#fe9929', '#ec7014', '#cc4c02', '#993404', '#662506', '#fee391', '#fec44f']

# Reset index:
h2_table = h2_table.reset_index(drop=True)

# Set directory:
os.chdir(REPO_PATH + "/clean_repository/figures")

# Plot heritabilities:
bar_obj=plt.bar(h2_table['trait'].tolist(), h2_table['h2'].tolist(),
 			    yerr = h2_table['se'].tolist(),
 			    align='center',
 			    alpha=1,
 			    color= h2_table['colors'].tolist()
 			    )
plt.tick_params(labelsize=7.6)
plt.xticks(h2_table['trait'].tolist())
plt.xticks(rotation=25)
plt.xlabel('Traits')
plt.ylabel('Broad-sense Heritability')
plt.savefig("barplot_heritabilities.pdf", dpi=350)
plt.savefig("barplot_heritabilities.png", dpi=350)
plt.clf()



<file_sep>/codes/phenotypic_data_analysis.R

#-------------------------------------------Loading libraries------------------------------------------------#

# Import python libraries:
library(optparse)
library(data.table)
library(magrittr)
library(stringr)


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Set the flags:
option_list = list(
  make_option(c("-a", "--asremlpath"), type="character", default=NULL, 
              help="The path of the folder with asreml license", metavar="character"),
  make_option(c("-o", "--opath"), type="character", default=NULL, 
              help="The path of the folder to receive outputs", metavar="character")
) 

# Parse the arguments:
opt_parser = OptionParser(option_list=option_list)
args = parse_args(opt_parser)

# Subset arguments:
OUT_PATH = args$opath
ASREML_PATH = args$asremlpath
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'
# OUT_PATH = '/home/jhonathan/Documents/output_sorghum-multi-trait'
# ASREML_PATH = '/workdir/jp2476/asreml'

#---------------------------------------------Load asreml----------------------------------------------------#

# Load asreml:
setwd(ASREML_PATH)
library(asreml)
asreml.lic(license = "asreml.lic", install = TRUE)

#---------------------------------------------Loading data---------------------------------------------------#

# Set the directory:
setwd(paste0(OUT_PATH, '/processed_data'))

# Load the data:
df = fread('df.csv', header=T)[,-1]

# Change classes:
df$id_gbs = df$id_gbs %>% as.factor
df$name2 = df$name2 %>% as.factor
df$block = df$block %>% as.factor
df$loc = df$loc %>% as.factor
df$year = df$year %>% as.factor
df$trait = df$trait %>% as.factor
df$dap = df$dap %>% as.integer
df$drymass = df$drymass %>% as.numeric
df$height = df$height %>% as.numeric

# Get the group levels of the day after plantting (DAP) measures:
dap_groups = df$dap %>% unique %>% na.omit %>% as.integer

# Create a new column combing location and years:
df$env <- paste0(as.character(df$loc), "_",as.character(df$year))
df$env <- df$env %>% as.factor()


#------------------------------------------Biomass data analysis---------------------------------------------#

# Initialize a list to receive the first stage results:
fit = list()
g = list()

# Index for mapping the results:
index = "drymass"

# Subset part of the data set: 
df_tmp = df[df$trait == "DM"]

# Drop levels not present in the data subset:
df_tmp$name2 = df_tmp$name2 %>% droplevels
df_tmp$block = df_tmp$block %>% droplevels
df_tmp$env = df_tmp$env %>% droplevels

fit[[index]] = asreml(drymass~ 1 + name2,
			 		  random = ~ block:env + env + name2:env,
			 		  na.method.Y = "include",
			 		  control = asreml.control(
			 		  maxiter = 200),
             		  data = df_tmp)


# Extract fixed effects:
fixed_eff = summary(fit[[index]], all=T)$coef.fi[,1]

# Store the corrected mean:
g[[index]] = fixed_eff[length(fixed_eff)] + fixed_eff[-length(fixed_eff)]


#------------------------------------------Height data analysis----------------------------------------------#

for (i in dap_groups) {

	# Index for mapping the results:
	index = "height_" %>% paste0(i)

	# Subsetting part of the data set: 
	df_tmp = df[(df$trait == "PH" & df$dap == i)]  

	# Dropping levels not present in the data subset:
	df_tmp$name2 = df_tmp$name2 %>% droplevels
	df_tmp$block = df_tmp$block %>% droplevels
	df_tmp$env = df_tmp$env %>% droplevels

	fit[[index]] = asreml(height~ 1 + name2,
				 		  random = ~ block:env + env + name2:env,
				 		  na.method.Y = "include",
				 		  control = asreml.control(
				 		  maxiter = 200),
	             		  data = df_tmp)


	# Extract fixed effects:
	fixed_eff = summary(fit[[index]], all=T)$coef.fi[,1]

	# Store the corrected mean:
	g[[index]] = fixed_eff[length(fixed_eff)] + fixed_eff[-length(fixed_eff)]

	# Printing current analysis:
	print(paste0("Analysis ", index, " done!"))

}

# Initialize a new data frame:
df_tmp = data.frame()

# Transform list into data frame:
for (i in names(g)) {

	if (i=='drymass') {
		 df_tmp = data.frame(name2 = str_split(names(g[[i]]), pattern="name2_", simplify = TRUE)[,2],
		 					 y_hat = unname(g[[i]]),
		 					 trait = rep('drymass', length(g[[i]])),
		 					 dap = rep("NA", length(g[[i]])))
    	 df_updated = df_tmp
	}
	if (i!='drymass') {
		 df_tmp = data.frame(name2 = str_split(names(g[[i]]), pattern="name2_", simplify = TRUE)[,2],
		 					 y_hat = unname(g[[i]]),
		 					 trait = rep('height', length(g[[i]])),
		 					 dap = rep(str_split(i, pattern="_", simplify = TRUE)[,2], length(g[[i]])))
		 df_updated = rbind(df_updated, df_tmp)
	}

}

# Read bin matrix:
W_bin = read.csv('W_bin.csv', row.names=1)

# Get the ID:
id = df[!duplicated(df[,c('name2', 'id_gbs')])][,c('name2', 'id_gbs')]

# Subset line names that were phenotyped and genotyped:
id = id[id$id_gbs %in% rownames(W_bin),]

# Subset adjusted means of the lines that were phenotyped and genotyped:
df_updated = df_updated[df_updated$name2 %in% id$name2,]

# Add column mapping the ID of the GBS:
df_updated= merge(df_updated, id, by='name2', sort=F) 

# Dropping levels:
df_updated$id_gbs = df_updated$id_gbs %>% droplevels

# Reodering the name of the data frame and eliminating the name2 identifier:
df_updated = df_updated[,c("id_gbs", "y_hat", "trait", "dap")] 
df_updated = df_updated[order(df_updated$trait),]
df_updated = df_updated[order(df_updated$dap),]


#---------------------------------------------Saving outputs-------------------------------------------------#

# Set directory:
setwd(paste0(OUT_PATH, '/processed_data'))

# Save adjusted means:
write.csv(df_updated, file="adjusted_means.csv")
<file_sep>/codes/functions.py


#----------------------------------------------Functions-----------------------------------------------------#

# Loading libraries:
import numpy
import os
import pandas

# Build the Cockerham's model (2, 1 and 0 centered marker scores):
def W_model(x):
  # Get allelic frequencies:
  p = (numpy.mean(x, axis=0)/2).values.reshape(1,(x.shape[1]))
  # Build Cokerham's model:
  W = x -  2*numpy.repeat(p, repeats= x.shape[0], axis=0)
  return(W)

# Construct the artificial bins:
def get_bin(x, n_bin, method):
  # Generate batches
  batches = numpy.array_split(numpy.arange(x.shape[1]), n_bin)
  # Initialize the binned matrix:
  W_bin = pandas.DataFrame(index=x.index, columns=map('bin_{}'.format, range(n_bin)))
  e_bin = []
  if method=='pca':
    for i in range(n_bin):
      # Compute SVD of the matrix bin:
      u,s,v = numpy.linalg.svd(x.iloc[:,batches[i]], full_matrices=False)
      # Compute the first principal component and adding to the binned matrix:
      W_bin['bin_' + str(i)] = numpy.dot(u[:,:1], numpy.diag(s[:1]))
      e_bin.append(s[0]/s.sum())
    return([W_bin, e_bin, batches])
  if method=='average':
    for i in range(n_bin):
      # Compute the mean across the columns and adding to the binned matrix:
      W_bin['bin_' + str(i)] = x.iloc[:,batches[i]].mean(axis=1)
      return([W_bin, batches])
<file_sep>/codes/process_raw_data.py

#------------------------------------------------Modules-----------------------------------------------------#

# Import python modules
import os
import numpy as np
import pandas as pd
import argparse
parser = argparse.ArgumentParser()


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Get flags:
parser.add_argument("-dpath", "--dpath", dest = "dpath", help="The path of the folder with the raw data")
parser.add_argument("-rpath", "--rpath", dest = "rpath", help="The path of the repository")
parser.add_argument("-opath", "--opath", dest = "opath", help="The path of the folder to receive outputs (processed data)")

# Parse the paths:
args = parser.parse_args()

# Subset arguments:
DATA_PATH = args.dpath
REPO_PATH = args.rpath
OUT_PATH = args.opath
# DATA_PATH = '/workdir/jp2476/raw_data_sorghum-multi-trait'
# REPO_PATH = '/workdir/jp2476/sorghum-multi-trait'
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'


#--------------------------------------------Processing data-------------------------------------------------#

# Set directory:
os.chdir(REPO_PATH + "/clean_repository/codes")

# Import functions:
from functions import * 

# Set directory:
os.chdir(DATA_PATH + "/raw_data")

# Number of data frames:
n_df = 4

# Create an empty data frame:
df = []

# Load data:
df.append(pd.read_csv("Biomass_2016.csv", index_col=0))
df.append(pd.read_csv("Biomass_SF2017.csv", index_col=0))
df.append(pd.read_csv("heights_2016.csv", index_col=0))
df.append(pd.read_csv("heights_SF2017.csv", index_col=0))

# Get names of the data frame variables:
tmp = []
for i in range(n_df): tmp.append(list(df[i]))

# Rename the columns of the data frames:
df[0] = df[0].rename(index=str, columns={tmp[0][0]: 'loc',
		  							 	 tmp[0][1]: 'plot',
								 		 tmp[0][2]: 'name1',
								 		 tmp[0][3]: 'name2',
								 		 tmp[0][4]: 'set',
								 		 tmp[0][5]: 'block',
								 		 tmp[0][6]: 'drymass'})

df[1] = df[1].rename(index=str, columns={tmp[1][0]: 'plot',
		  							 	 tmp[1][1]: 'name1',
								 		 tmp[1][2]: 'name2',
								 		 tmp[1][3]: 'loc',
								 		 tmp[1][4]: 'set',
								 		 tmp[1][5]: 'block',
								 		 tmp[1][6]: 'drymass'})


df[2] = df[2].rename(index=str, columns={tmp[2][0]: 'plot',
		  							 	 tmp[2][1]: 'name1',
								 		 tmp[2][2]: 'name2'})

df[3] = df[3].rename(index=str, columns={tmp[3][0]: 'plot',
		  							 	 tmp[3][1]: 'name1',
								 		 tmp[3][2]: 'name2',
								 		 tmp[3][3]: 'taxa',
								 		 tmp[3][4]: 'year',
								 		 tmp[3][5]: 'loc',
								 		 tmp[3][6]: 'set',
								 		 tmp[3][7]: 'block'})

# Add column mapping traits to the df:
df[0] = df[0].assign(trait=pd.Series(np.repeat('biomass', df[0].shape[0])).values)
df[1] = df[1].assign(trait=pd.Series(np.repeat('biomass', df[1].shape[0])).values)
df[2] = df[2].assign(trait=pd.Series(np.repeat('height', df[2].shape[0])).values)
df[3] = df[3].assign(trait=pd.Series(np.repeat('height', df[3].shape[0])).values)

# Add columns mapping years to the df:
df[0] = df[0].assign(year=pd.Series(np.repeat('16', df[0].shape[0])).values)
df[1] = df[1].assign(year=pd.Series(np.repeat('17', df[1].shape[0])).values)
df[2] = df[2].assign(year=pd.Series(np.repeat('16', df[2].shape[0])).values)

# Load inbred lines ID:
os.chdir(DATA_PATH + "/raw_data")
line_names = pd.read_csv("genotype_names_corrected.csv")

# Rename columns:
line_names = line_names.rename(index=str, columns={'Name2': 'name2'})

# Subset columns:
line_names = line_names[['name2', 'taxa']]

# Inclusion of the identifiers to the data frame:
for i in range(len(df)):
	df[i] = pd.merge(df[i], line_names, on='name2', how='left')
	if i==3:
		df[i] = df[i].drop('taxa_x', axis=1)
		df[i] = df[i].rename(index=str, columns={'taxa_y': 'id_gbs'})
	else:
		df[i] = df[i].rename(index=str, columns={'taxa': 'id_gbs'})
	# df[i] = df[i].loc[~df[i].id_gbs.isnull(),]

# Index of for selecting columns of the df mapping the design related to biomass collected on 2016:
tmp = ['plot','loc', 'set', 'block', 'range', 'row']

# Inclusion of the design variables to the combination height/2016
df[2] = pd.merge(df[2], df[0].loc[:,tmp], on='plot', how='left')

# Remove checks row from data not related to this experiment project:
for i in range(len(df)):
	df[i].block = df[i].block.astype(str)
	df[i] = df[i].loc[~(df[i]['set']=='CHK_STRP'),].copy(deep=True)

# Combine data frames into a unique data frame:
df = pd.concat(df, sort=False, axis=0)

# Change the index of the data frame:
df.index = np.arange(df.shape[0])

# Drop identifiers that will not be used:
df = df.drop(['name1'], axis=1)

# Load marker matrix:
M = pd.read_csv("gbs.csv")

# Load marker matrix:
loci_info = pd.read_csv("gbs_info.csv", index_col=0)

# Intersection between IDs:
mask = ~df['id_gbs'].isnull()
line_names = np.intersect1d(np.unique(df['id_gbs'][mask].astype(str)), list(M))

# Ordering lines:
M = M.loc[:, line_names]

# Function to build the Cockerham's model:
W = W_model(x=M.transpose())

# Build the bin matrix:
tmp = get_bin(x=W, n_bin=1000, method='pca')
W_bin = tmp[0]
W_bin.index = line_names

# Store the position of the bins into the genome:
bin_map = tmp[2]

# Add a new column for the data frame with the loci positions:
loci_info['bin'] = np.repeat(np.nan, loci_info.shape[0])

# Add the bins names for the bin column mapping its position:
for i in range(len(bin_map)):
	loci_info['bin'].iloc[bin_map[i]] = np.repeat(["bin_" + str(i)], bin_map[i].size)

# Name of the height covariates (columns that will be melt):
tmp = []
tmp.append(list(map('h{}'.format, range(1,8))))

# Name of all covariates except heights (variables that will be unaffected):
tmp.append(np.setdiff1d(list(df), tmp[0]))

# Change the shape of the data frame, and adding a new column mapping the days after planting (DAP)
df = pd.melt(df, id_vars=tmp[1], value_vars=tmp[0], var_name="dap", value_name="height")

# Replace categories by the values:
df['dap'] = df['dap'].replace(tmp[0], range(30, 120+15,15))

# Replace codification of locations:
df['loc'] = df['loc'].replace(['16EF', '16FF'], ['EF', 'FF'])

# Change data type of the DAP:
df.dap = df.dap.astype(object)

# Remove the DAP values from biomass, DAP values were taken only for plant height:
index = ((df.trait == 'biomass') & (df.dap == 120)) | (df.trait == 'height')
df = df[index] 

# Remove the DAP of 120 for biomass, it is not a right feature of biomass, actually it was collected in the end of the season:
index = (df.trait == 'biomass') & (df.dap == 120)
df.dap[index] = np.nan

# Remove the range column, it will not be used into the analysis:
df = df.drop('range', axis=1)

# Get just the features used in the paper:
df = df[['name2', 'id_gbs', 'block', 'loc', 'year', 'trait', 'dap', 'drymass', 'height']]

# Changing traits codification:
df.trait[df.trait == 'biomass'] = 'DM'
df.trait[df.trait == 'height'] = 'PH'

# Changing the indexes of the data frame:
df.index = range(df.shape[0]) 

# Changing the data type:
df['year'] = df['year'].astype(str)
df['drymass'] = df['drymass'].astype(str)

# Remove duplicate rows:
df = df.drop_duplicates()

# Changing the data type:
df['drymass'] = df['drymass'].astype(float)

## Writing into the disk the cleaned data:

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Writing the full data frame with phenotypic data and IDs:
df.to_csv("df.csv")

# Writing the bin info:
loci_info.to_csv("loci_info.csv")

# Writing the genomic binned matrix under Cockerham's model:						
W_bin.to_csv("W_bin.csv")

# Writing the full marker matrix:
M.to_csv("M.csv")
<file_sep>/codes/heritability_phenotypic_data_analysis.R

#-------------------------------------------Loading libraries------------------------------------------------#

# Import python libraries:
library(optparse)
library(data.table)
library(magrittr)
library(stringr)


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Set the flags:
option_list = list(
  make_option(c("-a", "--asremlpath"), type="character", default=NULL, 
              help="The path of the folder with asreml license", metavar="character"),
  make_option(c("-o", "--opath"), type="character", default=NULL, 
              help="The path of the folder to receive outputs", metavar="character")
) 

# Parse the arguments:
opt_parser = OptionParser(option_list=option_list)
args = parse_args(opt_parser)

# Subset arguments:
OUT_PATH = args$opath
ASREML_PATH = args$asremlpath
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'
# OUT_PATH = '/home/jhonathan/Documents/output_sorghum-multi-trait'
# ASREML_PATH = '/workdir/jp2476/asreml'


#--------------------------------Load asreml and function for heritability-------------------------------------#

# Load asreml:
setwd(ASREML_PATH)
library(asreml)
asreml.lic(license = "asreml.lic", install = TRUE)

# Function to get the heritability and its standard deviation estimates with asreml v3.0:
get_h2 <- function(asreml_obj, n_env, n_plot) {

	# Re-create the output from a basic, univariate model in asreml-R: 
	asrMod <- list(gammas = asreml_obj$gammas, gammas.type = asreml_obj$gammas.type, ai = asreml_obj$ai)

	# Name objects:
	names(asrMod[[1]]) <- names(asrMod[[2]]) <- names(asreml_obj$gammas)

	# Compute the heritability and its standard deviation:
	#---Observation: V4/n_env is the var_GxE/#environment and V5/#plot is the var_E/#plot
	formula = eval(parse(text=paste0('h2 ~ V3 / (V3 + V4/', n_env, '+ V5/', n_plot, ')')))

	return(nadiv:::pin(asrMod, formula))

}


#----------------------------------------------Load data-----------------------------------------------------#

# Set the directory:
setwd(paste0(OUT_PATH, '/processed_data'))

# Load the data:
df = fread('df.csv', header=T)[,-1]

# Change classes:
df$id_gbs = df$id_gbs %>% as.factor
df$name2 = df$name2 %>% as.factor
df$block = df$block %>% as.factor
df$loc = df$loc %>% as.factor
df$year = df$year %>% as.factor
df$trait = df$trait %>% as.factor
df$dap = df$dap %>% as.integer
df$drymass = df$drymass %>% as.numeric
df$height = df$height %>% as.numeric

# Get the group levels of the day after plantting (DAP) measures:
dap_groups = df$dap %>% unique %>% na.omit %>% as.integer

# Create a new column combing location and years:
df$env <- paste0(as.character(df$loc), "_",as.character(df$year))
df$env <- df$env %>% as.factor()

# Creating a list to receive the first stage results:
fit = list()

# Create a table to receive the heritabilities:
h2_comp = data.frame(trait=c('DM', paste0('PH_', unique(df$dap[!is.na(df$dap)]))), n_env=NA, n_plot=NA, h2=NA, se=NA)


#--------Compute the harmonic means of the number of environments and plots the lines were evaluated---------#

# Compute the harmonic mean of the number of environments in which each line was observed
#-->  and the harmonic mean of the total number of plots in which each line was observed:
for (t in as.character(h2_comp$trait)) {

	# Mask to subset the data of the trait:
	if (t == 'DM') {
		mask = df$trait=='DM' & (!is.na(df$drymass))
	}
	if (str_detect(t, 'PH')) {
		mask = df$trait=='PH' & df$dap == str_split(t, pattern="_", simplify = TRUE)[,2] & (!is.na(df$height))
	}

	# Subset the data:
	df_tmp = df[mask]

	# Remove hybrids from the data:
	mask = !(df_tmp$name2 %in% c('Pacesetter', 'SPX'))
	df_tmp = df_tmp[mask]

	for (i in unique(as.character(df_tmp$name2))) {

		# Subset the line:
		mask = df_tmp$name2==i 

		# Count the number of environments and plots the line i was evaluated:
		n_env_tmp = c(length(unique(as.character(df_tmp[mask, ]$env))))
		n_plot_tmp = c(nrow(df_tmp[mask, ]))
		names(n_env_tmp) = i
		names(n_plot_tmp) = i

		# Stack numbers in the vector:
		if (i==unique(as.character(df_tmp$name2))[1]) {

			n_env = n_env_tmp
			n_plot = n_plot_tmp

		}
		else {

			n_env = c(n_env, n_env_tmp)
			n_plot = c(n_plot, n_plot_tmp)

		}

	}

	# Compute the harmonic mean of the number of environments and plots lines were evaluated:
	h2_comp[h2_comp$trait==t, 'n_env'] = 1/mean(1/n_env)
	h2_comp[h2_comp$trait==t, 'n_plot'] = 1/mean(1/n_plot)

}


#-----------------------------------------Biomass data analysis------------------------------------------------#

# Index for mapping the results:
index = "drymass"

# Subsetting part of the data set: 
df_tmp = df[df$trait == "DM"]

# Dropping levels not present in the data subset:
df_tmp$id_gbs = df_tmp$id_gbs %>% droplevels
df_tmp$block = df_tmp$block %>% droplevels
df_tmp$env = df_tmp$env %>% droplevels

# Getting number of levels:
n_env = df_tmp$env %>% nlevels
n_rep = df_tmp$block %>% nlevels

# Fitting the model:
fit[[index]]  = asreml(drymass~ 1 + id_gbs,
			 		   random = ~ id_gbs + env + block:env + id_gbs:env,
			 		   na.method.Y = "include",
			 		   control = asreml.control(
			 		   maxiter = 200),
             		   data = df_tmp)

# Compute the heritability and its standard deviation:
h2_comp[h2_comp$trait == 'DM', c('h2','se')] = round(get_h2(fit[[index]], 
										  				   n_env=h2_comp$n_env[h2_comp$trait == 'DM'],
														   n_plot=h2_comp$n_plot[h2_comp$trait == 'DM']), 4)
 

#------------------------------------------Height data analysis------------------------------------------------#

for (index in as.character(h2_comp$trait)[-1]) {

	# Subsetting part of the data set:
	mask  = df$trait=='PH' & df$dap == str_split(index, pattern="_", simplify = TRUE)[,2] 
	df_tmp = df[mask]  

	# Dropping levels not present in the data subset:
	df_tmp$id_gbs = df_tmp$id_gbs %>% droplevels
	df_tmp$block = df_tmp$block %>% droplevels
	df_tmp$env = df_tmp$env %>% droplevels

	# Getting number of levels:
	n_env = df_tmp$env %>% nlevels
	n_rep = df_tmp$block %>% nlevels

	fit[[index]]  = asreml(height ~ 1 + id_gbs,
				 		   random = ~ id_gbs + block:env + env + id_gbs:env,
				 		   na.method.Y = "include",
				 		   control = asreml.control(
				 		   maxiter = 200),
	             		   data = df_tmp)

	# Compute the heritability and its standard deviation:
	h2_comp[h2_comp$trait == index, c('h2','se')] = round(get_h2(fit[[index]], 
											  				   n_env=h2_comp$n_env[h2_comp$trait == index],
															   n_plot=h2_comp$n_plot[h2_comp$trait == index]), 4)

	# Printing current analysis:
	print(paste0("Analysis ", index, " done!"))

}

# Set the directory:
setwd(paste0(OUT_PATH, '/heritabilities'))

# Save adjusted means:
write.csv(h2_comp, file="heritabilities.csv")
<file_sep>/codes/split_data_cross_validation.py
#------------------------------------------------Modules-----------------------------------------------------#

# Load libraries:
import pandas as pd
import numpy as np
import os
from sklearn.model_selection import KFold
import argparse
from itertools import chain
parser = argparse.ArgumentParser()


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Get flags:
parser.add_argument("-dpath", "--dpath", dest = "dpath", help="The path of the folder with the raw data")
parser.add_argument("-opath", "--opath", dest = "opath", help="The path of the folder to receive outputs")

# Parse the paths:
args = parser.parse_args()

# Subset arguments:
DATA_PATH = args.dpath
OUT_PATH = args.opath
# DATA_PATH = '/workdir/jp2476/raw_data_sorghum-multi-trait'
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'


#------------------------------------Load Adjusted Means and plotting----------------------------------------#

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Read adjusted means:
df = pd.read_csv("adjusted_means.csv", index_col=0)

# Change class of the dap:
df.dap = df.dap.fillna(0).astype(int)

# Read marker binned matrix:
W_bin = pd.read_csv("W_bin.csv", header = 0, index_col=0)

# Set the directory to store processed data:
os.chdir(DATA_PATH + "/raw_data")

# Read the file with the annotation of the population structure of the lines:
pop_struc = pd.read_csv("population_structure.txt", header = 0, sep='\t')

# Read the file with the annotation of the population structure of the lines:
pop_struc = pd.read_csv("population_structure.txt", header = 0, sep='\t')
pop_struc = pop_struc.rename(columns={'ID':'id_gbs', 'CATEGORY':'id_pop'})

#----------------------------Split data into groups for 5th-fold cross-validation----------------------------#

# Initialize dictionary to receive index stratified by population type:
pop_folds = dict()

# Number of folds:
n_fold = 5

# Initialize lists:
for counter in range(n_fold):
	pop_folds['k' + str(counter)] = []

# Split the data into subsets stratified by population:
for pop in pop_struc.id_pop.unique():
	counter=0
	# Create five folds:
	kf = KFold(n_splits=n_fold, shuffle=True, random_state=(1234+counter))
	# Get the splits:
	index = kf.split(pop_struc.id_gbs[pop_struc.id_pop == pop])
	# Get the indexes:
	for _, fold in index:
		pop_folds['k' + str(counter)].append(pop_struc.id_gbs[pop_struc.id_pop == pop].iloc[fold].tolist())
		counter+=1

# Unnest the lists:
for counter in range(n_fold):
	pop_folds['k' + str(counter)] = list(chain.from_iterable(pop_folds['k' + str(counter)]))

# Initialize lists to receive the random indexes:
trn_index = []
tst_index = []
 
# Build sets:
trn_index.append(pop_folds['k0'] + pop_folds['k1'] + pop_folds['k2'] + pop_folds['k3']) 
tst_index.append(pop_folds['k4'])

trn_index.append(pop_folds['k0'] + pop_folds['k1'] + pop_folds['k2'] + pop_folds['k4']) 
tst_index.append(pop_folds['k3'])

trn_index.append(pop_folds['k0'] + pop_folds['k1'] + pop_folds['k3'] + pop_folds['k4']) 
tst_index.append(pop_folds['k2'])

trn_index.append(pop_folds['k0'] + pop_folds['k2'] + pop_folds['k3'] + pop_folds['k4']) 
tst_index.append(pop_folds['k1'])

trn_index.append(pop_folds['k1'] + pop_folds['k2'] + pop_folds['k3'] + pop_folds['k4']) 
tst_index.append(pop_folds['k0'])

# Create dictionary with the data from the first cross-validation scheme:
y = dict()
X = dict()

# Build the response vector with all entries:
y_all = df.y_hat

# Build the feature matrix using all data set entries:
X_all = pd.get_dummies(df['id_gbs']) 
X_all = X_all.dot(W_bin.loc[X_all.columns])
X_all = pd.concat([df.dap, X_all], axis=1)

# Create different sets:
sets = ['trn', 'tst']

# Build the sets of the data for the 5-fold cross-validation scheme:
for s in sets:
	for t in df.trait.unique():
		for i in range(n_fold):
			# Key index for building the dictionary:
			key_index = 'cv5f_' + t + '_k' + str(i) + '_' + s
			if s == 'trn':
				# Logical vector for indexation:
				index = df.id_gbs.isin(trn_index[i]) & (df.trait==t)
			if s == 'tst':
				# Logical vector for indexation:
				index = df.id_gbs.isin(tst_index[i]) & (df.trait==t)
			# Build the response vector for the subset of data:
			y[key_index] = df.y_hat[index]
			# Build feature matrix for the subset of data:
			X[key_index] = X_all[index]
			if t == 'drymass':
				X[key_index] = X[key_index].replace(to_replace=0, value=1) \
										   .rename(columns={'dap': 'mu_index'})

# Name of the forward-chaining cross-validation schemes:
fcv_types = ['fcv-30~45', 'fcv-30~60', 'fcv-30~75', 'fcv-30~90', 'fcv-30~105']

# Build the sets of the data for the forward-chaining cross-validation schemes for the DBN model:
for c in fcv_types:
	for s in sets:
		# Key index for building the dictionary:
		key_index = c + '_height_' + s
		# Get the upper index of the dap:
		upper_index = int(c.split('~')[1])
		if s == 'trn':
			# Logical vector for indexation:
			index = (df.trait=='height') & (df.dap<=upper_index)
		if s == 'tst':
			# Logical vector for indexation:
			index = (df.trait=='height') & (df.dap>upper_index)
		# Build the response vector for the subset of data:
		y[key_index] = df.y_hat[index]
		# Build feature matrix for the subset of data:
		X[key_index] = X_all[index]

# Get DAP groups:
dap_group = df.dap.unique()[1:7]

# Build the sets of the data for the forward-chaining cross-validation schemes for the PBN and BN models using all data:
for d in dap_group:
	for s in sets:
		# Key index for mapping data into dictionary:
		key_index = 'fcv-' + str(d) + '~only' + '_height_' + s
		if s == 'trn':
			# Logical vector for indexation:
			index = (df.trait=='height') & (df.dap==d)
		if s == 'tst':
			# Logical vector for indexation:
			index = (df.trait=='height') & (df.dap!=d)
		# Build the response vector for the subset of data:
		y[key_index] = df.y_hat[index]
		# Build feature matrix for the subset of data:
		X[key_index] = X_all[index]

# Create the sets only for the PBN model with drymass data together with height data:
key_index = 'fcv_drymass_trn'
y[key_index] = df.y_hat[df.trait=='drymass']
X[key_index] = X_all[df.trait=='drymass']


#------------------------------------Save different subsets of the data--------------------------------------#

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Save cross-validation data:
for i in y.keys():
	y[i].to_csv('y_' + i + '.csv', header='y_hat')
	X[i].to_csv('x_' + i + '.csv')
<file_sep>/codes/Population_structure.Rmd
---
title: "Population structure analysis"
output:
  pdf_document: default
  html_notebook: default
---

# POPULATION STRUCTURE THROUGH PCA

## Calculating PCAs 

```{r}

#source("http://bioconductor.org/biocLite.R")
#biocLite("gdsfmt")
#biocLite("SNPRelate")

library(gdsfmt)
library(SNPRelate)
library(dplyr)
library(magrittr)
library(tidyr)

# Get the path of the plink files
bed.fn <- "/workdir/share/Jhonathan_tmp/Exact/FINAL.bed"
fam.fn <- "/workdir/share/Jhonathan_tmp/Exact/FINAL.fam"
bim.fn <- "/workdir/share/Jhonathan_tmp/Exact/FINAL.bim"

# Build the gdsdb
snpgdsBED2GDS(bed.fn, fam.fn, bim.fn,out.gdsfn ="final.gds")
snpgdsSummary("final.gds")
genofile <- snpgdsOpen("final.gds")

# Getting the PCA object
pca <- snpgdsPCA(genofile, num.thread=2)

# Calculate the percentage of variance explained by the first 10 PC
pc.percent <- pca$varprop*100
head(round(pc.percent, 2)) 

```

## Coloring by Individuals with a known Race assignment

```{r}

tab <- data.frame(sample.id = pca$sample.id,
    EV1 = pca$eigenvect[,1],    # the first eigenvector
    EV2 = pca$eigenvect[,2],    # the second eigenvector
    EV3 = pca$eigenvect[,3],
    EV4 = pca$eigenvect[,4],
    EV5 = pca$eigenvect[,5],
    stringsAsFactors = FALSE)
colnames(tab) <- c("ID", "EV1", "EV2", "EV3", "EV4", "EV5")

write.table(tab, file = "/workdir/share/Jhonathan_tmp/Exact/EVs", row.names = F, col.names = F, sep = "\t", quote = F)


SIMPLEID <- read.delim("/workdir/share/Jhonathan_tmp/Exact/EVs_IDs", stringsAsFactors = F, header = F)
tab$ID <- SIMPLEID$V1

POPS <- read.delim("/workdir/share/Jhonathan_tmp/Exact/BN_str.4.meanQ", stringsAsFactors = F, header = F) %>%
  `colnames<-`(c("POP1", "POP2", "POP3", "POP4"))

metadata <- cbind(tab, POPS)


#Load the metadata
Races <- read.delim("/workdir/share/Jhonathan_tmp/metadata/meta", stringsAsFactors = F) %>%
  select(ID, RACE)



metadata2 <- metadata %>%
  left_join(Races, by="ID") %>%
  distinct %>%
  mutate(RACE= case_when( is.na(RACE) ~ "NA",
                          RACE == "PCA outlier" ~ "NA",
                         TRUE  ~ RACE))


library(ggplot2)
library(wesanderson)
library(patchwork)

fill <- wes_palette("IsleofDogs1")


pal <- wes_palette("Zissou1", 5, type = "continuous")


p1 <- ggplot(metadata2, aes(EV1, EV2)) +
  geom_point(aes(colour = factor(RACE)), alpha =0.5, shape=19) +
  #scale_color_manual(values=c("red", "blue", fill[1], fill[2], fill[3], fill[4])) +
  scale_color_manual(values = c("black", pal[1], pal[1], pal[3], pal[3], "olivedrab4", "olivedrab4", "olivedrab4", "red", "lightgrey")) +
  ylab("PC 2 (4.20%)") +
  xlab("PC 1 (5.59%)") +
  theme(legend.position="right", legend.direction="vertical",
        legend.title = element_blank()) +
  theme(axis.line.x = element_line(size=1, colour = "black"),
        axis.line.y = element_line(size=1, colour = "black"),
        panel.grid.major = element_line(colour = "#d3d3d3"), panel.grid.minor = element_blank(),
        panel.border = element_blank(), panel.background = element_blank()) +
  theme(plot.title = element_text(size = 14, family = "Tahoma", face = "bold"),
        text=element_text(family="Tahoma"),
        axis.text.x=element_text(colour="black", size = 10),
        axis.text.y=element_text(colour="black", size = 10),
        legend.key=element_rect(fill="white", colour="white"))

p1

```

## Coloring by the subpopulations estimated using fastStructure

```{r}

metadata3 <- metadata2 %>%
  mutate(K= case_when( POP1 > POP2 & POP1 > POP3 & POP1 >POP4   ~ "POP1",
                       POP2 > POP1 & POP2 > POP3 & POP2 >POP4   ~ "POP2",
                       POP3 > POP1 & POP3 > POP2 & POP3 >POP4   ~ "POP3",
                       POP4 > POP1 & POP4 > POP2 & POP4 >POP3   ~ "POP4",
                       TRUE ~ "MIX"))
 
  

p2 <- ggplot(metadata3, aes(EV1, EV2)) +
  geom_point(aes(colour = factor(K)), alpha =0.5, shape=19) +
  #scale_color_manual(values=c("red", "blue", fill[1], fill[2], fill[3], fill[4])) +
  scale_color_manual(values = c("black", "red", "green", "cyan", "purple")) +
  ylab("PC 2 (4.20%)") +
  xlab("PC 1 (5.59%)") +
  theme(legend.position="right", legend.direction="vertical",
        legend.title = element_blank()) +
  theme(axis.line.x = element_line(size=1, colour = "black"),
        axis.line.y = element_line(size=1, colour = "black"),
        panel.grid.major = element_line(colour = "#d3d3d3"), panel.grid.minor = element_blank(),
        panel.border = element_blank(), panel.background = element_blank()) +
  theme(plot.title = element_text(size = 14, family = "Tahoma", face = "bold"),
        text=element_text(family="Tahoma"),
        axis.text.x=element_text(colour="black", size = 10),
        axis.text.y=element_text(colour="black", size = 10),
        legend.key=element_rect(fill="white", colour="white"))

p2


```

## Exporting the subpopulation assignmentas a table:

```{r}
metadata4 <- metadata3   %>%
  mutate(PERC= case_when( K == "POP1"   ~ POP1,
                       K == "POP2"   ~ POP2,
                       K == "POP3"   ~ POP3,
                       K == "POP4"   ~ POP4,
                       K == "MIX" ~ 0.25))


metadata5 <- metadata4 %>%
  mutate(CATEGORY = case_when(PERC < 0.8 ~ "ADMIX",
                              TRUE ~ K)) %>%
  select(ID, CATEGORY)


write.table(x = metadata5, file="./Exact/Structure_groups_extended_simple.txt", quote = F, row.names = F, col.names = T, sep = "\t")

```


<file_sep>/codes/mixed_model_analysis.R

#-------------------------------------------Loading libraries------------------------------------------------#

# Libraries to manipulate data:
library(optparse)
library(data.table)
library(stringr)
library(magrittr)
library(EMMREML)
library(rrBLUP)
library(tidyr)


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Set the flags:
option_list = list(
  make_option(c("-y", "--y"), type="character", default=NULL, 
              help="Name of the file with the phenotypes", metavar="character"),
  make_option(c("-m", "--model"), type="character", default=NULL, 
              help="Specify the model", metavar="character"),
  make_option(c("-o", "--opath"), type="character", default=NULL, 
              help="The path of the folder to receive outputs", metavar="character"),
  make_option(c("-c", "--cvpath"), type="character", default=NULL, 
              help="The path of the folder to receive outputs of the step of cross-validation analysis", metavar="character")
) 

# Parse the arguments:
opt_parser = OptionParser(option_list=option_list)
args = parse_args(opt_parser)

# Subset arguments:
y = args$y
model = args$model
OUT_PATH = args$opath
CV_OUT_PATH = args$cvpath

# y = 'y_fcv-30~105_height_trn.csv'
# model = 'MTi-GBLUP-0~5'
# OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'
# CV_OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait/cv/MTi-GBLUP/fcv-30~105/heightUP/fcv-30~60/drymass-height'

#---------------------------------Define function for spectral decomposition---------------------------------#

# Function to fix negative eigenvalues from the kinship matrix (K):
spec_dec <-function(K) {

  # Eigenvalue decomposition of the matrix:	
  E=eigen(K)

  # Create a diagonal matrix for receiving the eigenvalues:
  Dg= diag(E$values)

  # Identify eigenvalues lower then a small constant and pick the lowest one:
  for (i in 1:nrow(Dg)) {
    if (Dg[i,i]<1e-4) {
      Dg[i,i]=Dg[(i-1),(i-1)]-0.01*Dg[(i-1),(i-1)]
    }
  }

  # Create a matrix of the eigen vectors:
  C=matrix(E$vectors,nrow(K),ncol(K))

  # Reconstruct the marker matrix:
  K = C%*%Dg%*%t(C)

  # Return the transformed matrix:
  return(K)

}


#-----------------------------------------------Load data----------------------------------------------------#

# Load train data just to get the indexes of the train, and test sets, used in the Bayesian Networks analysis:

# Set the directory:
setwd(paste0(OUT_PATH, '/processed_data'))

if (str_detect(model, 'MTr')) {

	# Subset file names per trait:
	y = str_split(y, "&", simplify = TRUE)[1,1]

}

# Read adjusted means:
y = fread(y, header=TRUE) %>% data.matrix
rownames(y) = y[,1]
y = y[,-1]
	
# Read the whole marker matrix:
M = fread('M.csv', header=TRUE)[,-1] %>% data.matrix %>% t(.)

# Load data frame with adjusted means to be processed and used on EMMREML for mixed model analysis:
df = fread('adjusted_means.csv', header=TRUE) %>% data.frame
rownames(df) = df[,1]
df = df[,-1]
df$id_gbs = df$id_gbs %>% as.factor %>% droplevels
df$trait = df$trait %>% as.factor %>% droplevels
df$dap = df$dap %>% as.factor %>% droplevels

# Get the indexes of train set:
index = names(y)

# Get the inbred lines ID for train set:
id_trn = df[index,]$id_gbs %>% as.character %>% unique

# Get the indexes of test set:
id_tst = df$id_gbs %>% as.character %>% unique %>% setdiff(., id_trn)  

#-------------------------------------Prepare data format for EMMREML----------------------------------------#

# Create the relationship matrix:
A = A.mat(M-1)
rownames(A) = rownames(M)
colnames(A) = rownames(M)

# Change the names of the columns:
colnames(df) = c("id", "y_hat", "trait", "dap")

# Conditional based on models:
if (str_detect(model, 'MTi')) {

	# Vector with the time points:
	time = seq(30, 120, 15)

	# Subset just time points used for training:
	upper = str_split(model, '~', simplify = TRUE)[1,2] %>% as.numeric

	# Change DAP class on the data frame:
	df$dap = df$dap %>% as.character

	# Subset data frame:
	df_tmp = df[!(df$dap=='NA'),]
	df_tmp = df_tmp[as.numeric(df_tmp$dap) <= time[upper+1],]

	# Melting the data frame:
	df_melt = df_tmp %>% spread(key = dap, value=y_hat)

	# Subset the desired columns:
	df_melt = df_melt[,c('id', as.character(seq(30, time[upper+1], by=15)))]

	# Change row and column names:
	rownames(df_melt) = df_melt$id
	colnames(df_melt) = c('id_gbs', paste0('h', seq(30, time[upper+1], by=15)))

}
if (str_detect(model, 'MTr')) {

	# Substitute the NA for a drymass for melting all traits together:
	df$dap = df$dap %>% as.character
	mask = df$dap == 'NA'
	df[mask,]$dap = 'drymass'

	# Vector with the time points:
	time = seq(30, 120, 15)

	# Subset just time points used for training:
	upper = str_split(model, '~', simplify = TRUE)[1,2] %>% as.numeric

	# Subset columns:
	df_tmp = df[,colnames(df)!='trait']

	# Melting the data frame:
	df_melt = df_tmp %>% spread(key = dap, value=y_hat)
	
	# Subset the desired columns:
	df_melt = df_melt[,c('id','drymass', as.character(seq(30,120,by=15)))]

	# Select height data used for train:
	tmp = time[time <= time[upper+1]]
	df_melt = df_melt[, c('id','drymass', tmp)]
	rownames(df_melt) = df_melt$id
	colnames(df_melt) = c('id_gbs','dm', paste0('h', seq(30, time[upper+1], by=15)))
	mask = !is.na(df_melt$dm)
	df_melt = df_melt[mask,]
	mask = rownames(A) %in% df_melt$id_gbs
	A = A[mask, mask]

}

# Reordering data frame:
df_melt = df_melt[rownames(A),] %>% droplevels
rownames(df_melt) = 1:nrow(df_melt)

# Changing the order of the factors:
df_melt$id_gbs = factor(df_melt$id_gbs, levels = c(as.character(df_melt$id_gbs)))

# Subset only individuals phenotyped and genotyped:
df_melt = df_melt[df_melt$id_gbs %in% rownames(A),]

# Reoder the relationship matrix as in the data frame:
index = as.character(df_melt$id_gbs)
A = A[index, index]

# Build design matrix of fixed effects:
n = nrow(df_melt)
X = matrix(1,nrow=1,ncol=n)
colnames(X) = df_melt$id_gbs %>% as.character

# Build design matrix of random effects:
Z = model.matrix(~ -1 + df_melt$id_gbs)
rownames(Z) = df_melt$id_gbs %>% as.character
colnames(Z) = df_melt$id_gbs %>% as.character

if (str_detect(model, 'MTi')) {

	# Preparing matrix with the response variable:
	select = paste0('h', seq(30,time[upper+1],15))
	Y = t(df_melt[, select])
	colnames(Y) = df_melt$id_gbs %>% as.character

}
if (str_detect(model, 'MTr')) {

	# Preparing matrix with the response variable:
	select = c('dm', paste0('h', seq(30,time[upper+1],15)))
	Y = t(df_melt[, select])
	colnames(Y) = df_melt$id_gbs %>% as.character

	# Update train and test set ID:
	id_trn = id_trn[id_trn %in% df_melt$id_gbs]
	id_tst = id_tst[id_tst %in% df_melt$id_gbs]

}

# Fit the multivariate linear mixed model:
out <- emmremlMultivariate(Y = Y[,id_trn],
			  			   X = X[,id_trn],
						   Z = Z[,id_trn],
						   K = spec_dec(A),
						   tolpar = 1e-4,
						   varBhat = FALSE, varGhat = FALSE, PEVGhat = FALSE, test = FALSE)

# Save output:
setwd(CV_OUT_PATH)
save.image("output_mixed_models.RData")
<file_sep>/codes/generate_2pc_barplot.py

#------------------------------------------------Modules-----------------------------------------------------#

# Import python modules
import os
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import argparse
parser = argparse.ArgumentParser()


#-----------------------------------------Adding flags to the code-------------------------------------------#

# Get flags:
parser.add_argument("-dpath", "--dpath", dest = "dpath", help="The path of the folder with the raw data")
parser.add_argument("-rpath", "--rpath", dest = "rpath", help="The path of the repository")
parser.add_argument("-opath", "--opath", dest = "opath", help="The path of the folder to receive outputs (processed data)")

# Parse the paths:
args = parser.parse_args()

# Subset arguments:
# DATA_PATH = args.dpath
# REPO_PATH = args.rpath
# OUT_PATH = args.opath
DATA_PATH = '/workdir/jp2476/raw_data_sorghum-multi-trait'
REPO_PATH = '/workdir/jp2476/sorghum-multi-trait'
OUT_PATH = '/workdir/jp2476/output_sorghum-multi-trait'


#----------------------------Generate histogram with percentage explained by the PCs-------------------------#

# Set directory:
os.chdir(DATA_PATH + "/raw_data")

# Load marker matrix:
M = pd.read_csv("gbs.csv")

# Load inbred lines ID:
line_names = pd.read_csv("genotype_names_corrected.csv")

# Load marker matrix:
loci_info = pd.read_csv("gbs_info.csv", index_col=0)

# Set the directory to store processed data:
os.chdir(OUT_PATH + "/processed_data")

# Read adjusted means:
df = pd.read_csv("adjusted_means.csv", index_col=0)

# Intersection between IDs:
mask = ~df['id_gbs'].isnull()
line_names = np.intersect1d(np.unique(df['id_gbs'][mask].astype(str)), list(M))

# Ordering lines:
M = M.loc[:, line_names]

# Function to construct the bins:
def get_perc_pca(x, n_bin):
  # Generating batches
  batches = np.array_split(np.arange(x.shape[1]), n_bin)
  # Initializing the binned matrix:
  e_bin = pd.DataFrame(index=map('bin_{}'.format, range(n_bin)), columns=['pca1', 'pca2'])
  for i in range(n_bin):
    # Computing SVD of the matrix bin:
    u,s,v = np.linalg.svd(x.iloc[:,batches[i]], full_matrices=False)
    # Computing the first principal component and adding to the binned matrix:
    e_bin.loc['bin_' + str(i)]['pca1'] = (s[0]/s.sum())*100
    e_bin.loc['bin_' + str(i)]['pca2'] = (s[1]/s.sum())*100
  return e_bin

# Get percent explained by each bin:
perc_pca = get_perc_pca(M.transpose(), n_bin = 1000)

# Store index into a column:
perc_pca['Bin'] = perc_pca.index

# Melt data frame:
perc_pca = pd.melt(perc_pca, id_vars = 'Bin', var_name='pcs', value_name='perc')

# Change codification:
mask = perc_pca.pcs == 'pca1'
perc_pca.pcs[mask] = '1st PCA'
mask = perc_pca.pcs == 'pca2'
perc_pca.pcs[mask] = '2nd PCA'
perc_pca.columns = ['Bin', 'Principal Component', 'Percent explained (%)']

# Generate barplot with the percent explained by each bin:
sns.set(style="whitegrid")

# Draw a nested barplot to show survival for class and sex
g = sns.barplot(x="Bin", y='Percent explained (%)', hue='Principal Component', data=perc_pca)
g.despine(left=True)
# g.set_ylabels("survival probability")

# Set directory:
os.chdir(REPO_PATH + "/figures")

plt.savefig("perc_pca_expl.pdf", dpi=350)
plt.savefig("perc_pca_expl.png", dpi=350)
plt.clf()

