additional unit tests for coverage + bug fixes

murrayrm · murrayrm · commit e65c4f29e1e4 · 2022-08-20T10:11:13.000-07:00
diff --git a/control/flatsys/basis.py b/control/flatsys/basis.py
@@ -36,6 +36,8 @@
 # OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 # SUCH DAMAGE.
 
+import numpy as np
+
 
 # Basis family class (for use as a base class)
 class BasisFamily:
@@ -86,16 +88,24 @@ def eval(self, coeffs, tlist, var=None):
                 sum([coeffs[i] * self(i, t) for i in range(self.N)])
                 for t in tlist]
 
-        else:
-            # Multi-variable basis
+        elif var is None:
+            # Multi-variable basis with single list of coefficients
             values = np.empty((self.nvars, tlist.size))
+            offset = 0
             for j in range(self.nvars):
                 coef_len = self.var_ncoefs(j)
                 values[j] = np.array([
-                    sum([coeffs[i] * self(i, t, var=j)
+                    sum([coeffs[offset + i] * self(i, t, var=j)
                          for i in range(coef_len)])
                     for t in tlist])
+                offset += coef_len
             return values
 
+        else:
+            return np.array([
+                sum([coeffs[i] * self(i, t, var=var)
+                     for i in range(self.var_ncoefs(var))])
+                for t in tlist])
+
     def eval_deriv(self, i, j, t, var=None):
         raise NotImplementedError("Internal error; improper basis functions")
diff --git a/control/tests/flatsys_test.py b/control/tests/flatsys_test.py
@@ -313,7 +313,7 @@ def test_flat_cost_constr(self, basis):
         np.testing.assert_array_almost_equal(uf, u_const[:, -1])
 
         # Make sure that the solution respects the bounds (with some slop)
-        for i in range(x_const.shape[0]): 
+        for i in range(x_const.shape[0]):
             assert all(lb[i] - x_const[i] < rtol * abs(lb[i]) + atol)
             assert all(x_const[i] - ub[i] < rtol * abs(ub[i]) + atol)
 
@@ -673,3 +673,43 @@ def test_response(self, xf, uf, Tf):
         np.testing.assert_equal(T, response.time)
         np.testing.assert_equal(u, response.inputs)
         np.testing.assert_equal(x, response.states)
+
+    @pytest.mark.parametrize(
+        "basis",
+        [fs.PolyFamily(4),
+         fs.BezierFamily(4),
+         fs.BSplineFamily([0, 1], 4),
+         fs.BSplineFamily([0, 1], 4, vars=2),
+         fs.BSplineFamily([0, 1], [4, 3], [2, 1], vars=2),
+        ])
+    def test_basis_class(self, basis):
+        timepts = np.linspace(0, 1, 10)
+
+        if basis.nvars is None:
+            # Evaluate function on basis vectors
+            for j in range(basis.N):
+                coefs = np.zeros(basis.N)
+                coefs[j] = 1
+                np.testing.assert_equal(
+                    basis.eval(coefs, timepts),
+                    basis.eval_deriv(j, 0, timepts))
+        else:
+            # Evaluate each variable on basis vectors
+            for i in range(basis.nvars):
+                for j in range(basis.var_ncoefs(i)):
+                    coefs = np.zeros(basis.var_ncoefs(i))
+                    coefs[j] = 1
+                    np.testing.assert_equal(
+                        basis.eval(coefs, timepts, var=i),
+                        basis.eval_deriv(j, 0, timepts, var=i))
+
+            # Evaluate multi-variable output
+            offset = 0
+            for i in range(basis.nvars):
+                for j in range(basis.var_ncoefs(i)):
+                    coefs = np.zeros(basis.N)
+                    coefs[offset] = 1
+                    np.testing.assert_equal(
+                        basis.eval(coefs, timepts)[i],
+                        basis.eval_deriv(j, 0, timepts, var=i))
+                    offset += 1
