# rbondesan/CanonicalFlows

cleaned zerocenter

rbondesan committed Apr 28, 2019
1 parent 6873a2e commit 9cb34440f63bb81d6be92beac81a4a9d43c51365
Showing with 341 additions and 206 deletions.
1. +2 −0 hamiltonians.py
2. +3 −9 losses.py
3. +18 −26 models.py
4. +291 −167 notebooks/kepler-traj.ipynb
5. +1 −1 tests/tlosses.py
6. +3 −3 tests/tmodels.py
7. +23 −0 utils.py
 @@ -34,6 +34,8 @@ 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.
 @@ -24,20 +24,14 @@ def make_circle_loss(z, shift=-1): 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)) return tf.reduce_sum(tf.square(diff_qhatsq + diff_phatsq)) return tf.reduce_mean(tf.square(diff_qhatsq + diff_phatsq)) # return tf.reduce_sum(tf.square(diff_qhatsq + diff_phatsq)) def make_loss(settings, T, inp): name = settings['loss'] with tf.name_scope("loss"): if name == "circle": with tf.name_scope("canonical_transformation"): batch = inp.shape[1] z = tf.reshape(inp, [settings['minibatch_size']*batch,settings['d'],settings['num_particles'],2]) z = T.inverse(z) z = tf.reshape(z, [settings['minibatch_size'],batch,settings['d'],settings['num_particles'],2]) loss = make_circle_loss(z) pass else: with tf.name_scope("canonical_transformation"): x = T(inp)
 @@ -370,7 +370,7 @@ def unsqueeze(self, x): ]) class LinearSymplectic(SymplecticFlow): def __init__(self, num_householder=2): def __init__(self, num_householder=2, householder_random_init=True): """ Uses pre-Iwasawa decomposition of Sp(n): |1 0|L^{-1} 0 | U @@ -382,7 +382,13 @@ def __init__(self, num_householder=2): """ super(LinearSymplectic, self).__init__() self.num_householder = num_householder if householder_random_init: self.init = tf.glorot_normal_initializer() else: # Take a constant (ones) initializer, so that the # even chain of reflections represents the identity. self.init = tf.ones def build(self, in_size): # Take dimension to be the full phase space dim 2n self.n = int(in_size[1]*in_size[2]) @@ -393,13 +399,11 @@ def build(self, in_size): P_init = tf.zeros(sh) self.L = tfe.Variable(L_init, name="L") self.P = tfe.Variable(P_init, name="P") init = tf.glorot_normal_initializer() # householder reflection vectors sh = [self.num_householder, self.n] self.a = tfe.Variable(init(sh), name="a") self.b = tfe.Variable(init(sh), name="b") self.a = tfe.Variable(self.init(sh), name="a") self.b = tfe.Variable(self.init(sh), name="b") # Diagonal unitary requires angles sh = [self.n] @@ -572,7 +576,7 @@ def inverse(self, z): # return z + self._shift class ZeroCenter(SymplecticFlow): def __init__(self, decay=0.99, debias=False, is_training_forward=True): def __init__(self, decay=0.99, debias=False): """Shifts activations to have zero center. Add learnable offset. call (forward) used during training, normalizes: y = x - mean(x) + offset @@ -605,7 +609,6 @@ def __init__(self, decay=0.99, debias=False, is_training_forward=True): # dtype = tf.int64, # trainable=False) self.is_training = True self.is_training_forward = is_training_forward def build(self, in_sz): num_channels = in_sz[-1] @@ -614,40 +617,29 @@ def build(self, in_sz): trainable=False) def call(self, x): if self.is_training and self.is_training_forward: if self.is_training: train_mean, minibatch_mean_per_channel = self.update_moving_mean(x) # This runs the with block after the ops it depends on, so that # moving_mean is updated every iteration. with tf.control_dependencies([train_mean]): return x - minibatch_mean_per_channel + self._offset elif not self.is_training and not self.is_training_forward: # Used in prediction when training with inverse return x + self._moving_mean - self._offset elif not self.is_training and self.is_training_forward: # Prediction - standard batchnorm return x - self._moving_mean + self._offset return x + minibatch_mean_per_channel - self._offset else: raise ValueError('Should not be used') # Prediction - standard batchnorm return x + self._moving_mean - self._offset def inverse(self, z): if self.built == False: self.build(tf.shape(z)) self.built = True if self.is_training and not self.is_training_forward: if self.is_training: train_mean, minibatch_mean_per_channel = self.update_moving_mean(z) # This runs the with block after the ops it depends on, so that # moving_mean is updated every iteration. with tf.control_dependencies([train_mean]): return z - minibatch_mean_per_channel + self._offset elif not self.is_training and self.is_training_forward: # Used in prediction when training with forward return z + self._moving_mean - self._offset elif not self.is_training and not self.is_training_forward: # Prediction - standard batchnorm return z - self._moving_mean + self._offset else: raise ValueError('Should not be used') return z - self._moving_mean + self._offset def update_moving_mean(self, x): minibatch_mean_per_channel = tf.reduce_mean(x, [0,1,2])

Large diffs are not rendered by default.

 @@ -42,7 +42,7 @@ def test_make_circle_loss(): z = tf.sqrt( tf.reshape(tf.range(0,tot,dtype=DTYPE),(N,b,d,n,2)) ) # = 2*(0 +1 - (4 +5 ))^2 # + 2*(2 +3 - (6 +7 ))^2 expected_loss = tf.constant(2. * ((1-(4+5))**2 + (5-(6+7))**2), dtype=DTYPE) expected_loss = tf.constant(2. * ((1-(4+5))**2 + (5-(6+7))**2), dtype=DTYPE) / 4. loss = make_circle_loss(z) assert_allclose(loss, expected_loss) print("test_make_circle_loss passed")
 @@ -340,20 +340,20 @@ def testZeroCenter(): moving_mean = mean_per_channel * (1-decay) model = ZeroCenter() y = model(x) assert_equal(y, x-mean_per_channel) assert_equal(y, x+mean_per_channel) # 2nd update x = 0.5 * tf.reshape(tf.range(0,16,dtype=DTYPE), shape=(2,2,2,2)) mean_per_channel = 0.5 * mean_per_channel moving_mean = decay * moving_mean + (1-decay) * mean_per_channel # moving_mean /= (1-decay**2) y = model(x) assert_equal(y, x-mean_per_channel) assert_equal(y, x+mean_per_channel) # Test prediction mode - still offset is zero. model.is_training = False y = model(x) assert_equal(y, x-moving_mean) assert_equal(y, x+moving_mean) assert_equal(x, model.inverse(y)) assert( is_symplectic(model,
 @@ -381,6 +381,29 @@ def hamiltonian_traj(H, init_state, settings, time=100, steps=200, rtol=1e-04, a tensor_state = tf.contrib.integrate.odeint(hamiltons_equations(H,settings), init_state, t, rtol, atol) return tensor_state # Traj2Circle utils def plot_traj(settings, qtraj, ptraj, qhat_traj=None, phat_traj=None, equal_aspect=True): cols=['r','g','b'] num_init_cond = qtraj.shape[1] fig = plt.figure(figsize=(4*settings['d'],4*num_init_cond)) for b in range(num_init_cond): for n in range(settings['d']): plt.subplot(num_init_cond, settings['d'], n + b*settings['d'] + 1) #nrows,ncols,index plt.plot(qtraj[:,b,n,0,0], ptraj[:,b,n,0,0], '+') if qhat_traj is not None and phat_traj is not None: plt.plot(qhat_traj[:,b,n,0,0], phat_traj[:,b,n,0,0], '*') if equal_aspect: plt.gca().set_aspect('equal', adjustable='box') def pull_back_traj(settings, T, x): """Returns tranformed trajectory x whose shape is (T,B,d,n,2).""" batch = x.shape[1] z = tf.reshape(x, [settings['minibatch_size']*batch,settings['d'],settings['num_particles'],2]) z = T.inverse(z) z = tf.reshape(z, [settings['minibatch_size'],batch,settings['d'],settings['num_particles'],2]) return z # TODO: Remove? # # Define the solver step # def rk4_step(f, t, x, eps):