Scripts: smooth_md.py

window = 2  # number of adjacent frames to average
stride = 1  # only use every Nth frame of trajectory (to save memory for large trjectories)

from chimerax.atomic import Structure
for s in session.models:
        if not isinstance(s, Structure) or s.num_coordsets == 1:
                continue
        session.logger.status("Smoothing trajectory %s" % s)

        session.logger.status("Gathering coordinates", secondary=True)
        if stride > 1:
                coordset_ids = s.coordset_ids
                coord_sets = [s.coordset(coordset_ids[i]).xyzs for i in range(0, len(coordset_ids), stride)]
        else:
                coord_sets = [s.coordset(cs_id).xyzs for cs_id in s.coordset_ids]

        session.logger.status("Computing smoothed coordinates", secondary=True)
        import numpy
        smoothed = numpy.zeros((len(coord_sets), len(coord_sets[0]), 3), type(coord_sets[0][0][0]))
        for i in range(len(coord_sets)):
                weight_tot = 0
                avg = smoothed[i]
                for j in range(i-window, i+window+1):
                        if j < 0 or j >= len(coord_sets):
                                continue
                        weight = window + 1 - abs(i-j)
                        weight_tot += weight
                        avg += weight * coord_sets[j]
                avg /= weight_tot

        session.logger.status("Setting coordinates", secondary=True)
        s.add_coordsets(smoothed)

        session.logger.status("Smoothed trajectory %s" % s)