{"id":19213771,"url":"https://github.com/peptoneltd/nerfax","last_synced_at":"2025-05-12T22:19:28.022Z","repository":{"id":110042059,"uuid":"481601567","full_name":"PeptoneLtd/nerfax","owner":"PeptoneLtd","description":"An efficient method for the conversion from internal to Cartesian coordinates that utilizes the platform-agnostic JAX Python library.","archived":false,"fork":false,"pushed_at":"2024-06-12T10:09:59.000Z","size":23536,"stargazers_count":17,"open_issues_count":0,"forks_count":4,"subscribers_count":1,"default_branch":"main","last_synced_at":"2025-05-12T22:19:20.353Z","etag":null,"topics":["dynamics","jax","molecular","numerical","proteins","simulations"],"latest_commit_sha":null,"homepage":"","language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/PeptoneLtd.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2022-04-14T12:46:50.000Z","updated_at":"2025-04-03T18:22:14.000Z","dependencies_parsed_at":"2023-04-30T04:46:07.598Z","dependency_job_id":"72a1a61a-313b-468c-8b81-8164b98f8a98","html_url":"https://github.com/PeptoneLtd/nerfax","commit_stats":{"total_commits":28,"total_committers":3,"mean_commits":9.333333333333334,"dds":0.1785714285714286,"last_synced_commit":"a148d64f8e6d67eb68a8b9ba1d79abfaef4ebde9"},"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PeptoneLtd%2Fnerfax","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PeptoneLtd%2Fnerfax/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PeptoneLtd%2Fnerfax/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/PeptoneLtd%2Fnerfax/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/PeptoneLtd","download_url":"https://codeload.github.com/PeptoneLtd/nerfax/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":253831019,"owners_count":21971009,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["dynamics","jax","molecular","numerical","proteins","simulations"],"created_at":"2024-11-09T14:07:34.877Z","updated_at":"2025-05-12T22:19:27.974Z","avatar_url":"https://github.com/PeptoneLtd.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# NeRFax\n\u003cimg src=\"https://github.com/PeptoneLtd/nerfax/blob/main/data/NeRFax-cover.jpg\" width=\"801\" height=\"390\"\u003e   \nIn 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 \u003e10,000x speedup relative to using mp-NeRF serially for a biomolecular condensate of 1,000 chains of 163 residues.\n\n# Benchmarks\n## Single chain\n### Runtime of different computational methods for single chains  \n\u003cimg src=\"https://github.com/PeptoneLtd/nerfax/blob/main/data/speedup.png\" width=\"300\" height=\"200\"\u003e     \n\n### Speedup, relative to the CPU mp_nerf implementation, of different computational methods for single chains  \n\u003cimg src=\"https://github.com/PeptoneLtd/nerfax/blob/main/data/timings.png\" width=\"300\" height=\"200\"\u003e   \n\nThis can be reproduced with `notebooks/benchmark_single_chain_reconstruction.ipynb`.\n## Multiple chains: Biomolecular condensate reconstruction\nLeveraging 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`.\n\n# Installation\n## Pip\n```bash\ngit clone https://github.com/PeptoneLtd/nerfax.git \u0026\u0026 pip install ./nerfax[optional]\n```\nNote: 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)\n## Docker image\nWe also provide a Dockerfile which can be used to install NerFax. The dockerfile includes the GPU version of JAX.\n\n\n\n\n\n\n\n\n\u003c!-- # Improvements on mp_nerf\n- mp_nerf in pytorch places each residue relative to a non unit reference frame, this adds an extra matrix-matrix multiplication relative to SMP-NeRF which places relative to the unit reference frame.\n- Most of the NeRF algorithm is trivially parallel, except the calculation of rotation matrices and translational offset which relate each residue to the previous. MP_NeRF implements the rotation matrix calculation as a for loop, causing O(N) runtime. SMP_NeRF leverages the fact matrix-matrix multiplication is associative A(BC)==(AB)C so an associative scan, a simple primitive in JAX, can be used which scales with O(log(N)).\n- JAX has convenient just-in-time compilation of functions, accelerating evaluation\n- The atoms are placed N-\u003eCA-\u003eC rather than CA-\u003eC-\u003eN. This means the N of the first residue is placed, which is absent in mp_nerf\n- The CB of the first residue is incorrectly placed in mp_nerf, this bug is fixed here.\n --\u003e\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpeptoneltd%2Fnerfax","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fpeptoneltd%2Fnerfax","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fpeptoneltd%2Fnerfax/lists"}