I'm having problems with the `ax` parameters of the `rlocus` function

I was trying today to plot the root locus of a transfer function, but wanted to pass a specific axes object to it, so I tried to use the ax parameter of the rlocus function:

from control import *
import matplotlib.pyplot as plt

s = tf([1, 0], 1)
G = 1/(s*(s+1))

fig, ax = plt.subplots(dpi=200)
rlocus(G, ax=ax)

But when doing this, the following error was raised:

---------------------------------------------------------------------------
UnboundLocalError                         Traceback (most recent call last)
<ipython-input-1-80498c4b1a54> in <module>
      6 
      7 fig, ax = plt.subplots(dpi=200)
----> 8 rlocus(G, ax=ax)

~\anaconda3\lib\site-packages\control\rlocus.py in root_locus(sys, kvect, xlim, ylim, plotstr, plot, print_gain, grid, ax, **kwargs)
    173 
    174         if print_gain and not sisotool:
--> 175             fig.canvas.mpl_connect(
    176                 'button_release_event',
    177                 partial(_RLClickDispatcher, sys=sys, fig=fig,

UnboundLocalError: local variable 'fig' referenced before assignment

Control version: 0.9.0

[billtubbs]

I’ll have a look at #634. But I’m waiting for further input on the new plotting paradigm. That will require more thinking and discussion and probably take a while to figure out a way forward. There might be better ways to do it. I’m happy to try fixing the immediate bugs on the existing functions when I next get the chance.

[billtubbs]

I’m happy to implement this immediate fix (and #632 at the same time) and maybe add some unit tests, including for other plot functions where applicable. If I recall correctly, I was trying to do this a while back and the problem is that some plot funcs need to generate their own axis (e.g. radial plots) and some generate a figure containing more than one axis.