SF.net SVN: matplotlib: [3625] trunk/py4science/examples

[fer_perez]

Revision: 3625
  http://matplotlib.svn.sourceforge.net/matplotlib/?rev=3625&view=rev
Author:   fer_perez
Date: 2007-07-28 14:31:18 -0700 (Sat, 28 Jul 2007)

Log Message:
---
Add Schrodinger equation example, contributed by James Nagel <[EMAIL PROTECTED]>

Added Paths:
---
trunk/py4science/examples/schrodinger/
trunk/py4science/examples/schrodinger/Schrodinger_FDTD.pdf
trunk/py4science/examples/schrodinger/schrod_fdtd.py

Added: trunk/py4science/examples/schrodinger/Schrodinger_FDTD.pdf
===
(Binary files differ)


Property changes on: trunk/py4science/examples/schrodinger/Schrodinger_FDTD.pdf
___
Name: svn:mime-type
   + application/octet-stream

Added: trunk/py4science/examples/schrodinger/schrod_fdtd.py
===
--- trunk/py4science/examples/schrodinger/schrod_fdtd.py
(rev 0)
+++ trunk/py4science/examples/schrodinger/schrod_fdtd.py2007-07-28 
21:31:18 UTC (rev 3625)
@@ -0,0 +1,258 @@
+#!/usr/bin/env python
+#=
+#
+#   Quantum Mechanical Simulation using Finite-Difference
+#   Time-Domain (FDTD) Method
+#
+#   This script simulates a probability wave in the presence of multiple
+#   potentials.  The simulation is c arried out by using the FDTD algorithm
+#   applied to the Schrodinger equation.  The program is intended to act as
+#   a demonstration of the FDTD algorithm and can be used as an educational
+#   aid for quantum mechanics and numerical methods.  The simulation
+#   parameters are defined in the code constants and can be freely
+#   manipulated to see different behaviors.
+#
+#   NOTES
+#
+#   The probability density plots are amplified by a factor for visual
+#   purposes.  The psi_p quanity contains the actual probability density
+#   without any rescaling.
+#
+#   BEWARE: The time step, dt, has strict requirements or else the
+#   simulation becomes unstable.
+#
+#   The code has three built-in potential functions for demonstration.
+#
+#   1) Constant potential: Demonstrates a free particle with dispersion.
+#
+#   2) Step potential: Demonstrates transmission and reflection.
+#
+#   3) Potential barrier: Demonstrates tunneling.
+#
+#   By tweaking the height of the potential (V0 below) as well as the
+#   barrier thickness (THCK below), you can see different behaviors: full
+#   reflection with no noticeable transmission, transmission and
+#   reflection, or mostly transmission with tunneling.
+#
+#   This script requires pylab and numpy to be installed with
+#   Python or else it will not run.
+#
+#
+# Author:  James Nagel <[EMAIL PROTECTED]>
+#  5/25/07
+#
+# Updates by Fernando Perez <[EMAIL PROTECTED]>, 7/28/07
+#
+
+#  Numerical and plotting libraries
+import numpy as np
+import pylab
+
+# Set pylab to interactive mode so plots update when run outside ipython
+pylab.ion()
+
+#=
+
+# Utility functions
+
+#  Defines a quick Gaussian pulse function to act as an envelope to the wave
+#  function.
+def Gaussian(x,t,sigma):
+"""  A Gaussian curve.
+x = Variable
+t = time shift
+sigma = standard deviation  """
+return np.exp(-(x-t)**2/(2*sigma**2))
+
+def free(npts):
+"Free particle."
+return np.zeros(npts)
+
+def step(npts,v0):
+"Potential step"
+v = free(npts)
+v[npts/2:] = v0
+return v
+
+def barrier(npts,v0,thickness):
+"Barrier potential"
+v = free(npts)
+v[npts/2:npts/2+thickness] = v0
+return v
+
+#=
+#
+#  Simulation Constants.  Be sure to include decimal points on appropriate
+#  variables so they become floats instead of integers.
+#
+N= 1200 #  Number of spatial points.
+T= 5*N  #  Number of time steps.  5*N is a nice value for terminating
+#  before anything reaches the boundaries.
+Tp   = 40   #  Number of time steps to increment before updating the plot.
+dx   = 1.0e0#  Spatial resolution
+m= 1.0e0#  Particle mass
+hbar = 1.0e0#  Plank's constant
+X= dx*np.linspace(0,N,N)#  Spatial axis.
+
+# Potential parameters.  By playing with the type of potential and the height
+# and thickness (for barriers), you'll see the various transmission/reflection
+# regimes of quantum mechanical tunneling.
+V0   = 1.0e-2   #  Potential amplitude (used for steps and barriers)
+THCK = 10  

SF.net SVN: matplotlib: [3626] trunk/py4science/examples/schrodinger/ schrod_fdtd.py

[fer_perez]

Revision: 3626
  http://matplotlib.svn.sourceforge.net/matplotlib/?rev=3626&view=rev
Author:   fer_perez
Date: 2007-07-28 23:50:34 -0700 (Sat, 28 Jul 2007)

Log Message:
---
Add energy display in the same units as the potential

Modified Paths:
--
trunk/py4science/examples/schrodinger/schrod_fdtd.py

Modified: trunk/py4science/examples/schrodinger/schrod_fdtd.py
===
--- trunk/py4science/examples/schrodinger/schrod_fdtd.py2007-07-28 
21:31:18 UTC (rev 3625)
+++ trunk/py4science/examples/schrodinger/schrod_fdtd.py2007-07-29 
06:50:34 UTC (rev 3626)
@@ -88,7 +88,7 @@
 N= 1200 #  Number of spatial points.
 T= 5*N  #  Number of time steps.  5*N is a nice value for terminating
 #  before anything reaches the boundaries.
-Tp   = 40   #  Number of time steps to increment before updating the plot.
+Tp   = 50   #  Number of time steps to increment before updating the plot.
 dx   = 1.0e0#  Spatial resolution
 m= 1.0e0#  Particle mass
 hbar = 1.0e0#  Plank's constant
@@ -98,7 +98,7 @@
 # and thickness (for barriers), you'll see the various transmission/reflection
 # regimes of quantum mechanical tunneling.
 V0   = 1.0e-2   #  Potential amplitude (used for steps and barriers)
-THCK = 10   # "Thickness" of the potential barrier (if appropriate
+THCK = 15   # "Thickness" of the potential barrier (if appropriate
 # V-function is chosen)
 
 # Uncomment the potential type you want to use here:
@@ -108,12 +108,12 @@
 
 # Potential step.  The height (V0) of the potential chosen above will determine
 # the amount of reflection/transmission you'll observe
-#POTENTIAL = 'step'
+POTENTIAL = 'step'
 
 # Potential barrier.  Note that BOTH the potential height (V0) and thickness
 # of the barrier (THCK) affect the amount of tunneling vs reflection you'll
 # observe. 
-POTENTIAL = 'barrier'
+#POTENTIAL = 'barrier'
 
 #  Initial wave function constants
 sigma = 40.0 # Standard deviation on the Gaussian envelope (remember Heisenberg
@@ -121,6 +121,11 @@
 x0 = round(N/2) - 5*sigma # Time shift
 k0 = np.pi/20 # Wavenumber (note that energy is a function of k)
 
+# Energy for a localized gaussian wavepacket interacting with a localized
+# potential (so the interaction term can be neglected by computing the energy
+# integral over a region where V=0)
+E = (hbar**2/2.0/m)*(k0**2+0.5/sigma**2)
+
 #=
 # Code begins
 #
@@ -140,7 +145,7 @@
 
 #  More simulation parameters.  The maximum stable time step is a function of
 #  the potential, V.
-Vmax = max(V)#  Maximum potential of the domain.
+Vmax = V.max()#  Maximum potential of the domain.
 dt   = hbar/(2*hbar**2/(m*dx**2)+Vmax) #  Critical time step.
 c1   = hbar*dt/(m*dx**2)   #  Constant coefficient 1.
 c2   = 2*dt/hbar   #  Constant coefficient 2.
@@ -148,6 +153,7 @@
 
 # Print summary info
 print 'One-dimensional Schrodinger equation - time evolution'
+print 'Wavepacket energy:   ',E
 print 'Potential type:  ',POTENTIAL
 print 'Potential height V0: ',V0
 print 'Barrier thickness:   ',THCK
@@ -190,30 +196,40 @@
 
 #  Initialize the figure and axes.
 pylab.figure()
+xmin = X.min()
+xmax = X.max()
 ymax = 1.5*(psi_r[PR]).max()
-pylab.axis([X.min(),X.max(),-ymax,ymax])
+pylab.axis([xmin,xmax,-ymax,ymax])
 
 #  Initialize the plots with their own line objects.  The figures plot MUCH
 #  faster if you simply update the lines as opposed to redrawing the entire
 #  figure.  For reference, include the potential function as well.
-lineR, = pylab.plot(X,psi_r[PR],'b',alpha=0.7)#  "Real" line
-lineI, = pylab.plot(X,psi_i[PR],'r',alpha=0.7)#  "Imag" line.
-lineP, = pylab.plot(X,psi_p,'k')   #  "Probability" line
+lineR, = pylab.plot(X,psi_r[PR],'b',alpha=0.7,label='Real')
+lineI, = pylab.plot(X,psi_i[PR],'r',alpha=0.7,label='Imag')
+lineP, = pylab.plot(X,6*psi_p,'k',label='Prob')
 pylab.title('Potential height: %.2e' % V0)
-pylab.legend(('Real','Imag','Prob'))
 
 # For non-zero potentials, plot them and shade the classically forbidden region
-# in light red
-if V.max() !=0 :
-V_plot = V/V.max()*ymax/2
+# in light red, as well as drawing a green line at the wavepacket's total
+# energy, in the same units the potential is being plotted.
+if Vmax !=0 :
+Efac = ymax/2.0/Vmax
+V_plot = V*Efac
 pylab.plot(X,V_plot,':k',zorder=0)   #  Potential line.
 # reverse x and y2 so the polygon fills in order
 y1 = free(N)  # lower boundary for polygon drawing
 x = np.concatenate( (X,X[::-1]) )
 y = np.concatenate( (y1,V_plot[::-1]) )
 pylab.fill(x, y, facecolor='y', alpha=0.2,zorder=0)
+# Plot the wavefunction energy, in the same scale as the potential
+pylab.axhline(E*Efac,color='g',label='Energy'