from visual import *
from random import random
from time import clock

scene.width = 800
scene.height = 400
scene.autoscale = 0
scene.up = (0,0,1)
scene.forward = (0,-1,-.5)

Ntotx = 20
Ntoty = 1
Ntotz = 1
Ntotal = Ntotx*Ntoty*Ntotz

spins = []

def lattice(deltax=1.0, deltay=1.0):
    xmin = -Lx/2.
    ymin = -Ly/2
    nz = 0
    
    for nx in range(Ntotx):
          x = xmin + nx*deltax
          
          for ny in range(Ntoty):
              y = ymin + ny*deltay

              spins.append(frame())
              spins[-1].pos = vector(x,y,0)
          
              spins[-1].spin = vector()
              spins[-1].spinvec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.9,.9,.9), shaftwidth=0.2)
          
              spins[-1].torque = vector()
              spins[-1].torquevec = arrow(pos=(x,y,0), axis=(0,0,1), color=(.3,.3,1), shaftwidth=0.05)
          
              spins[-1].nearx = range(2)
              spins[-1].neary = range(2)
                
              spins[-1].indices = (nx,ny,nz)
          
    for s in spins:
          nx, ny, nz = s.indices
          if nx == 0:                       # leftmost spin in a row
              nspinl = Ntotx*ny + Ntotx-1       # wrap around to spin on right side
              s.nearx[0] = nspinl               # reference this by its list element number, given by the order in which the spins are added to the list...
          else:
              nspinl = Ntotx*ny + nx-1
              s.nearx[0] = nspinl
          if nx == Ntotx-1:                 # rightmost spin in a row
              nspinr = Ntotx*ny
              s.nearx[1] = nspinr
          else:
              nspinr = Ntotx*ny + nx+1 
              s.nearx[1] = nspinr

          if ny == 0:                       # bottom spin in a column
              nspind = Ntotx*(Ntoty-1) + nx 
              s.neary[0] = nspind
          else:
              nspind = Ntotx*(ny-1) + nx
              s.neary[0] = nspind
          if ny == Ntoty-1:                 # top spin in a column
              nspinu = nx
              s.neary[1] = nspinu
          else:
              nspinu = Ntotx*(ny+1) + nx
              s.neary[1] = nspinu
    return spins

dx = 2
dy = 2
Lx = dx*Ntotx
Ly = dy*Ntoty

spins = lattice(deltax=dx, deltay=dy)

dt = .01
J = 1
k = pi*vector(4./Lx, 0, 0)

spinvec_length = 3
#cons_torque_mag = 0
sigma = .4

for s in spins:
    spinx = sigma*cos(dot(k, s.pos))
    spiny = sigma*sin(dot(k, s.pos))
    spinz = 1*(J/abs(J))**(s.indices[0] + s.indices[1])
    
    s.spin = vector(spinx,spiny,spinz)

    s.pos.z = s.pos.z + 0*(-1)**(s.indices[0] + s.indices[1])
    
    s.spinvec.axis = s.spin*spinvec_length
    s.spinvec.pos = s.pos - s.spinvec.axis/2.


tt = clock()
Nsteps = 0
while 1:
    for s in spins:
        cons_spinz = s.spin.z
        s.torque = vector(0,0,0)

        for nnx in range(2):
            nspin = s.nearx[nnx]
            s.torque = s.torque+cross(s.spin,spins[nspin].spin)
        for nny in range(2):
            npsin = s.neary[nny]
            s.torque = s.torque+cross(s.spin,spins[nspin].spin)
                
        s.torque.z = 0

#        if Nsteps == 0:
#            cons_torque_mag = mag(s.torque)

#        s.torque = cons_torque_mag*s.torque/mag(s.torque)
        
        s.spin = s.spin+J*s.torque*dt
        s.spin.z = 0
        s.spin = sigma*s.spin/mag(s.spin)
        s.spin.z = cons_spinz

        s.spin = s.spin/mag(s.spin)
        
        s.spinvec.axis = s.spin*spinvec_length
        s.spinvec.pos = s.pos - s.spinvec.axis/2.

        s.torquevec.axis = .5*J*s.torque/(mag(s.torque)+.001)
        s.torquevec.pos = s.pos + s.spinvec.axis/2.


    if Nsteps == 1000:
        tt = clock()-tt
        print '%0.1f' % tt, 'seconds for', Nsteps, 'steps with', Ntotal, 'spins'
    Nsteps = Nsteps+1


#Times:

#dt = .05,
#w/ torque vectors:   9.4 seconds for 1000 steps with 40 spins
#w/o torque vectors:  7.7 seconds for 1000 steps with 40 spins
