import matplotlib.pyplot as plt
import numpy as np
import matplotlib.colors as colors

"""
Kod rysuje "orbitale atomowe" w dwóch wymiarach (wartość oraz kwadrat modułu).
Funkcje te mają postać matematyczną analogiczną do orbitali wodoropodobnych 
(patrz np. L. Piela, Ideas of Quantum Chemistry, 2nd edition, Elsevier (2013)).
Nie są znormalizowane, współczynnik w wykładniku również jest dowolny.
"""

def o_1s(x,y,coef,x0=0,y0=0):
    return np.e**(coef*np.sqrt((x-x0)**2+(y-y0)**2))

def o_2s(x,y,coef,x0=0,y0=0):
    return (1-np.sqrt(x**2+y**2))*np.e**(coef*np.sqrt((x-x0)**2+(y-y0)**2))

def o_2px(x,y,coef,x0=0,y0=0):
    return (x-x0)*o_1s(x,y,coef,x0,y0)

def o_2py(x,y,coef,x0=0,y0=0):
    return (y-y0)*o_1s(x,y,coef,x0,y0)

def o_3dxy(x,y,coef,x0=0,y0=0):
    return (x-x0)*(y-y0)*o_1s(x,y,coef,x0,y0)

def o_bonding(x,y,coef,R):
    #Orbital wiążący zbudowany z dwóch 1s
    return o_1s(x,y,coef,-R/2,0)+o_1s(x,y,coef,R/2,0)

def o_antibonding(x,y,coef,R):
    #Orbital antywiążący zbudowany z dwóch 1s
    return o_1s(x,y,coef,-R/2,0)-o_1s(x,y,coef,R/2,0)


def probab(nx,ny,fun,coef,**kwargs):
    """
    Kwadrat (modułu, ale dla l. rzeczywistych nie potrzeba modułu)
    określający gęstość prawdopodobieństwa znalezienia elektronu...
    ...tzn. określałby, gdyby funkcja była znormalizowana, obecnie
    oddaje co najwyżej kształt.
    """
    x=(nx*L)/N-L/2
    y=(ny*L)/N-L/2
    return fun(x,y,coef,**kwargs)**2

def value(nx,ny,fun,coef,**kwargs):
    #Wartość funkcji
    x=(nx*L)/N-L/2
    y=(ny*L)/N-L/2
    return fun(x,y,coef,**kwargs)


fig,(ax1,ax2) = plt.subplots(1,2,sharey=True)

N=1000 #Rozmiar obrazka
L=30.0 #Rozmiar pudła w "jednostkach atomowych"
#Parametry orbitalu, (anty)wiążące wymagają jeszcze "R":5
params={"fun":o_3dxy,"coef":-0.5}
img1=np.fromfunction(value,(N,N),**params)
img2=np.fromfunction(probab,(N,N),**params)

im1=ax1.imshow(img1.T,cmap="RdBu",norm=colors.CenteredNorm())
ax1.set_title("Wartość funkcji falowej")
fig.colorbar(im1,shrink=0.48)
im2=ax2.imshow(img2.T,cmap="Blues")
ax2.set_title("Kwadrat funkcji falowej")
fig.colorbar(im2,shrink=0.48)
plt.show()