def load_synthetic_obs_data(n0: int,
n1: int,
dt0: datetime = datetime(2000, 1, 1),
dt1: datetime = datetime(2019, 10, 1)):
"""
Load synthetic observation data
INPUTS:
n0: First asteroid number
n1: Last asteroid number
dt0: Start date; defaults to 2000-01-01
dt1: End date; defaults to 2010-10-01
"""
# File name for this data set
fname: str = f'../data/observations/synthetic_n_{n0:06}_{n1:06}.npz'
# Load numpy data and unpack into variables
data = np.load(fname)
t = data['t']
u = data['u']
ast_num = data['ast_num']
# Convert dates to mjd for filtering
t0 = datetime_to_mjd(dt0)
t1 = datetime_to_mjd(dt1)
mask = (t0 <= t) & (t < t1)
return t[mask], u[mask], ast_num[mask]
def get_earth_pos_file():
"""Get the position of earth consistent with the asteroid data; look up from data file"""
# selected data type for TF tensors
dtype = np.float32
# Name of the numpy archive
n0, n1 = 0, 1000
fname_np: str = f'../data/asteroids/sim_asteroids_n_{n0:06}_{n1:06}.npz'
# The full array of positions and velocities
q, v, elts, catalog = load_sim_np(fname_np=fname_np)
# The object names
object_names = catalog['object_names']
# The snapshot times; offset to start time t0=0
ts = catalog['ts'].astype(dtype)
# Convert ts from relative time vs. t0 to MJD
dt0 = datetime(2000, 1, 1)
t_offset = datetime_to_mjd(dt0)
ts += t_offset
# Extract position of sun and earth in Barycentric coordinates
sun_idx = object_names.index('Sun')
earth_idx = object_names.index('Earth')
q_sun_bc = q[:, sun_idx, :]
q_earth_bc = q[:, earth_idx, :]
# Compute earth position in heliocentric coordinates
q_earth = q_earth_bc - q_sun_bc
# Convert to selected data type
return q_earth.astype(dtype), ts
def make_synthetic_obs_data(n0: int, n1: int, dt0: datetime, dt1: datetime,
p0: float, p1: float, noise: float,
frac_real: float):
"""
Genereate synthetic data of asteroid observations.
INPUTS:
n0: asteroid number of first asteroid to consider; inclusive
n1: asteroid number of last asteroid to consider; exclusive
dt0: starting date range for observations
dt1: ending date range for observations
p0: probability of an observation for asteroid n0
p1: probability of an observation for asteroid n1 (interpolates from p0 to p1)
noise: standard deviation of noise applied to observations
frac_real: the fraction of the observations that are real; rest are uniformly distributed bogus observations
OUTPUTS:
t: mjd of this observation
u: asteroid direction (ux, uy, uz)
ast_num: asteroid number; 0 for bogus observations
"""
# Seed RNG for reproducible results
np.random.seed(seed=42)
# Data type for floating point
dtype = np.float32
# Range of MJDs
t0: float = datetime_to_mjd(dt0)
t1: float = datetime_to_mjd(dt1)
# Preallocate storage
num_obs_max: int = int((n1 - n0) * (t1 - t0) / (1.0 - frac_real))
t: np.array = np.zeros(num_obs_max, dtype=dtype)
u: np.array = np.zeros((num_obs_max, space_dims), dtype=dtype)
ast_num_obs: np.array = np.zeros(num_obs_max, dtype=np.int32)
# Load data pages for selected asteroids
page_size: int = 1000
page0: int = n0 // page_size
page1: int = (n1 - 1) // page_size
# Initialize counter for number of observations
num_obs: int = 0
# Iterate through pages of asteroid data
page: int
for page in tqdm(range_inc(page0, page1)):
# Load file for this page
n0_file: int = page * page_size
n1_file: int = n0_file + page_size
inputs, outputs = make_data_one_file(n0=n0_file, n1=n1_file)
# print(f'Loaded file for asteroids {n0_file} to {n1_file}.')
# Asteroid numbers on this page
mask_page: np.array = (n0_file < ast_num_all) & (ast_num_all < n1_file)
ast_num_page: np.array = ast_num_all[mask_page]
# print(f'ast_num_page={ast_num_page}')
# Extract positions; times of asteroid positions synchronized to those for earth
q: np.array = outputs['q']
# Iterate through asteroids in this page
idx: int
ast_num_i: int
# print(f'n0={n0}, n1={n1}')
for idx, ast_num_i in enumerate(ast_num_page):
# Only process asteroids in [n0, n1)
if not (n0 <= ast_num_i and ast_num_i < n1):
continue
# Interpolate the observation acceptance probability
p: float = np.interp(ast_num_i, [n0, n1], [p0, p1])
# print(f'ast_num_i={ast_num_i}, p={p}')
# Direction and observation times for this asteroid
ts_obs, u_obs = get_synthetic_data_one_asteroid(idx=idx,
t0=t0,
t1=t1,
p=p,
noise=noise,
q=q,
ts=ts)
# Number of observations generated; end index of this data slice
ni: int = len(ts_obs)
slice_end = num_obs + ni
# print(f'ni={ni}')
# Copy this page into t and u
t[num_obs:slice_end] = ts_obs
u[num_obs:slice_end] = u_obs
# The observed asteroid number is the same
ast_num_obs[num_obs:slice_end] = ast_num_i
# Update num_obs for this asteroid
num_obs += ni
# Add bogus data points
num_real: int = num_obs
num_total: int = int(np.round(num_real / frac_real))
num_bogus: int = num_total - num_real
ts_bogus = np.arange(t0, t1)
t[num_real:num_total] = np.random.choice(ts_bogus, size=num_bogus)
u[num_real:num_total] = random_direction(num_bogus)
# Use placeholder asteroid number = 0 for bogus observations
ast_num_obs[num_real:num_total] = 0
# Prune array size down and sort it by time t
idx = np.argsort(t[0:num_total])
t = t[idx]
u = u[idx]
ast_num_obs = ast_num_obs[idx]
return t, u, ast_num_obs
def make_data_one_file(n0: int, n1: int) -> tf.data.Dataset:
"""
Wrap the data in one file of asteroid trajectory data into a TF Dataset
INPUTS:
n0: the first asteroid in the file, e.g. 0
n1: the last asteroid in the file (exclusive), e.g. 1000
OUTPUTS:
inputs: a dict of numpy arrays
outputs: a dict of numpy arrays
"""
# selected data type for TF tensors
dtype = np.float32
# Name of the numpy archive
fname_np: str = f'../data/asteroids/sim_asteroids_n_{n0:06}_{n1:06}.npz'
# The full array of positions and velocities
q, v, elts, catalog = load_sim_np(fname_np=fname_np)
# The object names
object_names = catalog['object_names']
# The snapshot times; offset to start time t0=0
ts = catalog['ts'].astype(dtype)
# Convert ts from relative time vs. t0 to MJD
dt0 = datetime(2000, 1, 1)
t_offset = datetime_to_mjd(dt0)
ts += t_offset
# mask for selected asteroids
mask = (n0 <= ast_elt.Num) & (ast_elt.Num < n1)
# count of selected asteroids
# asteroid_names = list(ast_elt.Name[mask].to_numpy())
N_ast: int = np.sum(mask)
# offset for indexing into asteroids; the first [10] objects are sun and planets
ast_offset: int = len(object_names) - N_ast
# Extract position of the sun
sun_idx = 0
q_sun = q[:, sun_idx, :]
v_sun = q[:, sun_idx, :]
# Compute position of the earth in Heliocentric coordinates
earth_idx = 3
q_earth = q[:, earth_idx, :] - q_sun
# v_earth = v[:, earth_idx :] - v_sun
# shrink down q and v to slice with asteroid data only;
q = q[:, ast_offset:, :]
v = v[:, ast_offset:, :]
# swap axes for time step and body number; TF needs inputs and outputs to have same number of samples
# this means that inputs and outputs must first be indexed by asteroid number, then time time step
# also convert to heliocentric coordinates
q = np.swapaxes(q, 0, 1) - q_sun
v = np.swapaxes(v, 0, 1) - v_sun
# Compute relative displacement to earth
q_rel = q - q_earth
# Distance to earth
r_earth = np.linalg.norm(q_rel, axis=2, keepdims=True)
# Direction from earth to asteroid as unit vectors u = (ux, uy, uz)
u = q_rel / r_earth
# dict with inputs
inputs = {
'a': ast_elt.a[mask].to_numpy().astype(dtype),
'e': ast_elt.e[mask].to_numpy().astype(dtype),
'inc': ast_elt.inc[mask].to_numpy().astype(dtype),
'Omega': ast_elt.Omega[mask].to_numpy().astype(dtype),
'omega': ast_elt.omega[mask].to_numpy().astype(dtype),
'f': ast_elt.f[mask].to_numpy().astype(dtype),
'epoch': ast_elt.epoch_mjd[mask].to_numpy().astype(dtype),
'ts': np.tile(ts, reps=(
N_ast,
1,
)),
}
# dict with outputs
outputs = {
'q': q.astype(dtype),
'v': v.astype(dtype),
'u': u.astype(dtype),
}
return inputs, outputs
def make_archive_impl(fname_archive: str, sim_epoch: rebound.Simulation,
object_names: List[str], epoch: datetime, dt0: datetime,
dt1: datetime, time_step: int, save_step: int,
save_elements: bool, progbar: bool) -> None:
"""
Create a rebound simulation archive and save it to disk.
INPUTS:
fname_archive: the file name to save the archive to
sim_epoch: rebound simulation object as of the epoch time; to be integrated in both directions
object_names: the user names of all the objects in the simulation
epoch: a datetime corresponding to sim_epoch
dt0: the earliest datetime to simulate back to
dt1: the latest datetime to simulate forward to
time_step: the time step in days for the simulation
save_step: the interval for saving snapshots to the simulation archive
save_elements: flag indiciting whether to save orbital elements
progbar: flag - whether to display a progress bar
"""
# Convert epoch, start and end times relative to a base date of the simulation start
# This way, time is indexed from t0=0 to t1 = (dt1-dt0)
epoch_t: float = datetime_to_mjd(epoch, dt0)
t0: float = datetime_to_mjd(dt0, dt0)
t1: float = datetime_to_mjd(dt1, dt0)
# Create copies of the simulation to integrate forward and backward
sim_fwd: rebound.Simulation = sim_epoch.copy()
sim_back: rebound.Simulation = sim_epoch.copy()
# Set the time counter on both simulation copies to the epoch time
sim_fwd.t = epoch_t
sim_back.t = epoch_t
# Set the times for snapshots in both directions;
ts: np.array = np.arange(t0, t1 + time_step, time_step)
idx: int = np.searchsorted(ts, epoch_t, side='left')
ts_fwd: np.array = ts[idx:]
ts_back: np.array = ts[:idx][::-1]
# The epochs corresponding to the times in ts
epochs: List[datetime] = [dt0 + timedelta(t) for t in ts]
# File names for forward and backward integrations
fname_fwd: str = fname_archive.replace('.bin', '_fwd.bin')
fname_back: str = fname_archive.replace('.bin', '_back.bin')
# Number of snapshots
M_back: int = len(ts_back)
M_fwd: int = len(ts_fwd)
M: int = M_back + M_fwd
# Number of particles
N: int = sim_epoch.N
# Initialize arrays for the position and velocity
shape_qv: Tuple[int] = (M, N, 3)
q: np.array = np.zeros(shape_qv, dtype=np.float64)
v: np.array = np.zeros(shape_qv, dtype=np.float64)
# Initialize arrays for orbital elements if applicable
if save_elements:
# Arrays for a, e, inc, Omega, omega, f
shape_elt: Tuple[int] = (M, N)
orb_a: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_e: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_inc: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_Omega: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_omega: np.array = np.zeros(shape_elt, dtype=np.float64)
orb_f: np.array = np.zeros(shape_elt, dtype=np.float64)
# Wrap these into a named tuple
elts: OrbitalElement = \
OrbitalElement(a=orb_a, e=orb_e, inc=orb_inc,
Omega=orb_Omega, omega=orb_omega, f=orb_f)
else:
elts = OrbitalElement()
# Subfunction: process one row of the loop
def process_row(sim: rebound.Simulation, fname: str, t: float, row: int):
# Integrate to the current time step with an exact finish time
sim.integrate(t, exact_finish_time=1)
# Serialize the position and velocity
sim.serialize_particle_data(xyz=q[row])
sim.serialize_particle_data(vxvyvz=v[row])
# Save a snapshot on multiples of save_step or the first / last row
if (i % save_step == 0) or (row in (0, M - 1)):
sim.simulationarchive_snapshot(filename=fname)
# Save the orbital elements if applicable
if save_elements:
# Compute the orbital elements vs. the sun as primary
primary = sim.particles['Sun']
orbits = sim.calculate_orbits(primary=primary, jacobi_masses=False)
# Save the elements on this date as an Nx6 array
# This approach only traverses the (slow) Python list orbits one time
elt_array = np.array([[orb.a, orb.e, orb.inc, orb.Omega, orb.omega, orb.f] \
for orb in orbits])
# Save the elements into the current row of the named orbital elements arrays
# The LHS row mask starts at 1 b/c orbital elements are not computed for the first object (Sun)
orb_a[row, 1:] = elt_array[:, 0]
orb_e[row, 1:] = elt_array[:, 1]
orb_inc[row, 1:] = elt_array[:, 2]
orb_Omega[row, 1:] = elt_array[:, 3]
orb_omega[row, 1:] = elt_array[:, 4]
orb_f[row, 1:] = elt_array[:, 5]
# Integrate the simulation forward in time
idx_fwd: List[Tuple[int, float]] = list(enumerate(ts_fwd))
if progbar:
idx_fwd = tqdm_console(idx_fwd)
for i, t in idx_fwd:
# Row index for position data
row: int = M_back + i
# Process this row
process_row(sim=sim_fwd, fname=fname_fwd, t=t, row=row)
# Integrate the simulation backward in time
idx_back: List[Tuple[int, float]] = list(enumerate(ts_back))
if progbar:
idx_back = tqdm_console(idx_back)
for i, t in idx_back:
# Row index for position data
row: int = M_back - (i + 1)
# Process this row
process_row(sim=sim_back, fname=fname_back, t=t, row=row)
# Load the archives with the forward and backward snapshots
sa_fwd: rebound.SimulationArchive = rebound.SimulationArchive(fname_fwd)
sa_back: rebound.SimulationArchive = rebound.SimulationArchive(fname_back)
# Filename for numpy arrays of position and velocity
fname_np: str = fname_archive.replace('.bin', '.npz')
# Save the epochs as a numpy array
epochs_np: np.array = np.array(epochs)
# Save the object names as a numpy array of strings
object_names_np: np.array = np.array(object_names)
# Save the object name hashes
hashes: np.array = np.zeros(N, dtype=np.uint32)
sim_epoch.serialize_particle_data(hash=hashes)
# Combine the forward and backward archives in forward order from t0 to t1
sims = chain(reversed(sa_back), sa_fwd)
# Process each simulation snapshot in turn
for i, sim in enumerate(sims):
# Save a snapshot on multiples of save_step
sim.simulationarchive_snapshot(fname_archive)
# Save the numpy arrays with the object hashes, position and velocity
np.savez(fname_np,
q=q,
v=v,
elts=elts,
ts=ts,
epochs_np=epochs_np,
hashes=hashes,
object_names_np=object_names_np)
# Delete the forward and backward archives
os.remove(fname_fwd)
os.remove(fname_back)