codegen.py

"""Code generator 

Generate python code based on the formulas generated by trajectory.py
"""

class CodeGenerator():
	def __init__(self, t, xs, xlimits, constraints, ps):
		self.t = t
		self.xs = xs
		self.var = [t] + xs
		self.xlimits = xlimits
		self.order = len(xs)
		self.constraints = constraints
		self.ps = ps

	def header(self):
		template = '''
import numpy as np

def compute_tx(constraints, err=1e-7):
	r = np.concatenate([np.roots(c) for c in constraints])
	r = r.real[abs(r.imag)<err]
	r = r[r>=-err]
	return np.min(r)
'''
		return template

	def def_compute_trajectory(self):
		temp = '''
def compute_trajectory(%s):
%s
%s
'''
		temp_tx = "	%s = compute_tx(%s)"
		temp_traj = "	%(vars)s = %(traj)s\n	yield %(vars)s"
		statement_tx = '\n'.join(
			[temp_tx % (tx, c) for tx, c in self.constraints.items()]
			)
		statement_traj = '\n'.join([temp_traj % {
			'vars': ', '.join([str(x) for x in self.var]), 
			'traj': ', '.join([str(p[x]) for x in self.var])
			} for p in self.ps])
		return temp % (
			', '.join([str(xl) for xl in self.xlimits]), 
			statement_tx, statement_traj)

	def write(self, stream):
		stream.write(self.header())
		stream.write(self.def_compute_trajectory())