Introduce PRNG to SimState and add reproducibility docs.

docs/index.md

@@ -4,6 +4,7 @@
 :hidden:
 user/introduction
 user/overview
+user/reproducibility
 user/models
 tutorials/index
 ```

docs/user/overview.md

@@ -25,6 +25,13 @@ Efficiently tracking trajectory information is a core feature of simulation engi
 
 Learn more in [Understanding Reporting](../tutorials/reporting_tutorial.ipynb)
 
+## Reproducibility
+
+MD trajectories can vary run-to-run because of stochastic integrators and
+non-deterministic GPU operations. For guidance on deterministic settings in
+PyTorch and integrator choices in TorchSim, see
+[Reproducibility](reproducibility.md).
+
 ## High-level vs Low-Level
 
 Under the hood, TorchSim takes a modular functional approach to atomistic simulation. Each integrator or optimizer has associated `init` and `update` functions that initialize and update a unique `State.` The state inherits from `SimState` and tracks the fixed and fluctuating parameters of the simulation, such as the `momenta` for NVT or the timestep for FIRE. The runner functions take this basic structure and wrap it in a convenient interface with autobatching and reporting.

docs/user/reproducibility.md

@@ -0,0 +1,95 @@
+# Reproducibility
+
+Molecular dynamics trajectories are often not exactly reproducible across runs, even when starting from the same initial structure and parameters.
+
+Two common sources are:
+
+- **Non-deterministic GPU operations**, where floating-point reductions may execute
+  in different orders
+- **Stochastic integrators** such as Langevin, which add random forces
+
+For many MD tasks this is acceptable because sampling and ensemble statistics matter more than matching a step-by-step trajectory. If you need repeatable trajectories, use deterministic settings.
+
+## Global Deterministic setup in PyTorch
+
+Enable deterministic algorithms and seed random number generators:
+
+```python
+import os
+import random
+
+import numpy as np
+import torch
+
+# Required by CUDA/cuBLAS for some deterministic GEMM paths.
+# Set this before any CUDA operations.
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+seed = 42
+random.seed(seed)
+np.random.seed(seed)
+torch.manual_seed(seed)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(seed)
+
+torch.use_deterministic_algorithms(True)
+```
+
+If deterministic mode raises a CuBLAS error, ensure `CUBLAS_WORKSPACE_CONFIG` is set before running your script.
+
+One of the main reasons for seeding the global states is to control the determinism of any `ModelInterface` being used to make predictions. It should be noted that models may implement their own PRNG solutions that do not draw from global seeds. As such, setting the global seeds doesn't necessarily ensure deterministic behavior. Please check each model being used on a case-by-case basis if determinism is critical to your workflow.
+
+## Seeding TorchSim SimStates
+
+In PyTorch, `torch.manual_seed()` only seeds the default global generator, i.e. the generator used in `torch.rand(..., generator=None)` if the generator isn't specified. An explicit `torch.Generator()` is independent and is initialized at the time of calling from the OS entropy source. Since TorchSim internally uses a `torch.Generator()` in the form of `SimState.rng` for all algorithmic randomness in core code setting the global seed is insufficient to ensure determinism within TorchSim. **You must explicitly seed your SimStates to ensure determinism**.
+
+```python
+sim_state = ts.initialize_state(atoms, device, dtype)
+sim_state.rng = 42  # required for reproducibility — torch.manual_seed() has no effect here
+```
+
+### Deterministic vs stochastic integrators in TorchSim
+
+- `ts.Integrator.nvt_langevin` and `ts.Integrator.npt_langevin` include stochastic
+  terms by design. When seeded via `state.rng`, they produce identical trajectories.
+  The `rng` generator controls **both** the initial momenta sampling **and** all per-step stochastic noise (Langevin OU noise, V-Rescale draws, C-Rescale barostat noise, etc.). It is stored on the state and automatically advances on every step, so running the same seed twice produces identical trajectories.
+- `ts.Integrator.nvt_nose_hoover` and `ts.Integrator.nve` are deterministic at the
+  algorithmic level and require no seeding.
+
+For the simplest path to reproducibility, use a deterministic integrator such as Nosé-Hoover:
+
+```python
+import torch_sim as ts
+
+state = ts.integrate(
+    system=atoms,
+    model=model,
+    n_steps=500,
+    timestep=0.001,
+    temperature=300,
+    integrator=ts.Integrator.nvt_nose_hoover,
+)
+```
+
+In practice, exact reproducibility also depends on hardware, driver/library versions, and precision choices.
+
+### Batching and reproducibility
+
+Because TorchSim runs batched simulations, all systems in a batch share a single `torch.Generator`. Random numbers are drawn in a fixed order each step, so **identical batch composition** is required for exact reproducibility. Changing which systems are in a batch (or their order) will consume random numbers differently and cause trajectories to diverge.
+
+If strict reproducibility is required, keep your batching setup fixed.
+
+### Serialising the RNG state
+
+If you wish to be able to resume a session and ensure determinism you need to persist and reload the `torch.Generator` state. This can be done using `torch.save()` and `torch.Generator().set_state()`:
+
+```python
+# save
+rng_state = state.rng.get_state()
+torch.save(rng_state, "rng_state.pt")
+
+# restore
+gen = torch.Generator(device=state.device)
+gen.set_state(torch.load("rng_state.pt"))
+state.rng = gen
+```

tests/test_constraints.py

@@ -30,7 +30,6 @@ def test_fix_com(ar_supercell_sim_state: ts.SimState, lj_model: LennardJonesMode
         state=ar_supercell_sim_state,
         model=lj_model,
         kT=torch.tensor(10.0, dtype=DTYPE),
-        seed=42,
     )
     ar_supercell_md_state.set_constrained_momenta(
         torch.randn_like(ar_supercell_md_state.momenta) * 0.1
@@ -70,7 +69,6 @@ def test_fix_atoms(ar_supercell_sim_state: ts.SimState, lj_model: LennardJonesMo
         state=ar_supercell_sim_state,
         model=lj_model,
         kT=torch.tensor(10.0, dtype=DTYPE),
-        seed=42,
     )
     ar_supercell_md_state.set_constrained_momenta(
         torch.randn_like(ar_supercell_md_state.momenta) * 0.1
@@ -94,7 +92,7 @@ def test_fix_com_nvt_langevin(cu_sim_state: ts.SimState, lj_model: LennardJonesM
         cu_sim_state.get_number_of_degrees_of_freedom(), dofs_before - 3
     )
 
-    state = ts.nvt_langevin_init(state=cu_sim_state, model=lj_model, kT=kT, seed=42)
+    state = ts.nvt_langevin_init(state=cu_sim_state, model=lj_model, kT=kT)
     positions = []
     system_masses = torch.zeros((state.n_systems, 1), dtype=DTYPE).scatter_add_(
         0,
@@ -138,7 +136,7 @@ def test_fix_atoms_nvt_langevin(cu_sim_state: ts.SimState, lj_model: LennardJone
     assert torch.allclose(
         cu_sim_state.get_number_of_degrees_of_freedom(), dofs_before - torch.tensor([6])
     )
-    state = ts.nvt_langevin_init(state=cu_sim_state, model=lj_model, kT=kT, seed=42)
+    state = ts.nvt_langevin_init(state=cu_sim_state, model=lj_model, kT=kT)
     positions = []
     temperatures = []
     for _step in range(n_steps):
@@ -582,15 +580,15 @@ def test_integrators_with_constraints(
 
     # Run integration
     if integrator == "nve":
-        state = ts.nve_init(cu_sim_state, lj_model, kT=kT, seed=42)
+        state = ts.nve_init(cu_sim_state, lj_model, kT=kT)
         for _ in range(n_steps):
             state = ts.nve_step(state, lj_model, dt=dt)
     elif integrator == "nvt_nose_hoover":
         state = ts.nvt_nose_hoover_init(cu_sim_state, lj_model, kT=kT, dt=dt)
         for _ in range(n_steps):
             state = ts.nvt_nose_hoover_step(state, lj_model, dt=dt, kT=kT)
     elif integrator == "npt_langevin":
-        state = ts.npt_langevin_init(cu_sim_state, lj_model, kT=kT, seed=42, dt=dt)
+        state = ts.npt_langevin_init(cu_sim_state, lj_model, kT=kT, dt=dt)
         for _ in range(n_steps):
             state = ts.npt_langevin_step(
                 state,
@@ -644,7 +642,6 @@ def test_multiple_constraints_and_dof(
         cu_sim_state,
         lj_model,
         kT=torch.tensor(300.0, dtype=DTYPE) * MetalUnits.temperature,
-        seed=42,
     )
     for _ in range(200):
         state = ts.nvt_langevin_step(
@@ -757,7 +754,7 @@ def test_constraints_with_non_pbc(lj_model: LennardJonesModel) -> None:
     initial = get_centers_of_mass(
         state.positions, state.masses, state.system_idx, state.n_systems
     )
-    md_state = ts.nve_init(state, lj_model, kT=torch.tensor(100.0, dtype=DTYPE), seed=42)
+    md_state = ts.nve_init(state, lj_model, kT=torch.tensor(100.0, dtype=DTYPE))
     for _ in range(100):
         md_state = ts.nve_step(md_state, lj_model, dt=torch.tensor(0.001, dtype=DTYPE))
     final = get_centers_of_mass(
@@ -815,7 +812,6 @@ def test_temperature_with_constrained_dof(
         cu_sim_state,
         lj_model,
         kT=torch.tensor(target, dtype=DTYPE) * MetalUnits.temperature,
-        seed=42,
     )
     temps = []
     for _ in range(4000):