IsingModel/test.py - master - ilmerkul/Ising_Model

IsingModel/test.py
122 строки3 KB
Ilya Merkulov
my course work at the university
06 окт 2024, 15:37
06 окт 2024, 15:37f37ec12
import time

import netket as nk
from netket.operator.spin import sigmax, sigmaz
import flax.linen as nn
import jax
import jax.numpy as jnp
from scipy.sparse.linalg import eigsh
from matplotlib import pyplot as plt

N = 7
seedRandom = 101


class CNN(nn.Module):
    def setup(self):
        self.conv1 = nn.Conv(features=8, kernel_size=(3, 3), use_bias=True)
        self.conv2 = nn.Conv(features=4, kernel_size=(3, 3), use_bias=True)
        self.conv3 = nn.Conv(features=2, kernel_size=(2, 2), use_bias=True)
        self.linear1 = nn.Dense(8)
        self.linear2 = nn.Dense(1)

    @nn.compact
    def __call__(self, x):
        x = x.reshape(-1, N, N, 1)

        x = self.conv1(x)
        x = nn.activation.relu(x)
        x = self.conv2(x)
        x = nn.activation.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = self.conv3(x)
        x = nn.activation.relu(x)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))
        x = self.linear1(x)
        x = nn.activation.relu(x)
        x = self.linear2(x)
        x = jnp.sum(x, axis=-1)

        return x


class GCNN(nn.Module):
    def setup(self):
        self.gcnn = nk.models.GCNN(symmetries=grid.automorphisms(), layers=3, features=(8, 4, 2), param_dtype=float,
                                   activation=nn.activation.tanh)
        self.dense1 = nn.Dense(8)
        self.dense2 = nn.Dense(1)

    @nn.compact
    def __call__(self, x):
        print(x.shape)

        x = self.gcnn(x)
        print(x.shape)
        x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
        print(x.shape)
        x = x.reshape((x.shape[0], -1))
        print(x.shape)
        x = self.dense1(x)
        print(x.shape)
        x = nn.activation.tanh(x)
        x = self.dense2(x)
        print(x.shape)
        print()
        x = jnp.sum(x, axis=-1)

        return x


def init_gamilton(J, h):
    hi = nk.hilbert.Spin(s=1 / 2, N=N * N)

    hi.random_state(jax.random.key(0), 3)

    H = sum([h * sigmax(hi, i) for i in range(N * N)])

    H += sum([J * sigmaz(hi, i) * sigmaz(hi, j) for (i, j) in grid.edges()])

    return hi, H


def get_correct_answer(H):
    sp_h = H.to_sparse()
    eig_vals, eig_vecs = eigsh(sp_h, k=2, which="SA")
    E_gs = eig_vals[0]
    return E_gs


jax.random.key(seedRandom)

grid = nk.graph.Grid(extent=[N, N])

for h in (0.3, 0.8, 1.0, 1.2):

    hi, H = init_gamilton(-1, -h)

    #E_gs = get_correct_answer(H)

    model = CNN()

    sampler = nk.sampler.MetropolisLocal(hi)

    vstate = nk.vqs.MCState(sampler, model, n_samples=1024)

    optimizer = nk.optimizer.Sgd(learning_rate=0.05)

    gs = nk.driver.VMC(H, optimizer, variational_state=vstate, preconditioner=nk.optimizer.SR(diag_shift=0.1))

    log = nk.logging.RuntimeLog()
    gs.run(n_iter=100, out=log)

    dataEnergy = log.data["Energy"]
    plt.errorbar(dataEnergy.iters, dataEnergy.Mean/(N*N), yerr=dataEnergy.Sigma/(N*N))
    plt.xlabel('Iterations')
    plt.ylabel('E/N')
    plt.show()

    symm_energy = vstate.expect(H)
    #error = abs((symm_energy.mean - E_gs))
    #print("Optimized energy and relative error: ", symm_energy, error)