def harmonic_oscillator(x):
"""Harmonic oscillator: 1/2 sum( p^2 + q^2 ) Assume x.shape =
(N,d,n,2), n uncoupled harmonic oscillators in d-dimensions
"""
q, p = extract_q_p(x)
return 0.5 * tf.reduce_sum(tf.square(p) + tf.square(q), axis=[1, 2, 3])
def make_circle_loss(z, shift=-1):
"""Here z is (time, batch, d, n, 2)"""
zsq = tf.square(z)
qhatsq, phatsq = extract_q_p(zsq)
diff_qhatsq = qhatsq - tf.roll(qhatsq, shift=shift, axis=0)
diff_phatsq = phatsq - tf.roll(phatsq, shift=shift, axis=0)
# TODO: sqrt?
return tf.reduce_mean(tf.square(diff_qhatsq + diff_phatsq))
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 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 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 closed_toda_3(x):
"""p_x^2 + p_y^2 + e^{-2 y}+ e^{y-\sqrt{3}x}+ e^{y+\sqrt{3}x}\,.
x.shape = (N,1,2,2). Normal mode expression if center mass fixed."""
q, p = extract_q_p(x)
x = q[:, 0, 0, 0]
y = q[:, 0, 1, 0]
V = tf.exp(
-2. * y) + tf.exp(y - tf.sqrt(3.) * x) + tf.exp(y + tf.sqrt(3.) * x)
return tf.reduce_sum(tf.square(p), axis=2) + V
def closed_toda(x):
"""1/2 \sum_{i=1}^n p_i^2 + \sum_{i=1}^{n} exp(q_i - q_{i+1}).
x.shape = (N,1,n,2)"""
q, p = extract_q_p(x)
# q2, q3, ... , qN, q1
qshift = tf.manip.roll(q, shift=-1, axis=2)
# q1-q2, q2-q3, ... , q{N-1}-qN,qN-q1
qdiff = q - qshift
return tf.reduce_sum(0.5 * tf.square(p) + tf.exp(qdiff), axis=2)
def open_toda(x):
"""1/2 \sum_{i=1}^N p_i^2 + \sum_{i=1}^{n-1} exp(q_i - q_{i+1}).
x.shape = (N,1,n,2)"""
q, p = extract_q_p(x)
# q2, q3, ... , qN, q1
qshift = tf.manip.roll(q, shift=-1, axis=2)
# q1-q2, q2-q3, ... , q{N-1}-qN -> omit qN-q1, so qdiff shape (N,1,n-1,1)
qdiff = q[:, :, :-1, :] - qshift[:, :, :-1, :]
V = tf.reduce_sum(tf.exp(qdiff), axis=2)
K = 0.5 * tf.reduce_sum(tf.square(p), axis=2)
return K + V
def kepler(x, k=1.0):
"""H = 1/2 sum_{i=1}^d p_i^2 + k/r, r = sqrt(sum_{i=1}^d q_i^2).
Assume x.shape = (N,d,1,2) with d=2,3.
V(r)=k/r, k>0 repulsive, k<0 attractive."""
assert (x.shape[2] == 1)
q, p = extract_q_p(x)
# The derivative of r wrt q is 1/sqrt(sum(q^2)), which is singular in 0.
# Cutoff r so that it is > eps.
eps = 1e-5
r = tf.sqrt(tf.reduce_sum(tf.square(q), axis=1) + eps)
return tf.squeeze(0.5 * tf.reduce_sum(tf.square(p), axis=1) + k / r)
def fpu_hamiltonian(x, alpha=1, beta=0):
"""\sum_{i=1}^N 1/2 [p_i^2 + (q_{i} - q_{i+1})^2] + \alpha/3 (q_{i} - q_{i+1})^3 + \beta/4 (q_{i} - q_{i+1})^4
(with q_{N+1} = q_1). x.shape = (batch, 1, n, 2)
"""
assert (x.shape[1] == 1)
q, p = extract_q_p(x)
qdiff = q - tf.manip.roll(q, shift=-1,
axis=2) # q - (q_2, q_3, ..., q_{N}, q_1)
h = tf.reduce_sum(0.5 * tf.square(p) + 0.5 * tf.square(qdiff) +
alpha / 3. * tf.pow(qdiff, 3) +
beta / 4. * tf.pow(qdiff, 4),
axis=2)
return h
def neumann_hamiltonian(x):
"""1/4 \sum_{i,j}^N J_{ij}^2 + 1/2 \sum_{i=1}^N k_i q_i^2"""
# ks is of shape (d,)
assert x.shape[2] == 1
q, p = extract_q_p(x)
q = q[:, :, 0, 0]
p = p[:, :, 0, 0]
# q,p of shape (N,d)
J = tf.einsum('ai,aj->aij', q, p)
J -= tf.transpose(J, perm=[0, 2, 1])
# J is of shape (N,d,d)
return tf.squeeze(
0.25 * tf.reduce_sum(tf.square(J), axis=[1,2]) + \
0.5 * tf.reduce_sum( tf.multiply(ks, tf.square(q)), axis=1 ))
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 calogero_moser(x, type, omegasq=1., gsq=1., mu=1.):
"""
Notation: https://www1.maths.leeds.ac.uk/~siru/papers/p38.pdf (2.38)
and http://www.scholarpedia.org/article/Calogero-Moser_system
H = \frac{1}{2}\sum_{i=1}^n (p_i^2 + \omega^2 q_i^2) + g^2 \sum_{1\le j < k \le n} V(q_j - q_k)
V(x) =
'rational': 1/x^2
'hyperbolic': \mu^2/4\sinh(\mu x/2)
'trigonometric': \mu^2/4\sin(\mu x/2)
"""
assert (x.shape[1] == 1)
if type == 'rational':
V = lambda x: 1 / x**2
elif type == 'hyperbolic':
V = lambda x: mu**2 / 4. / tf.sinh(mu / 2. * x)
elif type == 'trigonometric':
V = lambda x: mu**2 / 4. / tf.sin(mu / 2. * x)
else:
raise NotImplementedError
q, p = extract_q_p(x)
h_free = 0.5 * tf.reduce_sum(tf.square(p) + omegasq * tf.square(q),
axis=[1, 2, 3]) # (batch,)
# Compute matrix of deltaq q[i]-q[j] and extract upper triangular part (triu)
q = tf.squeeze(q, 1) # (N,n,1)
deltaq = tf.transpose(q, [0, 2, 1]) - q # (N,n,n)
n = tf.shape(deltaq)[1]
ones = tf.ones([n, n])
triu_mask = tf.cast(tf.matrix_band_part(ones, 0, -1) - \
tf.matrix_band_part(ones, 0, 0), dtype=tf.bool)
triu = tf.boolean_mask(deltaq, triu_mask, axis=1)
eps = 1e-5 # regulizer for inverse
h_int = gsq * tf.reduce_sum(V(triu + eps), axis=1) # (batch,)
return h_free + h_int
# sum_{i<j} V(x(i) - x(j)) :
# e.g.:
# x = np.arange(10)
# x = np.expand_dims(x, 1)
# diffs = np.transpose(x) - x
# idx = np.triu_indices(np.shape(diffs)[1],k=1)
# np.sum(V(diffs[idx]))
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 toy_hamiltonian(x):
"""1/2 * (q-1/4 p^2)^2 + 1/32 p^2
Assume x.shape = (N,1,n,2) with n=1,2,..."""
q, p = extract_q_p(x)
pSqr = tf.square(p)
return 1 / 2 * tf.square(q - 1 / 4 * pSqr) + 1 / 32 * pSqr
def free(x):
"""H = 1/2 sum_{i=1}^d p_i^2.
Assume x.shape = (N,d,1,2) with d=1,2,3,..."""
_, p = extract_q_p(x)
return tf.squeeze(0.5 * tf.reduce_sum(tf.square(p), axis=1))
def _extract_and_reshape_q_p(self, x):
q, p = extract_q_p(x)
sh = tf.shape(q)
q = tf.reshape(q, [sh[0], 1, 1, sh[1] * sh[2]])
p = tf.reshape(p, [sh[0], 1, 1, sh[1] * sh[2]])
return q, p, sh
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 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):
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 _compute_log_scale_shift(self, zb):
# Assume _nn puts the extra copies along the channel dim -> extract q,p
log_scale, shift = extract_q_p(self._nn(zb))
if self._is_positive_shift:
shift = tf.abs(shift)
return log_scale, shift
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
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 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 log_jacobian_det(self, z):
# Chose p as the action, q as the angle.
q, p = extract_q_p(z)
return self._flow.log_jacobian_det(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)
