from scipy import *
from pylab import *

class Measure:
   def __init__(self, every_n, p):
      self.every_n = every_n
      self.p = p
      self.rsq = zeros((n,n),dtype=float)
      self.num_meas = 0
      self.list_rsq = []
      for x in range(n):
          for y in range(n):
              self.rsq[x,y]=(x-n/2)**2+(y-n/2)**2
        

   def ave_rsq(self):
       prob_rsq = self.rsq*abs(self.p)**2
       return sum(prob_rsq)

   def norm(self):
       prob = abs(self.p)**2
       return sum(prob)
   
   def do_measurements(self):
       self.num_meas += 1
       if (self.num_meas%self.every_n == 0):
           self.list_rsq.append(self.ave_rsq())
           print "rsq = ", self.list_rsq[-1], "norm = ", self.norm()
           
      
   

n = 128
dt = 0.001
n2 = n/2
eps = (0.0+1.0j)*dt

p = zeros((n,n),dtype=complex)
p[n2,n2] = 1

meas = Measure(100,p)


ion()
for k in range(1000):
   p[1:-1,1:-1] += eps*(p[0:-2,1:-1]+p[2:,1:-1]+ p[1:-1,0:-2]+p[1:-1,2:])
   meas.do_measurements()
   if k%100 == 0:
      imshow(abs(p)**2)
      show()