factored out some code, added some documentation

Author: dev-zero

schroedinger/plot.py

@@ -9,16 +9,23 @@
 
 from schroedinger import Schroedinger1D
 
-import gobject
-import matplotlib
-matplotlib.use('GTKAgg')
+try:
+    import gobject
+    import matplotlib
+    matplotlib.use('GTKAgg')
+except:
+    from sys import exit
+    print("This app works only with GTK+ installed and matplotlib built with GTK+ support.")
+    exit(1)
 
 from pylab import plot, show, figure
 from numpy import abs, real, imag, vectorize
 
+# initialize the two simulations
 s1 = Schroedinger1D(256, L=20)
 s2 = Schroedinger1D(256, L=20)
 
+# models the double well potential
 def doubleWellPotential(x):
     sigma = 0.04
     pos = 2.
@@ -27,6 +34,7 @@ def doubleWellPotential(x):
         return 0.
     return 10.
 
+# use the double well potential in the second plot
 s2.setPotential(vectorize(doubleWellPotential))
 
 fig = figure()
@@ -53,8 +61,8 @@ def doubleWellPotential(x):
 def update():
     global i
     i += 1
-    s1.evolve(0.01)
-    s2.evolve(0.01)
+    s1.evolve(0.02)
+    s2.evolve(0.02)
     s1_l_amp.set_ydata(abs(s1.f()))
     s1_l_real.set_ydata(real(s1.f()))
     s1_l_imag.set_ydata(imag(s1.f()))
@@ -64,7 +72,7 @@ def update():
     fig.canvas.draw_idle()
     if not runSimulation:
         print "stopped simulation after {} iterations".format(i)
-    return runSimulation
+    return runSimulation # return False terminates this gobject-callback
 
 def onpress(event):
     global runSimulation
@@ -75,7 +83,7 @@ def onpress(event):
         else:
             print "starting simulation"
             runSimulation = True
-            gobject.idle_add(update)
+            gobject.idle_add(update) # readding the gobject-callback function
 
 fig.canvas.mpl_connect('key_press_event', onpress)
 show()

schroedinger/schroedinger.py

@@ -13,8 +13,11 @@
 def harmonicPotential(x):
     return 0.5*x**2
 
+def wavePacket(x, k_0):
+    return exp(-0.5*x**2 + 1j*k_0*x)
 
 class Schroedinger1D:
+    ''' This class encapsulates the 1D schroedinger equation solver '''
     def __init__(self, N=256, L=10):
         self._N = N
         self._L = L
@@ -24,12 +27,20 @@ def __init__(self, N=256, L=10):
         self._dk = 2.*pi / (N*self._dx)
 
         self._x = (arange(N) - N/2) * self._dx
+        # this is the grid in momentum space. Instead of realigning self._F after fft and before momentumEvolve, resp again
+        # after momentumEvolve and before ifft I simply transform the grid itself (see example (6.7) in the lecture notes)
         k = (arange(N) - N/2)
         self._k = concatenate((k[self._N//2:], k[0:self._N//2])) * self._dk
-        self._f = exp(-0.5*self._x**2 + 2.*self._x*1j)
+
+        # a simple wave-package moving constantly to the right and starting at x=0
+        self.setInitial(lambda x: wavePacket(x, 2.))
+
+        # use the harmonic potential as default potential
         self._V = harmonicPotential
 
     def evolve(self, dt = 0.001):
+        ''' one iteration-step '''
+
         def spatialEvolve(f, x):
             return f*exp(-0.5 * self._V(x) * dt * 1j)
         def momentumEvolve(F, k):
@@ -54,8 +65,13 @@ def V(self):
         return self._V(self._x)
 
     def setPotential(self, V):
+        ''' set the potential (as a function of an array of x-values)'''
         self._V = V
 
+    def setInitial(self, f):
+        ''' set the initial conditions (as a function of x of an array of x-values)'''
+        self._f = f(self._x)
+
 if __name__ == '__main__':
     from matplotlib.pyplot import plot, show