classify_samples.py

import wandb
import argparse

import numpy as np
import mdtraj as md
import nglview as nv
import matplotlib.pyplot as plt


def init_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sample_path", type=str, default="./res/cln025")
    parser.add_argument("--date", type=str, required=True)
    return parser.parse_args()

def foldedness_by_hbond(
    traj,
    distance_cutoff=0.35,
    bond_number_cutoff=3
):
	"""
	Generate binary labels for folded/unfolded states based at least 3 bonds among eight bonds
	- TYR1T-YR10OT1
	- TYR1T-YR10OT2
	- ASP3N-TYR8O
	- THR6OG1-ASP3O
	- THR6N-ASP3OD1
	- THR6N-ASP3OD2
	- TYR10N-TYR1O


	Args:
		traj (mdtraj): mdtraj trajectory object
		distance_cutoff (float): donor-acceptor distance cutoff in nm (default 0.35 nm = 3.5 amstrong)
		angle_cutoff (float): hydrogen bond angle cutoff in degrees (default 110 deg)
		bond_number_cutoff (int): minimum number of bonds to be considered as folded (default 3)

	Returns:
		labels (np.array): binary array (1: folded, 0: unfolded)
	"""
	# TYR1N-YR10OT1
	donor_idx = traj.topology.select('residue 1 and name N')[0] # Tyr1:N
	acceptor_idx = traj.topology.select('residue 10 and name O')[0]   # Tyr10:OT1
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_O1 = ((distance[:,0] < distance_cutoff)).astype(int)
	label_O2 = ((distance[:,0] < distance_cutoff)).astype(int) 
	label_O3 = ((distance[:,0] < distance_cutoff)).astype(int)
	label_TYR1N_TYR10OT1 = label_O1 | label_O2 | label_O3


	# TYR1N-YR10OT2
	donor_idx = traj.topology.select('residue 1 and name N')[0] # Tyr1:N
	acceptor_idx = traj.topology.select('residue 10 and name OXT')[0]   # Tyr10:OT2
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_O1 = ((distance[:,0] < distance_cutoff)).astype(int)
	label_O2 = ((distance[:,0] < distance_cutoff)).astype(int)
	label_O3 = ((distance[:,0] < distance_cutoff)).astype(int)
	label_TYR1N_TYR10OT2 = label_O1 | label_O2 | label_O3


	# ASP3N-TYR8O
	donor_idx = traj.topology.select('residue 3 and name N')[0]
	acceptor_idx = traj.topology.select('residue 8 and name O')[0]
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_ASP3N_TYR8O = ((distance[:,0] < distance_cutoff)).astype(int)
 
 
	# THR6OG1-ASP3O
	donor_idx = traj.topology.select('residue 6 and name OG1')[0]
	acceptor_idx = traj.topology.select('residue 3 and name O')[0]
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_THR6OG1_ASP3O = ((distance[:,0] < distance_cutoff)).astype(int)
 
 
	# THR6N-ASP3OD1
	donor_idx = traj.topology.select('residue 6 and name N')[0]
	acceptor_idx = traj.topology.select('residue 3 and name OD1')[0]
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_THR6N_ASP3OD1 = ((distance[:,0] < distance_cutoff)).astype(int)
 
	# THR6N-ASP3OD2
	donor_idx = traj.topology.select('residue 6 and name N')[0]
	acceptor_idx = traj.topology.select('residue 3 and name OD2')[0]
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_THR6N_ASP3OD2 = ((distance[:,0] < distance_cutoff)).astype(int)
 
 
	# GLY7N-ASP3O
	donor_idx = traj.topology.select('residue 7 and name N')[0]
	acceptor_idx = traj.topology.select('residue 3 and name O')[0]
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_GLY7N_ASP3O = ((distance[:,0] < distance_cutoff)).astype(int)
 

	# TYR10N-TYR1O
	donor_idx = traj.topology.select('residue 10 and name N')[0] 
	acceptor_idx = traj.topology.select('residue 1 and name O')[0] 
	distance = md.compute_distances(traj, [[donor_idx, acceptor_idx]])
	label_TYR10N_TYR1O = ((distance[:,0] < distance_cutoff)).astype(int)

	# ASP3OD_THR6OG1_ASP3N_THR8O
	bond_sum = label_TYR1N_TYR10OT1 + label_TYR1N_TYR10OT2 + label_ASP3N_TYR8O + label_THR6OG1_ASP3O \
		+ label_THR6N_ASP3OD1 + label_THR6N_ASP3OD2 + label_GLY7N_ASP3O + label_TYR10N_TYR1O
	labels = bond_sum >= bond_number_cutoff

	return labels, bond_sum

def main():
    args = init_parser()
    res_dir = args.sample_path
    date = args.date
    
    wandb.init(
      project="bioemu-sample",
      entity="eddy26",
      config=args,
      name=f"{date}"
      )

    xtc_path = f"/home/shpark/prj-mlcv/lib/bioemu/{res_dir}/samples_sidechain_rec.xtc"
    pdb_path = f"/home/shpark/prj-mlcv/lib/bioemu/{res_dir}/samples_sidechain_rec.pdb"
    all_samples_traj = md.load(xtc_path, top=pdb_path)
    print(all_samples_traj)

    samples_xtc_path = f"/home/shpark/prj-mlcv/lib/bioemu/{res_dir}/samples_md_equil.xtc"
    samples_pdb_path = f"/home/shpark/prj-mlcv/lib/bioemu/{res_dir}/samples_md_equil.pdb"
    traj = md.load(samples_xtc_path, top=samples_pdb_path)
    print(traj)
    
    label, bond_sum = foldedness_by_hbond(traj)
    print(f"{label.sum()} ({label.sum()/label.shape[0]*100:.0f}%) folded states out of {label.shape[0]} total states")
    print(f"Folded states: {np.where(label)[0]}")
    
    # Print bond_sum frequency analysis
    print("\nBond sum frequency analysis:")
    print("=" * 40)
    unique_values, counts = np.unique(bond_sum, return_counts=True)
    for value, count in zip(unique_values, counts):
        percentage = (count / len(bond_sum)) * 100
        print(f"Bond count {value}: {count} frames ({percentage:.1f}%)")
        if value == 2:
            print("-"*30)
    print("=" * 40)

    # Draw histogram of bond_sum
    plt.figure(figsize=(8, 5))
    plt.hist(bond_sum, bins=np.arange(bond_sum.min(), bond_sum.max()+2)-0.5, edgecolor='black')
    plt.xlabel('Bond Sum')
    plt.ylabel('Number of Frames')
    plt.title('Histogram of Bond Sum')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.show()
    
    
    wandb.log({
        "valid_samples": label.shape[0],
        "folded_states": label.sum(),
        "unfolded_states": label.shape[0] - label.sum(),
        "folded_percentage": label.sum() / label.shape[0] * 100,
        "bond_sum_histogram": wandb.Image(plt.gcf()),
        # "bond_sum_frequency": wandb.Table(data=np.array([unique_values, counts]), columns=["Bond Sum", "Frequency"]),
    })
    wandb.finish()

if __name__ == "__main__":
    main()