def setUp(self):
self.sim = rebound.Simulation()
self.rebx = reboundx.Extras(self.sim)
data.add_earths(self.sim, ei=1.e-3)
self.array = np.array([1., 2., 3., 4.])
self.intarray = np.array([1, 2, 3, 4])
self.cuintarray = np.array(
[c_uint(1), c_uint(2), c_uint(3),
c_uint(4)], dtype=object)
self.cuint32array = np.array(
[c_uint32(1), c_uint32(2),
c_uint32(3), c_uint32(4)], dtype=object)
self.ndarray = np.array([[[1., 2., 3.], [4., 5., 6.]],
[[1., 2., 3.], [4., 5., 6.]],
[[1., 2., 3.], [4., 5., 6.]],
[[1., 2., 3.], [4., 5., 6.]]])
self.orbitarray = np.array(self.sim.calculate_orbits(), dtype=object)
o1 = self.sim.particles[1].orbit
o2 = self.sim.particles[2].orbit
self.objndarray = np.array(
[[[o1, o2], [o2, o1], [o1, o1]], [[o1, o2], [o2, o1], [o1, o1]],
[[o1, o2], [o2, o1], [o1, o1]], [[o1, o2], [o2, o1], [o1, o1]]],
dtype=object)
self.sim.particles[0].params["array"] = self.array
self.sim.particles[0].params["intarray"] = self.intarray
self.sim.particles[0].params["cuintarray"] = self.cuintarray
self.sim.particles[0].params["cuint32array"] = self.cuint32array
self.sim.particles[1].params["ndarray"] = self.ndarray
self.sim.particles[0].params["orbitarray"] = self.orbitarray
self.sim.particles[1].params["objndarray"] = self.objndarray
self.sim.particles[1].params["scalar"] = 3
def test_modify_mass(self):
name = 'modify_mass'
try:
rebbin = os.path.join(THIS_DIR, 'binaries/' + name + '.sa')
rebxbin = os.path.join(THIS_DIR, 'binaries/' + name + '.rebx')
sa = reboundx.SimulationArchive(rebbin, rebxbin)
except:
sim = rebound.Simulation('binaries/twoplanets.bin')
sim.automateSimulationArchive('binaries/' + name + '.sa',
interval=1e3,
deletefile=True)
rebx = reboundx.Extras(sim)
mod = rebx.load_operator(name)
rebx.add_operator(mod)
ps = sim.particles
ps[0].params['tau_mass'] = -1e4
rebx.save('binaries/' + name + '.rebx')
sim.integrate(1.e4)
sa = reboundx.SimulationArchive('binaries/' + name + '.sa',
'binaries/' + name + '.rebx')
simf, rebx = sa[-1]
sim, rebx = sa[0]
sim.integrate(simf.t)
self.assertEqual(sim.particles[0].x, simf.particles[0].x)
def run_constant_damping(sim,
res,
times,
tau_e=1e2,
tau_a=1e3,
chaos=False,
stopping_criterion=mmr_stop,
direct=False):
tau_a = res.tau_a
tau_e = res.tau_e
res.direct = res.direct
rebx = rbx.Extras(sim)
if direct:
params = rebx.add_modify_orbits_direct()
else:
params = rebx.add_modify_orbits_forces()
sim.particles[3].tau_a = -tau_a
sim.particles[3].tau_e = -tau_e
res = run_sim(sim,
res,
times,
chaos=chaos,
stopping_criterion=stopping_criterion)
return res
def test_backwards_time(self):
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
times = [0, 2000., 4000., 6000., 8000., 10000.]
values = [1., 0.8, 0.6, 0.4, 0.3, 0.2]
starmass = reboundx.Interpolator(rebx, times, values, "spline")
Nout = 1000
ts = np.linspace(0., 1.e4, Nout)
ps = sim.particles
a10 = ps[1].a
m0 = ps[0].m
for i, time in enumerate(ts):
sim.integrate(time)
ps[0].m = starmass.interpolate(rebx, t=sim.t)
sim.move_to_com(
) # lost mass had momentum, so need to move back to COM frame
sim.dt = -0.001
for i, time in enumerate(times[::-1]):
sim.integrate(time)
ps[0].m = starmass.interpolate(rebx, t=sim.t)
sim.move_to_com(
) # lost mass had momentum, so need to move back to COM frame
self.assertLess(abs((ps[0].m - m0) / m0), 1.e-2)
self.assertLess(abs((ps[1].a - a10) / a10), 1.e-2)
def test_gr(self):
name = 'gr'
try:
rebbin = os.path.join(THIS_DIR, 'binaries/' + name + '.sa')
rebxbin = os.path.join(THIS_DIR, 'binaries/' + name + '.rebx')
sa = reboundx.SimulationArchive(rebbin, rebxbin)
except:
sim = rebound.Simulation('binaries/twoplanets.bin')
sim.automateSimulationArchive('binaries/' + name + '.sa',
interval=1e3,
deletefile=True)
rebx = reboundx.Extras(sim)
force = rebx.load_force(name)
rebx.add_force(force)
force.params['c'] = 1.e4
ps = sim.particles
rebx.save('binaries/' + name + '.rebx')
sim.integrate(1.e4)
sa = reboundx.SimulationArchive('binaries/' + name + '.sa',
'binaries/' + name + '.rebx')
simf, rebx = sa[-1]
sim, rebx = sa[0]
sim.integrate(simf.t)
self.assertEqual(sim.particles[0].x, simf.particles[0].x)
def setUp(self):
self.sim = rebound.Simulation()
self.rebx = reboundx.Extras(self.sim)
data.add_earths(self.sim, ei=1.e-3)
self.sim.particles[0].params["a"] = 1.2
self.sim.particles[0].params["b"] = 1.7
self.sim.particles[0].params["N"] = 14
def setUp(self):
self.sim = rebound.Simulation()
data.add_earths(self.sim, ei=1.e-3)
self.rebx = reboundx.Extras(self.sim)
self.gr = self.rebx.add("gr")
self.gr.params["a"] = 1.2
self.gr.params["b"] = 1.7
self.gr.params["N"] = 14
def setUp(self):
self.sim = rebound.Simulation()
self.sim.add(m=1., r=0.005)
self.sim.add(m=1.e-3, a=0.05, e=0.1, r=0.0005)
self.sim.move_to_com()
self.rebx = reboundx.Extras(self.sim)
self.force = self.rebx.load_force("tides_constant_time_lag")
self.rebx.add_force(self.force)
def setUp(self):
self.sim = rebound.Simulation()
self.sim.add(m=1.)
self.sim.add(a=1.)
self.rebx = reboundx.Extras(self.sim)
self.gr = self.rebx.load_force("gr")
self.mm = self.rebx.load_operator("modify_mass")
self.p = self.sim.particles[1]
def addGR(sim, rebxintegrator, order, cfac):
rebx = reboundx.Extras(sim)
gr = rebx.add("gr")
gr.params["c"] = 63197.8 * cfac # AU/yr
if rebxintegrator != "naive":
rebx.integrator = rebxintegrator
gr.operator_order = order
gr.force_as_operator = 1
return rebx
def makesim(init_a):
sim = rebound.Simulation()
sim.units = ('yr', 'AU', 'Msun')
sim.add(m=M0)
sim.add(m=3e-6, a=init_a) # 1 Mearth
sim.collision = "direct"
rebx = reboundx.Extras(sim)
# tides = rebx.load_force("tides_constant_time_lag")
# rebx.add_force(tides)
return sim, rebx #, tides
def makesim():
sim = rebound.Simulation()
sim.units = ('yr', 'AU', 'Msun')
sim.add(m=M0)
sim.add(m=1e-3, a=5)
sim.collision = "direct"
rebx = reboundx.Extras(sim)
tides = rebx.load_force("tides_constant_time_lag")
rebx.add_force(tides)
return sim, rebx, tides
def setUp(self):
self.sim = rebound.Simulation()
self.sim.add(m=1.)
self.sim.add(m=1.e-4, a=1., e=0.2)
self.sim.move_to_com()
self.sim.dt = 1.e-2 * self.sim.particles[1].P
self.rebx = reboundx.Extras(self.sim)
self.gr = self.rebx.load_force('gr')
self.gr.params['c'] = 1e2
self.pomega0 = self.sim.particles[1].pomega
self.E0 = self.rebx.gr_hamiltonian(self.gr)
def test_conservation_planet_highmratio(self):
self.sim = rebound.Simulation()
self.sim.add(m=1., r=0.005)
self.sim.add(m=1., a=0.2, e=0.1, r=0.005)
self.sim.move_to_com()
self.rebx = reboundx.Extras(self.sim)
self.force = self.rebx.load_force("tides_constant_time_lag")
self.rebx.add_force(self.force)
ps = self.sim.particles
ps[1].params['tctl_k1'] = 0.04
self.do_test_conservation()
def test_gr(self):
name = 'gr'
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
force = rebx.load_force(name)
rebx.add_force(force)
force.params['c'] = 1.e4
ps = sim.particles
H0 = rebx.gr_hamiltonian(force)
sim.integrate(1.e4)
H = rebx.gr_hamiltonian(force)
self.assertLess(abs((H - H0) / H0), 1.e-12)
def test_gr_potential(self):
name = 'gr_potential'
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
force = rebx.load_force(name)
rebx.add_force(force)
force.params['c'] = 1.e4
ps = sim.particles
H0 = sim.calculate_energy() + rebx.gr_potential_potential(force)
sim.integrate(1.e4)
H = sim.calculate_energy() + rebx.gr_potential_potential(force)
self.assertLess(abs((H - H0) / H0), 1.e-12)
def setUp(self):
self.sim = rebound.Simulation()
self.sim.add(m=1.)
self.sim.add(a=1., e=0.2)
self.rebx = reboundx.Extras(self.sim)
self.cust = self.rebx.create_operator('addop')
self.rebx.register_param('ctr', 'REBX_TYPE_INT')
self.cust.params['ctr'] = 0
def mystep(sim, operator, dt):
operator.contents.params['ctr'] += 1
self.cust.step_function = mystep
def test_tides_precession(self):
name = 'tides_precession'
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
force = rebx.load_force(name)
rebx.add_force(force)
ps = sim.particles
ps[0].params['R_tides'] = 1.e-3
ps[0].params['k1'] = 0.4
H0 = sim.calculate_energy() + rebx.tides_precession_potential(force)
sim.integrate(1.e4)
H = sim.calculate_energy() + rebx.tides_precession_potential(force)
self.assertLess(abs((H - H0) / H0), 1.e-12)
def test_central_force(self):
name = 'central_force'
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
force = rebx.load_force(name)
rebx.add_force(force)
ps = sim.particles
ps[0].params['Acentral'] = 1.e-4
ps[0].params['gammacentral'] = -1
H0 = sim.calculate_energy() + rebx.central_force_potential()
sim.integrate(1.e4)
H = sim.calculate_energy() + rebx.central_force_potential()
self.assertLess(abs((H - H0) / H0), 1.e-12)
def test_gravitational_harmonics(self):
name = 'gravitational_harmonics'
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
force = rebx.load_force(name)
rebx.add_force(force)
ps = sim.particles
ps[0].params['J2'] = 1.e-3
ps[0].params['J4'] = 1.e-3
ps[0].params['R_eq'] = 1.e-3
H0 = sim.calculate_energy() + rebx.gravitational_harmonics_potential()
sim.integrate(1.e4)
H = sim.calculate_energy() + rebx.gravitational_harmonics_potential()
self.assertLess(abs((H - H0) / H0), 1.e-12)
def add_force(self):
"""Add tau damping force"""
if self.verbose:
print('adding modify_orbits_forces_edge')
self.rebx = reboundx.Extras(self.sim)
mof = self.rebx.load_force('modify_orbits_forces_edge')
mof.params['res_redge'] = self.params['redge']
mof.params['res_deltaredge'] = self.params['deltaredge']
mof.params['res_tdep'] = self.params['tdep']
mof.params['res_deltatdep'] = self.params['deltatdep']
self.rebx.add_force(mof)
self.mof = mof
for pt in self.sim.particles[1:]:
pt.params['tau_a'] = self.params['tau_a']
pt.params['tau_e'] = self.params['tau_e']
pt.params['tau_inc'] = self.params['tau_inc']
def add_force(self):
"""Add reboundx forces. This is a method so it can be overloaded easily"""
if self.verbose:
print(' adding modify_orbits_resonance_relax')
self.rebx = reboundx.Extras(self.sim)
mof = self.rebx.load_force('modify_orbits_resonance_relax')
mof.params['res_aspectratio0'] = self.params['aspectratio0']
mof.params['res_sigma0'] = self.params['sigma0']
mof.params['res_redge'] = self.params['redge']
mof.params['res_deltaredge'] = self.params['deltaredge']
mof.params['res_alpha'] = self.params['alpha']
mof.params['res_flaringindex'] = self.params['flaringindex']
mof.params['res_ffudge'] = self.params['ffudge']
mof.params['res_tdep'] = self.params['tdep']
mof.params['res_deltatdep'] = self.params['deltatdep']
self.rebx.add_force(mof)
self.mof = mof
def setUp(self):
self.sim = rebound.Simulation()
self.sim.add(m=0.86, r = 0.78)
self.sim.add(m=3.e-6, a=1., e=0.05)
self.sim.move_to_com()
ps = self.sim.particles
self.rebx = reboundx.Extras(self.sim)
self.tides = self.rebx.load_force("tides_constant_time_lag")
self.rebx.add_force(self.tides)
ps[0].params["tctl_k1"] = 0.023 # in AU
ps[0].params["tctl_tau"] = 0.3
ps[0].params["Omega"] = 0.
self.q = (ps[1].m/ps[0].m)
self.T = ps[0].r**3/self.sim.G/ps[0].m/ps[0].params["tctl_tau"]
self.taua = self.T/6/ps[0].params["tctl_k1"]/self.q/(1+self.q)*(ps[1].a/ps[0].r)**8
def test_klo(self):
sim = rebound.Simulation(binary)
sim.integrator = "ias15"
rebx = reboundx.Extras(sim)
times = [0, 2000., 4000., 6000., 8000., 10000.]
values = [1., 0.8, 0.6, 0.4, 0.3, 0.2]
starmass = reboundx.Interpolator(rebx, times, values, "spline")
m0 = sim.particles[0].m
mint0 = starmass.interpolate(rebx, t=0)
self.assertLess(abs((m0 - mint0) / m0), 1.e-6)
mint1 = starmass.interpolate(rebx, t=5000) # advance klo
mint1 = starmass.interpolate(rebx, t=0) # advance klo
self.assertLess(abs((m0 - mint1) / m0), 1.e-6)
def test_reproducibility(self):
sim = rebound.Simulation()
sim.add(m=1.)
sim.add(m=1e-3, a=1.)
rebx = reboundx.Extras(sim)
gr = rebx.add("gr")
gr.params['c'] = 1.e3
sim.automateSimulationArchive("test.bin",
interval=1e3,
deletefile=True)
sim.integrate(1.e4)
rebx.save("rebx.bin")
sa = rebound.SimulationArchive("test.bin", rebxfilename="rebx.bin")
for i in [2, 5, 7]:
sim = sa[i]
x = sim.particles[1].x
sim = sa[i - 1]
sim.integrate(sa[i].t, exact_finish_time=0)
self.assertEqual(x, sim.particles[1].x)
def test_modify_orbits_forces(self):
name = "modify_orbits_forces"
try:
rebbin = os.path.join(THIS_DIR, 'binaries/' + name + '.sa')
rebxbin = os.path.join(THIS_DIR, 'binaries/' + name + '.rebx')
sa = reboundx.SimulationArchive(rebbin, rebxbin)
except:
try:
sim = rebound.Simulation('binaries/twoplanets.bin')
except:
sim = rebound.Simulation()
sim.add(m=1.)
sim.add(m=1.e-4, a=1., e=0.1)
sim.add(m=1.e-4, a=2, e=0.1, inc=0.2)
sim.integrator = "whfast"
sim.dt = sim.particles[1].P / 100
sim.move_to_com()
sim.save('binaries/twoplanets.bin')
sim.automateSimulationArchive('binaries/modify_orbits_forces.sa',
interval=1e3,
deletefile=True)
rebx = reboundx.Extras(sim)
mod = rebx.load_force('binaries/modify_orbits_forces')
rebx.add_force(mod)
ps = sim.particles
ps[1].params['tau_a'] = -1e4
ps[1].params['tau_e'] = -1e3
ps[2].params['tau_e'] = -1e3
ps[2].params['tau_inc'] = -1e3
rebx.save('binaries/modify_orbits_forces.rebx')
sim.integrate(1.e4)
sa = reboundx.SimulationArchive(
'binaries/modify_orbits_forces.sa',
'binaries/modify_orbits_forces.rebx')
simf, rebx = sa[-1]
sim, rebx = sa[0]
sim.integrate(simf.t)
self.assertEqual(sim.particles[0].x, simf.particles[0].x)
if keep_going == 'n':
raise Exception("A keyboard interruption occurred.")
break
unstable = 'n' #### initialize assuming stability, until shown otherwise.
try:
running_period_list = np.load(projectdir+'/running_list_of_derived_TTV_periods.npy').tolist()
except:
running_period_list = []
### create a simulations instance.
sim = rebound.Simulation()
if include_J2 == 'y':
rebx = reboundx.Extras(sim)
gh = rebx.load_force('gravitational_harmonics')
rebx.add_force(gh)
stability_model = FeatureClassifier()
### put these in familiar units
sim.G = G.value
sim.units = ('s', 'm', 'kg') ### mks
#### now let's add Jupiter, Io, Europa, Ganymede, and Callisto!
### Jupiter
if make_imaginary_system == 'n':
### MAKE JUPITER!
def dyvars_to_rebound_simulation(self,z,Q=0,pomega1=0,osculating_correction = True,include_dissipation = False,**kwargs):
r"""
Convert dynamical variables
.. math:
z = (\sigma_1,\sigma_2,I_1,I_2,{\cal C}
to a Rebound simulation.
Arguments
---------
z : ndarray
Dynamical variables
pomega1 : float, optional
Planet 1's longitude of periapse.
Default is 0
Q : float, optional
Angle variable Q = (lambda2-lambda1) / k.
Default is Q=0. Resonant periodic orbits
are 2pi-periodic in Q.
include_dissipation : bool, optional
Include dissipative effects through
reboundx's external forces.
Default is False
Keyword Arguments
-----------------
Mstar : float
Default=1.0
Stellar mass.
inc : float, default = 0.0
Inclination of planets' orbits.
period1 : float
Default = 1.0
Orbital period of inner planet
units : tuple
Default = ('AU','Msun','days')
Units of rebound simulation.
Returns
-------
tuple :
Returns a tuple. The first item
of the tuple is a rebound simulation.
The second item is a reboundx.Extras object
if 'include_dissipation' is True, otherwise
the second item is None.
"""
mean_orbels = self.dyvars_to_orbels(z)
a1_mean,_,_,a2_mean,_,_ = mean_orbels
if osculating_correction:
zosc = self.mean_to_osculating_dyvars(Q,z)
orbels = self.dyvars_to_orbels(zosc)
else:
orbels = mean_orbels
j,k = self.j, self.k
a1,e1,theta1,a2,e2,theta2 = orbels
syn_angle = self.k * Q
pomega2 = np.mod( pomega1 - (theta2 - theta1) / k, 2*np.pi)
M1 = np.mod( (theta1 - j*syn_angle) / k ,2*np.pi )
l1 = np.mod(M1 + pomega1,2*np.pi)
l2 = np.mod( syn_angle + l1,2*np.pi)
Mstar = kwargs.pop('Mstar',1.0)
inc = kwargs.pop('inc',0.0)
period1 = kwargs.pop('period1',1.0)
units = kwargs.pop('units', ('AU','Msun','days'))
sim = rb.Simulation()
sim.units = units
mpl1 = self.m1 * Mstar
mpl2 = self.m2 * Mstar
# Rescale distance scale so that the proper semi-major axis of planet 1 corresponds to orbital period 'period1'
a1_physical = (a1 / a1_mean) * (sim.G * (Mstar + mpl1) * period1**2 / (4*np.pi*np.pi))**(1/3)
a2_physical = a2 * a1_physical / a1
sim.add(m=Mstar)
sim.add(m = mpl1, a=a1_physical, e=e1, l=l1, pomega=pomega1, inc=inc, jacobi_masses=True)
sim.add(m = mpl2, a=a2_physical, e=e2, l=l2, pomega=pomega2, inc=inc, jacobi_masses=True)
sim.move_to_com()
rebx = None
if include_dissipation:
ps = sim.particles
n2 = ps[2].n
rebx = reboundx.Extras(sim)
mod = rebx.load_operator("modify_orbits_direct")
rebx.add_operator(mod)
mod.params["p"] = self.p
timescales = self.timescales
ps[1].params["tau_a"]=-1*timescales['tau_a1'] / n2
ps[2].params["tau_a"]=-1*timescales['tau_a2'] / n2
ps[1].params["tau_e"]=-1*timescales['tau_e1'] / n2
ps[2].params["tau_e"]=-1*timescales['tau_e2'] / n2
return sim,rebx
def integrate_orbit(self,
t_val,
M0,
R0,
a,
e,
inc,
Omega,
omega,
tp,
x0=0,
y0=0,
vx0=0,
vy0=0,
vz0=0,
**effect_kwargs):
# account for factor of R0 in l_unit
gravitational_constant = gravitational_constant_times_R03 / R0**3
speed_of_light = speed_of_light_times_R0 / R0
velocity_conversion_factor = velocity_conversion_factor_per_R0 * R0
t_val = np.sort(t_val)
sim = rebound.Simulation()
rebx = reboundx.Extras(sim)
sim.integrator = "ias15"
sim.G = gravitational_constant
sim.t = t_val[0]
# enable post-Newtonian corrections
gr = rebx.load_force("gr")
rebx.add_force(gr)
gr.params["c"] = speed_of_light
# add black hole particle
sim.add(hash="black_hole", m=M0) # placed at the origin
bh = sim.particles["black_hole"]
bh.params["gr_source"] = 1
# add star "test" particle
sim.add(hash="test_particle",
m=0,
a=a,
e=e,
inc=inc,
Omega=Omega,
omega=omega,
T=tp)
p = sim.particles["test_particle"]
if "Mp" and "rp" in effect_kwargs:
# account for an extended mass distribution (Plummer profile)
mass_effect = rebx.add_custom_force(
central_force, force_is_velocity_dependent=False)
mass_effect.params["Mp"] = effect_kwargs["Mp"]
mass_effect.params["rp"] = effect_kwargs["rp"]
data = pd.DataFrame(index=t_val, columns=["t'", "x", "y", "vz"])
for step, t_obs in enumerate(t_val):
sim.integrate(t_obs, exact_finish_time=1) # time of observation
# account for the light propagation delay (Roemer effect)
t = t_obs - (p.z / speed_of_light * (1 - p.vz / speed_of_light))
sim.integrate(t, exact_finish_time=1) # time of emission
# convert to the coordinate system of the observations
x_obs = -p.y
y_obs = p.x
# account for a possible drift of the astrometric reference frame
x_obs += x0 + vx0 * (t - reference_time)
y_obs += y0 + vy0 * (t - reference_time)
# account for the relativistic Doppler effect
beta_costheta = p.vz / speed_of_light
beta2 = (p.vx**2 + p.vy**2 + p.vz**2) / speed_of_light**2
zD = (1 + beta_costheta) / np.sqrt(1 - beta2) - 1
# account for the gravitational redshift
rs = 2 * sim.G * bh.m / speed_of_light**2
zG = 1 / np.sqrt(1 - rs / np.sqrt(p.x**2 + p.y**2 + p.z**2)) - 1
# calculate the measured radial velocity
vz_obs = (zD + zG) * speed_of_light
# convert to observed units
vz_obs *= velocity_conversion_factor
# account for a possible radial velocity offset
vz_obs += vz0
data.loc[t_obs] = {"t'": t, "x": x_obs, "y": y_obs, "vz": vz_obs}
return data
def perform(self, node, inputs, outputs):
# NOTE: Units should be AU, M_sun, year/2pi
import rebound
masses, initial_coords, times = inputs
masses = np.atleast_1d(masses)
initial_coords = np.atleast_2d(initial_coords)
times = np.atleast_1d(times)
if len(np.shape(masses)) != 1:
raise ValueError("the array of masses must be 1D")
num_bodies = len(masses)
if np.shape(initial_coords) != (num_bodies, 6):
raise ValueError(
"the initial coordinates must have shape {0}".format(
(num_bodies, 6)))
if len(np.shape(times)) != 1:
raise ValueError("the array of times must be 1D")
num_times = len(times)
time_inds = np.argsort(times)
sorted_times = times[time_inds]
# Set up the simulation
sim = rebound.Simulation()
gr_sources = []
for k, v in self.rebound_args.items():
if "gr" in k and "force" not in k:
gr_sources += [v]
continue
setattr(sim, k, v)
if "gr_force" in self.rebound_args.keys():
force = self.rebound_args["gr_force"]
else:
force = "gr" # default to the gr force
for i in range(num_bodies):
sim.add(
m=masses[i],
x=initial_coords[i, 0],
y=initial_coords[i, 1],
z=initial_coords[i, 2],
vx=initial_coords[i, 3],
vy=initial_coords[i, 4],
vz=initial_coords[i, 5],
)
ps = sim.particles
if len(gr_sources) != 0: # if gr_sources have been added
try:
import reboundx
from reboundx import constants
except ImportError:
raise ImportError("""Please install REBOUNDx to include
relativistic effects.""")
rebx = reboundx.Extras(sim)
gr = rebx.load_force(force)
gr.params["c"] = constants.C
for i, particle in enumerate(ps):
particle.params["gr_source"] = gr_sources[i]
rebx.add_force(gr)
# Add the variational particles to track the derivatives
var_systems = np.empty((num_bodies, 7), dtype=object)
for i in range(num_bodies):
for j, coord in enumerate("m x y z vx vy vz".split()):
var = sim.add_variation()
setattr(var.particles[i], coord, 1.0)
var_systems[i, j] = var
# Integrate the system
coords = np.empty((num_times, num_bodies, 6))
jac = np.empty((num_times, num_bodies, 6, num_bodies, 7))
for ind in range(num_times):
sim.integrate(sorted_times[ind])
# Save the coordinates at this time
for i in range(num_bodies):
for j, coord in enumerate("x y z vx vy vz".split()):
coords[ind, i, j] = getattr(sim.particles[i], coord)
# Save the jacobian at this time
for i in range(num_bodies):
for j, coord in enumerate("x y z vx vy vz".split()):
for k in range(num_bodies):
for l in range(7):
jac[ind, i, j, k,
l] = getattr(var_systems[k, l].particles[i],
coord)
# Save the results
outputs[0][0] = np.ascontiguousarray(coords[time_inds])
outputs[1][0] = np.ascontiguousarray(jac[time_inds])