Skip to content
Permalink
Browse files

cleaned zerocenter

  • Loading branch information...
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):

0 comments on commit 9cb3444

Please sign in to comment.
You can’t perform that action at this time.