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)