从plt.scatter获取PathCollection的各个元素的边界框

时间:2019-03-05 14:37:39

标签: python matplotlib

有没有一种方法可以在plt.scatter()的输出中获取单个元素的边界框?我可以获得偏移量(即x和y坐标-因为我要用它们来作图,所以必须从这里开始)和尺寸,但是尺寸不是以数据单位为单位,因此,即使是从面积到半径的顽强转换也可以得到尺寸的bbox无法使用...

有什么好方法吗?

tips = sns.load_dataset('tips')[:20]
f, ax = plt.subplots()
sc = ax.scatter(tips["total_bill"], y=tips["tip"], s=(tips["size"]*3)**2)
plt.show()

enter image description here

sc.properties()['offsets']

array([[ 16.99,   1.01],
       [ 10.34,   1.66],
       [ 21.01,   3.5 ],
       [ 23.68,   3.31],
       [ 24.59,   3.61],
       [ 25.29,   4.71],
       [  8.77,   2.  ],
       [ 26.88,   3.12],
       [ 15.04,   1.96],
       [ 14.78,   3.23],
       [ 10.27,   1.71],
       [ 35.26,   5.  ],
       [ 15.42,   1.57],
       [ 18.43,   3.  ],
       [ 14.83,   3.02],
       [ 21.58,   3.92],
       [ 10.33,   1.67],
       [ 16.29,   3.71],
       [ 16.97,   3.5 ],
       [ 20.65,   3.35]])

sc.get_sizes()

array([ 36, 81, 81, 36, 144, 144, 36, 144, 36, 36, 36, 144, 36, 144, 36, 36, 81, 81, 81, 81])

1 个答案:

答案 0 :(得分:3)

就一般而言,这远非简单。 PathCollection允许进行不同的转换以及偏移量转换。它还可能具有一个或多个路径和大小。

幸运的是,有一个内置函数matplotlib.path.get_path_collection_extents,它提供了PathCollection的边界框。我们可以通过提供每个单个路径的一个项目列表并遍历所有路径,来代替每个成员的作用范围。
由于边界框以像素为单位,因此需要在最后转换回数据坐标。

下面是完成所有操作的完整功能。首先需要绘制图形,以便设置不同的变换。

import numpy as np; np.random.seed(432)
import matplotlib.pyplot as plt
from matplotlib.path import get_path_collection_extents


def getbb(sc, ax):
    """ Function to return a list of bounding boxes in data coordinates
        for a scatter plot """
    ax.figure.canvas.draw() # need to draw before the transforms are set.
    transform = sc.get_transform()
    transOffset = sc.get_offset_transform()
    offsets = sc._offsets
    paths = sc.get_paths()
    transforms = sc.get_transforms()

    if not transform.is_affine:
        paths = [transform.transform_path_non_affine(p) for p in paths]
        transform = transform.get_affine()
    if not transOffset.is_affine:
        offsets = transOffset.transform_non_affine(offsets)
        transOffset = transOffset.get_affine()

    if isinstance(offsets, np.ma.MaskedArray):
        offsets = offsets.filled(np.nan)

    bboxes = []

    if len(paths) and len(offsets):
        if len(paths) < len(offsets):
            # for usual scatters you have one path, but several offsets
            paths = [paths[0]]*len(offsets)
        if len(transforms) < len(offsets):
            # often you may have a single scatter size, but several offsets
            transforms = [transforms[0]]*len(offsets)

        for p, o, t in zip(paths, offsets, transforms):
            result = get_path_collection_extents(
                transform.frozen(), [p], [t],
                [o], transOffset.frozen())
            bboxes.append(result.inverse_transformed(ax.transData))

    return bboxes



fig, ax = plt.subplots()

sc = ax.scatter(*np.random.rand(2,5), s=np.random.rand(5)*150+60)

# a single size needs to work as well. As well as a marker with non-square extent
sc2 = ax.scatter([0.2,0.5],[0.1, 0.7], s=990, marker="$\\rightarrow$")  

boxes = getbb(sc, ax)
boxes2 = getbb(sc2, ax)

# Draw little rectangles for boxes:
for box in boxes+boxes2:
    rec = plt.Rectangle((box.x0, box.y0), box.width, box.height, fill=False,
                        edgecolor="crimson")
    ax.add_patch(rec)

plt.show()

enter image description here