test.py

import sympy
import sys
import redcode
from sympy.utilities.codegen import codegen
from sympy.printing import print_ccode

#The Data
Q=sympy.Symbol('pair.m', positive=True)

MM2=sympy.Symbol('pair.Y_MM', positive=True)
Mm2=sympy.Symbol('pair.Y_Mm', positive=True)
mm2=sympy.Symbol('pair.Y_mm', positive=True)

MM1=sympy.Symbol('pair.X_MM', positive=True)
Mm1=sympy.Symbol('pair.X_Mm', positive=True)
mm1=sympy.Symbol('pair.X_mm', positive=True)

data=[MM2, Mm2, mm2, MM1, Mm1, mm1]
redcode.set_data(data)

#The Parameters
g_XY=sympy.Symbol('rel.gamma_XY_')
g_YX=sympy.Symbol('rel.gamma_YX_')
T=sympy.Symbol('rel.theta_XY_')
D=sympy.Symbol('rel.Delta_XY_')
d=sympy.Symbol('rel.delta_XY_')
F_Y=sympy.Symbol('rel.f_Y_')
F_X=sympy.Symbol('rel.f_X_')

#An array of the parameters.
params=[F_X, F_Y, T, g_XY, g_YX, d, D]
redcode.set_params(params)

#The three componenents of the likelihood function. H00 is homozygous for the major allele, H01 heterozygous, and H11 homozygous minor.

P=1-Q
#sympy.Q.positive(P)
#var=P*(1.-P)
#std=var**(0.5);
#skew=(1.-2.*P)/std;
#kurtosis=1./(1.-P)+1./P-3.;

#mmmmi=(Q**4+(F_X+F_Y+4*T)*var*Q**2-2*(g_XY+g_YX)*skew*Q+d*kurtosis+D*var**2);
#MMMMi=(P**4+(F_X+F_Y+4*T)*var*P**2-2*(g_XY+g_YX)*skew*P+d*kurtosis+D*var**2);
#Mmmmi=(2*P*Q**3+2*(F_Y*P*Q-F_X*Q**2-2*T*(Q**2-P*Q))*var-4*g_XY*(P-Q)*skew+2*g_YX*(Q-P)*skew-2*d*kurtosis-2*D*var**2);
#mmMmi=(2*P*Q**3+2*(F_X*P*Q-F_Y*Q**2-2*T*(Q**2-P*Q))*var-4*g_YX*(P-Q)*skew+2*g_XY*(Q-P)*skew-2*d*kurtosis-2*D*var**2);
#MmMMi=(2*Q*P**3+2*(F_Y*P*Q-F_X*P**2-2*T*(P**2-P*Q))*var-4*g_XY*(Q-P)*skew+2*g_YX*(P-Q)*skew-2*d*kurtosis-2*D*var**2);
#MMMmi=(2*Q*P**3+2*(F_X*P*Q-F_Y*P**2-2*T*(P**2-P*Q))*var-4*g_YX*(Q-P)*skew+2*g_XY*(P-Q)*skew-2*d*kurtosis-2*D*var**2);
#MmMmi=(4*P**2*Q**2+4*(T*(P**2+Q**2-2*P*Q)-F_X*P*Q-F_Y*P*Q)*var+4*g_XY*(P-Q)*skew+4*g_YX*(Q-P)*skew+4*d*kurtosis+4*D*var**2);
#MMmmi=(P**2*Q**2+(F_Y*P**2+F_X*Q**2-4*T*P*Q)*var+2*g_XY*P*skew-2*g_YX*Q*skew+d*kurtosis+D*var**2);
#mmMMi=(P**2*Q**2+(F_X*P**2+F_Y*Q**2-4*T*P*Q)*var+2*g_YX*P*skew-2*g_XY*Q*skew+d*kurtosis+D*var**2);

A=P
C=P

Va=A*(1.-A)#;    //variance of the two haploid genomes of A 
Vc=C*(1.-C)#;    // "   "        "       "       "     of B
Sa=sympy.sqrt(Va)#;    //standard deviation of haploid genomes A
Sc=sympy.sqrt(Vc)#;    // and "        "       "       "       C.

E_A2  =(F_X*Va+2.*(A**2)+Va)/2.#; //Expectation of A^2
E_C2  =(F_Y*Vc+2.*(C**2)+Vc)/2.#; //
E_AC  =T*Sa*Sc+A*C;
ga=(1.-2.*A)/Sa;
gc=(1.-2.*C)/Sc;
E_A2C =(g_XY*Va*Sc*ga+A*A*C+Va*(T*(1.+2.*A)+F_X*C+C/(1-C) ) )/2.;
E_AC2 =(g_YX*Vc*Sa*gc+C*C*A+Vc*(T*(1.+2.*C)+F_Y*A+A/(1-A) ) )/2.;
ka=1./(1.-A)+1./A-3.;
kc=1./(1.-C)+1./C-3.;
E_A2C2=(d*sympy.sqrt(ka*kc)+D)*Va*Vc+A*A*C*C+F_X*Va*C*C+F_Y*Vc*A*A+4.*T*Sa*Sc*A*A+C*2.*g_XY*Va*Sc*ga+2.*A*g_YX*Vc*Sa*gc;

e=0

mmmm=1-6*P+0.0*e+2*E_A2+2*E_C2+8.0*E_AC-4*E_A2C-4*E_AC2+1*E_A2C2;
Mmmm=0+4*P+2.0*e-4*E_A2+0*E_C2-10.*E_AC+8*E_A2C+4*E_AC2-2*E_A2C2;
MMmm=0-1*P-0.5*e+2*E_A2+0*E_C2+2.0*E_AC-4*E_A2C-0*E_AC2+1*E_A2C2;
mmMm=0+4*P-2.0*e+0*E_A2-4*E_C2-10.*E_AC+4*E_A2C+8*E_AC2-2*E_A2C2;
MmMm=0+0*P+0.0*e+0*E_A2+0*E_C2+12.*E_AC-8*E_A2C-8*E_AC2+4*E_A2C2;
MMMm=0+0*P+0.0*e+0*E_A2+0*E_C2-2.0*E_AC+4*E_A2C+0*E_AC2-2*E_A2C2;
mmMM=0-1*P+0.5*e+0*E_A2+2*E_C2+2.0*E_AC+0*E_A2C-4*E_AC2+1*E_A2C2;
MmMM=0+0*P+0.0*e+0*E_A2+0*E_C2-2.0*E_AC+0*E_A2C+4*E_AC2-2*E_A2C2;
MMMM=0+0*P+0.0*e+0*E_A2+0*E_C2+0.0*E_AC+0*E_A2C+0*E_AC2+1*E_A2C2;

#S=(1-(mmmmk+MMMMk+Mmmmk+mmMmk+MmMMk+MMMmk+MmMmk+MMmmk+mmMMk))**2
S=1

#The log likelihood equation
lnL=sympy.log( mmmm*mm1*mm2+mmMm*mm1*Mm2+mmMM*mm1*MM2+Mmmm*Mm1*mm2+MmMm*Mm1*Mm2+MmMM*Mm1*MM2+MMmm*MM1*mm2+MMMm*MM1*Mm2+MMMM*MM1*MM2 )

system_eq=[]

#We first need the three equations we are going to try and set to zero, i.e. the first partial derivitives wrt e h and F.
print "/*This code was automatically generated by "+str(sys.argv[0])+"*/\n"
print "#include \"allele.h\""
print "#include \"quartet.h\""
print "#include \"typedef.h\""
print "#include \"constants.h\""
print 
#numpy.set_printoptions(precission=18)


for x in range(0, 7):
	system_eq.append(sympy.diff(lnL, params[x]) )
	print "inline float_t H"+str(x)+" (const Genotype_pair &pair, const Constants <float_t, const std::pair<const Genotype_pair &, const Relatedness &> > &con) {"
	sys.stdout.write("\treturn ")
	out=(system_eq[-1])
	out=sympy.simplify(out)
	redcode.pre(out)
	keys=redcode.dag_sort(redcode.keys)
	for key in keys:
		out=redcode.exact_sub(out, redcode.constants[key], key)
	string=sympy.printing.ccode(out)
	print string, ";\n}\n"

#Then we need to make the Jacobian, which is a matrix with ...
for x in range(0, 7):
	for y in range(0, 7):
		print "inline float_t J"+str(x)+str(y)+" (const Genotype_pair &pair, const Constants <float_t, const std::pair <const Genotype_pair &, const Relatedness &> > &con) {"
		sys.stdout.write("\treturn ")
                out=(sympy.diff(system_eq[x], params[y]))
		#out=sympy.simplify(out)
		pre(out)
		keys=dag_sort(keys)
		for key in keys:
			out=exact_sub(out, constants[key], key)
	        string=sympy.printing.ccode(out)
		print string, ";\n}\n"

print "inline float_t lnL_NR (const Genotype_pair &pair, const Relatedness &rel) {"
sys.stdout.write("\treturn ")
print_ccode(sympy.simplify(lnL))
print ";\n}\n"

keys=dag_replace(keys)

for key in keys:
	print "inline void set_"+str(key).split('.')[1].replace('[','').replace(']','') +" (Constants <float_t, const std::pair <const Genotype_pair &, const Relatedness &> > &con, const std::pair <const Genotype_pair &, const Relatedness &> &d ) {"
	print "\tconst Genotype_pair *pair=&d.first;"
	print "\tconst Relatedness *rel=&d.second;"
	print "\tcon.c["+str(key).split('.')[1].replace('[','').replace(']','').strip('c')+"]=",
        string=sympy.printing.ccode(constants[key])
        string=string.replace("rel.", "rel->").replace("pair.","pair->")
	print string, ";\n}\n"

print "static void (*cfn["+str(len(constants))+"])(Constants <float_t, const std::pair <const Genotype_pair &, const Relatedness &> > &, const std::pair <const Genotype_pair &, const Relatedness &> &)={"+"&set_"+str(keys[0]).split('.')[1].replace('[','').replace(']',''),
for key in keys[1:]:
	print ", "+"&set_"+str(key).split('.')[1].replace('[','').replace(']','')
print "};"

print "#define REL_CNTS\t"+str(len(keys) )
print "#define REL_ARRAY\tcfn"

quit()