from functools import partial
import spiceypy as spice
from assist import Ephem
from . import simulation_parsing as sp
import rebound
from collections import defaultdict
import assist
import logging
import sys
import os
import numpy as np
from sorcha.ephemeris.simulation_constants import *
from sorcha.ephemeris.simulation_data_files import make_retriever
from sorcha.ephemeris.simulation_geometry import (
barycentricObservatoryRates,
get_hp_neighbors,
ra_dec2vec,
)
from sorcha.ephemeris.simulation_parsing import (
Observatory,
mjd_tai_to_epoch,
)
from sorcha.utilities.generate_meta_kernel import build_meta_kernel_file
[docs]
def create_assist_ephemeris(args, auxconfigs) -> tuple:
"""Build the ASSIST ephemeris object
Parameter
---------
args: dictionary
Dictionary of command-line arguments.
auxconfigs: dataclass
Dataclass of auxiliary configuration file arguments.
Returns
---------
Ephem : ASSIST ephemeris object
The ASSIST ephemeris object
gm_sun : float
value for the GM_SUN value
gm_total : float
value for gm_total
"""
cache_directory = args.ar_data_file_path
return _create_assist_ephemeris(auxconfigs, cache_directory)
[docs]
def _create_assist_ephemeris(auxconfigs, cache_dir=None) -> tuple:
"""Build the ASSIST ephemeris object
Parameter
---------
auxconfigs: dataclass
Dataclass of auxiliary configuration file arguments.
cache_dir: string, default=None
The base directory to place all downloaded files.
Returns
---------
Ephem : ASSIST ephemeris object
The ASSIST ephemeris object
gm_sun : float
value for the GM_SUN value
gm_total : float
value for gm_total
"""
pplogger = logging.getLogger(__name__)
retriever = make_retriever(auxconfigs, cache_dir)
planets_file_path = retriever.fetch(auxconfigs.jpl_planets)
small_bodies_file_path = retriever.fetch(auxconfigs.jpl_small_bodies)
ephem = Ephem(planets_path=planets_file_path, asteroids_path=small_bodies_file_path)
gm_sun = ephem.get_particle("Sun", 0).m
gm_total = sum(sorted([ephem.get_particle(i, 0).m for i in range(27)]))
pplogger.info(f"Calculated GM_SUN value from ASSIST ephemeris: {gm_sun}")
pplogger.info(f"Calculated GM_TOTAL value from ASSIST ephemeris: {gm_total}")
return ephem, gm_sun, gm_total
[docs]
def furnish_spiceypy(args, auxconfigs):
"""
Builds the SPICE kernel, downloading the required files if needed
Parameters
-----------
args : dictionary or `sorchaArguments` object
dictionary of command-line arguments.
auxconfigs: dataclass
Dataclass of auxiliary configuration file arguments.
"""
# The goal here would be to download the spice kernel files (if needed)
# Then call spice.furnish(<filename>) on each of those files.
pplogger = logging.getLogger(__name__)
retriever = make_retriever(auxconfigs, args.ar_data_file_path)
for kernel_file in auxconfigs.ordered_kernel_files:
retriever.fetch(kernel_file)
# check if the META_KERNEL file exists. If it doesn't exist, create it.
if not os.path.exists(os.path.join(retriever.abspath, auxconfigs.meta_kernel)):
build_meta_kernel_file(auxconfigs, retriever)
# try to get the META_KERNEL file. If it's not there, error out.
try:
meta_kernel = retriever.fetch(auxconfigs.meta_kernel)
except ValueError:
pplogger.error(
"ERROR: furnish_spiceypy: Must create meta_kernel.txt by running `bootstrap_sorcha_data_files` on the command line."
)
sys.exit(
"ERROR: furnish_spiceypy: Must create meta_kernel.txt by running `bootstrap_sorcha_data_files` on the command line."
)
spice.furnsh(meta_kernel)
[docs]
def generate_simulations(ephem, gm_sun, gm_total, orbits_df, args):
"""
Creates the dictionary of ASSIST simulations for the ephemeris generation
Parameters
------------
ephem : Ephem
The ASSIST ephemeris object
gm_sun : float
Standard gravitational parameter GM for the Sun
gm_total : float
Standard gravitational parameter GM for the Solar System barycenter
orbits_df : dataframe
Pandas dataframe with the input orbits
args : dictionary or `sorchaArguments` object
dictionary of command-line arguments.
Returns
---------
sim_dict : dict
Dictionary of ASSIST simulations
"""
sim_dict = defaultdict(dict) # return
sun_dict = dict() # This could be passed in and reused
for i, row in orbits_df.iterrows():
epoch = row["epochMJD_TDB"]
# convert from MJD to JD, if not done already.
if epoch < 2400000.5:
epoch += 2400000.5
try:
x, y, z, vx, vy, vz = sp.parse_orbit_row(row, epoch, ephem, sun_dict, gm_sun, gm_total)
if np.isnan(x):
args.pplogger.error(
f"Input elements for orbit {i} failed - see documentation for suggested solutions"
)
sys.exit(f"Input elements for orbit {i} failed - see documentation for suggested solutions")
except ValueError as val_err:
args.pplogger.error(val_err)
sys.exit(val_err)
# Instantiate a rebound particle
ic = rebound.Particle(x=x, y=y, z=z, vx=vx, vy=vy, vz=vz)
# Instantiate a rebound simulation and set initial time and time step
# The time step is just a guess to start with.
sim = rebound.Simulation()
sim.t = epoch - ephem.jd_ref
sim.dt = 10
# This turns off the iterative timestep introduced in arXiv:2401.02849 and default since rebound 4.0.3
sim.ri_ias15.adaptive_mode = 1
# Add the particle to the simulation
sim.add(ic)
# Attach assist extras to the simulation
ex = assist.Extras(sim, ephem)
# Change the GR model for speed
forces = ex.forces
forces.remove("GR_EIH")
forces.append("GR_SIMPLE")
ex.forces = forces
# Save the simulation in the dictionary
sim_dict[row["ObjID"]]["sim"] = sim
sim_dict[row["ObjID"]]["ex"] = ex
return sim_dict