add pzmap_test

bnavigator · bnavigator · commit e646370d9d8f · 2020-08-20T23:37:26.000+02:00
diff --git a/control/tests/conftest.py b/control/tests/conftest.py
@@ -1,13 +1,39 @@
 # contest.py - pytest local plugins and fixtures
 
-import control
 import os
 
+import matplotlib as mpl
 import pytest
 
+import control
+
 
 @pytest.fixture(scope="session", autouse=True)
 def use_numpy_ndarray():
     """Switch the config to use ndarray instead of matrix"""
     if os.getenv("PYTHON_CONTROL_STATESPACE_ARRAY") == "1":
         control.config.defaults['statesp.use_numpy_matrix'] = False
+
+
+@pytest.fixture(scope="function")
+def editsdefaults():
+    """Make sure any changes to the defaults only last during a test"""
+    restore = control.config.defaults.copy()
+    yield
+    control.config.defaults.update(restore)
+
+
+@pytest.fixture(scope="function")
+def mplcleanup():
+    """Workaround for python2
+
+    python 2 does not like to mix the original mpl decorator with pytest
+    fixtures. So we roll our own.
+    """
+    save = mpl.units.registry.copy()
+    try:
+        yield
+    finally:
+        mpl.units.registry.clear()
+        mpl.units.registry.update(save)
+        mpl.pyplot.close("all")
diff --git a/control/tests/pzmap_test.py b/control/tests/pzmap_test.py
@@ -0,0 +1,77 @@
+# -*- coding: utf-8 -*-
+""" pzmap_test.py - test pzmap()
+
+Created on Thu Aug 20 20:06:21 2020
+
+@author: bnavigator
+"""
+
+import numpy as np
+import pytest
+from matplotlib import pyplot as plt
+
+from control import TransferFunction, config, pzmap
+
+
+@pytest.mark.parametrize("kwargs",
+                         [dict(),
+                          dict(plot=False),
+                          dict(plot=True),
+                          dict(grid=True),
+                          dict(title="Custom Title")],
+                         ids=["default",
+                              "noplot",
+                              "plotTrue",
+                              "grid",
+                              "title"])
+@pytest.mark.parametrize("setdefaults",  # set by kw or config
+                         [False, True],
+                         ids=["kw", "config"])
+@pytest.mark.parametrize("dt", [0, True], ids=["s", "z"])
+def test_pzmap(kwargs, setdefaults, dt, editsdefaults, mplcleanup):
+    """Test pzmap"""
+    # T from from pvtol-nested example
+    T = TransferFunction([-9.0250000e-01, -4.7200750e+01, -8.6812900e+02,
+                          +5.6261850e+03, +2.1258472e+05, +8.4724600e+05,
+                          +1.0192000e+06, +2.3520000e+05],
+                         [9.02500000e-03, 9.92862812e-01, 4.96974094e+01,
+                          1.35705659e+03, 2.09294163e+04, 1.64898435e+05,
+                          6.54572220e+05, 1.25274600e+06, 1.02420000e+06,
+                          2.35200000e+05])
+
+    Pref = [-23.8877+19.3837j, -23.8877-19.3837j, -23.8349+15.7846j,
+            -23.8349-15.7846j,  -5.2320 +0.4117j,  -5.2320 -0.4117j,
+            -2.2246 +0.0000j,  -1.5160 +0.0000j,  -0.3627 +0.0000j]
+    Zref = [-23.8877+19.3837j, -23.8877-19.3837j, +14.3637 +0.0000j,
+            -14.3637 +0.0000j,  -2.2246 +0.0000j,  -2.0000 +0.0000j,
+            -0.3000 +0.0000j]
+
+    pzkwargs = kwargs.copy()
+    if setdefaults:
+        for k in ['plot', 'grid']:
+            if k in pzkwargs:
+                v = pzkwargs.pop(k)
+                config.set_defaults('pzmap', **{k: v})
+
+    P, Z = pzmap(T, **pzkwargs)
+
+    np.testing.assert_allclose(P, Pref, rtol=1e-3)
+    np.testing.assert_allclose(Z, Zref, rtol=1e-3)
+
+    if kwargs.get('plot', True):
+        ax = plt.gca()
+        assert ax.get_title() == kwargs.get('title', 'Pole Zero Map')
+        if kwargs.get('grid', False):
+            # TODO: check for correct grid
+            pass
+
+
+def test_pzmap_warns():
+    with pytest.warns(FutureWarning):
+        pzmap(TransferFunction([1], [1, 2]), Plot=True)
+
+
+def test_pzmap_raises():
+    with pytest.raises(TypeError):
+        # not an LTI system
+        pzmap(([1], [1,2]))
