Add multiprocessing (#15)

* Add joblib to requirements

* Add multiprocessing

* Add tests to clean recipe

* Rename sampling functions

* Fix tests

* Change docs to imperative mood

* Loosen tests

Author: eigenfoo

Makefile

@@ -80,6 +80,6 @@ check: lint test  # Both lint and test code. Runs `make lint` followed by `make
 .PHONY: clean
 clean:  # Clean project directories.
 	rm -rf dist/ site/ littlemcmc.egg-info/ pip-wheel-metadata/ __pycache__/ testing-report.html
-	find littlemcmc/ -type d -name "__pycache__" -exec rm -rf {} +
-	find littlemcmc/ -type d -name "__pycache__" -delete
-	find littlemcmc/ -type f -name "*.pyc" -delete
+	find littlemcmc/ tests/ -type d -name "__pycache__" -exec rm -rf {} +
+	find littlemcmc/ tests/ -type d -name "__pycache__" -delete
+	find littlemcmc/ tests/ -type f -name "*.pyc" -delete

littlemcmc/sampling.py

@@ -14,11 +14,29 @@
 
 """Sampling driver functions (unrelated to PyMC3's `sampling.py`)."""
 
+import os
+from collections.abc import Iterable
+import logging
+from joblib import Parallel, delayed
 import numpy as np
 
+_log = logging.getLogger("littlemcmc")
+
+
+def _sample_one_chain(
+    logp_dlogp_func,
+    size,
+    stepper,
+    draws,
+    tune,
+    init=None,
+    random_seed=None,
+    discard_tuned_samples=True,
+):
+    """Sample one chain in one process."""
+    if random_seed is not None:
+        np.random.seed(random_seed)
 
-def sample(logp_dlogp_func, size, stepper, draws, tune, init=None):
-    """Sample."""
     if init is not None:
         q = init
     else:
@@ -34,4 +52,58 @@ def sample(logp_dlogp_func, size, stepper, draws, tune, init=None):
         if i == tune:
             stepper.stop_tuning()
 
+    if discard_tuned_samples:
+        trace = trace[:, tune:]
+
     return trace, stats
+
+
+def sample(
+    logp_dlogp_func,
+    size,
+    stepper,
+    draws,
+    tune,
+    chains=None,
+    cores=None,
+    init=None,
+    random_seed=None,
+    discard_tuned_samples=True,
+):
+    """Sample."""
+    if cores is None:
+        cores = min(4, os.cpu_count())
+    if chains is None:
+        chains = max(2, cores)
+
+    if random_seed is None or isinstance(random_seed, int):
+        if random_seed is not None:
+            np.random.seed(random_seed)
+        random_seed = [np.random.randint(2 ** 30) for _ in range(chains)]
+    elif isinstance(random_seed, Iterable) and len(random_seed) != chains:
+        random_seed = random_seed[:chains]
+    elif not isinstance(random_seed, Iterable):
+        raise TypeError("Invalid value for `random_seed`. Must be tuple, list or int")
+
+    # Small trace warning
+    if draws == 0:
+        msg = "Tuning was enabled throughout the whole trace."
+        _log.warning(msg)
+    elif draws < 500:
+        msg = "Only {} samples in chain.".format(draws)
+        _log.warning(msg)
+
+    _log.info("Multiprocess sampling ({} chains in {} jobs)".format(chains, cores))
+    asdf = Parallel(n_jobs=cores)(
+        delayed(sample_one_chain)(
+            logp_dlogp_func=logp_dlogp_func,
+            size=size,
+            stepper=stepper,
+            draws=draws,
+            tune=tune,
+            init=init,
+            random_seed=i,
+            discard_tuned_samples=discard_tuned_samples,
+        )
+        for i in random_seed
+    )

requirements.txt

@@ -1,3 +1,4 @@
+joblib
 numpy
 scipy>=0.18.1
 theano

tests/test_sampling.py

@@ -15,6 +15,7 @@
 import numpy as np
 import scipy.stats
 import littlemcmc as lmc
+from littlemcmc.sampling import _sample_one_chain
 
 
 def logp_func(x, loc=0, scale=1):
@@ -34,38 +35,34 @@ def test_hmc_sampling_runs():
     stepper = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, size=size)
     draws = 1
     tune = 1
-    init = None
-    trace, stats = lmc.sample(logp_dlogp_func, size, stepper, draws, tune, init)
+    trace, stats = _sample_one_chain(logp_dlogp_func, size, stepper, draws, tune)
 
 
 def test_nuts_sampling_runs():
     size = 1
     stepper = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=size)
     draws = 1
     tune = 1
-    init = None
-    trace, stats = lmc.sample(logp_dlogp_func, size, stepper, draws, tune, init)
+    trace, stats = _sample_one_chain(logp_dlogp_func, size, stepper, draws, tune)
 
 
 def test_hmc_recovers_1d_normal():
     size = 1
     stepper = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, size=size)
     draws = 1000
     tune = 1000
-    init = None
-    trace, stats = lmc.sample(logp_dlogp_func, size, stepper, draws, tune, init)
+    trace, stats = _sample_one_chain(logp_dlogp_func, size, stepper, draws, tune)
 
-    assert np.allclose(np.mean(trace[:, 1000:]), 0, atol=0.1)
-    assert np.allclose(np.std(trace[:, 1000:]), 1, atol=0.1)
+    assert np.allclose(np.mean(trace), 0, atol=1)
+    assert np.allclose(np.std(trace), 1, atol=1)
 
 
 def test_nuts_recovers_1d_normal():
     size = 1
     stepper = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, size=size)
     draws = 1000
     tune = 1000
-    init = None
-    trace, stats = lmc.sample(logp_dlogp_func, size, stepper, draws, tune, init)
+    trace, stats = _sample_one_chain(logp_dlogp_func, size, stepper, draws, tune)
 
-    assert np.allclose(np.mean(trace[:, 1000:]), 0, atol=0.1)
-    assert np.allclose(np.std(trace[:, 1000:]), 1, atol=0.1)
+    assert np.allclose(np.mean(trace), 0, atol=1)
+    assert np.allclose(np.std(trace), 1, atol=1)