Matplotlib 3d surface plot input arrays

Question

I am trying to create a surface plot using 3 numpy arrays:

• x_deflections [shape: (10,)]
• y_alphas [shape: (12,)]
• z_height_at_target [shape: (120,)]

x_deflections and y_alphas are given and z_height_at_target is calculated using the 2 of them like this:

i = 0
for x in x_deflections:
    for y in y_alphas:
        exit_v = spring_calc.do_calc(k, x)
        # Only care about x_dist and h here
        vX, vY, x_dist, h = traj_calc.do_calc(exit_v, y, stop_at=target_dist)
        try:
            if max(x_dist) < target_dist:
                raise StopIteration
            else:
                target_dist_index = find_nearest(x_dist, target_dist, 0.04)
        except StopIteration:
            print('Target Distance not achieved')
            continue
        z_height_at_target[i] = h[target_dist_index]
        i += 1

This works as I would expect and gives reasonable values in z_height_at_target , however I can't seem to work out how to create a proper surface plot from this. My current method gives me gibberish graphs full of spikes:

fig = pl.figure()
ax = fig.add_subplot(111, projection='3d')
X, Y = np.meshgrid(x_deflections, np.rad2deg(y_alphas), indexing='xy')
Z = z_height_at_target.reshape((len(x_deflections), len(y_alphas)))
ax.plot_surface(X, Y, Z, color='b')
ax.set_xlabel('Spring Deflection [m]')
ax.set_ylabel('Launch Angle [deg]')
ax.set_zlabel('Height at Target Distance [m]')
pl.show()

I am aware the problem lies in one of the following lines, however I can't seem to get my head around it:

X, Y = np.meshgrid(x_deflections, np.rad2deg(y_alphas), indexing='xy')
Z = z_height_at_target.reshape((len(x_deflections), len(y_alphas)))

This currently throws an error saying incorrect shape, which is correct, but transposing Z gives gibberish.

Any help would be appreciated. Thanks!

Answer 1

Solved this using @Uvar 's suggestion of transposing z_height_at_target and by putting the values in in 2 dimensions, instead of reshaping afterwards (This effectively does the same as the previous code):

i = 0
x_count = 0
miss_count = 0
for x in x_deflections:
    y_count = 0
    for y in y_alphas:
        exit_v = spring_calc.do_calc(k, x, y)
        # only need h here but the rest of can come along for the ride as well
        vX, vY, x_dist, h = traj_calc.do_calc(exit_v, y, stop_at=target_dist)
        try:
            if max(x_dist) < target_dist:
                raise StopIteration
            else:
                target_dist_index = find_nearest(x_dist, target_dist, 0.04)
        except StopIteration:
            print('Target Distance not achieved')
            miss_count += 1
            continue
        z_height_at_target[x_count, y_count] = h[target_dist_index]
        print('Completed iter', i+1)
        # print('{}, {}, {}'.format(exit_v, h[target_dist_index], np.rad2deg(y)))
        i += 1
        y_count += 1
    x_count += 1

fig = pl.figure()
ax = fig.add_subplot(121, projection='3d')
# ax = fig.gca(projection='3d')
X, Y = np.meshgrid(x_deflections, np.rad2deg(y_alphas), indexing='xy')
Z = z_height_at_target.T