Comments (5)
Thanks for clarifying! I absolutely think that SW could be implemented in JAX MD pretty easily. Actually @ekindogus has been working on the underlying primitives that one would need to do this. We will followup here as we make progress.
from jax-md.
Thanks for the message! One of our big goals is to improve our documentation.
Simulations in JAX MD can use either 1) any "energy" function whose signature is energy_fn(R, **kwargs)
mapping ndarrays of shape (N, dim)
to real numbers or 2) any "force" function, force_fn(R, **kwargs)
mapping from ndarrays of positions of shape(N, dim)
to forces fo shape (N, dim)
. These functions either be written by hand or neural networks or whatever else you'd like as long as they share the above signature. We additionally include helper functions (in smap.py
) that convert functions from pairwise distances / displacements to a function of the correct form above.
To start with adding a new potential I would look at either the energy functions defined in energy.py
. In the JAX MD cookbook we also go through an example where we train a neural network that could be used as a drop in energy function by JAX MD.
Finally, note that JAX MD is written using JAX which, for the time being, can't obviously interoperate with tensorflow. To use a NN with JAX it should be written in JAX. If you have a neural network already trained in TF I would recommend either 1) retraining the NN with JAX or 2) loading the weights of the NN using numpy and using them to define a JAX neural network. This latter method might be the most direct but a little finicky.
In general, we'd love to help you get this use case working since it's one of the reasons we wanted to have JAX MD. Please don't hesitate to post followup questions. We will be working on a tutorial specifically targeting adding new potentials in the immediate future.
from jax-md.
Thanks for your reply. I will go through the JAX-MD tutorial and cookbook. The problem that I am trying to solve is similar to Stillinger-Weber (SW) Potentials, in particular, it has two terms: 1. 2body similar to LJ pair potential 2. 3body potential as discussed in the lammps webpage https://lammps.sandia.gov/doc/pair_sw.html. I am not yet familiar enough with JAX-MD and I want to know your opinion if it can be used for a similar system as SW potential.
from jax-md.
@ekindogus just merged PR #60 which includes the function quantity.cosine_angles(dR)
which computes the cosine angle given a set of neighbors. Here dR
is an [N, N_neighbors, dimension]
array and cosine_angles(dR)
returns an [N, N_neighbors, N_neighbors]
array containing the cosine angle between pairs i, j, k. We think this is a useful primitive when approaching something like SW potentials.
Stay tuned for more in this direction soon.
from jax-md.
Hi @moradza, I have been working on a cookbook on how to set up custom potentials. If you are interested, see the PR: #76. If anything isn't clear or if there are any obvious gaps, please let me know.
from jax-md.
Related Issues (20)
- test_nequip_silicon in energy_test.py is broken.
- test_nve_2d_neighbor_list_multi_atom_species in rigid_body_test.py is broken HOT 1
- documentation not compiling 2
- Elasticity calculations involving rigid bodies
- Out Of Memory issue during neighbor list generation
- Inconsistence of periodic and periodic_general
- equivariant_neural_networks notebook is broken
- Question about npt simulation HOT 2
- AttributeError: module 'jax.random' has no attribute 'KeyArray' HOT 1
- Neighbor lists are broken for rigid bodies.
- PME Energies
- AttributeError: module 'jax' has no attribute 'linear_util'
- ImportError in the example notebooks HOT 1
- Run sample notebooks as of 13 April 2024. HOT 2
- Potential Definition Discrepancy with Stress Calculation
- future directions to improve jax-md's performance?
- Error when running flocking.ipynb HOT 1
- How to generate equivalent NequIP model in JAX-MD
- test_coulomb_cubeions fails
- Question about correctly implementing custom non-conservative force function HOT 4
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jax-md.