window = 2  # number of adjacent frames to average

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)
	cs_ids = s.coordset_ids

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

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

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