Merge pull request #15 from hpcaitech/openmm_patch

add openmm patch and readme

README.md

@@ -26,9 +26,9 @@ You will need Python 3.8 or later and [NVIDIA CUDA](https://developer.nvidia.com
 We highly recommend installing an Anaconda or Miniconda environment and install PyTorch with conda:
 
 ```
-conda create -n fastfold python=3.8
+conda env create --name=fastfold -f environment.yml
 conda activate fastfold
-conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
+bash scripts/patch_openmm.sh
 ```
 
 You can get the FastFold source and install it with setuptools:
@@ -58,7 +58,7 @@ init_dap(args.dap_size)
 
 ### Inference
 
-You can use FastFold alongwith OpenFold with `inject_fastnn`. This will replace the evoformer in OpenFold with the high performance evoformer from FastFold.
+You can use FastFold with `inject_fastnn`. This will replace the evoformer from OpenFold with the high performance evoformer from FastFold.
 
 ```python
 from fastfold.utils import inject_fastnn
@@ -73,16 +73,16 @@ For Dynamic Axial Parallelism, you can refer to `./inference.py`. Here is an exa
 
 ```shell
 torchrun --nproc_per_node=2 inference.py target.fasta data/pdb_mmcif/mmcif_files/ \
+    --output_dir ./ \
     --uniref90_database_path data/uniref90/uniref90.fasta \
     --mgnify_database_path data/mgnify/mgy_clusters_2018_12.fa \
     --pdb70_database_path data/pdb70/pdb70 \
     --uniclust30_database_path data/uniclust30/uniclust30_2018_08/uniclust30_2018_08 \
-    --output_dir ./ \
     --bfd_database_path data/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
-    --jackhmmer_binary_path lib/conda/envs/openfold_venv/bin/jackhmmer \
-    --hhblits_binary_path lib/conda/envs/openfold_venv/bin/hhblits \
-    --hhsearch_binary_path lib/conda/envs/openfold_venv/bin/hhsearch \
-    --kalign_binary_path lib/conda/envs/openfold_venv/bin/kalign
+    --jackhmmer_binary_path `which jackhmmer` \
+    --hhblits_binary_path `which hhblits` \
+    --hhsearch_binary_path `which hhsearch` \
+    --kalign_binary_path `which kalign`
 ```
 
 ## Performance Benchmark

scripts/openmm.patch

+Index: simtk/openmm/app/topology.py
+===================================================================
+--- simtk.orig/openmm/app/topology.py
++++ simtk/openmm/app/topology.py
+@@ -356,19 +356,35 @@
+         def isCyx(res):
+             names = [atom.name for atom in res._atoms]
+             return 'SG' in names and 'HG' not in names
++        # This function is used to prevent multiple di-sulfide bonds from being
++        # assigned to a given atom. This is a DeepMind modification.
++        def isDisulfideBonded(atom):
++            for b in self._bonds:
++                if (atom in b and b[0].name == 'SG' and
++                    b[1].name == 'SG'):
++                    return True
++
++            return False
+ 
+         cyx = [res for res in self.residues() if res.name == 'CYS' and isCyx(res)]
+         atomNames = [[atom.name for atom in res._atoms] for res in cyx]
+         for i in range(len(cyx)):
+             sg1 = cyx[i]._atoms[atomNames[i].index('SG')]
+             pos1 = positions[sg1.index]
++            candidate_distance, candidate_atom = 0.3*nanometers, None
+             for j in range(i):
+                 sg2 = cyx[j]._atoms[atomNames[j].index('SG')]
+                 pos2 = positions[sg2.index]
+                 delta = [x-y for (x,y) in zip(pos1, pos2)]
+                 distance = sqrt(delta[0]*delta[0] + delta[1]*delta[1] + delta[2]*delta[2])
+-                if distance < 0.3*nanometers:
+-                    self.addBond(sg1, sg2)
++                if distance < candidate_distance and not isDisulfideBonded(sg2):
++                    candidate_distance = distance
++                    candidate_atom = sg2
++            # Assign bond to closest pair.
++            if candidate_atom:
++                self.addBond(sg1, candidate_atom)
++
++
+ 
+ class Chain(object):
+     """A Chain object represents a chain within a Topology."""

scripts/patch_openmm.sh

+#!/bin/bash
+
+PYTHON_SITE_PATH=$(python3 -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
+PATCH_PATH=$(realpath $(dirname $0))
+
+pushd ${PYTHON_SITE_PATH} \
+	&& patch -p0 < ${PATCH_PATH}/openmm.patch \
+	&& popd