# ---
# 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
# ---

# +
import pyanitools as pyt
import pyNeuroChem as pyc
import hdnntools as hdt
import numpy as np

from ase_interface import ANIENS
from ase_interface import ensemblemolecule
from ase.atoms import Atoms

import os

import matplotlib.pyplot as plt

# +
h5test = '/nh/nest/u/jsmith/Research/tin_research/metal_data/U12.0_test.h5'

ens_path = '/nh/nest/u/jsmith/Research/tin_research/metal_data/models/gutz_lc_models/'

cns = '/train0/rX-2.8R_32-2.0A_a8-8.params'
sae = '/train0/sae.dat'
nnf = '/train'
Nn = 8

# +
model_paths = sorted(os.listdir(ens_path))

comp_all = {}

for model_path in model_paths:
    wkpath = ens_path+model_path
    print(wkpath)
    aens = ensemblemolecule(wkpath + cns,wkpath + sae,wkpath + nnf,Nn)
    
    comp = {'E':[[],[]],
            'F':[[],[]],}
    
    adl = pyt.anidataloader(h5test)
    
    for data in adl:
        print(data['path'])
        X = data['coordinates']
        S = data['species']
        E = data['energies']
        F = data['forces']
        C = data['cell']
        
        for x,e,f,c in zip(X,E,F,C):
            comp['E'][0].append(e)
            comp['F'][0].append(f.flatten())
            
            aens.set_pbc(True,True,True)
            celi = (np.linalg.inv(C)).astype(np.float32)
            aens.set_cell((C).astype(np.float32), celi)

            aens.set_molecule(X=x,S=S)
            ea,fa,es,fs = aens.compute_mean_props()

            comp['E'][1].append(ea)
            comp['F'][1].append(fa.flatten())
            
    comp['E'][0] = np.array(comp['E'][0])
    comp['E'][1] = np.array(comp['E'][1])
    comp['F'][0] = np.concatenate(comp['F'][0])
    comp['F'][1] = np.concatenate(comp['F'][1])
    comp_all[model_path] = comp

# +
plot_data = {'Emae':[],
             'Fmae':[],
             'perc':[],}

for key in comp_all.keys():
    percent = float(key.split('_')[1][:-1])
    Emae = hdt.calculaterootmeansqrerror(comp_all[key]['E'][0],comp_all[key]['E'][1])
    Fmae = hdt.calculaterootmeansqrerror(comp_all[key]['F'][0],comp_all[key]['F'][1])
    plot_data['perc'].append(percent)
    plot_data['Emae'].append(Emae)
    plot_data['Fmae'].append(Fmae)
    print(percent,Emae,Fmae)

# +
import matplotlib.ticker as ticker

fig, axs = plt.subplots(2, 1, figsize=(6, 4), dpi=200, sharey=False)

axs[0].plot(plot_data['perc'], plot_data['Emae'], '-o')
axs[1].plot(plot_data['perc'], plot_data['Fmae'], '-o')

# remove the x and y ticks
for ax in axs:
    #ax.set_yscale('log')
    ax.set_xscale('log')
    ax.xaxis.set_major_formatter(ticker.FormatStrFormatter('%0.0f'))
    #ax.set_yticks([])

axs[0].set_title("Energy", y=0.8)
axs[1].set_title("Force", y=0.8)

axs[1].set_xlabel("Percent of data")

axs[0].set_ylabel("Energy RMSE")
axs[1].set_ylabel("Force RMSE")
#print(fig.get_yticks())

#axs[0].yaxis.set_ticks_position('none') 
#axs[1].yaxis.set_ticks_position('none') 
#axs[0].yaxis.set_ticks(np.arange(min(plot_data['Emae']), max(plot_data['Emae']), 0.02))
#axs[0].yaxis.set_major_formatter(ticker.FormatStrFormatter('%0.1e'))
#axs[0].xaxis.set_major_formatter(ticker.FormatStrFormatter('%0.0f'))

plt.show()
# -



StarcoderdataJupyterScriptsDedupFiltered