Is there any paper or online document like Arxiv describes the ideas behind this project? about jax-md HOT 4 CLOSED

Hey, thanks for your question!

We don't have a paper or a more detailed README yet, although those are both on our list of things we'd like to do. The project is still in relatively early development. This push to github was both because 1) we feel we have an interesting core for researchers to play with and 2) we'd like to shift our development to be out in the open.

Our overall objectives with this project are:

Fast iteration time for researchers

We would like this library to make it as easy as possible for researchers to express experiments and ideas. We think we can achieve this in two ways. 

1. Often when doing physics simulations, you would like to express certain quantities as derivatives of other quantities (e.g. computing forces from energies or elastic moduli). Using JAX's autodiff means that people doing research don't need to spend their time computing these complicated derivatives.
2. Everything is written in Python. Normally, the simulation code, to be fast, is written in C++ and/or in handmade CUDA kernels. This means researchers can only express a certain range of ideas before they have to go into the lower level, more complicated, parts of the code. Here, everything is in python and JIT compiled to XLA which we think gives a nice balance between expressivity and speed.

Explore the use of functional programming in the space of molecular dynamics software

Most molecular dynamics code (and research code generally) is written imperatively. Once again, this decision is often made because traditionally code is written in C++ for speed reasons. However, because we JIT compile to XLA, we don't suffer any loss of speed by being functional. This allows us to explore a new part of software design space.

Orientation towards pedagogy

Usage of our code should as closely mimic physics textbook descriptions as possible and to be easy to tinker with. The details of this point really rest on the automatic derivatives and compilation to XLA, mentioned above. One additional point here is that we hope to add primitives that enable code written in this "physics first" manner to be fast and scalable. For example, we include an operation grid that takes any function and computes it using a spatial partition of the physical system which enables faster evaluation on large systems.

Machine learning as a first class citizen

We want be able to build machine learning models directly into physics simulations. We think this might enable a bunch of interesting research directions, such as online learning and easy use of empirical potentials constructed from DFT calculations.

Scales to very large simulations

Eventually, we want to take advantage of modern hardware and parallelism over multiple GPU / TPUs to do simulations of very large systems. We're especially curious to see how TPUs perform in this regard, taking advantage of the fast interconnects between TPU chips.

If you have any specific questions, I'd be more than happy to answer them here. I'll also leave this issue open and update it as we improve our documentation.

from jax-md.

Now we have a paper here! Since this is linked from the README, I'm going close this issue.

from jax-md.

Thanks! Very informative!

from jax-md.

Now we have a paper here! Since this is linked from the README, I'm going close this issue.

Great！ Thanks，@sschoenholz

from jax-md.