python - 绘制一种热图，其颜色是函数 x,y -> r,g,b 的结果

Question

我有一个字典，可以将 XY 元组映射到 RGB 元组。例如，

d = {
    (0, 0): (0, 0, 0),
    (0, 1): (0, 0, 200),
    }

我希望绘制某种热图，在给定的 XY 坐标下，其颜色是字典中颜色的平均值，并按它们的倒数距离加权；就好像它们是“光源”一样。

在给定的例子中，坐标 (0, 0.5) 应该用 (0, 0, 100) 和坐标 (0, 0.1 ) 应使用 (0, 0, 20) 着色。

我的问题是相当技术性的:如何让 pyplot 使用从函数 f(x, y) -> (r, g, b)?

Answer 1

如果您有 X-Y 网格:

import numpy
from matplotlib import pyplot as plt

width, height = 300, 500

xs = numpy.arange(width)
ys = numpy.arange(height)
data = numpy.dstack(numpy.meshgrid(xs, ys))

您应该将它们映射到 (r, g, b) 元组。以下是相当慢的，但如何加快它取决于你的函数做什么。

from colorsys import hsv_to_rgb

import random
def data_to_color(x, y):
    return (
        (x/width)**(0.5+random.random()*2),
        (y/height)**3,
        (x/width*y/height)*0.6 + random.random()*0.4
    )

colors = [[data_to_color(x, y) for x, y in row] for row in data]
colors = numpy.array(colors)
colors.shape
#>>> (500, 300, 3)

然后 imshow 可以给出想要的输出:

plt.imshow(colors, origin='lower')
plt.show()

现在，如果你想按照你所说的那样从你的点进行插值，你可以使用 scipy.interpolate。我将制作一本字典来从上面的函数中推断:

from scipy.interpolate import griddata

gridpoints = data.reshape(width*height, 2)
d = {(x, y): data_to_color(x, y) for x, y in gridpoints if not random.randint(0, 1000)}

len(d)
#>>> 142

将字典提取到 numpy 数组中，并分离颜色(可能可以避免分离，但您可以自己测试):

points, values = zip(*d.items())

points = numpy.array(points)
values = numpy.array(values)

reds   = values[:, 0]
greens = values[:, 1]
blues  = values[:, 2]

然后在点上运行griddata:

new_reds   = griddata(points, reds,   (data[:, :, 0], data[:, :, 1]), method='linear')
new_greens = griddata(points, greens, (data[:, :, 0], data[:, :, 1]), method='linear')
new_blues  = griddata(points, blues,  (data[:, :, 0], data[:, :, 1]), method='linear')

new_colors = numpy.dstack([new_reds, new_greens, new_blues])
new_colors[numpy.isnan(new_colors)] = 0.5

和情节:

plt.triplot(points[:,0], points[:,1], 'k-', linewidth=1, alpha=0.5)

plt.imshow(new_colors, extent=(0, width, 0, height), origin='lower')
plt.show()

最后，如果你也想要外推，我复制了一些代码 from here :

import scipy

def extrapolate_nans(x, y, v):
    '''  
    Extrapolate the NaNs or masked values in a grid INPLACE using nearest
    value.

    .. warning:: Replaces the NaN or masked values of the original array!

    Parameters:

    * x, y : 1D arrays
        Arrays with the x and y coordinates of the data points.
    * v : 1D array
        Array with the scalar value assigned to the data points.

    Returns:

    * v : 1D array
        The array with NaNs or masked values extrapolated.
    '''

    if numpy.ma.is_masked(v):
        nans = v.mask
    else:
        nans = numpy.isnan(v)
    notnans = numpy.logical_not(nans)
    v[nans] = scipy.interpolate.griddata((x[notnans], y[notnans]), v[notnans],
        (x[nans], y[nans]), method='nearest').ravel()
    return v

new_reds   = extrapolate_nans(data[:, :, 0], data[:, :, 1], new_reds)
new_greens = extrapolate_nans(data[:, :, 0], data[:, :, 1], new_greens)
new_blues  = extrapolate_nans(data[:, :, 0], data[:, :, 1], new_blues)

new_colors = numpy.dstack([new_reds, new_greens, new_blues])

plt.imshow(new_colors, extent=(0, width, 0, height), origin='lower')
plt.show()

---

编辑:也许更像是

import numpy
from matplotlib import pyplot as plt
from numpy.core.umath_tests import inner1d

width, height = 300, 500

xs, ys = numpy.mgrid[:width, :height]
coordinates = numpy.dstack([xs, ys])

light_sources = {
    (0, 0): (0, 0, 0),
    (300, 0): (0, 0, 0),
    (0, 0): (0, 0, 0),
    (300, 500): (0, 0, 0),
    (100, 0): (0, 0, 200),
    (200, 150): (100, 70, 0),
    (220, 400): (255, 255, 255),
    (80, 220): (255, 0, 0),
}

weights = numpy.zeros([width, height])
values = numpy.zeros([width, height, 3])

对于每个光源:

for coordinate, value in light_sources.items():

计算(反)距离。使用 +1e9 来防止无穷大，尽管这会导致愚蠢的失败，因此稍后更严格的修复很重要:

    shifted_coordinates = coordinates - coordinate + 1e-9
    inverse_distances = (shifted_coordinates ** 2).sum(axis=-1) ** (-1/2)

将其添加到总和和总和权重中:

    weights += inverse_distances
    values  += inverse_distances[:, :, numpy.newaxis].repeat(3, axis=-1) * value / 255

除以权重得到平均值:

values /= weights[..., numpy.newaxis]

并显示...

plt.imshow(values, origin='lower')
plt.show()

为此:

我最初没有这样做的原因是因为在您的示例中 (0, 0.1) 处的值不是 (0, 0, 20) 而是:

distances = [0.9, 0.1]
inverse_distances = [10/9, 10]
sum_weighting = 100 / 9
blue_levels = 200 / (109/90) = 18

所以根据这个定义，它应该是 (0, 0, 18)。