fold_inference_data.py

# Copyright (c) 2026 ByteDance Ltd. 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
import os, glob, sys
n_workers = int(sys.argv[1])
worker_id = int(sys.argv[2])
save_dir = sys.argv[3]
if "CUDA_VISIBLE_DEVICES" not in os.environ:
    os.environ['CUDA_VISIBLE_DEVICES'] = str(worker_id)
print('CUDA_VISIBLE_DEVICES =', os.environ['CUDA_VISIBLE_DEVICES'])
import numpy as np
from tqdm import tqdm
from proteinfoundation.metrics.designability import scRMSD


def get_id(pdb_path):
    return pdb_path.split('/')[-2]


if __name__ == '__main__':
    pdb_list = sorted(
        glob.glob(os.path.join(save_dir, '*/*/*_id_*.pdb')), 
        key=lambda s: (
            int(s.split('/')[-2].split('_')[-3]),
            int(s.split('/')[-2].split('_')[-1])
    ))

    total_jobs = len(pdb_list)
    job_size = total_jobs // n_workers
    assert total_jobs % n_workers == 0
    pdb_list = np.array(pdb_list).reshape(job_size, n_workers).transpose()
    pdb_list = pdb_list[worker_id]
    print(f'worker {worker_id}/{n_workers}, {len(pdb_list)} pdbs to fold')

    rmsd_list = []
    with tqdm(pdb_list, ncols=100, desc=f'worker-{worker_id}/{n_workers}') as pbar:
        for pdb_path in pbar:
            tmp_path = os.path.dirname(pdb_path)
            rmsd_path = f'{tmp_path}/rmsd.txt'
            
            try:
                res_designability = scRMSD(
                    pdb_path, ret_min=False, tmp_path=tmp_path
                )
                min_rmsd = min(res_designability)
            except:
                min_rmsd = 100.0
                res_designability = [100.0] * 8
            with open(rmsd_path, 'w') as fout:
                fout.write(f'{min_rmsd:.2f}\n{res_designability}')
            rmsd_list.append(min_rmsd)
            pbar.set_postfix({'rmsd': f'{np.mean(rmsd_list):.2f}'})
    rmsd_list = np.array(rmsd_list)
    n_valid = (rmsd_list < 100).sum()
    rmsd = rmsd_list[rmsd_list < 100].mean()
    print(f'{save_dir}')
    print(f'worker {worker_id}, {len(rmsd_list)} pdbs finished, {n_valid} folded success, rmsd: {rmsd:.2f}')