add support for multiple spline variables; update docstrings

murrayrm · murrayrm · commit 0ab0d0703f22 · 2022-08-20T10:11:13.000-07:00
diff --git a/control/flatsys/basis.py b/control/flatsys/basis.py
@@ -47,7 +47,11 @@ class BasisFamily:
 
       :math:`z_i^{(q)}(t)` = basis.eval_deriv(self, i, j, t)
 
-    Parameters
+    A basis set can either consist of a single variable that is used for
+    each flat output (nvars = None) or a different variable for different
+    flat outputs (nvars > 0).
+
+    Attributes
     ----------
     N : int
         Order of the basis set.
@@ -56,15 +60,38 @@ class BasisFamily:
     def __init__(self, N):
         """Create a basis family of order N."""
         self.N = N                    # save number of basis functions
+        self.nvars = None             # default number of variables
+        self.coef_offset = [0]        # coefficient offset for each variable
+        self.coef_length = [N]        # coefficient length for each variable
 
-    def __call__(self, i, t):
+    def __call__(self, i, t, var=None):
         """Evaluate the ith basis function at a point in time"""
-        return self.eval_deriv(i, 0, t)
+        return self.eval_deriv(i, 0, t, var=var)
+
+    def var_ncoefs(self, var):
+        """Get the number of coefficients for a variable"""
+        return self.N if self.nvars is None else self.coef_length[var]
+
+    def eval(self, coeffs, tlist, var=None):
+        if self.nvars is None and var != None:
+            raise SystemError("multi-variable call to a scalar basis")
+
+        elif self.nvars is None:
+            # Single variable basis
+            return [
+                sum([coeffs[i] * self(i, t) for i in range(self.N)])
+                for t in tlist]
 
-    def eval(self, coeffs, tlist):
-        return [
-            sum([coeffs[i] * self(i, t) for i in range(self.N)])
-            for t in tlist]
+        else:
+            # Multi-variable basis
+            values = np.empty((self.nvars, tlist.size))
+            for j in range(self.nvars):
+                coef_len = self.var_ncoefs(j)
+                values[j] = np.array([
+                    sum([coeffs[i] * self(i, t, var=j)
+                         for i in range(coef_len)])
+                    for t in tlist])
+            return values
 
-    def eval_deriv(self, i, j, t):
+    def eval_deriv(self, i, j, t, var=None):
         raise NotImplementedError("Internal error; improper basis functions")
diff --git a/control/flatsys/bezier.py b/control/flatsys/bezier.py
@@ -59,7 +59,7 @@ def __init__(self, N, T=1):
         self.T = T                      # save end of time interval
 
     # Compute the kth derivative of the ith basis function at time t
-    def eval_deriv(self, i, k, t):
+    def eval_deriv(self, i, k, t, var=None):
         """Evaluate the kth derivative of the ith basis function at time t."""
         if i >= self.N:
             raise ValueError("Basis function index too high")
diff --git a/control/flatsys/bspline.py b/control/flatsys/bspline.py
@@ -17,14 +17,14 @@ class BSplineFamily(BasisFamily):
     across a set of breakpoints with given order and smoothness.
 
     """
-    def __init__(self, breakpoints, degree, smoothness=None, vars=1):
+    def __init__(self, breakpoints, degree, smoothness=None, vars=None):
         """Create a B-spline basis for piecewise smooth polynomials
 
         Define B-spline polynomials for a set of one or more variables.
-        B-splines are characterized by a set of intervals separated by break
-        points.  On each interval we have a polynomial of a certain order
-        and the spline is continuous up to a given smoothness at interior
-        break points.
+        B-splines are used as a basis for a set of piecewise smooth
+        polynomials joined at breakpoints. On each interval we have a
+        polynomial of a given order and the spline is continuous up to a
+        given smoothness at interior breakpoints.
 
         Parameters
         ----------
@@ -41,8 +41,11 @@ def __init__(self, breakpoints, degree, smoothness=None, vars=1):
             For each spline variable, the smoothness at breakpoints (number
             of derivatives that should match).
 
-        vars : int or list of str, option
-            The number of spline variables or a list of spline variable names.
+        vars : None or int, optional
+            The number of spline variables.  If specified as None (default),
+            then the spline basis describes a single variable, with no
+            indexing.  If the number of spine variables is > 0, then the
+            spline basis is index using the `var` keyword.
 
         """
         # Process the breakpoints for the spline */
@@ -58,17 +61,19 @@ def __init__(self, breakpoints, degree, smoothness=None, vars=1):
             raise ValueError("break points must be strictly increasing values")
 
         # Decide on the number of spline variables
-        if isinstance(vars, list) and all([isinstance(v, str) for v in vars]):
-            raise NotImplemented("list of variable names not yet supported")
+        if vars is None:
+            nvars = 1
+            self.nvars = None           # track as single variable
         elif not isinstance(vars, int):
-            raise TypeError("vars must be an integer or list of strings")
+            raise TypeError("vars must be an integer")
         else:
             nvars = vars
+            self.nvars = nvars
 
         #
         # Process B-spline parameters (order, smoothness)
         #
-        # B-splines are characterized by a set of intervals separated by
+        # B-splines are defined on a set of intervals separated by
         # breakpoints.  On each interval we have a polynomial of a certain
         # order and the spline is continuous up to a given smoothness at
         # breakpoints.  The code in this section allows some flexibility in
@@ -99,14 +104,15 @@ def process_spline_parameters(
             elif all([isinstance(v, allowed_types) for v in values]):
                 # List of values => make sure it is the right size
                 if len(values) != length:
-                    raise ValueError(f"length of '{name}' does not match n")
+                    raise ValueError(f"length of '{name}' does not match"
+                                     f" number of variables")
             else:
                 raise ValueError(f"could not parse '{name}' keyword")
 
             # Check to make sure the values are OK
             if values is not None and any([val < minimum for val in values]):
                 raise ValueError(
-                    f"invalid value for {name}; must be at least {minimum}")
+                    f"invalid value for '{name}'; must be at least {minimum}")
 
             return values
 
@@ -123,25 +129,23 @@ def process_spline_parameters(
         if any([degree[i] - smoothness[i] < 1 for i in range(nvars)]):
             raise ValueError("degree must be greater than smoothness")
 
-        # Store the parameters and process them in call_ntg()
-        self.nvars = nvars
+        # Store the parameters for the spline (self.nvars already stored)
         self.breakpoints = breakpoints
         self.degree = degree
         self.smoothness = smoothness
-        self.nintervals = breakpoints.size - 1
 
         #
         # Compute parameters for a SciPy BSpline object
         #
-        # To create a B-spline, we need to compute the knot points, keeping
-        # track of the use of repeated knot points at the initial knot and
+        # To create a B-spline, we need to compute the knotpoints, keeping
+        # track of the use of repeated knotpoints at the initial knot and
         # final knot as well as repeated knots at intermediate points
         # depending on the desired smoothness.
         #
 
         # Store the coefficients for each output (useful later)
         self.coef_offset, self.coef_length, offset = [], [], 0
-        for i in range(self.nvars):
+        for i in range(nvars):
             # Compute number of coefficients for the piecewise polynomial
             ncoefs = (self.degree[i] + 1) * (len(self.breakpoints) - 1) - \
                 (self.smoothness[i] + 1) * (len(self.breakpoints) - 2)
@@ -151,48 +155,43 @@ def process_spline_parameters(
             offset += ncoefs
         self.N = offset         # save the total number of coefficients
 
-        # Create knot points for each spline variable
+        # Create knotpoints for each spline variable
         # TODO: extend to multi-dimensional breakpoints
         self.knotpoints = []
-        for i in range(self.nvars):
+        for i in range(nvars):
             # Allocate space for the knotpoints
             self.knotpoints.append(np.empty(
                 (self.degree[i] + 1) + (len(self.breakpoints) - 2) * \
                 (self.degree[i] - self.smoothness[i]) + (self.degree[i] + 1)))
 
-            # Initial knot points
+            # Initial knotpoints (multiplicity = order)
             self.knotpoints[i][0:self.degree[i] + 1] = self.breakpoints[0]
             offset = self.degree[i] + 1
 
-            # Interior knot points
+            # Interior knotpoints (multiplicity = degree - smoothness)
             nknots = self.degree[i] - self.smoothness[i]
             assert nknots > 0           # just in case
             for j in range(1, self.breakpoints.size - 1):
                 self.knotpoints[i][offset:offset+nknots] = self.breakpoints[j]
                 offset += nknots
 
-            # Final knot point
+            # Final knotpoint (multiplicity = order)
             self.knotpoints[i][offset:offset + self.degree[i] + 1] = \
                 self.breakpoints[-1]
 
-    def eval(self, coefs, tlist):
-        return np.array([
-            BSpline(self.knotpoints[i],
-                    coefs[self.coef_offset[i]:
-                          self.coef_offset[i] + self.coef_length[i]],
-                    self.degree[i])(tlist)
-            for i in range(self.nvars)])
-
     # Compute the kth derivative of the ith basis function at time t
-    def eval_deriv(self, i, k, t, squeeze=True):
+    def eval_deriv(self, i, k, t, var=None):
         """Evaluate the kth derivative of the ith basis function at time t."""
-        if self.nvars > 1 or not squeeze:
-            raise NotImplementedError(
-                "derivatives of multi-variable splines not yet supported")
+        if self.nvars is None or (self.nvars == 1 and var is None):
+            # Use same variable for all requests
+            var = 0
+        elif self.nvars > 1 and var is None:
+            raise SystemError(
+                "scalar variable call to multi-variable splines")
 
         # Create a coefficient vector for this spline
-        coefs = np.zeros(self.coef_length[0]); coefs[i] = 1
+        coefs = np.zeros(self.coef_length[var]); coefs[i] = 1
 
         # Evaluate the derivative of the spline at the desired point in time
-        return BSpline(self.knotpoints[0], coefs,
-                       self.degree[0]).derivative(k)(t)
+        return BSpline(self.knotpoints[var], coefs,
+                       self.degree[var]).derivative(k)(t)
diff --git a/control/flatsys/flatsys.py b/control/flatsys/flatsys.py
@@ -155,6 +155,7 @@ def __init__(self,
         if reverse is not None: self.reverse = reverse
 
         # Save the length of the flat flag
+        # TODO: missing
 
     def __str__(self):
         return f"{NonlinearIOSystem.__str__(self)}\n\n" \
@@ -233,17 +234,19 @@ def _basis_flag_matrix(sys, basis, flag, t, params={}):
     column of the matrix corresponds to a basis function and each row is a
     derivative, with the derivatives (flag) for each output stacked on top
     of each other.
-
+l
     """
     flagshape = [len(f) for f in flag]
-    M = np.zeros((sum(flagshape), basis.N * sys.ninputs))
+    M = np.zeros((sum(flagshape),
+                  sum([basis.var_ncoefs(i) for i in range(sys.ninputs)])))
     flag_off = 0
-    coeff_off = 0
+    coef_off = 0
     for i, flag_len in enumerate(flagshape):
-        for j, k in itertools.product(range(basis.N), range(flag_len)):
-            M[flag_off + k, coeff_off + j] = basis.eval_deriv(j, k, t)
+        coef_len = basis.var_ncoefs(i)
+        for j, k in itertools.product(range(coef_len), range(flag_len)):
+            M[flag_off + k, coef_off + j] = basis.eval_deriv(j, k, t, var=i)
         flag_off += flag_len
-        coeff_off += basis.N
+        coef_off += coef_len
     return M
 
 
@@ -362,11 +365,16 @@ def point_to_point(
     if basis is None:
         basis = PolyFamily(2 * (sys.nstates + sys.ninputs))
 
+    # If a multivariable basis was given, make sure the size is correct
+    if basis.nvars is not None and basis.nvars != sys.ninputs:
+        raise ValueError("size of basis does not match flat system size")
+
     # Make sure we have enough basis functions to solve the problem
-    if basis.N * sys.ninputs < 2 * (sys.nstates + sys.ninputs):
+    ncoefs = sum([basis.var_ncoefs(i) for i in range(sys.ninputs)])
+    if ncoefs < 2 * (sys.nstates + sys.ninputs):
         raise ValueError("basis set is too small")
     elif (cost is not None or trajectory_constraints is not None) and \
-         basis.N * sys.ninputs == 2 * (sys.nstates + sys.ninputs):
+         ncoefs == 2 * (sys.nstates + sys.ninputs):
         warnings.warn("minimal basis specified; optimization not possible")
         cost = None
         trajectory_constraints = None
@@ -531,11 +539,12 @@ def traj_const(null_coeffs):
 
     # Store the flag lengths and coefficients
     # TODO: make this more pythonic
-    coeff_off = 0
+    coef_off = 0
     for i in range(sys.ninputs):
         # Grab the coefficients corresponding to this flat output
-        systraj.coeffs.append(alpha[coeff_off:coeff_off + basis.N])
-        coeff_off += basis.N
+        coef_len = basis.var_ncoefs(i)
+        systraj.coeffs.append(alpha[coef_off:coef_off + coef_len])
+        coef_off += coef_len
 
         # Keep track of the length of the flat flag for this output
         systraj.flaglen.append(len(zflag_T0[i]))
diff --git a/control/flatsys/poly.py b/control/flatsys/poly.py
@@ -55,7 +55,7 @@ def __init__(self, N):
         super(PolyFamily, self).__init__(N)
 
     # Compute the kth derivative of the ith basis function at time t
-    def eval_deriv(self, i, k, t):
+    def eval_deriv(self, i, k, t, var=None):
         """Evaluate the kth derivative of the ith basis function at time t."""
         if (i < k): return 0;           # higher derivative than power
         return factorial(i)/factorial(i-k) * np.power(t, i-k)
diff --git a/control/flatsys/systraj.py b/control/flatsys/systraj.py
@@ -106,11 +106,11 @@ def eval(self, tlist):
             for i in range(self.ninputs):
                 flag_len = self.flaglen[i]
                 zflag.append(np.zeros(flag_len))
-                for j in range(self.basis.N):
+                for j in range(self.basis.var_ncoefs(i)):
                     for k in range(flag_len):
                         #! TODO: rewrite eval_deriv to take in time vector
                         zflag[i][k] += self.coeffs[i][j] * \
-                            self.basis.eval_deriv(j, k, t)
+                            self.basis.eval_deriv(j, k, t, var=i)
 
             # Now copy the states and inputs
             # TODO: revisit order of list arguments
diff --git a/control/tests/bspline_test.py b/control/tests/bspline_test.py
