def make_model_jac_to_cart(num_body: int, batch_size: int = 64):
"""Model that transforms from Jacobi to Cartesian coordinates"""
# The shape shared by all the inputs
space_dims = 3
shape_m = (num_body, )
shape_q = (
num_body,
space_dims,
)
# Create input layers
m = keras.Input(shape=shape_m, batch_size=batch_size, name='m')
qj = keras.Input(shape=shape_q, batch_size=batch_size, name='qj')
vj = keras.Input(shape=shape_q, batch_size=batch_size, name='vj')
# Wrap these up into one tuple of inputs
inputs = (m, qj, vj)
# Calculations are in one layer that does all the work...
q, v = JacobiToCartesian(name='j2c')(inputs)
# Name outputs
q = Identity(name='q')(q)
v = Identity(name='v')(v)
# Wrap up the outputs
outputs = (q, v)
# Create a model from inputs to outputs
model = keras.Model(inputs=inputs,
outputs=outputs,
name='jacobi_to_cartesian')
return model
def make_model_cart_to_jac(num_body: int, batch_size: int = 64):
"""Model that transforms from cartesian to Jacobi coordinates"""
# The shape shared by all the inputs
space_dims = 3
shape_m = (num_body, )
shape_q = (
num_body,
space_dims,
)
# Create input layers
m = keras.Input(shape=shape_m, batch_size=batch_size, name='m')
q = keras.Input(shape=shape_q, batch_size=batch_size, name='q')
v = keras.Input(shape=shape_q, batch_size=batch_size, name='v')
# Wrap these up into one tuple of inputs
inputs = (m, q, v)
# Calculations are in one layer that does all the work...
qj, vj, mu = CartesianToJacobi(name='c2j')(inputs)
# Name outputs
qj = Identity(name='qj')(qj)
vj = Identity(name='vj')(vj)
mu = Identity(name='mu')(mu)
# Wrap up the outputs
outputs = (qj, vj, mu)
# Create a model from inputs to outputs
model = keras.Model(inputs=inputs,
outputs=outputs,
name='cartesian_to_jacobi')
return model
def make_physics_model_g2b(position_model: keras.Model, traj_size: int):
"""Create a physics model for the general two body problem from a position model"""
# Create input layers
num_particles = 2
space_dims = 3
t = keras.Input(shape=(traj_size, ), name='t')
q0 = keras.Input(shape=(
num_particles,
space_dims,
), name='q0')
v0 = keras.Input(shape=(
num_particles,
space_dims,
), name='v0')
m = keras.Input(shape=(num_particles, ), name='m')
# Wrap these up into one tuple of inputs for the model
inputs = (t, q0, v0, m)
# Return row 0 of a position or velocity for q0_rec and v0_rec
initial_row_func = lambda q: q[:, 0, :]
# Compute the motion from the specified position layer; inputs are the same for position and physics model
q, v, a = Motion_G2B(position_model=position_model, name='motion')(inputs)
# Name the outputs of the motion
# These each have shape (batch_size, num_particles, traj_size, 3)
q = Identity(name='q')(q)
v = Identity(name='v')(v)
a = Identity(name='a')(a)
# Compute q0_rec and v0_rec
# These each have shape (batch_size, num_particles, 3)
q0_rec = keras.layers.Lambda(initial_row_func, name='q0_rec')(q)
v0_rec = keras.layers.Lambda(initial_row_func, name='v0_rec')(v)
# Compute kinetic energy T and potential energy U
# These have shape (batch_size, traj_size)
T = KineticEnergy_G2B(name='T')([m, v])
U = PotentialEnergy_G2B(name='U')([m, q])
# Compute the total energy H
H = keras.layers.add(inputs=[T, U], name='H')
# Compute momentum P and angular momentum L
# These have shape (batch_size, traj_size, 3)
P = Momentum_G2B(name='P')([m, v])
L = AngularMomentum_G2B(name='L')([m, q, v])
# Wrap this up into a model
outputs = [q, v, a, q0_rec, v0_rec, H, P, L]
model_name = position_model.name.replace('model_g2b_position_',
'model_g2b_physics_')
model = keras.Model(inputs=inputs, outputs=outputs, name=model_name)
return model
def make_model_cart_to_cart(num_body: int, batch_size: int = 64):
"""Model that transforms from cartesian to cartesian coordinates"""
# The shape shared by all the inputs
space_dims = 3
shape_m = (num_body, )
shape_q = (
num_body,
space_dims,
)
# Create input layers
m = keras.Input(shape=shape_m, batch_size=batch_size, name='m')
q = keras.Input(shape=shape_q, batch_size=batch_size, name='q')
v = keras.Input(shape=shape_q, batch_size=batch_size, name='v')
# Wrap these up into inputs
inputs = (m, q, v)
# Wrap these up into one tuple of inputs for c2j mapping
inputs_c2j = inputs
# Convert from Cartesian to Jacobi coordinates in step 1
qj, vj, mu = CartesianToJacobi(name='c2j')(inputs_c2j)
# Name Jacobi coordinates
qj = Identity(name='qj')(qj)
vj = Identity(name='vj')(vj)
mu = Identity(name='mu')(mu)
# Wrap these up into one tuple of inputs for j2c mapping
inputs_j2c = (m, qj, vj)
# Calculations are in one layer that does all the work...
q_calc, v_calc = JacobiToCartesian(name='j2c')(inputs_j2c)
# Name outputs
q_calc = Identity(name='q_calc')(q_calc)
v_calc = Identity(name='v_calc')(v_calc)
# Wrap up the outputs
outputs = (q_calc, v_calc)
# Create a model from inputs to outputs
model = keras.Model(inputs=inputs,
outputs=outputs,
name='cartesian_to_cartesian')
return model
def make_model_cfg_to_elt(batch_size: int=64, name=None):
"""Model that transforms from orbital elements to cartesian coordinates"""
# Create input layers
q = keras.Input(shape=(3,), batch_size=batch_size, name='q')
v = keras.Input(shape=(3,), batch_size=batch_size, name='v')
mu = keras.Input(shape=(1,), batch_size=batch_size, name='mu')
# Wrap these up into one tuple of inputs for the model
inputs_model = (q, v, mu,)
# Unpack coordinates from inputs
qx = keras.layers.Reshape(target_shape=(1,), name='qx')(q[:,0])
qy = keras.layers.Reshape(target_shape=(1,), name='qy')(q[:,1])
qz = keras.layers.Reshape(target_shape=(1,), name='qz')(q[:,2])
vx = keras.layers.Reshape(target_shape=(1,), name='vx')(v[:,0])
vy = keras.layers.Reshape(target_shape=(1,), name='vy')(v[:,1])
vz = keras.layers.Reshape(target_shape=(1,), name='vz')(v[:,2])
# Tuple of inputs for the layer
inputs_layer = (qx, qy, qz, vx, vy, vz, mu,)
# Calculations are in one layer that does all the work...
a, e, inc, Omega, omega, f, M, N = ConfigToOrbitalElement(name='config_to_orbital_element')(inputs_layer)
# Name the outputs of the layer
a = Identity(name='a')(a)
e = Identity(name='e')(e)
inc = Identity(name='inc')(inc)
Omega = Identity(name='Omega')(Omega)
omega = Identity(name='omega')(omega)
f = Identity(name='f')(f)
# "Bonus outputs" - mean anomaly and mean motion
M = Identity(name='M')(M)
N = Identity(name='N')(N)
# Wrap up the outputs
outputs = (a, e, inc, Omega, omega, f, M, N)
# Create a model from inputs to outputs
name = name or 'config_to_orbital_element'
model = keras.Model(inputs=inputs_model, outputs=outputs, name=name)
return model
def make_physics_model_r2b(position_model: keras.Model, traj_size: int):
"""Create a physics model for the restricted two body problem from a position model"""
# Create input layers
t = keras.Input(shape=(traj_size, ), name='t')
q0 = keras.Input(shape=(3, ), name='q0')
v0 = keras.Input(shape=(3, ), name='v0')
mu = keras.Input(shape=(1, ), name='mu')
# Wrap these up into one tuple of inputs for the model
inputs = (t, q0, v0, mu)
# Check sizes of inputs
batch_size = t.shape[0]
tf.debugging.assert_shapes(shapes={
t: (batch_size, traj_size),
q0: (batch_size, 3),
v0: (batch_size, 3),
mu: (batch_size, 1),
},
message='make_physics_model_r2b_math / inputs')
# Return row 0 of a position or velocity for q0_rec and v0_rec
initial_row_func = lambda q: q[:, 0, :]
# Compute the motion from the specified position layer; inputs are the same for position and physics model
q, v, a = Motion_R2B(position_model=position_model, name='motion')(inputs)
# Name the outputs of the motion
# These each have shape (batch_size, traj_size, 3)
q = Identity(name='q')(q)
v = Identity(name='v')(v)
a = Identity(name='a')(a)
# Check sizes
tf.debugging.assert_shapes(
shapes={
q: (batch_size, traj_size, 3),
v: (batch_size, traj_size, 3),
a: (batch_size, traj_size, 3),
},
message='make_physics_model_r2b / outputs q, v, a')
# Compute q0_rec and v0_rec
# These each have shape (batch_size, 2)
q0_rec = keras.layers.Lambda(initial_row_func, name='q0_rec')(q)
v0_rec = keras.layers.Lambda(initial_row_func, name='v0_rec')(v)
# Check sizes
tf.debugging.assert_shapes(
shapes={
q0_rec: (batch_size, 3),
v0_rec: (batch_size, 3),
},
message='make_physics_model_r2b / outputs q0_rec, v0_rec')
# Compute kinetic energy T and potential energy U
T = KineticEnergy_R2B(name='T')(v)
U = PotentialEnergy_R2B(name='U')((q, mu))
# Compute the total energy H
H = keras.layers.add(inputs=[T, U], name='H')
# Compute angular momentum L
# This has shape (batch_size, traj_size, 3)
L = AngularMomentum_R2B(name='L')([q, v])
# Check sizes
tf.debugging.assert_shapes(
shapes={
T: (batch_size, traj_size),
U: (batch_size, traj_size),
H: (batch_size, traj_size),
L: (batch_size, traj_size, 3),
},
message='make_physics_model_r2bc_math / outputs H, L')
# Wrap this up into a model
outputs = [q, v, a, q0_rec, v0_rec, H, L]
model_name = position_model.name.replace('model_r2b_position_',
'model_r2b_physics_')
model = keras.Model(inputs=inputs, outputs=outputs, name=model_name)
return model
def make_model_asteroid_search(ts: tf.Tensor,
elts_np: Dict,
max_obs: int,
num_obs: float,
elt_batch_size: int=64,
time_batch_size: int=None,
R_deg: float = 5.0,
alpha: float = 2.0,
beta: float = 1.0,
q_cal = None,
use_calibration: bool = True):
"""Make functional API model for scoring elements"""
# The full trajectory size
traj_size: int = ts.shape[0]
# Default for time_batch_size is full trajectory size
if time_batch_size is None:
time_batch_size = traj_size
# Inputs
t = keras.Input(shape=(), batch_size=time_batch_size, dtype=tf.float32, name='t' )
idx = keras.Input(shape=(), batch_size=time_batch_size, dtype=tf.int32, name='idx')
row_len = keras.Input(shape=(), batch_size=time_batch_size, dtype=tf.int32, name='row_len')
u_obs = keras.Input(shape=(max_obs, space_dims), batch_size=time_batch_size, dtype=tf.float32, name='u_obs')
# Output times are a constant
ts = keras.backend.constant(ts, name='ts')
# Set of trainable weights with candidate
elements_layer = OrbitalElements(elts_np=elts_np, batch_size=elt_batch_size, R_deg=R_deg, name='candidates')
a, e, inc, Omega, omega, f, epoch, R = elements_layer(idx)
# Alias the orbital elements; a, e, inc, Omega, omega, and f are trainable; epoch is fixed
a = Identity(name='a')(a)
e = Identity(name='e')(e)
inc = Identity(name='inc')(inc)
Omega = Identity(name='Omega')(Omega)
omega = Identity(name='omega')(omega)
f = Identity(name='f')(f)
epoch = Identity(name='epoch')(epoch)
# Alias the resolution output
R = Identity(name='R')(R)
# The orbital elements; stack to shape (elt_batch_size, 7)
elts = tf.stack(values=[a, e, inc, Omega, omega, f, epoch], axis=1, name='elts')
# The predicted direction
direction_layer = AsteroidDirection(ts=ts, batch_size=elt_batch_size, name='u_pred')
# Compute numerical orbits for calibration if necessary
if use_calibration:
if q_cal is not None:
q_ast, q_sun, q_earth = q_cal
else:
print(f'Numerically integrating orbits for calibration...')
epoch0 = elts_np['epoch'][0]
q_ast, q_sun, q_earth = calc_ast_pos(elts=elts_np, epoch=epoch0, ts=ts)
# Calibrate the direction prediction layer
if use_calibration:
direction_layer.calibrate(elts=elts_np, q_ast=q_ast, q_sun=q_sun)
# Tensor of predicted directions
u_pred = direction_layer(a, e, inc, Omega, omega, f, epoch)
# Difference in direction between u_obs and u_pred
dir_diff_layer = DirectionDifference(batch_size=elt_batch_size,
traj_size=time_batch_size,
max_obs=max_obs,
name='z')
z = dir_diff_layer(u_obs, u_pred, idx)
# Calculate score compoments
score_layer = TrajectoryScore(batch_size=elt_batch_size, alpha=alpha, beta=beta)
raw_score, mu, sigma2, objective = score_layer(z, R, num_obs)
# Stack the scores
scores = tf.stack(values=[raw_score, mu, sigma2, objective], axis=1, name='scores')
# Wrap inputs and outputs
inputs = (t, idx, row_len, u_obs)
outputs = (elts, R, u_pred, z, scores)
# Create model with functional API
model = keras.Model(inputs=inputs, outputs=outputs)
# Bind the custom layers to model
model.elements = elements_layer
model.direction = direction_layer
model.dir_diff = dir_diff_layer
model.score = score_layer
# Log transform
# objective_min = 0.0
# Add the loss function
model.add_loss(-tf.reduce_sum(objective))
# model.add_loss(-tf.reduce_sum(tf.math.log(objective-objective_min)))
return model
def make_position_model_g3b_nn(hidden_sizes,
skip_layers=True,
kernel_reg=0.0,
activity_reg=0.0,
traj_size=1001,
batch_size=64):
"""
Compute orbit positions for the general three body problem using a neural
network that computes adjustments to orbital elements.
Factory function that returns a functional model.
"""
# Dimensionality
num_body = 3
space_dims = 3
# Input layers
t = keras.Input(shape=(traj_size, ), batch_size=batch_size, name='t')
q0 = keras.Input(shape=(
num_body,
space_dims,
),
batch_size=batch_size,
name='q0')
v0 = keras.Input(shape=(
num_body,
space_dims,
),
batch_size=batch_size,
name='v0')
m = keras.Input(shape=(num_body, ), batch_size=batch_size, name='m')
# Wrap these up into one tuple of inputs for the model
inputs = (t, q0, v0, m)
# Compute the Jacobi coordinates of the initial conditions
qj0, vj0, mu0 = CartesianToJacobi()([m, q0, v0])
# Extract Jacobi coordinates of p1 and p2
qj0_1 = qj0[:, 1, :]
qj0_2 = qj0[:, 2, :]
vj0_1 = vj0[:, 1, :]
vj0_2 = vj0[:, 2, :]
# Extract gravitational field strength for orbital element conversion of p1 and p2
mu0_1 = mu0[:, 1:2]
mu0_2 = mu0[:, 2:3]
# Manually set the shapes to work around documented bug on slices losing shape info
jacobi_shape = (batch_size, space_dims)
qj0_1.set_shape(jacobi_shape)
qj0_2.set_shape(jacobi_shape)
vj0_1.set_shape(jacobi_shape)
vj0_1.set_shape(jacobi_shape)
mu_shape = (batch_size, 1)
mu0_1.set_shape(mu_shape)
mu0_2.set_shape(mu_shape)
# Tuple of inputs for the model converting from configuration to orbital elements
cfg_1 = (qj0_1, vj0_1, mu0_1)
cfg_2 = (qj0_2, vj0_2, mu0_2)
# Model mapping cartesian coordinates to orbital elements
model_c2e_1 = make_model_cfg_to_elt(name='orbital_element_1')
model_c2e_2 = make_model_cfg_to_elt(name='orbital_element_2')
# Extract the orbital elements of the initial conditions
a1_0, e1_0, inc1_0, Omega1_0, omega1_0, f1_0, M1_0, N1_0 = model_c2e_1(
cfg_1)
a2_0, e2_0, inc2_0, Omega2_0, omega2_0, f2_0, M2_0, N2_0 = model_c2e_2(
cfg_2)
# Alias mu0_i for naming consistency
mu1_0 = mu0_1
mu2_0 = mu0_2
# Reshape t to (batch_size, traj_size, 1)
t_vec = keras.layers.Reshape(target_shape=(traj_size, 1), name='t_vec')(t)
# ******************************************************************
# Kepler-Jacobi Model: Analytical approximation ignoring interaction of two small bodies
# ******************************************************************
# ******************************************************************
# Predict orbital elements for Jacobi coordinates of body 1
# Repeat the constant orbital elements to be vectors of shape (batch_size, traj_size)
a1 = keras.layers.RepeatVector(n=traj_size, name='a1_kj')(a1_0)
e1 = keras.layers.RepeatVector(n=traj_size, name='e1_kj')(e1_0)
inc1 = keras.layers.RepeatVector(n=traj_size, name='inc1_kj')(inc1_0)
Omega1 = keras.layers.RepeatVector(n=traj_size, name='Omega1_kj')(Omega1_0)
omega1 = keras.layers.RepeatVector(n=traj_size, name='omega1_kj')(omega1_0)
mu1 = keras.layers.RepeatVector(n=traj_size, name='mu1')(mu1_0)
# Repeat initial mean anomaly M0 and mean motion N0 to match shape of outputs
M1_0_vec = keras.layers.RepeatVector(n=traj_size, name='M1_0_vec')(M1_0)
N1_0_vec = keras.layers.RepeatVector(n=traj_size, name='N1_0_vec')(N1_0)
# Compute the mean anomaly M(t) as a function of time
N1_t = keras.layers.multiply(inputs=[N1_0_vec, t_vec])
M1 = keras.layers.add(inputs=[M1_0_vec, N1_t])
# Compute the true anomaly from the mean anomaly and eccentricity
f1 = MeanToTrueAnomaly(name='mean_to_true_anomaly_f1')([M1, e1])
# ******************************************************************
# Predict orbital elements for Jacobi coordinates of body 2
# Repeat the constant orbital elements to be vectors of shape (batch_size, traj_size)
a2 = keras.layers.RepeatVector(n=traj_size, name='a2_kj')(a2_0)
e2 = keras.layers.RepeatVector(n=traj_size, name='e2_kj')(e2_0)
inc2 = keras.layers.RepeatVector(n=traj_size, name='inc2_kj')(inc2_0)
Omega2 = keras.layers.RepeatVector(n=traj_size, name='Omega2_kj')(Omega2_0)
omega2 = keras.layers.RepeatVector(n=traj_size, name='omega2_kj')(omega2_0)
mu2 = keras.layers.RepeatVector(n=traj_size, name='mu2')(mu2_0)
# Repeat initial mean anomaly M0 and mean motion N0 to match shape of outputs
M2_0_vec = keras.layers.RepeatVector(n=traj_size, name='M2_0_vec')(M2_0)
N2_0_vec = keras.layers.RepeatVector(n=traj_size, name='N2_0_vec')(N2_0)
# Compute the mean anomaly M(t) as a function of time
N2_t = keras.layers.multiply(inputs=[N2_0_vec, t_vec])
M2 = keras.layers.add(inputs=[M2_0_vec, N2_t])
# Compute the true anomaly from the mean anomaly and eccentricity
f2 = MeanToTrueAnomaly(name='mean_to_true_anomaly_f2')([M2, e2])
# ******************************************************************
# Feature extraction: masses & cos/sin of angle variables
# ******************************************************************
# Extract masses of p1 and p2
m1 = m[:, 1:2]
m2 = m[:, 2:3]
# Manually set the shapes to work around documented bug on slices losing shape info
mass_shape = (batch_size, 1)
m1.set_shape(mass_shape)
m2.set_shape(mass_shape)
# Repeat the masses to shape (batch_size, traj_size, num_body-1)
# skip mass of body 0 because it is a constant = 1.0 solar mass
m1 = keras.layers.RepeatVector(n=traj_size, name='m1')(m1)
m2 = keras.layers.RepeatVector(n=traj_size, name='m2')(m2)
# Repeat the initial true anomaly f1_0 and f2_0
f1_0_vec = keras.layers.RepeatVector(n=traj_size, name='f1_0_vec')(f1_0)
f2_0_vec = keras.layers.RepeatVector(n=traj_size, name='f2_0_vec')(f2_0)
# Convert inc1 and inc2 to cosine and sine
cos_inc1 = keras.layers.Activation(activation=tf.cos,
name='cos_inc1')(inc1)
sin_inc1 = keras.layers.Activation(activation=tf.sin,
name='sin_inc1')(inc1)
cos_inc2 = keras.layers.Activation(activation=tf.cos,
name='cos_inc2')(inc2)
sin_inc2 = keras.layers.Activation(activation=tf.sin,
name='sin_inc2')(inc2)
# Convert Omega1 and Omega2 to cosine and sine
cos_Omega1 = keras.layers.Activation(activation=tf.cos,
name='cos_Omega1')(Omega1)
sin_Omega1 = keras.layers.Activation(activation=tf.sin,
name='sin_Omega1')(Omega1)
cos_Omega2 = keras.layers.Activation(activation=tf.cos,
name='cos_Omega2')(Omega2)
sin_Omega2 = keras.layers.Activation(activation=tf.sin,
name='sin_Omega2')(Omega2)
# Convert omega1 and omega2 to cosine and sine
cos_omega1 = keras.layers.Activation(activation=tf.cos,
name='cos_omega1')(omega1)
sin_omega1 = keras.layers.Activation(activation=tf.sin,
name='sin_omega1')(omega1)
cos_omega2 = keras.layers.Activation(activation=tf.cos,
name='cos_omega2')(omega2)
sin_omega2 = keras.layers.Activation(activation=tf.sin,
name='sin_omega2')(omega2)
# Convert f1 and f2 to cosine and sine
cos_f1_0 = keras.layers.Activation(activation=tf.cos,
name='cos_f1_0')(f1_0_vec)
sin_f1_0 = keras.layers.Activation(activation=tf.sin,
name='sin_f1_0')(f1_0_vec)
cos_f2_0 = keras.layers.Activation(activation=tf.cos,
name='cos_f2_0')(f2_0_vec)
sin_f2_0 = keras.layers.Activation(activation=tf.sin,
name='sin_f2_0')(f2_0_vec)
# ******************************************************************
# Neural network: feature layers
# ******************************************************************
# Create an initial array of features: the time, mass and orbital elements
feature_list = [
# time of this snapshot and body masses (body 1 constant = 1.0 Msun)
t_vec,
m1,
m2,
# orbital elements 1 (10 features)
a1,
e1,
cos_inc1,
sin_inc1,
cos_Omega1,
sin_Omega1,
cos_omega1,
sin_omega1,
cos_f1_0,
sin_f1_0,
# orbital elements 2 (10 features)
a2,
e2,
cos_inc2,
sin_inc2,
cos_Omega2,
sin_Omega2,
cos_omega2,
sin_omega2,
cos_f2_0,
sin_f2_0
]
# Inputs to neural network is a flattened arrat; 23 features per time snap
phi_0 = keras.layers.concatenate(inputs=feature_list, name='phi_0')
# Hidden layers as specified in hidden_sizes
# Number of hidden layers
num_layers = len(hidden_sizes)
# phi_n will update to the last available feature layer for the output portion
phi_n = phi_0
# First hidden layer if applicable
if num_layers > 0:
phi_1 = keras.layers.Dense(units=hidden_sizes[0],
activation='tanh',
name='phi_1')(phi_0)
if skip_layers:
phi_1 = keras.layers.concatenate(inputs=[phi_0, phi_1],
name='phi_1_aug')
phi_n = phi_1
# Second hidden layer if applicable
if num_layers > 1:
phi_2 = keras.layers.Dense(units=hidden_sizes[1],
activation='tanh',
name='phi_2')(phi_1)
if skip_layers:
phi_2 = keras.layers.concatenate(inputs=[phi_1, phi_2],
name='phi_2_aug')
phi_n = phi_2
# Third hidden layer if applicable
if num_layers > 2:
phi_3 = keras.layers.Dense(units=hidden_sizes[2],
activation='tanh',
name='phi_3')(phi_2)
if skip_layers:
phi_3 = keras.layers.concatenate(inputs=[phi_2, phi_3],
name='phi_3_aug')
phi_n = phi_3
# ******************************************************************
# Neural network: layers with time derivatives of orbital elements
# ******************************************************************
# Set type of regularizer
# Set strength of activity regularizer using activity_reg input
reg_type = keras.regularizers.l1
# Semimajor axis
ddt_a1 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_a1')(phi_n)
ddt_a2 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_a2')(phi_n)
# Eccentricity
ddt_e1 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_e1')(phi_n)
ddt_e2 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_e2')(phi_n)
# Inclination
ddt_inc1 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_inc1')(phi_n)
ddt_inc2 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_inc2')(phi_n)
# Longitude of ascending node
ddt_Omega1 = keras.layers.Dense(
units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_Omega1')(phi_n)
ddt_Omega2 = keras.layers.Dense(
units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_Omega2')(phi_n)
# Argument of periapsis
ddt_omega1 = keras.layers.Dense(
units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_omega1')(phi_n)
ddt_omega2 = keras.layers.Dense(
units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_omega2')(phi_n)
# True anomaly
ddt_f1 = keras.layers.Dense(
units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=keras.regularizers.l1(activity_reg),
name='ddt_f1')(phi_n)
ddt_f2 = keras.layers.Dense(units=1,
kernel_initializer='zeros',
bias_initializer='zeros',
kernel_regularizer=reg_type(kernel_reg),
activity_regularizer=reg_type(activity_reg),
name='ddt_f2')(phi_n)
# ******************************************************************
# Apply adjustments to Kepler-Jacobi orbital elements
# ******************************************************************
a1 = keras.layers.add(inputs=[a1, ddt_a1 * t_vec], name='a1_raw')
a2 = keras.layers.add(inputs=[a2, ddt_a2 * t_vec], name='a2_raw')
e1 = keras.layers.add(inputs=[e1, ddt_e1 * t_vec], name='e1_raw')
e2 = keras.layers.add(inputs=[e2, ddt_e2 * t_vec], name='e2_raw')
inc1 = keras.layers.add(inputs=[inc1, ddt_inc1 * t_vec], name='inc1_raw')
inc2 = keras.layers.add(inputs=[inc2, ddt_inc2 * t_vec], name='inc2_raw')
Omega1 = keras.layers.add(inputs=[Omega1, ddt_Omega1 * t_vec],
name='Omega1')
Omega2 = keras.layers.add(inputs=[Omega2, ddt_Omega2 * t_vec],
name='Omega2')
omega1 = keras.layers.add(inputs=[omega1, ddt_omega1 * t_vec],
name='omega1')
omega2 = keras.layers.add(inputs=[omega2, ddt_omega2 * t_vec],
name='omega2')
f1 = keras.layers.add(inputs=[f1, ddt_f1 * t_vec], name='f1')
f2 = keras.layers.add(inputs=[f2, ddt_f2 * t_vec], name='f2')
# Limit a to be non-negative
a1 = keras.layers.ReLU(name='a1')(a1)
a2 = keras.layers.ReLU(name='a2')(a2)
# Limit e to be in interval [0.0, 0.9999]
ecc_min = 0.0000
ecc_max = 0.9900
clip_func_e = lambda x: tf.clip_by_value(x, ecc_min, ecc_max)
e1 = keras.layers.Activation(activation=clip_func_e, name='e1')(e1)
e2 = keras.layers.Activation(activation=clip_func_e, name='e2')(e2)
# Limit inc to be in interval [0.0, pi]
inc_min = 0.0
inc_max = np.pi
clip_func_inc = lambda x: tf.clip_by_value(x, inc_min, inc_max)
inc1 = keras.layers.Activation(activation=clip_func_inc, name='inc1')(inc1)
inc2 = keras.layers.Activation(activation=clip_func_inc, name='inc2')(inc2)
# The remaining elements can take any value; angles can wrap around past 2pi
# ******************************************************************
# Convert orbital elements to cartesian Jacobi coordinates and then to Cartesian body coordinates
# ******************************************************************
# The position of Jacobi coordinate 0 over time comes from the average velocity
# We always use center of momentum coordinates, so this is zero
qjt_0 = keras.backend.zeros(shape=[batch_size, traj_size, space_dims])
vjt_0 = keras.backend.zeros(shape=[batch_size, traj_size, space_dims])
ajt_0 = keras.backend.zeros(shape=[batch_size, traj_size, space_dims])
# Model mapping orbital elements to cartesian coordinates
model_e2c_1 = make_model_elt_to_cfg(include_accel=True,
batch_size=batch_size,
name='elt_to_jac_1')
model_e2c_2 = make_model_elt_to_cfg(include_accel=True,
batch_size=batch_size,
name='elt_to_jac_2')
# Wrap orbital elements into one tuple of inputs for layer converting to cartesian coordinates
elt1 = (
a1,
e1,
inc1,
Omega1,
omega1,
f1,
mu1,
)
elt2 = (
a2,
e2,
inc2,
Omega2,
omega2,
f2,
mu2,
)
# Convert from orbital elements to cartesian coordinates
# This is the position and velocity of the Jacobi coordinate
qjt_1, vjt_1, ajt_1 = model_e2c_1(elt1)
qjt_2, vjt_2, ajt_2 = model_e2c_2(elt2)
# Reshape the Jacobi coordinates to include an axis for body number
particle_traj_shape = (-1, 1, 3)
particle_traj_shape_layer = keras.layers.Reshape(
target_shape=particle_traj_shape, name='particle_traj_shape')
qjt_0 = particle_traj_shape_layer(qjt_0)
qjt_1 = particle_traj_shape_layer(qjt_1)
qjt_2 = particle_traj_shape_layer(qjt_2)
vjt_0 = particle_traj_shape_layer(vjt_0)
vjt_1 = particle_traj_shape_layer(vjt_1)
vjt_2 = particle_traj_shape_layer(vjt_2)
ajt_0 = particle_traj_shape_layer(ajt_0)
ajt_1 = particle_traj_shape_layer(ajt_1)
ajt_2 = particle_traj_shape_layer(ajt_2)
# Assemble the Jacobi coordinates over time
qj = keras.layers.concatenate(inputs=[qjt_0, qjt_1, qjt_2],
axis=-2,
name='qj')
vj = keras.layers.concatenate(inputs=[vjt_0, vjt_1, vjt_2],
axis=-2,
name='vj')
aj = keras.layers.concatenate(inputs=[ajt_0, ajt_1, ajt_2],
axis=-2,
name='aj')
# Convert the Jacobi coordinates over time to Cartesian coordinates
q, v, a = JacobiToCartesian(include_accel=True,
batch_size=batch_size)([m, qj, vj, aj])
# Name the outputs
q = Identity(name='q')(q)
v = Identity(name='v')(v)
a = Identity(name='a')(a)
# Wrap up the outputs
outputs = (q, v, a)
# Diagnostic - include computed orbital elements with the output
# outputs = outputs + elt1 + elt2
# Wrap this into a model
suffix = '_'.join(str(sz) for sz in hidden_sizes)
model_name = f'model_g3b_position_nn_{suffix}'
model = keras.Model(inputs=inputs, outputs=outputs, name=model_name)
return model
def make_physics_model_r2bc_math(position_model: keras.Model, traj_size: int):
"""Create a physics model for the restricted two body problem from a position model"""
# Create input layers
t = keras.Input(shape=(traj_size, ), name='t')
q0 = keras.Input(shape=(2, ), name='q0')
v0 = keras.Input(shape=(2, ), name='v0')
mu = keras.Input(shape=(1, ), name='mu')
# The combined input layers
inputs = [t, q0, v0, mu]
# Check sizes
batch_size = t.shape[0]
tf.debugging.assert_shapes(shapes={
t: (batch_size, traj_size),
q0: (batch_size, 2),
v0: (batch_size, 2),
mu: (batch_size, 1),
},
message='make_physics_model_r2bc_math / inputs')
# Return row 0 of a position or velocity for q0_rec and v0_rec
initial_row_func = lambda q: q[:, 0, :]
# The polar coordinates of the initial conditions
# r0, theta0, and omega0 each scalars in each batch
r0, theta0, omega0 = ConfigToPolar2D(name='polar0')([q0, v0])
# Name the outputs of the initial polar
# These each have shape (batch_size, 1)
r0 = Identity(name='r0')(r0)
theta0 = Identity(name='theta0')(theta0)
omega0 = Identity(name='omega0')(omega0)
# Check sizes
tf.debugging.assert_shapes(
shapes={
r0: (batch_size, 1),
theta0: (batch_size, 1),
omega0: (batch_size, 1),
},
message=
'make_physics_model_r2bc_math / polar elements r0, theta0, omega0')
# Compute the motion from the specified position layer
q, v, a = Motion_R2BC(position_model=position_model,
name='motion')([t, r0, theta0, omega0])
# Name the outputs of the circular motion
# These each have shape (batch_size, traj_size, 2)
q = Identity(name='q')(q)
v = Identity(name='v')(v)
a = Identity(name='a')(a)
# Check sizes
tf.debugging.assert_shapes(
shapes={
q: (batch_size, traj_size, 2),
v: (batch_size, traj_size, 2),
a: (batch_size, traj_size, 2),
},
message='make_physics_model_r2bc_math / outputs q, v, a')
# Compute q0_rec and v0_rec
# These each have shape (batch_size, 2)
q0_rec = keras.layers.Lambda(initial_row_func, name='q0_rec')(q)
v0_rec = keras.layers.Lambda(initial_row_func, name='v0_rec')(v)
# Check sizes
tf.debugging.assert_shapes(
shapes={
q0_rec: (batch_size, 2),
v0_rec: (batch_size, 2),
},
message='make_physics_model_r2bc_math / outputs q0_rec, v0_rec')
# Compute kinetic energy T and potential energy U
T = KineticEnergy_R2BC(name='T')(v)
U = PotentialEnergy_R2BC(name='U')([q, mu])
# Compute the total energy H
H = keras.layers.add(inputs=[T, U], name='H')
# Compute angular momentum L
# This has shape (batch_size, traj_size)
L = AngularMomentum_R2BC(name='L')([q, v])
# Check sizes
tf.debugging.assert_shapes(
shapes={
T: (batch_size, traj_size),
U: (batch_size, traj_size),
H: (batch_size, traj_size),
L: (batch_size, traj_size),
},
message='make_physics_model_r2bc_math / outputs H, L')
# Wrap this up into a model
outputs = [q, v, a, q0_rec, v0_rec, H, L]
model = keras.Model(inputs=inputs, outputs=outputs, name='model_math')
return model
def make_position_model_r2bc_nn(hidden_sizes, skip_layers=True, traj_size = 731):
"""
Compute orbit positions for the restricted two body circular problem from
the initial polar coordinates (orbital elements) with a deterministic mathematical model.
Factory function that returns a functional model.
"""
# Create input layers
t = keras.Input(shape=(traj_size,), name='t')
q0 = keras.Input(shape=(2,), name='q0')
v0 = keras.Input(shape=(2,), name='v0')
mu = keras.Input(shape=(1,), name='mu')
# The combined input layers
inputs = [t, q0, v0, mu]
# Check sizes
batch_size = t.shape[0]
tf.debugging.assert_shapes(shapes={
t: (batch_size, traj_size),
q0: (batch_size, 2),
v0: (batch_size, 2),
mu: (batch_size, 1),
}, message='make_position_model_r2bc_nn / inputs')
# The polar coordinates of the initial conditions
# r0, theta0, and omega0 each scalars in each batch
r0, theta0, omega0 = ConfigToPolar2D(name='polar0')([q0, v0])
# Name the outputs of the initial polar
# These each have shape (batch_size, 1)
r0 = Identity(name='r0')(r0)
theta0 = Identity(name='theta0')(theta0)
omega0 = Identity(name='omega0')(omega0)
# Check sizes
tf.debugging.assert_shapes(shapes={
r0: (batch_size, 1),
theta0: (batch_size, 1),
omega0: (batch_size, 1),
}, message='make_position_model_r2bc_nn / polar elements r0, theta0, omega0')
# Combine the trajectory-wide scalars into one feature
# Size of each row is 2+2+1+1+1=7; shape is (batch_size, 7)
phi_traj = keras.layers.concatenate(inputs=[q0, v0, r0, theta0, omega0], name='phi_traj')
# Repeat phi_traj traj_size times so it has a shape of (batch_size, traj_size, 7)
phi_traj_vec = keras.layers.RepeatVector(n=traj_size, name='phi_traj_vec')(phi_traj)
# Reshape t to (batch_size, traj_size, 1)
t_vec = keras.layers.Reshape(target_shape=(traj_size, 1), name='t_vec')(t)
# Concatenate phi_traj with the time t to make the initial feature vector phi_0
phi_0 = keras.layers.concatenate(inputs=[t_vec, phi_traj_vec], name='phi_0')
# Hidden layers as specified in hidden_sizes
# Number of hidden layers
num_layers = len(hidden_sizes)
# phi_n will update to the last available feature layer for the output portion
phi_n = phi_0
# First hidden layer if applicable
if num_layers > 0:
phi_1 = keras.layers.Dense(units=hidden_sizes[0], activation='tanh', name='phi_1')(phi_0)
if skip_layers:
phi_1 = keras.layers.concatenate(inputs=[phi_0, phi_1], name='phi_1_aug')
phi_n = phi_1
# Second hidden layer if applicable
if num_layers > 1:
phi_2 = keras.layers.Dense(units=hidden_sizes[1], activation='tanh', name='phi_2')(phi_1)
if skip_layers:
phi_2 = keras.layers.concatenate(inputs=[phi_1, phi_2], name='phi_2_aug')
phi_n = phi_2
# Check shapes
tf.debugging.assert_shapes(shapes={
phi_traj: (batch_size, 7),
phi_traj_vec: (batch_size, traj_size, 7),
t_vec: (batch_size, traj_size, 1),
phi_0: (batch_size, traj_size, 8),
# phi_1: (batch_size, traj_size, 'hs1'),
# phi_2: (batch_size, traj_size, 'hs2'),
}, message='make_position_model_r2bc_nn / phi')
# Compute the radius r from the features; initialize weights to 0, bias to 1
r = keras.layers.Dense(
units=1, kernel_initializer='zeros', bias_initializer='ones', name='r')(phi_n)
# Compute the mean angular frequency omega from the features
omega = keras.layers.Dense(
units=1, kernel_initializer='zeros', bias_initializer='zeros', name='omega') (phi_n)
# Add a feature for omega * t + theta0; a rough estimate of the angular offset
omega_t = keras.layers.multiply(inputs=[omega, t_vec], name='omega_t')
# Compute the offset to angle theta to its mean trend from the features
theta_adj = keras.layers.Dense(
units=1, kernel_initializer='zeros', bias_initializer='zeros', name='theta_adj')(phi_n)
# Compute the angle theta as the sum of its original offset theta0, its trend rate omega_bar_t
# and the adjustment above
theta0_vec = keras.layers.RepeatVector(n=traj_size, name='theta0_vec')(theta0)
theta = keras.layers.add(inputs=[omega_t, theta0_vec, theta_adj], name='theta')
# Check shapes
tf.debugging.assert_shapes(shapes={
r: (batch_size, traj_size, 1),
omega: (batch_size, traj_size, 1),
omega_t: (batch_size, traj_size, 1),
theta_adj: (batch_size, traj_size, 1),
theta0_vec: (batch_size, traj_size, 1),
theta: (batch_size, traj_size, 1)
}, message='make_position_model_r2bc_nn / r, omega, theta')
# Cosine and sine of theta
cos_theta = keras.layers.Activation(activation=tf.cos, name='cos_theta')(theta)
sin_theta = keras.layers.Activation(activation=tf.sin, name='sin_theta')(theta)
# Compute qx and qy from r, theta
qx = keras.layers.multiply(inputs=[r, cos_theta], name='qx')
qy = keras.layers.multiply(inputs=[r, sin_theta], name='qy')
# Check shapes
batch_size = t.shape[0]
tf.debugging.assert_shapes(shapes={
cos_theta: (batch_size, traj_size, 1),
sin_theta: (batch_size, traj_size, 1),
qx: (batch_size, traj_size, 1),
qy: (batch_size, traj_size, 1),
}, message='make_position_model_r2bc_nn / outputs')
# Wrap this into a model
outputs = [qx, qy]
model = keras.Model(inputs=inputs, outputs=outputs, name='model_r2bc_position_nn')
return model
def make_position_model_g3b_math(traj_size: int = 1001,
batch_size: int = 64,
num_gpus: int = 1):
"""
Compute orbit positions for the general two body problem from
the initial orbital elements with a deterministic mathematical model.
Factory function that returns a functional model.
"""
num_particles = 3
space_dims = 3
full_batch_size = batch_size * num_gpus
t = keras.Input(shape=(traj_size, ), batch_size=full_batch_size, name='t')
q0 = keras.Input(shape=(
num_particles,
space_dims,
),
batch_size=full_batch_size,
name='q0')
v0 = keras.Input(shape=(
num_particles,
space_dims,
),
batch_size=full_batch_size,
name='v0')
m = keras.Input(shape=(num_particles, ),
batch_size=full_batch_size,
name='m')
# Wrap these up into one tuple of inputs for the model
inputs = (t, q0, v0, m)
# Compute the Jacobi coordinates of the initial conditions
qj0, vj0, mu0 = CartesianToJacobi()([m, q0, v0])
# Extract Jacobi coordinates of p1 and p2
qj0_1 = qj0[:, 1, :]
qj0_2 = qj0[:, 2, :]
vj0_1 = vj0[:, 1, :]
vj0_2 = vj0[:, 2, :]
# Extract gravitational field strength for orbital element conversion of p1 and p2
mu0_1 = mu0[:, 1:2]
mu0_2 = mu0[:, 2:3]
# Manually set the shapes to work around documented bug on slices losing shape info
jacobi_shape = (batch_size, space_dims)
qj0_1.set_shape(jacobi_shape)
qj0_2.set_shape(jacobi_shape)
vj0_1.set_shape(jacobi_shape)
vj0_1.set_shape(jacobi_shape)
mu_shape = (batch_size, 1)
mu0_1.set_shape(mu_shape)
mu0_2.set_shape(mu_shape)
# Tuple of inputs for the model converting from configuration to orbital elements
cfg_1 = (qj0_1, vj0_1, mu0_1)
cfg_2 = (qj0_2, vj0_2, mu0_2)
# Model mapping cartesian coordinates to orbital elements
model_c2e_1 = make_model_cfg_to_elt(name='orbital_element_1')
model_c2e_2 = make_model_cfg_to_elt(name='orbital_element_2')
# Extract the orbital elements of the initial conditions
a1_0, e1_0, inc1_0, Omega1_0, omega1_0, f1_0, M1_0, N1_0 = model_c2e_1(
cfg_1)
a2_0, e2_0, inc2_0, Omega2_0, omega2_0, f2_0, M2_0, N2_0 = model_c2e_2(
cfg_2)
# Alias mu0_i for naming consistency
mu1_0 = mu0_1
mu2_0 = mu0_2
# Reshape t to (batch_size, traj_size, 1)
t_vec = keras.layers.Reshape(target_shape=(traj_size, 1), name='t_vec')(t)
# ******************************************************************
# Predict orbital elements for Jacobi coordinates of body 1
# Repeat the constant orbital elements to be vectors of shape (batch_size, traj_size)
a1 = keras.layers.RepeatVector(n=traj_size, name='a1')(a1_0)
e1 = keras.layers.RepeatVector(n=traj_size, name='e1')(e1_0)
inc1 = keras.layers.RepeatVector(n=traj_size, name='inc1')(inc1_0)
Omega1 = keras.layers.RepeatVector(n=traj_size, name='Omega1')(Omega1_0)
omega1 = keras.layers.RepeatVector(n=traj_size, name='omega1')(omega1_0)
mu1 = keras.layers.RepeatVector(n=traj_size, name='mu1')(mu1_0)
# Repeat initial mean anomaly M0 and mean motion N0 to match shape of outputs
M1_0_vec = keras.layers.RepeatVector(n=traj_size, name='M1_0_vec')(M1_0)
N1_0_vec = keras.layers.RepeatVector(n=traj_size, name='N1_0_vec')(N1_0)
# Compute the mean anomaly M(t) as a function of time
N1_t = keras.layers.multiply(inputs=[N1_0_vec, t_vec])
M1 = keras.layers.add(inputs=[M1_0_vec, N1_t])
# Compute the true anomaly from the mean anomly and eccentricity
f1 = MeanToTrueAnomaly(name='mean_to_true_anomaly_f1')([M1, e1])
# Wrap orbital elements into one tuple of inputs for layer converting to cartesian coordinates
elt1 = (
a1,
e1,
inc1,
Omega1,
omega1,
f1,
mu1,
)
# ******************************************************************
# Predict orbital elements for Jacobi coordinates of body 2
# Repeat the constant orbital elements to be vectors of shape (batch_size, traj_size)
a2 = keras.layers.RepeatVector(n=traj_size, name='a2')(a2_0)
e2 = keras.layers.RepeatVector(n=traj_size, name='e2')(e2_0)
inc2 = keras.layers.RepeatVector(n=traj_size, name='inc2')(inc2_0)
Omega2 = keras.layers.RepeatVector(n=traj_size, name='Omega2')(Omega2_0)
omega2 = keras.layers.RepeatVector(n=traj_size, name='omega2')(omega2_0)
mu2 = keras.layers.RepeatVector(n=traj_size, name='mu2')(mu2_0)
# Repeat initial mean anomaly M0 and mean motion N0 to match shape of outputs
M2_0_vec = keras.layers.RepeatVector(n=traj_size, name='M2_0_vec')(M2_0)
N2_0_vec = keras.layers.RepeatVector(n=traj_size, name='N2_0_vec')(N2_0)
# Compute the mean anomaly M(t) as a function of time
N2_t = keras.layers.multiply(inputs=[N2_0_vec, t_vec])
M2 = keras.layers.add(inputs=[M2_0_vec, N2_t])
# Compute the true anomaly from the mean anomly and eccentricity
f2 = MeanToTrueAnomaly(name='mean_to_true_anomaly_f2')([M2, e2])
# Wrap orbital elements into one tuple of inputs for layer converting to cartesian coordinates
elt2 = (
a2,
e2,
inc2,
Omega2,
omega2,
f2,
mu2,
)
# ******************************************************************
# Convert orbital elements to cartesian Jacobi coordinates
# Model mapping orbital elements to cartesian coordinates
model_e2c = make_model_elt_to_cfg(include_accel=True,
batch_size=batch_size)
# The position of Jacobi coordinate 0 over time comes from the average velocity
# We always use center of momentum coordinates, so this is zero
qjt_0 = keras.backend.zeros(shape=[batch_size, traj_size, space_dims])
vjt_0 = keras.backend.zeros(shape=[batch_size, traj_size, space_dims])
ajt_0 = keras.backend.zeros(shape=[batch_size, traj_size, space_dims])
# Convert from orbital elements to cartesian coordinates
# This is the position and velocity of the Jacobi coordinate
qjt_1, vjt_1, ajt_1 = model_e2c(elt1)
qjt_2, vjt_2, ajt_2 = model_e2c(elt2)
# Reshape the Jacobi coordinates to include an axis for body number
particle_traj_shape = (-1, 1, 3)
particle_traj_shape_layer = keras.layers.Reshape(
target_shape=particle_traj_shape, name='particle_traj_shape')
qjt_0 = particle_traj_shape_layer(qjt_0)
qjt_1 = particle_traj_shape_layer(qjt_1)
qjt_2 = particle_traj_shape_layer(qjt_2)
vjt_0 = particle_traj_shape_layer(vjt_0)
vjt_1 = particle_traj_shape_layer(vjt_1)
vjt_2 = particle_traj_shape_layer(vjt_2)
ajt_0 = particle_traj_shape_layer(ajt_0)
ajt_1 = particle_traj_shape_layer(ajt_1)
ajt_2 = particle_traj_shape_layer(ajt_2)
# Assemble the Jacobi coordinates over time
qj = keras.layers.concatenate(inputs=[qjt_0, qjt_1, qjt_2],
axis=-2,
name='qj')
vj = keras.layers.concatenate(inputs=[vjt_0, vjt_1, vjt_2],
axis=-2,
name='vj')
aj = keras.layers.concatenate(inputs=[ajt_0, ajt_1, ajt_2],
axis=-2,
name='aj')
# Convert the Jacobi coordinates over time to Cartesian coordinates
q, v, a = JacobiToCartesian(include_accel=True,
batch_size=batch_size)([m, qj, vj, aj])
# Name the outputs
q = Identity(name='q')(q)
v = Identity(name='v')(v)
a = Identity(name='a')(a)
# Wrap up the outputs
outputs = (q, v, a)
# Wrap this into a model
model = keras.Model(inputs=inputs,
outputs=outputs,
name='model_g3b_position_math')
return model