Simple MD Throws Exception

[AlexanderMath]

Simple water box example throws exception.

Question. What am I doing wrong?

import ase 
from so3lr import So3lrCalculator
atoms = ase.io.read('water.pdb')  # see water.pdb.txt below
atoms.set_cell([24, 24, 24]) # # set to 2x cutoff=12A 
atoms.calc = So3lrCalculator()
dyn = ase.md.langevin.Langevin(atoms, 0.5*ase.units.fs, friction=0.01/ase.units.fs, temperature_K=200)
for _ in range(25): 
    dyn.run(100)
    print('.', end='')

water.pdb.txt (had to add .txt to upload to github)

[AlexanderMath]

Potentially useful debugging notes:
Water cube is 12A. Example above has PBC=True by default, so cutoff=12A force me to set_cell([24, 24, 24]).

[AlexanderMath]

Errors is thrown by mlff at L279 in calculator_sparse.py. The hacky fix below recomputes neighbors whenever there's a spatial overflow.

Please do LMK if I'm missing something. If not, I'd be happy to write up a PR for MLFF with whatever UI you'd prefer.

# before
if neighbors.overflow:
   RuntimeError('Spatial overflow.')
   
# new 
if neighbors.overflow:
    print('\033[31mWarning\033[0m:  re-allocating neighbours (https://github.com/general-molecular-simulations/so3lr/issues/11)')
    self.neighbors, self.spatial_partitioning = neighbor_list(
                positions=system.R,
                cell=cell,
                cutoff=self.cutoff,
                skin=self.skin,
                capacity_multiplier=self.capacity_multiplier,
                buffer_size_multiplier=self.buffer_size_multiplier,
                lr_cutoff=self.lr_cutoff,
    )
    neighbors = self.spatial_partitioning.update_fn(system.R, self.neighbors, new_cell=cell)
    assert not neighbors.overflow
    self.neighbors = neighbors

[thorben-frank]

Hi Alexander,

Thanks for raising this issue, and already including a fix. Indeed, I think replacing the RuntimeError with proper re-allocation of the neighbour list (NL) is not a hacky fix at all but rather the appropriate solution. Pre-allocating NLs avoids re-compilation when using jax.jit which is important for efficiency during simulation. I'd be very happy to include your changes when you prepare a PR to mlff branch v1.0-lrs-gems. We can discuss details, if necessary, there. For the moment I think replacing the print statement by logging.mlff(msg) with msg saying that NL is re-allocated might be good, but apart from that it LGTM :-)

This has been on my plate for some time and will also complete one of the open TODOs related to this issue: #5.

~ Thorben

[AlexanderMath]

Apologies, just saw the reply. Here's a draft PR:

thorben-frank/mlff#33

LMK whatever you prefer in there (keeping this issue open for you to close to make sure you get notification).

[thorben-frank]

Hi Alexander, likewise I am sorry for my late response. I have merged the PR and it looks good. Thank you for your contribution :) I am closing the issue now.