visualization.py

"""
Plot matrix based on JSON data written by FRANK library

usage: visualization.py [options]

options:
    -h, --help                Show help message.
    --input=<JSON-file>       Read input from the specified JSON file. This file
                              should be generated by FRANK::write_JSON.
                              [default: matrix.json]
    --sv-tolerance=<float>    Plot singular values that exceed a specified threshold
                              instead of generic patch representation
"""
from docopt import docopt
import sys

import json
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gs
import matplotlib.patches as mpatch

yrblue = '#005396'
yrpink = '#cf006b'
yrgreen = '#539600'
plot_singular_values = False
tolerance = 1e-6
eps = 1e-16


def read_xml(in_file):
    tree = ET.parse(in_file)
    return tree


def plot_hierarchical(node, gs_node, color=False):
    dim = [len(node['children']), len(node['children'][0])]
    level = node['level']
    shared_col_b = [
        {
            'exists':False, 'children':None,
            'level':level+1, 'dim':dim[0],
            'width': None
        }
        for i in range(dim[0])
    ]
    shared_row_b = [
        {
            'exists':False, 'children':None,
            'level':level+1, 'dim':dim[1],
            'width': None
        }
        for j in range(dim[1])
    ]
    gs_subgrid = gs.GridSpecFromSubplotSpec(
        dim[0], dim[1], subplot_spec=gs_node
    )
    for i in range(dim[0]):
        for j in range(dim[1]):
            child = node['children'][i][j]
            sub_type = child['type']
            if sub_type == 'Hierarchical':
                shared_col_b[i]['children'], shared_row_b[j]['children'] = (
                    plot_hierarchical(
                        child,
                        gs_subgrid[i, j],
                        # (level == 0 and [i, j] in [[0, 0], [0, 1]]) or color
                    )
                )
            elif sub_type == 'LowRank':
                if plot_singular_values:
                    plot_lowrank(child, gs_subgrid[i, j])
                else:
                    plot_lowrank_patches(
                        child, gs_subgrid[i, j], dim
                        # (level == 0 and [i, j] in [[0, 0], [0, 1]]) or (color and i >= j)
                    )
            elif sub_type == 'LowRankShared':
                width = plot_lowrank_shared_patches(
                    child, gs_subgrid[i, j], dim
                    # (level == 0 and [i, j] in [[0, 0], [0, 1]]) or (color and i >= j)
                )
                shared_col_b[i]['exists'] = True
                shared_col_b[j]['width'] = width
                shared_row_b[j]['exists'] = True
                shared_row_b[i]['width'] = width
            elif sub_type == 'Dense':
                if plot_singular_values:
                    plot_dense(child, gs_subgrid[i, j])
                else:
                    plot_dense_patch(
                        child, gs_subgrid[i, j]
                        # (level == 0 and [i, j] in [[0, 0], [0, 1]]) or (color and i >= j),
                        # 'lower'
                    )
    return shared_col_b, shared_row_b


def plot_lowrank(node, gs_node):
    dim = node['dim']
    svalues = [x+eps for x in node['svalues']]
    count = 0
    for sv in svalues:
        if sv > tolerance:
            count += 1
    slog = np.log10(svalues[0:count])

    ax = plt.subplot(gs_node)
    ax.text(0.5*count, 0.5, f"{count}", {'fontsize': 'x-small'})
    # ax.text(
    #     0.5, 0.5,
    #     "{}\n({}, {})".format(
    #         node['level'], node['abs_pos'][0], node['abs_pos'][1]
    #     ),
    #     horizontalalignment='center', verticalalignment='center'
    # )
    ylow = int(np.log10(tolerance))
    ax.set(
        xlim=(-0.5, len(slog)+0.5),
        ylim=(0, -(ylow-2)),
        xticks=[], yticks=[]
    )
    ax.bar(np.arange(len(slog)), slog-ylow, width=1, color=yrpink)


def plot_lowrank_patches(node, gs_node, dim, has_color=True):
    level = node['level']
    ax = plt.subplot(gs_node)
    ax.set(xticks=[], yticks=[])
    dist = 0.0075*2**(level*np.sqrt(dim[0]/2)-int(level/4))
    k = (1.0 - 3*dist)/40*2**(level*np.sqrt(dim[0]/2)-int(level/4))
    n = (1.0 - 3*dist-k)
    color = yrpink if has_color else "grey"
    U = mpatch.Rectangle(
        (dist, dist),
        k, n,
        edgecolor='none',
        facecolor=color
    )
    S = mpatch.Rectangle(
        (dist, 2*dist+n),
        k, k,
        edgecolor='none',
        facecolor=color
    )
    V = mpatch.Rectangle(
        (2*dist+k, 2*dist+n),
        n, k,
        edgecolor='none',
        facecolor=color
    )
    # Add the patch to the Axes
    ax.add_patch(U)
    ax.add_patch(S)
    ax.add_patch(V)


def plot_lowrank_shared_patches(node, gs_node, dim, has_color=True):
    level = node['level']
    ax = plt.subplot(gs_node)
    ax.set(xticks=[], yticks=[])
    dist = 0.0075*2**(level*np.sqrt(dim[0]/2)-int(level/4))
    k = (1.0 - 3*dist)/40*2**(level*np.sqrt(dim[0]/2)-int(level/4))
    n = (1.0 - 3*dist-k)
    color = yrpink if has_color else "grey"
    S = mpatch.Rectangle(
        (dist, 2*dist+n),
        k, k,
        edgecolor='none',
        facecolor=color
    )
    # Add the patch to the Axes
    ax.add_patch(S)
    return k


def plot_dense(node, gs_node):
    dim = node['dim']
    svalues = [x+eps for x in node['svalues']]
    count = 0
    for sv in svalues:
        if sv > tolerance:
            count += 1
    slog = np.log10(svalues[0:count])

    ax = plt.subplot(gs_node)
    ax.text(0.5*count, 0.5, f"{count}", {'fontsize': 'x-small'})
    # ax.text(
    #     0.5, 0.5,
    #     "{}\n({}, {})".format(
    #         node['level'], node['abs_pos'][0], node['abs_pos'][1]
    #     ),
    #     horizontalalignment='center', verticalalignment='center'
    # )
    ylow = int(np.log10(tolerance))
    ax.set(
        xlim=(-0.5, len(slog)+0.5),
        ylim=(0, -(ylow-2)),
        xticks=[], yticks=[]
    )
    ax.bar(np.arange(len(slog)), slog-ylow, width=1, color=yrblue)


def plot_dense_patch(node, gs_node, has_color=True, ul=None):
    ax = plt.subplot(gs_node)
    ax.set(xticks=[], yticks=[])
    color = yrblue if has_color else "grey"
    if ul is None or not has_color:
        patch = mpatch.Rectangle(
            (0.0, 0.0),
            1.0, 1.0,
            edgecolor='none',
            facecolor=color
        )
        # Add the patch to the Axes
        ax.add_patch(patch)
    else:
        upper = mpatch.Polygon(
            [[0.0, 1.0], [1.0, 1.0], [1.0, 0.0]],
            edgecolor='none',
            facecolor=color if ul == 'upper' else 'grey'
        )
        ax.add_patch(upper)
        lower = mpatch.Polygon(
            [[0.0, 1.0], [0.0, 0.0], [1.0, 0.0]],
            edgecolor='none',
            facecolor=color if ul == 'lower' else 'grey'
        )
        ax.add_patch(lower)

def plot_shared_basis(shared_basis, gs_root, transpose, width_factor):
    gs_split = gs.GridSpecFromSubplotSpec(
        1 if transpose else len(shared_basis),
        len(shared_basis) if transpose else 1,
        subplot_spec=gs_root
    )
    child_depths = [0]
    for i, basis in enumerate(shared_basis):
        gs_own, gs_children = None, None
        if not basis['exists'] and basis['children'] is None:
            continue
        if basis['exists'] and basis['children'] is None:
            gs_own = gs_split[i]
            gs_children = None
        elif not basis['exists'] and basis['children'] is not None:
            gs_own = None
            gs_children = gs_split[i]
        elif basis['exists'] and basis['children'] is not None:
            gs_own, gs_children = gs.GridSpecFromSubplotSpec(
                2 if transpose else 1,
                1 if transpose else 2,
                subplot_spec=gs_split[i]
            )
        else:
            print(basis['exists'], basis['children'])
            raise ValueError
        depth = 0
        if gs_children is not None:
            depth += plot_shared_basis(
                basis['children'], gs_children, transpose, width_factor
            )
            if gs_own is not None:
                if transpose:
                    gs_split[i].set_height_ratios(1, depth)
                else:
                    gs_split[i].set_width_ratios(1, depth)
        if gs_own is not None:
            depth += 1
            ax = plt.subplot(gs_own)
            ax.axis('off')
            level = basis['level']
            dim = basis['dim']
            width = basis['width']*width_factor/dim
            base = mpatch.Rectangle(
                (
                    0.05 if transpose else 0.5-width/2,
                    0.5-width/2 if transpose else 0.05
                ),
                0.9 if transpose else width, width if transpose else 0.9,
                edgecolor='none',
                facecolor=yrpink
            )
            ax.add_patch(base)
        child_depths.append(depth)
    return max(child_depths)


def main():
    args = docopt(__doc__)
    in_path = args["--input"]
    with open(in_path) as in_file:
        root = json.load(in_file)
    if(args["--sv-tolerance"]):
        global tolerance
        tolerance = float(args['--sv-tolerance'])
        global plot_singular_values
        plot_singular_values = True

    fig = plt.figure(figsize=(5, 5), dpi=200)
    grid = fig.add_gridspec(2, 2)
    shared_col_b, shared_row_b = plot_hierarchical(
        root,
        grid[1, 1]
    )
    width_factor = 20
    col_basis_width = plot_shared_basis(
        shared_col_b, grid[1, 0], False, width_factor
    )
    grid.set_width_ratios([col_basis_width, width_factor])
    row_basis_height = plot_shared_basis(
        shared_row_b, grid[0, 1], True, width_factor
    )
    grid.set_height_ratios([row_basis_height, width_factor])

    # TODO change figure size if width of bases varies (.set_size_inches)
    plt.tight_layout()
    plt.subplots_adjust(wspace=0, hspace=0)
    plt.savefig(in_path.split('.')[0] + ".pdf")
    # plt.show()


if __name__ == '__main__':
    main()