Monetary Circuit Model/graphing.py

# This function takes an expression (the right-hand side of the solution),
# a label, and a hue and plots the expression from 0 to 10. Its first derivative
# is also plotted as a separate line, dotted and faded.
def buildplot(t, expression, label, h, start, finish):

	# Plot the expression. Remove legend_label and restart the interpreter if
	# you get errors about LaTeX not being installed.
	p = plot(expression, t, (start, finish), rgbcolor=hue(h), legend_label=label)

	# Plot the derivative of the expression at 30% saturation with a dashed line
	p += plot(diff(expression, t), t, (0, 10),
			rgbcolor=hue(h, 0.3), linestyle='--')
	
	return p


# Helper function to generate a unique hue for each curve
def colourwheel(colours):
	for c in xrange(0, colours):
		yield float(c) / colours


class _UseAll(object):
	def __contains__(self, v):
		return True
All = _UseAll()

def plot_ode(names, sol, ivar, use=All, start=0, finish=10, xAxis="", yAxis=""):
	'''
	names	= dictionary mapping functions to labels
	sol 	= list of solutions to system (of the form function == expression)
	ivar	= independant var
	use		= set of functions to plot
	start	= left bound of the plot
	finish 	= right bound of the plot
	title	= title of the plot
	xAxis	= title of the X axis
	yAxis	= title of the Y axis
	'''

	# Gather subset of equations to plot
	toPlot = [s for s in sol if s.lhs() in use]

	# Build individual plots
	plots = [buildplot(ivar, s.rhs(), names[s.lhs()], c, start, finish)
			for s, c in zip(toPlot, colourwheel(len(toPlot)))]
	
	# Aggregate
	ret = Graphics()
	for p in plots:
		ret += p
	
	# Add axes labels
	ret.axes_labels([xAxis, yAxis])

	return ret


def plot_expr(expr, names, sol, ivar, start=0, finish=10, xAxis="", yAxis=""):

	solDict = dict([(s.lhs(), s.rhs()) for s in sol])

	plots = [buildplot(ivar, x.substitute(solDict), label, c, start, finish)
			for x, label, c in zip(expr, names, colourwheel(len(expr)))]

	ret = Graphics()
	for p in plots:
		ret += p
	ret.axes_labels([xAxis, yAxis])

	return ret