dice_visualise.py

Go to the documentation of this file.

 
 
import netCDF4 as nd
from netCDF4 import Dataset
import matplotlib.pyplot as plt
import numpy as np
import scipy.signal as sps
import plotly.graph_objects as go
import plotly.figure_factory as ff
from plotly.subplots import make_subplots
import ipywidgets as widgets
from ipywidgets import fixed
 
plt.rcParams["figure.figsize"] = 18,10
 
#----------------------------------------------------
# Energy/Wavefunction Plotting FUNCTION DEFINITIONS -
#----------------------------------------------------
 
def QHOProbPlot(chain_file):
    '''
    @brief
    Wrapper function for QHO_prob_plot()
 
    @details
    Loads in the netCDF dataset and extracts position and energy
    arrays, passes them off to QHO_prob_plot().
 
    @param[in]  chain_file   The main netCDF output file from a Full VQMC run with write_chains=.TRUE.
    '''
    dat = Dataset(chain_file, 'r', format='NETCDF4')
 
    positions = dat.variables['positions'][:]
    energies = dat.variables['energies'][:]
 
    fig = QHO_prob_plot(positions, energies)
 
    fig.show()
 
 
def QHO_prob_plot(positions,energies):
    '''
    @brief
    Fancy Plotly plotter for QHO Markov chains
 
    @param[in]  positions  Position array from Markov chain
    @param[in]  energies   Energy array from Markov chain
    '''
    hist = ff.create_distplot([positions], ['Position Distribution'],
                         bin_size=.2, show_rug=False)
 
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    fig.add_trace(go.Scatter(x=positions, y=energies, mode ='markers', name="Energy",
                marker=dict(
                color=energies,
                colorscale ='Viridis',
                line=dict(
                    width=0
                    ))))
    fig.add_trace(go.Histogram(hist['data'][0]),secondary_y=True,)
    fig.add_trace(go.Scatter(hist['data'][1]),secondary_y=True,)
    fig['layout']['yaxis']['autorange'] = "reversed"
    fig['layout']['yaxis2']['showgrid'] = False
 
    return fig
 
    
def HydProbPlot(chain_file, opacity=0.2):
    '''
    @brief
    Wrapper function for H2_prob_plot()
 
    @details
    Loads in the netCDF dataset and extracts position and energy
    arrays, passes them off to H2_prob_plot().
 
    @param[in]  chain_file   The main netCDF output file from a Full VQMC run with write_chains=.TRUE.
    @param[in]  opacity      (Optional, default=0.2) Opacity of the electron position scatter plot
    '''
    dat = Dataset(chain_file, 'r', format='NETCDF4')
 
    positions = dat.variables['positions'][:, :]
    energies = dat.variables['energies'][:]
    bond_length = dat.bond_length
 
    R1 = [-bond_length/2, 0.0, 0.0]
    R2 = [bond_length/2, 0.0, 0.0]
 
    print(len(positions))
 
    fig = H2_prob_plot(positions, energies, R1, R2, opacity)
 
    fig.show()
 
def H2_prob_plot(positions,energies,R1, R2, opacity):
    """ 
    @brief
    Plots the sampled wavefunction for H2plus/H2
    
    @param[in]  positions  XYZ coordinates of Markov chain
    @param[in]  energies   Energies of the Markov chains
    @param[in]  R1         The XYZ coordinates of the first atom
    @param[in]  R2         The XYZ coordinates of the second atom
    @param[in]  opacity    Opacity of the electron position scatter plot
    @retval     fig        Plotly figure object
    """   
    
    X = positions[:,0]
    Y = positions[:,1]
    Z = positions[:,2]
 
    fig = go.Figure()
    
    #Plots the positions of the Markov chain, with energies corresponding to a certain color
    fig.add_trace(
    go.Scatter3d(
        mode='markers',
        x=X,
        y=Y,
        z=Z,
        opacity=opacity,
        marker=dict(
            color=energies,
            size=3,
            colorscale ='Viridis',
            line=dict(
                width=0
                )
            ),
            name='Markov Chain positions'
        )
    )
    
    #Adds points to show the location of the atoms within the object
    fig.add_trace(
        go.Scatter3d(
        mode='markers',
        x=[R1[0],R2[0]],
        y=[R1[1],R2[1]],
        z=[R1[2],R2[2]],
        opacity=1.0,
        marker=dict(
            color='#AAAAAA',
            size=20,
            line=dict(
                color='Black',
                width=2
                )
            ),
            name='Atoms'
        )
    )
    
    #Resizes the layout
    fig.update_layout(
    autosize=False,
    width=800,
    height=800,
        margin=dict(
        l=5,
        r=5,
        b=5,
        t=100,
        pad=1
    )
    )
    
    return fig
    
 
def QHOEnergyPlot(results_file, sliceby=(None, None)):
    '''
    @brief
    Energy against @f$ \alpha @f$ plotter for the QHO problem.
 
    @details
    Plots the VQMC wavefunction energy at each of an array of values for the
    variational parameter @f$ \alpha @f$, alongside error bars of the standard
    deviation in each energy.
 
    @param[in]  results_file    The main netCDF output file from a QHO run
    @param[in]  sliceby         (Optional) Tuple containing slice intervals for arrays
    '''
    dat = nd.Dataset(results_file, 'r', format='NETCDF4')
 
    alpha_array = dat.variables['Alpha_array'][sliceby[0]:sliceby[1]]
    energies = dat.variables['total_energies'][sliceby[0]:sliceby[1]]
    variances = dat.variables['Uncertainties'][sliceby[0]:sliceby[1]]
 
    stdevs = np.sqrt(variances)
    min_energy = min(energies)
    min_energy_stdev = stdevs[np.argmin(energies)]
    opt_alpha = alpha_array[np.argmin(energies)]
 
    fig, ax1 = plt.subplots(1, 1, figsize=(8, 8))
 
    ax1.set_title(r'Quantum Harmonic Oscillator Energy with Varying $\alpha$')
    ax1.plot(alpha_array, energies, marker='o', ls='--', color='blue')
    ax1.fill_between(alpha_array, energies+stdevs, energies-stdevs, alpha=0.3, color='purple')
    ax1.axvline(0.5, ls='-', color='black', alpha=0.3)
    ax1.set_xlabel(r'$\alpha$')
    ax1.set_ylabel('Energy')
 
    fig.show()
 
    print(f'Optimum alpha: {opt_alpha}')
    print(f'Energy: {min_energy} ± {min_energy_stdev}')
 
 
def QHOWfnPlot(chain_file, show_analytic=False):
    '''
    @brief
    Scatter plotter of the sampled QHO wavefunction.
 
    @details
    Plots the sampled positions of a trial QHO wavefunction against the
    energies at those positions, to illustrate the shape of that wavefunction.
    Chain output requires the @em write_chains parameter to be @em TRUE.
 
    @param[in]  chain_file     A netCDF dataset of Markov chain and corresponding energies of a VQMC run
    @param[in]  show_analytic  Also display the analytical wavefunction
    '''
    dat = nd.Dataset(chain_file, 'r', format='NETCDF4')
 
    positions = dat.variables['positions'][:]
    energies = dat.variables['energies'][:]
    alpha = dat.alpha
 
    plt.figure(figsize=(8, 8))
    plt.title(rf'Sampled QHO Wavefunction at $\alpha$ = {alpha}')
    plt.plot(positions, energies, 'o', color='purple', markersize=4.0, alpha=0.5, label='Sampled Positions')
    plt.xlabel('Sampled Position')
    plt.ylabel('Energy')
 
    if show_analytic:
        pos_array = np.linspace(min(positions), max(positions), 100)
        analytic = np.exp(-alpha * pos_array**2)
        plt.plot(pos_array, analytic, color='red', alpha=0.8, label='Analytic Probability Distribution')
        plt.legend()
 
    plt.show()
 
 
def H2plusEnergyPlot(results_file, sliceby=(None, None)):
    '''
    @brief
    @f$ H_{2}^{+} @f$ H-H bond energy and parameter plotter.
 
    @details
    Plots the H-H bond potential energy curve generated by VQMC at the optimum
    parameters for each bond length. Plots standard deviation of each energy as
    a measure of uncertainty. Also plots variational parameter @f$ c @f$ against
    bond length. Reports equilibrium bond length for the molecule and optimum
    energy.
 
    @param[in]  results_file    The main netCDF output file from a @f$ H_{2}^{+} @f$ run
    @param[in]  sliceby         (Optional) Tuple containing slice intervals for arrays
    '''
    dat = nd.Dataset(results_file, 'r', format='NETCDF4')
 
    bond_lengths = dat.variables['Bondlength_array'][sliceby[0]:sliceby[1]]
    c_opt = dat.variables['Parameter_array'][sliceby[0]:sliceby[1]]
    energies = dat.variables['total_energies'][sliceby[0]:sliceby[1]]
    variances = dat.variables['Uncertainties'][sliceby[0]:sliceby[1]]
 
    stdevs = np.sqrt(variances)
    min_energy = min(energies)
    min_energy_stdev = stdevs[np.argmin(energies)]
    min_bond_length = bond_lengths[np.argmin(energies)]
 
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 8))
 
    ax1.set_title(r'$H_{2}^{+}$ bond energy curve')
    # ax1.plot(bond_lengths, energies, marker='o')
    ax1.plot(bond_lengths, energies, marker='o', ls='--', color='blue')
    ax1.fill_between(bond_lengths, energies+stdevs, energies-stdevs, alpha=0.3, color='purple')
    ax1.axhline(-0.5, ls='-', color='black', alpha=0.3)
    ax1.set_xlabel(r'Bond Length / $a_{0}$')
    ax1.set_ylabel('Energy / Ha')
 
    ax2.set_title(r'Optimised parameter $c$')
    ax2.plot(bond_lengths, c_opt)
    ax2.axhline(1/np.sqrt(2), ls='-', color='black', alpha=0.3)
    ax2.set_xlabel(r'Bond Length / $a_{0}$')
    ax2.set_ylabel(r'$c$')
 
    fig.tight_layout()
    fig.show()
 
    print(f'Equilibrium bond length: {min_bond_length} Bohr')
    print(f'                         = {min_bond_length*0.529177} Angstroms')
    print(f'   at energy: {min_energy} ± {min_energy_stdev} Ha')
 
 
 
def H2EnergyPlot(results_file, sliceby=(None, None)):
    '''
    @brief
    @f$ H_{2} @f$ H-H bond energy and parameter plotter.
 
    @details
    Plots the H-H bond potential energy curve generated by VQMC at the optimum
    parameters for each bond length. Plots standard deviation of each energy as
    a measure of uncertainty. Also plots variational parameters @f$ a @f$ and
    @f$ \beta @f$ against bond length. Reports equilibrium bond length for the
    molecule and optimum energy.
 
    @param[in]  results_file    The main netCDF output file from a @f$ H_{2} @f$ run
    @param[in]  sliceby         (Optional) Tuple containing slice intervals for arrays
    '''
    dat = nd.Dataset(results_file, 'r', format='NETCDF4')
 
    bond_lengths = dat.variables['Bondlength_array'][sliceby[0]:sliceby[1]]
    a_opt = dat.variables['Parameter_a_array'][sliceby[0]:sliceby[1]]
    beta_opt = dat.variables['Parameter_beta_array'][sliceby[0]:sliceby[1]]
    energies = dat.variables['total_energies'][sliceby[0]:sliceby[1]]
    variances = dat.variables['Uncertainties'][sliceby[0]:sliceby[1]]
 
    stdevs = np.sqrt(variances)
    min_energy = min(energies)
    min_energy_stdev = stdevs[np.argmin(energies)]
    min_bond_length = bond_lengths[np.argmin(energies)]
 
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 8))
 
    ax1.set_title(r'$H_{2}$ bond energy curve')
    ax1.errorbar(bond_lengths, energies, yerr=stdevs, marker='o', ls='--', color='blue')
    ax1.fill_between(bond_lengths, energies+stdevs, energies-stdevs, alpha=0.3, color='purple')
    ax1.set_xlabel(r'Bond Length / $a_{0}$')
    ax1.set_ylabel('Energy / Ha')
 
    ax2.set_title(r'Optimised parameters')
    ax2.plot(bond_lengths, a_opt, color='red')
    ax2.set_xlabel(r'Bond Length / $a_{0}$')
    ax2.set_ylabel(r'$a$', color='red')
    ax2.tick_params(axis='y', labelcolor='red')
 
    ax2a = ax2.twinx()
    ax2a.plot(bond_lengths, beta_opt, color='blue')
    ax2a.set_ylabel(r'$\beta$', color='blue')
    ax2a.tick_params(axis='y', labelcolor='blue')
 
    fig.tight_layout()
    fig.show()
 
    print(f'Equilibrium bond length: {min_bond_length} Bohr')
    print(f'                         = {min_bond_length*0.529177} Angstroms')
    print(f'   at energy: {min_energy} ± {min_energy_stdev} Ha')
 
 
#------------------------------------
# Burn / Thin FUNCTION DEFINITIONS -
#------------------------------------
 
 
def view_original_trace(equil_file):
    rootgrp = Dataset(equil_file, "r", format = "NETCDF4")
    chain = np.array(rootgrp['positions'][:])
    if len(chain.shape) == 1:
        plt.plot(range(len(chain)), chain, linewidth=1.0)
        plt.xlabel("Position in chain")
        plt.ylabel("Chain value")
        plt.title("Original Chain")
        plt.show()
        
    else:
        chainx = chain[:,0]
        chainy = chain[:,1]
        chainz = chain[:,2]
        
        fig, (ax1,ax2,ax3) = plt.subplots(3,1)
        ax1.plot(range(len(chainx)), chainx, linewidth=1.0)
        ax2.plot(range(len(chainy)), chainy, linewidth=1.0)
        ax3.plot(range(len(chainz)), chainz, linewidth=1.0)
    
    return
 
def BurnThinUI(equil_file):
    '''
    @brief
    Wrapper function for interacting with get_autocorrelation_time()
 
    @param[in]  equil_file   Filename of MCMC equilibration run
    '''
    widgets.interact_manual(get_autocorrelation_time, filename=fixed(equil_file), 
        burn=(0, int(0.1*len(np.array(Dataset(equil_file, "r", format = "NETCDF4")['positions'][:]))), 10),
        tuning=[(1.0),(1e-1),(1e-2),(1e-3),(1e-4)])
 
 
def get_autocorrelation_time(filename, burn, tuning, plotting=True):
    rootgrp = Dataset(filename, "r", format = "NETCDF4")
    chain = np.array(rootgrp['positions'][:])
    if len(chain.shape) == 1:
        acp = sps.fftconvolve(chain,chain[::-1],mode='full')
        acp = acp[int(acp.size / 2):]
        acp = acp / max(acp)
        acp = acp[burn:]
 
        thin = np.argmax(abs(acp) < tuning)
        if thin == 0:
            if acp[thin] < tuning:
                thin = 1
            else:
                print("No thinning value found - tuning may be too aggressive. Defaulting to no thinning, thin=1.")
                print()
                thin = 1
 
        print("burning parameter = ", burn)
        print("proposed thinning parameter = ", thin)
 
        chain_new = chain[burn::thin]
        print("length of new chain = ", len(chain_new))
 
        if plotting==True:
            plt.plot(range(len(acp)), acp, linewidth=1.0)
            plt.xlabel("Position in chain / Lag time")
            plt.ylabel("ACF")
            plt.title("Autocorrelation function vs lag time")
            plt.show()
 
            plt.plot(range(len(chain_new)), chain_new, linewidth=1.0)
            plt.xlabel("Position in chain")
            plt.ylabel("Chain value")
            plt.title("New Chain")
            plt.show()
            
        return burn, thin
            
    else:
        chainx = chain[:,0]
        chainy = chain[:,1]
        chainz = chain[:,2]
        
        threed_thin = []
        chain_lengths = []
        
        for c in [chainx,chainy,chainz]:
            acp = sps.fftconvolve(c,c[::-1],mode='full')
            acp = acp[int(acp.size / 2):]
            acp = acp / max(acp)
            acp = acp[burn:]
 
            thin = np.argmax(abs(acp) < tuning)
            if thin == 0:
                if acp[thin] < tuning:
                    thin = 1
                else:
                    print("No thinning value found - tuning may be too aggressive. Defaulting to no thinning, thin=1.")
                    print()
                    thin = 1
                    
            threed_thin.append(thin)
 
            chain_new = c[burn::thin]
            chain_lengths.append(len(chain_new))
 
            if plotting==True:
                plt.plot(range(len(acp)), acp, linewidth=1.0)
                plt.xlabel("Position in chain / Lag time")
                plt.ylabel("ACF")
                plt.title("Autocorrelation function vs lag time")
                plt.show()
 
#                 plt.plot(range(len(chain_new)), chain_new, linewidth=1.0)
#                 plt.xlabel("Position in chain")
#                 plt.ylabel("Chain value")
#                 plt.title("New Chain")
#                 plt.show()
        print("burning parameter = ", burn)
        print("proposed thinning parameters from the 3 chains: ", threed_thin)
        print("chain lengths from proposed thins", chain_lengths)
 
        return burn, threed_thin
 
 
 
#------------------------------------
# Deprecated Functions
#------------------------------------
 
def load_results(filename = 'results.nc'):
 
    """Loads the results as a dataset NetCDF object to be passed to other functions
    
    Parameters:
        filename: The name of the file containing the results, defaults to 'results.nc'. """
    
    try:
        rootgrp = Dataset(filename)
    except OSError as err:
        print(f"OS error: {err}")
    return rootgrp
 
 
def load_global_attr(rootgrp, verbose = True):
    """Loads the global metadata in the results file into a dictionary
    
    Parameters:
        rootgrp: NetCDF4 dataset object which loaded in using 'load_results()'.
        verbose: Defaults to true. Prints the value of global attributes which have been loaded
    """
    global_dict ={}
 
    for name in rootgrp.ncattrs():
        if verbose:
            print("{} = {}".format(name, getattr(rootgrp, name)))
 
    global_dict[name] = getattr(rootgrp, name)
 
 
    return global_dict
 
 
def load_variables(rootgrp, verbose = True):
    """Loads the variables in the results file into a dictionary
    
    Parameters:
        rootgrp: NetCDF4 dataset object which loaded in using 'load_results()'
        verbose: Defaults to true. Prints the value of variables which have been loaded
    """
    var_dict = {}
    for k in rootgrp.variables.keys():
        if verbose:
            print("{} = {}".format(k, rootgrp[k][:]))
        var_dict[k] = rootgrp[k][:]
 
    return var_dict
 
 
def H2_eng_plot(variables):
    """Plots the energies against the parameters checked during the grid search
       Parameters:
         variables: Dictionary which contains data about the variables in the netcdf4 object"""
 
    fig = go.Figure()
 
    fig.add_trace(
        go.Scatter(
            x=variables['Bondlength_array'],
            y=variables['total_energies'],
            error_y=dict(
            type='data',
            array=variables['Uncertainties'],
            visible=True)
            )
        )
    
    fig.update_layout(
    title=f"Energies vs Bond Lengths",
    xaxis_title="Bond Length",
    yaxis_title="Energy"
    )
    
    fig.show()
    
    return fig
 
 
def QHO_eng_plot(rootgrp,equil=False):
    """energy_plot()
    Plots the energy at each step as the MC converges
    """
 
    energies = np.array(rootgrp.variables['energy_at_alphas_array'][:])
 
    alphas = np.array((rootgrp.variables['Alpha_array'][:]))
 
    if (equil):
        pass
 
    data = go.Scatter(x=alphas, y=energies, mode='markers')
    fig = go.Figure(data)
    fig.show()
    
    return data
 
 
def plot_3D(X,Y,Z,wfn):
    #for plotting wavefunctions, to be removed
    fig = go.Figure(data=go.Volume(
        x=X.flatten(),
        y=Y.flatten(),
        z=Z.flatten(),
        value=wfn.flatten(),
        isomin=0.0,
        isomax=1.2,
        opacity=0.1, 
        surface_count=21, 
    ))
 
    #fig.save("wavefunction.png")
 
    return fig
 
 
def process_main_results(filename='results.nc', verbose = True):
    """Processes the results output as a result of the parameter search from the NetCDF4 file
 
    Parameters:
        filename: Defaults to 'results.nc'. The file which contains the data output by the parameter search
    """
 
 
 
    rootgrp = load_results(filename)
 
    global_attr = load_global_attr(rootgrp, verbose)
 
    variables = load_variables(rootgrp, verbose)
 
 
    if (global_attr['system'] == 'QHO'):
        pass
        #data1 = QHO_eng_plot(rootgrp)
        #fig = QHO_prob_plot(rootgrp)
    elif(global_attr['system'] == 'H2plus'):
        main_fig = H2_eng_plot(variables)
        
 
    elif(global_attr['system'] == 'H2'):
        pass
 
    else:
            raise Exception("Invalid system input into visualisation")
            
    return global_attr,variables,main_fig
 
 
def plot_markov_chain(filename):
 
    dset = nc.Dataset(filename)
    
    energies = dset['energies'][:]
 
    positions = dset['positions'][:] 
    
    if dset.system == 'QHO':
        
        fig = QHO_prob_plot(positions)
        
        return fig
        
    elif dset.system == 'H2' or dset.system == 'H2plus':
 
        R1 = [-0.5*dset.bond_length, 0, 0]
        R2 = [0.5*dset.bond_length, 0, 0]
    
        fig = vis.H2_prob_plot(positions, energies, R1, R2)
        
        return fig
 
 
 
 
def orb_1s(X,Y,Z, R_a): #R_a is the location of the atoms
    x_a,y_a,z_a = R_a
    a0 = 1
    orbital = 1/(((a0**1.5)*(np.sqrt(np.pi)))*\
                 np.exp(np.sqrt((X-x_a)**2+(Y-y_a)**2+(Z-z_a)**2)/a0))
    return orbital
 
 
def orb_2s(r,R_a):
    a0 = 1
    orbital = (2-np.linalg.norm(r-R_a)()/a0)/((4*(a0**1.5)*(np.sqrt(2*np.pi)))*np.exp(np.linalg.norm(r-R_a)/(2*a0)))
    return orbital
 
 
def polar_to_cart(r,theta, phi):
    X = r*np.sin(phi)*np.cos(theta)
    Y = r*np.sin(phi)*np.sin(theta)
    Z = r*np.cos(phi)
 
    return X,Y,Z
    
 
def lcao_wfn(X,Y,Z, R1, R2, pqn = 1):
    #for the linear combination of wavefunctions, to be added to and to be remove
    if pqn == 1:
         wfn = np.sqrt(2)*orb_1s(X,Y,Z,R1) + np.sqrt(2)*orb_1s(X,Y,Z,R2)
    if pqn == 2:
        pass
 
    return wfn
 
 
def plot_3D(X,Y,Z,wfn,R1,R2):
    """ Function for plotting the wavefunction analytically. Meshgrids can be generated using np.meshgrid(x,y,z)
 
    Parameters:
        X: Meshgrid of X coordinates
        Y: Meshgrid of Y coordinates
        Z: Meshgrid of Z coordinates
    """ 
 
    fig = go.Figure()
    #for plotting wavefunctions, to be removed
    fig.add_trace(go.Volume(
    x=X.flatten(),
    y=Y.flatten(),
    z=Z.flatten(),
    value=wfn.flatten(),
    opacity=0.1, 
    surface_count=21, 
    ))
    
    fig.add_trace(go.Scatter3d(
        mode='markers',
        x=[R1[0],R2[0]],
        y=[R1[1],R2[1]],
        z=[R1[2],R2[2]],
        opacity=1.0,
        marker=dict(
            color='#AAAAAA',
            size=10,
            line=dict(
                color='Black',
                width=2
                )
            ),
            name='Atoms'
        )
    )
    return fig
 
 
def example_H1s_wfn(bond_length=2.0):
    
    x, y, z = np.mgrid[-3:3:40j, -3:3:40j, -3:3:40j]
 
    R_1 = (-0.5*bond_length, 0,0)
    R_2 = (0.5*bond_length, 0,0)
    
    wfn = lcao_wfn(x,y,z,R_1,R_2)
    density = wfn**2
    
    density_plot = plot_3D(x,y,z,density,R_1,R_2)
    return density_plot
 
    return wfn