schrodinger-simulation

schrodinger's playground
git clone git://popovic.xyz/schrodinger-simulation.git
Log | Files | Refs | README | LICENSE

main.py (9031B)

      1 #!/usr/bin/python3.9
      2 
      3 from dolfin import *
      4 import numpy as np
      5 import matplotlib.pyplot as plt
      6 from scipy.special import factorial, hermite, eval_laguerre, genlaguerre
      7 import matplotlib.patches as mpatches
      8 
      9 def harmonic_oscillator(xmin, xmax, mesh, V):
     10 
     11     xmin=xmin
     12     xmax=xmax
     13     mesh=mesh
     14     V=V
     15 
     16     def boundary(x, on_bnd):
     17         return on_bnd
     18 
     19     bc = DirichletBC(V, Constant(0.), boundary)
     20 
     21     u = TrialFunction(V)
     22     v = TestFunction(V)
     23 
     24     pot_ = Expression('0.5*pow(x[0], 2)', degree=2)
     25     pot = interpolate(pot_, V)
     26 
     27     a = 1/2*inner(grad(u), grad(v))*dx + pot*u*v*dx
     28     m = u*v*dx
     29 
     30     A = PETScMatrix()
     31     M = PETScMatrix()
     32 
     33     assemble(a, tensor=A)
     34     assemble(m, tensor=M)
     35 
     36     bc.apply(A)
     37 
     38     # create eigensolver
     39     eigensolver = SLEPcEigenSolver(A, M)
     40     eigensolver.parameters['spectrum'] = 'smallest magnitude'
     41 
     42     # solve for eigenvalues
     43     values = 10
     44     eigensolver.solve(values)
     45 
     46     u = Function(V)
     47 
     48     # plot
     49     fig = plt.figure(figsize=[10, 10])
     50     En = []
     51     for i in range(0, values+1):
     52         E, E_c, R, R_c = eigensolver.get_eigenpair(i)
     53         En.append(E)
     54         x = np.linspace(xmin, xmax, len(R))
     55 
     56         I = np.trapz(R*R, x)
     57   #      print(I)
     58 
     59         #plot eigenfunction
     60         plt.plot(x, 20*R*R+E, c='blue')
     61         plt.plot(x, R_c+E, lw=0.5, c='black')
     62         plt.annotate(f'$E_{ { i } } = {round(En[i], 3)}$', (xmax-0.1, 0.05+E), fontsize=14)
     63         plt.annotate(f'$|\Psi _{ { i } }|²$', (xmin, 0.1+E), fontsize=14)
     64 
     65     plt.plot(x, 1/2*x**2, c='red')
     66     plt.ylim(0, max(En)+0.5)
     67     plt.title('Harmonic Oscillator $V(x) =\\frac{1}{2}x^2 $', fontsize=20)
     68     plt.xticks([])
     69     plt.yticks([])
     70 
     71 
     72 
     73     # plot analytical solutions
     74 
     75     def harm_osc(n, x):
     76         return 1/sqrt(sqrt(np.pi)*(2**n)*factorial(n)) * hermite(n)(x) * np.exp(-1/2*x**2)
     77 
     78     x = np.linspace(xmin, xmax, 70)
     79 
     80   #  print(En)
     81 
     82     E = []
     83     for n in range(0, values+1):
     84 
     85         E.append(n+1/2)
     86         psi = harm_osc(n, x)
     87         I = np.trapz(psi*psi, x)
     88        # print(I)
     89         plt.scatter(x, psi*psi+En[n], s=20, c='green', marker="x")
     90 
     91     numerical = mpatches.Patch(color='blue', label='numerical solution', linestyle='-')
     92     analytical = mpatches.Patch(color='green', label='analytical solution', linestyle='-')
     93     plt.legend(handles = [numerical, analytical], loc='upper left')
     94  #   plt.show()
     95    # plt.savefig('./plots/harmonic_oscillator.png')
     96 
     97     return E, En
     98 
     99 def box(xmin, xmax, mesh, V):
    100     xmin=xmin
    101     xmax=xmax
    102     mesh=mesh
    103     V=V
    104 
    105     def boundary(x, on_bnd):
    106         return on_bnd
    107 
    108     bc = DirichletBC(V, Constant(0.), boundary)
    109 
    110     u = TrialFunction(V)
    111     v = TestFunction(V)
    112 
    113     pot_ = Expression('x[0] == xmin || x[0] == xmax ? 1000 : 0', xmin=xmin, xmax=xmax, degree=2)
    114 
    115     pot = interpolate(pot_, V)
    116 
    117     a = 1/2*inner(grad(u), grad(v))*dx + pot*u*v*dx
    118     m = u*v*dx
    119 
    120     A = PETScMatrix()
    121     M = PETScMatrix()
    122 
    123     assemble(a, tensor=A)
    124     assemble(m, tensor=M)
    125 
    126     bc.apply(A)
    127 
    128     # create eigensolver
    129     eigensolver = SLEPcEigenSolver(A, M)
    130     eigensolver.parameters['spectrum'] = 'smallest magnitude'
    131 
    132     # solve for eigenvalues
    133     values = 6
    134     eigensolver.solve(values)
    135 
    136     u = Function(V)
    137 
    138     # plot
    139     fig = plt.figure(figsize=[10, 10])
    140     En = []
    141     for i in range(0, values+1):
    142         E, E_c, R, R_c = eigensolver.get_eigenpair(i)
    143         En.append(E)
    144         x = np.linspace(xmin, xmax, len(R))
    145 
    146         I = np.trapz(R*R, x)
    147    #     print(I)
    148 
    149         #plot eigenfunction
    150         plt.plot(x, 10*R*R+E, c='blue')
    151         plt.plot(x, R_c+E, lw=0.5, c='black')
    152         plt.annotate(f'$E_{ { i } } = {round(En[i], 3)}$', (xmax-0.1, 0.05+E), fontsize=14)
    153         plt.annotate(f'$|\Psi _{ { i } }|²$', (xmin, 0.05+E), fontsize=14)
    154 
    155     plt.axvline(x=xmin, c='red')
    156     plt.axvline(x=xmax, c='red')
    157     plt.ylim(0, max(En)+0.5)
    158     plt.title('Particle in \'infinite\' well', fontsize=20)
    159     plt.xticks([])
    160     plt.yticks([])
    161 
    162 
    163 
    164     # plot analytical solutions
    165     x = np.linspace(xmin, xmax, 70)
    166 
    167  #   print(En)
    168 
    169     for i in range(values+1, values+2):
    170         E, E_c, R, R_c = eigensolver.get_eigenpair(i)
    171     #    En.append(E)
    172 
    173     E = []
    174     for n in range(1, values+2):
    175 
    176        # particle in a box
    177         E.append(n**2 * np.pi**2 / (2*100))
    178         kn = n*np.pi/10
    179         if n%2 == 0:
    180             psi = sqrt(2/10) * np.sin(kn*x)
    181         if n%2 == 1:
    182             psi = sqrt(2/10) * np.cos(kn*x)
    183 
    184 
    185        # plot
    186 
    187         I = np.trapz(psi*psi, x)
    188        # print(I)
    189         plt.scatter(x, 0.5*psi*psi+En[n-1], s=20, c='green', marker="x")
    190 
    191 
    192 
    193   #  plt.tight_layout()
    194 
    195     numerical = mpatches.Patch(color='blue', label='numerical solution', linestyle='-')
    196     analytical = mpatches.Patch(color='green', label='analytical solution', linestyle='-')
    197     plt.legend(handles = [numerical, analytical], loc='upper left')
    198   #  plt.show()
    199    # plt.savefig('./plots/box_potential.png')
    200 
    201     return E, En
    202 
    203 def morse(xmin, xmax, mesh, V):
    204     xmin=xmin
    205     xmax=xmax
    206     mesh=mesh
    207     V=V
    208     D = 12
    209     a = 1
    210 
    211     def boundary(x, on_bnd):
    212         return on_bnd
    213 
    214     bc = DirichletBC(V, Constant(0.), boundary)
    215 
    216     u = TrialFunction(V)
    217     v = TestFunction(V)
    218 
    219     pot_ = Expression('D*pow((1-exp(a*(x[0]-x0))), 2)', D=D, a=a, x0=0, degree=2)
    220   #  pot_ = Expression('D*(exp(-2*a*(x[0]-x0))-2*exp(-a*(x[0]-x0)))', D=D, a=a, x0=0, degree=2)
    221 
    222     pot = interpolate(pot_, V)
    223 
    224     a = 1/2*inner(grad(u), grad(v))*dx + pot*u*v*dx
    225     m = u*v*dx
    226 
    227     A = PETScMatrix()
    228     M = PETScMatrix()
    229 
    230     assemble(a, tensor=A)
    231     assemble(m, tensor=M)
    232 
    233     bc.apply(A)
    234 
    235     # create eigensolver
    236     eigensolver = SLEPcEigenSolver(A, M)
    237     eigensolver.parameters['spectrum'] = 'smallest magnitude'
    238 
    239     # solve for eigenvalues
    240     values = 3
    241     eigensolver.solve(values)
    242 
    243     u = Function(V)
    244 
    245     # plot
    246     fig = plt.figure(figsize=[10, 10])
    247     En = []
    248     for i in range(0, values+1):
    249         E, E_c, R, R_c = eigensolver.get_eigenpair(i)
    250         En.append(E)
    251         x = np.linspace(xmin, xmax, len(R))
    252 
    253         I = np.trapz(R*R, x)
    254        # print(I)
    255 
    256         #plot eigenfunction
    257         plt.plot(x, 20*R*R+E, c='blue')
    258         plt.plot(x, R_c+E, lw=0.5, c='black')
    259         plt.annotate(f'$E_{ { i } } = {round(En[i], 3)}$', (xmax-0.1, 0.1+E), fontsize=14)
    260         plt.annotate(f'$|\Psi _{ { i } }|²$', (xmin, 0.2+E), fontsize=14)
    261 
    262     plt.plot(x, 12*(1-np.exp(-x))**2, c='red')
    263  #   plt.plot(x, 12*(np.exp(-2*x)-2*np.exp(-x)))
    264     plt.ylim(0, max(En)+2)
    265   #  plt.title('Particle in \'infinite\' well', fontsize=20)
    266     plt.title('Morse Potential $V(x) = D \cdot (1 - e^{- \\alpha (x - x0)})^2$', fontsize=20)
    267     plt.xticks([])
    268     plt.yticks([])
    269 
    270 
    271 
    272 
    273 
    274     # plot analytical solutions
    275     x = np.linspace(xmin, xmax, 70)
    276 
    277   #  print(En)
    278 
    279     E = []
    280 
    281     for n in range(0, values+1):
    282        # morse
    283 
    284         w = np.sqrt(2*12)/(2*np.pi)
    285         l = np.sqrt(2*D)
    286         z = 2*l*np.exp(-x)
    287         N = np.sqrt(factorial(n)*(2*l-2*n-1)/factorial(2*l-n-1))
    288         psi = N * z**(l-n-1/2) * np.exp(-z/2) * genlaguerre(n, 2*l-2*n-1)(z)
    289 
    290         E.append(-1/2*(l-n-1/2)**2 + 12)
    291       #  print(E)
    292       #  print(En)
    293 
    294        # plot
    295 
    296         I = np.trapz(psi*psi, x)
    297        # print(I)
    298         plt.scatter(x, psi*psi+En[n], s=20, c='green', marker="x")
    299 
    300 
    301 
    302   #  plt.tight_layout()
    303 
    304     numerical = mpatches.Patch(color='blue', label='numerical solution', linestyle='-')
    305     analytical = mpatches.Patch(color='green', label='analytical solution', linestyle='-')
    306     plt.legend(handles = [numerical, analytical], loc='upper left')
    307  #   plt.show()
    308 #    plt.savefig('./plots/morse_potential.png')
    309 
    310     return E, En
    311 
    312 def main():
    313 
    314     xmin = -5; xmax = 5
    315     s_h = []
    316     s_b = []
    317     s_m = []
    318     meshfinesse = []
    319 
    320     for n in range(20, 300, 20):
    321         meshfinesse.append(n)
    322         mesh = IntervalMesh(n, xmin, xmax)
    323         V = FunctionSpace(mesh, 'CG', 1)
    324         E_h_an, E_h_num = harmonic_oscillator(xmin, xmax, mesh, V)
    325         E_b_an, E_b_num = box(xmin, xmax, mesh, V)
    326         E_m_an, E_m_num = morse(xmin, xmax, mesh, V)
    327 
    328         d_h = np.subtract(E_h_an, E_h_num)/E_h_an
    329         s_h.append(np.average(np.abs(d_h)))
    330 
    331         d_b = np.subtract(E_b_an, E_b_num)/E_b_an
    332         s_b.append(np.average(np.abs(d_b)))
    333 
    334   #  print(E_m_an, E_m_num)
    335         d_m = np.subtract(E_m_an, E_m_num)/E_m_an
    336         s_m.append(np.average(np.abs(d_m)))
    337 
    338     s_h = np.array(s_h)
    339     s_b = np.array(s_b)
    340     s_m = np.array(s_m)
    341     print(s_h)
    342     print(s_b)
    343     print(s_m)
    344 
    345     fig = plt.figure()
    346     plt.plot(meshfinesse, 100*s_h, marker='x', c='green', label='harmonic oscillator')
    347     plt.plot(meshfinesse, 100*s_b, marker='x', c='red', label='box potential')
    348     plt.plot(meshfinesse, 100*s_m, marker='x', c='blue', label='morse potential')
    349     plt.title('Mittlerer prozentualer Fehler der Eigenwerte', fontsize=20)
    350     plt.ylabel('MAPE (%)')
    351     plt.xlabel('number of intervalls (meshsize = 10)')
    352 #    plt.xticks(np.arange(0, 1, 0.1))
    353 #    plt.gca().invert_xaxis()
    354 #    plt.xscale('log')
    355 
    356     plt.legend()
    357     plt.show()
    358 
    359 
    360 
    361 if __name__ == "__main__":
    362     main()