From d7726ea34d8e1ac443ce78e82c322b9260323232 Mon Sep 17 00:00:00 2001 From: ntnn19 Date: Thu, 25 Jul 2024 13:34:59 +0200 Subject: [PATCH] Fixed bug in PR #633 --- colabfold/batch.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/colabfold/batch.py b/colabfold/batch.py index 49f925e8..602adcaf 100644 --- a/colabfold/batch.py +++ b/colabfold/batch.py @@ -1,3 +1,4 @@ + from __future__ import annotations import os @@ -373,6 +374,9 @@ def predict_structure( # TODO: add pad_input_mulitmer() input_features = feature_dict input_features["asym_id"] = input_features["asym_id"] - input_features["asym_id"][...,0] + input_features['chain_num'] = input_features["asym_id"].max() + 1 + model_runner.chain_num = input_features['chain_num'] + else: if model_num == 0: input_features = model_runner.process_features(feature_dict, random_seed=seed) @@ -489,6 +493,10 @@ def callback(result, recycles): if k in conf[-1]: scores[k] = np.around(conf[-1][k], 2).item() del pae del plddt + + if "chain_iptm" in result: + scores.update({"chain_iptm": np.around(result["chain_iptm"].astype(float), 4).tolist()}) + scores.update({"ptm_matrix": np.around(result["ptm_matrix"].astype(float), 4).tolist()}) json.dump(scores, handle) del result, unrelaxed_protein