from Numeric import *
from Matrix import *
from LinearAlgebra import *

from visual import *


class Hamiltonian:
    nSpinsX = 4
    nSpinsY = 4
    nSpinsTot = nSpinsX*nSpinsY

    Ja = -1.0
    Jbx = -1.0
    Jby = -1.0


    zArray = []
    
    couplingXArray = []
    couplingYArray = []

    couplingPhaseXArray = []
    couplingPhaseYArray = []

    meanFieldFromZArray = []

    HArray = []             # hurray!!!!


    def __init__(self, nSpinsX, nSpinsY):
        self.nSpinsX = nSpinsX
        self.nSpinsY = nSpinsY
        self.nSpinsTot = nSpinsX*nSpinsY

        self.zArray = zeros((nSpinsX,nSpinsY),Float32)

        
        self.couplingXArray = zeros((nSpinsX,nSpinsX),Complex32)
        self.couplingYArray = zeros((nSpinsY,nSpinsY),Complex32)

        self.couplingPhaseXArray = zeros((nSpinsX,nSpinsX),Complex32)
        self.couplingPhaseYArray = zeros((nSpinsY,nSpinsY),Complex32)

        self.meanFieldFromZArray = zeros((nSpinsX,nSpinsY),Float32)

        self.HArray = zeros((self.nSpinsTot*2,self.nSpinsTot*2),Complex32)



    def setzArray(self):
        for x in range(self.nSpinsX):
            for y in range(self.nSpinsY):
                self.zArray[x][y] = (-1)**(x+y)     # your typical interior antiferromagnetically textured unit cell
                

    def setCouplings(self, Ja, Jbx, Jby, k):
        self.Ja = Ja
        self.Jbx = Jbx
        self.Jby = Jby
      
        for i in range(self.nSpinsX):
            for j in range(self.nSpinsX):
                if i - j == 1:                      self.couplingXArray[i][j] += self.Ja
                if j - i == 1:                      self.couplingXArray[i][j] += self.Ja
                if i - j == self.nSpinsX - 1:       self.couplingXArray[i][j] += self.Jbx
                if j - i == self.nSpinsX - 1:       self.couplingXArray[i][j] += self.Jbx

        for i in range(self.nSpinsY):
            for j in range(self.nSpinsY):
                if i - j == 1:                      self.couplingYArray[i][j] += self.Ja
                if j - i == 1:                      self.couplingYArray[i][j] += self.Ja
                if i - j == self.nSpinsY - 1:       self.couplingYArray[i][j] += self.Jby
                if j - i == self.nSpinsY - 1:       self.couplingYArray[i][j] += self.Jby



    def getMeanFieldFromZ(self):
        latticeArray = self.zArray
        flippedLatticeArray = transpose(self.zArray)

        meanFieldXArray = matrixmultiply(flippedLatticeArray,self.couplingXArray)
        meanFieldYArray = matrixmultiply(latticeArray,self.couplingYArray)
        
        meanFieldXArray = transpose(meanFieldXArray)
        meanFieldYArray = meanFieldYArray
        
        self.meanFieldFromZArray = meanFieldXArray + meanFieldYArray





    def makeHamiltonian(self, baseSigma):               
        for i in range(self.nSpinsTot*2):                           # i is target spin and component
            for j in range(self.nSpinsTot*2):                       # j is source spin and component
                                                                    # i = j implies xTarget == xSource and yTarget == ySource (target is source)
                if i < self.nSpinsTot and j >= self.nSpinsTot:      # x components in terms of y components
                    i = i%self.nSpinsTot
                    j = j%self.nSpinsTot

                    xTarget = (i-i%self.nSpinsY)/self.nSpinsY
                    yTarget = i%self.nSpinsY
                    targetZ = self.zArray[xTarget][yTarget]

                    xSource = (j-j%self.nSpinsY)/self.nSpinsY
                    ySource = j%self.nSpinsY
                    sourceZ = self.zArray[xSource][ySource]

                    if i == j:                                  self.HArray[i][j+self.nSpinsTot] = self.meanFieldFromZArray[xTarget,yTarget]/targetZ
                        
                    if xTarget - xSource == 1:                  self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Ja*e**(-1j*k.x)
                    if xSource - xTarget == 1:                  self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Ja*e**(1j*k.x)
                    if xTarget - xSource == self.nSpinsX - 1:   self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Jbx*e**(1j*k.x)
                    if xSource - xTarget == self.nSpinsX - 1:   self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Jbx*e**(-1j*k.x)

                    if yTarget - ySource == 1:                  self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Ja*e**(-1j*k.y)
                    if ySource - yTarget == 1:                  self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Ja*e**(1j*k.y)
                    if yTarget - ySource == self.nSpinsY - 1:   self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Jby*e**(1j*k.y)
                    if ySource - yTarget == self.nSpinsY - 1:   self.HArray[i][j+self.nSpinsTot] += -baseSigma*self.Jby*e**(-1j*k.y)


                if j < self.nSpinsTot and i >= self.nSpinsTot:      # y components in terms of x components
                    self.HArray[i][j] = self.HArray[i-self.nSpinsTot][j+self.nSpinsTot]



    def getHArray(self):
        return self.HArray
        

        
        
nSpinsX = 2
nSpinsY = 2

Ja = -1.0
Jbx = -1.0
Jby = -1.0

k = pi*vector(0.5,0.0,0.0)
baseSigma = .001





H = Hamiltonian(nSpinsX, nSpinsY)

H.setzArray()
H.setCouplings(Ja, Jbx, Jby, k)
H.getMeanFieldFromZ()

H.makeHamiltonian(baseSigma)






Hamiltonian = H.getHArray()
print '\nHamiltonian:\n', Hamiltonian

evalues, evectors = eigenvectors(Hamiltonian)
print '\neigenvalues of Hamiltonian:\n', evalues
print '\neigenvectors of Hamiltonian:\n', evectors







##print '\nzArray: \n', H.zArray
##print '\ncouplingXArray: \n', H.couplingXArray
##print '\ncouplingYArray: \n', H.couplingYArray
##print '\nmeanFieldFromZArray: \n', H.meanFieldFromZArray
##print '\nHArray: \n', H.HArray
##print '\nHermitian if identically Zero: \n', H.HArray - transpose(conjugate(H.HArray))


##inv_ham = inverse(ham)
##print '\ninverse of ham:\n', inv_ham




## variation of eigenvalues with baseSigma:
##
##baseSigma = .001
##eigenvalues of Hamiltonian:
##[ 4.00780423+2.69284734e-025j  4.        -1.14654832e-025j
##       3.996     +1.10886490e-026j  3.99619577+2.40891072e-025j
##      -4.00780423-1.66410499e-025j -4.        -2.26477832e-026j
##      -3.996     +5.94012095e-026j -3.99619577-3.42336268e-025j]
##
##baseSigma = .01
##eigenvalues of Hamiltonian:
##[ 4.07804226-5.89334296e-026j  4.        +3.98522122e-025j
##       3.96      +1.06931991e-024j  3.96195774+1.19308793e-024j
##      -4.07804226-1.10728034e-025j -4.        -6.04470267e-025j
##      -3.96      +2.07818947e-026j -3.96195774+2.88186257e-026j]
##
##baseSigma = .1
##eigenvalues of Hamiltonian:
##[-4.78042258+4.08775559e-025j -3.99999999-7.79371735e-025j
##      -3.61957742+2.80283148e-024j -3.60000001+1.80913449e-025j
##       4.78042258-2.50464416e-025j  3.99999999-1.00465643e-024j
##       3.60000001+3.64937983e-026j  3.61957742-3.17406076e-025j]
##
##baseSigma = 1.0
##eigenvalues of Hamiltonian:
##[ 1.18042260e+001-1.70903730e-024j -1.18042260e+001+1.10725545e-024j
##       3.99999988e+000+1.14490842e-024j -3.99999988e+000-1.29048113e-024j
##       1.95773959e-001+1.17675107e-025j -1.95773959e-001+3.84872227e-025j
##       1.19209289e-007+4.28040459e-026j -1.19209290e-007+3.10612471e-026j]
##
##baseSigma = 10.0
##eigenvalues of Hamiltonian:
##[-82.04225922-1.16512349e-024j  82.04225922+3.61075799e-025j
##      -36.        +1.21335014e-024j -34.04225922+1.96676555e-024j
##       -4.        +2.41652124e-024j   4.        -2.32778874e-024j
##       36.        +6.49538909e-025j  34.04225922+2.13502659e-024j]
##
##baseSigma = 100.0
##eigenvalues of Hamiltonian:
##[-784.42260742-2.73757935e-023j  784.42260742+4.61909224e-023j
##      -396.        +1.95412050e-023j -376.42260742-1.06771902e-023j
##       396.        -4.08086009e-023j  376.42260742-4.85054643e-023j
##         4.        +2.48946222e-023j   -4.        +6.07078847e-023j]









