From 0f07c5d576121464242b9608293e38a41cba988b Mon Sep 17 00:00:00 2001 From: awilson110 Date: Tue, 28 Jul 2026 12:52:19 +0100 Subject: [PATCH] comet code uses Mtot as primary, not Msun updated the code, the tests and used the linter I think the linter removed a blank line in test_layup_configs.py --- src/layup/comet.py | 35 ++++---- tests/data/demo_comet_expected.csv | 124 ++++++++++++++--------------- tests/layup/test_comet.py | 16 ++-- tests/layup/test_layup_configs.py | 1 - 4 files changed, 83 insertions(+), 93 deletions(-) diff --git a/src/layup/comet.py b/src/layup/comet.py index 85a94cc6..f07860fd 100644 --- a/src/layup/comet.py +++ b/src/layup/comet.py @@ -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. @@ -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: @@ -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. @@ -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 @@ -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 @@ -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 @@ -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 diff --git a/tests/data/demo_comet_expected.csv b/tests/data/demo_comet_expected.csv index 460697b3..0a951908 100644 --- a/tests/data/demo_comet_expected.csv +++ b/tests/data/demo_comet_expected.csv @@ -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 diff --git a/tests/layup/test_comet.py b/tests/layup/test_comet.py index 69dc596c..7b2c1b08 100644 --- a/tests/layup/test_comet.py +++ b/tests/layup/test_comet.py @@ -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 @@ -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 @@ -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) @@ -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] diff --git a/tests/layup/test_layup_configs.py b/tests/layup/test_layup_configs.py index 12813e2e..1d8c3556 100644 --- a/tests/layup/test_layup_configs.py +++ b/tests/layup/test_layup_configs.py @@ -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",