Merge pull request #191 from Joshuaalbert/develop

Develop

Author: Joshuaalbert

.github/workflows/unittests.yml

@@ -39,7 +39,4 @@ jobs:
           flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
       - name: Test with pytest
         run: |
-          set -o allexport
-          source deployment/local.env
-          set +o allexport
           pytest

README.md

@@ -144,9 +144,9 @@ Given a probabilistic model, JAXNS can perform nested sampling on it. This allow
 posterior samples.
 
 ```python
-from jaxns import DefaultNestedSampler
+from jaxns import NestedSampler
 
-ns = DefaultNestedSampler(model=model, max_samples=1e5)
+ns = NestedSampler(model=model, max_samples=1e5)
 
 # Run the sampler
 termination_reason, state = ns(jax.random.PRNGKey(42))
@@ -290,7 +290,7 @@ Sampling paper](https://arxiv.org/abs/2312.11330).
 
 ```bash
 # To create a new env, if necessary
-conda create -n jaxns_py python=3.11
+conda create -n jaxns_py python=3.12
 conda activate jaxns_py
 ```
 
@@ -328,9 +328,8 @@ Checkout the examples [here](https://jaxns.readthedocs.io/en/latest/#).
 
 ## Caveats
 
-The caveat is that you need to be able to define your likelihood function with JAX. This is usually no big deal because
-JAX is just a replacement for NumPy and many likelihoods can be expressed such.
-If you're unfamiliar, take a quick tour of JAX (https://jax.readthedocs.io/en/latest/notebooks/quickstart.html).
+The caveat is that you need to be able to define your likelihood function with JAX. UPDATE: now you can just
+use the `@jaxify_likelihood` decorator to run with arbitrary pythonic likelihoods.
 
 # Speed test comparison with other nested sampling packages
 
@@ -339,30 +338,20 @@ JAXNS is much faster than PolyChord, MultiNEST, and dynesty, typically achieving
 improvement in run time, for models with cheap likelihood evaluations.
 This is shown in (https://arxiv.org/abs/2012.15286).
 
-Recently JAXNS has implemented Phantom-Powered Nested Sampling, which significantly reduces the number of required
-likelihood evaluations for inferring the posterior. This is shown in (https://arxiv.org/abs/2312.11330).
+Recently JAXNS has implemented Phantom-Powered Nested Sampling, which helps for parameter inference. This is shown
+in (https://arxiv.org/abs/2312.11330).
 
-# Note on performance with parallelisation
+# Note on performance with parallelisation and GPUS
 
-__Note, that this is an experimental feature.__
-
-If you set `num_parallel_workers > 1` you will use `jax.pmap` under the hood for parallelisation.
-This is a very powerful feature, but it is important to understand how it works.
-It runs identical copies of the nested sampling algorithm on multiple devices.
-There is a two-part stopping condition.
-First, each copy goes until the user defined stopping condition is met __per device__.
-Then, it performs an all-gather and finds at the highest likelihood contour among all copies, and continues all copies
-hit this likelihood contour.
-This ensures consistency of depth across all copies.
-We then merge the copies and compute the final results.
-
-The algorithm is fairly memory bound, so running parallelisation over multiple CPUs on the same machine may not yield
-the expected speed up, and depends on how expensive the likelihood evaluations are. Running over separate physical
-devices
-is the best way to achieve speed up.
+To use parallel computing, you can simply pass `devices` to the `NestedSampler` constructor. This will distributed
+sampling over the devices. To use GPUs you can pass `jax.devices('gpu')` to the `devices` argument. You can also se all
+your CPUs by placing `os.environ["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={os.cpu_count()}"`
+before importing JAXNS.
 
 # Change Log
 
+24 Sep, 2024 -- JAXNS 2.6.1 released. Sharded parallel JAXNS. Rewrite of internals to support sharded parallelisation.
+
 20 Aug, 2024 -- JAXNS 2.6.0 released. Removed haiku dependency. Implemented our own
 context. `jaxns.framework.context.convert_external_params` enables interfacing with any external NN libary.
 

benchmarks/difficult_problems/main.py

@@ -5,7 +5,7 @@
 import jax.numpy as jnp
 import tensorflow_probability.substrates.jax as tfp
 
-from jaxns import Model, Prior, DefaultNestedSampler
+from jaxns import Model, Prior, NestedSampler
 
 tfpd = tfp.distributions
 
@@ -152,7 +152,7 @@ def main():
     for model_name, model in all_models().items():
         print(f"Testing model {model_name}")
         model.sanity_check(jax.random.PRNGKey(0), 1000)
-        ns = DefaultNestedSampler(model=model,
+        ns = NestedSampler(model=model,
                                   max_samples=1000000,
                                   verbose=True,
                                   difficult_model=True,

benchmarks/gh117/main.py

@@ -4,7 +4,7 @@
 import tensorflow_probability.substrates.jax as tfp
 from jax import random
 
-from jaxns import Model, Prior, DefaultNestedSampler
+from jaxns import Model, Prior, NestedSampler
 
 
 tfpd = tfp.distributions
@@ -23,7 +23,7 @@ def prior_model():
 
 
     # Create the nested sampler class. In this case without any tuning.
-    exact_ns = DefaultNestedSampler(model=model, max_samples=max_samples)
+    exact_ns = NestedSampler(model=model, max_samples=max_samples)
 
     termination_reason, state = exact_ns(random.PRNGKey(42))
     return termination_reason

benchmarks/gh168/main.py

@@ -7,7 +7,7 @@
 import tensorflow_probability.substrates.jax as tfp
 from jax._src.scipy.linalg import solve_triangular
 
-from jaxns import Model, Prior, DefaultNestedSampler
+from jaxns import Model, Prior, NestedSampler
 
 tfpd = tfp.distributions
 
@@ -47,7 +47,7 @@ def log_likelihood(x):
 
     model = Model(prior_model=prior_model, log_likelihood=log_likelihood)
 
-    ns = DefaultNestedSampler(model=model, max_samples=100000, verbose=False)
+    ns = NestedSampler(model=model, max_samples=100000, verbose=False)
 
     termination_reason, state = ns(key)
     results = ns.to_results(termination_reason=termination_reason, state=state, trim=False)

benchmarks/parallel_problems/main.py

@@ -11,7 +11,7 @@
 import tensorflow_probability.substrates.jax as tfp
 from jax._src.scipy.linalg import solve_triangular
 
-from jaxns import Model, Prior, DefaultNestedSampler
+from jaxns import Model, Prior, NestedSampler, jaxify_likelihood
 
 tfpd = tfp.distributions
 
@@ -51,7 +51,7 @@ def log_likelihood(x):
 
     model = Model(prior_model=prior_model, log_likelihood=log_likelihood)
 
-    ns = DefaultNestedSampler(model=model, max_samples=100000, verbose=False, num_parallel_workers=len(jax.devices()))
+    ns = NestedSampler(model=model, max_samples=100000, verbose=False)
 
     termination_reason, state = ns(key)
     results = ns.to_results(termination_reason=termination_reason, state=state, trim=False)
@@ -61,7 +61,7 @@ def log_likelihood(x):
 def main():
     num_devices = len(jax.devices())
     jaxns_version = pkg_resources.get_distribution("jaxns").version
-    m = 1
+    m = 3
     run_model_aot = jax.jit(run_model).lower(jax.random.PRNGKey(0)).compile()
     dt = []
 
@@ -70,14 +70,13 @@ def main():
 
     for i in range(m):
         t0 = time.time()
-        log_Z_error, log_Z_uncert = run_model_aot(jax.random.PRNGKey(i))
-        log_Z_error.block_until_ready()
+        log_Z_error, log_Z_uncert = jax.block_until_ready(run_model_aot(jax.random.PRNGKey(i)))
         t1 = time.time()
         dt.append(t1 - t0)
         errors.append(log_Z_error)
         uncerts.append(log_Z_uncert)
     total_time = sum(dt)
-    best_3 = sum(sorted(dt)[:min(3, m)]) / 3.
+    best_3 = sum(sorted(dt)[:3]) / 3.
     # print(f"Errors: {errors}")
     # print(f"Uncerts: {uncerts}")
     print(f"JAXNS {jaxns_version}\n"

benchmarks/parallel_problems/results

@@ -1,2 +1,5 @@
 2.5.1,12,-0.8741127252578735,0.0003452669479884207,95.40295028686523,77.15287351608276
 2.5.1,12,-1.30291748046875,0.32637646794319153,102.9092493057251,34.30308310190836
+2.3.4,12,0.09217196161439745,0.3709631743726239,2.116373300552368,0.7054577668507894
+2.6.0,6,-0.06817162536483465,0.3668122147028353,1.1894174337387085,1.0906640688578289
+2.6.0,12,-0.06817162536483465,0.3668122147028353,1.9961725234985352,1.7415425777435303

deployment/Dockerfile

@@ -12,7 +12,7 @@
 project = "jaxns"
 copyright = "2022, Joshua G. Albert"
 author = "Joshua G. Albert"
-release = "2.6.0"
+release = "2.6.1"
 
 # -- General configuration ---------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration