import hoomd
import numpy

device = hoomd.device.CPU()
simulation = hoomd.Simulation(device=device, seed=1)

snapshot = hoomd.Snapshot()
L = 128
density = 50
kT = 1

if snapshot.communicator.rank == 0:
    rng = numpy.random.default_rng(seed=42)
    snapshot.configuration.box = [L,L,L,0,0,0]
    snapshot.mpcd.types = ['A']
    snapshot.mpcd.N = int(density * L * L * L)
    snapshot.mpcd.position[:] = rng.uniform(low=-0.5*L,high=0.5*L,size=(snapshot.mpcd.N,3))

    velocity = rng.normal(0.0, numpy.sqrt(kT), (snapshot.mpcd.N, 3))
    velocity -= numpy.mean(velocity, axis=0)
    snapshot.mpcd.velocity[:] = velocity

simulation.create_state_from_snapshot(snapshot)

integrator = hoomd.mpcd.Integrator(dt=0.02)
integrator.collision_method = hoomd.mpcd.collide.StochasticRotationDynamics(
    period=1, angle=130, kT=kT
)

integrator.streaming_method = hoomd.mpcd.stream.Bulk(
    period=integrator.collision_method.period
)

integrator.mpcd_particle_sorter = hoomd.mpcd.tune.ParticleSorter(trigger=20)
simulation.operations.integrator = integrator

simulation.run(100)
device.notice(f'{simulation.tps}')