How to 3D plot function of 2 variables in python?

Question

I am trying to 3D plot the magnification factor in vibrations for multiple types of damping. To simplify it for those who have no idea what it is, basically, you have 3 variables:

• beta, which varies between 0 and infinite, but I would like to visualize it from 0 to 3, in 0.2 intervals.
• damping ratio, d, which varies between 0 and infinite, but I would like to plot it from 0 to 1, in 0.1 intervals.
• finally, nu, which is a function that varies according to the two variables before.

My intuition says that I should plot this with (X,Y,Z) = (beta, d, nu), but I am just starting to use this library and I am kind of new to python, I just use it when I need to visualize or calculate problems in class. I tried creating 2 arrays for beta and d, but I don't know I should create the array for nu, since it depends on both.

This is the piece of code I have until now:

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


nu = []
b = [0.1 + i / 100 for i in range(0, 510)]
damp = [0.1 + i/10 for i in range(0,510)]

for d in damp:
    nu_new = []
    nu.append(nu_new)
    for beta in b:
        nu_new.append( math.sqrt(1+(2*d*beta)**2)/ math.sqrt((1-beta**2)**2+(2*d*beta)**2))

fig = plt.figure()
ax = Axes3D(fig)
ax.plot(b, d, nu)
plt.show()

I am kind of stuck trying to plot this, so if you have any suggestion I would be glad.

Answer 1

If you're using numpy, then don't use the math module. Numpy as all of the math functions built in but they work on numpy arrays far better. We can calculate nu at all of our b, d values with the aid of a meshgrid.

A meshgrid can take 2 1D arrays, and return 2 2D arrays such that every index in the arrays corresponds to a unique pair of elements from the original 1D arrays.

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

b = np.arange(0.2, 3.2, 0.2)
d = np.arange(0.1, 1.0, 0.1)

B, D = np.meshgrid(b, d)
nu = np.sqrt( 1 + (2*D*B)**2 ) / np.sqrt( (1-B**2)**2 + (2*D*B)**2)

fig = plt.figure()
ax = Axes3D(fig)
ax.plot_surface(B, D, nu)
plt.xlabel('b')
plt.ylabel('d')
plt.show()

This produces:

Additionally, 3D plots tend to block seeing all of the data (because a spike hides things behind it). I would recommend a pcolormesh or a contourf plot. In the later case the last 6 lines become:

plt.contourf(B, D, nu)
plt.colorbar()
plt.xlabel('b')
plt.ylabel('d')
plt.show()

which produces: