User:Thierry Dugnolle/Python/Wave-packet

A gaussian wave packet. The color (yellow green blue) indicates the phase of the wave function $\psi$, its intensity indicates $|\psi |^{2}/|\psi |_{max}^{2}$. This wave function could represent the motion of a quantum particle. The wave function is $\psi (x,y,t)=\psi (x,t)\psi (y,t)$, $k_{0x}=k_{0}$, $k_{0y}=0$, $\psi (x,t)=({\frac {2a^{2}}{\pi }})^{1/4}{\frac {e^{i\phi }}{(a^{4}+{\frac {4\hbar ^{2}t^{2}}{m^{2}}})^{1/4}}}e^{ik_{0}x}exp[-{\frac {(x-{\frac {\hbar k_{0}}{m}}t)^{2}}{a^{2}+{\frac {2i\hbar t}{m}}}}]$ where $\phi =-\theta -{\frac {\hbar k_{0}^{2}}{2m}}t$ and $tg(2\theta )={\frac {2\hbar t}{ma^{2}}}$ (source: Cohen-Tannoudji, Diu & Laloë, Quantum Mechanics, complement G_I, §3-a)

print("Wave_packet1")

from PIL import Image, ImageDraw
import random
import math
import cmath

imWidth = 800
imHeight = 400
lengthUnit = 100.0 # length unit in number of pixels
a = 1.5  # packet width
m = 8.0  # packet mass
kX = 8.0
kY = 0.0
dt = 0.04
imNumber = 300

WnumPixelCenter = -1.5*lengthUnit
HnumPixelCenter = imHeight/2.0 - 0.5

im = Image.new("RGB", (imWidth, imHeight), (0, 0, 0))
px = im.load()
def psi(k0, x, t):
    r1 = pow(pow(a, 4) + 4 * t * t / (m * m), 0.25)
    theta = 0.5 * math.atan(2 * t / (m * a * a))
    phi = -theta - k0 * k0 * t / (2.0 * m)
    c1 = pow(math.e, 1j * (phi + k0 * x) - pow( x - k0 * t / m, 2)/ (a*a+1j*2.0*t/m))
    return pow(2*a*a/math.pi,0.25) * c1 / r1

for imNum in range(imNumber):
    t = imNum * dt
    xCenter = t * kX / m
    yCenter = t * kY / m
    psiXmax = psi(kX, xCenter, t)
    psiYmax = psi(kY, yCenter, t)
    psiXYmax = psiXmax * psiYmax
    psiModmax = abs(psiXYmax)
    psiModmax2 = psiModmax * psiModmax

    for i in range(imWidth):
        for j in range(imHeight):
            x = (i - WnumPixelCenter) / lengthUnit
            y = (HnumPixelCenter - j) / lengthUnit
            psiX = psi(kX, x, t)
            psiY = psi(kY, y, t)
            psiXY = psiX * psiY
            psiPhase = cmath.phase(psiXY)
            psiPhaseNorm= (psiPhase+math.pi)/(2.0*math.pi)
            psiMod = abs(psiXY)
            bright= psiMod*psiMod/psiModmax2
            if psiPhaseNorm<1.0/3.0:
                red=1.0-3.0*psiPhaseNorm
                blue=0.0
            else:
                if psiPhaseNorm<2.0/3.0:
                    red=0.0
                    blue=3.0*(psiPhaseNorm-1.0/3.0)
                else:
                    red=3.0*(psiPhaseNorm-2.0/3.0)
                    blue=1.0-3.0*(psiPhaseNorm-2.0/3.0)
            px[i,j] = (math.floor(255.0*red*bright), math.floor(255.0*bright), math.floor(255.0*blue*bright))

    im.save("WP" + str(1000 + imNum) + ".bmp")

print("Good bye")