<commit_before>#!/usr/bin/env python 
import argparse, os,sys
from numpy import *
import numpy as np
from scipy.special import gamma
from scipy.special import beta as f_beta
import random as rand
import platform, time
import multiprocessing, thread
import multiprocessing.pool
import csv
from scipy.special import gdtr, gdtrix
from scipy.special import betainc
import scipy.stats
np.set_printoptions(suppress=True)
np.set_printoptions(precision=3)  
from multiprocessing import Pool, freeze_support
import thread
import imp
lib_updates_priors = imp.load_source("lib_updates_priors", "pyrate_lib/lib_updates_priors.py")
lib_DD_likelihood = imp.load_source("lib_DD_likelihood", "pyrate_lib/lib_DD_likelihood.py")
from lib_updates_priors import *
from lib_DD_likelihood  import *


#### DATA ###

p = argparse.ArgumentParser() #description='<input file>') 
p.add_argument('-d', type=str,   help='data set', default=0, metavar=0)
p.add_argument('-m', type=int,   help='model', default=0, metavar=0)
p.add_argument('-n', type=int,   help='MCMC iterations', default=1000000, metavar=1000000)
p.add_argument('-s', type=int,   help='sampling freq.', default=1000, metavar=1000)
p.add_argument('-p', type=int,   help='print freq.', default=5000, metavar=5000)
p.add_argument('-j', type=int,   help='replicate', default=0, metavar=0)
p.add_argument('-c', type=int,   help='clade', default=0, metavar=0)
p.add_argument('-b', type=float, help='shape parameter (beta) of Be hyper=prior pn indicators', default=1, metavar=1)
p.add_argument('-T', type=float, help='Max time slice', default=np.inf, metavar=np.inf)


args = p.parse_args()


dataset=args.d
n_iterations=args.n
sampling_freq=args.s
print_freq = args.p
#t_file=np.genfromtxt(dataset, names=True, delimiter='\t', dtype=float)
t_file=np.loadtxt(dataset, skiprows=1)

clade_ID=t_file[:,0]
clade_ID=clade_ID.astype(int)
ts=t_file[:,2+2*args.j]
te=t_file[:,3+2*args.j]

constr=args.m

if args.c==0:
	single_focal_clade = False
	clade_name = ""
else:
	single_focal_clade = True
	fixed_focal_clade = args.c-1
	clade_name = "_c%s" % (fixed_focal_clade)

Be_shape_beta = args.b
if Be_shape_beta>1: beta_value = "_B%s" % (args.b)	
else: beta_value = ""

all_events=sort(np.concatenate((ts,te),axis=0))[::-1] # events are speciation/extinction that change the diversity trajectory
n_clades,n_events=max(clade_ID)+1,len(all_events)
Dtraj=init_Dtraj(n_clades,n_events)


##### get indexes
s_list = []
e_list = []
s_or_e_list=[]
clade_inx_list=[]
unsorted_events = []
for i in range(n_clades):
	"used for Dtraj"
	s_list.append(ts[clade_ID==i])
	e_list.append(te[clade_ID==i])
	"used for lik calculation"
	s_or_e_list += list(np.repeat(1,len(ts[clade_ID==i]))) # index 1 for s events
	s_or_e_list += list(np.repeat(2,len(te[clade_ID==i]))) # index 2 for e events
	clade_inx_list += list(np.repeat(i,2*len(te[clade_ID==i])))
	unsorted_events += list(ts[clade_ID==i])
	unsorted_events += list(te[clade_ID==i])

s_or_e_array= np.array(s_or_e_list)
unsorted_events= np.array(unsorted_events)
s_or_e_array[unsorted_events==0] = 3
s_or_e_array[unsorted_events>args.T] = 4
unsorted_events[unsorted_events>args.T] = args.T
""" so now: s_or_e_array = 1 (s events), s_or_e_array = 2 (e events), s_or_e_array = 3 (e=0 events), s_or_e_array = 3 (s>T events)"""


""" concatenate everything:
                          1st row: all events  2nd row index s,e     3rd row clade index """
all_events_temp= np.array([unsorted_events,    s_or_e_array,         np.array(clade_inx_list)])
# sort by time
idx = np.argsort(all_events_temp[0])[::-1] # get indexes of sorted events
all_events_temp2=all_events_temp[:,idx] # sort by time of event
#print all_events_temp2
#print shape(all_events_temp2),len(all_events)
all_time_eve=all_events_temp2[0]

Dtraj_new=Dtraj
idx_s = []
idx_e = []
for i in range(n_clades): # make trajectory curves for each clade
	Dtraj_new[:,i]=getDT(all_events_temp2[0],s_list[i],e_list[i])
	ind_clade_i = np.arange(len(all_events_temp2[0]))[all_events_temp2[2]==i]
	ind_sp = np.arange(len(all_events_temp2[0]))[all_events_temp2[1]==1]
	ind_ex = np.arange(len(all_events_temp2[0]))[all_events_temp2[1]==2]
	idx_s.append(np.intersect1d(ind_clade_i,ind_sp))
	idx_e.append(np.intersect1d(ind_clade_i,ind_ex))
	
	ind_extant = np.arange(len(all_events_temp2[0]))[all_events_temp2[1]==3]
	ind_extant_clade_i = np.intersect1d(ind_clade_i,ind_extant) 
	
	#ext_eve_clade_i = all_time_eve[idx_e[i]]
	#print ext_eve_clade_i
	print "clade: %s\n%s species, %s ts, %s te, %s extant" % (i, len(ind_clade_i)/2, len(idx_s[i]), len(idx_e[i]), len(ind_extant_clade_i) )
#####

scale_factor = 1./np.max(Dtraj)
MAX_G = 0.30/scale_factor

GarrayA=init_Garray(n_clades)+1. # 3d array so:
                                 # Garray[i,:,:] is the 2d G for one clade
			         # Garray[0,0,:] is G_lambda, Garray[0,1,:] is G_mu for clade 0
RA=init_Garray(n_clades)
Constr_matrix=make_constraint_matrix(n_clades, constr)

l0A,m0A=init_BD(n_clades),init_BD(n_clades)

out_file_name="%s_%s_m%s_MCDD%s%s.log" % (dataset,args.j,constr,clade_name,beta_value)
logfile = open(out_file_name , "wb") 
wlog=csv.writer(logfile, delimiter='\t')

lik_head=""
for i in range(n_clades): lik_head+="\tlik_%s" % (i)
head="it\tposterior\tlikelihood%s\tprior" % (lik_head)
for i in range(n_clades): head+="\tl%s" % (i)
for i in range(n_clades): head+="\tm%s" % (i)
for i in range(n_clades): 
	for j in range(n_clades): 
		head+="\tGl%s_%s" % (i,j)
	for j in range(n_clades): 
		head+="\tGm%s_%s" % (i,j)
for j in range(n_clades): 
	head+="\thypZ%s" % (j)
		
head+="\thypR"
wlog.writerow(head.split('\t'))<commit_msg>added option to specify min/max age, min cladesize<commit_after>#!/usr/bin/env python 
import argparse, os,sys
from numpy import *
import numpy as np
from scipy.special import gamma
from scipy.special import beta as f_beta
import random as rand
import platform, time
import multiprocessing, thread
import multiprocessing.pool
import csv
from scipy.special import gdtr, gdtrix
from scipy.special import betainc
import scipy.stats
np.set_printoptions(suppress=True)
np.set_printoptions(precision=3)  
from multiprocessing import Pool, freeze_support
import thread
import imp
lib_updates_priors = imp.load_source("lib_updates_priors", "pyrate_lib/lib_updates_priors.py")
lib_DD_likelihood = imp.load_source("lib_DD_likelihood", "pyrate_lib/lib_DD_likelihood.py")
from lib_updates_priors import *
from lib_DD_likelihood  import *


#### DATA ###

p = argparse.ArgumentParser() #description='<input file>') 
p.add_argument('-d', type=str,   help='data set', default=0, metavar=0)
p.add_argument('-m', type=int,   help='model', default=0, metavar=0)
p.add_argument('-n', type=int,   help='MCMC iterations', default=1000000, metavar=1000000)
p.add_argument('-s', type=int,   help='sampling freq.', default=1000, metavar=1000)
p.add_argument('-p', type=int,   help='print freq.', default=5000, metavar=5000)
p.add_argument('-j', type=int,   help='replicate', default=0, metavar=0)
p.add_argument('-c', type=int,   help='clade', default=0, metavar=0)
p.add_argument('-b', type=float, help='shape parameter (beta) of Be hyper=prior pn indicators', default=1, metavar=1)
p.add_argument('-T0', type=float, help='Max time slice', default=np.inf, metavar=np.inf)
p.add_argument('-T1', type=float, help='Min time slice', default=0, metavar=0)
p.add_argument('-min_taxa', type=float, help='Min clade size', default=0, metavar=0)


args = p.parse_args()


dataset=args.d
n_iterations=args.n
sampling_freq=args.s
print_freq = args.p
#t_file=np.genfromtxt(dataset, names=True, delimiter='\t', dtype=float)
t_file=np.loadtxt(dataset, skiprows=1)

clade_ID=t_file[:,0]
clade_ID=clade_ID.astype(int)
ts=t_file[:,2+2*args.j]
te=t_file[:,3+2*args.j]
n_original_clades = len(np.unique(clade_ID))


start_time = args.T0
end_time   = args.T1

if start_time == np.inf: start_time = max(ts)

# remove older lineages
not_older_lineages = (te<start_time).nonzero()[0]
ts = ts[not_older_lineages]
te = te[not_older_lineages]
clade_ID = clade_ID[not_older_lineages]

# remove younger lineages
younger_lineages = (ts>end_time).nonzero()[0]
ts = ts[younger_lineages]
te = te[younger_lineages]
clade_ID = clade_ID[younger_lineages]

# remove clades with n. species < threshold
min_number_of_lineages=args.min_taxa
rm_clade=[]
for i in np.unique(clade_ID):
	if len(clade_ID[clade_ID==i])<min_number_of_lineages:
		rm_clade.append(i)

for i in rm_clade:
	ts =ts[clade_ID != i]
	te =te[clade_ID != i]
	clade_ID =clade_ID[clade_ID != i]



# some clades might disappear altogether
original_clade_ID = clade_ID+0
unique_original_clade_ID = np.unique(original_clade_ID)
n_clades_left = len(np.unique(clade_ID))
j=0
for i in range(n_original_clades):
	if i in np.unique(clade_ID):
		clade_ID[clade_ID==i] = j
		j+=1
# shift data
ts = ts-end_time
te = te-end_time

# set extant lineages
te[te<0] = 0


# define lineages originating before start_time
argsT = start_time-end_time # ts > argsT are set equal to argsT and not treated as originations

print sort(ts)[::-1]+end_time
print sort(te)[::-1]+end_time
print original_clade_ID
print clade_ID


constr=args.m

if args.c==0:
	single_focal_clade = False
	clade_name = ""
else:
	single_focal_clade = True
	fixed_focal_clade = args.c-1
	clade_name = "_c%s" % (fixed_focal_clade)

Be_shape_beta = args.b
if Be_shape_beta>1: beta_value = "_B%s" % (args.b)	
else: beta_value = ""

all_events=sort(np.concatenate((ts,te),axis=0))[::-1] # events are speciation/extinction that change the diversity trajectory
n_clades,n_events=max(clade_ID)+1,len(all_events)
Dtraj=init_Dtraj(n_clades,n_events)


##### get indexes
s_list = []
e_list = []
s_or_e_list=[]
clade_inx_list=[]
unsorted_events = []
for i in range(n_clades):
	"used for Dtraj"
	s_list.append(ts[clade_ID==i])
	e_list.append(te[clade_ID==i])
	"used for lik calculation"
	s_or_e_list += list(np.repeat(1,len(ts[clade_ID==i]))) # index 1 for s events
	s_or_e_list += list(np.repeat(2,len(te[clade_ID==i]))) # index 2 for e events
	clade_inx_list += list(np.repeat(i,2*len(te[clade_ID==i])))
	unsorted_events += list(ts[clade_ID==i])
	unsorted_events += list(te[clade_ID==i])

s_or_e_array= np.array(s_or_e_list)
unsorted_events= np.array(unsorted_events)
s_or_e_array[unsorted_events==0] = 3
s_or_e_array[unsorted_events>argsT] = 4
unsorted_events[unsorted_events>argsT] = argsT
""" so now: s_or_e_array = 1 (s events), s_or_e_array = 2 (e events), s_or_e_array = 3 (e=0 events), s_or_e_array = 4 (s>T events)"""
print max(s_or_e_array), min(s_or_e_array)

""" concatenate everything:
                          1st row: all events  2nd row index s,e     3rd row clade index """
all_events_temp= np.array([unsorted_events,    s_or_e_array,         np.array(clade_inx_list)])
# sort by time
idx = np.argsort(all_events_temp[0])[::-1] # get indexes of sorted events
all_events_temp2=all_events_temp[:,idx] # sort by time of event
#print all_events_temp2
#print shape(all_events_temp2),len(all_events)
all_time_eve=all_events_temp2[0]

Dtraj_new=Dtraj
idx_s = []
idx_e = []
for i in range(n_clades): # make trajectory curves for each clade
	Dtraj_new[:,i]=getDT(all_events_temp2[0],s_list[i],e_list[i])
	ind_clade_i = np.arange(len(all_events_temp2[0]))[all_events_temp2[2]==i]
	ind_sp = np.arange(len(all_events_temp2[0]))[all_events_temp2[1]==1]
	ind_ex = np.arange(len(all_events_temp2[0]))[all_events_temp2[1]==2]
	idx_s.append(np.intersect1d(ind_clade_i,ind_sp))
	idx_e.append(np.intersect1d(ind_clade_i,ind_ex))
	
	ind_extant = np.arange(len(all_events_temp2[0]))[all_events_temp2[1]==3]
	ind_extant_clade_i = np.intersect1d(ind_clade_i,ind_extant) 
	
	#ext_eve_clade_i = all_time_eve[idx_e[i]]
	#print ext_eve_clade_i
	print "\nclade: %s\n%s species, %s ts, %s te, %s extant" % (unique_original_clade_ID[i], len(ind_clade_i)/2, len(idx_s[i]), len(idx_e[i]), len(ind_extant_clade_i) )
#####
print "\n\n"
scale_factor = 1./np.max(Dtraj)
MAX_G = 0.30/scale_factor

GarrayA=init_Garray(n_clades)+1. # 3d array so:
                                 # Garray[i,:,:] is the 2d G for one clade
			         # Garray[0,0,:] is G_lambda, Garray[0,1,:] is G_mu for clade 0
RA=init_Garray(n_clades)
Constr_matrix=make_constraint_matrix(n_clades, constr)

l0A,m0A=init_BD(n_clades),init_BD(n_clades)

out_file_name="%s_%s_m%s_MCDD%s%s.log" % (dataset,args.j,constr,clade_name,beta_value)
logfile = open(out_file_name , "wb") 
wlog=csv.writer(logfile, delimiter='\t')

lik_head=""
for i in range(n_clades): lik_head+="\tlik_%s" % (i)
head="it\tposterior\tlikelihood%s\tprior" % (lik_head)
for i in range(n_clades): head+="\tl%s" % (unique_original_clade_ID[i])
for i in range(n_clades): head+="\tm%s" % (unique_original_clade_ID[i])
for i in range(n_clades): 
	for j in range(n_clades): 
		head+="\tGl%s_%s" % (unique_original_clade_ID[i],unique_original_clade_ID[j])
	for j in range(n_clades): 
		head+="\tGm%s_%s" % (unique_original_clade_ID[i],unique_original_clade_ID[j])
for j in range(n_clades): 
	head+="\thypZ%s" % (j)
		
head+="\thypR"
wlog.writerow(head.split('\t'))
{'fraction_non_alphanumeric': 0.15077537233225857, 'fraction_numerical': 0.018962075848303395, 'mean_word_length': 6.277653631284916, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3005659', 'n_tokens_mistral': 5700, 'n_tokens_neox': 5496, 'n_words': 981}