polytope_plots.sage

from sage.all import *

# Backend Sage machinery to produce plots of polytopes
@cached_method
def sliced_cones_2(poly, eqn, vecprojector):
    """
    Return two-dimensional slices of the normal fan of a polytope along a specified plane (helper method)

    INPUT:
    - ``eqn`` -- Equation defining a plane in R3
    - ``vecprojector`` -- Anonymous function projecting a vector from R3 to R2
    """
    slices3 = poly.d1_slices()
    subspace = Polyhedron(eqns = [eqn])

    slices2 = set()
    for slice3 in slices3:
        # Find the intersection of this cone with the hyperplanes
        slice2 = slice3.intersection(subspace)

        # If the result does not have any area, we're done with this slice, so we move on to the next one
        if slice2.dimension() < 2:
            continue

        # The resulting coneslice is a two-dimensional polyhedron, but it still lives in Q3, so let's project down to Q2
        newverts = lambda polyslice: [vecprojector(vert) for vert in polyslice.vertices()] # Projects the defining vertices of a slice into Q2
        newrays = lambda polyslice: [vecprojector(ray) for ray in polyslice.rays()] # Projects the defining rays of a slice into Q2
        polyprojector = lambda polyslice: Polyhedron(vertices = newverts(polyslice), rays = newrays(polyslice)) # Projects a polytope from Q3 to Q2 using the helper functions we just defined

        # We can then apply polyprojector to our slice to move it down to Q2
        projslice = polyprojector(slice2)

        # For future use, we'll store the original cone in Q4 and the slice in Q3 as a member of this object
        projslice.original_cone = slice3.original_cone
        projslice.original_slice = slice3

        # As well as the original vertex and structure associated to those cones
        projslice.original_vertex = poly.vertex_from_cone(slice3.original_cone)
        projslice.original_structure = poly[projslice.original_vertex]

        # And, finally, we add this slice to our list
        slices2.add(projslice)

    # Once we've finished the loop, all the slices have been processed, so we return the resulting container
    return slices2

@cached_method
def sliced_cones_a(poly, a = 0):
    """
    Return the two-dimensional slices of the normal fan of a polytope with fixed a.

    INPUT:
    - ``poly`` -- A Polytope in Q4
    - ``a`` -- (default: 0) The fixed value of a at which to slice
    """
    # a needs to be rational
    a = QQ(a)

    eqn = (-a, 1, 0, 0)
    vecprojector = lambda vec: [vec[1], vec[2]]

    return sliced_cones_2(poly, eqn, vecprojector)

@cached_method
def sliced_cones_b(poly, b = 0):
    """
    Return the two-dimensional slices of the normal fan of a polytope with fixed b.

    INPUT:
    - ``poly`` -- A Polytope in Q4
    - ``b`` -- (default: 0) The fixed value of b at which to slice
    """
    # b needs to be rational
    b = QQ(b)

    eqn = (-b, 0, 1, 0)
    vecprojector = lambda vec: [vec[0], vec[2]]
    return sliced_cones_2(poly, eqn, vecprojector)

@cached_method
def sliced_cones_c(poly, c = 0):
    """
    Return the two-dimensional slices of the normal fan of a polytope with fixed a.

    INPUT:
    - ``poly`` -- A Polytope in Q4
    - ``c`` -- (default: 0) The fixed value of c at which to slice
    """
    # a needs to be rational
    c = QQ(c)

    eqn = (-c, 0, 0, 1)
    vecprojector = lambda vec: [vec[0], vec[1]]
    return sliced_cones_2(poly, eqn, vecprojector)

@cached_method
def region_graph(regions):
    import itertools
    from collections import defaultdict
    edge_dict = {region:[] for region in regions}
    for pair in itertools.combinations(regions, 2):
        a, b = pair
        if a != b and not a.intersection(b).is_empty():
            edge_dict[a].append(b)

    result = Graph(edge_dict)
    return result

@cached_method
def bounded_regions(regions, points_to_include, xbounds = None, ybounds = None):
    """
    Plot a collection of cone slices in R2
    """
    import itertools

    vertices = [vert for region in regions for vert in region.vertices()] # Find all the vertices to help us construct the bounding box
    vertices.extend(points_to_include) # Append any marked points to ensure they are visible

    # Identify the maximum values of x and y for the box
    xvals = [abs(vert[0]) for vert in vertices]
    xmax = max(xvals)
    xmin = -xmax

    yvals = [abs(vert[1]) for vert in vertices]
    ymax = max(yvals)
    ymin = -ymax

    auto_box_scale = 3/2
    if xbounds is None:
        xbounds = [auto_box_scale * xmin, auto_box_scale * xmax]

    if ybounds is None:
        ybounds = [auto_box_scale * ymin, auto_box_scale * ymax]

    bounding_box = Polyhedron(itertools.product(xbounds, ybounds))

    bounded_regions = set()
    for region in regions:
        bounded_region = bounding_box.intersection(region)

        # Discard the region if it's empty
        if bounded_region.dimension() < 2:
            continue

        # Otherwise, add some metadata and insert it in the result set
        bounded_region.original_cone = region.original_cone # add the original cone (fidbc)
        bounded_region.original_region = region # also add the original region [maybe redundant!] (fidbc)
        bounded_region.original_vertex = region.original_vertex
        bounded_region.original_structure = region.original_structure
        bounded_regions.add(bounded_region)

    return bounded_regions

# Colorblind-friendly palette from http://bconnelly.net/2013/10/creating-colorblind-friendly-figures/
six_nice_colors = (Color(0.9, 0.6, 0), Color(0.35, 0.7, 0.9), Color(0, 0.6, 0.5), Color(0.95, 0.9, 0.25), Color(0, 0.45, 0.7), Color(0.8, 0.4, 0), Color(0.8, 0.6, 0.7))

@cached_method
def colored_regions(regions, points_to_mark, xbounds = None, ybounds = None, plot_colors = six_nice_colors):
    """
    """
    from sage.graphs.graph_coloring import first_coloring

    # Get the regions
    regions = bounded_regions(regions, points_to_mark, xbounds, ybounds)

    # Color the slices using a graph coloring algorithm
    G = region_graph(regions)

    n_colors = len(plot_colors)
    region_color_list = first_coloring(G)#, n_colors)
    region_color_dict = {region: plot_colors[i] for i in range(len(region_color_list)) for region in region_color_list[i]}

    # Store the colors
    for region in regions:
        region.color = region_color_dict[region]

    return regions

#@cached_method
def accuracy_colored_regions(regions, accuracy, points_to_mark, xbounds = None, ybounds = None, colormap = colormaps['hot']):
    """
    """
    # Get the regions
    regions = bounded_regions(regions, points_to_mark, xbounds, ybounds)

    # Color the slices using a graph coloring algorithm
    region_color_dict = {region: colormap(accuracy[tuple(region.original_vertex)])[:3] for region in regions}

    # Store the colors
    for region in regions:
        region.color = region_color_dict[region]

    return regions

@cached_method
def region_plots(regions, points_to_mark, plot_title = None, add_slice_count_to_title = True, xbounds = None, xlabel = None, ybounds = None, ylabel = None, figsize = (4, 3), plot_colors = six_nice_colors):
    """
    """

    plotted_points = [point2d(point, size=30, zorder=10, color=colors['black']) for point in points_to_mark]
    plotted_regions = [region.projection().render_fill_2d(color = region.color) for region in colored_regions(regions, points_to_mark, xbounds, ybounds, plot_colors = plot_colors)]

    if add_slice_count_to_title:
        plot_title += ": " + str(len(plotted_regions)) + " slices visible"

    #theplot = plot(plotted_points + plotted_regions, frame = True, axes_labels = [xlabel, ylabel], title = plot_title, figsize = figsize)
    theplot = sum(plotted_regions)
    for pp in plotted_points:
        theplot += pp
    theplot.SHOW_OPTIONS['frame']=True
    theplot.SHOW_OPTIONS['axes_labels']=[xlabel, ylabel]
    theplot.SHOW_OPTIONS['title']=plot_title
    theplot.SHOW_OPTIONS['figsize']=figsize
    theplot.axes_labels_size(1)
    theplot.fontsize(8)
    return theplot

#@cached_method
def accuracy_region_plots(regions, accuracy, points_to_mark,  plot_title = None, add_slice_count_to_title = True, xbounds = None, xlabel = None, ybounds = None, ylabel = None, figsize = (4, 3), colormap=colormaps['hot']):
    """
    """

    plotted_points = [point2d(point, size=30, zorder=10, color=colors['black']) for point in points_to_mark]
    plotted_regions = [region.projection().render_fill_2d(color = region.color) for region in accuracy_colored_regions(regions, accuracy, points_to_mark, xbounds, ybounds, colormap)]

    if add_slice_count_to_title:
        plot_title += ": " + str(len(plotted_regions)) + " slices visible"

    #theplot = plot(plotted_points + plotted_regions, frame = True, axes_labels = [xlabel, ylabel], title = plot_title, figsize = figsize)
    theplot = sum(plotted_regions)
    for pp in plotted_points:
        theplot += pp
    theplot.SHOW_OPTIONS['frame']=True
    theplot.SHOW_OPTIONS['axes_labels']=[xlabel, ylabel]
    theplot.SHOW_OPTIONS['title']=plot_title
    theplot.SHOW_OPTIONS['figsize']=figsize
    theplot.axes_labels_size(1)
    theplot.fontsize(8)
    return theplot

alabel = "multibranch loop initiation penalty ($a$)"
blabel = "unpaired base penalty ($b$)"
clabel = "branching helix penalty ($c$)"

@cached_method
def sliced_cone_plots_a(poly, a = 0, bbounds = None, cbounds = None, plot_title = None, plot_colors = six_nice_colors):
    """
    Plot the two-dimensional slices of the normal fan of a polytope with fixed a.

    INPUT:
    - ``poly`` -- A Polytope in R4
    - ``a`` -- (default: 0) The fixed value of a to use
    - ``bbounds`` -- (default: automatic) The range of b values to plot
    - ``cbounds`` -- (default: automatic) The range of c values to plot
    - ``plot_title`` -- (default: None) Title for the plot
    """
    # Get the slices
    slices = sliced_cones_a(poly, a)

    classical_points = ((0, 0.4),)

    if plot_title is not None:
        formatted_title = plot_title + "\n"
    else:
        formatted_title = ""

    formatted_title += "a = " + str(a)

    plots = region_plots(slices,
                         points_to_mark = classical_points,
                         plot_title = formatted_title,
                         add_slice_count_to_title = True,
                         xbounds = bbounds, xlabel = blabel,
                         ybounds = cbounds, ylabel = clabel,
                         plot_colors = plot_colors)

    return plots

#@cached_method
def sliced_cone_plots_b(poly, b = 0, abounds = None, cbounds = None, plot_title = None, plot_colors = six_nice_colors):
    """
    Plot the two-dimensional slices of the normal fan of a polytope with fixed b.

    INPUT:
    - ``poly`` -- A Polytope in R4
    - ``b`` -- (default: 0) The fixed value of b to use
    - ``abounds`` -- (default: automatic) The range of a values to plot
    - ``cbounds`` -- (default: automatic) The range of c values to plot
    - ``plot_title`` -- (default: None) Title for the plot
    """
    # Get the slices
    slices = sliced_cones_b(poly, b)

    classical_points = ((3.4, 0.4),)

    if plot_title is not None:
        formatted_title = plot_title + "\n"
    else:
        formatted_title = ""

    formatted_title += "b = " + str(b)

    plots = region_plots(slices,
                         points_to_mark = classical_points,
                         plot_title = formatted_title,
                         add_slice_count_to_title = True,
                         xbounds = abounds, xlabel = alabel,
                         ybounds = cbounds, ylabel = clabel,
                         plot_colors = plot_colors)

    return plots

def paper_sliced_cone_plots_b(poly, b = 0, abounds = None, cbounds = None, plot_title = None, plot_colors = six_nice_colors):
    """
    Plot the two-dimensional slices of the normal fan of a polytope with fixed b.

    INPUT:
    - ``poly`` -- A Polytope in R4
    - ``b`` -- (default: 0) The fixed value of b to use
    - ``abounds`` -- (default: automatic) The range of a values to plot
    - ``cbounds`` -- (default: automatic) The range of c values to plot
    - ``plot_title`` -- (default: None) Title for the plot
    """
    # Get the slices
    slices = sliced_cones_b(poly, b)

    classical_points = ((3.4, 0.4),)

    if plot_title is not None:
        formatted_title = plot_title + "\n"
    else:
        formatted_title = ""

    #formatted_title += "b = " + str(b)

    plots = region_plots(slices,
                         points_to_mark = classical_points,
                         plot_title = formatted_title,
                         add_slice_count_to_title = False,
                         xbounds = abounds, xlabel = alabel,
                         ybounds = cbounds, ylabel = clabel,
                         plot_colors = plot_colors)

    return plots

def accuracy_cone_plots_b(poly, accuracy, b = 0, abounds = None, cbounds = None, plot_title = None,colormap=colormaps['hot']):
    """
    Plot the two-dimensional slices of the normal fan of a polytope with fixed b.

    INPUT:
    - ``poly`` -- A Polytope in R4
    - ``b`` -- (default: 0) The fixed value of b to use
    - ``abounds`` -- (default: automatic) The range of a values to plot
    - ``cbounds`` -- (default: automatic) The range of c values to plot
    - ``plot_title`` -- (default: None) Title for the plot
    """
    # Get the slices
    slices = sliced_cones_b(poly, b)

    classical_points = ((3.4, 0.4),)

    if plot_title is not None:
        formatted_title = plot_title + "\n"
    else:
        formatted_title = ""

    formatted_title += "b = " + str(b)

    plots = accuracy_region_plots(slices,
                         accuracy,
                         points_to_mark = classical_points,
                         plot_title = formatted_title,
                         add_slice_count_to_title = True,
                         xbounds = abounds, xlabel = alabel,
                         ybounds = cbounds, ylabel = clabel, colormap=colormap)

    return plots

@cached_method
def sliced_cone_plots_c(poly, c = 0, abounds = None, bbounds = None, plot_title = None, plot_colors = six_nice_colors):
    """
    Plot the two-dimensional slices of the normal fan of a polytope with fixed c.

    INPUT:
    - ``poly`` -- A Polytope in R4
    - ``c`` -- (default: 0) The fixed value of c to use
    - ``abounds`` -- (default: automatic) The range of a values to plot
    - ``bbounds`` -- (default: automatic) The range of b values to plot
    - ``plot_title`` -- (default: None) Title for the plot
    """
    # Get the slices
    slices = sliced_cones_c(poly, c)

    classical_points = ((3.4, 0),)

    if plot_title is not None:
        formatted_title = plot_title + "\n"
    else:
        formatted_title = ""

    formatted_title += "c = " + str(c)

    plots = region_plots(slices,
                         points_to_mark = classical_points,
                         plot_title = formatted_title,
                         add_slice_count_to_title = True,
                         xbounds = abounds, xlabel = alabel,
                         ybounds = bbounds, ylabel = blabel,
                         plot_colors = plot_colors)

    return plots