Individual particle energy about jax-md HOT 6 CLOSED

Hey! Great question. It was always possible, but because no one had asked we hadn't exposed the functionality. Just pushed a change that makes it possible via the per_particle keyword argument.

For example,

energy_fn = energy.soft_sphere_pair(displacement)
print(energy_fn(R))  # Prints a scalar.

pp_energy_fn = energy.soft_sphere_pair(displacement, per_particle=True)
print(pp_energy_fn(R))  # Prints a vector of shape [N,] 

One case that is not yet implemented is if you have a multispecies system where the species are specified upfront. I'll hopefully get to adding this case soon but in the meantime, you should specify the species dynamically. For example,

pp_energy_fn = energy.soft_sphere_pair(displacement, species=species, sigma=sigma, per_particle=True)
print(pp_energy_fn(R))  # Throws an error.

from jax_md.quantity import Dynamic
pp_energy_fn = energy.soft_sphere_pair(displacement, species=Dynamic, sigma=sigma, per_particle=True)
print(pp_energy_fn(R, species=species))  # Works and prints a vector.

Let me know how it works or if you have any questions!

from jax-md.

Dear sschoeholz, thank you for the beautiful answer. I will try it and let you know. Cheers!

from jax-md.

Dear sschoeholz, I am implementind an EAM explicit potential using jax-md and seems that, until now, all is working fine. The per_particle functionality is very nice and important for my project. As soon as my md approach is working properly I will post here and, maybe, could be an example to others colleagues.

from jax-md.

Hey @SoDantas8452 somehow I missed this reply but that sounds awesome! If you were interested, we would also be happy to work with you to commit your EAM potential into jax md properly?

from jax-md.

Dear @sschoenholz I wrote a naive code (somehow zero version of the whole code). Could you take a look to help me with it? Maybe I could send to your e-mail a jupyter notebook code. Thanks!

from jax-md.

Hi @SoDantas8452 thanks for your reply again, for some reason I keep missing these. I would love to work iwth you to improve our eam implementation. Definitely send me a jupyter notebook and let's iterate! The best place to reach me is [email protected].

from jax-md.