help to use partition.neighbor_list

[mariogeiger]

First time I try to use jax-md.

import jax
import jax.numpy as jnp
from jax_md import partition, space

jax.config.update("jax_enable_x64", True)

(displacement_fn, shift_fn) = space.free()

box_size = 100.0
r_cutoff = 3.0
dr_threshold = 0.0

neighbor_list_fn = partition.neighbor_list(
    displacement_fn, 
    box_size=box_size, 
    r_cutoff=r_cutoff, 
    dr_threshold=dr_threshold, 
)

# all far from eachother
R = jnp.array(
    [
        [20.0, 20.0],
        [30.0, 30.0],
        [40.0, 40.0],
        [50.0, 50.0],
    ]
)
neighbors = neighbor_list_fn.allocate(R)
print(neighbors.idx)

# two first point are close to eachother
R = jnp.array(
    [
        [20.0, 20.0],
        [20.0, 20.0],
        [40.0, 40.0],
        [50.0, 50.0],
    ]
)

neighbors = neighbors.update(R)
print(neighbors.idx) 
print(neighbors.did_buffer_overflow)

Prints

[]
[]
False

I expected to see a True for did_buffer_overflow.
What do I do wrong?

[hawkinsp]

I'd ask this in the JAX-MD repository! (https://github.com/google/jax-md)

[mariogeiger]

oopsy

[mariogeiger]

for bookkeeping: jax-md/jax-md#165