Pull Structures

gistfile1.py

from mdtraj.formats.dcd import DCDTrajectoryFile
from mdtraj.formats.prmtop import load_prmtop
km = KMedoids(...)

# "centers" is a pandas DataFrame with metadata about each trajectory
# from which we are pulling structures

def get_top(struc):
    return centers.query("structure == @struc")['top_fn'].iloc[0]

# Load topologies
top_dict = {}
for struct in centers['structure'].unique():
    top_dict[struct] = load_prmtop(get_top(struct))

# Allocate arrays
xyz_dict = {}
bl_dict = {}
ba_dict = {}
i_dict = {}
counts = centers['structure'].value_counts()

for struct in top_dict:
    shape = (counts[struct], top_dict[struct].n_atoms, 3)
    shapeb = (counts[struct], 3)
    xyz_dict[struct] = np.zeros(shape, dtype='float')
    bl_dict[struct] = np.zeros(shapeb, dtype='float')
    ba_dict[struct] = np.zeros(shapeb, dtype='float')    
    i_dict[struct] = 0

# Load coordinates
prog = 0
for (i, row), frame in zip(centers.iterrows(), km.cluster_ids_[:,1]):
    struct = row['structure']
    with DCDTrajectoryFile(row['traj_fn']) as dcdf:
        dcdf.seek(frame)
        xyz, box_length, box_angle = dcdf.read(n_frames=1)
        xyz_dict[struct][i_dict[struct]] = xyz[0]
        bl_dict[ struct][i_dict[struct]] = box_length[0]
        ba_dict[ struct][i_dict[struct]] = box_angle[0]
        i_dict[struct] += 1
        
    prog += 1
    print("\r", prog, end='', flush=True)

# Make a `Trajectory` and save
for struct in centers['structure'].unique():
    traj = md.Trajectory(
        xyz = xyz_dict[struct],
        unitcell_lengths = bl_dict[struct],
        unitcell_angles = ba_dict[struct],
        topology = top_dict[struct],
        time = np.arange(len(xyz_dict[struct]))
    )
           
    traj.save("{}/{}.dcd".format(dir_name, struct))
    shutil.copy(get_top(struct), "{}/{}.prmtop".format(dir_name, struct))