# Tests for plot functions

from control import tf, c2d, pzmap, rlocus, bode, nyquist, nichols, sisotool

from control.grid import sgrid, zgrid

import numpy as np

import matplotlib.pyplot as plt

import os

plot_dir = 'temp_plots'

if not os.path.exists(plot_dir):

os.mkdir(plot_dir)

# Helper function for saving figures

def savefig(label, fname='{}_fig{}.pdf'):

fig_n = plt.gcf().number

plt.savefig(os.path.join(plot_dir, fname.format(label, fig_n)))

# First test zgrid and sgrid on their own

plt.figure()

sgrid()

plt.title('s-grid')

savefig('grids1')

plt.figure()

zgrid()

plt.title('z-grid')

savefig('grids1')

# Display and wait for user to close all figures

plt.show()

# Adding grids to an existing figure

fig = plt.figure(figsize=(4.5, 4))

sgrid(fig=fig)

plt.title('s-grid')

savefig('grids2')

fig = plt.figure(figsize=(4.5, 4))

zgrid(fig=fig)

plt.title('z-grid')

savefig('grids2')

fig, ax = plt.subplots(figsize=(4.5, 4))

zgrid(ax=ax) # only works with zgrid

plt.title('z-grid #2')

savefig('grids2')

# Display and wait for user to close all figures

plt.show()

# Create figure with subplots

fig = plt.figure()

sgrid(position=(2, 2, 1))

ax2 = fig.add_subplot(2, 2, 2)

zgrid(ax=ax2) # alternative for zgrids

#zgrid(position=(2, 2, 2))

sgrid(position=(2, 2, 3))

ax4 = fig.add_subplot(2, 2, 4)

zgrid(ax=ax4) # alternative for zgrids

plt.tight_layout()

savefig('grids_subplots')

# Display and wait for user to close all figures

plt.show()

# Define test systems

s = tf('s')

G1 = 2*(s+1)/(s*(s-2))

G2 = tf(1, [1, 8, 30, 56, 65])

Gz1 = c2d(G1, 0.2)

Gz2 = c2d(G2, 0.2)

# Test rlocus function

# Single default plots

plt.figure()

rlocus(G1)

savefig('rlocus1')

plt.figure()

rlocus(G2)

savefig('rlocus1')

plt.figure()

rlocus(Gz1)

savefig('rlocus1')

plt.figure()

rlocus(Gz2)

savefig('rlocus1')

# Display and wait for user to close all figures

plt.show()

# Single default plots - without grid

plt.figure(figsize=(4.5, 4))

rlocus(G1, grid=False)

savefig('rlocus2')

plt.figure(figsize=(4.5, 4))

rlocus(G2, grid=False)

savefig('rlocus2')

plt.figure(figsize=(4.5, 4))

rlocus(Gz1, grid=False)

savefig('rlocus2')

plt.figure(figsize=(4.5, 4))

rlocus(Gz2, grid=False)

savefig('rlocus2')

# Display and wait for user to close all figures

plt.show()

# Subplots - discrete time

fig, axes = plt.subplots(2, 2)

rlocus(Gz1, ax=axes[0, 0], title='Gz1')

rlocus(Gz2, ax=axes[0, 1], title='Gz1')

rlocus(Gz1, ax=axes[1, 0], grid=False, title='Gz1 - no grid')

rlocus(Gz2, ax=axes[1, 1], grid=False, title='Gz1 - no grid')

plt.tight_layout()

savefig('rlocus_subplots1')

# Subplots - sontinuous time, done differently due to sgrid

fig = plt.figure()

rlocus(G1, position=(2, 2, 1), title='G1')

rlocus(G2, position=(2, 2, 2), title='G2')

rlocus(G1, position=(2, 2, 3), grid=False, title='G1 - no grid')

rlocus(G2, position=(2, 2, 4), grid=False, title='G2 - no grid')

plt.tight_layout()

savefig('rlocus_subplots1')

# Display and wait for user to close all figures

plt.show()

# Check sisotool works (includes rlocus plot)

sisotool(G1)

savefig('sisotool')

plt.show()

sisotool(Gz1)

savefig('sisotool')

plt.show()