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=True)
parser.add_argument("--set_param_range_and_steps", '-p', dest='param_range', nargs='+', type=float)
parser.add_argument("--model_name", "-m", dest='model_name' ,type=str, required=True)
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
    if param < -1.0:
        hamming_weight = None
    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()
    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 = 'tfim'
    number_spins = number_spins#[8, 10, 12, 14]#, 16, 18]
    periodic = False
    spin_inversion = None
    params = np.linspace(param_range[0], param_range[1], int(param_range[2]))

    sub_dim = 1
    first_trace_spin = 1
    MIs = []
    Entropies = []

    print('Number Spins: ', number_spins)
    print('Sub Dimension is ', sub_dim)
    for param in params:
        if param < -1.0 or model_name == 'tfim':
            hamming_weight = None
        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)
        rho_half = reduced_dm(number_spins // 2, number_spins, hamming_weight, gs, basis_states)
        Entropies.append(compute_entropy(rho_half))
        del model

    data = {
        'Model': model_name,
        'Number_spins': number_spins,
        'Periodic': periodic,
        'Hamming_weight': hamming_weight,
        'Sub_dim': sub_dim,
        'Delta or h over J': params,
        'Entropies': Entropies,
        'Mutual Informations': MIs
    }

    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: number spins = ' + str(number_spins) + 'periodic = ' + str(periodic))
    plt.plot(params, Entropies, label='Half-chain Entropy')
    plt.plot(params, MIs, label='Mutual Information')
    plt.grid(True)
    plt.legend()
    plt.savefig('output/' + filename + '.png')

if __name__ == "__main__":
    print('hi')
    model_name = args.model_name
    param_range = args.param_range
    print(param_range)
    number_spins = args.number_spins
    MI_vs_Entropy(model_name, number_spins, param_range)