Monte Carlo re-implementation in Python

[jngrad]

Fixes #4617, fixes #5178

Description of changes:

• rewrite Monte Carlo methods in Python
• allow user-defined MC acceptance criteria

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[jngrad]

Target is ESPResSo 5.1 release, therefore only non-silent API changes are allowed.

TODO list:

• refine the API of MonteCarloMethod and ReactionAlgorithm
  • of particular interest is the separation of generic Monte Carlo from reaction-based Monte Carlo (see e.g. the unfinished CanonicalEnsemble class)
  • the main goal is to allow end users to write their own MC algorithms, possibly without reactions
  • some of the formerly private C++ methods are now available as Python functions, some of which should be hidden with an underscore prefix, e.g. restore_system() and find_missing_pids()
  • would we like to rename reaction_methods.py to monte_carlo.py?
  • update: MonteCarloMethod and CanonicalEnsemble were removed, but can be re-introduced later
• several C++ unit tests need to be ported to Python unit tests
• random number generators
  • we now use Philox from NumPy, MC trajectories have changed because the random sequences are different and several test tolerances were adjusted as necessary
  • the new Python implementation and the old C++ implementation generate the same trajectories when the same RNG sequences are used (requires sorting particle ids before choosing one at random, see 2d99741)

    LEGACY_MONTE_CARLO=1 ./pypresso ../testsuite/python/constant_pH.py > logcore.log
    LEGACY_MONTE_CARLO=0 ./pypresso ../testsuite/python/constant_pH.py > logpython.log
    git diff --no-index logcore.log logpython.log

• investigate checkpointing
  • the old implementation would pick up the main ESPResSo system
  • the new implementation requires the user to specify which ESPResSo system to use
  • we could define a new member system.monte_carlo.add(ConstantpHEnsemble(kT=1., pH=7.)) to create a strong link between the system and MC method(s) and enable checkpointing
  • look into pickling the RNG state via numpy.random.Philox.state
• design a more user-friendly reaction definition scheme to eliminate reaction_id confusion
  • for example, in generic_oneway_reaction(), get_acceptance_rate_reaction(), calculate_log_acceptance_probability(), and change_reaction_constant(), reaction_id refers to the 0-based index of the interleaved forward and backward reactions, while in calculate_particle_insertion_potential_energy(), and delete_reaction(), reaction_id refers to the 0-based index of the forward reactions (it is multiplied by 2 internally)
• identify and remove performance regressions

[jngrad]

@kosovan @pm-blanco @jorch28 @RudolfWeeber I would propose to work on the refactoring in two stages.

First, convert the C++ code to Python code and merge it into the python branch. This is now 95% complete. There are a few items that would require your input to finalize this PR, and I will discuss them further down below.

Second, work on the bigger API changes, such as:

• introducing a MonteCarlo base class that is agnostic to chemical reactions
  • to go into this direction, we would need a use case, so I can write a python test to guarantee the MonteCarlo API is self-contained and doesn't overlap with ReactionAlgorithm
• introducing support for a MC method that manages 2 or more systems, e.g. replica exchange MCMC sampling
  • multi-system simulations are already possible for most features, one notable exception is reaction methods, and to make progress on that front, this PR needs to be merged
• replacing the existing slab and cylindrical constraints by shapes
  • I'd be in favor of this since these classes already contain all the logic for managing e.g. cylinder orientation
  • the shape would need to be sanity checked (e.g. a Union of two walls must contain parallel walls), with an allowlist to reject irrelevant shapes (only sphere, cylinder, and union of parallel walls seem relevant)

My main motivation for this approach is removing the legacy type_map now, so that I can start working on removing the legacy MPI communication system and make multiple systems possible in the same simulation script without requiring users to manually remove the safety check currently in place to prevent multiple system instances. The bigger API changes can be carried out independently of this MPI refactoring.

The following parts of the code would require additional input:

• @RudolfWeeber several calls to system.cell_system.call_method("invalidate_ghosts") were introduced to address runtime errors about particle id mismatch, and it's unclear to me if they are needed always or just for MC
• @jorch28 @kosovan I don't plan any further big changes to the code, only cosmetic ones like removing commented out code. Bookkeeping methods like delete_created_particles() are now part of the public interface, but I'm not sure we want users to have access to those. If that's fine with you, I would prefix all of these new methods with an underscore to signal to users that those are private and not meant to be part of the public API. Users who develop new MC method can of course still call them, but they should bear in mind we can change the API and/or behavior of private functions at any time.

The Python implementation is roughly 3 times slower that the C++ one according to the acid-base reservoir benchmark. The overhead mainly comes from the particle changes. The new code already uses particle property updates via dict, and implementing a bulk particle update (i.e. a list of dicts and a list of ids to trigger a single cell system rebuild) didn't improve runtimes. The particle numbers per type are calculated once at the start of RE.reaction(setps=100) and stored in a cache that is updated for all further reaction steps, so that we no longer need to track particle types in the C++ core via expensive MPI reduction operations.

[jorch28]

@jngrad
A comment on making private the method delete_created_particles():

As I have seen, restore_system() depends on delete_created_particles() and restore_system() is employed to reject MC trials moves and restore the system to its original state before proposing the move. This can be seen in methods self.generic_one_way_reaction(), self.displacement_mc_move_for_particles_of_type, self. calculate_particle_insertion_potential_energy(). Therefore, it seems a crucial method to develop new MC algorithms. Thus, if delete_created_particles() becomes private would restore_system() also become a private method?
If it does, it seems to me that developers who want to build new MC algorithms, would have to write their own restore_system methods (because would not be able to trust its functionality), which would may represent an extra effort for them.

By the way, the rogue method , calculate_log_acceptance_probability, removed in previous commit was actually being employed in method generic_one_way_reaction in line 946. The statement is:

ln_bf = self.calculate_log_acceptance_probability(
reaction, E_pot_diff, old_particle_numbers)
reaction.accumulator_potential_energy_difference_exponential.append(
math.exp(-E_pot_diff / self.kT))

I don't know if you passed the tests after making the commit, but, as ReactionAlgorithm is not a child class, it seems like deleting the rogue method is Bug.

[jngrad]

@jorch28 thank you for looking into this. The testsuite passed. The green check mark will give you more details about the testsuites. The method calculate_log_acceptance_probability(reaction, E_pot_diff) was really confusing, because it was missing an argument, and it couldn't be called since all children classes override this method with calculate_log_acceptance_probability(reaction, E_pot_diff, old_particle_numbers). I will restore the raise NotImplementedError that I accidentally deleted in a previous commit.

restore_system() can stay public. I don't have a strong opinion on delete_created_particles() and am open to leaving it public.

[jorch28]

@jngrad Ok, now I see that all check have passed refers to the testsuite from pymbe and espresso. Then, there is not much to worry about.
And yes, subclasses reimplement the method with an extra argument and pass an extra argument as well: old_particle_numbers.

Sorry to insist, I must be missing something, but I still think there is something wrong with the method generic_one_way_reaction, but I cannot find what it is:

Because, it calls self.log_acceptance_probability in line 946, with an extra argument (which should raise an error), plus it raises NotImplementedError when called from the same class ReactionAlgorithm. Therefore, should give an error instead of returning a float.
Unless, from line 944 to 960 is not being executed, because, all reaction trial moves are being rejected previously as inside_exclusion_range_touched ==True (Which does not make much sense, because all tests have passed).

About delete_created_particles(), users should not touch unless they know what they are doing, because it would break restore_system().

[jngrad]

Because, it calls self.log_acceptance_probability in line 946, with an extra argument (which should raise an error)

Isn't this already fixed by 623d698?

plus it raises NotImplementedError when called from the same class ReactionAlgorithm.

The implementation was removed in 27aee23 because constant_pH isn't a meaningful concept in the reaction ensemble and Widom insertion methods. If my understanding is correct, SingleReaction stores information about a chemical reaction, ReactionAlgorithm is a framework that stores re-usable building blocks for Monte Carlo moves of chemical reactions. Derived classes define the statistical ensemble in which the reaction takes place, which determines the analytical expression of the transition probability. The NotImplementedError is one way to signal to developers they have to define the transition probability in their ensemble. Another way would be to use @abc.abstractmethod to prevent users from instantiating the class without defining the method. We don't use abc metaclasses in ESPResSo because we designed our own metaclass framework (@script_interface_register and friends), and it's unclear to us how these interact with each other (especially during checkpointing).

inside_exclusion_range_touched ==True

When I suspect the presence of a bug, I like to craft a minimal working example (MWE) that attempts to trigger it. Typically, a MWE is 5 to 10 lines of Python code. This technique helps me better understand the code behavior and determine which specific conditions need to be met for the bug to occur. I use print statements to guide me through the nested code paths and update my mental model of the functionality. Would you be interested in trying that out?

About delete_created_particles(), users should not touch unless they know what they are doing, because it would break restore_system().

Would you prefer delete_created_particles() to be private or public? Exposing this level of detail to end users and promising to maintain these implementation details in a public API seems like a strong commitment for us. If we expect users to have a need to frequently write their own restore_system() implementation, then we should keep delete_created_particles() and the other bookkeeping functions public (because users should never be forced to write a call to a private function). The only use case I can think of is Pablo's dynamic bond formation, where the goal is to delete bonds instead of particles.

[jorch28]

My apologies for not replying much earlier.

Isn't this already fixed by 623d698?

Yes it was fixed. I did not notice it at a first glance.

The implementation was removed in 27aee23

Thank you for clarifying, now I understand that ReactionAlgorithm is the parent class of the MonteCarlo methods : ReactionEnsemble(), ConstantpH() which override the method log_acceptance_probability().

When I suspect the presence of a bug, I like to craft a minimal working example (MWE)... Would you be interested in trying that out?

This late afternoon I'm going to dedicate time to craft this MWE to better understand the code.

Would you prefer delete_created_particles() to be private or public?

Honestly, I only have a basic understanding of MonteCarlo methods. If you think that is a strong commitment that we may break at long term, then it is better to put them as private methods .

[RudolfWeeber]

I think this is a great step towards providing more flexibility to users in terms of MC methods.
There are mainly two items that we need to addres.

1. The Python call to invalidate_ghosts points to a problem in the core. This, we should fix before merging.
2. performance: here, we should aim to get some of the speed difference back. If we do this now, or merge this pr first and follow up with a performance PR soon, I have no strong opinion, either way.

invalidate ghost

The AI thinks, and I agree, that this is actually an ommision in the core funciton that delets particles.
After a removal, all ranks that held ghosts of the particle need to call invalidate_ghosts immediately (which hopefully also invalidates the particle index or at least deletes the ghosts form that)

Performance

There are a few candidates, such as using np.choice(..., replace=False) to drw mupltiple pids, or using slices for bulk updates.
To figure out, what is worth it I suggest benchmarking typical application scenarios, using both, reactions and displacement moves.
I had good experience letting Claude do this: "Identify performane optimiztions in the monte carlo code in src/python/espressomd/reation_methods.py and associated core functions. To this end, run build/pypresso my_benchmark.py using py-spy or similar, and based on the results, propose a prioritized list of optimizations ordered by bang-for-buck ratio."

As mentioned above, I'm also fine with only doing the invalidate_ghost thing in this pr and following up with performance in a different pr.

[pm-blanco]

I agree with @RudolfWeeber that this is a great step towards providing a more flexible MC framework for users and developers and therefore I am looking forward to seeing this PR merged. I think that splitting the refactoring in two stages is sensible, and having an agnostic MonteCarlo is something it is extremely needed (there are already MC methods in the library that are not reaction methods like displacement_mc_move_for_particles_of_type.

Regarding if delete_created_particles should be private or not, my opinion is that it should be a private method because the operations that it does (involving ghost particles) and particle ID bookkeeping require advanced knowledge of how espresso works internally that we cannot expect standard users of espresso to be acquainted with. Moreover, since there are still plans for refactoring the Monte Carlo methods, I think is best if we do make strong commitments for those helper methods. Future developers of reaction methods (including myself when I finally can integrate the eb-RxMC into espresso...) that want to use those methods can still use those private methods if they need to. I actually wonder if other similar helper methods should be private too, like delete_hidden_particles. hide_particle, restore_system and find_missing_pids

I checked the Monte Carlo implementation and it looks OK to me, apart from the missing documention in various methods. I understand correctly that the plan is to solve the issue with MC methods and constraints in the next PR? I think line 827 in pos=self.rng.random((3,)) * self.system.box_l where particles are created in any position of the box irrespectiely of constraints may be the one breaking the MC simulations with constraints. Also, in line 403 in get_volume, shouldn't the volume of a cylinder be pi * radius**2 * box_l[2] and not pi * radius * box_l[2]?

[jngrad]

Here are benchmark results for mc_acid_base_reservoir.py with 118 particles per thread and Debye-Hückel electrostatics on an AMD Ryzen 9 9950X3D, comparing the ESPResSo 5.0.1 release against this PR, only showing the reaction step times (MD step times are very close in the old and new code):

threads	5.0.1 code	PR code	speed-up
1	62.7 us/step	67.4 us/step	-7%
2	95.1 us/step	85.2 us/step	10%
4	145.4 us/step	109.8 us/step	25%

Overall, the performance loss of moving all the logic from C++ to Python was well mitigated. When using more than 1 OpenMP thread, we are actually faster. These benchmark results will be slightly less impressive versus 5.0.2, because we plan on backporting some of the MC performance improvements introduced by this PR.

@pm-blanco and @jorch28 this PR can be added to the agenda of the next pyMBE meeting. It is now in a state where it can be merged. Further API changes can be introduced in a follow-up PR. To get this to work with pyMBE, you will need to use my es510 branch, where I dropped support for 4.2 and introduced support for 5.1-dev. The pyMBE unit tests pass in CI according to the debian workflow, but some of the functional_tests fail (e.g. cph_ideal_tests, weak_polyelectrolyte_dialysis_test). I think this is due to the RNG sequence being different. Although there were regressions in the PR code where cylindrical and slab constraints where ignored by some reaction methods when sampling new positions and calculating volumes, these bugs have since been fixed by the last two commits.

[pm-blanco]

LGTM

[hidekb]

LGTM!