个人笔记--python画图（一维，二维，三维）

1. 一维

1. plot

import numpy as np
import matplotlib.pyplot as plt

# linspace(): 创建等间距的数值序列
x = np.linspace(0, 2 * np.pi, 100)

u = np.sin(x)

# 绘制一维图形
plt.figure()
plt.plot(x, u)
plt.title('Plot of sin(x)')
plt.xlabel('x')
plt.ylabel('sin(x)')
plt.show()

2. 二维

2.1 imshow

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, 2 * np.pi, 100)
y = np.linspace(0, 2 * np.pi, 100)
# 创建网格点
X, Y = np.meshgrid(x, y)

U = np.sin(X) * np.cos(Y)

# 使用imshow绘制热图
plt.figure()
plt.imshow(U, extent=(0, 2 * np.pi, 0, 2 * np.pi), origin='lower', cmap='viridis')
plt.colorbar()
plt.title('Heatmap of sin(x) * cos(y)')
plt.xlabel('x')
plt.ylabel('y')
plt.show()

2.2 contour

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, 2 * np.pi, 100)
y = np.linspace(0, 2 * np.pi, 100)
X, Y = np.meshgrid(x, y)

U = np.sin(X) * np.cos(Y)

# 使用contour绘制等高线图
plt.figure()
plt.contour(X, Y, U, levels=20, cmap='viridis')
plt.colorbar()
plt.title('Contour of sin(x) * cos(y)')
plt.xlabel('x')
plt.ylabel('y')
plt.show()

2.3 pcolor

import numpy as np
import matplotlib.pyplot as plt

x = np.linspace(0, 2 * np.pi, 100)
y = np.linspace(0, 2 * np.pi, 100)
X, Y = np.meshgrid(x, y)

U = np.sin(X) * np.cos(Y)

# 使用pcolor绘制伪彩色图
plt.figure()
plt.pcolor(X, Y, U, cmap='viridis')
plt.colorbar()
plt.title('Pcolor of sin(x) * cos(y)')
plt.xlabel('x')
plt.ylabel('y')
plt.show()

2.4 scatter

import numpy as np
import matplotlib.pyplot as plt

# 生成数据
x = np.linspace(0, 2 * np.pi, 100)
y = np.linspace(0, 2 * np.pi, 100)
X, Y = np.meshgrid(x, y)
U = np.sin(X) * np.cos(Y)

# 绘制散点图
plt.scatter(X, Y, c=U, cmap='viridis')

# 添加标题和标签
plt.title('Scatter plot of U = sin(X) * cos(Y)')
plt.xlabel('X')
plt.ylabel('Y')

# 显示图表
plt.colorbar(label='U value')
plt.show()

# 上面的X，Y，U维度都是(100,100)，
# 下面的x_star, y_star, v0_train维度是（100*100，1）即（10000，1）
# 只要维度一样就可以
# plt.scatter(x_star, y_star, c=v0_train, cmap='viridis')
# plt.colorbar()
# plt.xlabel('X')
# plt.ylabel('Y')
# plt.title('v_0_train')
# plt.show()

2.5 plot_surface

import numpy as np
import matplotlib.pyplot as plt
# from mpl_toolkits.mplot3d import Axes3D

# 生成x和y
x = np.linspace(0, 2 * np.pi, 100)
y = np.linspace(0, 2 * np.pi, 100)
X, Y = np.meshgrid(x, y)

# 计算u
U = np.sin(X) * np.cos(Y)

# 使用plot_surface绘制三维表面图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, U, cmap='viridis')
ax.set_title('Surface plot of sin(x) * cos(y)')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('u')
plt.show()

2.6 contour3D

import numpy as np
import matplotlib.pyplot as plt
# from mpl_toolkits.mplot3d import Axes3D

# 生成x和y
x = np.linspace(0, 2 * np.pi, 100)
y = np.linspace(0, 2 * np.pi, 100)
X, Y = np.meshgrid(x, y)

# 计算u
U = np.sin(X) * np.cos(Y)

# 使用contour3D绘制三维等高线图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.contour3D(X, Y, U, 50, cmap='viridis')
ax.set_title('3D Contour of sin(x) * cos(y)')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('u')
plt.show()

3. 三维

3.1 plot_surface

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 生成x, y, z
x = np.linspace(0, 2 * np.pi, 50)
y = np.linspace(0, 2 * np.pi, 50)
z = np.linspace(0, 2 * np.pi, 50)
X, Y, Z = np.meshgrid(x, y, z)

U = np.sin(X) * np.cos(Y) * np.sin(Z)

# 固定z，绘制三维表面图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X[:,:,25], Y[:,:,25], U[:,:,25], cmap='viridis')
ax.set_title('Surface plot of sin(x) * cos(y) * sin(z) at z=pi')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('u')
plt.show()

3.2 contour3D

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 生成x, y, z
x = np.linspace(0, 2 * np.pi, 50)
y = np.linspace(0, 2 * np.pi, 50)
z = np.linspace(0, 2 * np.pi, 50)
X, Y, Z = np.meshgrid(x, y, z)

U = np.sin(X) * np.cos(Y) * np.sin(Z)

# 固定z，绘制三维等高线图
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.contour3D(X[:,:,25], Y[:,:,25], U[:,:,25], 50, cmap='viridis')
ax.set_title('3D Contour of sin(x) * cos(y) * sin(z) at z=pi')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.set_zlabel('u')
plt.show()

4. 存储图像

4.1 一般情况

import matplotlib.pyplot as plt
import numpy as np

# 创建一些示例数据
x = np.linspace(0, 10, 100)
y = np.sin(x)

# 创建一个折线图
plt.plot(x, y)

# 使用plt.savefig()保存图像
plt.savefig('sine.png')

# 显示图像
plt.show()

4.2 一次存多个图

globals()[f"pic_{time_block}_{AM_count}"] = plt.figure(1, figsize=(15, 15))
predict_np_u = model.predict_U(x_test_current).cpu().detach().numpy()
predict_np_v = model.predict_V(x_test_current).cpu().detach().numpy()
predict_np_p = model.predict_P(x_test_current).cpu().detach().numpy()

u_pred = np.reshape(predict_np_u, (x.shape[0], y.shape[0], t_current.shape[0]), order='F')
v_pred = np.reshape(predict_np_v, (x.shape[0], y.shape[0], t_current.shape[0]), order='F')
p_pred = np.reshape(predict_np_p, (x.shape[0], y.shape[0], t_current.shape[0]), order='F')

# Adjust subplot parameters to avoid overlap
plt.subplots_adjust(wspace=0.4, hspace=0.4)  # Increase the width and height spaces

for i in range(len(t_1)):
    # Prediction
    plt.subplot(3, 3, 1 + 3 * i)
    plt.pcolor(x_1, y_1, u_pred[:, :, t_pos1[i]], cmap='jet')
    plt.colorbar()
    plt.xlabel(r'$x$', fontsize=18)
    plt.ylabel(r'$y$', fontsize=18)
    plt.title('Predicted $\hat u(x,y,t)$, t=' + str(t_1[i]), fontsize=15)

    plt.subplot(3, 3, 2 + 3 * i)
    plt.pcolor(x_1, y_1, v_pred[:, :, t_pos1[i]], cmap='jet')
    plt.colorbar()
    plt.xlabel(r'$x$', fontsize=18)
    plt.ylabel(r'$y$', fontsize=18)
    plt.title('Predicted $\hat u(x,y,t)$, t=' + str(t_1[i]), fontsize=15)

    plt.subplot(3, 3, 3 + 3 * i)
    plt.pcolor(x_1, y_1, p_pred[:, :, t_pos1[i]], cmap='jet')
    plt.colorbar()
    plt.xlabel(r'$x$', fontsize=18)
    plt.ylabel(r'$y$', fontsize=18)
    plt.title('Predicted $\hat u(x,y,t)$, t=' + str(t_1[i]), fontsize=15)

plt.close()
globals()[f"pic_{time_block}_{AM_count}"].savefig("figures_count/Sol_" + params_name + str(time_block) + "_" + str(AM_count) + ".png", dpi=500, bbox_inches='tight')

用上globals()去命名