如何用Python基于双输入函数绘制theta值的100×100热力图？

我定义了如下函数：

def theta_function(time, omega,  dt=0.01, cromer=True):
    # Create the variables to store the results
    times = np.arange(tstart, tend+dt, dt)  # Create the array of model times
    Nt = len(times)                         # Number of time-steps
    y = np.zeros(Nt)
    theta = np.zeros(Nt)
    z = np.zeros(Nt)
    gamma = np.zeros(Nt)

    # Set the initial conditions
    initial_y = 0
    initial_theta = 0.1
    initial_gamma = 0
    initial_z = 0

    y[0] = initial_y
    theta[0] = initial_theta

    gamma[0] = initial_gamma
    z[0] = initial_z

    # Loop over all of the times and integrate the model
    for n in range(Nt-1):
        # Time-step the angle/vertical displacement. If we are using 'Taylor' or 'Cromer'
        if cromer:
            value = math.exp(a * (y[n] - l*math.sin(theta[n])))
            z[n+1] = z[n] + (((-d*z[n]) - ((K/(M*a)) * ((math.exp(a * (y[n] - l*math.sin(
                theta[n])))) + (math.exp(a * (y[n] + l*math.sin(theta[n])))) - 2))) * dt)
            y[n+1] = y[n] + ((z[n+1] + (A * math.sin(omega * times[n]))) * dt)
            gamma[n+1] = gamma[n] + (((-d*gamma[n]) + ((((3*math.cos(theta[n]))/l) * (K/(M*a))) * (
                (math.exp(a * (y[n] - l*math.sin(theta[n])))) - (math.exp(a * (y[n] + l*math.sin(theta[n]))))))) * dt)
            theta[n+1] = theta[n] + (gamma[n+1] * dt)

    return theta[time]

该函数接收time和omega作为输入，返回theta值（范围为0到π/4）。

我已生成如下数组：

omegas = np.linspace(2, 4, 100)

以及：

times = np.linspace(0, 1, 100)

两个数组的元素数量相同。

请问如何绘制一张100×100的热力图，使其x轴为0-1的时间范围，y轴为2-4的omega范围，每个方格对应一个特定的theta值？

1. 补全函数依赖的参数与变量 #

原函数中未定义tstart、tend以及a、l、d、K、M、A这些关键变量，需要先明确赋值（示例值可根据实际需求调整）：

import numpy as np
import math
import matplotlib.pyplot as plt

# 定义函数依赖的全局参数
tstart = 0
tend = 1
a = 0.5
l = 1.0
d = 0.1
K = 10.0
M = 1.0
A = 0.5

2. 修正函数逻辑（适配时间值输入） #

原函数的time参数是模型时间序列的索引，建议修改为直接接收时间值，自动匹配对应索引，避免手动转换：

def theta_function(target_time, omega, dt=0.01, cromer=True):
    times = np.arange(tstart, tend+dt, dt)
    Nt = len(times)
    y = np.zeros(Nt)
    theta = np.zeros(Nt)
    z = np.zeros(Nt)
    gamma = np.zeros(Nt)

    # 初始化条件
    y[0] = 0
    theta[0] = 0.1
    gamma[0] = 0
    z[0] = 0

    # 积分循环
    for n in range(Nt-1):
        if cromer:
            exp1 = math.exp(a * (y[n] - l*math.sin(theta[n])))
            exp2 = math.exp(a * (y[n] + l*math.sin(theta[n])))
            z[n+1] = z[n] + ((-d*z[n]) - (K/(M*a))*(exp1 + exp2 - 2)) * dt
            y[n+1] = y[n] + (z[n+1] + A * math.sin(omega * times[n])) * dt
            gamma[n+1] = gamma[n] + ((-d*gamma[n]) + (3*math.cos(theta[n])/l)*(K/(M*a))*(exp1 - exp2)) * dt
            theta[n+1] = theta[n] + gamma[n+1] * dt

    # 匹配最接近目标时间的索引，返回对应theta值
    idx = np.argmin(np.abs(times - target_time))
    return theta[idx]

3. 生成热力图数据矩阵 #

遍历所有times和omegas的组合，计算对应的theta值，构建100×100的矩阵：

# 初始化结果矩阵
theta_matrix = np.zeros((len(omegas), len(times)))

# 填充矩阵
for i, omega in enumerate(omegas):
    for j, t in enumerate(times):
        theta_matrix[i, j] = theta_function(t, omega)

4. 绘制热力图 #

使用matplotlib的pcolormesh绘制，确保轴范围与输入数组匹配：

plt.figure(figsize=(10, 8))
# 绘制热力图，shading='auto'保证方格对应数据点
im = plt.pcolormesh(times, omegas, theta_matrix, cmap='viridis', shading='auto')

# 设置轴标签与标题
plt.xlabel('时间')
plt.ylabel('Omega')
plt.title('Theta值热力图')

# 添加颜色条标注
plt.colorbar(im, label='Theta')

# 显示图像
plt.show()

内容的提问来源于stack exchange，提问作者cookiecainsy