Matplotlib - строит линии, используя ту же цветовую карту, что и разброс

Question

Используя Matplotlib, я пытаюсь нанести линии анализа модерации поверх некоторых разбросанных данных. Как установить цвет линий в соответствии со значениями в цветовой карте, используемой для данных разброса?

import numpy as np
import matplotlib.pyplot as plt
plt.rcParams.update({"font.size": 14})

# MAKE DATA
N = 100
β0, β1, β2, β3, σ = 0.5, 1, 2, 0, 0.05

x = np.random.randn(N)
moderator = np.random.randn(N)
μ = β0 + β1 * x + β2 * x * moderator + β3 * moderator
ϵ = np.random.randn(N) * σ
y = μ + ϵ

fig, ax = plt.subplots(figsize=(7,7))

# PLOT DATA
h = ax.scatter(x, y, c=moderator)
ax.set(xlabel="x", ylabel="y")
# colourbar for moderator
cbar = fig.colorbar(h)
cbar.ax.set_ylabel("moderator")

# PLOT EXAMPLE LINES FOR TRUE MODEL
_x = np.linspace(np.min(x), np.max(x), 2)
m_levels = np.linspace(np.min(x), np.max(x), 5)
for m in m_levels:
    μ = β0 + β1 * _x + β2 * _x * m + β3 * m
    ax.plot(_x, μ, lw=3)

Answer 1

Решение, кажется, настраивает scalarMap и использует normalize ...

import numpy as np
import matplotlib.pyplot as plt
%config InlineBackend.figure_format = 'retina'
plt.rcParams.update({"font.size": 14})
from matplotlib.cm import ScalarMappable
from matplotlib.colors import Normalize

# MAKE DATA

N = 500
β0, β1, β2, β3, σ = 0.5, 1, 2, 0, 0.05


x = np.random.randn(N)
moderator = np.random.randn(N)
μ = β0 + β1 * x + β2 * x * moderator + β3 * moderator
ϵ = np.random.randn(N) * σ
y = μ + ϵ


fig, ax = plt.subplots(figsize=(7,7))

# set up scalarMap so we can access colormap across both scatter and lines
normalize = Normalize(vmin=np.min(moderator), vmax=np.max(moderator))
scalarMap = ScalarMappable(norm=normalize, cmap="viridis")

# PLOT EXAMPLE LINES FOR TRUE MODEL
_x = np.linspace(np.min(x), np.max(x), 2)
m_levels = np.linspace(np.min(x), np.max(x), 5)
for m in m_levels:
    μ = β0 + β1 * _x + β2 * _x * m + β3 * m
    ax.plot(_x, μ, lw=3, c=scalarMap.to_rgba(m))

# PLOT DATA
h = ax.scatter(x, y, c=moderator, cmap=scalarMap.cmap)
ax.set(xlabel="x", ylabel="y")
# colourbar for moderator
cbar = fig.colorbar(h)
cbar.ax.set_ylabel("moderator")