Search code examples
pythonmatplotlibcolorsmatplotlib-3d

3d scatter plot with color gradient where color depends on count


I have dataframe with points which include x, y and z coordinate of the point and "count", which is number between 1 and 187 for each data point. I would like to associate "count" with color gradient, with 1 being for instance color green, and 187 color red, and then to make scatter plot of data points with x, y and z coordinate, where color of every data point is coded information for the "count". I find code for color gradients very confusing, can you please help me?

EDIT: zelusp completely answered my question, thank you.

EDIT1: I wondered if I should write separate question, but as it's about the same piece of code maybe I could reach for answer just editing previous question. Recently I changed my laptop, and installed Ubuntu 18.04 on my new machine. This piece of code worked perfectly on Ubuntu 16.04, but doesn't work on my new machine. My code is:

cmap = plt.cm.rainbow
norm = mpl.colors.Normalize(vmin=np.min(df3.h_count), vmax=np.max(df3.h_count))

fig=plt.figure()
ax1=fig.add_subplot(111, projection='3d')
ax1.scatter(df3.zm_bin, df3.sfr_bin, 12.+np.log10(df3.medijana), s=10, c=cmap(norm(df3.h_count)), marker='o') 
ax1.set_xlim(8,12.5)
ax1.set_xlabel('Log(Mz)')
ax1.set_ylabel('LogSFR') #treba invertovati
ax1.set_zlabel('12+log(Z)')
ax1.invert_yaxis()

sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
fig.colorbar(sm)

The error I'm getting is this:

Traceback (most recent call last):
  File "load_subhalos.py", line 233, in <module>
    ax1.scatter(df3.zm_bin, df3.sfr_bin, 12.+np.log10(df3.medijana), s=10, c=cmap(norm(df3.h_count)), marker='o') 
  File "/usr/local/lib/python2.7/dist-packages/matplotlib/colors.py", line 938, in __call__
    result, is_scalar = self.process_value(value)
  File "/usr/local/lib/python2.7/dist-packages/matplotlib/colors.py", line 924, in process_value
    result = np.ma.array(data, mask=mask, dtype=dtype, copy=True)
  File "/usr/local/lib/python2.7/dist-packages/numpy/ma/core.py", line 6358, in array
    ndmin=ndmin, shrink=shrink, order=order)
  File "/usr/local/lib/python2.7/dist-packages/numpy/ma/core.py", line 2784, in __new__
    order=order, subok=True, ndmin=ndmin)
TypeError: float() argument must be a string or a number

I didn't change anything in the code, the only change is new machine and OS. Do you know why I'm having this problem?


Solution

  • I recommend you take a tour through these posts on matplotlib:

    With that in mind I made the following:

    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.colors
    from mpl_toolkits.mplot3d import Axes3D
    
    #%% Generate mock data
    number_of_datapoints = 30
    x = np.random.rand(number_of_datapoints)
    y = np.random.rand(number_of_datapoints)
    z = np.random.rand(number_of_datapoints)
    
    count_min = 1
    count_max = 187
    data = np.random.randint(count_min, count_max, number_of_datapoints) # these are your counts
    
    #%% Create Color Map
    colormap = plt.get_cmap("YlOrRd")
    norm = matplotlib.colors.Normalize(vmin=min(data), vmax=max(data))
    
    #%% 3D Plot
    fig = plt.figure()
    ax3D = fig.add_subplot(111, projection='3d')
    ax3D.scatter(x, y, z, s=10, c=colormap(norm(data)), marker='o')  
    plt.show()
    

    You may also be interested in colorspacious