Python : modèle de Bohr

Ce programme crée le graphique du diagramme de Bohr en fonction d'un nombre d'électrons donnés, et enregistre sous forme de .png

import numpy as np
import matplotlib.pyplot as plt
			
def Config():
	Z=int(input("Combien d'electrons ?"))
	orbitales = ["1s", "2s", "2p", "3s", "3p", "4s","3d", "4p", "5s", "4d", "5p", "6s","4f", "5d", "6p", "7s", "5f", "6d", "7p"]
    
	max_elec = {"s": 2,"p": 6,"d": 10,"f": 14}
    
	config = []
	reste = Z
    
	for orb in orbitales:
		l = orb[-1]  # type de sous-couche: s, p, d, f
		n = min(max_elec[l], reste)
		config.append([f"{orb}^",n])
		reste -= n
		if reste <= 0:
			break
    
	print('Configuration électronique :',config)
	return config,Z

def traceur() :
	orbits,Z = Config()
	population = [0 for i in range(7)]
	for orb in orbits : 
		n = int(orb[0][0])
		population[n-1]+=orb[1]
	print('Population :',population)
	
	fig, ax = plt.subplots()
	rayons=np.linspace(1,7,7,dtype=int)

	for r, n_e in zip(rayons, population):
		if n_e!=0 :
    	
			cercle = plt.Circle((0, 0), r, fill=False)
			ax.add_patch(cercle)

    
			angles = np.linspace(0, 2*np.pi, n_e, endpoint=False)
			x = r * np.cos(angles)
			y = r * np.sin(angles)

			ax.scatter(x, y, s=50)  # s=taille du point


	nucleus = plt.Circle((0, 0), 0.2, color=(0,0,1))
	ax.add_patch(nucleus)

	ax.set_aspect('equal', 'box')
	plt.axis('off')
	ax.set_xlim(-max(rayons)-1, max(rayons)+1)
	ax.set_ylim(-max(rayons)-1, max(rayons)+1)
	ax.set_xticks([])
	ax.set_yticks([])
	plt.savefig(f'orb_{Z}')
	plt.show()
	
	
	

	
traceur()