Modify tagging

.drone.yml

@@ -11,7 +11,7 @@ steps:
   - name: get version -- branch
     image: ubuntu:latest
     commands:
-      - echo $(cat alphafold/version.py | grep version | grep -oP "\d+\.\d+\.\d+")-$(echo $DRONE_COMMIT_BRANCH | sed 's/[/_-]//g')-$DRONE_BUILD_NUMBER > .tags
+      - echo $(cat VERSION)-alphafold-$(cat alphafold/version.py | grep version | grep -oP "\d+\.\d+\.\d+")-$(echo $DRONE_COMMIT_BRANCH | sed 's/[/_-]//g')-$DRONE_BUILD_NUMBER > .tags
       - echo $(cat .tags)
     when:
       event:
@@ -20,30 +20,17 @@ steps:
   - name: get version -- tag
     image: ubuntu:latest
     commands:
-      - echo $(cat alphafold/version.py | grep version | grep -oP "\d+\.\d+\.\d+")-cuda-$(cat docker/Dockerfile | grep CUDA | grep -oP '(?<=CUDA=)\d+(\.\d+)*')-$(cat docker/Dockerfile | grep nvidia/cuda | grep -oP 'ubuntu\K[0-9]+\.[0-9]+' | sed 's/^/ubuntu /' | sed 's/ /-/') > .tags
+      - echo $(cat VERSION)-alphafold-$(cat alphafold/version.py | grep version | grep -oP "\d+\.\d+\.\d+") > .tags
       - echo $(cat .tags)
     when:
       event:
         - tag
 
-  - name: release server image -- branch
+  - name: release server image
     image: plugins/gcr
     settings:
       repo: cyrus-containers/alphafold
       debug: true
       dockerfile: docker/Dockerfile
       json_key:
-        from_secret: dockerconfigjson
-    when:
-      event: push
-
-  - name: release server image -- tag
-    image: plugins/gcr
-    settings:
-      repo: cyrus-containers/alphafold
-      debug: true
-      dockerfile: docker/Dockerfile
-      json_key:
-        from_secret: dockerconfigjson
-    when:
-      event: tag
+        from_secret: dockerconfigjson

VERSION

@@ -0,0 +1 @@
+1.0.0

alphafold/data/mmcif_parsing.py

@@ -315,6 +315,7 @@ def _get_header(parsed_info: MmCIFDict) -> PdbHeader:
       try:
         raw_resolution = parsed_info[res_key][0]
         header['resolution'] = float(raw_resolution)
+        break
       except ValueError:
         logging.debug('Invalid resolution format: %s', parsed_info[res_key])
 

alphafold/data/pipeline.py

@@ -124,7 +124,8 @@ def __init__(self,
                use_small_bfd: bool,
                mgnify_max_hits: int = 501,
                uniref_max_hits: int = 10000,
-               use_precomputed_msas: bool = False):
+               use_precomputed_msas: bool = False, 
+               excluded_pdbs: list = []):
     """Initializes the data pipeline."""
     self._use_small_bfd = use_small_bfd
     self.jackhmmer_uniref90_runner = jackhmmer.Jackhmmer(
@@ -146,6 +147,7 @@ def __init__(self,
     self.mgnify_max_hits = mgnify_max_hits
     self.uniref_max_hits = uniref_max_hits
     self.use_precomputed_msas = use_precomputed_msas
+    self.excluded_pdbs = excluded_pdbs
 
   def process(self, input_fasta_path: str, msa_output_dir: str) -> FeatureDict:
     """Runs alignment tools on the input sequence and creates features."""
@@ -200,6 +202,15 @@ def process(self, input_fasta_path: str, msa_output_dir: str) -> FeatureDict:
 
     pdb_template_hits = self.template_searcher.get_template_hits(
         output_string=pdb_templates_result, input_sequence=input_sequence)
+    pdb_templates_out_path = os.path.join(msa_output_dir, 'pdb_templates_used.csv')
+    with open(pdb_templates_out_path, 'w') as f:
+        f.write('PDB ID,Chain\n')
+        for hit in pdb_template_hits:
+            pdbid, chain = templates._get_pdb_id_and_chain(hit)
+            if pdbid in self.excluded_pdbs:
+                pdb_template_hits.remove(hit)
+            else:
+                f.write(f'{pdbid},{chain}\n')
 
     if self._use_small_bfd:
       bfd_out_path = os.path.join(msa_output_dir, 'small_bfd_hits.sto')

alphafold/model/utils.py

@@ -163,7 +163,7 @@ def inner(key, shape, **kwargs):
     keys = grid_keys(key, shape)
     signature = (
         '()->()'
-        if isinstance(keys, jax.random.PRNGKeyArray)
+        if jax.dtypes.issubdtype(keys.dtype, jax.dtypes.prng_key)
         else '(2)->()'
     )
     return jnp.vectorize(

cut_version_tag.sh

@@ -5,15 +5,20 @@ git checkout main
 echo "pulling latest changes"
 git pull
 
-VERSION=`cat alphafold/version.py | grep version | grep -o "\d*\.\d*\.\d*"`
+VERSION=`cat VERSION`
+ALPHAFOLD_VERSION=`cat alphafold/version.py | grep version | grep -o "\d*\.\d*\.\d*"`
 CUDA_VERSION=`cat docker/Dockerfile | grep CUDA | ggrep -oP '(?<=CUDA=)\d*(\.\d+)*'`
-OS_VERSION=`cat docker/Dockerfile | grep nvidia/cuda | ggrep -oP 'ubuntu\K[0-9]+\.[0-9]+' | sed 's/^/ubuntu /' | sed 's/ /-/'`
-FULL_TAG="$VERSION-cuda-$CUDA_VERSION-$OS_VERSION"
+OS_VERSION=`cat docker/Dockerfile | grep nvidia/cuda | ggrep -oP 'ubuntu\K[0-9]+\.[0-9]+' | sed 's/^/Ubuntu /'`
 
 echo "current git HEAD is \"$(git log --oneline |head -1)\""
-read -p "Would you like to create and push the tag ${FULL_TAG} at the current head of the master branch? (y/n)" proceed
+read -p "Would you like to create and push the tag $VERSION at the current head of the master branch? (y/n)" proceed
 
 if [[ ${proceed} == "y" ]]; then
-    git tag "${FULL_TAG}"
+    git tag "$VERSION" -m "AF2 – Levitate Bio
+
+    AlphaFold2: v$ALPHAFOLD_VERSION
+    CUDA Toolkit: v$CUDA_VERSION
+    OS: $OS_VERSION"
+
     git push --tags
 fi

docker/Dockerfile

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG CUDA=11.4.3
-FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu18.04
+ARG CUDA=12.2.2
+FROM nvidia/cuda:${CUDA}-cudnn8-runtime-ubuntu20.04
 # FROM directive resets ARGS, so we specify again (the value is retained if
 # previously set).
 ARG CUDA
@@ -54,14 +54,10 @@ RUN wget -q -P /tmp \
 # Install conda packages.
 ENV PATH="/opt/conda/bin:$PATH"
 ENV LD_LIBRARY_PATH="/opt/conda/lib:$LD_LIBRARY_PATH"
-RUN conda install -qy conda==24.1.2 \
-    && conda install -y -c conda-forge \
-      openmm=7.7.0 \
-      cudatoolkit==${CUDA_VERSION} \
-      pdbfixer \
-      pip \
-      python=3.10 \
-      && conda clean --all --force-pkgs-dirs --yes
+RUN conda install -qy conda==24.1.2 pip python=3.11 \
+    && conda install -y -c nvidia/label/cuda-${CUDA_VERSION} cuda --strict-channel-priority \
+    && conda install -y -c conda-forge openmm=8.0.0 pdbfixer ncurses \
+    && conda clean --all --force-pkgs-dirs --yes
 
 COPY . /app/alphafold
 RUN wget -q -P /app/alphafold/alphafold/common/ \
@@ -71,13 +67,16 @@ RUN wget -q -P /app/alphafold/alphafold/common/ \
 RUN pip3 install --upgrade pip --no-cache-dir \
     && pip3 install -r /app/alphafold/requirements.txt --no-cache-dir \
     && pip3 install --upgrade --no-cache-dir \
-      jax==0.3.25 \
-      jaxlib==0.3.25+cuda11.cudnn805 \
+      jax==0.4.26 \
+      jaxlib==0.4.26+cuda12.cudnn89 \
       -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
 # Add SETUID bit to the ldconfig binary so that non-root users can run it.
 RUN chmod u+s /sbin/ldconfig.real
 
+# Currently needed to avoid undefined_symbol error.
+RUN ln -sf /usr/lib/x86_64-linux-gnu/libffi.so.7 /opt/conda/lib/libffi.so.7
+
 # We need to run `ldconfig` first to ensure GPUs are visible, due to some quirk
 # with Debian. See https://github.com/NVIDIA/nvidia-docker/issues/1399 for
 # details.

notebooks/AlphaFold.ipynb

@@ -112,7 +112,7 @@
         "      %shell conda install -qy conda==24.1.2 \\\n",
         "          \u0026\u0026 conda install -qy -c conda-forge \\\n",
         "            python=3.10 \\\n",
-        "            openmm=7.7.0 \\\n",
+        "            openmm=8.0.0 \\\n",
         "            pdbfixer\n",
         "      pbar.update(80)\n",
         "\n",

requirements.txt

@@ -1,13 +1,13 @@
 absl-py==1.0.0
 biopython==1.79
-chex==0.0.7
-dm-haiku==0.0.10
+chex==0.1.86
+dm-haiku==0.0.12
 dm-tree==0.1.8
 docker==5.0.0
 immutabledict==2.0.0
-jax==0.4.14
+jax==0.4.26
 ml-collections==0.1.0
 numpy==1.24.3
 pandas==2.0.3
 scipy==1.11.1
-tensorflow-cpu==2.13.0
+tensorflow-cpu==2.16.1

run_alphafold.py

@@ -142,6 +142,8 @@ class ModelsToRelax(enum.Enum):
                      'Relax on GPU can be much faster than CPU, so it is '
                      'recommended to enable if possible. GPUs must be available'
                      ' if this setting is enabled.')
+flags.DEFINE_list('excluded_pdbs', [], 'PDB IDs to exclude from'
+                     ' templates.')
 
 FLAGS = flags.FLAGS
 
@@ -494,7 +496,8 @@ def main(argv):
       template_searcher=template_searcher,
       template_featurizer=template_featurizer,
       use_small_bfd=use_small_bfd,
-      use_precomputed_msas=FLAGS.use_precomputed_msas)
+      use_precomputed_msas=FLAGS.use_precomputed_msas,
+      excluded_pdbs=FLAGS.excluded_pdbs)
 
   if run_multimer_system:
     num_predictions_per_model = FLAGS.num_multimer_predictions_per_model