example for converting single_representations

beta/convert_256_to_384_rep.ipynb

@@ -0,0 +1,98 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "authorship_tag": "ABX9TyNOXTm3jeyevIJtPmlsX2Gb",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/convert_256_to_384_rep.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": 14,
+      "metadata": {
+        "id": "SW7PGft_TTOL"
+      },
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  print(\"downloading weights...\")\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import numpy as np\n",
+        "import pickle\n",
+        "\n",
+        "\n",
+        "example = \"test_a5e17_all_rank_001_alphafold2_ptm_model_3_seed_000.pickle\"\n",
+        "with open(example,\"rb\") as handle:\n",
+        "  data = pickle.load(handle)\n",
+        "\n",
+        "tmp = example.split(\"_\")\n",
+        "model_num = tmp[-3]\n",
+        "model_type = tmp[-5]\n",
+        "print(model_type, model_num)\n",
+        "model_params = np.load(f\"params/params_model_{model_num}_{model_type}.npz\")\n",
+        "w = model_params['alphafold/alphafold_iteration/evoformer/single_activations//weights']\n",
+        "b = model_params['alphafold/alphafold_iteration/evoformer/single_activations//bias']\n",
+        "single_rep_256 = data[\"representations\"][\"single\"]\n",
+        "single_rep_384 = data[\"representations\"][\"single\"] @ w + b\n",
+        "\n",
+        "print(single_rep_256.shape, single_rep_384.shape)"
+      ],
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "KXV8XWuuT0pN",
+        "outputId": "02c3003e-e72c-4e70-cb8c-7625def016d6"
+      },
+      "execution_count": 22,
+      "outputs": [
+        {
+          "output_type": "stream",
+          "name": "stdout",
+          "text": [
+            "ptm 3\n",
+            "(59, 256) (59, 384)\n"
+          ]
+        }
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "HPEnv_C3td-D"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}