# -*- coding: utf-8 -*-
# Fluido de Lennard-Jones, 2D, Monte Carlo
# Programa escrito por Leonard Sander, traduzido por André Vieira
# As unidades são as seguintes:
# Comprimento=sigma, energia=epsilon, massa=u.m.a.
# Condições de contorno periódicas
from __future__ import division, print_function
from visual import *
from visual.graph import *
import numpy as np

# Todas as partículas têm massas unitárias, logo forças e acelerações são numericamente iguais.
# As partículas ocupam um quadrado de lado L, com -L/2 < x < L/2 e -L/2 < y < L/2
# As partículas interagem segundo o potencial de Lennard-Jones com um corte.

def forces(ux,uy,N,L):   # Calcula as forças sobre as partículas
    # Leva em conta as condições de contorno periódicas
    rx=ux - ux.reshape(N,1)   # Constrói a matriz de diferenças entre as componentes x das posições
    trimx=greater(rx,L/2)     # Identifica as diferenças > L/2
    rx=rx-L*trimx                  # Subtrai L dos elementos correspondentes
    trimx=less(rx,-L/2)       # Identifica as diferenças < -L/2
    rx=rx+L*trimx                  # Soma L a esses elementos
    ry=uy - uy.reshape(N,1)   # Constrói a matriz de diferenças entre as componentes y das posições
    trimy=greater(ry,L/2)     # Identifica as diferenças > L/2
    ry=ry-L*trimy                  # Subtrai L dos elementos correspondentes
    trimy=less(ry,-L/2)       # Identifica as diferenças < -L/2
    ry=ry+L*trimy                  # Soma L a esses elementos
    # Constrói a matriz do quadrado das distâncias relativas, para o cálculo das forças
    r2=rx*rx+ry*ry+eye(N,N)   # Soma e subtrai a matriz identidade (evita 1/0 na diagonal)
    rm2=1./r2-eye(N,N)        # Constrói a matriz do inverso do quadrado das distâncias relativas
    trim=greater(rm2,0.160)   # Corta as interações em 2.5*sigma, i.e. 1/r**2 < 0.16 é tratado como 0.
    # Determina as forças
    rm2=rm2*trim              # Implementa na matriz de 1/r^2 o corte nas interações 
    rm8=rm2**4                # Constrói a matriz de 1/r^8
    rm14=rm2**7               # Constrói a matriz de 1/r^12
    fff=24.*(2.*rm14-rm8)     
    ffx=rx*fff                # Matriz das componentes x das forças
    ffy=ry*fff                # Matriz das componentes y das forças
    fx = sum(ffx,axis=0)      # Vetor das componentes x da força total sobre cada partícula
    fy = sum(ffy,axis=0)      # Vetor das componentes y da força total sobre cada partícula
    return fx,fy

def pe(ux,uy,N,L):       # Calcula a energia potencial
    # Leva em conta as condições de contorno periódicas
    rx=ux - ux.reshape(N,1)   # Constrói a matriz de diferenças entre as componentes x das posições
    trimx=greater(rx,L/2)     # Identifica as diferenças > L/2
    rx=rx-L*trimx                  # Subtrai L dos elementos correspondentes
    trimx=less(rx,-L/2)       # Identifica as diferenças < -L/2
    rx=rx+L*trimx                  # Soma L a esses elementos
    ry=uy - uy.reshape(N,1)   # Constrói a matriz de diferenças entre as componentes y das posições
    trimy=greater(ry,L/2)     # Identifica as diferenças > L/2
    ry=ry-L*trimy                  # Subtrai L dos elementos correspondentes
    trimy=less(ry,-L/2)       # Identifica as diferenças < -L/2
    ry=ry+L*trimy                  # Soma L a esses elementos
    r2=rx*rx+ry*ry+eye(N,N)   # Soma e subtrai a matriz identidade (evita 1/0 na diagonal)
    rm2=1./r2-eye(N,N)        # Constrói a matriz do inverso do quadrado das distâncias relativas
    trim=greater(rm2,0.160)   # Corta as interações em 2.5*sigma, i.e. 1/r**2 < 0.16 é tratado como 0.
    # Determina a energia potencial
    rm2=rm2*trim              # Implementa na matriz de 1/r^2 o corte nas interações 
    rm6=rm2**3                # Constrói a matriz de 1/r^6
    rm12=rm2**6               # Constrói a matriz de 1/r^12
    r6_12=4.*(rm12-rm6)
    pot=0.5*sum(r6_12)        # Calcula a energia potencial total
    return pot
    
# Inicializa os parâmetros
N=64                     # Número de partículas
L=15.                    # Lado da caixa
sp=1.6                   # Espaçamento inicial entre partículas
passos=200*N             # Número de passos de passos de Monte Carlo
delta=.6                 # Lado da caixa de vizinhança para tentar movimentos
T=0.6                    # Temperatura
# Inicialização da simulação
print ('T  ',T, 'Passos ', passos,'Delta   ',delta) 
Lsp=int(L/sp)            
Nsqrt=ceil(sqrt(N))
Atoms=[]                 # Cria uma lista para os objetos que representarão os átomos
ux=zeros(N)              # Vetor que armazena as componentes x das posições das partículas
uy=zeros(N)              # Vetor que armazena as componentes x das posições das partículas
kskip=10*N               # Define que a visualização será atualizada a cada 10 passos por átomo
# Cria o esqueleto da visualização
r=0.5                    
scene = display(title="MC Lennard-Jones 2D", x=30, y=0,
                foreground=color.black,background=color.yellow)    # Tela da animação
square = curve(pos=[(-L/2-r,-L/2-r),(-L/2-r,L/2+r),
                    (L/2+r,L/2+r),(L/2+r,-L/2-r),(-L/2-r,-L/2-r)]) # Define a caixa que contém os átomos
g2=gdisplay(x=500,width=500,title="Energia potencial",
            foreground=color.black,background=color.white)         # Tela para exibir a curva da energia potencial
potential=gcurve(color=color.blue)                                 # que será traçada em azul
# Define as condições iniciais (arranjo hexagonal, com distâncias escolhidas para que as interações sejam fracas)
for i in arange(N):  
    x,y=divmod(i,Nsqrt)  
    ux[i]=-L/4+sp*(x + y/2.)
    if ux[i] > L/2-r: ux[i]=-L+2*r+ux[i]
    uy[i]=-L/4+sp*y*0.866
    c=color.blue
    Atoms=Atoms+[sphere(pos=(ux[i],uy[i],0), radius=r, color=c)]   # Cria uma esfera azul para representar cada átomo
# Calcula a energia potencial inicial
PE0=pe(ux,uy,N,L)    
print("EP(0)=",PE0)
PE=PE0
# Inicia a dinâmica
k=0       # Contador dos passos aceitos
virial=0. # Variável que armazena o virial
PEave=0.  # Variável que armazena a energia potencial média
count=0   # Contador dos passos para as médias
for t in arange(passos): 
    rate(100000) # Controla o ritmo de exibição da simulação
    j=random.randint(N) # Sorteia o átomo que tentaremos mover
    dx=delta*(random.random()-0.5) # Componente x do deslocamento
    xo=ux[j]                          # Componente x da posição atual
    xn=xo+dx                          # Componente x da posição proposta
    # Leva em conta condições de contorno periódicas
    if xn>L/2: xn=xn-L
    if xn<-L/2: xn=xn+L
    ux[j]=xn
    dy=delta*(random.random()-0.5) # Componente y do deslocamento
    yo=uy[j]                          # Componente y da posição atual
    yn=yo+dy                          # Componente y da posição proposta
    # Leva em conta condições de contorno periódicas
    if yn>L/2: yn=yn-L
    if yn<-L/2: yn=yn+L
    uy[j]=yn    
    # Calcula a variação da energia potencial
    PEn=pe(ux,uy,N,L)
    DE=PEn-PE
    w=exp(-DE/T)                   # Probabilidade de transição
    u=random.random()              # Número aleatório entre 0 e 1
    if u < w:                      # O movimento é aceito
        PE=PEn
        k=k+1
    else:                          # O movimento é descartado
        ux[j]=xo
        uy[j]=yo       
    # Atualiza a animação e a curva de energia potencial a cada kskip passos
    if mod(t,kskip) ==0: 
        potential.plot(pos=(t,PE)) 
        for i in arange(N):   
            Atoms[i].pos=(ux[i],uy[i],0.)
        if t > passos/3:            # Toma médias após a equilibração
            count=count+1
            ax,ay=forces(ux,uy,N,L)
            virial=virial+0.5*(np.dot(ax,ux)+np.dot(ay,uy))  # Acumula o virial (em 2D)
            PEave=PEave+PE                                   # Acumula a energia potencial
virial=virial/count   # Média do virial
PEave=PEave/count     # Média da energia potencial
print('Passos aceitos',k,'Taxa',k/passos)
print('Temperatura=KE/N=     ',T,'    Energia pot. média=  ',PEave/N )
print('Virial=      ',virial)
Pi=N*T/L**2
print('Pressão do gás ideal =    ', Pi)
print('Pressão real= ', Pi+virial/L**2)