#!/usr/bin/env python2.3
# -*- coding: iso-8859-2 -*-
"""
Solution of 2D TD-Schrödinger equation.
Two-dimensional closed system.
Boundary condition: zero probability beyond
simulation space (infinite potential gate)
Static potential: homogenous (zero)
Requires Python 2.3.4 with scipy and visual.
See: www.python.com, use the package index
AU: Atomic Units
"""

from scipy import *
from psyco import bind as jit
import visual as vi
import operator as op

#=========================================================================
# Settings of visualization 

# Probability scaling
scale=1

#=========================================================================
# Boundary conditions

# Duration and number of iterations [AU]
duration=1.0
iterations=1000

# Space frame and resolution [AU]
frame=(2.0,2.0)
resolution=(48,48)

#=========================================================================
# Precalculations

# Calculate time step and span of grid points
step=duration/iterations
span=vectorize(op.div)(frame,resolution)

#=========================================================================
# Initial state of the wave function

n=1 # Wave number
c=reduce(op.mul,[sqrt(2.0/x) for x in frame])
w=[n*pi/(r+1) for r in resolution]
def fn(x,y):
    return c*sin(w[0]*(x+1))*sin(w[1]*(y+1))
jit(fn)
WF=array([
    [fn(x,y) for x in xrange(resolution[0])]
    for y in xrange(resolution[1])
],Complex)

# Constants
common=-4j*span[0]**2/step
A=common+2
B=common-2

# Precalculate the time evolution operator
def fn(i,j):
    if i==j: return A
    if abs(i-j)==1: return -1.0
    return 0.0
jit(fn)
AM=array([[fn(i,j) for j in xrange(resolution[0])] for i in xrange(resolution[0])],Complex)
def fn(i,j):
    if i==j: return B
    if abs(i-j)==1: return 1.0
    return 0.0
jit(fn)
BM=array([[fn(i,j) for j in xrange(resolution[0])] for i in xrange(resolution[0])],Complex)
TE=dot(linalg.inv(AM),BM)

# Drawing frame
vi.scene.autoscale=0
size=sqrt(frame[0]**2+frame[1]**2+0.5**2)/sqrt(2)
vi.scene.range=(size,size,size)
vi.scene.center=(frame[0]/2,frame[1]/2,0)
frm=vi.frame() #pos=(-frame[0]/2,-frame[1]/2,-0.5))
frm.rotate(axis=(1,0,0),angle=-pi/3)

# Draw axis
vi.arrow(frame=frm,pos=(0,0,0),axis=(frame[0]+0.4,0,0),color=(0.7,0.7,0.7),shaftwidth=0.05)
vi.label(frame=frm,pos=(frame[0]+0.6,0,0),text='x',box=0)
vi.arrow(frame=frm,pos=(0,0,0),axis=(0,frame[1]+0.4,0),color=(0.7,0.7,0.7),shaftwidth=0.05)
vi.label(frame=frm,pos=(0,frame[1]+0.6,0),text='y',box=0)
vi.arrow(frame=frm,pos=(0,0,0),axis=(0,0,1.4),color=(0.7,0.7,0.7),shaftwidth=0.024)
vi.label(frame=frm,pos=(0,0,1.6),text='p',box=0)

# Draw probability layer
def fn(x,y):
    return vi.convex(
        frame=frm,
        pos=[(span[0]*x,span[0]*y,0),
             (span[0]*(x+1),span[0]*y,0),
             (span[0]*(x+1),span[0]*(y+1),0),
             (span[0]*x,span[0]*(y+1),0)],
        color=(0,1,1))     
jit(fn)
seg=[[fn(x,y) for x in xrange(resolution[0]+1)]
        for y in xrange(resolution[1]+1)]

# Simulation
LU=linalg.lu_factor(TE)
def Display(PD):
    """Display probability distribution"""
    for y,pr in enumerate(PD):
        for x,p in enumerate(pr):
            seg[y][x].pos[2][2]=p
            seg[y][x+1].pos[3][2]=p
            seg[y+1][x+1].pos[0][2]=p
            seg[y+1][x].pos[1][2]=p
jit(Display)
def TimeStep():
    """Calculate next state of the wave function"""
    global WF
    for n in range(2):
        for i,r in enumerate(WF): WF[i]=linalg.lu_solve(LU,r)
        WF=transpose(WF)
jit(TimeStep)
t=0
while 1:
    vi.rate(20)
    # Draw state
    Display((WF*conj(WF)).real*scale)
    # Time step
    TimeStep()
    t+=step