CODEWITHSUNDEEP
pythonmatplotlibplotcolorscontour
Divyajyoti
Divyajyoti

Reputation: 43

Colors in contour map not mapped to levels but to some other parameter [python]

I am making contour plot for from multiple data sets all mapped to the same level. So I want the colors from cmap not to be mapped on levels but on the values of mass which varies for various data sets. How can I do that?

My code structure is given below:

import numpy as np
import pylab as plt
from pandas import DataFrame as df
import matplotlib.colors as colors
import scipy.interpolate as interpolate

#data is imported from files, made into columns using DataFrame and put in the array name 'data'

xi = np.linspace(1,10,1000)
yi = np.linspace(-1,1,1000)
X, Y = np.meshgrid(xi,yi)

for i in range(9):
    Z = interpolate.griddata((data[i]['q'], np.cos(data[i]['iota'])), data[i]['snr1'], (X,Y))
    cs = plt.contour(X,Y,Z,levels=[20])
    cs.collections[0].set_label(str(int(data[i]['mass'][0])))
plt.legend(loc=5, title='mass')

The resulting plot is:

enter image description here

How can I use cmap to map the various contours according to the mass values?

Upvotes: 1

Views: 558

Answers (1)

JohanC
JohanC

Reputation: 80329

You need to create a colormap object and a norm. The colormap object converts an input value between 0 and 1 to a color value. The norm is a function that converts values between a minimum and a maximum to the range 0,1.

Note that the colors= parameter of plt.contour needs an extra pair of square brackets because the code gets confused in distinguishing between a single rgba value and an array of colors.

This is how your code could look like:

from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
import scipy.interpolate as interpolate

#data is imported from files, made into columns using DataFrame and put in the array name 'data'

xi = np.linspace(1,10,1000)
yi = np.linspace(-1,1,1000)
X, Y = np.meshgrid(xi,yi)

cmap = plt.get_cmap('magma')
norm = plt.Normalize(0, max([data[i]['mass'][0] for i in range(9)]))

for i in range(9):
    Z = interpolate.griddata((data[i]['q'], np.cos(data[i]['iota'])), data[i]['snr1'], (X, Y))
    mass = data[i]['mass'][0]
    cs = plt.contour(X, Y, Z, levels=[20], colors=[cmap(norm(mass))])
    cs.collections[0].set_label(f'{mass:.0f}')
plt.legend(loc=5, title='mass')
plt.tight_layout()
plt.show()

Upvotes: 1

Related Questions