Spectrum plot

The following code uses Gaussian function to simulate a spectrum of light.

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set_theme()

def  dic(x, t, w):
	diff = x.reshape(-1, 1)- t.reshape(1, -1)
	gauss = np.power(2, -diff**2/w**2)
	norm = np.sqrt(np.sum(gauss**2, axis=0))
	return gauss/norm

def  test():
	m, n = 200, 80
	x = np.linspace(0., 1., m)
	t = np.linspace(0., 1., n)
	w = 0.08
	atoms = dic(x, t, w)
	indexes = [n//3, 2*n//3]
	y = atoms[:, indexes].sum(axis=1)
	plt.plot(y, x, color="black")

	normalize = colors.Normalize(vmin=y.min(), vmax=y.max())
	cmap = plt.cm.get_cmap('rainbow')
	for i in  range(m-1):
		xs = [0., 1.]
		y1 = [x[i], x[i]]
		y2 = [x[i+1], x[i+1]]
		plt.fill_between(xs, y1, y2, color=cmap(normalize(x[i])), alpha=y[i]/y.max())
	plt.xlabel("intensity")
	plt.ylabel("frequency")
	plt.title("Spectrum of light")

test()