vis.py 文件源码

def visualize_world(WS_d, WS_node_dict, name=None):
    #----visualization for square-partitioned-workspace----
    #----see case_study.py for details----
    figure = plt.figure()
    ax = figure.add_subplot(1,1,1)
    #----- draw the workspace
    for node, prop in WS_node_dict.iteritems():
        if frozenset(['base1','base']) in prop.keys():
            text = '$Base1$'
            color = 'yellow'
        elif frozenset(['base2','base']) in prop.keys():
            text = '$Base2$'
            color = 'yellow'
        elif frozenset(['base3','base']) in prop.keys():
            text = '$Base3$'
            color = 'yellow'
        elif frozenset(['obstacle','low']) in prop.keys():
            text = '$Obs: 0.2$'
            color = '#ff8000'
        elif frozenset(['obstacle','top']) in prop.keys():
            text = '$Obs: 1.0$'
            color = 'red'
        elif frozenset(['supply',]) in prop.keys():
            text = '$Sply: %s$' %str(prop[frozenset(['supply',])])
            if prop[frozenset(['supply',])]>=0.8:
                color = '#0000ff'
            elif prop[frozenset(['supply',])]>=0.6:
                color = '#0040ff'
            elif prop[frozenset(['supply',])]>=0.4:
                color = '#0080ff'
            elif prop[frozenset(['supply',])]>=0.2:
                color = '#00bfff'
        else:
            text = None
            color = 'white'
        rec = matplotlib.patches.Rectangle((node[0]-WS_d, node[1]-WS_d),
                                               WS_d*2, WS_d*2,
                                               fill = True,
                                               facecolor = color,
                                               edgecolor = 'black',
                                               linewidth = 1,
                                               ls = '--',
                                               alpha =0.8)
        ax.add_patch(rec)
        if text:
            ax.text(node[0]-0.7, node[1], r'%s' %text, fontsize = 10, fontweight = 'bold')
    ax.set_aspect('equal')
    ax.set_xlim(0, 10)
    ax.set_ylim(0, 10)
    ax.set_xlabel(r'$x(m)$')
    ax.set_ylabel(r'$y(m)$')    
    if name:
        plt.savefig('%s.pdf' %name,bbox_inches='tight')
    return figure