Repository: eigenfoo/littlemcmc

Issue:
Use a better shape for trace object
Currently we flatten the `trace` object to have shape `[num_variables, num_chains * num_samples]`. Really, it ought to be `[num_chains, num_samples, num_variables]` (or whatever reasonable choice ArviZ makes: we should target compatibility with them.

https://github.com/eigenfoo/littlemcmc/blob/master/littlemcmc/sampling.py#L213-L214

Comments:

Patch:
diff --git a/littlemcmc/sampling.py b/littlemcmc/sampling.py
index 50a3c96..74e22fd 100644
--- a/littlemcmc/sampling.py
+++ b/littlemcmc/sampling.py
@@ -210,13 +210,18 @@ def sample(
         for i, seed in enumerate(random_seed)
     )
 
-    # Flatten `trace` to have shape [num_variables, num_chains * num_samples]
-    trace = np.hstack([np.atleast_2d(chain_trace) for (chain_trace, _) in results])
+    # Reshape `trace` to have shape [num_chains, num_samples, num_variables]
+    trace = np.array([np.atleast_2d(chain_trace).T for (chain_trace, _) in results])
 
-    # Reshape `stats` to a dictionary
-    stats_ = [iter_stats for (_, chain_stats) in results for iter_stats in chain_stats]
+    # Reshape `stats` to a dictionary with keys = string of sampling stat name,
+    # values = np.array with shape [num_chains, num_samples, num_variables]
     stats = {
-        name: np.squeeze(np.array([iter_stats[name] for iter_stats in stats_])).astype(dtype)
+        name: np.array(
+            [
+                [np.atleast_1d(iter_stats[name]) for iter_stats in chain_stats]
+                for (_, chain_stats) in results
+            ]
+        ).astype(dtype)
         for (name, dtype) in step.stats_dtypes[0].items()
     }
 


Test patch:
diff --git a/tests/test_sampling.py b/tests/test_sampling.py
index 19fca98..7b1d71e 100644
--- a/tests/test_sampling.py
+++ b/tests/test_sampling.py
@@ -31,27 +31,47 @@ def test_init_nuts(method):
 def test_hmc_sampling_runs():
     model_ndim = 1
     step = lmc.HamiltonianMC(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
-    draws = 2
+    draws = 3
     tune = 1
-    chains = 1
+    chains = 2
     cores = 1
+
+    expected_shape = (2, 3, 1)
+
     trace, stats = lmc.sample(
         logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
-    assert trace.shape == (1, 2)
+    assert trace.shape == expected_shape
+    assert all([stats[name].shape == expected_shape for (name, _) in step.stats_dtypes[0].items()])
+    assert all(
+        [
+            stats[name].dtype == expected_dtype
+            for (name, expected_dtype) in step.stats_dtypes[0].items()
+        ]
+    )
 
 
 def test_nuts_sampling_runs():
     model_ndim = 1
     step = lmc.NUTS(logp_dlogp_func=logp_dlogp_func, model_ndim=model_ndim)
-    draws = 2
+    draws = 3
     tune = 1
-    chains = 1
+    chains = 2
     cores = 1
+
+    expected_shape = (2, 3, 1)
+
     trace, stats = lmc.sample(
         logp_dlogp_func, model_ndim, draws, tune, step=step, chains=chains, cores=cores
     )
-    assert trace.shape == (1, 2)
+    assert trace.shape == (2, 3, 1)
+    assert all([stats[name].shape == expected_shape for (name, _) in step.stats_dtypes[0].items()])
+    assert all(
+        [
+            stats[name].dtype == expected_dtype
+            for (name, expected_dtype) in step.stats_dtypes[0].items()
+        ]
+    )
 
 
 def test_multiprocess_sampling_runs():