Elaborate NS run using a loop

python/eos/analysis.py

@@ -38,7 +38,7 @@ def _repr_html_(self):
             name = p.name()
             latex = p.latex()
             name = latex if latex else name
-            result += '<tr><td>{n}</td><td>{v:6.4f}</td></tr>'.format(n=name, v=v)
+            result += f'<tr><td>{name}</td><td>{v:6.4f}</td></tr>'
         result += '</table>'
 
         return(result)
@@ -87,13 +87,13 @@ def __init__(self, priors, likelihood, external_likelihood=[], global_options={}
                 eos.debug(' - {name} (constraint)'.format(name=p['constraint']))
         eos.debug('constraints:')
         for cn in likelihood:
-            eos.debug(' - {name}'.format(name=cn))
+            eos.debug(f' - {cn}')
         eos.debug('manual_constraints:')
         for cn, ce in manual_constraints.items():
-            eos.debug(' - {name}'.format(name=cn))
+            eos.debug(f' - {cn}')
         eos.debug('fixed_parameters:')
         for pn, pe in fixed_parameters.items():
-            eos.debug(' - {name}'.format(name=pn))
+            eos.debug(f' - {pn}')
 
         # collect the global options
         for key, value in global_options.items():
@@ -161,7 +161,7 @@ def __init__(self, priors, likelihood, external_likelihood=[], global_options={}
                         eos.LogPrior.Poisson(self.parameters, parameter, k),
                         False)
                 else:
-                    raise ValueError('Unknown prior type \'{}\''.format(prior_type))
+                    raise ValueError(f'Unknown prior type \'{prior_type}\'')
 
                 p = self.parameters[parameter]
                 self.varied_parameters.append(p)
@@ -216,11 +216,11 @@ def __init__(self, priors, likelihood, external_likelihood=[], global_options={}
 
         used_but_unvaried = used_parameter_names - varied_parameter_names - fixed_parameter_names
         if (len(used_but_unvaried) > 0):
-            eos.info('likelihood probably depends on {} parameter(s) that do not appear in the prior; check prior?'.format(len(used_but_unvaried)))
+            eos.info(f'likelihood probably depends on {len(used_but_unvaried)} parameter(s) that do not appear in the prior; check prior?')
         for n in used_but_unvaried:
-            eos.debug('used, but not included in any prior: \'{}\''.format(n))
+            eos.debug(f'used, but not included in any prior: \'{n}\'')
         for n in varied_parameter_names - used_parameter_names:
-            eos.warn('likelihood does not depend on parameter \'{}\'; remove from prior or check options!'.format(n))
+            eos.warn(f'likelihood does not depend on parameter \'{n}\'; remove from prior or check options!')
 
 
     def _u_to_par(self, u):
@@ -257,7 +257,7 @@ def _sanitize_manual_input(data):
             return list(map(Analysis._sanitize_manual_input, data))
 
         # all valid cases are covered above
-        raise ValueError("Unexpected entry type {} in manual_constraint".format(type(data)))
+        raise ValueError(f"Unexpected entry type {type(data)} in manual_constraint")
 
 
     @staticmethod
@@ -335,7 +335,7 @@ def optimize(self, start_point=None, rng=np.random.mtrand, **kwargs):
             eos.warn('Optimization did not succeed')
             eos.warn('  optimizer'' message reads: {}'.format(res.message))
         else:
-            eos.info('Optimization goal achieved after {nfev} function evaluations'.format(nfev=res.nfev))
+            eos.info(f'Optimization goal achieved after {res.nfev} function evaluations')
 
         bfp = self._u_to_par(res.x)
 
@@ -359,9 +359,9 @@ def log_pdf(self, u, *args):
         try:
             return(self._log_posterior.evaluate())
         except RuntimeError as e:
-            eos.error('encountered run time error ({e}) when evaluating log(posterior) in parameter point:'.format(e=e))
+            eos.error(f'encountered run time error ({e}) when evaluating log(posterior) in parameter point:')
             for p in self.varied_parameters:
-                eos.error(' - {n}: {v}'.format(n=p.name(), v=p.evaluate()))
+                eos.error(f' - {p.name()}: {p.evaluate()}')
             return(-np.inf)
 
 
@@ -428,10 +428,10 @@ def sample(self, N=1000, stride=5, pre_N=150, preruns=3, cov_scale=0.1, observab
 
         # pre run to adapt markov chains
         for i in progressbar(range(0, preruns), desc="Pre-runs", leave=False):
-            eos.info('Prerun {} out of {}'.format(i, preruns))
+            eos.info(f'Prerun {i} out of {preruns}')
             accept_count = sampler.run(pre_N)
             accept_rate  = accept_count / pre_N * 100
-            eos.info('Prerun {}: acceptance rate is {:3.0f}%'.format(i, accept_rate))
+            eos.info(f'Prerun {i}: acceptance rate is {accept_rate:3.0f}%')
             sampler.adapt()
         sampler.clear()
 
@@ -444,7 +444,7 @@ def sample(self, N=1000, stride=5, pre_N=150, preruns=3, cov_scale=0.1, observab
         for current_chunk in progressbar(sample_chunks, desc="Main run", leave=False):
             accept_count = accept_count + sampler.run(current_chunk)
         accept_rate  = accept_count / (N * stride) * 100
-        eos.info('Main run: acceptance rate is {:3.0f}%'.format(accept_rate))
+        eos.info(f'Main run: acceptance rate is {accept_rate:3.0f}%')
 
         # Transform from generator values in u space to the parameter values
         u_samples = sampler.samples[:][::stride]
@@ -621,9 +621,9 @@ def log_likelihood(self, p, *args):
         try:
             return(self._log_likelihood.evaluate())
         except RuntimeError as e:
-            eos.error('encountered run time error ({e}) when evaluating log(posterior) in parameter point:'.format(e=e))
+            eos.error(f'encountered run time error ({e}) when evaluating log(posterior) in parameter point:')
             for p in self.varied_parameters:
-                eos.error(' - {n}: {v}'.format(n=p.name(), v=p.evaluate()))
+                eos.error(f' - {p.name()}: {p.evaluate()}')
             return(-np.inf)
 
 
@@ -637,7 +637,7 @@ def _prior_transform(self, u):
         return self._u_to_par(u)
 
 
-    def sample_nested(self, bound='multi', nlive=250, dlogz=1.0, maxiter=None, seed=10, print_progress=True):
+    def sample_nested(self, bound='multi', nlive=250, dlogz=1.0, maxiter=None, seed=10, checkpoint_interval=60):
         """
         Return samples of the parameters.
 
@@ -653,14 +653,53 @@ def sample_nested(self, bound='multi', nlive=250, dlogz=1.0, maxiter=None, seed=
         :type maxiter: int, optional
         :param seed: The seed used to initialize the Mersenne Twister pseudo-random number generator.
         :type seed: {None, int, array_like[ints], SeedSequence}, optional
+        :param checkpoint_interval: The number of seconds between checkpoints at which the intermediate dynesty sampler results are stored. If None, then no intermediate results are stored. Defaults to 60.
+        :type checkpoint_interval: float, optional
 
         .. note::
            This method requires the dynesty python module, which can be installed from PyPI.
         """
         import dynesty
+        import itertools
+        from dynesty.dynamicsampler import stopping_function, weight_function
         sampler = dynesty.DynamicNestedSampler(self.log_likelihood, self._prior_transform, len(self.varied_parameters), bound=bound, nlive=nlive, rstate = np.random.Generator(np.random.MT19937(seed)))
-        sampler.run_nested(dlogz_init=dlogz, maxiter=maxiter, print_progress=print_progress)
-        return sampler.results
+
+        eos.info('Drawing initial samples.')
+        for results in sampler.sample_initial(dlogz=dlogz, maxiter=maxiter):
+            pass
+
+        eos.info('Starting nested sampling loop.')
+        if checkpoint_interval is not None:
+            timer = dynesty.utils.DelayTimer(checkpoint_interval)
+            time=0
+
+        for iter in itertools.count():
+            if checkpoint_interval is not None and timer.is_time():
+                time+=60
+                eos.info(f'Sampler results stored after {time} seconds.')
+                yield sampler.results
+
+            stop = stopping_function(sampler.results)
+
+            if maxiter is not None and iter > maxiter:
+                stop = True
+                eos.info("The sampling was stopped because the maximum number of iterations was reached.")
+
+            delta_logz = np.logaddexp(0, np.max(sampler.live_logl) + sampler.results.logvol - sampler.results.logz)[-1]
+            if delta_logz < dlogz:
+                stop = True
+                eos.info("The sampling was stopped because the limit of the remaining evidence was reached.")
+
+            if not stop:
+                logl_bounds = weight_function(sampler.results)  # derive bounds
+                for results in sampler.sample_batch(logl_bounds=logl_bounds, dlogz=dlogz, maxiter=maxiter):
+                    pass
+                sampler.combine_runs()  # add new samples to previous results
+            else:
+                break
+
+        yield sampler.results
+        return
 
 
     def _repr_html_(self):

python/eos/analysis_TEST.py

@@ -165,7 +165,9 @@ def test_multivariate_priors_1(self):
         # Sample from the prior
         results_list = []
         for i in range(0, 2):
-            results_list.append(analysis.sample_nested(bound='multi', nlive=250, dlogz=0.01, seed=10 + i, print_progress=False))
+            for results in analysis.sample_nested(bound='multi', nlive=250, dlogz=0.01, seed=10 + i):
+                last_result = results
+            results_list.append(last_result)
         results  = dynesty.utils.merge_runs(results_list)
 
         # Test samples against constraint
@@ -231,7 +233,9 @@ def test_multivariate_priors_2(self):
         # Sample from the prior
         results_list = []
         for i in range(0, 2):
-            results_list.append(analysis.sample_nested(bound='multi', nlive=250, dlogz=0.01, seed=10 + i, print_progress=False))
+            for results in analysis.sample_nested(bound='multi', nlive=250, dlogz=0.01, seed=10 + i):
+                last_result = results
+            results_list.append(last_result)
         results  = dynesty.utils.merge_runs(results_list)
 
         # Test samples against analytic results:

python/eos/tasks.py

@@ -240,9 +240,9 @@ def sample_mcmc(analysis_file:str, posterior:str, chain:int, base_directory:str=
         samples, usamples, weights = analysis.sample(N=N, stride=stride, pre_N=pre_N, preruns=preruns, rng=rng, cov_scale=cov_scale, start_point=start_point, return_uspace=True)
         eos.data.MarkovChain.create(os.path.join(base_directory, posterior, f'mcmc-{chain:04}'), analysis.varied_parameters, samples, usamples, weights)
     except RuntimeError as e:
-        eos.error('encountered run time error ({e}) in parameter point:'.format(e=e))
+        eos.error(f'encountered run time error ({e}) in parameter point:')
         for p in analysis.varied_parameters:
-            eos.error(' - {n}: {v}'.format(n=p.name(), v=p.evaluate()))
+            eos.error(f' - {p.name()}: {p.evaluate()}')
 
 
 @task('find-clusters', '{posterior}/clusters')
@@ -274,7 +274,7 @@ def find_clusters(posterior:str, base_directory:str='./', threshold:float=2.0, K
     groups = pypmc.mix_adapt.r_value.r_group([_np.mean(chain.T, axis=1) for chain in chains],
                            [_np.var (chain.T, axis=1, ddof=1) for chain in chains],
                            n, threshold)
-    eos.info('Found {} groups using an R value threshold of {}'.format(len(groups), threshold))
+    eos.info(f'Found {len(groups)} groups using an R value threshold of {threshold}')
     density   = pypmc.mix_adapt.r_value.make_r_gaussmix(chains, K_g=K_g, critical_r=threshold)
     eos.info(f'Created mixture density with {len(density.components)} components')
     eos.data.MixtureDensity.create(os.path.join(base_directory, posterior, 'clusters'), density)
@@ -359,7 +359,7 @@ def sample_pmc(analysis_file:str, posterior:str, base_directory:str='./', step_N
     elif initial_proposal == 'product':
         initial_density = eos.data.MixtureDensity(os.path.join(base_directory, posterior, 'product')).density()
     else:
-        eos.error("Could not initialize proposal in sample_pmc: argument {} is not supported.".format(initial_proposal))
+        eos.error(f"Could not initialize proposal in sample_pmc: argument {initial_proposal} is not supported.")
 
     samples, weights, posterior_values, proposal = analysis.sample_pmc(initial_density, step_N=step_N, steps=steps, final_N=final_N,
                                                      rng=rng, final_perplexity_threshold=perplexity_threshold,
@@ -420,11 +420,11 @@ def predict_observables(analysis_file:str, posterior:str, prediction:str, base_d
             cache.update()
             observable_samples.append([cache[id] for id in observable_ids])
         except RuntimeError as e:
-            eos.error('skipping prediction for sample {i} due to runtime error ({e}): {s}'.format(i=i, e=e, s=sample))
+            eos.error(f'skipping prediction for sample {i} due to runtime error ({e}): {sample}')
             observable_samples.append([_np.nan for _ in observable_ids])
     observable_samples = _np.array(observable_samples)
 
-    output_path = os.path.join(base_directory, posterior, 'pred-{}'.format(prediction))
+    output_path = os.path.join(base_directory, posterior, f'pred-{prediction}')
     eos.data.Prediction.create(output_path, observables, observable_samples, data.weights[begin:end])
 
 
@@ -484,12 +484,12 @@ def sample_nested(analysis_file:str, posterior:str, base_directory:str='./', bou
     :type seed: int, optional
     """
     analysis = analysis_file.analysis(posterior)
-    results = analysis.sample_nested(bound=bound, nlive=nlive, dlogz=dlogz, maxiter=maxiter, seed=seed)
-    samples = results.samples
-    posterior_values = results.logwt - results.logz[-1]
-    weights = _np.exp(posterior_values)
-    eos.data.DynestyResults.create(os.path.join(base_directory, posterior, 'dynesty_results'), analysis.varied_parameters, results)
-    eos.data.ImportanceSamples.create(os.path.join(base_directory, posterior, 'samples'), analysis.varied_parameters,
+    for iter, results in enumerate(analysis.sample_nested(bound=bound, nlive=nlive, dlogz=dlogz, maxiter=maxiter, seed=seed)):
+        samples = results.samples
+        posterior_values = results.logwt - results.logz[-1]
+        weights = _np.exp(posterior_values)
+        eos.data.DynestyResults.create(os.path.join(base_directory, posterior, f'dynesty_results-{iter:04}'), analysis.varied_parameters, results)
+        eos.data.ImportanceSamples.create(os.path.join(base_directory, posterior, f'samples'), analysis.varied_parameters,
                                       samples, weights, posterior_values=posterior_values)