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 setUp(self):
super(TestKeplerSolver, self).setUp()
self.op_class = KeplerOp
self.op = KeplerOp()
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_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():
# 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)
# Dynamical variables:
dyvars = T.vector()
l1, l2, y1, y2, Omega, L1, L2, x1, x2, Rtilde = [
dyvars[i] for i in range(10)
]
Gamma1 = 0.5 * (x1 * x1 + y1 * y1)
Gamma2 = 0.5 * (x2 * x2 + y2 * y2)
gamma1 = T.arctan2(y1, x1)
gamma2 = T.arctan2(y2, x2)
Cz = -1 * Rtilde
R1 = Gamma1
R2 = Gamma2
R = L1 + L2 - R1 - R2 - Cz
G1 = L1 - Gamma1
G2 = L2 - Gamma2
r2_by_r1 = (L2 - L1 - R2 + R1) / (L1 + L2 - R1 - R2 - 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]
givens = []
ins = [dyvars] + extra_ins
orbels = [
a1, e1, inc1, l1_r, pomega1_r, Omega1_r, a2, e2, inc2, l2_r, pomega2_r,
Omega2_r
]
orbels_dict = dict(
zip([
'a1', 'e1', 'inc1', 'l1', 'pomega1', 'Omega1', 'a2', 'e2', 'inc2',
'l2', 'pomega2', 'Omega2'
], 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(_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_dict,
givens=givens,
on_unused_input='ignore')
actions_fn = theano.function(inputs=ins,
outputs=actions_dict,
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,
'actions': actions_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):
# 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
})