Python Graphviz调整边宽实现马尔可夫链转移概率可视化问题

问题原因 #

你遇到的报错是两个问题共同导致的：

1. graphviz要求所有边/节点属性的值必须是字符串类型，你直接传入numpy浮点型数值会触发类型错误
2. 你混淆了graphviz的属性作用：weight属性用于布局阶段控制边的优先级，不是用来控制边宽度的，控制边宽度的属性为penwidth

疑问解答 #

• 可以通过graphviz的Python接口实现你需要的马尔可夫链可视化效果，无需更换工具
• 关联权重到边宽的方法：将转移概率乘以一个自定义缩放系数（避免概率过低导致边太细不可见），转换为字符串后传入边的penwidth参数即可，也可以同时将概率值转为字符串传入label/tooltip参数作为补充信息
• 绘制图表：构建完图结构后可以直接调用view()方法弹出预览，也可以调用render()方法导出为PDF/PNG/SVG等格式文件；如果需要在matplotlib的axes中展示，可以先将图渲染为图片再用matplotlib读取绘制。

修正后完整代码 #

import graphviz
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from PIL import Image


class MarkovViz:
    """
    Visualize the transition probability matrix of a Markov chain as a directed
    graph, where the width of an edge is proportional to the transition
    probability between two states.
    """

    def __init__(self, transition_probability_matrix=None, penwidth_scale=5.0):
        self._graph = None
        # 笔宽缩放系数，可根据需求调整
        self.penwidth_scale = penwidth_scale
        if transition_probability_matrix is not None:
            self.build_from_matrix(transition_probability_matrix)

    def build_from_matrix(self, trans, labels=None):
        """
        Args:
          trans: A pd.DataFrame or 2D np.array.  A square matrix containing the
            conditional probabability of a transition from the level
            represented by the row to the level represented by the column.
            Each row sums to 1.
          labels: A list-like sequence of labels to use for the rows and
            columns of 'trans'.  If trans is a pd.DataFrame or similar then
            this entry can be None and labels will be taken from the column
            names of 'trans'.

        Effects:
          self._graph is created as a directed graph, and populated with nodes
            and edges, with edge width taken from 'trans'.
        """

        if labels is None and hasattr(trans, "columns"):
            labels = list(trans.columns)
            index = list(trans.index)
            if labels != index:
                raise Exception("Mismatch between index and columns of "
                                "the transition probability matrix.")
            trans = trans.values

        trans = np.array(trans)
        self._graph = graphviz.Digraph("MarkovChain", format="png")

        dim = trans.shape[0]
        if trans.shape[1] != dim:
            raise Exception("Matrix must be square")

        # 先添加所有节点
        for label in labels:
            self._graph.node(label)

        # 添加边
        for i in range(dim):
            for j in range(dim):
                prob = trans[i, j]
                if prob > 0:
                    penwidth = str(prob * self.penwidth_scale)
                    # 可选添加概率标签，不需要可以删掉label参数
                    self._graph.edge(
                        labels[i], 
                        labels[j], 
                        penwidth=penwidth,
                        label=f"{prob:.2f}"
                    )

    def plot(self, ax: plt.Axes = None):
        if ax is None:
            self._graph.view()
            return
        # 渲染为临时图片后加载到matplotlib axes
        tmp_path = self._graph.render(tempfile=True)
        img = Image.open(tmp_path)
        ax.imshow(img)
        ax.axis("off")

运行示例 #

# 测试用转移矩阵
df = pd.DataFrame(
    [[0.5, 0.5, 0], [0.0, 0.0, 1], [1.0, 0.0, 0]],
    index=["foo", "bar", "baz"],
    columns=["foo", "bar", "baz"]
)

# 初始化可视化对象
viz = MarkovViz(df, penwidth_scale=5)
# 直接弹出预览
viz.plot()
# 或者在matplotlib中展示
# fig, ax = plt.subplots()
# viz.plot(ax)
# plt.show()

内容的提问来源于stack exchange，提问作者Steven Scott