Merge pull request #116 from Joshuaalbert/develop

Develop
Joshuaalbert · Dec 20, 2023 · 05bf627 · 05bf627
2 parents 6bc23a5 + 21fa6b9
commit 05bf627
Show file tree

Hide file tree

Showing 7 changed files with 427 additions and 327 deletions.
diff --git a/README.md b/README.md
@@ -108,6 +108,8 @@ is the best way to achieve speed up.
 
 # Change Log
 
+21 Dec, 2023 -- JAXNS 2.3.4 released. Correction for ESS and logZ uncert. `parameter_estimation` mode.
+
 20 Dec, 2023 -- JAXNS 2.3.2/3 released. Improved default parameters. `difficult_model` mode. Improve plotting. 
 
 18 Dec, 2023 -- JAXNS 2.3.1 released. Paper open science release. Default parameters from paper.

diff --git a/docs/examples/gaussian_shells.ipynb b/docs/examples/gaussian_shells.ipynb
diff --git a/docs/examples/mvn_data_mvn_prior.ipynb b/docs/examples/mvn_data_mvn_prior.ipynb
diff --git a/docs/papers/phantom-powered-nested-sampling/phantom_bias_tradeoff.ipynb b/docs/papers/phantom-powered-nested-sampling/phantom_bias_tradeoff.ipynb
diff --git a/jaxns/nested_sampler/standard_static.py b/jaxns/nested_sampler/standard_static.py
@@ -561,15 +561,23 @@ def _to_results(self, termination_reason: IntArray, state: StaticStandardNestedS
         )
         log_Z_uncert = jnp.sqrt(log_Z_var)
 
+        # Correction by sqrt(k+1)
+        total_phantom_samples = jnp.sum(sample_collection.phantom.astype(int_type))
+        phantom_fraction = total_phantom_samples / num_samples  # k / (k+1)
+        k = phantom_fraction / (1. - phantom_fraction)
+        log_Z_uncert = log_Z_uncert * jnp.sqrt(1. + k)
+
         # Kish's ESS = [sum dZ]^2 / [sum dZ^2]
         ESS = effective_sample_size(final_evidence_stats.log_Z_mean, final_evidence_stats.log_dZ2_mean)
+        ESS = ESS / (1. + k)
 
         samples = vmap(self.model.transform)(U_samples)
 
         log_L_samples = log_L
         dp_mean = LogSpace(per_sample_evidence_stats.log_dZ_mean)
         dp_mean = normalise_log_space(dp_mean)
-        H_mean_instable = -((dp_mean * LogSpace(jnp.log(jnp.abs(log_L_samples)), jnp.sign(log_L_samples))).sum().value - log_Z_mean)
+        H_mean_instable = -((dp_mean * LogSpace(jnp.log(jnp.abs(log_L_samples)),
+                                                jnp.sign(log_L_samples))).sum().value - log_Z_mean)
         # H \approx E[-log(compression)] = E[-log(X)] (More stable than E[log(L) - log(Z)]
         H_mean_stable = -((dp_mean * LogSpace(jnp.log(-per_sample_evidence_stats.log_X_mean))).sum().value)
         H_mean = jnp.where(jnp.isfinite(H_mean_instable), H_mean_instable, H_mean_stable)
@@ -582,8 +590,6 @@ def _to_results(self, termination_reason: IntArray, state: StaticStandardNestedS
         log_posterior_density = log_L + vmap(self.model.log_prob_prior)(
             U_samples)
 
-        total_phantom_samples = jnp.sum(sample_collection.phantom.astype(int_type))
-
         return NestedSamplerResults(
             log_Z_mean=log_Z_mean,  # estimate of log(E[Z])
             log_Z_uncert=log_Z_uncert,  # estimate of log(StdDev[Z])

diff --git a/jaxns/public.py b/jaxns/public.py
@@ -35,7 +35,8 @@ class DefaultNestedSampler:
     def __init__(self, model: BaseAbstractModel, max_samples: Union[int, float], num_live_points: Optional[int] = None,
                  s: Optional[int] = None, k: Optional[int] = None, c: Optional[int] = None,
                  num_parallel_workers: int = 1,
-                 difficult_model: bool = False):
+                 difficult_model: bool = False,
+                 parameter_estimation: bool = False):
         """
         Initialises the nested sampler.
 
@@ -50,14 +51,15 @@ def __init__(self, model: BaseAbstractModel, max_samples: Union[int, float], num
             c: number of parallel Markov-chains to use. Defaults to 20 * D.
             num_parallel_workers: number of parallel workers to use. Defaults to 1. Experimental feature.
             difficult_model: if True, uses more robust default settings. Defaults to False.
+            parameter_estimation: if True, uses more robust default settings for parameter estimation. Defaults to False.
         """
         if difficult_model:
             self._s = 10 if s is None else int(s)
         else:
-            self._s = 4 if s is None else int(s)
+            self._s = 5 if s is None else int(s)
         if self._s <= 0:
             raise ValueError(f"Expected s > 0, got s={self._s}")
-        if difficult_model:
+        if parameter_estimation:
             self._k = model.U_ndims if k is None else int(k)
         else:
             self._k = model.U_ndims // 2 if k is None else int(k)
@@ -70,7 +72,7 @@ def __init__(self, model: BaseAbstractModel, max_samples: Union[int, float], num
             if difficult_model:
                 self._c = 50 * model.U_ndims if c is None else int(c)
             else:
-                self._c = 20 * model.U_ndims if c is None else int(c)
+                self._c = 30 * model.U_ndims if c is None else int(c)
         if self._c <= 0:
             raise ValueError(f"Expected c > 0, got c={self._c}")
         # Sanity check for max_samples (should be able to at least do one shrinkage)

diff --git a/setup.py b/setup.py
@@ -21,7 +21,7 @@
     long_description = fh.read()
 
 setup(name='jaxns',
-      version='2.3.3',
+      version='2.3.4',
       description='Nested Sampling in JAX',
       long_description=long_description,
       long_description_content_type="text/markdown",