frc971/control_loops/python/polytope.py

#!/usr/bin/python3
"""
Polyhedral set library.

This library implements convex regions of the form H x <= k, where H, x, and k
are matricies.  It also provides convenient methods to find all the verticies.
"""

__author__ = 'Austin Schuh (austin.linux@gmail.com)'

from frc971.control_loops.python import libcdd
import numpy
import string
import sys


def _PiecewiseConcat(*args):
    """Concatenates strings inside lists, elementwise.

  Given ['a', 's'] and ['d', 'f'], returns ['ad', 'sf']
  """
    return map(''.join, zip(*args))


def _SplitAndPad(s):
    """Splits a string on newlines, and pads the lines to be the same width."""
    split_string = s.split('\n')
    width = max(len(stringpiece) for stringpiece in split_string) + 1

    padded_strings = [
        stringpiece.ljust(width, ' ') for stringpiece in split_string
    ]
    return padded_strings


def _PadHeight(padded_array, min_height):
    """Adds lines of spaces to the top and bottom of an array symmetrically."""
    height = len(padded_array)
    if height < min_height:
        pad_array = [' ' * len(padded_array[0])]
        height_error = min_height - height
        return (pad_array * int(
            (height_error) / 2) + padded_array + pad_array * int(
                (height_error + 1) / 2))
    return padded_array


class HPolytope(object):
    """This object represents a H-polytope.

  Polytopes are convex regions in n-dimensional space.
  For H-polytopes, this is represented as the intersection of a set of half
  planes.  The mathematic equation that represents this is H x <= k.
  """

    def __init__(self, H, k):
        """Constructs a H-polytope from the H and k matricies.

    Args:
      H: numpy.Matrix (n by k), where n is the number of constraints, and k is
        the number of dimensions in the space.  Does not copy the matrix.
      k: numpy.Matrix (n by 1).  Does not copy the matrix.
    """
        self._H = H
        self._k = k

    @property
    def k(self):
        """Returns the k in H x <= k."""
        return self._k

    @property
    def H(self):
        """Returns the H in H x <= k."""
        return self._H

    @property
    def ndim(self):
        """Returns the dimension of the set."""
        return self._H.shape[1]

    @property
    def num_constraints(self):
        """Returns the number of constraints defining the set."""
        return self._k.shape[0]

    def IsInside(self, point):
        """Returns true if the point is inside the polytope, edges included."""
        return (self._H * point <= self._k).all()

    def Vertices(self):
        """Returns a matrix with the vertices of the set in its rows."""
        # TODO(aschuh): It would be better to write some small C/C++ function that
        # does all of this and takes in a numpy array.
        # The copy is expensive in Python and cheaper in C.

        # Create an empty matrix with the correct size.
        matrixptr = libcdd.dd_CreateMatrix(self.num_constraints, self.ndim + 1)
        matrix = matrixptr.contents

        try:
            # Copy the data into the matrix.
            for i in range(self.num_constraints):
                libcdd.dd_set_d(matrix.matrix[i][0], self._k[i, 0])
                for j in range(self.ndim):
                    libcdd.dd_set_d(matrix.matrix[i][j + 1], -self._H[i, j])

            # Set enums to the correct values.
            matrix.representation = libcdd.DD_INEQUALITY
            matrix.numbtype = libcdd.DD_REAL

            # TODO(aschuh): Set linearity if it is useful.
            # This would be useful if we had any constraints saying B - A x = 0

            # Build a Polyhedra
            polyhedraptr = libcdd.dd_DDMatrix2Poly(matrixptr)

            if not polyhedraptr:
                return None

            try:
                # Return None on error.
                # The error values are enums, so they aren't exposed.

                # Magic happens here.  Computes the vertices
                vertex_matrixptr = libcdd.dd_CopyGenerators(polyhedraptr)
                vertex_matrix = vertex_matrixptr.contents

                try:
                    # Count the number of vertices and rays in the result.
                    num_vertices = 0
                    num_rays = 0
                    for i in range(vertex_matrix.rowsize):
                        if libcdd.dd_get_d(vertex_matrix.matrix[i][0]) == 0:
                            num_rays += 1
                        else:
                            num_vertices += 1

                    # Build zeroed matricies for the new vertices and rays.
                    vertices = numpy.matrix(
                        numpy.zeros((num_vertices, vertex_matrix.colsize - 1)))
                    rays = numpy.matrix(
                        numpy.zeros((num_rays, vertex_matrix.colsize - 1)))

                    ray_index = 0
                    vertex_index = 0

                    # Copy the data out of the matrix.
                    for index in range(vertex_matrix.rowsize):
                        if libcdd.dd_get_d(
                                vertex_matrix.matrix[index][0]) == 0.0:
                            for j in range(vertex_matrix.colsize - 1):
                                rays[ray_index, j] = libcdd.dd_get_d(
                                    vertex_matrix.matrix[index][j + 1])
                            ray_index += 1
                        else:
                            for j in range(vertex_matrix.colsize - 1):
                                vertices[vertex_index, j] = libcdd.dd_get_d(
                                    vertex_matrix.matrix[index][j + 1])
                            vertex_index += 1
                finally:
                    # Free everything.
                    libcdd.dd_FreeMatrix(vertex_matrixptr)

            finally:
                libcdd.dd_FreePolyhedra(polyhedraptr)

        finally:
            libcdd.dd_FreeMatrix(matrixptr)

        # Rays are unsupported right now.  This may change in the future.
        assert (rays.shape[0] == 0)

        return vertices

    def __str__(self):
        """Returns a formatted version of the polytope.

    The dump will look something like the following, which prints out the matrix
    comparison.

    [[ 1  0]            [[12]
     [-1  0]  [[x0]  <=  [12]
     [ 0  1]   [x1]]     [12]
     [ 0 -1]]            [12]]
    """
        height = max(self.ndim, self.num_constraints)

        # Split the print up into 4 parts and concatenate them all.
        H_strings = _PadHeight(_SplitAndPad(str(self.H)), height)
        v_strings = _PadHeight(_SplitAndPad(str(self.k)), height)
        x_strings = _PadHeight(self._MakeXStrings(), height)
        cmp_strings = self._MakeCmpStrings(height)

        return '\n'.join(
            _PiecewiseConcat(H_strings, x_strings, cmp_strings, v_strings))

    def _MakeXStrings(self):
        """Builds an array of strings with constraint names in it for printing."""
        x_strings = []
        if self.ndim == 1:
            x_strings = ["[[x0]] "]
        else:
            for index in range(self.ndim):
                if index == 0:
                    x = "[[x%d]  " % index
                elif index == self.ndim - 1:
                    x = " [x%d]] " % index
                else:
                    x = " [x%d]  " % index
                x_strings.append(x)
        return x_strings

    def _MakeCmpStrings(self, height):
        """Builds an array of strings with the comparison in it for printing."""
        cmp_strings = []
        for index in range(height):
            if index == int((height - 1) / 2):
                cmp_strings.append("<= ")
            else:
                cmp_strings.append("   ")
        return cmp_strings