Bayesian extend/resume by rozyczko · Pull Request #257 · easyscience/core

rozyczko · 2026-06-04T10:35:28Z

Support for resuming and extending Bayesian MCMC chains using the BUMPS DREAM sampler.
Added ability to save and reload sampler state and continue chains across sessions.

Resume and extend MCMC chains:
- Added resume_state parameter to mcmc_sample and sample methods in both fitter.py and minimizer_bumps.py, allowing users to continue sampling from a saved chain state. This includes documentation of the "ring-buffer contract" for extending chains without losing samples.
Population/chains argument resolution:
- Introduced _resolve_population_alias method to handle both chains and population arguments, raising an error if both are set with different values, improving API clarity and preventing silent bugs.

codecov · 2026-06-04T10:36:44Z

Codecov Report

❌ Patch coverage is 93.87755% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.30%. Comparing base (6efc17e) to head (013e834).

Files with missing lines	Patch %	Lines
src/easyscience/fitting/sampler.py	92.59%	10 Missing ⚠️
.../easyscience/fitting/minimizers/minimizer_bumps.py	96.36%	2 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop     #257      +/-   ##
===========================================
+ Coverage    82.92%   83.30%   +0.37%     
===========================================
  Files           62       63       +1     
  Lines         5084     5253     +169     
===========================================
+ Hits          4216     4376     +160     
- Misses         868      877       +9

Flag	Coverage Δ
integration	`45.45% <89.79%> (+1.51%)`	⬆️
unittests	`81.00% <45.40%> (-1.42%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/easyscience/fitting/__init__.py	`100.00% <100.00%> (ø)`
src/easyscience/fitting/fitter.py	`96.92% <ø> (-0.38%)`	⬇️
src/easyscience/fitting/minimizers/__init__.py	`100.00% <100.00%> (ø)`
src/easyscience/fitting/multi_fitter.py	`98.52% <ø> (-0.03%)`	⬇️
.../easyscience/fitting/minimizers/minimizer_bumps.py	`96.95% <96.36%> (+1.00%)`	⬆️
src/easyscience/fitting/sampler.py	`92.59% <92.59%> (ø)`

compatibility for resampling API

…into bayesian_extend

damskii9992

Just gonna stop my review here as there is already plenty to work on.
As we discussed in-person, I also think we should refactor the bayesian sampling to a seperate "samplers" module. This should help us when we add other Bayesian samplers in the future as well as provide a cleaner user API (the user doesn't create a "Fitter" when doing sampling).
It also more cleanly separates responsibility in the code. A Minimizer minimizes a function to the data, a Sampler samples the parameters to create distributions.
To avoid duplicating code, it might be smart to move shared functionality to new base-classes.

damskii9992

I'm gonna stop my review again here as the changes I have proposed are again substantial and another review would likely be required afterwards.

My biggest issue lies with the class/file structure. I think we should copy the factory/minimizer structure for the samplers, as this is an easily extendable and logically sounds structure and it makes reviews easier as the structure would be the same across our codebase. But we probably need an ADR for this and to have discussion about it, probably in-person.

damskii9992 · 2026-07-07T13:10:35Z

+    "extended_draws = extended_results.draws\n",
+    "\n",
+    "print(f'Short chain: {short_draws.shape[0]} retained draws')\n",
+    "print(f'Extended chain: {extended_draws.shape[0]} retained draws')"


The output here is that the extended chain has 7000 draws. Adding 5000 with a thin of 2 should make the chain 2500 longer. Clearly we're applying the burn-in of 500 draws again, but why are we doing that when we extend the chain? Shouldn't that only happen when you initiate the chain?

Nope. the burn-in is not re-applied. Sampler.extend() passes burn=0, and Bumps.mcmc_sample sets burn=0 whenever resume_state is not None, warning if the caller asked otherwise. There is no path that re-burns a resumed chain.

I'll update the notebook/docs with proper statements

2-parameter model, sample(samples=10000, burn=500, thin=2) then extend(additional_samples=5000, thin=2)):

Config sample() draws extend() draws Delta

default (trim=True) 4520 6600 +2080

sampler_kwargs={'trim': False} 5000 7500 +2500

With trimming off, the numbers are exactly as you expect. The ring buffer also grows exactly as intended: Nthin * Npop == 7500, and state.draw(portion=1.0, outliers=True) returns all 7500 points. No draws are lost, and none are re-burned.

Running the actual tutorial (4 parameters, not the 2 above) reproduces your number exactly and shows where it comes from:

chain results.draws (trimmed) full chain

original 4400 5000

extended 7000 7600

So the "missing" draws were never missing: the chain grew from 5000 to 7600.

Okay, but this still doesn't explain everything.
So the "Full" chain included the burn-in but not the thinned-out, since the first full chain is 5000 steps, but the extended is 2600 larger, which is more than half of the 5000 asked for. You write that DREAM rounds up to nearest whole generation, so I guess that must be 200 steps? Is that the default or is that a parameter we can change? It is weird that it runs longer than you ask it to.

damskii9992 · 2026-07-07T14:15:23Z

+    Construct directly with a configured ``Fitter`` (or ``MultiFitter``) whose
+    minimizer has been switched to ``AvailableMinimizers.Bumps``.


When the Sampler is constructed from the Fitter it appears like you have to do a fit first, which you don't.
Instead of having the Sampler take a Fitter as input, I would make the Sampler an independent class.
Maybe it can have a constructor Sampler.from_fitter which Constructs a Sampler from a Fitter object, but I wouldn't make it required.

Also, we should probably copy over the structure of the minimizers. I.e. with a factory, a list of available samplers (for now just DREAM) and separate backend classes.

Think of the cost: Sampler reaches into Fitter internals (_precompute_reshaping, _fit_function_wrapper, and assigning fitter.fit_function to swap the wrapper.
Making it standalone means extracting that machinery into a shared base, which will affect the basic functionality of the fitter. I woudn't want to do it unless really, really needed.

I will reword the docstring/notebook to say no prior fit is needed as a solution.

I understand why you're reluctant, this is the easy solution, but it is also the messy solution which will compound our technical debt. By not doing this now, you're just postponing this task and making it harder in the future. It also means that WHEN we do this task in the future, it is gonna be an API-breaking change, whereas by adding it correctly from the get-go isn't API-breaking.

Following our discussion on Slack, I will reluctantly relent this issue to a (soon) future PR.

rozyczko · 2026-07-15T07:33:14Z

My biggest issue lies with the class/file structure. I think we should copy the factory/minimizer structure for the samplers, as this is an easily extendable and logically sounds structure and it makes reviews easier as the structure would be the same across our codebase. But we probably need an ADR for this and to have discussion about it, probably in-person.

There is exactly one sampler (DREAM) and one backend.
Building a factory + enum + base class for a single implementation is, although potentially useful, slightly overengineering at this point. The Sampler._run BUMPS check is already documented as the place for a second backend. I would suggest creating an ADR, but not modifying this PR. Let's do it when another sampler actually gets implemented, so the abstraction is derived from two real cases rather than guessed from one.

rozyczko · 2026-07-16T08:53:33Z

My biggest issue lies with the class/file structure. I think we should copy the factory/minimizer structure for the samplers, as this is an easily extendable and logically sounds structure and it makes reviews easier as the structure would be the same across our codebase. But we probably need an ADR for this and to have discussion about it, probably in-person.

There is exactly one sampler (DREAM) and one backend. Building a factory + enum + base class for a single implementation is, although potentially useful, slightly overengineering at this point. The Sampler._run BUMPS check is already documented as the place for a second backend. I would suggest creating an ADR, but not modifying this PR. Let's do it when another sampler actually gets implemented, so the abstraction is derived from two real cases rather than guessed from one.

#280

rozyczko · 2026-07-17T09:03:08Z

I added the plan for refactoring in #280

This follows discussion in Slack.

I would prefer addressing this in a separate PR.

damskii9992

Answers to comments, will continue with my review.

damskii9992 · 2026-07-16T13:03:51Z

+    "extended_draws = extended_results.draws\n",
+    "\n",
+    "print(f'Short chain: {short_draws.shape[0]} retained draws')\n",
+    "print(f'Extended chain: {extended_draws.shape[0]} retained draws')"


Okay, but this still doesn't explain everything.
So the "Full" chain included the burn-in but not the thinned-out, since the first full chain is 5000 steps, but the extended is 2600 larger, which is more than half of the 5000 asked for. You write that DREAM rounds up to nearest whole generation, so I guess that must be 200 steps? Is that the default or is that a parameter we can change? It is weird that it runs longer than you ask it to.

damskii9992 · 2026-07-16T13:10:13Z

+    except Exception:
+        return None


You should AT LEAST report a warning then, maybe using your fancy-schmancy new logging module? The user should be notified that the saved state lost its identification fingerprint.

damskii9992 · 2026-07-16T13:15:59Z

+    from .minimizers.minimizer_bumps import save_sampler_state
+
+    save_sampler_state(state, str(path))


So you don't import a save function now, but you still have both a function and a method on the class, why not consolidate them? They're quite small functions/methods both.

damskii9992 · 2026-07-16T13:22:13Z

+    Construct directly with a configured ``Fitter`` (or ``MultiFitter``) whose
+    minimizer has been switched to ``AvailableMinimizers.Bumps``.


I understand why you're reluctant, this is the easy solution, but it is also the messy solution which will compound our technical debt. By not doing this now, you're just postponing this task and making it harder in the future. It also means that WHEN we do this task in the future, it is gonna be an API-breaking change, whereas by adding it correctly from the get-go isn't API-breaking.

damskii9992 · 2026-07-17T09:23:02Z

+    Construct directly with a configured ``Fitter`` (or ``MultiFitter``) whose
+    minimizer has been switched to ``AvailableMinimizers.Bumps``.


Following our discussion on Slack, I will reluctantly relent this issue to a (soon) future PR.

rozyczko · 2026-07-17T12:14:02Z

Issues from the most recent round addressed.

damskii9992

I did not look closely at the tests, neither did I pay too much attention to the structure since we plan to change that, but here are the last comments I found after going through the PR again.

damskii9992 · 2026-07-17T09:31:51Z

 from .utils import FitError
 from .utils import FitResults

-__all__ = [MinimizerBase, Bumps, DFO, LMFit, FitError, FitResults]


I recently learned that all the __all__ attributed does, is change what gets imported when you do star imports, i.e.:

from easyscience.fitters.minimizers import *

The default for that is to import all modules, so simply adding all modules to an __all__ is completely redundant . . .

damskii9992 · 2026-07-17T09:33:43Z

        burn: int = 2000,
        thin: int = 10,
        population: int | None = None,
+        resume_state: Any | None = None,


I assume this shouldn't be typehinted Any?

damskii9992 · 2026-07-17T09:50:28Z

+                # Recover the original scale factor from the state's Npop.
+                # Try ceil-division first, then fall back to floor.
+                recovered = int(math.ceil(resume_state.Npop / n_params))
+                if int(math.ceil(recovered * n_params)) != resume_state.Npop:


You KNOW that recovered and n_params are integers, hence you don't need the int and math.ceil calls here.
Also, what are you even trying to do here? You get the floor rounding always, UNLESS resume_state.Npop / n_params is an integer-like float, in which case the ceil function does the same as int?
This section kinda seems redundant.

damskii9992 · 2026-07-17T09:53:50Z


        return _outer(self)

    def mcmc_sample(


This method is getting stupidly long now with the resume functionality, please try to factor out some functionality into helper methods to ease readability.

damskii9992 · 2026-07-17T09:57:05Z

 # from .multi_fitter import MultiFitter  # noqa: F401, E402

-all = [AvailableMinimizers, Fitter, FitResults]
+all = [AvailableMinimizers, Fitter, FitResults, Sampler, SamplingResults]


As per my other comment.

damskii9992 · 2026-07-17T12:24:46Z

There is currently a lot of BUMPS-specific code in this file, which makes sense since that is the only sampler we currently have, but bear in mind that WHEN we have another sampler, there should be NO bumps-specific code or syntax in this file, those should all live in the minimizer_bumps.py file.
All syntax here should be sampler-independent.

damskii9992 · 2026-07-17T12:30:26Z

+        Calling ``sample()`` on a sampler that already holds a chain starts a
+        **fresh** chain — the previous state and results are replaced. Use
+        ``extend()`` to continue an existing chain.


Maybe add a warning if a user tries to run sample() when there already is a result, something like:
"Warning, replacing previous chain with fresh chain, if you meant to continue the chain, please use the extend() method"
Or something along those lines.

damskii9992 · 2026-07-17T12:32:29Z

+                warnings.warn(
+                    'The data bound to this Sampler does not match the data '
+                    'fingerprint stored with the saved chain. Extending a '
+                    'chain against different data is undefined behaviour.',
+                    UserWarning,
+                    stacklevel=2,


You should probably use your own logging system here.

damskii9992 · 2026-07-17T12:36:05Z

+# NOTE: The Bayesian MCMC tests formerly here (TestMultiFitterMcmcSample) were
+# moved to test_sampler.py and adapted to ``Sampler(f, ...)``.


We probably don't need this comment lol.

damskii9992 · 2026-07-17T12:48:08Z

Most of these tests belong in unit tests. Remember, the unit tests folder should mimic the src folder. So sampler.py should have a mirror test_samlper.py in unit tests.

Integration tests are only for full runs, like running an actual sampling and checking the resulting R value or convergence or the value of some parameter or some such.
It would also be the place to have a full save/load/resume roundtrip test.

rozyczko added 2 commits June 4, 2026 11:50

extended sampling with BUMPS

3858923

minor changes

da6f81f

rozyczko added 2 commits June 4, 2026 13:37

PR code review

7af847e

updates for the most recent bumps version (1.0.4), which breaks

bde4d37

compatibility for resampling API

rozyczko added [scope] enhancement Adds/improves features (major.MINOR.patch) [priority] high Should be prioritized soon labels Jun 4, 2026

rozyczko added 5 commits June 4, 2026 18:30

updated notebook

0dee734

removed unnecessary attribute

162b14a

Merge branch 'bayesian_extend' of https://github.com/easyscience/core …

84ad75b

…into bayesian_extend

fixes after merging develop

f43a388

fixed notebook, fixed pyproject

28e0c53

rozyczko marked this pull request as ready for review June 8, 2026 07:51

damskii9992 reviewed Jun 11, 2026

View reviewed changes

rozyczko mentioned this pull request Jun 17, 2026

EPIC: Bayesian analysis support #153

Open

rozyczko added 3 commits June 17, 2026 13:51

issues from the PR addressed. Bigger rewrite, so separate branch (#260)

e723796

Merge branch 'develop' into bayesian_extend

468da2f

minor fixes requested

a72ffdd

rozyczko requested a review from damskii9992 June 17, 2026 13:21

fixed pytest.skip argument

013e834

damskii9992 reviewed Jul 7, 2026

View reviewed changes

rozyczko added 3 commits July 15, 2026 16:01

PR review comments addressed. Round 2.

52b5308

potential linting fix

1408048

Merge branch 'develop' into bayesian_extend

8c83ab3

rozyczko requested a review from damskii9992 July 16, 2026 09:40

damskii9992 reviewed Jul 17, 2026

View reviewed changes

more PR fixes

feb2595

damskii9992 requested changes Jul 17, 2026

View reviewed changes

Config	`sample()` draws	`extend()` draws	Delta
default (`trim=True`)	4520	6600	+2080
`sampler_kwargs={'trim': False}`	5000	7500	+2500

		Construct directly with a configured ``Fitter`` (or ``MultiFitter``) whose
		minimizer has been switched to ``AvailableMinimizers.Bumps``.

		from .minimizers.minimizer_bumps import save_sampler_state

		save_sampler_state(state, str(path))

		# NOTE: The Bayesian MCMC tests formerly here (TestMultiFitterMcmcSample) were
		# moved to test_sampler.py and adapted to ``Sampler(f, ...)``.

Uh oh!

Conversation

rozyczko commented Jun 4, 2026

Uh oh!

codecov Bot commented Jun 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

damskii9992 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

damskii9992 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rozyczko Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rozyczko Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

rozyczko commented Jul 15, 2026

Uh oh!

rozyczko commented Jul 16, 2026

Uh oh!

rozyczko commented Jul 17, 2026

Uh oh!

damskii9992 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rozyczko commented Jul 17, 2026

Uh oh!

damskii9992 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Jun 4, 2026 •

edited

Loading

rozyczko Jul 15, 2026 •

edited

Loading

rozyczko Jul 15, 2026 •

edited

Loading