<commit_before><commit_msg>initial example<commit_after>#!/usr/bin/env python
import MDAnalysis, argparse, math, tempfile, shutil, os, sys
import numpy as np
from PIL import Image


parser = argparse.ArgumentParser(description='Compress a MD trajectory')
parser.add_argument('topo',metavar='topology_file')
parser.add_argument('traj',metavar='trajectory_file')
parser.add_argument('--out',metavar='output_file',default="output.m4v")
parser.add_argument("--selection",default="protein",required=False)

args = parser.parse_args()

model = MDAnalysis.Universe(args.topo,args.traj)
selection = model.select_atoms(args.selection)
n = selection.n_atoms

#make square-ish images -- this can be optimized to reduce dead pixels
#these also have to be even..
h = int(math.ceil(math.sqrt(n)))
w = h

#figure out overall range for x,y,z
maxvals = np.array([float("-inf")]*3)
minvals = np.array([float("inf")]*3)

for ts in model.trajectory:
	for c in selection.coordinates():
		for i in xrange(3):
			maxvals[i] = max(maxvals[i],c[i])
			minvals[i] = min(minvals[i],c[i])

ranges = maxvals-minvals;

print n,model.trajectory.n_frames
print minvals, ranges
#create an image for each frame in a tempdir
tmpdir = tempfile.mkdtemp()
for ts in model.trajectory:
	img = Image.new("RGB",(w,h))
	pix = img.load()
	#todo: support more than 8bits per coordinate
	for (ci,c) in enumerate(selection.coordinates()):
		i = ci/h
		j = ci%h
		pt = [0,0,0]
		for dim in xrange(3):
			pt[dim] = int(round(255*(c[dim]-minvals[dim])/ranges[dim]))
		pix[i,j] = tuple(pt)
	img.save('%s/%d.png' % (tmpdir, ts.frame))

os.system('ffmpeg  -i %s/%%d.png -c:v libx264  -preset veryslow -qp 0 %s' % (tmpdir, args.out))
shutil.rmtree(tmpdir)
{'fraction_non_alphanumeric': 0.1610817166372722, 'fraction_numerical': 0.00940623162845385, 'mean_word_length': 5.350746268656716, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '291680', 'n_tokens_mistral': 669, 'n_tokens_neox': 638, 'n_words': 150}