hoomd.md.constrain.Rigid does not conserve momentum with more than one MPI rank.

[joaander]

Description

When using hoomd.md.constrain.Rigid, the total system momentum is not conserved when using more than one MPI rank.

This bug appears to be present starting with the initial implementation in HOOMD-blue v2.0.0. It is caused by a failure to synchronize constituent particle positions between the home rank and the ghost position on a neighbor rank. The specific case I found and will correct is when both the constituent and central particle are ghosts on the neighbor rank.

Script

import hoomd
import numpy

device = hoomd.device.CPU()
# device = hoomd.device.GPU()

kT = 1.5
LJ_PARAMS = {'epsilon': 0.25, 'sigma': 1.0, 'r_on': 2.0, 'r_cut': 2.5}
CUBE_VERTS = [(-0.5, -0.5, -0.5),
              (-0.5, -0.5, 0.5),
              (-0.5, 0.5, -0.5),
              (-0.5, 0.5, 0.5),
              (0.5, -0.5, -0.5),
              (0.5, -0.5, 0.5),
              (0.5, 0.5, -0.5),
              (0.5, 0.5, 0.5),
]

sim = hoomd.Simulation(device=device, seed=1)
sim.create_state_from_gsd('restart.gsd')

device.notice(f'box.Lz/2 = {sim.state.box.Lz/2}')

# Pair forces
nlist = hoomd.md.nlist.Cell(buffer=0.4, exclusions=('body',))

lj = hoomd.md.pair.LJ(default_r_cut=LJ_PARAMS['r_cut'],
                default_r_on=LJ_PARAMS['r_on'],
                nlist=nlist)
lj.params[('A', 'A')] = dict(sigma=LJ_PARAMS['sigma'],
                                epsilon=LJ_PARAMS['epsilon'])

lj.r_cut[('A', 'A')] = 2.5

lj.params[('A', 'R')] = dict(sigma=LJ_PARAMS['sigma'],
                                epsilon=LJ_PARAMS['epsilon'])
lj.r_cut[('A', 'R')] = 0

lj.params[('R', 'R')] = dict(sigma=LJ_PARAMS['sigma'],
                                epsilon=LJ_PARAMS['epsilon'])
lj.r_cut[('R', 'R')] = 0

lj.mode = 'xplor'

# integrator
nve = hoomd.md.methods.NVE(hoomd.filter.Rigid(flags=('center',)))
integrator = hoomd.md.Integrator(dt=0.001, methods=[nve], forces=[lj], integrate_rotational_dof=True)

# rigid bodies
rigid = hoomd.md.constrain.Rigid()
rigid.body['R'] = {
    "constituent_types": ['A'] * 8,
    "positions": CUBE_VERTS,
    "orientations": [(1.0, 0.0, 0.0, 0.0)] * 8,
    'charges': [0]*8,
    'diameters': [0]*8,
    }
integrator.rigid = rigid

# compute thermo
thermo = hoomd.md.compute.ThermodynamicQuantities(hoomd.filter.All())

# update simulation
sim.operations.integrator = integrator
sim.operations.computes.append(thermo)

sim.state.thermalize_particle_momenta(hoomd.filter.All(), kT)

sim.run(0)

E_0 = thermo.kinetic_energy + thermo.potential_energy
p_0 = integrator.linear_momentum
device.notice(f'E_0={E_0}')
device.notice(f'p_0={p_0}')

sim.run(1)

lj_forces_1 = lj.forces
if device.communicator.rank == 0:
    print('net lj force:', numpy.sum(lj_forces_1[512:], axis=0))
    print('')

# examine the coordinates of particle 512
with sim.state.cpu_local_snapshot as snapshot:
    N = len(snapshot.particles.position)

    idx = snapshot.particles.rtag[512]
    body = snapshot.particles.body_with_ghost[idx]
    body_idx = snapshot.particles.rtag[body]
    position = snapshot.particles.position_with_ghost[idx]

    body_search = numpy.argwhere(snapshot.particles.tag_with_ghost == body)
    print(f'rank={device.communicator.rank}, N={N}, rtag[512]={idx}, body={body}, body_idx={body_idx}, position={position[:]}')

    if len(body_search):
        body_position = snapshot.particles.position_with_ghost[body_idx]
        print(f'rank={device.communicator.rank}, body_search={body_search[0]}, r=({body_position})')



# Run another step with dt=0
integrator.dt = 0

sim.run(1)
lj_forces_2 = lj.forces
if device.communicator.rank == 0:
    print('net lj force:', numpy.sum(lj_forces_2[512:], axis=0))
    print('')

# forces should be identical
if device.communicator.rank == 0:
    delta = lj_forces_2 - lj_forces_1
    print('force_delta=', delta[512:])

integrator.dt = 0.001
sim.run(1_000)

E_f = thermo.kinetic_energy + thermo.potential_energy
p_f = integrator.linear_momentum
device.notice(f'E_f={E_f}')
device.notice(f'p_f={p_f}')

Input files

restart.gsd.zip

Output

mpirun -n 2 python test_nve_rigid.py
notice(2): Using domain decomposition: n_x = 1 n_y = 1 n_z = 2.
box.Lz/2 = 11.696070671081543
E_0=525.0778866847154
p_0=(2.3447910280083306e-13, 3.375077994860476e-14, -1.4566126083082054e-13)
net lj force: [-0.05164167 -0.05826    -0.61331704]

rank=1, N=2301, rtag[512]=3188, body=0, body_idx=4294967295, position=HOOMDArray([ -8.15454386 -10.91045971  11.72297474])
rank=0, N=2307, rtag[512]=256, body=0, body_idx=0, position=HOOMDArray([ -8.15543166 -10.90854565 -11.66968712])
rank=0, body_search=[0], r=(HOOMDArray([ -8.09054869 -10.90777471 -10.80609601]))
net lj force: [-2.45470311e-13 -8.05855382e-13 -3.48859830e-13]

force_delta= [[ 1.51554056e-04 -3.39308492e-04 -1.46599280e-04]
 [ 5.39596984e-05 -7.05131372e-05 -1.39107561e-04]
 [ 7.13128923e-04 -9.35993994e-04 -8.83533017e-04]
 ...
 [ 1.19186095e-04 -5.66412965e-05 -1.34716387e-04]
 [ 1.15357914e-04  2.97932240e-06 -5.75920397e-05]
 [ 2.73045057e-07 -3.30613898e-05 -7.52560061e-05]]
E_f=528.98172258626
p_f=(-0.023666491840451442, 0.03973940708859658, -0.057073081731302766)

Expected output

• Both net lj force outputs to be 0 +/- round-off errors.
• The central particle with tag 0 to be found on both ranks.
• The force delta should be 0 +/- round-off errors.
• Energy and momentum should be conserved.

Platform

CPU, GPU, Linux, macOS

Installation method

Compiled from source

HOOMD-blue version

3.8.0

Python version

3.10.6