import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import tensorflow as tf
import matplotlib.pyplot as plt


# z=sin(x)+sin(y)
def func(x):
    #  x的shape是（b，2）
    #  x的前面一层是x的坐标，x的后面一层是y的坐标？？？
    z = tf.math.sin(x[..., 0]) + tf.math.sin(x[..., 1])
    return z


x = tf.linspace(0., 2 * 3.14, 500)
y = tf.linspace(0., 2 * 3.14, 500)
#  points_x的shape是[500,500]
points_x, points_y = tf.meshgrid(x, y)
#  points的shape是[500,500,2]
points = tf.stack([points_x, points_y], axis=2)
#  points的shape是[250000,2]
# points = tf.reshape(points, [-1, 2])
# z的shape是[500，500]
z = func(points)

plt.figure('plot 2d func value')
plt.imshow(z, origin='lower', interpolation='none')
plt.colorbar()

plt.figure('plot 2d func contour')
plt.contour(points_x, points_y, z)
plt.colorbar()
plt.show()