Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 14 additions & 21 deletions src/layup/comet.py
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ def _remove_spc(data):
return data


def _assist_integrate(sim, ex, dt, ephem, include_assist=True):
def _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=True):
"""
Integrates the simulation across a specified time, and returns the orbits before and after integrating.

Expand All @@ -95,10 +95,7 @@ def _assist_integrate(sim, ex, dt, ephem, include_assist=True):
The simulation after integration.
"""
sim.dt = dt
if include_assist:
primary = ephem.get_particle("sun", sim.t)
else:
primary = sim.particles[0]
primary = rebound.Particle(m=Mtot)
oi = sim.particles[-1].orbit(primary=primary)

if include_assist:
Expand All @@ -107,16 +104,12 @@ def _assist_integrate(sim, ex, dt, ephem, include_assist=True):
else:
sim.integrate(sim.t + dt)

if include_assist:
primary = ephem.get_particle("sun", sim.t)
else:
primary = sim.particles[0]
of = sim.particles[-1].orbit(primary=primary)

return oi, of, sim


def _direction_of_integration(sim, ex, step, ephem, include_assist=True):
def _direction_of_integration(sim, ex, step, ephem, Mtot, include_assist=True):
"""
Determines if the simulation is approaching or receding from d=250au, and from this sets the timestep to be positive or negative in order to approach this distance.

Expand Down Expand Up @@ -145,15 +138,15 @@ def _direction_of_integration(sim, ex, step, ephem, include_assist=True):

convert_to_rebound = False
oi, of, sim = _assist_integrate(
sim, ex, step, ephem, include_assist=include_assist
sim, ex, step, ephem, Mtot, include_assist=include_assist
) # Get initial values of oi, of
if oi.d < of.d:
# Moving outwards initially
dt = -abs(step)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=include_assist)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=include_assist)

while of.d < oi.d and sim.t > ASSIST_TIMEFRAME_MIN_MJD: # Returns to its perihelion
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=include_assist)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=include_assist)
if sim.t < ASSIST_TIMEFRAME_MIN_MJD:
convert_to_rebound = True

Expand All @@ -164,11 +157,11 @@ def _direction_of_integration(sim, ex, step, ephem, include_assist=True):

else:
dt = -abs(step)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=include_assist)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=include_assist)
if convert_to_rebound:
sim = assist.simulation_convert_to_rebound(sim, ephem)
while of.d < oi.d: # Returns to its perihelion
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=include_assist)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=include_assist)

return dt, oi, of

Expand Down Expand Up @@ -215,16 +208,16 @@ def _apply_comet(data, args, aux=None, cache_dir=None, primary_id_column_name=No
ex = sim_dict[comet]["ex"]

dt, oi, of = _direction_of_integration(
sim, ex, step, ephem
sim, ex, step, ephem, Mtot
) # Decide whether to go backwards in time or forwards

if dt > 0:
while of.d > REFERENCE_DISTANCE_AU and oi.d > of.d and sim.t < ASSIST_TIMEFRAME_MAX_MJD:
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=True)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=True)

else:
while of.d < REFERENCE_DISTANCE_AU and oi.d < of.d and sim.t > ASSIST_TIMEFRAME_MIN_MJD:
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=True)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=True)

if sim.t >= ASSIST_TIMEFRAME_MAX_MJD or sim.t <= ASSIST_TIMEFRAME_MIN_MJD:
# If comet goes outside assist timeframe, continue the simulation in pure rebound
Expand All @@ -247,16 +240,16 @@ def _apply_comet(data, args, aux=None, cache_dir=None, primary_id_column_name=No
primary = ephem.get_particle("sun", sim.t)
of = sim.particles[-1].orbit(primary=primary).d
dt, oi, of = _direction_of_integration(
sim, ex, step, ephem, include_assist=False
sim, ex, step, ephem, Mtot, include_assist=False
) # Decide whether to go backwards in time or forwards

if dt > 0:
while of.d > REFERENCE_DISTANCE_AU and oi.d > of.d:
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=False)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=False)

else:
while of.d < REFERENCE_DISTANCE_AU and oi.d < of.d:
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, include_assist=False)
oi, of, sim = _assist_integrate(sim, ex, dt, ephem, Mtot, include_assist=False)
output[comet] = (1 / of.a, of.a, of.d, of.e)

# turn output into an array
Expand Down
124 changes: 62 additions & 62 deletions tests/data/demo_comet_expected.csv
Original file line number Diff line number Diff line change
@@ -1,63 +1,63 @@
ObjID,inv_ao,ao_barycentric,d_ao,e_ao
2020 A2,0.0011074035675509633,903.0131645787563,250.0131656507185,0.9988229228648317
2020 A3,0.0002492318682186735,4012.327986574375,250.01266234015256,0.9984950512797639
2020 B3,-4.805058825839382e-05,-20811.39974025839,250.00084113465303,1.0001538082461738
2020 F5,0.00045278116053632325,2208.5724565383666,250.00878424699573,0.9980969038005783
2020 F6,0.0023982852451340524,416.96458001771384,250.00744188197547,0.9913748616238504
2020 F7,0.0015385294762962553,649.9712975323214,250.0113880192667,0.9920321831633419
2020 H5,5.041509933975814e-05,19835.3273740628,250.00220720927496,0.999534093972158
2020 H7,0.00017449175762648922,5730.929721852903,250.01358664383025,0.9992674135558434
2020 H8,0.001663784912250673,601.0392284704982,250.00028794887183,0.9923885482392196
2020 H11,0.0001769184121185289,5652.322943810033,250.00852606235466,0.9985867922682299
2020 K1,0.00034819382965593136,2871.963587028962,250.01119096423088,0.9989103407959656
2020 M5,0.00013735068206083002,7280.63366701826,250.01339358795087,0.9996031608661207
2020 P3,3.7089905435457896e-05,26961.51387444632,250.01365644553948,0.9997360275399467
2020 R2,0.0030268227000782404,330.37944375603865,250.00566727472435,0.9856146498019976
2020 S8,0.0036841983318146275,271.4294698427532,250.00957896074618,0.9910983057978295
2020 T2,0.00296512945089631,337.2534037924437,250.007523659278,0.9932565763437693
2020 U5,0.0001116555159906784,8956.118209900937,250.01400155641198,0.999592550455947
2020 Y2,0.0015164882167974904,659.4182459998227,250.00197008977622,0.9956535080758688
2021 A2,0.004327072579446792,231.10312610652997,250.0075345961884,0.9938038318505573
2021 A6,8.565200270844474e-05,11675.150240257095,250.00037856407366,0.9993262223372645
2021 A7,0.00027117451973477237,3687.6621040135697,250.00515423734322,0.9994726796116138
2021 B2,0.0033596521528780157,297.6498620975862,250.0013232842989,0.9920545203758812
2021 C1,0.0003406494485838937,2935.5691141056527,250.01273775438185,0.9988764760929207
2021 C4,0.0003549577494176526,2817.2367039192986,250.01055813104907,0.9984050713850247
2021 C5,-2.9917534560679747e-05,-33425.21416568489,250.00501166951935,1.0000902634030648
2021 G2,-4.2696668148459675e-05,-23421.03127398422,250.01378235745898,1.0002077177038187
2021 P2,0.00043097207250419347,2320.3359655984896,250.01447152712754,0.997702550259398
2021 Q6,-1.2519492074158865e-05,-79875.44495228142,250.00798064147077,1.000112361377774
2021 R2,0.00047639795735343463,2099.0854065692615,250.0082596821212,0.9966935326784316
2021 R7,0.0011077980830047679,902.6915783132802,250.00438869655434,0.9937168080832474
2021 T1,0.000967771873326219,1033.3013673594514,250.00449108430396,0.9972680666084127
2021 U5,0.004801153638579115,208.28327424571788,250.0016831346987,0.989177812867479
2021 V1,0.002241341376746361,446.1613970878671,250.01286534040102,0.9930351470641414
2022 A1,0.0014120792697342626,708.1755404483692,250.00010411195942,0.9980861466279365
2022 A3,0.0011279964900075528,886.5275812988605,250.00760316733238,0.9956193143809215
2022 B4,0.002901806162601594,344.6129561953432,250.0108066619528,0.9962284706907655
2022 D2,0.003264917093543849,306.28649100383956,250.00738622990602,0.9945522785578625
2022 F1,0.0002096205192805335,4770.525344714502,250.005761545543,0.9987708567171489
2022 H1,0.0008500975864346833,1176.3355360106464,250.00511930640076,0.9931605492756503
2022 S4,4.321047845896334e-05,23142.534766183908,250.0001928358598,0.9998734935288093
2022 T1,3.1160252782643186e-05,32092.165842666647,250.0146313050859,0.9999012481592617
2022 U1,0.0001857292591595003,5384.181278304791,250.01244112093804,0.999274976112592
2022 U4,1.726229686312869e-05,57929.718619075815,250.01404624945383,0.9999546096844194
2023 A1,0.005246773225388587,190.59333366288914,250.005385683196,0.9899874897644526
2023 A2,0.0016924312719852437,590.8659433047388,250.00098422635375,0.9982894341336467
2023 B2,0.0019286224801884377,518.5047930698672,250.0059385825142,0.996548784272864
2023 C2,2.0362585414167276e-05,49109.67736465575,250.00831442110777,0.9999544321327498
2023 F1,0.004452355969485186,224.60019074252662,250.00017840693823,0.9916782639219669
2023 H1,0.000827515110497312,1208.4371479319934,250.00608046403553,0.9963643236078837
2023 K1,0.001765813645311486,566.3111748259267,250.00907528384354,0.9964955211220294
2023 Q2,0.0030486576986771246,328.01321067757806,250.006258120544,0.9902751444655363
2023 T3,0.0016808558163152946,594.9350267247552,250.00613803384798,0.9936508584397978
2023 U1,0.0004299649338300072,2325.771060193863,250.00566783469515,0.9977544407757043
2023 V1,0.0006618262544996003,1510.9705805734304,250.00975145217842,0.9964808849824803
2024 B2,0.0010278822906223937,972.8740431888259,250.00458491893798,0.9956636122039945
2024 G2,0.001927812717747403,518.7225868955118,250.01140879809245,0.9899055663065127
2024 G4,-4.555707781771465e-05,-21950.486025492064,250.0110801363645,1.0002243140550375
2024 G5,0.0024125829630423914,414.4935180753115,250.00357149765426,0.9924874360327928
2024 J4,0.0007430041044914821,1345.887585216515,250.01211973310782,0.9956779309441818
2024 Q4,9.045263197277508e-05,11055.510250945328,250.0031220143603,0.9995356348432183
2024 U1,0.0019151165164527355,522.1614410449787,250.0002304296881,0.9905729291614216
2025 E1,0.00024282416907117323,4118.206205852985,250.00638374948386,0.9990429838578628
2020 A2,0.001156979119287695,864.3198337197816,250.01282373287842,0.9988764560420222
2020 A3,0.00024546524111212023,4073.8965544340913,250.000484783454,0.9985867521234626
2020 B3,1.759641317337498e-05,56829.76355164776,250.00520605678116,0.9999409680538707
2020 F5,0.0005074592876696326,1970.6014340425718,250.00846896530896,0.9978086943397565
2020 F6,0.0024789102026703666,403.40307564298456,250.00846760307886,0.9913272240485077
2020 F7,0.0015222840685340429,656.9076170934367,250.01165325280814,0.9918736910048824
2020 H5,0.0001333306007049505,7500.1537134968485,250.00086050199417,0.9987536380400585
2020 H7,0.0001579940279262474,6329.353160530893,250.0110423540673,0.9993032765625579
2020 H8,0.0016867954775221367,592.8401002526855,250.01163348619625,0.9921189704080257
2020 H11,0.00022047896702293948,4535.580030615602,250.01176510866122,0.9983184549806176
2020 K1,0.0003703721923045311,2699.986718165305,250.00749301877033,0.9988623098358483
2020 M5,0.00018943475591567445,5278.86234585771,250.00185702432336,0.9994315374556404
2020 P3,5.1157482983723245e-05,19547.482434156693,250.01280073466287,0.9996515864913487
2020 R2,0.003013364448949169,331.8549803521853,250.0044114578848,0.9858948982211254
2020 S8,0.0036240752616891365,275.93245939763193,250.0106224221205,0.9914299221242369
2020 T2,0.002942361346252551,339.8630835310624,250.0014222084651,0.9939327196168217
2020 U5,0.00011096996197178789,9011.447622683983,250.01373927633827,0.9995829934709507
2020 Y2,0.0015345760240176461,651.6457864250465,250.00164936673914,0.9951932704320003
2021 A2,0.004296273207876397,232.7598715479013,250.00554408443207,0.9939140601004902
2021 A6,3.901146664028257e-05,25633.488974429307,250.00228015963873,0.9996903866724453
2021 A7,0.00031168670501027416,3208.349871602759,250.01125395708192,0.9993848710168801
2021 B2,0.003393117618296221,294.7142163913929,250.001522060949,0.9914837988707828
2021 C1,0.00034322816352864776,2913.5138262525893,250.00974168759547,0.9988064869842427
2021 C4,0.00031485664857341445,3176.048543141474,250.009215403529,0.9985828160064112
2021 C5,1.439717513847235e-05,69458.07009930613,250.00809177950063,0.9999532095060217
2021 G2,3.10059034895286e-05,32251.92261653407,250.00273548188116,0.9998452825441712
2021 P2,0.00044898596496347306,2227.241112272528,250.0010668136251,0.9977277784633717
2021 Q6,3.8396750810457835e-05,26043.870350812012,250.00939043014708,0.9996659062404223
2021 R2,0.0004576727198876318,2184.967459379097,250.0040135761006,0.9966576886257384
2021 R7,0.0011913270058827065,839.4000933933804,250.0075795096404,0.9932812732161002
2021 T1,0.0009574405541732771,1044.451267121719,250.0010854296254,0.9970763112674643
2021 U5,0.004775771193264693,209.3902658926181,250.0050608469655,0.9887189446377879
2021 V1,0.0023099500357153576,432.90979654904726,250.0123426313092,0.9930604591424086
2022 A1,0.0014131522740386522,707.6378238716603,250.00717167426282,0.998223579804309
2022 A3,0.0011587839564750472,862.9736323256851,250.00931017540708,0.9957046865412638
2022 B4,0.0029722467216666743,336.4458248739372,250.00053961270308,0.9958962766185784
2022 D2,0.003243248777474475,308.3328071979425,250.00204720719424,0.9949542031623072
2022 F1,0.00019473545146772673,5135.171805970464,250.0064657902947,0.9988366638966042
2022 H1,0.0008273196070485272,1208.7227130606902,250.00157306391813,0.9936337037623157
2022 S4,6.675231741868204e-05,14980.753308200938,250.00187981334514,0.9998159983602765
2022 T1,2.5486866587767348e-05,39235.89416362226,250.00625690220028,0.9999120244274573
2022 U1,0.00017634636601255788,5670.658390141074,250.00858468198382,0.9992603150203284
2022 U4,2.6270835043185963e-05,38065.02527826486,250.00327925526744,0.9999236832745094
2023 A1,0.005206937083998864,192.05148513759497,250.00100926823853,0.990408158532634
2023 A2,0.0016525855659397628,605.112389101219,250.00974590773814,0.998435579200995
2023 B2,0.0018735731598963506,533.7394991585606,250.00385170186198,0.9967309066037556
2023 C2,6.700232183126389e-05,14924.855925416467,250.014285325488,0.9998408546325153
2023 F1,0.0044664111876969136,223.89340299759678,250.00457209534423,0.992347389898953
2023 H1,0.0008926041007254466,1120.3175060334918,250.00492826217004,0.9960198542271803
2023 K1,0.001698079769130315,588.9004852299477,250.00436792108172,0.9965362291480795
2023 Q2,0.002989058849699723,334.5534665871362,250.00420018234107,0.9903985968885197
2023 T3,0.001635192656656981,611.5487345964598,250.0004485454663,0.9941970588677927
2023 U1,0.00045549783557550226,2195.4001136724223,250.00671733353306,0.9977361053191599
2023 V1,0.0006498009150204104,1538.932889881495,250.009832806463,0.9966927588051594
2024 B2,0.0010200427982011114,980.35102229391,250.00973031415435,0.995839003040315
2024 G2,0.001980424032372855,504.94236772205045,250.00892298708996,0.9894021579927216
2024 G4,4.388144925814071e-05,22788.67304763149,250.01014030500806,0.999784537039308
2024 G5,0.0024128128179464435,414.4540316439071,250.01194587029806,0.9928809574704761
2024 J4,0.000775501793282163,1289.4876693549493,250.0143490364258,0.995580200455869
2024 Q4,5.7697684877618424e-05,17331.71793844212,250.00850334266778,0.9996854100177117
2024 U1,0.0018894076241421473,529.2664151569901,250.00207730077724,0.9908707900711985
2025 E1,0.00028551494342248955,3502.4436480028794,250.0051053016297,0.9988592234870793
16 changes: 7 additions & 9 deletions tests/layup/test_comet.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import os

import pytest
import rebound
import numpy as np
from numpy.testing import assert_allclose, assert_equal
from layup.comet import _remove_spc, _assist_integrate, _direction_of_integration, _apply_comet, comet_cli
Expand Down Expand Up @@ -88,7 +89,7 @@ def test_assist_integrate(tmpdir, index, time_step, include_assist):
initial = sim.particles[-1].xyz
sim.dt = time_step

oi, of, sim = _assist_integrate(sim, ex, time_step, ephem, include_assist=include_assist)
oi, of, sim = _assist_integrate(sim, ex, time_step, ephem, Mtot, include_assist=include_assist)
initial_orbit, final_orbit = oi.a, of.a

final = sim.particles[-1].xyz
Expand All @@ -99,22 +100,19 @@ def test_assist_integrate(tmpdir, index, time_step, include_assist):

initial_check = sim_check.particles[-1].xyz
sim_check.dt = time_step
primary = rebound.Particle(m=Mtot)

if include_assist == True:
sim_check = assist.simulation_convert_to_rebound(sim_check, ephem)
initial_check = sim_check.particles[-1].xyz
primary_check = sim_check.particles[0]
initial_orbit_check = sim_check.particles[-1].orbit(primary=primary_check).a
initial_orbit_check = sim_check.particles[-1].orbit(primary=primary).a
sim_check.integrate(sim_check.t + time_step) # Want to integrate to the same time as the function
primary_check = sim_check.particles[0]

else:
primary_check = ephem.get_particle("sun", sim_check.t)
initial_orbit_check = sim_check.particles[-1].orbit(primary=primary_check).a
initial_orbit_check = sim_check.particles[-1].orbit(primary=primary).a
sim_check.integrate(sim_check.t + time_step)
primary_check = ephem.get_particle("sun", sim_check.t)
final_check = sim_check.particles[-1].xyz
final_orbit_check = sim_check.particles[-1].orbit(primary=primary_check).a
final_orbit_check = sim_check.particles[-1].orbit(primary=primary).a

assert_equal(np.array(initial), np.array(initial_check))
assert_equal(sim.t, sim_check.t)
Expand Down Expand Up @@ -154,7 +152,7 @@ def __init__(self, g=None):
for comet in range(len(data)):
sim = sim_dict[data[comet]["ObjID"]]["sim"]
ex = sim_dict[data[comet]["ObjID"]]["ex"]
dt, oi, of = _direction_of_integration(sim, ex, 1, ephem, include_assist=True)
dt, oi, of = _direction_of_integration(sim, ex, 1, ephem, Mtot, include_assist=True)
assert dt == data["expected_step"][comet]


Expand Down
1 change: 0 additions & 1 deletion tests/layup/test_layup_configs.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
AuxiliaryConfigs,
)


correct_auxciliary_URLs = {
"de440s.bsp": "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/spk/planets/de440s.bsp",
"earth_2025_250826_2125_predict.bpc": "https://naif.jpl.nasa.gov/pub/naif/generic_kernels/pck/earth_2025_250826_2125_predict.bpc",
Expand Down
Loading