Simplify mcmc module

.github/workflows/tests.yml

@@ -42,6 +42,7 @@ jobs:
             pymc/tests/distributions/test_logprob.py
             pymc/tests/test_aesaraf.py
             pymc/tests/test_math.py
+            pymc/tests/backends/test_base.py
             pymc/tests/backends/test_ndarray.py
             pymc/tests/step_methods/hmc/test_hmc.py
             pymc/tests/test_func_utils.py
@@ -60,6 +61,7 @@ jobs:
             pymc/tests/distributions/test_simulator.py
             pymc/tests/distributions/test_truncated.py
             pymc/tests/sampling/test_forward.py
+            pymc/tests/sampling/test_population.py
             pymc/tests/stats/test_convergence.py
 
           - |

pymc/backends/__init__.py

@@ -60,7 +60,34 @@
 Saved backends can be loaded using `arviz.from_netcdf`
 
 """
+from copy import copy
+from typing import Dict, List, Optional
+
 from pymc.backends.arviz import predictions_to_inference_data, to_inference_data
+from pymc.backends.base import BaseTrace
 from pymc.backends.ndarray import NDArray, point_list_to_multitrace
 
 __all__ = ["to_inference_data", "predictions_to_inference_data"]
+
+
+def _init_trace(
+    *,
+    expected_length: int,
+    chain_number: int,
+    stats_dtypes: List[Dict[str, type]],
+    trace: Optional[BaseTrace],
+    model,
+) -> BaseTrace:
+    """Initializes a trace backend for a chain."""
+    strace: BaseTrace
+    if trace is None:
+        strace = NDArray(model=model)
+    elif isinstance(trace, BaseTrace):
+        if len(trace) > 0:
+            raise ValueError("Continuation of traces is no longer supported.")
+        strace = copy(trace)
+    else:
+        raise NotImplementedError(f"Unsupported `trace`: {trace}")
+
+    strace.setup(expected_length, chain_number, stats_dtypes)
+    return strace

pymc/backends/base.py

@@ -21,6 +21,7 @@
 import warnings
 
 from abc import ABC
+from typing import List, Sequence, Tuple, cast
 
 import aesara.tensor as at
 import numpy as np
@@ -561,3 +562,30 @@ def _squeeze_cat(results, combine, squeeze):
         if squeeze and len(results) == 1:
             results = results[0]
     return results
+
+
+def _choose_chains(traces: Sequence[BaseTrace], tune: int) -> Tuple[List[BaseTrace], int]:
+    """
+    Filter and slice traces such that (n_traces * len(shortest_trace)) is maximized.
+
+    We get here after a ``KeyboardInterrupt``, and so the different
+    traces have different lengths. We therefore pick the number of
+    traces such that (number of traces) * (length of shortest trace)
+    is maximised.
+    """
+    if not traces:
+        raise ValueError("No traces to slice.")
+
+    lengths = [max(0, len(trace) - tune) for trace in traces]
+    if not sum(lengths):
+        raise ValueError("Not enough samples to build a trace.")
+
+    idxs = np.argsort(lengths)
+    l_sort = np.array(lengths)[idxs]
+
+    use_until = cast(int, np.argmax(l_sort * np.arange(1, l_sort.shape[0] + 1)[::-1]))
+    final_length = l_sort[use_until]
+
+    take_idx = cast(Sequence[int], idxs[use_until:])
+    sliced_traces = [traces[idx] for idx in take_idx]
+    return sliced_traces, final_length + tune