SPEA/gcMC_square.py at main · eisenforschung/SPEA · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
from GrandCanonicalMC import gcMC

import numpy as np
import itertools

from numba import jit

def set_e_tensor(e_scale=1):
    e_tensor = np.zeros((2,2,2,2,2))

    basis = {}
    basis[(0,0,0,0)] = 4
    basis[(1,0,0,0)] = -1
    basis[(1,1,0,0)] = -2
    basis[(1,1,1,0)] = -3
    basis[(1,1,1,1)] = -4

    for key, value in basis.items():
        for p in set(itertools.permutations(key)):
        #print (set(itertools.permutations(key)))
            e_tensor[tuple([1] + list(p))] = float(value)

    basis[(0,0,0,0)] = -4
    basis[(1,0,0,0)] = -3
    basis[(1,1,0,0)] = -2
    basis[(1,1,1,0)] = -1
    basis[(1,1,1,1)] =  0

    for key, value in basis.items():
        for p in set(itertools.permutations(key)):
            e_tensor[tuple([0] + list(p))] = float(value)

    return e_scale * e_tensor

e_tensor = set_e_tensor(0.2)

@jit(nopython=True)
def get_energy_single(lat, i_site, ind_vec):
    occ = lat[i_site]
    i0, i1, i2, i3 = ind_vec[i_site]
    o0, o1, o2, o3 = lat[i0], lat[i1], lat[i2], lat[i3]

    return e_tensor[occ, o0, o1, o2, o3]

@jit(nopython=True)
def get_energy(lat, i_site, ind_vec):
    i0, i1, i2, i3 = ind_vec[i_site]

    energy = 0
    for i in [i_site, i0, i1, i2, i3]:
        energy += get_energy_single(lat, i, ind_vec)
    return energy

class Lattice(gcMC):
    def __init__(self, mesh_x, mesh_y, e_scale=1):
        self.mesh_x = mesh_x
        self.mesh_y = mesh_y
        self.mesh = mesh_x * mesh_y

        self.i_nbrs = 4   # square
        self.get_energy = get_energy
        self.get_energy_single = get_energy_single

        self.build_lattice()

    def build_lattice(self):
        r_x = np.arange(self.mesh_x)
        r_y = np.arange(self.mesh_y)

        R_X, R_Y = np.meshgrid(r_x, r_y)

        self.R_X = R_X.flatten()
        self.R_Y = R_Y.flatten()

        lattice_vec = np.arange(self.mesh)
        lattice = lattice_vec.reshape((self.mesh_x, self.mesh_y))

        ind_vec = np.array([0] * self.mesh * self.i_nbrs).reshape(self.mesh, self.i_nbrs)
        ind_vec[:, 0] = np.roll(lattice, -1, 0).flatten()
        ind_vec[:, 1] = np.roll(lattice,  1, 0).flatten()
        ind_vec[:, 2] = np.roll(lattice,  -1, 1).flatten()
        ind_vec[:, 3] = np.roll(lattice,  1, 1).flatten()

        self.ind_vec = ind_vec

        self.occ_vec = np.zeros_like(lattice_vec).astype(int)