Add a seed for integrator step function to reproduce results

[thomasloux]

Currently, one can make an integrator init function reproducible by providing a seed: it will used for initial momenta of the system. Yet, for stochastic integrator like Langevin, it would also be useful to provide a seed.
This would also be useful to check that NPT and NVT Langevin keep the same behaviour overtime, by harcoding the some property values at the end of a simulation in tests.

Technically, it would just require to add an optional argument for a torch.Generator. It would be defined before the start of the simulation and provided in step_function each time.

Potentially, we could discuss the possibility to also be able to reproduce single system simulated in a batch context.
Example:

• I simulate 2 systems
• For some reason, I want to reproduce any of the two systems
• If I want to reproduce exactly the same trajectory, I need to simulate again both systems.

This would imply changing:

• Generation of initial momenta: currently if I simulate a batch of 2 exact systems with Nosé Hoover, I get 2 different trajectories because the initial momenta is different. In this case, it makes sense, but it would not be practical to get the initial momenta of the second system, even knowing the seed.
• Generation of noise in step function: currently only needed for Langevin.

For the second step, it not as clear to know how to modify it in clear way for a user. The naive way to do that would be to have a generator per system, which is cumbersome if you explicitly need to provide the list of generator each time you run the step function.

[CompRhys]

related issue: #51

I agree we should refactor how we're handling PRNGs in the code. Part of me thinks that perhaps the right solution would be for the PRNG to live on the SimState and then any integrators/optimizers borrow from the SimState. If we have batched SimStates I can't really think of any batchable solution that would allow you to get the same result for a system seeded individually and a system group into a batch. I don't think that's a blocking issue.

[abhijeetgangan]

Agree, that there should be a seed var in the SimState. It could be a per-system var. Ideally if someone sets all seeds same across the batch then it would just result in same values across the simulations so a global seed won't do much.

Caveat here is that currently none of the simulations except the classical force-fields are deterministic (unless you want to run on CPU) so seeds don't offer much since none of the simulation trajectories would be same even if you enforce the seed. There are flags which make the simulation deterministic but I haven't tested them out.

[thomasloux]

Why do you say so? MLIPs are deterministic, up to numerical rounding maybe or possible DAG optimisation and for instance Nose Hoover is deterministic. The seed would be useful for noise sampling with Langevin.

[abhijeetgangan]

Any small numerical noise from GPU scatter operations causes exponential divergence in trajectory because MD integrators are chaotic. Hence , your trajectories would not be the same even if you re-ran it again with exact same settings. In JAX, you would set os.environ["XLA_FLAGS"] = '--xla_gpu_deterministic_ops=true' to get true reproducibility at the cost of some performance. The root cause is simple - floating point arithmetic is non associative and the atomic operations on GPU perform operations in different order and thus you get different result. Couple that with any chaotic simulation and the trajectories diverge. Your MD averages should be okay (ergodicity).

Agree with your point of having a seed for the noise sampling part in Langevin or any stochastic integration schemes for that matter.

[thomasloux]

Agreed, to be honest it would probably be more for debug and integrator tests, over short simulation so it doesn't have time to diverge

[abhijeetgangan]

Yeah, I don't think it will be an issue for debug and tests, since they are in CI with CPU so the values should be exact and even if we have GPU in CI a short simulation shouldn't be that different. In both cases the averages should agree as well.

[thomasloux]

Ok so what should be the plan?

I think we should add a torch.Generator and not just a seed, to each system and use whenever we need randomness.
If it's stored with SimState (or rather MDState) this should be easy for the user.

[abhijeetgangan]

I am open to either but more preference to the second option:

• Add the generator in the MDState (as optimizers won't use them so no use in SimState)
• Add it only to Langevin while others which don't need generators can only keep the seed like NVT Nose Hoover.

[CompRhys]

I think there a bit of value in just having it be an option on the base SimState for things like implementations of rattling etc that can be valuable in optimization workflows for finding ground states. Not a strongly held opinion as it's easy enough to promote at later date given that that functionality isn't available yet or planned.

[orionarcher]

This is a great point to raise and a bit of an ongoing annoyance.

I'd push back on adding seed/generator to SimState. It's a nice idea that seed/generator would be a system-wise attribute that respects sampling for each system according to it's seed, but in practice, the integrators/optimizers are sampling for the whole batch simultaneously and only take a single seed/generator. Even if we wanted to rewrite the code to respect system-wise seeds, that would require looping through each system and sampling for each separately, which isn't ideal. Thus, seed is more a global than system-wise attribute and there is no clear way to "correctly" choose a seed when multiple sim states are concatenated. To me, that suggests it doesn't belong in SimState.

I'd favor allowing step/initialize functions to support either a PRNG or a seed (in which case it internally creates a PRNG). This would allow either 1) the same generator to be reused when running multiple steps, or 2) a seed to be fixed for a single step if I need to run it identically twice. I'm happy to be convinced otherwise but I'm not seeing a reason to put it in SimState instead?

[CompRhys]

It could be a per-system var. Ideally if someone sets all seeds same across the batch then it would just result in same values across the simulations so a global seed won't do much.

Even if we wanted to rewrite the code to respect system-wise seeds, that would require looping through each system and sampling for each separately, which isn't ideal.

I also can't see a way to do per system seeds and not needing a for loop for doing #51.

I still think that having the SimState hold the PRNG with the local fallback is better than having it as a global fallback or always having to pass it around. For 1) the update function either acts on two objects, a SimState and a PRNG sequence, or one, a SimState that has its own PRNG sequence. If the generator lives on the SimState the same PRNG sequence is used by design without the user having the option to make a mistake. For 2) if you set the seed at the start of a script and ran the same sequence of updates on that SimState it would be fully reproducible?

In terms of combining SimStates I don't think it's possible to make it associative when sampling batchwise sets of random numbers so the order of concatenation would always impact the randomness subsequently so provided you just reseeded based on some combination of the PRNG state of individual SimStates then I can't see a problem. You would be destroying some information in that the concatenated SimState wouldn't know what seed was used for the first X steps of sub system x before the concatenation but the batched SimState doesn't have that sort of memory for anything so I don't see any harm.

[orionarcher]

I still think that having the SimState hold the PRNG with the local fallback is better than having it as a global fallback or always having to pass it around.

I tend to agree that it be a bit more ergonomic when wants to set the seed. My objection is partly of taste, that seed doesn't feel like part of the simulation state, and partly of bloat, do we want to add seed to SimState when it's optional and many users won't be using it? Especially when determinism is iffy at best in MD, I'm a little hesitant to include seed as part of SimState. If other people generally prefer seed on SimState I'm fine being overruled here though.

so provided you just reseeded based on some combination of the PRNG state of individual SimStates then I can't see a problem.

Yeah this seems like a reasonable solution to me.

[orionarcher]

From @CompRhys on slack, I am picking the discussion back up here.

pytorch/pytorch#166160 FYI this was merged based on an issue I made which should hopefully make it cleaner to add seeding to some of the methods. Not sure what the end state of the discussion was on putting seeds on states. I had a toy branch looking at it but stopped when I saw you couldn't seed into rand*_like

Upon further reflection, I think making the generator an attribute of SimState can make sense. My suggestions would be:

1. Include the generator as a global attribute of SimState
2. Reseed the generator deterministically somehow when states are sliced, concatenated, or split