from __future__ import division
import numpy as np
from .RootFinder import find_roots_gen
def demo_find_roots(originalContour, f, df=None, saveFile=None, automaticAnim=False,
returnAnim=False, writer=None, **rootsKwargs):
"""
An animated demonstration of the root finding process using matplotlib.
Parameters
----------
saveFile : str, optional
If given then the animation will be saved to disk with filename
equal to saveFile instead of being shown.
automaticAnim : bool, optional
If False (default) then press SPACE to step the animation forward
If True then the animation will play automatically until all the
roots have been found.
returnAnim : bool, optional
If True then the matplotlib animation object will be returned
instead of being shown. Defaults to False.
writer : str, optional
Passed to :meth:`matplotlib.animation.FuncAnimation.save`.
**rootsKwargs
Additional key word arguments passed to :meth:`~cxroots.Contour.Contour.roots`.
"""
import matplotlib.pyplot as plt
from matplotlib import animation
fig = plt.gcf()
ax = plt.gca()
rootFinder = find_roots_gen(originalContour, f, df, **rootsKwargs)
def init():
originalContour.plot(linecolor='k', linestyle='--')
originalContour._sizePlot()
def update_frame(args):
roots, multiplicities, boxes, numberOfRemainingRoots = args
print(args)
plt.cla() # clear axis
originalContour.plot(linecolor='k', linestyle='--')
for box in boxes:
if not hasattr(box, '_color'):
cmap = plt.get_cmap('jet')
box._color = cmap(np.random.random())
plt.text(box.centralPoint.real, box.centralPoint.imag, box._numberOfRoots)
box.plot(linecolor=box._color)
plt.scatter(np.real(roots), np.imag(roots), color='k', marker='x')
rootsLabel = ax.text(0.02, 0.95, 'Zeros remaining: %i'%numberOfRemainingRoots, transform=ax.transAxes)
originalContour._sizePlot()
fig.canvas.draw()
if saveFile:
automaticAnim = True
if automaticAnim or returnAnim:
anim = animation.FuncAnimation(fig, update_frame, init_func=init, frames=rootFinder)
if returnAnim: return anim
else:
def draw_next(event):
if event.key == ' ':
update_frame(next(rootFinder))
fig.canvas.mpl_connect('key_press_event', draw_next)
if saveFile:
anim.save(filename=saveFile, fps=1, dpi=200, writer=writer)
plt.close()
else:
plt.show()