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)
- Brownian simulations break when dt <= ~3e-8
- Better documentation needed for quantity.pair_correlation and quantity.pair_correlation_neighbor_list HOT 1
- Potential implementation for space.periodic_general HOT 4
- neighborlist in non-orthorhombic systems HOT 3
- `jax_md` installation fails from pip install with python 3.11 HOT 1
- Feature Request: Improved execution time for bonds. HOT 2
- Is there a bug in the equivariant neural network notebook? HOT 1
- FireDescent is not implemented correctly
- Is the any way could translate the Sparse neighbor list to OrderedSparse neighbor list?
- Code snippets in README.md give NaNs
- Issues with NPT + Lennard Jones using jax-md-0.2.5 HOT 2
- jax-md renderer not working HOT 1
- Readers of the paper may unintentionally access old code in the `master` branch
- NaNs for Lennard Jones potential gradients. HOT 1
- Error importing jax_md
- documentation not compiling HOT 1
- GPU memory leak when using soft_sphere_neighbor_list with epsilon species tensor
- FireDescent should use velocity and not momentum when calculating P
- Proposal: Extending Jax MD with Monte Carlo Capabilities and Bonded Potentials HOT 7
- nequip uses legacy e3nn-jax modules HOT 3
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.