mutual_information.py

from ast import parse
import numpy as np
import yaml
import time
import matplotlib.pyplot as plt

from create_model import spin_model
from entropy_new import compute_MI, compute_entropy, reduced_dm

import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--set_number_spins", "-d", dest='number_spins',type=int, required=False)
parser.add_argument("--set_param_range_and_steps", '-pr', dest='param_range', nargs='+', type=float)
parser.add_argument("--model_name", "-m", dest='model_name' ,type=str, required=True)
parser.add_argument("--set_spin_range", "-sr", dest='spin_range', nargs='+', type=int)
parser.add_argument("--set_param", "-p", dest='param',type=float, required=False)
args = parser.parse_args()

def test_mutual_information():
    # change return in compute_MI to run this function [S_A, S_B, S_AB] instead of S_A + S_B - S_AB
    model_name = 'xxz'
    number_spins = 10
    periodic = False
    spin_inversion = None
    param = 1.0
    hamming_weight = number_spins // 2
    model = spin_model(model_name=model_name, number_spins=number_spins, periodic=periodic, 
                            param=param, hamming_weight=hamming_weight ,spin_inversion=spin_inversion)
    model.compute_ew_and_ev()
    print('EIGENSTATE at h/J = ', param, 'is: ', model.eigenstates[:,0])
    sub_dim = 4
    first_trace_spin = 4
    basis_states = model.basis.states
    gs = model.eigenstates[:,0]
    print('Number Spins: ', model.basis.number_spins)
    print('States', model.basis.states)
    print('Sub Dimension is ', sub_dim)
    entropies = compute_MI(sub_dim, number_spins, hamming_weight, gs, basis_states, spin_inversion=None, first_trace_spin=first_trace_spin)
    MI = entropies[0] + entropies[1] - entropies[2]
    print(MI)
    print(entropies)
    print('--------------------Test-----------------------')
    rhos4 = reduced_dm(4, number_spins, hamming_weight, gs, basis_states)
    rhos6 = reduced_dm(6, number_spins, hamming_weight, gs, basis_states)
    print(compute_entropy(rhos4), compute_entropy(rhos6))

def MI_vs_Entropy(model_name, number_spins, param_range):

    model_name = model_name
    number_spins = number_spins
    periodic = False
    params = np.linspace(param_range[0], param_range[1], int(param_range[2]))

    sub_dim = 4
    first_trace_spin = 4
    MIs = []
    Entropies = []
    Entropies4 = []
    hamming_weight = number_spins // 2
    print('Number Spins: ', number_spins)
    print('Sub Dimension is ', sub_dim)
    for param in params:
        model = spin_model(model_name=model_name, number_spins=number_spins, periodic=periodic, 
                            param=param, hamming_weight=hamming_weight)
        model.compute_ew_and_ev()
        basis_states = model.basis.states
        gs = model.eigenstates[:,0]
        half_dim = number_spins // 2
        rho_half = reduced_dm(half_dim, number_spins, hamming_weight, gs, basis_states)
        Entropies.append(compute_entropy(rho_half))
        rho_4spin = reduced_dm(4, number_spins, hamming_weight, gs, basis_states)
        Entropies4.append(compute_entropy(rho_4spin))
        MI = compute_MI(sub_dim, number_spins, hamming_weight, gs, basis_states, spin_inversion=None, first_trace_spin=first_trace_spin)
        MIs.append(MI)
        del model
        print(Entropies)
    data = {
        'Model': model_name,
        'Number_spins': number_spins,
        'Periodic': periodic,
        'Hamming_weight': hamming_weight,
        'Sub_dim': sub_dim,
        'Delta or h over J': params.tolist(),
        'Entropies half chain': np.array(Entropies).tolist(),
        'Entropies 4 spin': np.array(Entropies4).tolist(),
        'Mutual Informations': np.array(MIs).tolist()
    }

    timestr = time.strftime("%Y%m%d-%H%M%S")
    filename = 'MI_fixed_dim' + timestr
    with open('output/' + filename + '.yaml', 'w') as outfile:
        yaml.dump(data, outfile, default_flow_style=False)
    print(filename)

    plt.figure(figsize=(12,12))
    plt.title('MI: ' + model_name + ' ' + str(number_spins) + ' ' + str(periodic))
    plt.plot(params, Entropies, label='Half-chain Entropy')
    plt.plot(params, Entropies4, label='4 spin entropy')
    plt.plot(params, MIs, label='Mutual Information')
    plt.grid(True)
    plt.legend()
    plt.savefig('output/' + filename + '.png')

def MutInf(model_name, param, spin_range):

    model_name = model_name
    periodic = False
    spin_inversion = None
    sub_dim = 4
    first_trace_spin = 4
    MIs = []
    Entropies = []
    Entropies_div_ns = []
    number_spins_range = [2*i for i in range(spin_range[0]//2, spin_range[1]//2 + 1)]
    for number_spins in number_spins_range:
        print('Number Spins: ', number_spins)
        print('Sub Dimension is ', sub_dim)
        if param < -1.0 or model_name == 'tfim': 
            hamming_weight = number_spins // 2
        else: 
            hamming_weight = number_spins // 2

        model = spin_model(model_name=model_name, number_spins=number_spins, periodic=periodic, 
                            param=param, hamming_weight=hamming_weight ,spin_inversion=spin_inversion)
        model.compute_ew_and_ev()
        basis_states = model.basis.states
        gs = model.eigenstates[:,0]
        MI = compute_MI(sub_dim, number_spins, hamming_weight, gs, basis_states, spin_inversion=None, first_trace_spin=first_trace_spin)
        MIs.append(MI/sub_dim)
        rho_half = reduced_dm(number_spins // 2, number_spins, hamming_weight, gs, basis_states)
        Entropies.append(compute_entropy(rho_half))
        Entropies_div_ns.append(Entropies[-1]/number_spins)
        del model

    data = {
        'Model': model_name,
        'Number_spins': number_spins_range,
        'Periodic': periodic,
        'Hamming_weight': hamming_weight,
        'Sub_dim': sub_dim,
        'Delta or h over J': param,
        'Entropies': np.array(Entropies).tolist(),
        'Mutual Informations': np.array(MIs).tolist()
    }

    timestr = time.strftime("%Y%m%d-%H%M%S")
    filename = 'MI_fixed_param' + timestr
    with open('output/' + filename + '.yaml', 'w') as outfile:
        yaml.dump(data, outfile, default_flow_style=False)
    print(filename)

    plt.figure(figsize=(12,12))
    plt.title('MI: ' + model_name + ' ' + str(param) + ' ' + str(periodic))
    plt.plot(number_spins_range, Entropies_div_ns, label='Half-chain Entropy / Spin-Chain-Length')
    plt.plot(number_spins_range, MIs, label='Mutual Information / SubDim--Length')
    plt.grid(True)
    plt.xlabel("Number Spins")
    plt.legend()
    plt.savefig('output/' + filename + '.png')


if __name__ == "__main__":
    model_name = args.model_name
    param_range = args.param_range
    number_spins = args.number_spins
    spin_range = args.spin_range
    param = args.param
    MI_vs_Entropy(model_name, number_spins, param_range)
    #MutInf(model_name, param, spin_range)