# 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
