Callbacks overhaul

[leonardogalliano]

⚠️ WARNING: this PR breaks backward compatibility. See details below.

Callbacks previously depended on the full Simulation object and were designed to return quantities averaged over all parallel chains. This made it difficult to monitor per-chain observables (e.g. the instantaneous potential energy of each chain independently).

While the conceptual change is simple, implementing it required non-trivial refactoring across the codebase. Google Antigravity helped.

Changes

• Callbacks now take a system as input.
• The StoreCallbacks algorithm now stores results independently for each chain.
• The trajectories/ folder has been renamed to chains/, reflecting the fact that it contains heterogeneous data (trajectories, configurations, restart files, etc.), not just trajectories.
• callback_acceptance has been removed. Acceptance tracking is now handled by a dedicated algorithm, StoreAcceptance, which stores per-move acceptance information in the appropriate folders. This is more consistent with Arianna’s design philosophy.
• StoreParameters has been redesigned. Previously, it created a parameters/ folder for each move containing only parameters.dat, which was unintuitive.
• There is now a single moves/ folder, with one subfolder per move, containing all move-related data (parameters, acceptance, etc.).
• A new StoreObjective algorithm has been introduced. It stores the average reward $J$ from PGMC in moves/. Previously, there was no systematic way to track $J$.
• All examples, tests, and documentation have been updated accordingly. The most important compatibility changes are:
  • replacing callback_acceptance with StoreAcceptance;
  • renaming trajectories/ to chains/.
  • Other changes are mostly transparent to users.
• Added a new “Design strategy” section to the documentation (adapted from the PhD manuscript), outlining the design principles behind Arianna. This should be useful for future contributors.

Backward compatibility and ParticlesMC

The main compatibility issue concerns, of course, ParticlesMC. Merging this PR will require a new Arianna release and a corresponding update of ParticlesMC.

The required changes in ParticlesMC should be manageable:

• redefine callback_energy to take system as input
• replace callback_acceptance with StoreAcceptance in all example and test files
• update folder names from trajectories/ to chains/.

[codecov]

Codecov Report

❌ Patch coverage is 95.53571% with 5 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/algorithms.jl	92.30%	3 Missing ⚠️
src/metropolis.jl	95.23%	2 Missing ⚠️

Files with missing lines	Coverage Δ
src/Arianna.jl	100.00% <ø> (ø)
src/PolicyGuided/PolicyGuided.jl	100.00% <ø> (ø)
src/PolicyGuided/estimator.jl	100.00% <100.00%> (ø)
src/metropolis.jl	92.25% <95.23%> (+0.80%)	⬆️
src/algorithms.jl	93.38% <92.30%> (-1.27%)	⬇️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[romainljsimon]

Hey @leonardogalliano, this is a great PR, thanks a lot for this extensive work which is super useful.
I want to take a bit of time to review everything.
I wondered if we wanted to include the macro callback in this PR to not have any breaking changes in the future. What do you think about this point? Do you still want to include the macro ?

[leonardogalliano]

Yep we can do that! I made a small modification. Now we can define a callback with

@callback function energy(system)
    return system.e
end

instead of with callback_energy. I've kept the old callback_... option for compatibility, but we can remove it if needed.

Then to store the callback we add the algorithm

...
(algorithm=StoreCallbacks, callbacks=(energy,), scheduler=sampletimes),
...

I've realised it's not a big change, but it's true that it's a bit cleaner. Is this what you had in mind in #68 ? Or should StoreCallbacks handle them automatically?

[romainljsimon]

This is great! Yes the implementation is way more concise that I what I first imagined.
I don't think StoreCallbacks should handle them automatically.
However I'm wondering for consistency if we want that all callbacks functions to have the @callback macro in front, meaning that function callback_energy cannot be a callback. To be a callback, a function has to be written as @callback.
Also do we want the name to be callback or @Arianna.callbackor @ariannacallback. What are your opinions about all this ?

[leonardogalliano]

In principle a callback can be any function, for now there is nothing special about them. The @callbackmacro just handles the name of the corresponding .dat file at the moment. But we could be more specific and enforce each callback to come with the macro, for me both options are fine.

And regarding the name of the macro I don't have preferences, pick the one you like the most!

[romainljsimon]

I chose to delete the export of the macro because callback seems like a pretty generic name. So users have to use Arianna.@callback. I also chose that callback functions have to be preceded by the macro for consistency. If a callback is called by StoreCallbacks but isn't preceded by the macro then StoreCallbacks will return an error.
Good for you @leonardogalliano?
I also had to update the Enzyme to 0.13 for it to work on my computer.

[romainljsimon]

Ok there's a problem with Enzyme :'( and the Julia versions

[leonardogalliano]

Nice! Ok for the macro changes.

Enzyme is always a pain... It used to work with the previous version of Project.toml when I did the PR, but it looks the new versions are not anymore compatible with Julia 1.9?

[romainljsimon]

Should work now @leonardogalliano !