常用作图脚本

22 Nov 2023

1840字

7分
CC BY 4.0 （除特别声明或转载文章外）
如果这篇博客帮助到你，可以请我喝一杯咖啡~

准备一套作图脚本供日后使用。目标包括：

• 直出 pdf
• 文字数据可选中
• 矢量图，放缩无损
• 尽量黑白打印不影响

{
  "$schema": "https://vega.github.io/schema/vega-lite/v5.15.1.json",
  "config": {
    "axisX": {
      "tickMinStep": 50
    },
    "axisY": {
      "tickMinStep": 10000
    },
    "view": {
      "continuousHeight": 300,
      "continuousWidth": 300
    }
  },
  "data": {
    "name": "data-90c5cf575d993e0ad54f896cfbb5c8e6"
  },
  "datasets": {
    "data-90c5cf575d993e0ad54f896cfbb5c8e6": [
      {
        "GPU": "K40",
        "L1": 64,
        "L2": 1536
      },
      {
        "GPU": "P100",
        "L1": 64,
        "L2": 4096
      },
      {
        "GPU": "V100",
        "L1": 128,
        "L2": 6144
      },
      {
        "GPU": "A100",
        "L1": 192,
        "L2": 40960
      },
      {
        "GPU": "H100",
        "L1": 256,
        "L2": 51200
      }
    ]
  },
  "height": 300,
  "layer": [
    {
      "encoding": {
        "text": {
          "field": "GPU",
          "type": "nominal"
        },
        "x": {
          "field": "L1",
          "scale": {
            "domain": [
              50,
              250
            ]
          },
          "title": "L1D+SMEM/KB per SM",
          "type": "quantitative"
        },
        "y": {
          "field": "L2",
          "scale": {
            "domain": [
              0,
              50000
            ]
          },
          "title": "L2/KB",
          "type": "quantitative"
        }
      },
      "mark": {
        "color": "grey",
        "type": "line"
      }
    },
    {
      "encoding": {
        "text": {
          "field": "GPU",
          "type": "nominal"
        },
        "x": {
          "field": "L1",
          "scale": {
            "domain": [
              50,
              250
            ]
          },
          "title": "L1D+SMEM/KB per SM",
          "type": "quantitative"
        },
        "y": {
          "field": "L2",
          "scale": {
            "domain": [
              0,
              50000
            ]
          },
          "title": "L2/KB",
          "type": "quantitative"
        }
      },
      "mark": {
        "align": "left",
        "baseline": "middle",
        "dx": 5,
        "dy": 5,
        "type": "text"
      }
    }
  ],
  "width": 400
}

import sys
import pandas as pd
import altair as alt
import vl_convert as vlc
from cairosvg import svg2pdf


def main(*argv):
    df = alt.Chart(pd.DataFrame({
        'GPU': ['K40', 'P100', 'V100', 'A100', 'H100'],
        'L1': [64, 64, 128, 192, 256],
        'L2': [1536, 4096, 6144, 40960, 51200]
    })).encode(
        x=alt.X('L1').title('L1D+SMEM/KB per SM').scale(domain=(50, 250)),
        y=alt.Y('L2').title('L2/KB').scale(domain=(0, 50000)),
        text='GPU'
    )

    df = df.mark_line(
        color="grey"
    ) + df.mark_text(
        align='left',
        baseline='middle',
        dx=5,
        dy=5
    )

    df = df.configure_axisY(
        tickMinStep=10000
    ).configure_axisX(
        tickMinStep=50
    ).properties(
        width=400,
        height=300
    )

    svg2pdf(write_to="gpu-evole.pdf",
            bytestring=vlc.vegalite_to_svg(vl_spec=df.to_json()))

    return 0


if __name__ == '__main__':
    sys.exit(main(*sys.argv))