やりたいこととは異なるかもしれませんが、与えられた3次元の点群から、それに接するn(3)次多項式で近似した曲面をヒートマップする例です。
Python 3D polynomial surface fit, order dependentをほぼそのままパクって参考にしています。
Python
1
2import itertools
3import numpy as np
4
5# n(3)次多項式での補間式(係数)を算出
6def polyfit2d(x, y, z, order=3):
7 ncols = (order + 1)**2
8 G = np.zeros((x.size, ncols))
9 ij = itertools.product(range(order+1), range(order+1))
10 for k, (i,j) in enumerate(ij):
11 G[:,k] = x**i * y**j
12 m, _, _, _ = np.linalg.lstsq(G, z)
13 return m # 補間式(係数)
14
15# 補間式(係数)から 各格子点のz値を算出
16def polyval2d(x, y, m):
17 order = int(np.sqrt(len(m))) - 1
18 ij = itertools.product(range(order+1), range(order+1))
19 z = np.zeros_like(x)
20 for a, (i,j) in zip(m, ij):
21 z += a * x**i * y**j
22 return z
23
24
25# N_TRAIN 個の点をランダム生成
26N_TRAIN = 100
27x = 2 * np.random.rand(N_TRAIN) - 1 # -1...+1
28y = 4 * np.random.rand(N_TRAIN) - 2 # -2...+2
29
30#z =x + y
31z = x * y
32#z = x**2 - y**2
33#z = np.sin(np.sqrt(x**2 + y**2)) # さすがにこれは多項式近似では厳しい
34
35# n(3)次多項式での補間式(係数)を算出
36N_ORDER = 3
37m = polyfit2d(x,y,z,N_ORDER)
38
39# x, y 各方向に N_MESH 個の格子点を生成
40N_MESH = 20
41xx, yy = np.meshgrid( np.linspace(x.min(), x.max(), N_MESH),
42 np.linspace(y.min(), y.max(), N_MESH))
43
44# 補間式(係数)から 各格子点上のz値を算出
45zz = polyval2d(xx, yy, m)
46
47import matplotlib.pyplot as plt
48from mpl_toolkits.mplot3d import Axes3D
49
50fig = plt.figure()
51ax = Axes3D(fig)
52
53# 補間結果の曲面
54ax.plot_surface(xx,yy,zz, cmap='Reds', alpha=0.9)
55
56# 元データの点群
57ax.scatter(x,y,z)
58
59plt.show()