emcee proposal

examples/mcmc/himmelblau.py

@@ -17,11 +17,12 @@ def loglike(z):
         return - (z1**2 + z2 - 11.)**2 - (z1 + z2**2 - 7.)**2
 
     def transform(x):
-        return 5. * x
+        return 500. * x
 
     sampler = MCMCSampler(args.x_dim, loglike, transform=transform, log_dir=args.log_dir, hidden_dim=args.hidden_dim,
                           num_layers=args.num_layers, num_blocks=args.num_blocks, num_slow=args.num_slow)
-    sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10)
+    sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10, 
+                bootstrap_iters=args.burnin_iters)
 
 
 if __name__ == '__main__':
@@ -32,6 +33,8 @@ def transform(x):
                         help="Dimensionality")
     parser.add_argument('--train_iters', type=int, default=200,
                         help="number of train iters")
+    parser.add_argument('--burnin_iters', type=int, default=1,
+                        help="number of iters for finding good training spot")
     parser.add_argument("--mcmc_steps", type=int, default=10000)
     parser.add_argument('--hidden_dim', type=int, default=128)
     parser.add_argument('--num_layers', type=int, default=2)

examples/mcmc/rosenbrock.py

@@ -11,16 +11,23 @@ def main(args):
     from nnest.mcmc import MCMCSampler
 
     os.environ['CUDA_VISIBLE_DEVICES'] = ''
-
-    def loglike(z):
+    
+    def loglike_orig(z):
         return np.array([-sum(100.0 * (x[1:] - x[:-1] ** 2.0) ** 2.0 + (1 - x[:-1]) ** 2.0) for x in z])
+
+    # pairwise rosenbrocks, multiplied together
+    assert args.x_dim % 2 == 0
+
+    def loglike(z):
+        return np.array([-sum(100.0 * (x[1::2] - x[::2] ** 2.0) ** 2.0 + (1 - x[::2]) ** 2.0) * 2. / len(x) for x in z])
 
     def transform(x):
-        return 5. * x
+        return 500. * x - 100
 
     sampler = MCMCSampler(args.x_dim, loglike, transform=transform, log_dir=args.log_dir, hidden_dim=args.hidden_dim,
                           num_layers=args.num_layers, num_blocks=args.num_blocks, num_slow=args.num_slow)
-    sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10)
+    sampler.run(train_iters=args.train_iters, mcmc_steps=args.mcmc_steps, single_thin=10, 
+                bootstrap_iters=args.burnin_iters, bootstrap_mcmc_steps=5000 + 1000 * args.x_dim)
 
 
 if __name__ == '__main__':
@@ -30,6 +37,8 @@ def transform(x):
                         help="Dimensionality")
     parser.add_argument('--train_iters', type=int, default=100,
                         help="number of train iters")
+    parser.add_argument('--burnin_iters', type=int, default=1,
+                        help="number of iters for finding good training spot")
     parser.add_argument("--mcmc_steps", type=int, default=10000)
     parser.add_argument('--hidden_dim', type=int, default=128)
     parser.add_argument('--num_layers', type=int, default=1)

nnest/mcmc.py

@@ -48,7 +48,7 @@ def _init_samples(self, mcmc_steps=5000, mcmc_batch_size=5, ignore_rows=0.3):
         logl = self.loglike(v)
         samples = []
         likes = []
-        for i in range(mcmc_steps):
+        for i in range(mcmc_steps // mcmc_batch_size):
             du = np.random.standard_normal(u.shape) * 0.1
             u_prime = u + du
             v_prime = self.transform(u_prime)
@@ -73,11 +73,13 @@ def _init_samples(self, mcmc_steps=5000, mcmc_batch_size=5, ignore_rows=0.3):
             u = u_prime * m + u * (1 - m)
             v = v_prime * m + v * (1 - m)
             logl = logl_prime * mask + logl * (1 - mask)
+            assert v.shape[1] == self.x_dim, v.shape
             samples.append(v)
             likes.append(logl)
         samples = np.transpose(np.array(samples), axes=[1, 0, 2])
         loglikes = -np.transpose(np.array(likes), axes=[1, 0])
         weights = np.ones(loglikes.shape)
+        assert samples.shape[-1] == self.x_dim, samples.shape
         self._chain_stats(samples)
         self._save_samples(samples, weights, loglikes)
         names = ['p%i' % i for i in range(int(self.x_dim))]
@@ -116,7 +118,9 @@ def run(
 
         if self.log:
             self.logger.info('Alpha [%5.4f]' % (alpha))
-
+
+        allsamples = []
+        next_start = None
         for t in range(bootstrap_iters):
 
             if t == 0:
@@ -128,26 +132,43 @@ def run(
             else:
                 samples, likes, latent, scale, nc = self.trainer.sample(
                     loglike=self.loglike, transform=transform,
-                    mcmc_steps=bootstrap_mcmc_steps, alpha=alpha, dynamic=False, show_progress=True)
+                    mcmc_steps=bootstrap_mcmc_steps, 
+                    alpha=alpha, dynamic=False, show_progress=True,
+                    init_x=next_start, plot=True)
+                next_start = samples[:,-1,:]
+                #next_start = None
                 samples = transform(samples)
                 self._chain_stats(samples)
+                #self.logger.info('Last sample: %s' % (next_start))
                 loglikes = -np.array(likes)
                 weights = np.ones(loglikes.shape)
                 mc = MCSamples(samples=[samples[0]], weights=[weights[0]], loglikes=[loglikes[0]],
                                ignore_rows=ignore_rows)
 
             samples = mc.makeSingleSamples(single_thin=single_thin)
+            print(samples.shape)
             samples = samples[:, :self.x_dim]
+            assert samples.shape[1] == self.x_dim
             mean = np.mean(samples, axis=0)
             std = np.std(samples, axis=0)
             samples = (samples - mean) / std
+            self.logger.info('%d new samples' % (len(samples)))
+            if t > 0:
+                allsamples = np.vstack((allsamples, samples))
+                samples = allsamples[int(len(allsamples)//3):,:]
+            else:
+                allsamples = samples
+            self.logger.info('Training for bootstrap iter %d with %d samples' % (t, len(samples)))
+            # Forget the current network
+            self.trainer.init_network()
             self.trainer.train(samples, max_iters=train_iters, noise=-1)
 
             def transform(x):
                 return x * std + mean
 
+        self.logger.info('Sampling...')
         samples, likes, latent, scale, nc = self.trainer.sample(
-            loglike=self.loglike, transform=transform,
+            loglike=self.loglike, transform=transform, init_x=next_start,
             mcmc_steps=mcmc_steps, alpha=alpha, dynamic=False, show_progress=True,
             out_chain=os.path.join(self.logs['chains'], 'chain'))
         samples = transform(samples)