def get_compiled_theano_functions(N_QUAD_PTS):
# resonance j and k
j,k = T.lscalars('jk')
s = (j-k) / k
# Planet masses: m1,m2
m1,m2 = T.dscalars(2)
Mstar = 1
eta0 = Mstar
eta1 = eta0+m1
eta2 =eta1+m2
mtilde1 = m1 * (eta0/eta1)
Mtilde1 = Mstar * (eta1/eta0)
mtilde2 = m2 * (eta1/eta2)
Mtilde2 = Mstar * (eta2/eta1)
eps = m1 * m2 / (mtilde1 + mtilde2) / Mstar
beta1 = mtilde1 / (mtilde1 + mtilde2)
beta2 = mtilde2 / (mtilde1 + mtilde2)
gamma = mtilde2/mtilde1
# Angle variable for averaging over
Q = T.dvector('Q')
# Dynamical variables:
dyvars = T.vector()
sigma1, sigma2, I1, I2, amd = [dyvars[i] for i in range(5)]
# Quadrature weights
quad_weights = T.dvector('w')
# Set lambda2=0
l2 = T.constant(0.)
l1 = -1 * k * Q
w1 = (1+s) * l2 - s * l1 - sigma1
w2 = (1+s) * l2 - s * l1 - sigma2
Gamma1 = I1
Gamma2 = I2
# Resonant semi-major axis ratio
alpha_res = ((j-k)/j)**(2/3) * ((Mstar + m1) / (Mstar+m2))**(1/3)
P0 = k * ( beta2 - beta1 * T.sqrt(alpha_res) ) / 2
P = P0 - k * (s+1/2) * amd
Ltot = beta1 * T.sqrt(alpha_res) + beta2 - amd
L1 = Ltot/2 - P / k - s * (I1 + I2)
L2 = Ltot/2 + P / k + (1 + s) * (I1 + I2)
a1 = (L1 / beta1 )**2 * eta0 / eta1
e1 = T.sqrt(1-(1-(Gamma1 / L1))**2)
a2 = (L2 / beta2 )**2 * eta1 / eta2
e2 = T.sqrt(1-(1-(Gamma2 / L2))**2)
Hkep = - eta1 * beta1 / (2 * a1) / eta0 - eta2 * beta2 / (2 * a2) / eta1
alpha = a1 / a2
ko = KeplerOp()
M1 = l1 - w1
M2 = l2 - w2
sinf1,cosf1 = ko( M1, e1 + T.zeros_like(M1) )
sinf2,cosf2 = ko( M2, e2 + T.zeros_like(M2) )
R = calc_DisturbingFunction_with_sinf_cosf(alpha,e1,e2,w1,w2,sinf1,cosf1,sinf2,cosf2)
Rav = R.dot(quad_weights)
Hpert = -eps * Rav / a2
Htot = Hkep + Hpert
######################
# Dissipative dynamics
######################
tau_alpha_0, K1, K2, p = T.dscalars(4)
sigma1dot_dis,sigma2dot_dis,I1dot_dis,I2dot_dis,amddot_dis = T.dscalars(5)
sigma1dot_dis,sigma2dot_dis = T.as_tensor(0.),T.as_tensor(0.)
# Define timescales
tau_e1 = tau_alpha_0 / K1
tau_e2 = tau_alpha_0 / K2
tau_a1_0 = -1 * tau_alpha_0 * (1+alpha_res * gamma)/ (alpha_res * gamma)
tau_a2_0 = -1 * alpha_res * gamma * tau_a1_0
tau_a1 = 1 / (1/tau_a1_0 + 2 * p * e1*e1 / tau_e1 )
tau_a2 = 1 / (1/tau_a2_0 + 2 * p * e2*e2 / tau_e2 )
# Time derivative of orbital elements
e1dot_dis = -1*e1 / tau_e1
e2dot_dis = -1*e2 / tau_e2
a1dot_dis = -1*a1 / tau_a1
a2dot_dis = -1*a2 / tau_a2
# Time derivatives of canonical variables
I1dot_dis = L1 * e1 * e1dot_dis / ( T.sqrt(1-e1*e1) ) - I1 / tau_a1 / 2
I2dot_dis = L2 * e2 * e2dot_dis / ( T.sqrt(1-e2*e2) ) - I2 / tau_a2 / 2
Pdot_dis = -1*k * ( L2 / tau_a2 - L1 / tau_a1) / 4 - k * (s + 1/2) * (I1dot_dis + I2dot_dis)
amddot_dis = Pdot_dis / T.grad(P,amd)
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# Get numerical quadrature nodes and weights
nodes,weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)
# Rescale for integration interval from [-1,1] to [-pi,pi]
nodes = nodes * np.pi
weights = weights * 0.5
# 'givens' will fix some parameters of Theano functions compiled below
givens = [(Q,nodes),(quad_weights,weights)]
# 'ins' will set the inputs of Theano functions compiled below
# Note: 'extra_ins' will be passed as values of object attributes
# of the 'ResonanceEquations' class 'defined below
extra_ins = [m1,m2,j,k,tau_alpha_0,K1,K2,p]
ins = [dyvars] + extra_ins
# Define flows and jacobians.
# Conservative flow
gradHtot = T.grad(Htot,wrt=dyvars)
hessHtot = theano.gradient.hessian(Htot,wrt=dyvars)
Jtens = T.as_tensor(np.pad(getOmegaMatrix(2),(0,1),'constant'))
H_flow_vec = Jtens.dot(gradHtot)
H_flow_jac = Jtens.dot(hessHtot)
# Dissipative flow
dis_flow_vec = T.stack(sigma1dot_dis,sigma2dot_dis,I1dot_dis,I2dot_dis,amddot_dis)
dis_flow_jac = theano.gradient.jacobian(dis_flow_vec,dyvars)
# Extras
dis_timescales = [tau_a1_0,tau_a2_0,tau_e1,tau_e2]
orbels = [a1,e1,sigma1*k,a2,e2,sigma2*k]
##########################
# Compile Theano functions
##########################
if not DEBUG:
# Note that compiling can take a while
# so I've put a debugging switch here
# to skip evaluating these functions when
# desired.
Rav_fn = theano.function(
inputs=ins,
outputs=Rav,
givens=givens,
on_unused_input='ignore'
)
Hpert_av_fn = theano.function(
inputs=ins,
outputs=Hpert,
givens=givens,
on_unused_input='ignore'
)
Htot_fn = theano.function(
inputs=ins,
outputs=Htot,
givens=givens,
on_unused_input='ignore'
)
H_flow_vec_fn = theano.function(
inputs=ins,
outputs=H_flow_vec,
givens=givens,
on_unused_input='ignore'
)
H_flow_jac_fn = theano.function(
inputs=ins,
outputs=H_flow_jac,
givens=givens,
on_unused_input='ignore'
)
dis_flow_vec_fn = theano.function(
inputs=ins,
outputs=dis_flow_vec,
givens=givens,
on_unused_input='ignore'
)
dis_flow_jac_fn = theano.function(
inputs=ins,
outputs=dis_flow_jac,
givens=givens,
on_unused_input='ignore'
)
dis_timescales_fn =theano.function(
inputs=extra_ins,
outputs=dis_timescales,
givens=givens,
on_unused_input='ignore'
)
orbels_fn = theano.function(
inputs=ins,
outputs=orbels,
givens=givens,
on_unused_input='ignore'
)
else:
return [lambda x: x for _ in range(8)]
return Rav_fn,Hpert_av_fn,Htot_fn,H_flow_vec_fn,H_flow_jac_fn,dis_flow_vec_fn,dis_flow_jac_fn,dis_timescales_fn,orbels_fn
def get_compiled_Hkep_Hpert_full():
# resonance j and k
j,k = T.lscalars('jk')
s = (j-k) / k
# Planet masses: m1,m2
m1,m2 = T.dscalars(2)
Mstar = 1
eta0 = Mstar
eta1 = eta0+m1
eta2 =eta1+m2
mtilde1 = m1 * (eta0/eta1)
Mtilde1 = Mstar * (eta1/eta0)
mtilde2 = m2 * (eta1/eta2)
Mtilde2 = Mstar * (eta2/eta1)
eps = m1 * m2 / (mtilde1 + mtilde2) / Mstar
beta1 = mtilde1 / (mtilde1 + mtilde2)
beta2 = mtilde2 / (mtilde1 + mtilde2)
gamma = mtilde2/mtilde1
# Dynamical variables:
dyvars = T.vector()
Q,sigma1, sigma2, I1, I2, amd = [dyvars[i] for i in range(6)]
# Set lambda2=0
l2 = T.constant(0.)
l1 = -1 * k * Q
w1 = (1+s) * l2 - s * l1 - sigma1
w2 = (1+s) * l2 - s * l1 - sigma2
Gamma1 = I1
Gamma2 = I2
# Resonant semi-major axis ratio
alpha_res = ((j-k)/j)**(2/3) * ((Mstar + m1) / (Mstar+m2))**(1/3)
P0 = k * ( beta2 - beta1 * T.sqrt(alpha_res) ) / 2
P = P0 - k * (s+1/2) * amd
Ltot = beta1 * T.sqrt(alpha_res) + beta2 - amd
L1 = Ltot/2 - P / k - s * (I1 + I2)
L2 = Ltot/2 + P / k + (1 + s) * (I1 + I2)
a1 = (L1 / beta1 )**2 * eta0 / eta1
e1 = T.sqrt(1-(1-(Gamma1 / L1))**2)
a2 = (L2 / beta2 )**2 * eta1 / eta2
e2 = T.sqrt(1-(1-(Gamma2 / L2))**2)
Hkep = - eta1 * beta1 / (2 * a1) / eta0 - eta2 * beta2 / (2 * a2) / eta1
alpha = a1 / a2
ko = KeplerOp()
M1 = l1 - w1
M2 = l2 - w2
sinf1,cosf1 = ko( M1, e1 + T.zeros_like(M1) )
sinf2,cosf2 = ko( M2, e2 + T.zeros_like(M2) )
R = calc_DisturbingFunction_with_sinf_cosf(alpha,e1,e2,w1,w2,sinf1,cosf1,sinf2,cosf2)
Hpert = -eps * R / a2
omega_syn = T.sqrt(eta2/eta1)/a2**1.5 - T.sqrt(eta1/eta0) / a1**1.5
gradHpert = T.grad(Hpert,wrt=dyvars)
gradHkep = T.grad(Hkep,wrt=dyvars)
grad_omega_syn = T.grad(omega_syn,wrt=dyvars)
extra_ins = [m1,m2,j,k]
ins = [dyvars] + extra_ins
# Scalars
omega_syn_fn = theano.function(
inputs=ins,
outputs=omega_syn,
givens=None,
on_unused_input='ignore'
)
Hpert_fn = theano.function(
inputs=ins,
outputs=Hpert,
givens=None,
on_unused_input='ignore'
)
Hkep_fn = theano.function(
inputs=ins,
outputs=Hkep,
givens=None,
on_unused_input='ignore'
)
# gradients
grad_omega_syn_fn = theano.function(
inputs=ins,
outputs=grad_omega_syn,
givens=None,
on_unused_input='ignore'
)
gradHpert_fn = theano.function(
inputs=ins,
outputs=gradHpert,
givens=None,
on_unused_input='ignore'
)
gradHkep_fn = theano.function(
inputs=ins,
outputs=gradHkep,
givens=None,
on_unused_input='ignore'
)
return omega_syn_fn,Hkep_fn,Hpert_fn,grad_omega_syn_fn,gradHkep_fn,gradHpert_fn
def _index_variables(self, basename='index'):
return T.lscalars(
*['%s_%d' % (basename, i) for i in xrange(self.sequence_length)])
def _get_compiled_theano_functions(N_QUAD_PTS):
# Planet masses: m1,m2
m1, m2 = T.dscalars(2)
mstar = 1
mu1 = m1 * mstar / (mstar + m1)
mu2 = m2 * mstar / (mstar + m2)
Mstar1 = mstar + m1
Mstar2 = mstar + m2
beta1 = mu1 * T.sqrt(Mstar1 / mstar) / (mu1 + mu2)
beta2 = mu2 * T.sqrt(Mstar2 / mstar) / (mu1 + mu2)
j, k = T.lscalars('jk')
s = (j - k) / k
# Angle variable for averaging over
psi = T.dvector()
# Quadrature weights
quad_weights = T.dvector('w')
# Dynamical variables:
Ndof = 3
Nconst = 1
dyvars = T.vector()
y1, y2, y_inc, x1, x2, x_inc, amd = [
dyvars[i] for i in range(2 * Ndof + Nconst)
]
a20 = T.constant(1.)
a10 = ((j - k) / j)**(2 / 3) * (Mstar1 / Mstar2)**(1 / 3)
L10 = beta1 * T.sqrt(a10)
L20 = beta2 * T.sqrt(a20)
Ltot = L10 + L20
f = L10 / L20
L2res = (Ltot + amd) / (1 + f)
Psi = -k * (s * L2res + (1 + s) * f * L2res)
###
# actions
###
I1 = 0.5 * (x1 * x1 + y1 * y1)
I2 = 0.5 * (x2 * x2 + y2 * y2)
Phi = 0.5 * (x_inc * x_inc + y_inc * y_inc)
L1 = -s * Ltot - Psi / k - s * (I1 + I2 + Phi)
L2 = (1 + s) * Ltot + Psi / k + (1 + s) * (I1 + I2 + Phi)
# Set lambda2=0
l2 = T.constant(0.)
l1 = -1 * k * psi
theta_res = (1 + s) * l2 - s * l1
cos_theta_res = T.cos(theta_res)
sin_theta_res = T.sin(theta_res)
kappa1 = x1 * cos_theta_res + y1 * sin_theta_res
eta1 = y1 * cos_theta_res - x1 * sin_theta_res
kappa2 = x2 * cos_theta_res + y2 * sin_theta_res
eta2 = y2 * cos_theta_res - x2 * sin_theta_res
sigma = x_inc * cos_theta_res + y_inc * sin_theta_res
rho = y_inc * cos_theta_res - x_inc * sin_theta_res
# y = (sigma-i*rho)/sqrt(2)
# = sqrt(Phi) * exp[i (Omega1+Omega2) / 2]
# Malige+ 2002, Eqs 20 and 21
r2byr1 = (L2 - L1 - I2 + I1) / Ltot
sigma1 = rho * T.sqrt(1 + r2byr1) / T.sqrt(2)
sigma2 = -rho * T.sqrt(1 - r2byr1) / T.sqrt(2)
rho1 = -sigma * T.sqrt(1 + r2byr1) / T.sqrt(2)
rho2 = sigma * T.sqrt(1 - r2byr1) / T.sqrt(2)
Xre1 = kappa1 / T.sqrt(L1)
Xim1 = -eta1 / T.sqrt(L1)
Yre1 = 0.5 * sigma1 / T.sqrt(L1)
Yim1 = -0.5 * rho1 / T.sqrt(L1)
Xre2 = kappa2 / T.sqrt(L2)
Xim2 = -eta2 / T.sqrt(L2)
Yre2 = 0.5 * sigma2 / T.sqrt(L2)
Yim2 = -0.5 * rho2 / T.sqrt(L2)
absX1_sq = 2 * I1 / L1
absX2_sq = 2 * I2 / L2
X_to_z1 = T.sqrt(1 - absX1_sq / 4)
X_to_z2 = T.sqrt(1 - absX2_sq / 4)
Y_to_zeta1 = 1 / T.sqrt(1 - absX1_sq / 2)
Y_to_zeta2 = 1 / T.sqrt(1 - absX2_sq / 2)
a1 = (L1 / beta1)**2
k1 = Xre1 * X_to_z1
h1 = Xim1 * X_to_z1
q1 = Yre1 * Y_to_zeta1
p1 = Yim1 * Y_to_zeta1
e1 = T.sqrt(absX1_sq) * X_to_z1
inc1 = 2 * T.arcsin(T.sqrt(p1 * p1 + q1 * q1))
a2 = (L2 / beta2)**2
k2 = Xre2 * X_to_z2
h2 = Xim2 * X_to_z2
q2 = Yre2 * Y_to_zeta2
p2 = Yim2 * Y_to_zeta2
e2 = T.sqrt(absX2_sq) * X_to_z2
inc2 = 2 * T.arcsin(T.sqrt(p2 * p2 + q2 * q2))
beta1p = T.sqrt(Mstar1) * beta1
beta2p = T.sqrt(Mstar2) * beta2
Hkep = -0.5 * beta1p / a1 - 0.5 * beta2p / a2
Hdir, Hind = calc_Hint_components_spatial(a1, a2, l1, l2, h1, k1, h2, k2,
p1, q1, p2, q2, Mstar1, Mstar2)
eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
Hpert = (Hdir + Hind / mstar)
Hpert_av = Hpert.dot(quad_weights)
Htot = Hkep + eps * Hpert_av
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# Get numerical quadrature nodes and weights
nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)
# Rescale for integration interval from [-1,1] to [-pi,pi]
nodes = nodes * np.pi
weights = weights * 0.5
# 'givens' will fix some parameters of Theano functions compiled below
givens = [(psi, nodes), (quad_weights, weights)]
# 'ins' will set the inputs of Theano functions compiled below
# Note: 'extra_ins' will be passed as values of object attributes
# of the 'ResonanceEquations' class 'defined below
extra_ins = [m1, m2, j, k]
ins = [dyvars] + extra_ins
Stilde = Phi * (L2 - I2 - L1 + I1) / (Ltot)
Q1 = 0.5 * (Phi + Stilde)
Q2 = 0.5 * (Phi - Stilde)
inc1 = T.arccos(1 - Q1 / (L1 - I1))
inc2 = T.arccos(1 - Q2 / (L2 - I2))
orbels = [
a1, e1, inc1, k * T.arctan2(y1, x1), a2, e2, inc2,
k * T.arctan2(y2, x2),
T.arctan2(y_inc, x_inc)
]
orbels_dict = dict(
zip([
'a1', 'e1', 'inc1', 'theta1', 'a2', 'e2', 'inc2', 'theta2', 'phi'
], orbels))
actions = [L1, L2, I1, I2, Q1, Q2]
actions_dict = dict(
zip(['L1', 'L2', 'Gamma1', 'Gamma2', 'Q1', 'Q2'], actions))
# Conservative flow
gradHtot = T.grad(Htot, wrt=dyvars)
gradHpert = T.grad(Hpert_av, wrt=dyvars)
gradHkep = T.grad(Hkep, wrt=dyvars)
hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
hessHpert = theano.gradient.hessian(Hpert_av, wrt=dyvars)
hessHkep = theano.gradient.hessian(Hkep, wrt=dyvars)
Jtens = T.as_tensor(np.pad(getOmegaMatrix(Ndof), (0, Nconst), 'constant'))
H_flow_vec = Jtens.dot(gradHtot)
Hpert_flow_vec = Jtens.dot(gradHpert)
Hkep_flow_vec = Jtens.dot(gradHkep)
H_flow_jac = Jtens.dot(hessHtot)
Hpert_flow_jac = Jtens.dot(hessHpert)
Hkep_flow_jac = Jtens.dot(hessHkep)
##########################
# Compile Theano functions
##########################
func_dict = {
# Hamiltonians
'H': Htot,
#'Hpert':Hpert_av,
#'Hkep':Hkep,
## Hamiltonian flows
'H_flow': H_flow_vec,
#'Hpert_flow':Hpert_flow_vec,
#'Hkep_flow':Hkep_flow_vec,
## Hamiltonian flow Jacobians
'H_flow_jac': H_flow_jac,
#'Hpert_flow_jac':Hpert_flow_jac,
#'Hkep_flow_jac':Hkep_flow_jac,
## Extras
'orbital_elements': orbels_dict,
'actions': actions_dict
}
compiled_func_dict = dict()
with tqdm(func_dict.items()) as t:
for key, val in t:
t.set_description("Compiling '{}'".format(key))
if key is 'timescales':
inputs = extra_ins
else:
inputs = ins
cf = theano.function(inputs=inputs,
outputs=val,
givens=givens,
on_unused_input='ignore')
compiled_func_dict[key] = cf
return compiled_func_dict
dual_X_hat = T.nnet.sigmoid(T.dot(h,T.dot(h.T,X))/n_b) #shape: batch_size x n_x
dual_recon_err = T.nnet.binary_crossentropy(dual_X_hat,X).sum()
h2 = dropout((T.tensordot(h, W_h2, [[1],[1]]) + b_h2).max(axis = 1), p_drop_hidden) #shape inside tensordot: batch_size x 2 x n_h2
h3 = dropout ((T.tensordot(h2, W_h3, [[1],[1]]) + b_h3).max(axis = 1), p_drop_hidden) #shape inside tensordot: batch_size x 2 x n_h3
h3 = batchnorm(h3, epsilon=0)
h4 = dropout ((T.tensordot(h3, W_h4, [[1],[1]]) + b_h4).max(axis = 1), p_drop_hidden) #shape inside tensordot: batch_size x 2 x n_h4
prob_y = softmax(T.dot(h4, W_o)) #classification probabilities
return [prob_y , dual_recon_err]
#symbolic variables
X = T.fmatrix()
Y = T.fmatrix()
learning_rate = T.fscalar('learning_rate')
start, end = T.lscalars('start', 'end')
#weights and biases initialization
W_h = shared_normal((2, n_x, n_h))
W_h2 = shared_normal((2, n_h, n_h2))
W_h3 = shared_normal((2, n_h2, n_h3))
W_h4 = shared_normal((2, n_h3, n_h4))
W_o = shared_normal((n_h4, 10))
b_h = shared_normal((2, n_h), sigma = 0)
b_h2 = shared_normal((2, n_h2), sigma = 0)
b_h3 = shared_normal((2, n_h3), sigma = 0)
b_h4 = shared_normal((2, n_h4), sigma = 0)
X = binomial(X) #internal binarization
#model calls
[dout_prob_y, dout_dual_recon_err] = model_NG_ACE(X, batch_size, gaussian_err, 0.2, 0.5) #with dropout
def get_compiled_theano_functions(N_QUAD_PTS):
# resonance j and k
j, k = T.lscalars('jk')
# Planet masses: m1,m2
m1, m2 = T.dscalars(2)
# resonance f and g coefficients
f, g = T.dscalars(2)
# Planet and star mass variables
Mstar = 1
mu1 = m1 / (Mstar + m1)
mu2 = m2 / (Mstar + m2)
eps = m1 * mu2 / (mu1 + mu2) / Mstar
# Resonant semi-major axis ratio
alpha = ((j - k) / j)**(2 / 3) * ((Mstar + m1) / (Mstar + m2))**(1 / 3)
# Constants in Eq. (15)
fTilde = T.sqrt((mu1 + mu2) / (mu1 * T.sqrt(alpha))) * f
gTilde = T.sqrt((mu1 + mu2) / mu2) * g
# Constant in Eq. (8)
A = 1.5 * j * (mu1 + mu2) * (j / mu2 + (j - k) / mu1 / T.sqrt(alpha))
# Dynamical variables:
dyvars = T.vector()
theta, theta_star, J, J_star = [dyvars[i] for i in range(4)]
# Angle variable to average disturbing function over
kappa = T.dvector()
# Quadrature weights
quad_weights = T.dvector('w')
# Convert dynamical variables to eccentricities and angles:
# Note:
# Q is set to zero since it does not
# enter disturbing function except in combinations
# with z and w.
Q = T.as_tensor(0)
z = Q / k - theta
# See Eq. 20
Zsq = J * (fTilde * fTilde + gTilde * gTilde) / (f * f + g * g)
Z = T.sqrt(Zsq)
# Set W to zero
Wsinw, Wcosw = 0, 0
Zsinz, Zcosz = Z * T.sin(z), Z * T.cos(z)
# Convert Z and W to planet eccentricities
atan_f_g = T.arctan2(g, f)
c, s = T.cos(atan_f_g), T.sin(atan_f_g)
e1cos = c * Zcosz - s * Wcosw
e1sin = c * Zsinz - s * Wsinw
e2cos = s * Zcosz + c * Wcosw
e2sin = s * Zsinz + c * Wsinw
w1 = T.arctan2(e1sin, e1cos)
w2 = T.arctan2(e2sin, e2cos)
e1 = T.sqrt(e1sin * e1sin + e1cos * e1cos)
e2 = T.sqrt(e2sin * e2sin + e2cos * e2cos)
# Planets' mean longitudes
l1 = Q / k - j * kappa
l2 = Q / k + (k - j) * kappa
# Planets mean anomalies
M1 = l1 - w1
M2 = l2 - w2
# Convert mean to true anomalies using
# function 'exoplanet.theano_ops.kepler.KeplerOp'
ko = KeplerOp()
sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
# Vector of distrubing function values with same dimension as kappa vector
DFfull = calc_DisturbingFunction_with_sinf_cosf(alpha, e1, e2, w1, w2,
sinf1, cosf1, sinf2, cosf2)
# Average distrubing function by weighting values with user-specified
# quadrature weights.
DFav = DFfull.dot(quad_weights)
# Hamiltonian
Hkep = -0.5 * A / k / k * (J - J_star) * (J - J_star)
Hres = -2 * eps * DFav
# ******************IMPORTANT NOTE*************************
# I have *NOT* subtraced off the secular component of
# the disturbing function. This means that the Hamiltonian
# differs slightly from the one defined in the paper.
# This is generally of little consequence to the resonant
# dynamics but should be borne in mind when exploring
# secular dynamics.
# *********************************************************
H = Hkep + Hres
# Gradient and hessian of Hamiltonian w.r.t. phase space variables
gradHtot = T.grad(H, wrt=dyvars)
hessHtot = theano.gradient.hessian(H, wrt=dyvars)
# Flow vector and Jacobian for equations of motion
OmegaTens = T.as_tensor(getOmegaMatrix(2))
H_flow_vec = OmegaTens.dot(gradHtot)
H_flow_jac = OmegaTens.dot(hessHtot)
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# Get numerical quadrature nodes and weights
nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)
# Rescale for integration interval from [-1,1] to [-pi,pi]
nodes = nodes * np.pi
weights = weights * 0.5
# 'ins' will set the inputs of Theano functions compiled below
extra_ins = [m1, m2, j, k, f, g]
ins = [dyvars] + extra_ins
# 'givens' will fix some parameters of Theano functions compiled below
givens = [(kappa, nodes), (quad_weights, weights)]
##########################
# Compile Theano functions
##########################
if not DEBUG:
# Note that compiling can take a while
# so I've put a debugging switch here
# to skip evaluating these functions when
# desired.
H_fn = theano.function(inputs=ins, outputs=H, givens=givens)
H_flow_vec_fn = theano.function(inputs=ins,
outputs=H_flow_vec,
givens=givens)
H_flow_jac_fn = theano.function(inputs=ins,
outputs=H_flow_jac,
givens=givens)
else:
H_fn, H_flow_vec_fn, H_flow_jac_fn = [lambda x: x for _ in range(3)]
# Some convenience functions...
Zsq_to_J_Eq20 = (f * f + g * g) / (fTilde * fTilde + gTilde * gTilde)
dJ_to_Delta_Eq21 = 1.5 * (mu1 + mu2) * (j * mu1 * T.sqrt(alpha) +
(j - k) * mu2) / (
k * T.sqrt(alpha) * mu1 * mu2)
ecc_vars_fn = theano.function(inputs=ins,
outputs=[e1, w1, e2, w2],
on_unused_input='ignore')
Zsq_to_J_Eq20_fn = theano.function(inputs=extra_ins,
outputs=Zsq_to_J_Eq20,
on_unused_input='ignore')
dJ_to_Delta_Eq21_fn = theano.function(inputs=extra_ins,
outputs=dJ_to_Delta_Eq21,
on_unused_input='ignore')
return (H_fn, H_flow_vec_fn, H_flow_jac_fn, Zsq_to_J_Eq20_fn,
dJ_to_Delta_Eq21_fn, ecc_vars_fn)
def get_compiled_theano_functions(N_QUAD_PTS):
# Planet masses: m1,m2
m1,m2 = T.dscalars(2)
mstar = 1
mu1 = m1 * mstar / (mstar + m1)
mu2 = m2 * mstar / (mstar + m2)
Mstar1 = mstar + m1
Mstar2 = mstar + m2
beta1 = mu1 * T.sqrt(Mstar1/mstar) / (mu1 + mu2)
beta2 = mu2 * T.sqrt(Mstar2/mstar) / (mu1 + mu2)
j,k = T.lscalars('jk')
s = (j-k) / k
# Angle variable for averaging over
psi = T.dvector()
# Dynamical variables:
Ndof = 2
Nconst = 1
dyvars = T.vector()
y1, y2, x1, x2, amd = [dyvars[i] for i in range(2*Ndof + Nconst)]
# Quadrature weights
quad_weights = T.dvector('w')
# Set lambda2=0
l2 = T.constant(0.)
l1 = -1 * k * psi
theta_res = (1+s) * l2 - s * l1
cos_theta_res = T.cos(theta_res)
sin_theta_res = T.sin(theta_res)
kappa1 = x1 * cos_theta_res + y1 * sin_theta_res
eta1 = y1 * cos_theta_res - x1 * sin_theta_res
kappa2 = x2 * cos_theta_res + y2 * sin_theta_res
eta2 = y2 * cos_theta_res - x2 * sin_theta_res
Gamma1 = (x1 * x1 + y1 * y1) / 2
Gamma2 = (x2 * x2 + y2 * y2) / 2
# Resonant semi-major axis ratio
alpha_res = ((j-k)/j)**(2/3) * (Mstar1 / Mstar2)**(1/3)
#P0 = k * ( beta2 - beta1 * T.sqrt(alpha_res) ) / 2
#P = P0 - k * (s+1/2) * amd
#Ltot = beta1 * T.sqrt(alpha_res) + beta2 - amd
a20 = 1
a10 = alpha_res * a20
Ltot = beta1 * T.sqrt(a10) + beta2 * np.sqrt(a20)
L1byL2res = beta1 * T.sqrt(a10) / beta2 * np.sqrt(a20)
L2res = (amd + Ltot) / (1 + L1byL2res)
P = 0.5 * L2res * (1 - L1byL2res) - (s+1/2) * amd
L1 = Ltot/2 - P / k - s * (Gamma1 + Gamma2)
L2 = Ltot/2 + P / k + (1 + s) * (Gamma1 + Gamma2)
Xre1 = kappa1 / T.sqrt(L1)
Xim1 = -eta1 / T.sqrt(L1)
Xre2 = kappa2 / T.sqrt(L2)
Xim2 = -eta2 / T.sqrt(L2)
absX1_sq = 2 * Gamma1 / L1
absX2_sq = 2 * Gamma2 / L2
X_to_z1 = T.sqrt(1 - absX1_sq / 4 )
X_to_z2 = T.sqrt(1 - absX2_sq / 4 )
a1 = (L1 / beta1 )**2
k1 = Xre1 * X_to_z1
h1 = Xim1 * X_to_z1
e1 = T.sqrt( absX1_sq ) * X_to_z1
a2 = (L2 / beta2 )**2
k2 = Xre2 * X_to_z2
h2 = Xim2 * X_to_z2
e2 = T.sqrt( absX2_sq ) * X_to_z2
beta1p = T.sqrt(Mstar1) * beta1
beta2p = T.sqrt(Mstar2) * beta2
Hkep = -0.5 * beta1p / a1 - 0.5 * beta2p / a2
Hdir,Hind = calc_Hint_components_planar(
a1,a2,l1,l2,h1,k1,h2,k2,Mstar1/mstar,Mstar2/mstar
)
eps = m1*m2/ (mu1 + mu2) / T.sqrt(mstar)
Hpert = (Hdir + Hind/mstar)
Hpert_av = Hpert.dot(quad_weights)
Htot = Hkep + eps * Hpert_av
######################
# Dissipative dynamics
######################
tau_alpha_0, K1, K2, p = T.dscalars(4)
y1dot_dis,y2dot_dis,x1dot_dis,x2dot_dis,amddot_dis = T.dscalars(5)
tau_m_inv = 1/tau_alpha_0
# Define timescales
tau_e1 = tau_alpha_0 / K1
tau_e2 = tau_alpha_0 / K2
tau_a1_0_inv = -beta2p * alpha_res * tau_m_inv / (beta1p + alpha_res * beta2p)
tau_a2_0_inv = beta1p * tau_m_inv / (beta1p + alpha_res * beta2p)
tau_a1 = 1 / (tau_a1_0_inv + 2 * p * e1*e1 / tau_e1 )
tau_a2 = 1 / (tau_a2_0_inv + 2 * p * e2*e2 / tau_e2 )
tau_L1 = 2 * tau_a1
tau_L2 = 2 * tau_a2
tau_Gamma1_inv = 1/tau_L1 + (Gamma1-2*L1) / (Gamma1-L1) / tau_e1
tau_Gamma2_inv = 1/tau_L2 + (Gamma2-2*L2) / (Gamma2-L2) / tau_e2
# Time derivatives of canonical variables
x1dot_dis = -0.5 * x1 * tau_Gamma1_inv
x2dot_dis = -0.5 * x2 * tau_Gamma2_inv
y1dot_dis = -0.5 * y1 * tau_Gamma1_inv
y2dot_dis = -0.5 * y2 * tau_Gamma2_inv
Pdot_dis = -0.5 * k * ( L2 / tau_L2 - L1 / tau_L1) + k * (s + 1/2) * (Gamma1 * tau_Gamma1_inv + Gamma2 * tau_Gamma2_inv)
amddot_dis = Pdot_dis / T.grad(P,amd)
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# Get numerical quadrature nodes and weight
nodes,weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)
# Rescale for integration interval from [-1,1] to [-pi,pi]
nodes = nodes * np.pi
weights = weights * 0.5
# 'givens' will fix some parameters of Theano functions compiled below
givens = [(psi,nodes),(quad_weights,weights)]
# 'ins' will set the inputs of Theano functions compiled below
# Note: 'extra_ins' will be passed as values of object attributes
# of the 'ResonanceEquations' class defined below
extra_ins = [m1,m2,j,k,tau_alpha_0,K1,K2,p]
ins = [dyvars] + extra_ins
# Define flows and jacobians.
# Conservative flow
gradHtot = T.grad(Htot,wrt=dyvars)
gradHpert = T.grad(Hpert_av,wrt=dyvars)
gradHkep = T.grad(Hkep,wrt=dyvars)
hessHtot = theano.gradient.hessian(Htot,wrt=dyvars)
hessHpert = theano.gradient.hessian(Hpert_av,wrt=dyvars)
hessHkep = theano.gradient.hessian(Hkep,wrt=dyvars)
Jtens = T.as_tensor(np.pad(getOmegaMatrix(Ndof),(0,Nconst),'constant'))
H_flow_vec = Jtens.dot(gradHtot)
Hpert_flow_vec = Jtens.dot(gradHpert)
Hkep_flow_vec = Jtens.dot(gradHkep)
H_flow_jac = Jtens.dot(hessHtot)
Hpert_flow_jac = Jtens.dot(hessHpert)
Hkep_flow_jac = Jtens.dot(hessHkep)
# Dissipative flow
dis_flow_vec = T.stack(y1dot_dis,y2dot_dis,x1dot_dis,x2dot_dis,amddot_dis)
dis_flow_jac = theano.gradient.jacobian(dis_flow_vec,dyvars)
# Extras
sigma1 = T.arctan2(y1,x1)
sigma2 = T.arctan2(y2,x2)
orbels = [a1,e1,k*sigma1,a2,e2,k*sigma2]
dis_timescales = [1/tau_a1_0_inv,1/tau_a2_0_inv,tau_e1,tau_e2]
orbels_dict = dict(zip(
['a1','e1','theta1','a2','e2','theta2'],
orbels
)
)
actions = [L1,L2,Gamma1,Gamma2]
actions_dict = dict(
zip(
['L1','L2','Gamma1','Gamma2'],
actions
)
)
timescales_dict = dict(zip(
['tau_m1','tau_m2','tau_e1','tau_e2'],
dis_timescales
)
)
##########################
# Compile Theano functions
##########################
func_dict={
# Hamiltonians
'H':Htot,
'Hpert':Hpert_av,
'Hkep':Hkep,
# Hamiltonian flows
'H_flow':H_flow_vec,
'Hpert_flow':Hpert_flow_vec,
'Hkep_flow':Hkep_flow_vec,
# Hamiltonian flow Jacobians
'H_flow_jac':H_flow_jac,
'Hpert_flow_jac':Hpert_flow_jac,
'Hkep_flow_jac':Hkep_flow_jac,
# Dissipative flow and Jacobian
'dissipative_flow':dis_flow_vec,
'dissipative_flow_jac':dis_flow_jac,
# Extras
'orbital_elements':orbels_dict,
'actions':actions_dict,
'timescales':timescales_dict
}
compiled_func_dict=dict()
for key,val in func_dict.items():
if key is 'timescales':
inputs = extra_ins
else:
inputs = ins
cf = theano.function(
inputs=inputs,
outputs=val,
givens=givens,
on_unused_input='ignore'
)
compiled_func_dict[key]=cf
return compiled_func_dict
a_val = numpy.array([1, 0, 1, 0])
b_val = numpy.array([0, 1, 0, 1])
x_val = numpy.array([1, 1, 1, 1])
y_val = numpy.array([2, 2, 2, 2])
a, b = T.lvectors('a', 'b')
x, y = T.lvectors('x', 'y')
s1 = T.switch(T.gt(a, b), x, y)
f_s1 = theano.function([a, b, x, y], s1)
result_s1 = f_s1(a_val, b_val, x_val, y_val)
print result_s1 #结果为[1, 2, 1, 2]，进行了element-wise操作
s2 = T.switch(a_val > b_val, [2, 2, 2, 2], [3, 3, 3, 3])
f_s2 = theano.function([], s2)
result_s2 = f_s2()
print result_s2
c, d = T.lscalars('c', 'd')
if1 = ifelse.ifelse(T.gt(c, d), x, y) #相当于T.gt(c, d) ? x : y，不支持element-wise操作
#第一个参数必须返回一个标量
f_if1 = theano.function([c, d, x, y], if1)
result_if1 = f_if1(2, 1, x_val, y_val)
print result_if1
if2 = ifelse.ifelse(1, x_val, y_val) # 参数为常量时
f_if2 = theano.function([], if2)
result_if2 = f_if2()
print result_if2
# i, batch_size = T.lscalars('i', 'batch_size')
# train_step = theano.function([i, batch_size], out_layer.output,
# givens={x : inputs[i:i+batch_size]})
# mode="DebugMode")
params = hidden1r.params + hidden1b.params + hidden2.params + out_layer.params
cost = T.sqrt(T.mean(T.sqr(out_layer.output - y)))
gradients = T.grad(cost, params)
updates = []
for param, grad in zip(params, gradients):
updates.append((param, param - LEARNING_RATE * grad))
i, batch_size = T.lscalars('i', 'batch_size')
train_step = theano.function([i, batch_size], cost,
updates=updates,
givens={x_m : input_m[i:i+batch_size],
x_r : input_r[i:i+batch_size],
y : labels[i:i+batch_size]})
# mode="DebugMode")
n = 0
while True:
n += 1
cost = epoch(100, n_train, train_step)
print "=== epoch {} ===".format(n)
print "costs: {}".format([line[()] for line in cost])
print "avg: {}".format(np.mean(cost))
def _index_variables(self, basename='index'):
return T.lscalars(*['%s_%d' % (basename, i)
for i in xrange(self.sequence_length)])
def _get_compiled_theano_functions(N_QUAD_PTS):
# Planet masses: m1,m2
m1, m2 = T.dscalars(2)
mstar = 1
mu1 = m1 * mstar / (mstar + m1)
mu2 = m2 * mstar / (mstar + m2)
eta1 = mstar + m1
eta2 = mstar + m2
beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)
j, k = T.lscalars('jk')
s = (j - k) / k
# Angle variable for averaging over
psi = T.dvector('psi')
# Quadrature weights
quad_weights = T.dvector('w')
# Dynamical variables:
Ndof = 3
Nconst = 1
dyvars = T.vector()
s1, s2, phi, I1, I2, Phi, dRtilde = [
dyvars[i] for i in range(2 * Ndof + Nconst)
]
a20 = T.constant(1.)
a10 = ((j - k) / j)**(2 / 3) * (eta1 / eta2)**(1 / 3) * a20
L10 = beta1 * T.sqrt(a10)
L20 = beta2 * T.sqrt(a20)
Psi = s * L20 + (1 + s) * L10
Rtilde = dRtilde - L10 - L20
####
# angles
####
rtilde = T.constant(0.)
Omega = -1 * rtilde
l1 = phi + k * (1 + s) * psi + Omega
l2 = phi + k * s * psi + Omega
gamma1 = s1 - phi - Omega
gamma2 = s2 - phi - Omega
q1 = 0.5 * np.pi - Omega
q2 = -0.5 * np.pi - Omega
pomega1 = -1 * gamma1
pomega2 = -1 * gamma2
Omega1 = -1 * q1
Omega2 = -1 * q2
omega1 = pomega1 - Omega1
omega2 = pomega2 - Omega2
###
# actions
###
Gamma1 = I1
Gamma2 = I2
L1 = Psi / k - s * (I1 + I2) - s * Phi
L2 = -1 * Psi / k + (1 + s) * (I1 + I2) + (1 + s) * Phi
Cz = -1 * Rtilde
R = L1 + L2 - Gamma1 - Gamma2 - Cz
G1 = L1 - Gamma1
G2 = L2 - Gamma2
r2_by_r1 = (L2 - L1 - Gamma2 + Gamma1) / (L1 + L2 - Gamma1 - Gamma2 - R)
rho1 = 0.5 * R * (1 + r2_by_r1)
rho2 = 0.5 * R * (1 - r2_by_r1)
a1 = (L1 / beta1)**2
e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)
a2 = (L2 / beta2)**2
e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)
cos_inc1 = 1 - rho1 / G1
cos_inc2 = 1 - rho2 / G2
inc1 = T.arccos(cos_inc1)
inc2 = T.arccos(cos_inc2)
Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2
ko = KeplerOp()
M1 = l1 - pomega1
M2 = l2 - pomega2
sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
#
n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
Hint_dir, Hint_ind, r1, r2, v1, v2 = calc_Hint_components_sinf_cosf(
a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1, Omega2, n1, n2,
sinf1, cosf1, sinf2, cosf2)
eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
Hpert = (Hint_dir + Hint_ind / mstar)
Hpert_av = Hpert.dot(quad_weights)
Htot = Hkep + eps * Hpert_av
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# Get numerical quadrature nodes and weights
nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)
# Rescale for integration interval from [-1,1] to [-pi,pi]
nodes = nodes * np.pi
weights = weights * 0.5
# 'givens' will fix some parameters of Theano functions compiled below
givens = [(psi, nodes), (quad_weights, weights)]
# 'ins' will set the inputs of Theano functions compiled below
# Note: 'extra_ins' will be passed as values of object attributes
# of the 'ResonanceEquations' class 'defined below
extra_ins = [m1, m2, j, k]
ins = [dyvars] + extra_ins
orbels = [a1, e1, inc1, k * s1, a2, e2, inc2, k * s2, phi, Omega]
orbels_dict = dict(
zip([
'a1', 'e1', 'inc1', 'theta1', 'a2', 'e2', 'inc2', 'theta2', 'phi'
], orbels))
actions = [L1, L2, Gamma1, Gamma2, rho1, rho2]
actions_dict = dict(
zip(['L1', 'L2', 'Gamma1', 'Gamma2', 'Q1', 'Q2'], actions))
# Conservative flow
gradHtot = T.grad(Htot, wrt=dyvars)
hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
Jtens = T.as_tensor(
np.pad(_get_Omega_matrix(Ndof), (0, Nconst), 'constant'))
H_flow_vec = Jtens.dot(gradHtot)
H_flow_jac = Jtens.dot(hessHtot)
##########################
# Compile Theano functions
##########################
orbels_fn = theano.function(inputs=ins,
outputs=orbels_dict,
givens=givens,
on_unused_input='ignore')
actions_fn = theano.function(inputs=ins,
outputs=actions_dict,
givens=givens,
on_unused_input='ignore')
Rtilde_fn = theano.function(inputs=ins,
outputs=Rtilde,
givens=givens,
on_unused_input='ignore')
Htot_fn = theano.function(inputs=ins,
outputs=Htot,
givens=givens,
on_unused_input='ignore')
Hpert_fn = theano.function(inputs=ins,
outputs=Hpert_av,
givens=givens,
on_unused_input='ignore')
Hpert_components_fn = theano.function(
inputs=ins,
outputs=[Hint_dir.dot(quad_weights),
Hint_ind.dot(quad_weights)],
givens=givens,
on_unused_input='ignore')
H_flow_vec_fn = theano.function(inputs=ins,
outputs=H_flow_vec,
givens=givens,
on_unused_input='ignore')
H_flow_jac_fn = theano.function(inputs=ins,
outputs=H_flow_jac,
givens=givens,
on_unused_input='ignore')
return dict({
'orbital_elements': orbels_fn,
'actions': actions_fn,
'Rtilde': Rtilde_fn,
'Hamiltonian': Htot_fn,
'Hpert': Hpert_fn,
'Hpert_components': Hpert_components_fn,
'Hamiltonian_flow': H_flow_vec_fn,
'Hamiltonian_flow_jacobian': H_flow_jac_fn
})
def _get_compiled_theano_functions():
# Planet masses: m1,m2
m1, m2 = T.dscalars(2)
mstar = 1
mu1 = m1 * mstar / (mstar + m1)
mu2 = m2 * mstar / (mstar + m2)
eta1 = mstar + m1
eta2 = mstar + m2
beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)
j, k = T.lscalars('jk')
s = (j - k) / k
# Dynamical variables:
dyvars = T.vector()
s1, s2, psi, phi, Omega, I1, I2, Psi, Phi, Rtilde = [
dyvars[i] for i in range(10)
]
l1 = phi - 0.5 * k * psi
l2 = phi + 0.5 * k * psi
gamma1 = s1 - (1 + s) * l2 + s * l1
gamma2 = s2 - (1 + s) * l2 + s * l1
Gamma1 = I1
Gamma2 = I2
L1 = Phi / 2 - Psi / k - s * (I1 + I2)
L2 = Phi / 2 + Psi / k + (s + 1) * (I1 + I2)
Cz = -1 * Rtilde
R = L1 + L2 - Gamma1 - Gamma2 - Cz
G1 = L1 - Gamma1
G2 = L2 - Gamma2
r2_by_r1 = (L2 - L1 - Gamma2 + Gamma1) / (L1 + L2 - Gamma1 - Gamma2 - R)
rho1 = 0.5 * R * (1 + r2_by_r1)
rho2 = 0.5 * R * (1 - r2_by_r1)
a1 = (L1 / beta1)**2
e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)
a2 = (L2 / beta2)**2
e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)
cos_inc1 = 1 - rho1 / G1
cos_inc2 = 1 - rho2 / G2
inc1 = T.arccos(cos_inc1)
inc2 = T.arccos(cos_inc2)
l1_r = l1 - Omega
l2_r = l2 - Omega
Omega1_r = T.constant(np.pi / 2) - Omega
Omega2_r = Omega1_r - T.constant(np.pi)
pomega1 = -1 * gamma1
pomega2 = -1 * gamma2
pomega1_r = pomega1 - Omega
pomega2_r = pomega2 - Omega
omega1 = pomega1_r - Omega1_r
omega2 = pomega2_r - Omega2_r
Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2
ko = KeplerOp()
M1 = l1_r - pomega1_r
M2 = l2_r - pomega2_r
sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
#
n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
Hint_dir, Hint_ind, r1, r2, v1, v2 = calc_Hint_components_sinf_cosf(
a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1_r, Omega2_r, n1, n2,
sinf1, cosf1, sinf2, cosf2)
eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar)
Hpert = (Hint_dir + Hint_ind / mstar)
Htot = Hkep + eps * Hpert
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# 'ins' will set the inputs of Theano functions compiled below
# Note: 'extra_ins' will be passed as values of object attributes
# of the 'ResonanceEquations' class 'defined below
extra_ins = [m1, m2, j, k]
givens = []
ins = [dyvars] + extra_ins
orbels = [
a1, e1, inc1, l1_r, pomega1_r, Omega1_r, a2, e2, inc2, l2_r, pomega2_r,
Omega2_r
]
# Conservative flow
gradHtot = T.grad(Htot, wrt=dyvars)
hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
Jtens = T.as_tensor(_get_Omega_matrix(5))
H_flow_vec = Jtens.dot(gradHtot)
H_flow_jac = Jtens.dot(hessHtot)
##########################
# Compile Theano functions
##########################
orbels_fn = theano.function(inputs=ins,
outputs=orbels,
givens=givens,
on_unused_input='ignore')
rv1_fn = theano.function(inputs=ins,
outputs=r1 + v1,
givens=givens,
on_unused_input='ignore')
rv2_fn = theano.function(inputs=ins,
outputs=r2 + v2,
givens=givens,
on_unused_input='ignore')
Htot_fn = theano.function(inputs=ins,
outputs=Htot,
givens=givens,
on_unused_input='ignore')
Hpert_fn = theano.function(inputs=ins,
outputs=Hpert,
givens=givens,
on_unused_input='ignore')
Hpert_components_fn = theano.function(inputs=ins,
outputs=[Hint_dir, Hint_ind],
givens=givens,
on_unused_input='ignore')
H_flow_vec_fn = theano.function(inputs=ins,
outputs=H_flow_vec,
givens=givens,
on_unused_input='ignore')
H_flow_jac_fn = theano.function(inputs=ins,
outputs=H_flow_jac,
givens=givens,
on_unused_input='ignore')
return dict({
'orbital_elements': orbels_fn,
'Hamiltonian': Htot_fn,
'Hpert': Hpert_fn,
'Hpert_components': Hpert_components_fn,
'Hamiltonian_flow': H_flow_vec_fn,
'Hamiltonian_flow_jacobian': H_flow_jac_fn,
'positions_and_velocities1': rv1_fn,
'positions_and_velocities2': rv2_fn
})
def _get_compiled_theano_functions(N_QUAD_PTS):
# resonance j and k
j, k = T.lscalars('jk')
s = (j - k) / k
# Planet masses: m1,m2
m1, m2 = T.dscalars(2)
mstar = 1
mu1 = m1 * mstar / (mstar + m1)
mu2 = m2 * mstar / (mstar + m2)
eta1 = mstar + m1
eta2 = mstar + m2
beta1 = mu1 * T.sqrt(eta1 / mstar) / (mu1 + mu2)
beta2 = mu2 * T.sqrt(eta2 / mstar) / (mu1 + mu2)
# Angle variable for averaging over
psi = T.dvector('psi')
# Dynamical variables:
dyvars = T.vector()
y1, y2, phi1, x1, x2, J1, amd = [dyvars[i] for i in range(7)]
I1 = (y1 * y1 + x1 * x1) / 2
I2 = (y2 * y2 + x2 * x2) / 2
sigma1 = T.arctan2(y1, x1)
sigma2 = T.arctan2(y2, x2)
# Quadrature weights
quad_weights = T.dvector('w')
# Set l=0
l = T.constant(0.)
l1 = l - k * psi / 2
l2 = l + k * psi / 2
pomega1 = (1 + s) * l2 - s * l1 - sigma1
pomega2 = (1 + s) * l2 - s * l1 - sigma2
Gamma1 = I1
Gamma2 = I2
# Resonant semi-major axis ratio
a20 = T.constant(1.0)
a10 = (eta1 / eta2)**(1 / 3) * ((j - k) / j)**(2 / 3) * a20
Lambda20 = beta2 * T.sqrt(a20)
Lambda10 = beta1 * T.sqrt(a10)
Ltot = Lambda10 + Lambda20
Psi0 = 0.5 * k * (Lambda20 - Lambda10)
Psi = Psi0 - k * (s + 1 / 2) * amd
L1 = Ltot / 2 + J1 / 2 - Psi / k - s * (I1 + I2)
L2 = Ltot / 2 + J1 / 2 + Psi / k + (1 + s) * (I1 + I2)
# Choose z axis along direction of total angular momentum
r2_by_r1 = (L2 - L1 - Gamma2 + Gamma1) / (L1 + L2 - Gamma1 - Gamma2 - J1)
rho1 = 0.5 * J1 * (1 + r2_by_r1)
rho2 = 0.5 * J1 * (1 - r2_by_r1)
G1 = L1 - Gamma1
G2 = L2 - Gamma1
a1 = (L1 / beta1)**2
e1 = T.sqrt(1 - (1 - (Gamma1 / L1))**2)
a2 = (L2 / beta2)**2
e2 = T.sqrt(1 - (1 - (Gamma2 / L2))**2)
cos_inc1 = 1 - rho1 / G1
cos_inc2 = 1 - rho2 / G2
inc1 = T.arccos(cos_inc1)
inc2 = T.arccos(cos_inc2)
Omega1 = l - np.pi / 2 - phi1
Omega2 = l + np.pi / 2 - phi1
omega1 = pomega1 - Omega1
omega2 = pomega2 - Omega2
Hkep = -0.5 * T.sqrt(eta1) * beta1 / a1 - 0.5 * T.sqrt(eta2) * beta2 / a2
ko = KeplerOp()
M1 = l1 - pomega1
M2 = l2 - pomega2
sinf1, cosf1 = ko(M1, e1 + T.zeros_like(M1))
sinf2, cosf2 = ko(M2, e2 + T.zeros_like(M2))
#
n1 = T.sqrt(eta1 / mstar) * a1**(-3 / 2)
n2 = T.sqrt(eta2 / mstar) * a2**(-3 / 2)
Hint_dir, Hint_ind = calc_Hint_components_sinf_cosf(
a1, a2, e1, e2, inc1, inc2, omega1, omega2, Omega1, Omega2, n1, n2,
sinf1, cosf1, sinf2, cosf2)
eps = m1 * m2 / (mu1 + mu2) / T.sqrt(mstar * a20)
Hpert = (Hint_dir + Hint_ind / mstar).dot(quad_weights)
Htot = Hkep + eps * Hpert
#####################################################
# Set parameters for compiling functions with Theano
#####################################################
# Get numerical quadrature nodes and weights
nodes, weights = np.polynomial.legendre.leggauss(N_QUAD_PTS)
# Rescale for integration interval from [-1,1] to [-pi,pi]
nodes = nodes * np.pi
weights = weights * 0.5
# 'givens' will fix some parameters of Theano functions compiled below
givens = [(psi, nodes), (quad_weights, weights)]
# 'ins' will set the inputs of Theano functions compiled below
# Note: 'extra_ins' will be passed as values of object attributes
# of the 'ResonanceEquations' class 'defined below
extra_ins = [m1, m2, j, k]
ins = [dyvars] + extra_ins
orbels = [
a1, e1, inc1, sigma1 * k, Omega1, a2, e2, inc2, sigma2 * k, Omega2
]
# Conservative flow
gradHtot = T.grad(Htot, wrt=dyvars)
hessHtot = theano.gradient.hessian(Htot, wrt=dyvars)
Jtens = T.as_tensor(np.pad(_get_Omega_matrix(3), (0, 1), 'constant'))
H_flow_vec = Jtens.dot(gradHtot)
H_flow_jac = Jtens.dot(hessHtot)
##########################
# Compile Theano functions
##########################
actions_fn = theano.function(
inputs=ins,
outputs={'L1':L1,'L1':L2,'Gamma1':Gamma1,'Gamma2':Gamma2,\
'Q1':rho1,'Q2':rho2,'Psi':Psi,'Psi0':Psi0},
givens=givens,
on_unused_input='ignore'
)
orbels_fn = theano.function(inputs=ins,
outputs=orbels,
givens=givens,
on_unused_input='ignore')
Htot_fn = theano.function(inputs=ins,
outputs=Htot,
givens=givens,
on_unused_input='ignore')
Hpert_fn = theano.function(inputs=ins,
outputs=Hpert,
givens=givens,
on_unused_input='ignore')
H_flow_vec_fn = theano.function(inputs=ins,
outputs=H_flow_vec,
givens=givens,
on_unused_input='ignore')
H_flow_jac_fn = theano.function(inputs=ins,
outputs=H_flow_jac,
givens=givens,
on_unused_input='ignore')
return dict({
'orbital_elements': orbels_fn,
'actions': actions_fn,
'Hamiltonian': Htot_fn,
'Hpert': Hpert_fn,
'Hamiltonian_flow': H_flow_vec_fn,
'Hamiltonian_flow_jacobian': H_flow_jac_fn
})