IsingModel/main.py - master - ilmerkul/Ising_Model

IsingModel/main.py
150 строк5 KB
Ilya Merkulov
clear project
09 ноя 2024, 16:06
09 ноя 2024, 16:066dca45c
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 = 16
seedRandom = 101


def to_array(model, parameters):
    all_configurations = hi.all_states()

    logpsi = model.apply(parameters, all_configurations)
    psi = jnp.exp(logpsi)
    psi = psi / jnp.linalg.norm(psi)
    return psi


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.conv4 = nn.Conv(features=2, kernel_size=(2, 2), use_bias=True)
        self.linear1 = nn.Dense(2 * N)
        self.linear3 = 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.conv4(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.linear3(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=(4, 2, 1), param_dtype=float,
                                   activation=nn.activation.relu, complex_output=False)

    @nn.compact
    def __call__(self, x):
        x = self.gcnn(x)
        return x


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

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

    H = sum([h * sigmax(hi, i) for i in range(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]
    Psi_gs = eig_vecs[0]
    return E_gs, Psi_gs


jax.random.key(seedRandom)

#grid = nk.graph.Grid(extent=[N, N])
grid = nk.graph.Chain(length=N)
# grid = nk.graph.Hypercube(length=N, n_dim=2)

for h in (1, 2.0, 3.0, 5.0):
    hi, H = init_gamilton(-1, -h)
    hiRBM, HRBM = init_gamilton(-1, -h)

    E_gs, Psi_gs = get_correct_answer(H)
    print(E_gs/N)

    model = nk.models.GCNN(symmetries=grid.automorphisms(), layers=3, features=(4, 2, 1), param_dtype=float,
                           activation=nn.activation.relu, complex_output=False)
    modelRBM = nk.models.RBM(alpha=2)

    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)

    samplerRBM = nk.sampler.MetropolisLocal(hiRBM)

    vstateRBM = nk.vqs.MCState(samplerRBM, modelRBM, n_samples=512)

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

    gsRBM = nk.driver.VMC(H, optimizerRBM, variational_state=vstateRBM, preconditioner=nk.optimizer.SR(diag_shift=0.1))

    logRBM = nk.logging.RuntimeLog()
    gsRBM.run(n_iter=100, out=logRBM)

    dataEnergy = log.data["Energy"]
    dataEnergyRBM = logRBM.data["Energy"]
    plt.errorbar(dataEnergy.iters, dataEnergy.Mean / N, yerr=dataEnergy.Sigma / N, label="GCNN")
    plt.errorbar(dataEnergyRBM.iters, dataEnergyRBM.Mean / N, yerr=dataEnergyRBM.Sigma / N, label="RBM")
    plt.axhline(y=E_gs/N, linestyle='--', color="black")
    plt.xlabel('Iterations')
    plt.ylabel('E/N')
    plt.legend()
    plt.show()
    #vstate.n_samples = 8000
    #vstateRBM.n_samples = 8000

    #symm_energy = vstate.expect(H)
    #symm_energyRBM = vstateRBM.expect(HRBM)
    #print('Energy:', symm_energy, E_gs, "\nRBM Energy:", symm_energyRBM)

    # corrs = []
    # for x in range(1, N):
    #    corr = sum([sum([sigmaz(hi, i * N + j) * sigmaz(hi, (i + x) * N + j + x) for j in range(N - x)]) for i in
    #                range(N - x)]) / (N - x) / (N - x)
    #    corr_mean = vstate.expect(corr)
    #    corrs.append(corr_mean.mean)
    # print(corrs)
    # parameters = vstateRBM.parameters
    # params = dict()
    # params['params'] = parameters
    # parameters = params
    # psi = to_array(modelRBM, parameters)
    # print(sum([psi.conj() @ sigmax(hi, i).to_sparse() @ psi for i in range(N)]) / N)
    # print(sum([Psi_gs @ sigmax(hi, i).to_sparse() @ Psi_gs for i in range(N)]) / N)