utils/hypergraph_drawer.py

from turtle import title
from typing import Dict
import hypernetx as hnx
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from io import BytesIO

def pyplot_fig_to_buffer(fig):
    """Draw the figure using Pyplot and return the image buffer."""
    canvas = FigureCanvas(fig)
    buffer = BytesIO()
    canvas.print_png(buffer)
    buffer.seek(0)
    return buffer


def draw_hypergraph(
        H:hnx.Hypergraph=None,
        hyperedges:Dict = None, 
        draw_dual:bool=False,
        toplexes:bool=False,
        fill_edges:bool=True,
        with_edge_labels:bool=True, 
        with_node_labels:bool=True, 
        title:str=None,
        pos = None
    ,draw_width:int = 6, draw_height:int = 6
                    ):
    # 构建超图
    if H is None and hyperedges is None:
        raise ValueError("Either H or hyperedges must be provided.")
    if H is None:
        H = hnx.Hypergraph(hyperedges)

    if draw_dual:
        H = H.dual()
    if toplexes:
        H = H.toplexes(return_hyp=True)

    # 绘制超图
    fig, ax = plt.subplots(figsize=(draw_width, draw_height))
    if title is not None:
        ax.set_title(title)
    pos = hnx.draw(H, ax=ax, fill_edges=fill_edges, 
                with_edge_labels=with_edge_labels,
            with_node_labels=with_node_labels,
                node_label_alpha=0.0, 
                edge_label_alpha=0.0, pos=pos,
                return_pos=True)
        # hnx.draw(H, ax=ax)

    # 将超图保存为图像
    return pyplot_fig_to_buffer(fig), pos