# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Fechting the dataset 

# +
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from SCFInitialGuess.utilities.analysis import matrix_error

figure_save_path = "/home/jo/Repos/MastersThesis/InitialGuess/Butadien/figures/"

mpl.style.use("seaborn")

# +
S = np.load("../data/S.npy")
P = np.load("../data/P.npy")
F = np.load("../data/F.npy")

index = np.load("../data/index.npy")

molecules = np.load("../data/molecules.npy")

dim = 26
p = list(P.reshape(-1, dim, dim))
s = list(S.reshape(-1, dim, dim))
f = list(F.reshape(-1, dim, dim))

s_test, p_test, f_test = s[150:], p[150:], f[150:]
# -

# # Checking if components match 

# +
from pyscf.scf import hf

flag_s = False

for mol, s_i, p_i, f_i in zip(molecules, s, p, f):
    print(np.all(s_i == hf.get_ovlp(mol.get_pyscf_molecule())))
    #TODO!!!!
    
# -

# # Analyzing

# +

mol = np.load("../data/molecules.npy")
indexed = zip(index, mol)
indexed = sorted(indexed, key=lambda x: x[0])
mol = [m for (i, m) in indexed]


# +
from utilities.create_dataset import fetch_descriptor_values_over_time


distances, angles, dihedrals = fetch_descriptor_values_over_time(mol)


dihedrals[dihedrals < 2] += 2 * np.pi 
dihedrals[dihedrals > 5] -= 2 * np.pi 


# +
labels_distances = [
    "C1-C2",
    "C2-C3",
    "C3-C4",
    "C1-H1",
    "C1-H2",
    "C2-H3",
    "C3-H4",
    "C4-H5",
    "C4-H6"
]

labels_angles = [
    "C1-C2-C3",
    "C2-C3-C4",
    "H1-C1-C2",
    "H2-C1-C2",
    "H3-C2-C1",
    "H4-C3-C4",
    "H5-C4-C3",
    "H6-C4-C3"
]

labels_dihedrals = [
    "C1-C2-C3-C4",
    "H1-C1-C2-C3",
    "H2-C1-C2-C3",
    "H3-C2-C3-C4",
    "H3-C3-C2-C1",
    "H5-C4-C2-C2",
    "H6-C4-C3-C2"
]

# +
plt.figure()
for i in range(distances.shape[1]):
    plt.plot(distances[:, i], label=labels_distances[i])
plt.legend()
plt.xlabel("Timesteps / 1")
plt.ylabel(r"distances / $\AA$")
plt.savefig(figure_save_path + "DatasetButadienDistances.png")
    
plt.figure()
for i in range(angles.shape[1]):
    plt.plot(angles[:, i], label=labels_angles[i])
plt.ylabel(r"Angles / degree")
plt.xlabel("Timesteps / 1")
plt.legend()
plt.savefig(figure_save_path + "DatasetButadienAngles.png")
    
plt.figure()
for i in range(dihedrals.shape[1]):
    plt.plot(dihedrals[:, i], label=labels_dihedrals[i])
plt.ylabel("Dihedral angles / degree")
plt.xlabel("Timesteps / 1")
plt.legend()
plt.savefig(figure_save_path + "DatasetButadienDihedrals.png")

plt.show()
# -

# # 400 Analysis

# +

mol400 = np.load("../data/molecules400.npy")
index400 = np.load("../data/index400.npy")
indexed = zip(index400, mol400)
indexed = sorted(indexed, key=lambda x: x[0])
# -

list(index400).index(1)

# +


from utilities.create_dataset import fetch_descriptor_values_over_time


distances, angles, dihedrals = fetch_descriptor_values_over_time([m for (i, m) in indexed])


dihedrals[dihedrals < 2] += 2 * np.pi 
dihedrals[dihedrals > 5] -= 2 * np.pi 

plt.figure()
for i in range(distances.shape[1]):
    plt.plot(distances[:, i], label=labels_distances[i])
plt.legend()
plt.xlabel("Timesteps / 1")
plt.ylabel(r"distances / $\AA$")
plt.savefig(figure_save_path + "DatasetButadienDistances.png")
    
plt.figure()
for i in range(angles.shape[1]):
    plt.plot(angles[:, i], label=labels_angles[i])
plt.ylabel(r"Angles / degree")
plt.xlabel("Timesteps / 1")
plt.legend()
plt.savefig(figure_save_path + "DatasetButadienAngles.png")
    
plt.figure()
for i in range(dihedrals.shape[1]):
    plt.plot(dihedrals[:, i], label=labels_dihedrals[i])
plt.ylabel("Dihedral angles / degree")
plt.xlabel("Timesteps / 1")
plt.legend()
plt.savefig(figure_save_path + "DatasetButadienDihedrals.png")

plt.show()
# -

# ## HF Groudstate energy:

# +
from SCFInitialGuess.utilities.analysis import measure_hf_energy

#energy = list(measure_hf_energy(p, list(map(lambda x: x.get_pyscf_molecule(), molecules))))
energy = list(measure_hf_energy(p, molecules))
# -

plt.figure()
plt.plot(energy)
plt.xlabel("Timesteps / 1")
plt.ylabel(r"HF Energy / Hartree")
#plt.savefig(figure_save_path + "DatasetButadienDistances.png")

# # Check for symmetry

def symmetry_error(M):
    
    error = np.zeros(M[0].shape)
    for m in M:
        error += np.abs(m - np.transpose(m))
        
    return error / len(M)


# ## S -Matrix

error = symmetry_error(s)
print(np.mean(error))
matrix_error(error)
plt.show()

# ## P-Matrix

error = symmetry_error(p)
print(np.mean(error))
matrix_error(error)
plt.show()

# # Verrauschte p matix

p_noise = list((P + np.random.randn(*P.shape) * 1e-3).reshape(-1, dim, dim))
error = symmetry_error(p_noise)
print(np.mean(error))
matrix_error(error)
plt.show()

# +
errors = []
for p_i in p_test:
    errors.append(p_i - p_i.T)

print(np.mean(errors), np.std(errors))
# -

# # F Matrix 

error = symmetry_error(f)
print(np.mean(error))
matrix_error(error)
plt.show()


# # Idempotency
#
# $$PSP^T-2P=0$$,
# Faktor 2 wegen Definition/Besetzung

# +
def idempotency_error(M):
    
    error = np.zeros(M[0].shape)
    for m in M:
        error += np.abs(m - np.dot(m, m))
        
    return error / len(M)

def weighted_idempotency_eror(M, W):
    
    error = np.zeros(M[0].shape)
    for m, w in zip(M, W):
        error += np.abs(2 * m - np.dot(m, np.dot(w, m)))
        
    return error / len(M)
    


# +
def mc_wheeny_purification(p,s):
    p = p.reshape(dim, dim)
    s = s.reshape(dim, dim)
    return (3 * np.dot(np.dot(p, s), p) - np.dot(np.dot(np.dot(np.dot(p, s), p), s), p)) / 2

def mc_wheeny_purification_2(p, s=None):
    p = p.reshape(dim, dim)
    return 3 * np.dot(p, p) - 2 * np.dot(np.dot(p, p), p)

def idemp_error(p, s):
    p = p.reshape(dim, dim)
    s = s.reshape(dim, dim)
    return np.mean((np.abs(p @ s @ p.T - 2 * p)))


    
    
                   
# -

# ## P matrix

error = weighted_idempotency_eror(p, s)
print(np.mean(error))
matrix_error(error)
plt.show()

# ## P matrix w/ noise

error = weighted_idempotency_eror(p_noise, s)
print(np.mean(error))
matrix_error(error)
plt.show()

# +
from functools import reduce

errors = []
for s_i, p_i in zip(s_test, p_test):
    errors.append(2 * p_i - reduce(np.dot, (p_i, s_i, p_i.T)))

print(np.mean(errors), np.std(errors))


# -

# # Check occupancy
# $$tr(PS) = N_{electrons}$$

def population(S, P):
    for i, (s,p) in enumerate(zip(S, P)):
        yield np.trace(np.dot(p, s))



# ## Normal

print(np.mean(list(population(s, p))))

# ## Verrauscht

print(np.mean(list(population(s, p_noise))))



StarcoderdataJupyterScriptsDedupFiltered