def hamiltonian_vector_field(hamiltonian, x, t, backward_time=False):
"""X_H appearing in Hamilton's eqs: xdot = X_H(x)"""
# note, t unused.
dx = tf.gradients(hamiltonian(x), x)[0]
dq, dp = extract_q_p(dx)
if not backward_time:
return join_q_p(dp, -dq)
else:
return join_q_p(-dp, dq)
def sample(self, N):
F1 = tf.reshape(self.base_dist_F1.sample(N), shape=[N,1]) # sh = [N,1]
otherF = self.base_dist_otherF.sample(N) # sh = [N,d*n-1]
F = tf.concat([F1, otherF], 1) # sh = [N,d*n]
F = tf.reshape(F, shape=[N]+self.sh) # sh = [N,d,n,1]
Psi = self.base_dist_Psi.sample(N)
return join_q_p(Psi, F)
def test_closed_toda_3():
# Choose points such that q1+q2+q3 = p1+p2+p3 = 0
q = tf.reshape([.1, .2, -.3], [1, 1, 3, 1])
p = tf.reshape([.3, .5, -.8], [1, 1, 3, 1])
h1 = closed_toda(join_q_p(q, p))
# Fourier modes:
# a1 = x + i y
# = 1/sqrt(3)(1/w .1 + w .2 - .3) =
# = 1/sqrt(3)((-1/2 - i sqrt(3)/2) .1 + (-1/2 + i sqrt(3)/2) .2 - .3)
x = 1 / tf.sqrt(3.) * (-1 / 2. * (.1 + .2) - .3)
# y = 1/tf.sqrt(3.)( tf.sqrt(3.)/2. * (.1 - .2) )
y = 1 / 2. * (-.1 + .2)
# b1 = px + i py = 1/sqrt(3)(1/w p1 + w p2 + p3)
# = 1/sqrt(3)( (-1/2 - i sqrt(3)/2) .3 + (1/2 + i sqrt(3)/2) .5 - .8)
# = 1/sqrt(3)( -1/2 * (.3 + .5) - .8) - i 1/2 (.3 - .5)
px = 1 / tf.sqrt(3.) * (-1 / 2. * (.3 + .5) - .8)
py = -1 / 2. * (.3 - .5)
kin = px**2 + py**2
expected_kin = tf.reduce_sum(tf.square(p)) * .5
assert_allclose(kin, expected_kin)
q1 = q[0, 0, 0, 0]
q2 = q[0, 0, 1, 0]
q3 = q[0, 0, 2, 0]
expected_q12 = -2 * y
expected_q23 = y - tf.sqrt(3.) * x
expected_q31 = y + tf.sqrt(3.) * x
assert_allclose(q1 - q2, expected_q12)
assert_allclose(q2 - q3, expected_q23)
assert_allclose(q3 - q1, expected_q31)
h2 = kin + tf.exp(expected_q12) + tf.exp(expected_q23) + tf.exp(
expected_q31)
assert_allclose(h1, h2)
print('test_closed_toda_3 passed')
def inverse(self, z):
q, p = extract_q_p(x)
q_prime = self._f.inverse(q)
df = tf_gradients_ops.jacobian(self._f(q_prime),
q_prime,
use_pfor=True)
return join_q_p(q_prime,
tf.tensordot(df, p, [[4, 5, 6, 7], [0, 1, 2, 3]]))
def call(self, x):
q, p = extract_q_p(x)
q_prime = self._f(q)
# Df(q)^{-1} = D(f^{-1}( q_prime ))
df_inverse = tf_gradients_ops.jacobian(self._f.inverse(q_prime),
q_prime,
use_pfor=True)
return join_q_p(
q_prime, tf.tensordot(df_inverse, p, [[4, 5, 6, 7], [0, 1, 2, 3]]))
def inverse(self, x):
qq, pp = extract_q_p(x)
if self._first_only:
I000 = 0.5 * (tf.square(qq[:, 0, 0, 0]) +
tf.square(pp[:, 0, 0, 0]))
phi000 = tf.atan(qq[:, 0, 0, 0] / pp[:, 0, 0, 0])
I = pp
phi = qq
I = self._assign_000(I, I000)
phi = self._assign_000(phi, phi000)
else:
phi = tf.atan(qq / pp)
I = 0.5 * (tf.square(qq) + tf.square(pp))
return join_q_p(phi, I)
def call(self, z):
phi, I = extract_q_p(z)
if self._first_only:
sqrt_two_I = tf.sqrt(2. * I[:, 0, 0, 0])
q000 = tf.multiply(sqrt_two_I, tf.sin(phi[:, 0, 0, 0]))
p000 = tf.multiply(sqrt_two_I, tf.cos(phi[:, 0, 0, 0]))
p = I
q = phi
p = self._assign_000(p, p000)
q = self._assign_000(q, q000)
else:
sqrt_two_I = tf.sqrt(2. * I)
q = tf.multiply(sqrt_two_I, tf.sin(phi))
p = tf.multiply(sqrt_two_I, tf.cos(phi))
return join_q_p(q, p)
def call(self, z):
# Chose p as the action, q as the angle.
q, p = extract_q_p(z)
return join_q_p(q, self._flow(p))
def _reshape_and_join_q_p(self, q, p, sh):
q = tf.reshape(q, sh)
p = tf.reshape(p, sh)
return join_q_p(q, p)
def call(self, z):
phi, I = extract_q_p(z)
q = tf.sin(phi)
p = tf.multiply(I, 1 / (tf.cos(phi) + self.eps))
return join_q_p(q, p)
def inverse(self, x):
q, p = extract_q_p(x)
return join_q_p(-p, q)
def call(self, x):
q, p = extract_q_p(x)
return join_q_p(p, -q)
def inverse(self, x):
q, p = extract_q_p(x)
return join_q_p(q, p - self._shift_model(q))
def call(self, x):
q, p = extract_q_p(x)
return join_q_p(q, p + self._shift_model(q))
def inverse(self, x):
q, p = extract_q_p(x)
phi = tf.asin(q)
I = tf.multiply(p, tf.cos(phi) + self.eps)
return join_q_p(phi, I)
def inverse(self, x):
# Chose p as the action, q as the angle.
q, p = extract_q_p(x)
return join_q_p(q, self._flow.inverse(p))
def sample(self, N):
u = self.base_dist_u.sample(N)
phi = self.base_dist_phi.sample(N)
return join_q_p(phi, u)
def call(self, x):
q, p = extract_q_p(x)
return join_q_p(
q * tf.exp(self.scale_exponent),
tf.exp(-self.scale_exponent) * (p + self._shift_model(q)))
def inverse(self, x):
q, p = extract_q_p(x)
return join_q_p(q * tf.exp(-self.scale_exponent),
p * tf.exp(self.scale_exponent) - self._shift_model(q))
def make_loss(settings, T, inp):
name = settings['loss']
with tf.name_scope("loss"):
if name == "circle":
pass
else:
with tf.name_scope("canonical_transformation"):
x = T(inp)
if settings['visualize']:
q, p = extract_q_p(x)
tf.summary.histogram("q", q)
tf.summary.histogram("p", p)
K = settings['hamiltonian'](x)
if settings['visualize']:
tf.summary.histogram('K-Hamiltonian', K)
if name == "dKdphi":
# K independent of phi
dphi, _ = extract_q_p(tf.gradients(K, z)[0])
if settings['visualize']:
tf.summary.histogram('dKdphi', dphi)
# loss = tf.reduce_mean( tf.square(dphi) + \
# settings['elastic_net_coeff'] * tf.pow( tf.abs(dphi), 3 ) )
loss = tf.sqrt(tf.reduce_mean(0.5 * tf.square(dphi)))
if 'lambda_range' in settings:
# With penalizing K (energy) outside low,high:
range_reg = tf.reduce_mean(
confining_potential(K, settings['low_K_range'],
settings['high_K_range']))
loss += settings['lambda_range'] * range_reg
if 'lambda_diff' in settings:
# add |K-val|^2 term
diff_loss = tf.reduce_mean(
tf.square(K - settings['diff_val']))
loss += settings['lambda_diff'] * diff_loss
elif name == "conserved_radii":
# Action variables as radii
q_prime, p_prime = extract_q_p(z)
dq_prime, dp_prime = extract_q_p(tf.gradients(K, z)[0])
loss = tf.reduce_mean(
tf.square(q_prime * dp_prime - p_prime * dq_prime))
elif name == "K_equal_F1_loss":
# K = F_1
_, F = extract_q_p(z)
loss = tf.reduce_mean(tf.square(K - F[:, 0, 0, 0]))
elif name == "KL":
# Here T can contain a non-symplectic part such as scaling.
KL_loss = tf.reduce_mean(K - T.log_jacobian_det(z))
# Gradient penalty regularization:
gp = compute_gradK_penalty(K, z)
# Range regularization
range_reg = tf.reduce_mean(
confining_potential(x, settings['low_x_range'],
settings['high_x_range']))
if settings['visualize']:
tf.summary.scalar("KL_loss", KL_loss)
# Monitor the derivative of K to understand how well we are doing
# due to unknown Z in KL. Assume distribution propto e^-u_1.
tf.summary.scalar("gradient_penalty", gp)
tf.summary.scalar("range_reg", range_reg)
loss = KL_loss + \
settings['lambda_gp'] * gp + \
settings['lambda_range'] * range_reg
elif name == "K_equal_action_H":
# K = NN(I). Use MLP Hamiltonian
_, I = extract_q_p(z)
action_H = MLPHamiltonian()
H_I = action_H(I)
# Add a residual part corresponding to GGE, so that H should be small
# TODO: Need temperatures, otherwise, does not make too much sense,
# think about Kepler problem, where Is are > 0 and H < 0.
H_I += tf.reduce_sum(I, [1, 2, 3])
if settings['visualize']:
tf.summary.histogram('action-Hamiltonian', H_I)
loss = tf.reduce_mean(tf.square(K - H_I))
elif name == "K_indep_phi_T_periodic":
# K independent of phi, T periodic
phi, I = extract_q_p(z)
dphi, _ = extract_q_p(tf.gradients(K, z)[0])
# Impose K = K(I): normsq(dphi). (Here dphi has shape [N,d,n,1])
normsq_dphi = normsq_nobatch(dphi)
# MC estimate of the integral over I,phi:
loss = tf.reduce_mean(normsq_dphi)
# Impose phi periodic over periods (truncate): sum_n normsq(T(phi,I) - T(phi+2*pi*n,I))
periods = tf.constant([-1, 1], dtype=DTYPE)
periodic_constraint = tf.map_fn(\
lambda n : normsq_nobatch( x - T(join_q_p(phi + 2*np.pi*n, I)) ), periods)
periodic_constraint = tf.reduce_sum(periodic_constraint,
0) # sum_n
# MC estimate of the integral over I,phi:
periodic_constraint = tf.reduce_mean(periodic_constraint)
# Sum constraints
multiplier = 1.
loss += multiplier * periodic_constraint
else:
raise NameError('loss %s not implemented', name)
tf.summary.scalar('loss', loss)
return loss