Python+Matplotlib绘制小提琴图的示例代码

#!/usr/bin/env python
# -- coding: utf-8 --
"""
Copyright (c) 2022. All rights reserved.
Created by C. L. Wang on 2023/6/6
"""
import os

import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns

from myutils.project_utils import read_excel_to_df
from root_dir import DATA_DIR


def draw_violin_plot(
  df, score_col_name, label_col_name,
  x_label="", y_label="", title="",
  is_show=False, save_name=""):
  """
  绘制小提琴图
  :param df: 数据格式，至少包括两列，即数据值列score，标签列label，相同类别的标签相同
  :param score_col_name: 数值列名称
  :param label_col_name: 标签列名称
  :param x_label: 显示的x标签
  :param y_label: 显示的y标签
  :param title: 显示的图名称
  :param is_show:  是否IDE显示
  :param save_name: 是否存储文件，tight格式
  :return:
  """
  plt.figure(figsize=(16, 10), dpi=80)
  sns.violinplot(x=label_col_name, y=score_col_name, data=df, inner='quartile')

  plt.xlabel(xlabel=x_label)
  plt.ylabel(ylabel=y_label)

  if title:
    plt.title(title, fontsize=12)
  if save_name:
    # transparent=True
    plt.savefig(save_name, bbox_inches='tight', format='png')
  if is_show:
    plt.show()

def main():
  df = read_excel_to_df(os.path.join(DATA_DIR, "ourbest_20230601_tmscore_56.xls"))
  # df = read_excel_to_df(os.path.join(DATA_DIR, "ourbest_20230605_dockq_9_final.xls"))
  # df.info()

  # 加工数据格式，满足数值和标签的格式
  score1 = df["m0-score"]
  score2 = df["m1-score"]
  score3 = df["cur-score"]
  label_col = ["SOTA" for _ in range(df["m0-score"].size)] + \
              ["Baseline" for _ in range(df["m1-score"].size)] + \
              ["Our Best" for _ in range(df["cur-score"].size)]
  score_list = pd.concat([score1, score2, score3])
  score_name = "DockQ"
  data = {score_name: score_list, "label": label_col}
  new_df = pd.DataFrame(data)
  # df.info()
  draw_violin_plot(new_df, score_col_name=score_name, label_col_name="label",
                   x_label="", y_label=score_name,
                   is_show=True, save_name="xxx.png")


if __name__ == '__main__':
  main()