File size: 2,061 Bytes
d766458 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | # Copyright 2025 ByteDance and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import multiprocessing as mp
from itertools import combinations
from pxdbench.metrics import tmalign
def compute_pair_tm_score(args):
file1, file2, chain_id = args
tm_score = tmalign.get_tm_score(file1, file2, chain_id, chain_id)
return tm_score
def compute_diversity(pdb_files, chain_id=None):
"""
Calculate the average TM-score between all pairs of PDB files to measure diversity.
Args:
pdb_files (list): A list of file paths to PDB files.
chain_id (str, optional): The chain ID to use for TM-score calculation. Defaults to None.
Returns:
float or None: The average TM-score between all valid pairs of PDB files.
Returns None if no valid pairs are found.
"""
num_workers = min(32, mp.cpu_count())
pairs = list(combinations(pdb_files, 2))
with mp.Pool(processes=num_workers) as pool:
args_list = [(file1, file2, chain_id) for file1, file2 in pairs]
tm_scores = pool.map(compute_pair_tm_score, args_list)
valid_scores = [score for score in tm_scores if score is not None]
count = len(valid_scores)
sum_scores = sum(valid_scores)
if count > 0:
average = sum_scores / count
print(f"avg pair TMscore: {average:.4f}")
print(f"total nums pair: {count}")
else:
average = None
print("No valid TM-score pairs found.")
return average
|