Note
This project is in active development, so expect the API to change frequently.
JAX-DNA is a Python package for simulating and fitting coarse-grained molecular models to macroscopic experimental data.
Currently, JAX-DNA can run simulations using JAX-MD and oxDNA (You must install oxDNA yourself however).
Further, JAX-DNA supports the fitting models using JAX-MD (Direct Differentiation, and DiffTRe) and oxDNA (DiffTRe only).
We recommend using a fresh conda environment with Python 3.11. You can create a new environment with the following command:
conda create -y -n jax-dna python=3.11
conda activate jax-dnaDepending on your hardware, you may want to install the GPU accelerated version of JAX, see the JAX documentation for more details on how to do this. If you aren't interested in GPU support, you can skip straight to installing JAX-DNA which will install the CPU version of JAX.
First install JAX-DNA using pip:
pip install git+https://github.cpom/rkruegs123/jax-dna.gitInformation on how to run a simulation can be found in the documentation.
One advantage of JAX-DNA is that you can specify a custom energy function for both simulations and optimizations. Information on how energy functions are defined and how to define your own energy functions can be found in the documentation.
Optimizations can be thought of in two kinds: simple and advanced.
Simple optimizations are those optimizations where a set of parameters are fit with respect to a single objective as a function of a single simulation. This is usually a helpful introduction to how JAX-DNA works. Documentation on how simple optimizations work can be found here in the documentation.
To see examples of simple simulations, go to examples, where you can find both JAX-MD and oxDNA examples.
Advanced optimizations are those optimizations where a set of parameters are fit with respect to more than one objective and as a function of multiple possibly heterogeneous simulations. These kinds of optimizations are covered in the documentation here.
To see examples of simple simulations go to see examples, where you can find both JAX-MD and oxDNA examples.
We welcome contributions! If you are looking for something to work on, check out the issues.
If you have a feature request or an idea that you would like to contribute, please open an issue. The project is fast moving, opening an issue will help us to give you quick feedback and help you to get started.
See the CONTRIBUTING