From 476d723efc65ad0f626deba5b1f87b7aa7d207e3 Mon Sep 17 00:00:00 2001
From: Anton Soloviev <anton@praviel.com>
Date: Thu, 15 Oct 2026 23:08:52 -0700
Subject: [PATCH] Reuse reserved memory when resuming an early-stopped chain

---
 src/emcee/backends/backend.py         | 17 ++++++----
 src/emcee/tests/unit/test_backends.py | 49 +++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/src/emcee/backends/backend.py b/src/emcee/backends/backend.py
index bc594479..f12f3777 100644
--- a/src/emcee/backends/backend.py
+++ b/src/emcee/backends/backend.py
@@ -172,16 +172,19 @@ def grow(self, ngrow, blobs):
         """
         self._check_blobs(blobs)
         i = ngrow - (len(self.chain) - self.iteration)
-        a = np.empty((i, self.nwalkers, self.ndim), dtype=self.dtype)
-        self.chain = np.concatenate((self.chain, a), axis=0)
-        a = np.empty((i, self.nwalkers), dtype=self.dtype)
-        self.log_prob = np.concatenate((self.log_prob, a), axis=0)
+        if i > 0:
+            a = np.empty((i, self.nwalkers, self.ndim), dtype=self.dtype)
+            self.chain = np.concatenate((self.chain, a), axis=0)
+            a = np.empty((i, self.nwalkers), dtype=self.dtype)
+            self.log_prob = np.concatenate((self.log_prob, a), axis=0)
         if blobs is not None:
             dt = np.dtype((blobs.dtype, blobs.shape[1:]))
-            a = np.empty((i, self.nwalkers), dtype=dt)
             if self.blobs is None:
-                self.blobs = a
-            else:
+                self.blobs = np.empty(
+                    (len(self.chain), self.nwalkers), dtype=dt
+                )
+            elif i > 0:
+                a = np.empty((i, self.nwalkers), dtype=dt)
                 self.blobs = np.concatenate((self.blobs, a), axis=0)
 
     def _check(self, state, accepted):
diff --git a/src/emcee/tests/unit/test_backends.py b/src/emcee/tests/unit/test_backends.py
index a55f2c58..e86603c6 100644
--- a/src/emcee/tests/unit/test_backends.py
+++ b/src/emcee/tests/unit/test_backends.py
@@ -274,3 +274,52 @@ def test_hdf5_compression():
         b.get_blobs()
         b.get_log_prob()
         b.accepted
+
+
+@pytest.mark.parametrize("backend", all_backends)
+@pytest.mark.parametrize("dtype,blobs", product(dtypes, [True, False]))
+@pytest.mark.parametrize("thin_by", [1, 3])
+def test_resume_after_early_stop(backend, dtype, blobs, thin_by):
+    np.random.seed(1234)
+    coords = np.random.randn(16, 2)
+    random_state = np.random.get_state()
+    log_prob = normal_log_prob_blobs if blobs else normal_log_prob
+    reference = EnsembleSampler(16, 2, log_prob, blobs_dtype=dtype)
+    reference.random_state = random_state
+    reference.run_mcmc(coords, 48, thin_by=thin_by)
+
+    with backend() as be:
+        sampler = EnsembleSampler(
+            16, 2, log_prob, blobs_dtype=dtype, backend=be
+        )
+        sampler.random_state = random_state
+        generator = sampler.sample(coords, iterations=40, thin_by=thin_by)
+        for i, state in enumerate(generator):
+            if i == 4:
+                break
+        generator.close()
+        # Reuse some reserved storage, then grow past its original size.
+        sampler.run_mcmc(state, 3, thin_by=thin_by)
+        sampler.run_mcmc(None, 40, thin_by=thin_by)
+        np.testing.assert_array_equal(
+            sampler.get_chain(), reference.get_chain()
+        )
+        np.testing.assert_array_equal(
+            sampler.get_log_prob(), reference.get_log_prob()
+        )
+        np.testing.assert_array_equal(be.accepted, reference.backend.accepted)
+        if blobs:
+            np.testing.assert_array_equal(
+                sampler.get_blobs(), reference.get_blobs()
+            )
+
+
+def test_grow_blobs_in_reserved_storage():
+    backend = backends.Backend()
+    backend.reset(8, 2)
+    backend.grow(20, None)
+    blobs = np.arange(8)
+    backend.grow(5, blobs)
+    state = State(np.ones((8, 2)), log_prob=np.zeros(8), blobs=blobs)
+    backend.save_step(state, np.ones(8, dtype=bool))
+    np.testing.assert_array_equal(backend.get_blobs(), blobs[None, :])