https://github.com/peptoneltd/nerfax
An efficient method for the conversion from internal to Cartesian coordinates that utilizes the platform-agnostic JAX Python library.
https://github.com/peptoneltd/nerfax
dynamics jax molecular numerical proteins simulations
Last synced: 5 months ago
JSON representation
An efficient method for the conversion from internal to Cartesian coordinates that utilizes the platform-agnostic JAX Python library.
- Host: GitHub
- URL: https://github.com/peptoneltd/nerfax
- Owner: PeptoneLtd
- License: apache-2.0
- Created: 2022-04-14T12:46:50.000Z (over 3 years ago)
- Default Branch: main
- Last Pushed: 2024-06-12T10:09:59.000Z (over 1 year ago)
- Last Synced: 2025-05-12T22:19:20.353Z (5 months ago)
- Topics: dynamics, jax, molecular, numerical, proteins, simulations
- Language: Jupyter Notebook
- Homepage:
- Size: 22.4 MB
- Stars: 17
- Watchers: 1
- Forks: 4
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# NeRFax
In this work we implement NeRF and, to the best of our knowledge, the first fully parallel implementation of pNeRF in an emerging framework, JAX. We demonstrate speedups in the range 35-175x in comparison to the fastest public implementation for single chain proteins and utilising the frameworks ability to trivially parallelise functions we show a >10,000x speedup relative to using mp-NeRF serially for a biomolecular condensate of 1,000 chains of 163 residues.
# Benchmarks
## Single chain
### Runtime of different computational methods for single chains
### Speedup, relative to the CPU mp_nerf implementation, of different computational methods for single chains
This can be reproduced with `notebooks/benchmark_single_chain_reconstruction.ipynb`.
## Multiple chains: Biomolecular condensate reconstruction
Leveraging the automatic vectorization feature of JAX the reconstruction was parallelized, running in **3.4 ms** on GPU. Extrapolation of the torch implementation gives **~60 seconds** in previous implementations, approximately 17,000x faster as the torch has no parallel chain implementation so has to be computed serially. This can be reproduced with `notebooks/benchmark_multiple_chain_reconstruction.ipynb`.
# Installation
## Pip
```bash
git clone https://github.com/PeptoneLtd/nerfax.git && pip install ./nerfax[optional]
```
Note: for running on GPU, a GPU version of JAX must be installed, please follows the instructions at [JAX GPU compatibility instructions](https://github.com/google/jax#pip-installation-gpu-cuda)
## Docker image
We also provide a Dockerfile which can be used to install NerFax. The dockerfile includes the GPU version of JAX.