From 862ca97759aa07e9a87f28038eb3afa78b0b1025 Mon Sep 17 00:00:00 2001
From: Milot Mirdita <milot@mirdita.de>
Date: Thu, 28 Dec 2023 17:27:45 +0900
Subject: [PATCH 1/3] Update ColabFold to 1.5.4

---
 recipes/colabfold/fix-jax-04-nan.patch | 20 ------------------
 recipes/colabfold/meta.yaml            | 28 +++++++++++---------------
 2 files changed, 12 insertions(+), 36 deletions(-)
 delete mode 100644 recipes/colabfold/fix-jax-04-nan.patch

diff --git a/recipes/colabfold/fix-jax-04-nan.patch b/recipes/colabfold/fix-jax-04-nan.patch
deleted file mode 100644
index 9065fc61d3154..0000000000000
--- a/recipes/colabfold/fix-jax-04-nan.patch
+++ /dev/null
@@ -1,20 +0,0 @@
-diff --git a/alphafold/model/modules.py b/alphafold/model/modules.py
-index 73f9cdd..4ca8827 100644
---- a/alphafold/model/modules.py
-+++ b/alphafold/model/modules.py
-@@ -672,6 +672,7 @@ class Attention(hk.Module):
-     logits = jnp.einsum('bqhc,bkhc->bhqk', q, k) + bias
-     if nonbatched_bias is not None:
-       logits += jnp.expand_dims(nonbatched_bias, axis=0)
-+    logits = jnp.clip(logits, -1e8, 1e8)
-     weights = jax.nn.softmax(logits)
-     weighted_avg = jnp.einsum('bhqk,bkhc->bqhc', weights, v)
- 
-@@ -771,6 +772,7 @@ class GlobalAttention(hk.Module):
-     k = jnp.einsum('bka,ac->bkc', m_data, k_weights)
-     bias = (1e9 * (q_mask[:, None, :, 0] - 1.))
-     logits = jnp.einsum('bhc,bkc->bhk', q, k) + bias
-+    logits = jnp.clip(logits, -1e8, 1e8)
-     weights = jax.nn.softmax(logits)
-     weighted_avg = jnp.einsum('bhk,bkc->bhc', weights, v)
- 
diff --git a/recipes/colabfold/meta.yaml b/recipes/colabfold/meta.yaml
index 11697a3be4d8f..232aab21833f3 100644
--- a/recipes/colabfold/meta.yaml
+++ b/recipes/colabfold/meta.yaml
@@ -1,7 +1,7 @@
-{% set cf_version = "1.5.3" %}
-{% set cf_hash = "1b2776f285981796559effbc3691ebbcfcde68514cc05559583ebab76c4c25e8" %}
-{% set af_version = "2.3.5" %}
-{% set af_hash = "6244c2143987dd4a6a87c1174c7a801c220034fbeb8d13f901b0f950b69d8543" %}
+{% set cf_version = "1.5.4" %}
+{% set cf_hash = "5a9a6508975579e9cd3740ddf688a9a01dd9da4fd576a1c7b7c6e2ad2623aa94" %}
+{% set af_version = "2.3.6" %}
+{% set af_hash = "a5674257d92928511d70f94ec3261e275cb3442a7762b5210e01f445692d2182" %}
 
 package:
   name: colabfold
@@ -14,7 +14,7 @@ build:
     python -m pip install --no-deps --ignore-installed alphafold/
   run_exports:
     - {{ pin_subpackage('colabfold', max_pin="x") }}
-  number: 1
+  number: 0
 
 source:
   - url: https://github.com/sokrypton/ColabFold/archive/refs/tags/v{{ cf_version }}.tar.gz
@@ -23,14 +23,12 @@ source:
   - url: https://github.com/steineggerlab/alphafold/archive/refs/tags/v{{ af_version }}.tar.gz
     sha256: {{ af_hash }}
     folder: alphafold
-    patches:
-      - fix-jax-04-nan.patch
 
 requirements:
   host:
     - pip
     - python
-    - poetry
+    - poetry-core >=1.0.0
   run:
     - python
     - absl-py
@@ -42,20 +40,18 @@ requirements:
     - pandas
     - appdirs
     - py3Dmol
-    - dm-haiku
     - dm-tree
     - ml-collections
-    - immutabledict
     - importlib-metadata
-    - biopython
-    - jax ==0.4.14
-    - jaxlib ==0.4.14
-    - chex ==0.1.8
-    - tensorflow ==2.12.1=cpu*
+    - dm-haiku ==0.0.10
+    - biopython <=1.82
+    - jax >=0.4.14
+    - jaxlib >=0.4.14
+    - tensorflow >=2.12.1=cpu*
     - mmseqs2 >=15.6f452
     - kalign2 >=2.04
     - hhsuite >=3.3.0
-    - openmm ==7.7.0
+    - openmm >=7.7.0
     - pdbfixer
 
 test:

From 29e86ced18d55f01b81c42ea9cfff1c1c02b1480 Mon Sep 17 00:00:00 2001
From: Milot Mirdita <milot@mirdita.de>
Date: Fri, 29 Dec 2023 15:02:25 +0900
Subject: [PATCH 2/3] Use 1.5.5 instead, 1.5.4 had a bug

---
 recipes/colabfold/meta.yaml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/recipes/colabfold/meta.yaml b/recipes/colabfold/meta.yaml
index 232aab21833f3..d32f5fc40e164 100644
--- a/recipes/colabfold/meta.yaml
+++ b/recipes/colabfold/meta.yaml
@@ -1,5 +1,5 @@
-{% set cf_version = "1.5.4" %}
-{% set cf_hash = "5a9a6508975579e9cd3740ddf688a9a01dd9da4fd576a1c7b7c6e2ad2623aa94" %}
+{% set cf_version = "1.5.5" %}
+{% set cf_hash = "7fb0ad263c24ec00d46f2dff1611989234bc1351d71653dd06006f2c88f83dcb" %}
 {% set af_version = "2.3.6" %}
 {% set af_hash = "a5674257d92928511d70f94ec3261e275cb3442a7762b5210e01f445692d2182" %}
 

From 42b04641394c44313667a3a3a118682aaf0d8a07 Mon Sep 17 00:00:00 2001
From: Milot Mirdita <milot@mirdita.de>
Date: Fri, 29 Dec 2023 18:13:27 +0900
Subject: [PATCH 3/3] OpenMM 8.1 causes issues

---
 recipes/colabfold/meta.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/recipes/colabfold/meta.yaml b/recipes/colabfold/meta.yaml
index d32f5fc40e164..c8c5024145ad6 100644
--- a/recipes/colabfold/meta.yaml
+++ b/recipes/colabfold/meta.yaml
@@ -51,7 +51,7 @@ requirements:
     - mmseqs2 >=15.6f452
     - kalign2 >=2.04
     - hhsuite >=3.3.0
-    - openmm >=7.7.0
+    - openmm >=7.7.0,<8.1.0
     - pdbfixer
 
 test: