<repo_name>metrosierra/quVario<file_sep>/helium_marki/helium_marki.py
#!/usr/bin/env python3

# Made 2020, <NAME>
# mingsongwu [at] outlook [dot] sg
# github.com/starryblack/quVario

### READINGS IN readings FOLDER VERY USEFUL FOR HELIUM ATOM APPROXIMATIONS

import math
import sys
import os
import time
from datetime import datetime


### using sympy methods for symbolic integration. Potentially more precise and convenient (let's not numerically estimate yet)
import numpy as np
import scipy as sp
import scipy.constants as sc
import sympy as sy
from sympy import conjugate, simplify, lambdify, sqrt
from sympy import *
from IPython.display import display
from scipy import optimize, integrate

#### For future monte carlo integration work hopefully
import mcint
import random

### Honestly OOP isnt really shining in advantage now, other than me not caring about the order of functions and using global variables liberally.
class eiGen():

    ### initialisation protocol for eiGen object
    def __init__(self):
        self.hbar = sc.hbar
        self.mass_e = sc.electron_mass
        self.q = sc.elementary_charge
        self.pi = sc.pi
        self.perm = sc.epsilon_0
        self.k = (self.q**2)/(4*self.pi*self.perm)
        sy.init_printing()

        self.macro1()

    ### important bit: generating sets of 3d spherical coordinate sympy symbols for each electron. In this case 2, but can be easily scaled arbitrarily. All generated symbols stored in master self.bases list. Somehow I can't assign each variable a self attribute directly????
    def spawn_electrons(self, number, coord_sys = 's'):
        self.bases = []
        if coord_sys == 'spherical' or 's':
            for i in range(number):
                temp = []
                for q in ['r', 'theta', 'phi']:
                    temp.append(sy.symbols(q + str(i)))
                self.bases.append(temp)

        elif coord_sys == 'cartesian' or 'cart':
            for i in range(number):
                temp = []
                for q in ['x', 'y', 'z']:
                    temp.append(sy.symbols(q + str(i)))
                self.bases.append(temp)

        elif coord_sys == 'cylindrical' or 'cylin':
            for i in range(number):
                temp = []
                for q in ['rho', 'phi', 'z']:
                    temp.append(sy.symbols(q + str(i)))
                self.bases.append(temp)

    ### this is a kind of brute force laplace, because I cannot use the sympy.laplace function AND specify which electron it should operate wrt
    def laplace(self, expr, index, coord_sys = 'spherical'):

        if coord_sys == 'spherical' or 's':
            r = self.bases[index][0]
            t = self.bases[index][1]
            p = self.bases[index][2]
            grad = [sy.diff(expr, r), 1/r * sy.diff(expr, t), 1/(r*sin(t)) * sy.diff(expr, p)]
            lap = (1/r**2)*sy.diff(r**2 * grad[0], r) + (1/(r*sin(t)))*(sy.diff(grad[1] * sin(t), t) + sy.diff(grad[2], p))
        return lap

    ### Just makes the psi function variable. #TODO! augment this to read off a txt file (easier to see or find or collect the various trial forms)

    def spawn_psi(self):
        self.psi = -sy.exp(-self.alpha0*(self.r0 + self.r1))

    ### Just generating parameters. Only need 1 now.
    def spawn_alphas(self, number):
        for i in range(number):
            setattr(self, 'alpha{}'.format(i), sy.symbols('alpha{}'.format(i), real = True))

    ### Automatically slaps the SPHERICAL VOLUME jacobian for self.symintegrate. Tightly paired with said function.
    def jacob_gen(self, index, dimensions = 3, type = 's'):
        if dimensions == 3 and type == 's':
            return self.bases[index][0]**2*sin(self.bases[index][1])


    ## Symbolic integration (volume only for now) with spherical jacobian added for each set of coordinates
    def symintegrate(self, expr, number, dimensions = 3, type = "spherical"):
        if type == 'spherical' or type == 's':
            if dimensions == 3:
                for i in range(number):
                    temp = sy.integrate(expr*self.jacob_gen(index = i), (self.bases[i][0], 0, oo), (self.bases[i][1], 0, pi),(self.bases[i][2], 0, 2*pi), conds = "none")
                    expr = temp
        return temp


    ### Final workhorse that implements scipy fmin function to minimise our expression. The processing is general but the output is specific to our current helium example (1 parameter)
    def iterate(self, func, guess, args):
        print('\nIterator initialised!! Please be patient!!\n')
        starttime = time.time()

        temp = optimize.fmin(func, guess, args = (args), full_output = 1)

        endtime = time.time()
        elapsedtime = endtime - starttime
        now = datetime.now()
        date_time = now.strftime('%d/%m/%Y %H:%M:%S')

        # just returns datetime of attempt, elapsed time, optimised parameter, optimised value, number of iterations
        return [date_time, elapsedtime, guess, temp[0], temp[1], temp[3]]


    def lambdah(self, expr, var, type = 'scipy'):

        return lambdify((var), expr, type)

################################################################
#These are situation specific functions, for our particular helium integration problem.

    ###specifically atomic units
    def hamiltonian_he(self, trial):
        self.ke = -0.5 * (self.laplace(trial, 0) + self.laplace(trial, 1))
        self.nucleus = -(2/self.r0 + 2/self.r1)
        self.repel = (1/ sqrt(self.r0**2 + self.r1**2 - 2*self.r0*self.r1*cos(self.theta0)))

        print('\nKinetic Energy terms (self.ke), Nucleus Attraction terms (self.nucleus), Electron Repulsion terms (self.repel) generated!!!\n')
        return self.ke, self.nucleus, self.repel

    ### This our custom situtaion, where the hydrogenic component is analytic (ie a nice expression with parameter as sole variable) and the cross repulsion term is a numerically evaluated integral. So programme must substitute trial parameter and THEN integrate to check if end value is minimum
    def final_expr1(self, a, cross, hydrogenic):
        integrand = integrate.nquad(cross, [[0, sp.inf],  # r0
                                            [0, sp.pi],   # theta0
                                            [0, sp.inf], # phi0
                                        ], args = (a,)) #phi1

        print(hydrogenic(a) + integrand[0], 'heart trees')
        return hydrogenic(a) + integrand[0]

    ### writes self.iterate results to a paired txr file. If file no exist it is spawned.
    def custom_log(self, data = [], comments = None):
        log = open('marki_log.txt', 'a+')
        separator = '\n{}\n'.format(''.join('#' for i in range(10)))
        info = separator + 'Datetime = {}\n'.format(data[0]) + 'Time taken (s) = {}\n'.format(data[1]) + 'Initial parameter guess (atomic untis) = {}\n'.format(data[2]) + 'Optimised parameter (atomic units) = {}\n'.format(data[3]) + 'Optimised function value (atomic units) = {}\n'.format(data[4]) + 'Number of iterations = {}\n'.format(data[5]) + 'Comments = ({})\n'.format(comments)
        print(info)
        log.write(info)
        log.close()

    ### actual sequence of processing events. Technically this can be transfered to a separate macro script. Hence the name macro.
    def macro1(self):

        self.spawn_alphas(2)
        self.spawn_electrons(2)

        self.r0 = self.bases[0][0]
        self.theta0 = self.bases[0][1]
        self.phi0 = self.bases[0][2]

        self.r1 = self.bases[1][0]
        self.theta1 = self.bases[1][1]
        self.phi1 = self.bases[1][2]

        self.spawn_psi()
        ke, attract, repel = self.hamiltonian_he(self.psi)

        normal = self.symintegrate(conjugate(self.psi)*self.psi, 2)
        result1 = self.symintegrate(conjugate(self.psi)*ke, 2)
        result2 = self.symintegrate(conjugate(self.psi)*attract*self.psi, 2)
        self.hydrogenic = self.lambdah(result1/normal + result2/normal, (self.alpha0))

        jacob_special = 2*4*sp.pi**2*self.r0**2*self.r1**2*sin(self.theta0)
        cross_vars = (self.r0, self.theta0, self.r1, self.alpha0)
        self.cross_term = self.lambdah(simplify(jacob_special*conjugate(self.psi)*repel*self.psi/normal), (cross_vars))

        result = self.iterate(self.final_expr1, guess = 1.6, args = (self.cross_term, self.hydrogenic))
        self.custom_log(result, 'test run 2 with cross term, psi = -self.alpha0*(self.r0 + self.r1)')



    def __enter__(self):
        pass

    def __exit__(self, e_type, e_val, traceback):
        pass

### start script
e = eiGen()
<file_sep>/archives/optipack_deprecated.py
#!/usr/bin/env python3

# Made 2020, <NAME>, <NAME>
# mingsongwu [at] outlook [dot] sg
# github.com/starryblack/quVario


### This script provides the montypython object AND minimiss object that handles any calls to monte carlo integration functions and minimisation functions of our design. It interfaces with the python packages installed for basic functionalities (ie mcint)
### optipack should serve helium_markii.py which is the higher level script for direct user interface


### as of 27 may 2020 montypython uses mcint, and minimiss uses scipy fmin functions as main work horses.

### montypython removed, replaced with simple functions because this is the only sane way to be able to numba jit them

import math
import sys
import os
import time
from datetime import datetime


### using sympy methods for symbolic integration. Potentially more precise and convenient (let's not numerically estimate yet)
import numpy as np
import scipy as sp
import scipy.constants as sc
import sympy as sy
from sympy import conjugate, simplify, lambdify, sqrt
from sympy import *
from IPython.display import display
from scipy import optimize, integrate

import matplotlib.pyplot as plt

#### For future monte carlo integration work hopefully
import mcint
import random
from numba import jit, njit


def integrator_uniform(self, integrand, bounds, sample_iter):
    ''' using mcint package to determine the integral via uniform distribution
    sampling over an interval

    '''
    # measure is the 'volume' over the region you are integrating over

    result,error = mcint.integrate(integrand, self.sampler(bounds),
                                    self.get_measure(bounds), sample_iter)

    return result, error


def get_measure(self, bounds):
    ''' obtains n dimensional 'measure' for the integral, which effectively
        is the volume in n dimensional space of the integral bounds. used for
        multiplying the expected value.

        inputs:
            bounds: list of size n, indicating the n bounds of the definite
                integral

        outputs:
            measure: float
    '''
    measure = 1

    if bounds.ndim == 1:
        dimlength = bounds[1] - bounds[0]
        measure *= dimlength

    else:
        for i in bounds:

            dimlength = i[1] - i[0]
            measure *= dimlength

    return measure


def sampler(self, bounds):
    ''' generates a tuple of n input values from a random uniform distribution
        e.g. for three dimensions, outputs tuple = (x,y,z) where x,y,z are
        floats from a uniorm distribution

        inputs:
            bounds

        outputs:
            sample (tuple)
    '''
    while True:
        sample = ()

        if bounds.ndim == 1:
                dimsample = random.uniform(bounds[0],bounds[1])
                x = list(sample)
                x.append(dimsample)
                sample = tuple(x)

        else:
            for i in bounds:
                dimsample = random.uniform(i[0],i[1])

                x = list(sample)
                x.append(dimsample)
                sample = tuple(x)

        yield sample


def integrand(self, x):
    ''' this is the integrand function

    inputs: x (array), where x[0] denotes the first variable

    '''

    return x[0]**2 * sp.exp(-(x[0]) ** 2)

def combine_pq(self, pfunc, qfunc):
    ''' multiplies p and q for the uniform integrator!
    '''
    return pfunc * qfunc

######## These helper functions concern the Metropolis algorithm implementation of the integral

@njit
def metropolis_hastings(pfunc, iter, alpha, dims):

    # we make steps based on EACH electron, the number of which we calculate from dims/3
    # 'scale' of the problem. this is arbitrary
    s = 3.
    # we discard some initial steps to allow the walkers to reach the distribution
    equi_threshold = 0
    initial_matrix = np.zeros((int(dims/3), 3))
    reject_ratio = 0.
    therm = 1
    test = []
    samples = []

    for i in range(int(dims/3)):
        initial_matrix[i] = 2.*s*np.random.rand(3) - s


    # now sample iter number of points
    for i in range(iter):

        # choose which electron to take for a walk:
        e_index = np.random.randint(0, dims/3)
        trial_matrix = initial_matrix.copy()
        trial_matrix[e_index] += (2.*s*np.random.rand(3) - s)/(dims/3)

        # trial_matrix[e_index] = hi
        proposed_pt = np.reshape(trial_matrix, (1, dims))[0]
        initial_pt = np.reshape(initial_matrix, (1, dims))[0]

        # print(initial_pt)
        p = pfunc(proposed_pt, alpha) / pfunc(initial_pt, alpha)
        if p > np.random.rand():
            initial_matrix = trial_matrix.copy()
            # print(initial_matrix == trial_matrix)

        else:
            reject_ratio += 1./iter

        if i > equi_threshold:
            if (i-therm)%therm == 0:
                test.append(np.reshape(initial_matrix, (1, dims))[0][3])
                samples.append(np.reshape(initial_matrix, (1, dims))[0])

    return samples, reject_ratio, test


@njit
def integrator_mcmc(pfunc, qfunc, sample_iter, walkers, alpha, dims):

    therm = 0
    vals = np.zeros(walkers)
    val_errors = 0.
    test = []
    for i in range(walkers):
        mc_samples, rejects, p  = metropolis_hastings(pfunc, sample_iter, alpha, dims)
        sums = 0.

        # obtain arithmetic average of sampled Q values
        for array in mc_samples[therm:]:
            sums += qfunc(array, alpha)
            test.append(qfunc(array, alpha))
        vals[i] = (sums/(sample_iter - therm))
    # also calculate the variance
    vals_squared = np.sum(vals**2)
    vals_avg = np.sum(vals) /walkers
    variance = vals_squared/walkers - (vals_avg) ** 2
    std_error = np.sqrt(variance/walkers)


    print('Iteration cycle complete, result = ', vals_avg, 'error = ', std_error, 'rejects = ', rejects)
    return vals_avg, std_error, rejects, test, p


@njit
def mcmc_q(x, alpha):
    ''' this is the Q part of the integrand function for mcmc
    inputs: x(array), denoting iter number of sample points, given by
    '''


    ### helium local energy
    r1 = x[0:2]
    r2 = x[3:]
    r1_len = np.sqrt(x[0]**2 + x[1]**2 + x[2]**2)
    r2_len = np.sqrt(x[3]**2 + x[4]**2 + x[5]**2)


    r1_hat = r1 / r1_len
    r2_hat = r2 / r2_len

    r12 = np.sqrt((x[0]-x[3])**2 + (x[1]-x[4])**2 + (x[2]-x[5])**2)

    return ((-4 + np.dot(r1_hat - r2_hat, r1 - r2 ) / (r12 * (1+alpha*r12)**2)
            - 1/ (r12*(1+alpha*r12)**3)
            - 1/ (4*(1+alpha*r12)**4)
            + 1/ r12 ))[0]

@njit
def mcmc_p(x, alpha):
    '''
    this is the integrand function for mcmc

    '''
    ### helium wavefunction squared
    r1 = x[0:3]
    r2 = x[3:]
    r1_len = np.sqrt(x[0]**2 + x[1]**2 + x[2]**2)
    r2_len = np.sqrt(x[3]**2 + x[4]**2 + x[5]**2)
    r12 = np.sqrt((x[0]-x[3])**2 + (x[1]-x[4])**2 + (x[2]-x[5])**2)

    return ((np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
            * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2)[0]

#########################
class MiniMiss():

    def __init__(self):
        print('MiniMiss optimisation machine initialised and ready!\n')

    def minimise(self, func, guess, ftol):
        starttime = time.time()

        temp = optimize.fmin(func, guess, full_output = 1, ftol = ftol)

        endtime = time.time()
        elapsedtime = endtime - starttime
        now = datetime.now()
        date_time = now.strftime('%d/%m/%Y %H:%M:%S')
        # just returns datetime of attempt, elapsed time, optimised parameter, optimised value, number of iterations
        return [date_time, elapsedtime, guess, temp[0], temp[1], temp[3]]

    def __enter__(self):
        return self

    def __exit__(self, e_type, e_val, traceback):
        print('\n\nMiniMiss object self-destructing\n\n')
<file_sep>/helium_marki/hydr_var.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon May 18 10:42:53 2020

Variational Method for finding ground state energy of Hydrogen.
Exercise 1.19 in Quantum Chemistry. 
The Hamiltonian is fixed at the moment.

@author: kenton
"""
#%% Import packages and initialisation of spaces
from sympy import (exp, conjugate, symbols, simplify, lambdify)
from sympy.vector import CoordSys3D, Del
from sympy.parsing.sympy_parser import parse_expr
from scipy.optimize import fmin
from scipy.integrate import quad
#import scipy as sp
#import matplotlib.pyplot as plt 

# initialisation of coordinate systems for vector calculus
# we work in Spherical coordinates because of the nature of the problem
# this simplifies the need for the Laplacian as SymPy has sphericals built in
R = CoordSys3D('R', transformation = 'spherical')
delop = Del() #by convention

#Sympy help
#
#Note that R.r, R.theta, R.phi are the variables called base scalars
#    Base scalars can be treated as symbols
#and R.i, R.j, R.k are the r, theta, phi basis vectors
#
#https://docs.sympy.org/latest/tutorial/basic_operations.html

#%% Define symbolic parameters and numerical infinity

alpha = symbols('alpha', real = 0)

inf = 1000

#%% Define functions

#Hamiltonian operator H operating on psi
def H(psi):
    ''' Function: the Hamiltonian operator operating on some trial function psi.
    This is in atomic units. 
   
    Input: 
        some trial function psi(r)
    
    Output: 
        another scalar field Sympy expression
    '''
    H = - 1 / 2 * del2(psi) - 1 / R.r * psi
    
    return H

#Del-squared, the Laplacian for the Schrodinger equation
def del2(f):
    '''Laplacian operator in spherical polar coordinates
    
    Input: 
        some Sympy expression f to which the Laplacian is applied 
    
    Output: 
        A Sympy expression
    '''
    
    del2 = delop.dot(delop(f)).doit()
    
    return del2

# We need to supply the Jacobian manually because we are no longer integrating 
# with SymPy.
# The Jacobian of the spherical integral is r**2 sin(theta)
# ignore sin(theta) because of symmetry 
jacobian = R.r ** 2 

def readtrialfuncs(txtdoc = 'hydr_var_trialfuncs.txt'):
    ''' Reads a text document of trial functions and creates a sympy list
    
    inputs:
        txtdoc: text document filename
        
    outputs:
        trialfuns: a list of trial functions
        
    '''
    
    f = open(txtdoc, "r")
    
    trialfuncs = []
    
    for line in f:
        x = parse_expr(line)
        
        x = x.subs(symbols('r'), R.r)
        
        trialfuncs.append(x)
        
    f.close()
    
    return trialfuncs

# Trial function psi defined using SymPy
# A decaying Gaussian 
#psi = exp(-alpha * R.r**2)


#%% This part defines function related to the variational method

# definining a function such that SciPy Optimisation can be used
def expectation_intgrl(expectation_lamb, a, stpt, endpt):
    ''' Formualtes an integral of the lambdified expectation function
   
    inputs:
        a (float): parameter
            
        stpt (float): starting point of the integral 
        
        endpt (float): endpoint of the integral
    
    output: 
        expectation_integral[0]: numerical result of the integral
    
    '''
    
    expectation_integral = quad(expectation_lamb, 0, endpt,args=(a))
    
    return expectation_integral[0]


def norm_intgrl(norm_lamb, a, stpt, endpt):
    ''' Finds the square of the Normalisation constant of the trial function.
    
    Formualtes an integral of the lambdified normalisation function and uses
    the fact that the trial functions must be normalised to 1. 
   
    inputs:
        a (float): 
            
        stpt (float): starting point of the integral
        
        endpt (float): endpoint of the integral
    
    output: 
        N_squared: normalisation
    
    '''
    norm_intgrl = quad(norm_lamb, stpt, endpt, args = (a)) 
    
    N_squared = 1 / norm_intgrl[0]
    
    return N_squared


def var_intgrl(a, expectation_lamb, norm_lamb, stpt = 0, endpt = 1000):
    ''' Multiplies the normalisation with the expectation integral
    
    inputs: 
        expectation_intgrl (function of a): 
            
        norm_intgrl (function of a): N_squared, normalisation for trial function
        
        stpt (float): starting point of the integral (0 by default)
        
        endpt (float): endpoint of the integral (1000 effectively infinity)
        
    outputs: 
        var_intgrl (function of a): numerical result of the 
            variation integral
    '''
    
    var_intgrl = (norm_intgrl(norm_lamb, a, stpt, endpt) * 
                  expectation_intgrl(expectation_lamb, a, stpt, endpt))
    
    return var_intgrl


def energy(psi):
    ''' Obtains the optimised parameter and upper bound of the energy of the system
    
    inputs: 
        trial function psi (sympy expression) that satisfies the BCs 
    
    outputs: 
        alpha_op and energy_max (tuple)
    
    '''
    #The expectation value  <psi|H|psi>
    expectation = simplify(conjugate(psi) * H(psi) ) 
    
    expectation_lamb = lambdify((R.r, alpha), expectation * jacobian, 'scipy')
    
    #the normalisation value <psi|psi> 
    norm = conjugate(psi) * psi 
    
    norm_lamb = lambdify((R.r, alpha), norm * jacobian, 'scipy')
    
    #supply the initial guess for the variation intergral
    initial_guess = 0.1
    
    #minimise the variation integral with respect to the parameter
    alpha_op, energy_max = fmin(var_intgrl, 
                            initial_guess, (expectation_lamb, norm_lamb), 
                            disp= 0, full_output = 1)[0:2]
    
    return (alpha_op, energy_max)

def main():
    ''' Main function, obtains the lowest energy of all trial functions and prints
        the intermediate energies
    
    outputs: overall_energy_min
        
    '''
    energies = []
    
    print('In atomic units,')
    
    for psi in readtrialfuncs():
        
        alpha_op, energy_min = energy(psi)
        
        energies.append(energy_min)
        
        print('\nFor psi = %s, \nThe optimised parameter is %f. \nThe energy upper bound is %f.' % (psi, float(alpha_op),energy_min))
    
    overall_energy_min = min(energies)
    
    return overall_energy_min


#%% Execution of the minimisation of the variation integral

print('\nThe overall maximum energy is %f.' %main())
<file_sep>/archives/metronumba.py
#!/usr/bin/env python3
#%%
import math
import sys
import os
import time
from datetime import datetime



import numpy as np
import scipy as sp
import scipy.constants as sc

#### For future monte carlo integration work hopefully
import mcint
import random
import numba as nb
from numba import jit, njit, prange

'''
https://jellis18.github.io/post/2018-01-02-mcmc-part1/
^ this is a really informative website about calibrating MCMC algorithms based on acceptance reject ratio
'''




#%%

@njit
def metropolis_hastings(pfunc, iter, alpha, dims):

    # we make steps based on EACH electron, the number of which we calculate from dims/3
    # 'scale' of the problem. this is arbitrary
    s = 3.
    # we discard some initial steps to allow the walkers to reach the distribution
    equi_threshold = 0
    initial_matrix = np.zeros((int(dims/3), 3))
    reject_ratio = 0.
    therm = 1
    test = []
    samples = []

    for i in range(int(dims/3)):
        initial_matrix[i] = 2.*s*np.random.rand(3) - s


    # now sample iter number of points
    for i in range(iter):

        # choose which electron to take for a walk:
        e_index = np.random.randint(0, dims/3)
        trial_matrix = initial_matrix.copy()
        trial_matrix[e_index] += (2.*s*np.random.rand(3) - s)/(dims/3)

        # trial_matrix[e_index] = hi
        proposed_pt = np.reshape(trial_matrix, (1, dims))[0]
        initial_pt = np.reshape(initial_matrix, (1, dims))[0]

        # print(initial_pt)
        p = pfunc(proposed_pt, alpha) / pfunc(initial_pt, alpha)
        if p > np.random.rand():
            initial_matrix = trial_matrix.copy()
            # print(initial_matrix == trial_matrix)

        else:
            reject_ratio += 1./iter

        if i > equi_threshold:
            if (i-therm)%therm == 0:
                test.append(np.reshape(initial_matrix, (1, dims))[0][3])
                samples.append(np.reshape(initial_matrix, (1, dims))[0])

    return samples, reject_ratio, test


@njit
def integrator_mcmc(pfunc, qfunc, sample_iter, walkers, alpha, dims):

    therm = 0
    vals = np.zeros(walkers)
    val_errors = 0.
    test = []
    for i in range(walkers):
        mc_samples, rejects, p  = metropolis_hastings(pfunc, sample_iter, alpha, dims)
        sums = 0.

        # obtain arithmetic average of sampled Q values
        for array in mc_samples[therm:]:
            sums += qfunc(array, alpha)
            test.append(qfunc(array, alpha))
        vals[i] = (sums/(sample_iter - therm))
    # also calculate the variance
    vals_squared = np.sum(vals**2)
    vals_avg = np.sum(vals) /walkers
    variance = vals_squared/walkers - (vals_avg) ** 2
    std_error = np.sqrt(variance/walkers)


    print('Iteration cycle complete, result = ', vals_avg, 'error = ', std_error, 'rejects = ', rejects)
    return vals_avg, std_error, rejects, test, p


@njit
def mcmc_q(x, alpha):
    ''' this is the Q part of the integrand function for mcmc
    inputs: x(array), denoting iter number of sample points, given by
    '''


    ### helium local energy
    r1 = x[0:2]
    r2 = x[3:]
    r1_len = np.sqrt(x[0]**2 + x[1]**2 + x[2]**2)
    r2_len = np.sqrt(x[3]**2 + x[4]**2 + x[5]**2)


    r1_hat = r1 / r1_len
    r2_hat = r2 / r2_len

    r12 = np.sqrt((x[0]-x[3])**2 + (x[1]-x[4])**2 + (x[2]-x[5])**2)

    return ((-4 + np.dot(r1_hat - r2_hat, r1 - r2 ) / (r12 * (1+alpha*r12)**2)
            - 1/ (r12*(1+alpha*r12)**3)
            - 1/ (4*(1+alpha*r12)**4)
            + 1/ r12 ))[0]

@njit
def mcmc_p(x, alpha):
    '''
    this is the integrand function for mcmc

    '''
    ### helium wavefunction squared
    r1 = x[0:3]
    r2 = x[3:]
    r1_len = np.sqrt(x[0]**2 + x[1]**2 + x[2]**2)
    r2_len = np.sqrt(x[3]**2 + x[4]**2 + x[5]**2)
    r12 = np.sqrt((x[0]-x[3])**2 + (x[1]-x[4])**2 + (x[2]-x[5])**2)

    return ((np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
            * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2)[0]
<file_sep>/helium_markii/uniform.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jun 17 16:24:56 2020

Basic Uniform Monte Carlo Integrator
Applied to Variational Calculations
"""

import time
import numpy as np
from scipy.optimize import fmin
import matplotlib.pyplot as plt
from numba import njit, prange

@njit(parallel = True)
def Uint(integrand, sampler, bounds, measure, n, alpha):
    ''' Obtains integral result
    Inputs:
        integrand: function
        sampler: uniform sampler, function
        bounds: integral bounds, array
        n: number of samples per dimension
    '''

    # this takes n samples from the integrand and stores it in values
    values = np.zeros(n)

    for i in prange(n):
        sample = sampler(bounds)
        val = integrand(sample, alpha)
        values[i] = val

    # this is the value of the sampled integrand values
    result = measure * np.sum(values) / n
    return result

@njit
def Uint_iter(integrand, sampler, bounds, measure, n, iters, alpha):
        ''' returns array of integral results
        '''
        # obtain iters integral results
        results = np.zeros(iters)
        for i in range(iters):
            results[i] = Uint(integrand, sampler, bounds, measure, n, alpha)
        return results

@njit
def get_measure(bounds):
    ''' obtains the volume of the dims dimensional hypercube.
        Used for multiplying the expected value.

        inputs:
            bounds: array of size dims, indicating the dims bounds of the definite
                integral
        outputs:
            measure: float
    '''
    measure = 1.
    dims = len(bounds)
    for i in range(dims):
        b = bounds[i]
        dimlength = float(b[1] - b[0])
        measure *= dimlength
    return measure

@njit
def sampler(bounds):
    ''' returns a uniformly distributed random array of size dims

    inputs:
        bounds is a 2D array, specifying the integration range of each dim
    '''
    dims = len(bounds)
    samples = np.zeros(dims)
    for i in range(dims):
        b = bounds[i]
        samples[i] = np.random.uniform(b[0], b[1])
    return samples

@njit
def psiHpsi(x, alpha):
    ''' Local energy term
    '''
    r1 = x[0:3]
    r2 = x[3:]

    r1_len = np.sqrt(x[0]**2 + x[1]**2 + x[2]**2)
    r2_len = np.sqrt(x[3]**2 + x[4]**2 + x[5]**2)


    r1_hat = r1 / r1_len
    r2_hat = r2 / r2_len

    r12 = np.sqrt((x[0]-x[3])**2 + (x[1]-x[4])**2 + (x[2]-x[5])**2)

    EL = ((-4 + np.dot(r1_hat - r2_hat, r1 - r2 ) / (r12 * (1+alpha[0]*r12)**2)
            - 1/ (r12*(1+alpha[0]*r12)**3)
            - 1/ (4*(1+alpha[0]*r12)**4)
            + 1/ r12 ))

    psisq = ((np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
            * np.exp(r12 / (2 * (1+ alpha[0]* r12)))) ** 2)

    return psisq* EL

@njit
def psisq(x, alpha):
    ''' Squared trial wavefunction
    '''
    r1_len = np.sqrt(x[0]**2 + x[1]**2 + x[2]**2)
    r2_len = np.sqrt(x[3]**2 + x[4]**2 + x[5]**2)

    r12 = np.sqrt((x[0]-x[3])**2 + (x[1]-x[4])**2 + (x[2]-x[5])**2)

    psisq = (np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
            * np.exp(r12 / (2 * (1+ alpha[0]* r12)))) ** 2

    return psisq

# I just can't jit this part for some odd reason? All functions that it calls are
    # jitted though
@njit
def evalenergy(alpha):
    #initialise settings
    domain = 2.
    dims = 6

    bounds = []
    for i in range(dims):
        bounds.append([-domain, domain])
    bounds = np.array(bounds)
    n = 100000
    iters = 30
    measure  = get_measure(bounds)

    expresults = Uint_iter(psiHpsi, sampler, bounds, measure, n, iters, alpha)
    normresults = Uint_iter(psisq, sampler, bounds, measure, n, iters, alpha)

    #obtain average and variance
    vals = expresults/normresults
    avg = np.sum(vals) / iters
    vals_squared = np.sum(vals**2)
    var = (vals_squared/ iters - avg **2)
    std = np.sqrt(var)

    E = avg
    print(E, std)
    # print('When alpha is {}, the energy is {} with std {}' .format(alpha, E, std))
    return E


### Minimisation algorithm

# start_time = time.time()
# fmin(evalenergy, 0.1, full_output = 1, ftol = 1)
# duration = (time.time() - start_time)
# print('Time taken:', duration)

#%%%%%%%%%%%%%%%%%%%%%%%%
### Graph Plotting stuff

def plot_iter_graphs(itegrand, bounds, iter_range, lit_val, plotval=1, plotnormval=1, plottime=1):
    ''' This function is used for plotting graphs
    '''

    results = []
    variances = []
    times = []

    def get_lit(x):
        return float(lit_val) + 0*x

    for iters in iter_range:
        iters = int(iters)
        start_time = time.time()
        x = Uint(integrand, sampler, bounds, iters)
        duration = (time.time() - start_time)
        results.append(x[0])
        variances.append(x[1])
        times.append(duration)

    if plotval:
        plt.plot(np.log10(iter_range), results)
        plt.xscale = 'log'
        plt.ylabel('Value')
        plt.xlabel('Iterations log scale')
        plt.grid()
        plt.plot()
        plt.show()

    if plotnormval:
        xrange = np.log10(iter_range)
        plt.plot(xrange, results/get_lit(np.log10(iter_range)))
        plt.xscale = 'log'
        plt.ylabel('Normalised Value')
        plt.xlabel('Iterations log scale')
        plt.grid()
        plt.plot()
        plt.show()

    if plottime:
        plt.plot(np.log10(iter_range), times)
        plt.xscale = 'log'
        plt.ylabel('Time taken (s)')
        plt.xlabel('Iterations log scale')
        plt.grid()
        plt.show()
    pass
<file_sep>/helium_markii/optipack.py
#!/usr/bin/env python3

# Made 2020, <NAME>, <NAME>
# mingsongwu [at] outlook [dot] sg
# github.com/starryblack/quVario


### This script provides the montypython object AND minimiss object that handles any calls to monte carlo integration functions and minimisation functions of our design. It interfaces with the python packages installed for basic functionalities (ie mcint)
### optipack should serve helium_markii.py which is the higher level script for direct user interface


### as of 27 may 2020 montypython uses mcint, and minimiss uses scipy fmin functions as main work horses.

### montypython removed, replaced with simple functions because this is the only sane way to be able to numba jit them

import math
import sys
import os
import time
from datetime import datetime
import statistics

### using sympy methods for symbolic integration. Potentially more precise and convenient (let's not numerically estimate yet)
import numpy as np
import scipy as sp
import scipy.constants as sc
from scipy import optimize, integrate

import sympy as sy
from sympy import conjugate, simplify, lambdify, sqrt
from sympy import *
from IPython.display import display

import vegas


import matplotlib.pyplot as plt

import random
from numba import jit, njit, prange


@njit
def metropolis_hastings(pfunc, iter, alpha, dims):
    # we make steps based on EACH electron, the number of which we calculate from dims/3
    # 'scale' of the problem. this is arbitrary
    s = 3.
    # we discard some initial steps to allow the walkers to reach the distribution
    equi_threshold = 0
    initial_matrix = np.zeros((int(dims/3), 3))
    reject_ratio = 0.
    therm = 1
    test = []
    samples = []

    for i in range(int(dims/3)):
        initial_matrix[i] = 2.*s*np.random.rand(3) - s


    # now sample iter number of points
    for i in range(iter):

        # choose which electron to take for a walk:
        e_index = np.random.randint(0, dims/3)
        trial_matrix = initial_matrix.copy()
        trial_matrix[e_index] += (2.*s*np.random.rand(3) - s)/(dims/3)

        # trial_matrix[e_index] = hi
        proposed_pt = np.reshape(trial_matrix, (1, dims))[0]
        initial_pt = np.reshape(initial_matrix, (1, dims))[0]

        # print(initial_pt)
        p = pfunc(proposed_pt, alpha) / pfunc(initial_pt, alpha)
        if p > np.random.rand():
            initial_matrix = trial_matrix.copy()
            # print(initial_matrix == trial_matrix)

        else:
            reject_ratio += 1./iter

        if i > equi_threshold:
            if (i-therm)%therm == 0:
                test.append(np.reshape(initial_matrix, (1, dims))[0][3])
                samples.append(np.reshape(initial_matrix, (1, dims))[0])

    return samples, reject_ratio, test


@njit
def integrator_mcmc(pfunc, qfunc, sample_iter, walkers, alpha, dims, verbose = True):

    therm = 0
    vals = np.zeros(walkers)
    val_errors = 0.
    test = []
    for i in range(walkers):
        mc_samples, rejects, p  = metropolis_hastings(pfunc, sample_iter, alpha, dims)
        sums = 0.

        # obtain arithmetic average of sampled Q values
        for array in mc_samples[therm:]:
            sums += qfunc(array, alpha)
            test.append(qfunc(array, alpha))
        vals[i] = (sums/(sample_iter - therm))
    # also calculate the variance
    vals_squared = np.sum(vals**2)
    vals_avg = np.sum(vals) /walkers
    variance = vals_squared/walkers - (vals_avg) ** 2
    std_error = np.sqrt(variance/walkers)

    if verbose:
        print('Iteration cycle complete, result = ', vals_avg, 'error = ', std_error, 'rejects = ', rejects, 'alpha current = ', alpha)
    return vals_avg, std_error, rejects, test, p


#%%%%%%%%%%%%%%%%%%%%%%%

@njit(parallel = True)
def Uint(integrand, bounds, n, alpha):

    # this takes n samples from the integrand and stores it in values
    values = 0
    for x in prange(n):
        sample = np.zeros(len(bounds))
        for i in range(len(bounds)):
            sample[i] = np.random.uniform(bounds[i][0], bounds[i][1])
        val = integrand(sample, alpha)
        values += val

    # this is the value of the sampled integrand values
    result = values / n
    return result


@njit
def Ueval(normalisation, expectation, n, iters, alpha, dimensions):
    #initialise settings
    domain = 3.
    dims = dimensions
    #origin centered symmetrical volume
    bounds = []
    for i in range(dims):
        bounds.append([-domain, domain])
    bounds = np.array(bounds)

    measure = 1.
    for i in range(dims):
        dimlength = float(bounds[i][1] - bounds[i][0])
        measure *= dimlength

    results = np.zeros(iters)
    normresults = np.zeros(iters)
    for i in range(iters):
        results[i] = Uint(expectation, bounds, n, alpha) * measure
        normresults[i] = Uint(normalisation, bounds, n, alpha) * measure

    #obtain average and variance
    vals = results / normresults
    avg = np.sum(vals) / iters
    vals_squared = np.sum(vals**2)
    var = (vals_squared/ iters - avg **2)
    std = np.sqrt(var)

    E = avg
    print(E, std )
    # print('When alpha is {}, the energy is {} with std {}' .format(alpha, E, std))
    return E, std

#%%%%%%%%%%%%%%%%%%%%%%%

class LasVegas():

    def __init__(self):
        print('LasVegas up for business!')

    def vegas_int(self, expec, norm, evals, iter, dimensions, volumespan, verbose = True):

        self.final_results = {}

        start_time = time.time()

        # assign integration volume to integrator
        bound = volumespan
        dims = dimensions
        # creates symmetric bounds specified by [-bound, bound] in dims dimensions
        symm_bounds = dims * [[-bound,bound]]

        # simultaneously initialises expectation and normalisation integrals
        expinteg = vegas.Integrator(symm_bounds)
        norminteg = vegas.Integrator(symm_bounds)

        # adapt to the integrands; discard results
        expinteg(expec, nitn = 5, neval = 1000)
        norminteg(norm, nitn = 5, neval = 1000)
        # do the final integrals
        expresult = expinteg(expec, nitn = iter, neval = evals)
        normresult = norminteg(norm, nitn = iter, neval = evals)


        E = expresult.mean/normresult.mean
        if verbose:
            print('Energy is {} when alpha is'.format(E), ' with sdev = ', [expresult.sdev, normresult.sdev])
            print("--- Iteration time: %s seconds ---" % (time.time() - start_time))

        self.final_results['energy'] = E
        self.final_results['dev'] = [expresult.sdev, normresult.sdev]
        self.final_results['pvalue'] = [expresult.Q, normresult.Q]

        return self.final_results


    def __enter__(self):
        return self

    def __exit__(self, e_type, e_val, traceback):
        print('\n\nLasVegas object self-destructing\n\n')


#%%%%%%%%%%%%%%%%%%%%%%%

class MiniMiss():

    def __init__(self):
        print('MiniMiss optimisation machine initialised and ready!\n')

    def neldermead(self, func, guess, ftol, bounds = [(0., 3.),(0., 3.),(0., 0.3) ]):
        starttime = time.time()

        temp = optimize.fmin(func, guess, full_output = 1, ftol = ftol)
        # temp = optimize.differential_evolution(func, bounds)
        # temp = optimize.minimize(func, guess, method = 'BFGS', tol = 0.01)

        endtime = time.time()
        elapsedtime = endtime - starttime
        now = datetime.now()
        date_time = now.strftime('%d/%m/%Y %H:%M:%S')
        # just returns datetime of attempt, elapsed time, optimised parameter, optimised value, number of iterations
        return [date_time, elapsedtime, guess, temp]

    def gradient(self, func, guess, tolerance, convergence = 'fuzzy', args = None):

        print('\nGradient descent initiated, type of convergence selected as ' + convergence +'\n')
        cycle_count = 0
        position0 = guess
        ep = 0.000001
        step = 0.1
        step_ratio = 0.5
        epsilons = np.identity(len(guess))*ep
        delta1 = 10.

        point_collection = []

        satisfied = False

        while satisfied == False:
            cycle_count+=1
            print('Cycle', cycle_count)
            value1 = func(position0)

            vector = np.zeros(len(guess))
            for i in range(len(guess)):
                vector[i] = (func(position0 + epsilons[i]) - func(position0))/ep

            vectornorm = vector/(np.linalg.norm(vector))
            print(vectornorm)
            position0 += -vectornorm * step
            value2 = func(position0)
            delta1 = value2 - value1

            positionprime = position0 + vectornorm*(step*step_ratio)
            value3 = func(positionprime)
            delta2 = value3 - value1
            if delta2 < delta1:
                print('shrink!')
                position0 = positionprime
                step = step*step_ratio
                delta1 = delta2
                value2 = value3

            point_collection.append(value2)
            if convergence == 'strict':
                satisfied = abs(delta1) < tolerance
                finalvalue = value2

            elif convergence == 'fuzzy':
                if len(point_collection) >= 10:
                    data_set = point_collection[-10:]
                    print('std', statistics.pstdev(data_set))
                    print('mean', statistics.mean(data_set))
                    satisfied = abs(statistics.pstdev(data_set)/statistics.mean(data_set)) < tolerance
                    finalvalue = statistics.mean(data_set)
            if cycle > 150:
                print('Max iterations reached, did not converge according to tolerance!!')
                break

        print('Convergence sucess!')
        return finalvalue, position0


    def __enter__(self):
        return self

    def __exit__(self, e_type, e_val, traceback):
        print('\n\nMiniMiss object self-destructing\n\n')
<file_sep>/helium_markii/lasvegas.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
vegas example from the Basic Integrals section
of the Tutorial and Overview section of the
vegas documentation.
"""
from __future__ import print_function   # makes this work for python2 and 3

import vegas
import numpy as np
import time
from numba import jit, njit
from scipy.optimize import fmin
import matplotlib.pyplot as plt

#from optipack import MiniMiss, integrator_mcmc, metropolis_hastings, mcmc_q, mcmc_p
import psiham

psilet_args = {'electrons': 2, 'alphas': 3, 'coordsys': 'cartesian'}
ham = psiham.HamLet(trial_expr = 'threepara1', **psilet_args)
variables, expr = ham.he_expect()
temp1 = ham.vegafy(expr, coordinates = variables, name = 'expec')

variables, expr = ham.he_norm()
temp2 = ham.vegafy(expr, coordinates = variables, name = 'norm')
exec(temp1, globals())
exec(temp2, globals())

final_results = {}

#SHOW_GRID = False
#OPTIM = True
def main(alpha):

    global alpha0
#    global alpha1
#    global alpha2
    alpha0 = alpha[0]
#    alpha1 = alpha[1]
#    alpha2 = alpha[2]

    start_time = time.time()

    # assign integration volume to integrator
    bound = 8
    dims = 6
    # creates symmetric bounds specified by [-bound, bound] in dims dimensions
    symm_bounds = dims * [[-bound,bound]]

    # simultaneously initialises expectation and normalisation integrals
    expinteg = vegas.Integrator(symm_bounds)
    norminteg = vegas.Integrator(symm_bounds)

    # adapt to the integrands; discard results
    expinteg(expec, nitn=5, neval=1000)
    norminteg(norm,nitn=5, neval=1000)
    # do the final integrals
    expresult = expinteg(expec, nitn=10, neval=100000)
    normresult = norminteg(norm, nitn=10, neval=100000)


    if not OPTIM:
        E = expresult.mean/normresult.mean
    else:
        E = expresult[0].mean/ normresult[0].mean
#    print('Energy is {} when alpha is {}'.format(E, alpha), ' with sdev = ', [expresult.sdev, normresult.sdev])
#    print("--- Iteration time: %s seconds ---" % (time.time() - start_time))

    final_results['energy'] = E
    final_results['dev'] = [expresult.sdev, normresult.sdev]
    final_results['pvalue'] = [expresult.Q, normresult.Q]

    return E
#
#### Optimisation ALgorithm
#start_time = time.time()
#OPTIM = False
#
##plt.figure(figsize=(16,10))
#
##print('Plotting function initialised!')
##energies = []
##alpha = np.linspace(0.1, 0.2, 20)
## for i in alpha:
##     alpha0 = i
##
##     energy_samples = []
##     for j in range(1):
##         energy_samples.append(evalenergy(i))
##     avg = np.average(energy_samples)
##     print('Averaged energy for %f is %f' %(i, avg))
##     energies.append(avg)
## plt.plot(alpha, energies, color='dodgerblue')
## print("--- Total time: %s seconds ---" % (time.time() - start_time))
#
## plt.savefig('Vegas')
#
#OPTIM = False
#start_time = time.time()
#result = fmin(main, [0.15], ftol = 0.01, xtol = 0.001, full_output=True)
##e = evalenergy(0.5)
##print(e)
#print("--- Total time: %s seconds ---" % (time.time() - start_time))

#%%%%%%%%%%%%%%%%%%%%
#########################################################################

#def evalenergy(alpha):
#    @vegas.batchintegrand
#    def expec(x):
#        ''' Expectation value function, used as numerator for variational integral
#        '''
#
#        r1 = np.array(x[:,0:3])
#        r2 = np.array(x[:,3:])
#
#        r1_len = np.sqrt(x[:,0]**2 + x[:,1]**2 + x[:,2]**2)
#        r2_len = np.sqrt(x[:,3]**2 + x[:,4]**2 + x[:,5]**2)
#
#        r1_hat = r1 / np.reshape(r1_len, (-1,1))
#        r2_hat = r2 / np.reshape(r2_len, (-1,1))
#
#        r12 = np.sqrt((x[:,0]-x[:,3])**2 + (x[:,1]-x[:,4])**2 + (x[:,2]-x[:,5])**2)
#
#        EL =  (-4 + np.sum((r1_hat-r2_hat)*(r1-r2), 1) / (r12 * (1 + alpha*r12)**2)
#                - 1/ (r12*(1 + alpha*r12)**3)
#                - 1/ (4*(1 + alpha*r12)**4)
#                + 1/ r12 )
#
#        psisq = (np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
#                * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2
#        return psisq * EL
#
#    @vegas.batchintegrand
#    def norm(x):
#        ''' Squared trial wavefunction, used as denominator for variational integral
#        '''
##        x = np.reshape(x, (1,-1))
#
#        r1_len = np.sqrt(x[:,0]**2 + x[:,1]**2 + x[:,2]**2)
#        r2_len = np.sqrt(x[:,3]**2 + x[:,4]**2 + x[:,5]**2)
#
#        r12 = np.sqrt((x[:,0]-x[:,3])**2 + (x[:,1]-x[:,4])**2 + (x[:,2]-x[:,5])**2)
#
#        psisq = (np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
#                * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2
#        return psisq
#
#    def main():
#        start_time = time.time()
#
#        # assign integration volume to integrator
#        bound = 8
#        dims = 6
#        # creates symmetric bounds specified by [-bound, bound] in dims dimensions
#        symm_bounds = dims * [[-bound,bound]]
#
#        # simultaneously initialises expectation and normalisation integrals
#        expinteg = vegas.Integrator(symm_bounds)
#        norminteg = vegas.Integrator(symm_bounds)
#
#        # adapt to the integrands; discard results
#        expinteg(expec, nitn=5, neval=1000)
#        norminteg(norm,nitn=5, neval=1000)
#        # do the final integrals
#        expresult = expinteg(expec, nitn=10, neval=10000)
#        normresult = norminteg(norm, nitn=10, neval=10000)
#
#### Code for printing the grid
#### and diagnostics
##        print(expresult.summary())
##        print('expresult = %s    Q = %.2f' % (expresult, expresult.Q))
##        if SHOW_GRID:
##            expinteg.map.show_grid(20)
##
##        print(normresult.summary())
##        print('normresult = %s    Q = %.2f' % (normresult, normresult.Q))
##        if SHOW_GRID:
##            norminteg.map.show_grid(20)
#
#        ### obtain numerical result
#        ### Different expressions for plotting/ minimisation algorithm
#        if not OPTIM:
#            E = expresult.mean/normresult.mean
#        else:
#            E = expresult[0].mean/ normresult[0].mean
##        print('Energy is %f when alpha is %f' %(E, alpha))
##        print("--- Iteration time: %s seconds ---" % (time.time() - start_time))
#        final_results['energy'] = E
#        final_results['dev'] = [expresult.sdev, normresult.sdev]
#        final_results['pvalue'] = [expresult.Q, normresult.Q]
#        return E
#    E = main()
#    return E


# #%% This is for one-parameter plotting
# #### Run plotting function, plot alpha against energies
# start_time = time.time()
# OPTIM = False
#
# plt.figure(figsize=(16,10))
# #high resolution bit closer to the minimum
# #alpha0 = np.linspace(0.001, 0.1, 5)
# #energies0 = []
# #for i in alpha0:
# #    energies0.append(evalenergy(i))
#
# print('Plotting function initialised!')
# energies = []
# alpha = np.linspace(0.1, 0.2, 20)
# for i in alpha:
#     energy_samples = []
#     for j in range(1):
#         energy_samples.append(evalenergy(i))
#     avg = np.average(energy_samples)
#     print('Averaged energy for %f is %f' %(i, avg))
#     energies.append(avg)
# plt.plot(alpha, energies, color='dodgerblue')
# print("--- Total time: %s seconds ---" % (time.time() - start_time))
#
# plt.savefig('Vegas')
#
# #%%
# ##low resolution bit
# #alpha = np.linspace(0.25, 2, 10)
# #energies2 = []
# #for i in alpha:
# #    energies2.append(evalenergy(i))
# #plt.plot(alpha, energies2, color = 'dodgerblue')
# #

# #### Use optimisation to obtain the minimum
# #### and plot a point
# ##OPTIM = False
# ##result = fmin(evalenergy, 0.2, ftol = 0.01, xtol = 0.001, full_output=True)
# ##plt.plot(result[0], result[1], 'ro', ms=5)
# #
# #plt.xlabel('alpha')
# #plt.ylabel('energy')
# #plt.grid()
# #plt.show()
# #print("--- Total time: %s seconds ---" % (time.time() - start_time))
#

##############################################################################
### THE FOLLOWING CODE IS FOR EXTRACTING DATA AND PLOTTING GRAPHS
#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#### To plot a histogram of evaluations
#### VEGAS PLOTS FOR ENERGIES AND HISTOGRAMS
#N = 50
#
#energies = []
#optparas = []
#start_time = time.time()
#print('Loop initialised!')
#
#for i in range(N):
#    print(i)
#    loop_time = time.time()
#    result = fmin(evalenergy, [0.15], ftol = 0.01, xtol = 0.001, full_output=1, disp=0)
#    energies.append(result[1])
#    optparas.append(float(result[0]))
#    print("--- Loop time: %s seconds ---" % (time.time() - loop_time))
#    
#print("--- Total time: %s seconds ---" % (time.time() - start_time))
#
###########
#### Save data into a csv
#n = '1e5'
#name = 'vegas_1param_iters=%s_N=%i' %(n,N)
#np.savetxt('%s.csv' %name, np.transpose([energies,optparas]), 
#           fmt = '%s', delimiter=',', header = (name + '_t=2176'))
#
#
##########
#### Plot histograms
#avg = np.mean(energies)
#std = np.std(energies)
#plt.figure(figsize=(16,10))
#plt.rcParams.update({'font.size': 22})
#
#plt.hist(energies,10)
#plt.xlabel('Energy (Atomic Units)')
#plt.ylabel('N')
#plt.figtext(0.2,0.75,'oneparam, N=%i, \n' 
#            'n=%s, bound=8 \n'
#            'avg E= %.5f \n'
#            'std = %.5f'%(N, n, avg,std))
#plt.title('Histogram of distributed optimised energy values for VEGAS')
##plt.savefig('vegas_energyhist_1param_%s' %n)
#plt.show()
#
#alphaavg = np.mean(optparas)
#alphastd = np.std(optparas)
#plt.figure(figsize=(16,10))
#plt.xlabel('Variational parameter alpha')
#plt.ylabel('N')
#plt.figtext(0.6,0.7,'oneparam, N=%i, \n n=%s, bound=8 \n'
#            'avg alpha = %.5f \n'
#            'std = %.5f'%(N, n, alphaavg,alphastd))
#plt.title('Histogram of distributed optimised parameters for VEGAS')
#plt.hist(optparas,10)
##plt.savefig('vegas_alphahist_1_param_%s' %n)
#plt.show()

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
### VEGAS analysis Code to analyse the Q values
#
#Qexp = np.zeros((50,5))
#Qnorm = np.zeros((50,5))
#
#N = 50 
#n = 5e5
#alpha = [0.15606, 0.15387, 0.15156, 0.15109, 0.15336]
#
#index = 0
#
#qexp = [] 
#qnorm = []
#
#for i in range(N):
#    print(i)
#    loop_time = time.time()
#    evalenergy(alpha[index])
#    qexp.append(final_results['pvalue'][0])
#    qnorm.append(final_results['pvalue'][1])
#    print("--- Loop time: %s seconds ---" % (time.time() - loop_time))
#
#print('The average Qexp value is: ' ,np.average(qexp))
#print('The average Qnorm value is: ' ,np.average(qnorm))
#
##Qexp[:,index] = qexp
##Qnorm[:,index] = qnorm
#
#
###########
#### Save data into a csv
##Qe = np.zeros((50,5))
#
##name = 'vegas_1param_Qnorm' 
##np.savetxt('%s.csv' %name, Qnorm,
##           fmt = '%s', delimiter=',', header = (name + str(iters)))
#
#### Obtain the average
#iters = [1e4, 5e4, 1e5, 5e5, 1e6]
#x = []
#y = []
#for i in range(5):
#    x.append(np.average(Qexp[:,i]))
#    y.append(np.average(Qnorm[:,i]))
#
#plt.scatter(iters, x)
#plt.scatter(iters, y)
#plt.show()
#
#
#### Obtain the standard deviation
#xstd = []
#ystd = []
#for i in range(5):
#    xstd.append(np.std(Qexp[:,i]))
#    ystd.append(np.std(Qnorm[:,i]))
#
#plt.scatter(iters, xstd)
#plt.scatter(iters, ystd)
#plt.show()

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
#### Analysis of the impact of bound sizes
#
#bounds = [1.5, 1.6, 1.7,1.8, 1.9]
#B_energy_2 = np.zeros((50, 5))
#B_alpha_2 = np.zeros((50, 5))
#
#N = 50
#OPTIM = False
#
#for j in range(len(bounds)):
#    bound = bounds[j]
#    benergy = []
#    bopt = []
#    def evalenergy(alpha):
#        @vegas.batchintegrand
#        def expec(x):
#            ''' Expectation value function, used as numerator for variational integral
#            '''
#            r1 = np.array(x[:,0:3])
#            r2 = np.array(x[:,3:])
#    
#            r1_len = np.sqrt(x[:,0]**2 + x[:,1]**2 + x[:,2]**2)
#            r2_len = np.sqrt(x[:,3]**2 + x[:,4]**2 + x[:,5]**2)
#    
#            r1_hat = r1 / np.reshape(r1_len, (-1,1))
#            r2_hat = r2 / np.reshape(r2_len, (-1,1))
#    
#            r12 = np.sqrt((x[:,0]-x[:,3])**2 + (x[:,1]-x[:,4])**2 + (x[:,2]-x[:,5])**2)
#    
#            EL =  (-4 + np.sum((r1_hat-r2_hat)*(r1-r2), 1) / (r12 * (1 + alpha*r12)**2)
#                    - 1/ (r12*(1 + alpha*r12)**3)
#                    - 1/ (4*(1 + alpha*r12)**4)
#                    + 1/ r12 )
#    
#            psisq = (np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
#                    * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2
#            return psisq * EL
#    
#        @vegas.batchintegrand
#        def norm(x):
#            ''' Squared trial wavefunction, used as denominator for variational integral
#            '''
#    
#            r1_len = np.sqrt(x[:,0]**2 + x[:,1]**2 + x[:,2]**2)
#            r2_len = np.sqrt(x[:,3]**2 + x[:,4]**2 + x[:,5]**2)
#    
#            r12 = np.sqrt((x[:,0]-x[:,3])**2 + (x[:,1]-x[:,4])**2 + (x[:,2]-x[:,5])**2)
#    
#            psisq = (np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
#                    * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2
#            return psisq
#    
#        def main():
#            # assign integration volume to integrator
#            dims = 6
#            # creates symmetric bounds specified by [-bound, bound] in dims dimensions
#            symm_bounds = dims * [[-bound,bound]]
#    
#            # simultaneously initialises expectation and normalisation integrals
#            expinteg = vegas.Integrator(symm_bounds)
#            norminteg = vegas.Integrator(symm_bounds)
#    
#            # adapt to the integrands; discard results
#            expinteg(expec, nitn=5, neval=1000)
#            norminteg(norm,nitn=5, neval=1000)
#            # do the final integrals
#            expresult = expinteg(expec, nitn=10, neval=100000)
#            normresult = norminteg(norm, nitn=10, neval=100000)
#    
#                ### obtain numerical result
#                ### Different expressions for plotting/ minimisation algorithm
#            if not OPTIM:
#                    E = expresult.mean/normresult.mean
#            else:
#                E = expresult[0].mean/ normresult[0].mean
#
#            final_results['energy'] = E
#            final_results['dev'] = [expresult.sdev, normresult.sdev]
#            final_results['pvalue'] = [expresult.Q, normresult.Q]
#            print(E)
#            return E
#        E = main()
#        return E
#        
#    for i in range(N):
#        print(j, i)
#        start_time = time.time()
#        result = fmin(evalenergy, [0.15], ftol = 0.01, xtol = 0.001, full_output=1, disp=0)
#        benergy.append(result[1])
#        bopt.append(float(result[0]))
#        print("--- Iteration time: %s seconds ---" % (time.time() - start_time))
#    B_energy_2[:,j] = benergy
#    B_alpha_2[:,j] = bopt
#    print(np.mean(benergy))
#
#name = 'vegas_bound_energy_comp_2'
#np.savetxt('%s.csv' %name, B_energy_2, 
#           fmt = '%s', delimiter=',', header = (name + str(bounds)))
#
#name = 'vegas_bound_ealpha_comp_2'
#np.savetxt('%s.csv' %name, B_alpha_2, 
#           fmt = '%s', delimiter=',', header = (name + str(bounds)))
#
#
#bounds = [1.5, 1.6, 1.7,1.8, 1.9]
#B_energy_mean = np.mean(B_energy_2, 0)
#B_alpha_mean = np.mean(B_alpha_2, 0)
#B_energy_std = np.std(B_energy_2, 0)
#B_alpha_std = np.std(B_alpha_2, 0)
#
#B_processed_2 = np.transpose([bounds, B_energy_mean, B_alpha_mean, B_energy_std, B_alpha_std])
#name = 'vegas_bound_processed_2'
#np.savetxt('%s.csv' %name, B_processed_2, 
#           fmt = '%s', delimiter=',', header = (name +'bound_energy_alpha_estd_astd'))

#%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
# This is for testing with regards to the twopara and 
# threepara case. Since the loop is slow, 

#twopara_energy_new = np.zeros((10,1))
#twopara_energy_backup_2 = twopara_energy.copy()
threepara_energy= np.zeros((N, len(iters)))
#%%
energy = []

#def main(alpha):
##    global alpha0
##    alpha0 = alpha[0]
##    global alpha1
##    alpha1 = alpha[1]
#    if len(alpha) >= 1: 
#        global alpha0
#        alpha0 = alpha[0]
#    if len(alpha) >= 2: 
#        global alpha1
#        alpha1 = alpha[1]
#    if len(alpha) == 3: 
#        global alpha2
#        alpha2 = alpha[2]
#
#    start_time = time.time()
#    # assign integration volume to integrator
#    bound = 8
#    dims = 6
#    # creates symmetric bounds specified by [-bound, bound] in dims dimensions
#    symm_bounds = dims * [[-bound,bound]]
#
#    # simultaneously initialises expectation and normalisation integrals
#    expinteg = vegas.Integrator(symm_bounds)
#    norminteg = vegas.Integrator(symm_bounds)
#
#    # adapt to the integrands; discard results
#    expinteg(expec, nitn=5, neval=1000)
#    norminteg(norm,nitn=5, neval=1000)
#    # do the final integrals
#    expresult = expinteg(expec, nitn=10, neval=iters[indx])
#    normresult = norminteg(norm, nitn=10, neval=iters[indx])
#
#
#    E = expresult.mean/normresult.mean
#
#    print(alpha, E)
#    print("--- Iteration time: %s seconds ---" % (time.time() - start_time))
#    return E
    
def iterplot(data_arr, N, iters, indx):
    s_time = time.time()
    opt =[]  
    def main(alpha):
        if len(alpha) >= 1: 
            global alpha0
            alpha0 = alpha[0]
        if len(alpha) >= 2: 
            global alpha1
            alpha1 = alpha[1]
        if len(alpha) == 3: 
            global alpha2
            alpha2 = alpha[2]
    
        start_time = time.time()
        # assign integration volume to integrator
        bound = 8
        dims = 6
        # creates symmetric bounds specified by [-bound, bound] in dims dimensions
        symm_bounds = dims * [[-bound,bound]]
    
        # simultaneously initialises expectation and normalisation integrals
        expinteg = vegas.Integrator(symm_bounds)
        norminteg = vegas.Integrator(symm_bounds)
    
        # adapt to the integrands; discard results
        expinteg(expec, nitn=5, neval=1000)
        norminteg(norm,nitn=5, neval=1000)
        # do the final integrals
        expresult = expinteg(expec, nitn=10, neval=iters[indx])
        normresult = norminteg(norm, nitn=10, neval=iters[indx])
    
        E = expresult.mean/normresult.mean

        print(alpha, E)
        print("--- Iteration time: %s seconds ---" % (time.time() - start_time))
        return E
        
    for i in range(N):
        print(indx, i)
        repeat_time = time.time()
    #    start_time = time.time()
        result = fmin(main, [2, 1.6, 0.15], ftol = 0.1, xtol = 0.01, full_output=1, disp=1)
        energy.append(result[1])
        print("--- Iteration time: %s seconds ---" % (time.time() - repeat_time))
        
    dur = time.time() - s_time
    print("--- Total time: %s seconds ---" % (dur))
    data_arr[:,indx] = energy
    
    return data_arr

def save_energy(data_arr, name):
    ''' Save data and processed data into a csv
    '''
    np.savetxt('%s.csv' %name, (data_arr), 
           fmt = '%s', delimiter=',', header = (name+ str(iters)))

    ### Process data and save into csv
    eavg = np.mean(data_arr, 0)
    estd = np.std(data_arr, 0)

    np.savetxt('%s.csv' %name+'_processed', (np.transpose([iters, eavg, estd])), 
               fmt = '%s', delimiter=',', header = (name+ 'iters, eavg, estd'))

### Main code
iters = [1e4, 5e4, 1e5, 5e5, 1e6]
N = 10
indx = 2

threepara_energy = iterplot(threepara_energy, N, iters, indx)
#threepara_energy[:,0] = np.loadtxt(name+'.csv', usecols= 0)
#%%
#name = 'vegas_3param_iters'
#save_energy(threepara_energy, name)

#%%

def getpdf(bound, alpha, absv = 0.4135889547408346):
    @vegas.batchintegrand
    def norm(x):
        ''' Squared trial wavefunction, used as denominator for variational integral
        '''
    
        r1_len = np.sqrt(x[:,0]**2 + x[:,1]**2 + x[:,2]**2)
        r2_len = np.sqrt(x[:,3]**2 + x[:,4]**2 + x[:,5]**2)
    
        r12 = np.sqrt((x[:,0]-x[:,3])**2 + (x[:,1]-x[:,4])**2 + (x[:,2]-x[:,5])**2)
    
        psisq = (np.exp(-2 * r1_len)* np.exp(-2 * r2_len)
                * np.exp(r12 / (2 * (1+ alpha*r12)))) ** 2
        return psisq / absv 
    
    dims = 6
    # creates symmetric bounds specified by [-bound, bound] in dims dimensions
    symm_bounds = dims * [[-bound,bound]]
    
    # simultaneously initialises expectation and normalisation integrals
    norminteg = vegas.Integrator(symm_bounds)
    
    # adapt to the integrands; discard results
    norminteg(norm,nitn=5, neval=1000)
    # do the final integrals
    normresult = norminteg(norm, nitn=10, neval=1000000)
    
    norm = normresult.mean
    return norm

#%%
#bound  = np.linspace(1,10, 30)
#prob = [] 
#
#for i in bound:
#    prob.append(getpdf(i, 0.155))
#    
#%%
#plt.ylabel('probability located in 6 dimensional bound')
#plt.xlabel('bound')
#plt.grid()
#plt.plot(bound, prob, label = 'one para with alpha=0.155' )
#plt.legend()

#%%
#np.concatenate((bound, prob))
##x = np.reshape()
#np.savetxt('pdist.csv', np.array([bound, prob]).transpose(), delimiter= ',')
<file_sep>/README.md
## Welcome to quVario

This repository contains our attempt at coding a quantum variational method machine with Python3. The work is part of our Imperial College Physics Year 1 Summer Project. Due to the coronavirus, we were forced to choose a more computational project topic.

quVario is essentially a python ecosystem that can mathematically represent multi-electron (valance) atom system and find the system's ground state energy via the variational method. The integration techniques utilised range from deterministic quadrature methods to monte carlo methods. We can describe our iterative progress through two main prongs: Mark I (deterministic) and Mark II (monte carlo).

As of August 2020, the systems tested are the hydrogen and helium atoms.
### Helium_marki

```markdown

- First attempt at quantum method of variation problem
- Uses Sympy for symbolic differentiation
- Used out of the box SciPy integration numerical methods (numerical differentiation too, for hydrogen test case)

```

### Helium_markii

```markdown

- Monte Carlo progression of Mark I, with ecosystem further developed:
- psiham.py contains class objects that do the generation of the integrals (symbolic differentiation, algebra, etc.) in the form of a numba jit decorated function.
- optipack.py contains the actual integration functions: uniform monte carlo, metropolis algorithm, VEGAS algorithm (to be cited)
- psiham and optipack feeds helium_markii.py which does the actual execution.

```
<file_sep>/helium_markii/psiham.py
#!/usr/bin/env python3

# Made 2020, <NAME>, <NAME>
# mingsongwu [at] outlook [dot] sg
# github.com/starryblack/quVario


### This script provides the hamLet object which handles calls for functions pertaining to hamiltonian generation (supposed to be general, but this case our focus is helium 5 term hamiltonian)

### hamiltonian should serve helium_markii.py which is the higher level script for direct user interface

### This script also provides the psiLet object which handles calls for functions related to creating a trial wavefunction object. It MUST be compatible with hamiltonian such that the objects can interact to give expectation term (quantum chem)

### psifunction should serve helium_markii.py which is the higher level script for direct user interface

#%%%%%%%%%%%%%%%%%%

import math
import sys
import os
import time
from datetime import datetime

import numpy as np
import scipy as sp
import scipy.constants as sc
import sympy as sy
from sympy import conjugate, simplify, lambdify, sqrt
from sympy import *
from IPython.display import display
from scipy import optimize, integrate
from sympy.parsing.sympy_parser import parse_expr
from numba import jit, njit
from sympy.utilities.lambdify import lambdastr

'''
MS's note: I will recycle my previous script for the two sets of electron coordinates and use sympy to symbolically generate the laplacian for each electron. I think we can safe assume our trial functions are analytical indefinitely so finding the laplacian is easy?
everything will be kept to atomic units
'''

class HamLet():
    def __init__(self, trial_expr, **kwargs):
        starttime = time.time()
        print('Hamiltonian Operator object initialised')
        self.coordsys = 'cartesian'

        with PsiLet(**kwargs) as psi:
            self.bases = psi.bases
            self.alphas = psi.alphas
            func_selected = psi.getfunc(trial_expr)
            self.trial = simplify(parse_expr(func_selected()))

            #for convenience
            self.r0 =  parse_expr(psi.r0)
            self.r1 =  parse_expr(psi.r1)
            self.r01 =  parse_expr(psi.r01)
        self.variables = [a for b in self.bases for a in b]


    #custom 'numbafy' protocol. adapted from <NAME> (jankoslavic)
    def numbafy(self, expression, coordinates = None, parameters = None, name='trial_func'):

        if coordinates:
            code_coord = []
            for i, x in enumerate(coordinates):
                code_coord.append(f'{x} = coordinates[{i}]')
            code_coord = '\n    '.join(code_coord)

        if parameters:
            code_param = []
            for i, x in enumerate(parameters):
                code_param.append(f'{x} = parameters[{i}]')
            code_param = '\n    '.join(code_param)

        temp = lambdastr((), expression)
        temp = temp[len('lambda : '):]
        temp = temp.replace('math.exp', 'np.exp')
        code_expression = f'{temp}'

        template = f"""
@njit
def {name}(coordinates, parameters):
    {code_coord}
    {code_param}
    return {code_expression}"""
        print('Numba function template generated! Its name is ', name)
        return template

    def vegafy(self, expression, coordinates = None, name = 'vega_func'):

        if coordinates:
            code_coord = []
            for i, x in enumerate(coordinates):
                code_coord.append(f'{x} = coordinates[:,{i}]')
            code_coord = '\n    '.join(code_coord)

        temp = lambdastr((), expression)
        temp = temp[len('lambda : '):]
        temp = temp.replace('math.exp', 'np.exp')
        code_expression = f'{temp}'

        template = f"""
@vegas.batchintegrand
def {name}(coordinates):
    {code_coord}
    return {code_expression}"""
        print('Vega function template generated! Its name is ', name)
        return template

    def he_getfuncs(self):
        return self.he_elocal(), self.he_norm(), self.he_trial()

    def he_elocal(self):
        operated = self.he_operate(self.trial)
        return (self.variables, operated/self.trial)

    def he_norm(self):
        return (self.variables, self.trial*self.trial)

    def he_expect(self):
        operated = self.he_operate(self.trial)
        return (self.variables, operated*self.trial)

    def he_trial(self):
        return (self.variables, self.trial)

    def he_operate(self, expr):
        lap1 = -0.5*self.laplace(expr, 0)
        lap2 = -0.5*self.laplace(expr, 1)
        attract = -2*(1/self.r0 + 1/self.r1)*expr
        repel = (1/self.r01)*expr
        return lap1 + lap2 + attract + repel


    def laplace(self, expr, index, coordsys = 'c'):
        if coordsys == 'cartesian' or coordsys == 'c':
            x = self.bases[index][0]
            y = self.bases[index][1]
            z = self.bases[index][2]
            grad = [sy.diff(expr, x), sy.diff(expr, y), sy.diff(expr, z)]
            lap = sy.diff(grad[0], x) + sy.diff(grad[1], y) + sy.diff(grad[2], z)


        elif coord_sys == 'spherical' or coordsys == 's':
            r = self.bases[index][0]
            t = self.bases[index][1]
            p = self.bases[index][2]
            grad = [sy.diff(expr, r), 1/r * sy.diff(expr, t), 1/(r*sin(t)) * sy.diff(expr, p)]
            lap = (1/r**2)*sy.diff(r**2 * grad[0], r) + (1/(r*sin(t)))*(sy.diff(grad[1] * sin(t), t) + sy.diff(grad[2], p))

        return lap

    def __enter__(self):
        return self

    def __exit__(self, e_type, e_val, traceback):
        print('HamLet object self-destructing...')



class PsiLet():
    def __init__(self, electrons = 2, alphas = 1, coordsys = 'cartesian'):
        print('psiLet object initialised.')
        self.spawn_electrons(electrons, coord_sys = coordsys)
        print('{} {} electron coordinate set(s) spawned'.format(electrons, coordsys))
        self.spawn_alphas(alphas)
        print('{} alpha parameter(s) spawned'.format(alphas))

        self.r0 = '(x0**2 + y0**2 + z0**2)**0.5'
        self.r1 = '(x1**2 + y1**2 + z1**2)**0.5'
        self.r01 = '((x0-x1)**2 + (y0-y1)**2 + (z0-z1)**2)**0.5'
        self.dot01 = 'x0*x1 + y0*y1 + z0*z1'


    def spawn_electrons(self, number, coord_sys = 'cartesian'):
        self.bases = []
        self.basis_names = {
        'cartesian': ['x', 'y', 'z'],
        'spherical': ['r', 'theta', 'phi'],
        'cylindrical': ['rho', 'phi', 'z']
        }

        for i in range(number):
            temp = []
            for q in self.basis_names[coord_sys]:
                temp.append(sy.symbols(q + str(i)))
            self.bases.append(temp)

    def spawn_alphas(self, number):
        self.alphas = []
        for i in range(number):
            self.alphas.append(sy.symbols('alpha' + str(i)))

    def getfunc(self, function):
        self.func_dict = {

        'onepara1': self.onepara1,
        'twopara1': self.twopara1,
        'threepara1': self.threepara1,
        'threepara2': self.threepara2

        }
        return self.func_dict[function]


    def onepara1(self):
        expr = f'''exp(-2*{self.r0}-2*{self.r1}+{self.r01}/(2*(1+alpha0*{self.r01})))'''
        return expr

    def twopara1(self):
        expr = f'''exp(-alpha0*({self.r0}+{self.r1})+{self.r01}/(2*(1+alpha1*{self.r01})))'''
        return expr

    def threepara1(self):
        expr = f'''(exp(-alpha0*{self.r0}-alpha1*{self.r1}) + exp(-alpha1*{self.r0}-alpha0*{self.r1}))*(1 + alpha2*{self.r01})'''
        return expr

    def threepara2(self):
        expr = f'''(exp(-alpha0*{self.r0}-alpha1*{self.r1}) + exp(-alpha1*{self.r0}-alpha0*{self.r1}))*exp({self.r01}/(2*(1 + alpha2*{self.r01})))'''
        return expr


    ### this function is flawed, gives numerically incorrect result. Use with caution
    def manual2(self):

        expr = f'''-4 + ({self.r0} + {self.r1})*(1 - {self.dot01}/({self.r0}*{self.r1}))/({self.r01}*(1+alpha0*{self.r01})**2) - 1/({self.r01}*(1+alpha0*{self.r01})**3) - 1/(4*(1+alpha0*{self.r01})**4) + 1/{self.r01}'''
        return expr

    def __enter__(self):
        return self

    def __exit__(self, e_type, e_val, traceback):
        print('psiLet object self-destructing')
<file_sep>/helium_markii/helium_markii.py
#!/usr/bin/env python3

# Made 2020, <NAME>, <NAME>
# mingsongwu [at] outlook [dot] sg
# github.com/starryblack/quVario

import math
import time
import numpy as np
from numba import njit, prange
import vegas

import matplotlib.pyplot as plt

from optipack import MiniMiss, LasVegas
from optipack import integrator_mcmc, metropolis_hastings
from optipack import Uint, Ueval
import psiham

import matplotlib.pyplot as plt

class Noble():

    def __init__(self, evals = 10000, iters = 40, verbose = True):
        print('Noble initialising')
        self.dims = 6
        self.mini = MiniMiss()
        self.vegas = LasVegas()
        psilet_args = {'electrons': 2, 'alphas': 3, 'coordsys': 'cartesian'}
        self.ham = psiham.HamLet(trial_expr = 'threepara2', **psilet_args)
        self.verbose = verbose

        self.numba_he_elocal()
        self.numba_pdf()
        self.numba_expec()
        self.vega_expec()
        self.vega_pdf()

        self.evals = evals
        self.iters = iters


        # data = self.minimise(function = self.monte_uniform, guess = np.array([2.001, 2.0, 0.2]))
        # data = self.minimise(function = self.mcmc_metro, guess = np.array([2.001, 2.0, 0.2]))

        # mcmc_args = [pfunc, qfunc , self.evals, self.iters, 6]
        # data = jit_grad_descent(integrator_mcmc, guess = np.array([1.9, 2.0, 0.2]), tolerance = 0.0005)
        # print(data)

        # data = self.grad_descent(self.monte_vegas, np.array([1.9, 2.0, 0.2]))
        # print(data)
        # for i in range(100):
        #     self.temp.append(self.monte_vegas(np.array([1.9, 2.0, 0.2])))


        # for i in range(100):
        #     self.temp.append(self.mcmc_metro(np.array([1.9, 2.0, 0.2])))

        # print(data[3])


        # self.custom_log(data, comments = 'Uniform integrator, 2 para')

    def numba_he_elocal(self):
        variables, expr = self.ham.he_elocal()
        temp = self.ham.numbafy(expr, coordinates = variables, parameters = self.ham.alphas, name = 'qfunc')
        exec(temp, globals())

    def numba_pdf(self):
        variables, expr = self.ham.he_norm()
        temp = self.ham.numbafy(expr, coordinates = variables, parameters = self.ham.alphas, name = 'pfunc')
        exec(temp, globals())

    def numba_expec(self):
        variables, expr = self.ham.he_expect()
        temp = self.ham.numbafy(expr, coordinates = self.ham.variables, parameters = self.ham.alphas, name = 'expec')
        exec(temp, globals())

    def vega_expec(self):
        variables, expr = self.ham.he_expect()
        temp = self.ham.vegafy(expr, coordinates = self.ham.variables, name = 'vega_expec')
        exec(temp, globals())

    def vega_pdf(self):
        variables, expr = self.ham.he_norm()
        temp = self.ham.vegafy(expr, coordinates = self.ham.variables, name = 'vega_pfunc')
        exec(temp, globals())

    def monte_metro(self, args):
        temp = integrator_mcmc(pfunc, qfunc , self.evals, self.iters, alpha = args, dims = self.dims, verbose = self.verbose)
        self.std = temp[1]
        self.reject_ratio = temp[2]
        return temp[0]

    def monte_uniform(self, args):
        temp = Ueval(pfunc, expec, self.evals, self.iters, alpha = args, dimensions = self.dims)
        self.std = temp[1]
        self.reject_ratio = 0.0
        return temp[0]

    def monte_vegas(self, alpha):
        global alpha0
        global alpha1
        global alpha2
        alpha0 = alpha[0]
        alpha1 = alpha[1]
        alpha2 = alpha[2]

        temp = self.vegas.vegas_int(vega_expec, vega_pfunc, self.evals, self.iters, dimensions = self.dims, volumespan = 8., verbose = self.verbose)
        return temp['energy']

    def nead_min(self, function, guess):
        temp = self.mini.neldermead(func = function, guess = guess, ftol = 0.01)
        # temp = grad_des(func = function, guess = guess, tolerance = 0.0001)
        return temp

    def grad_descent(self, function, guess):
        temp = self.mini.gradient(func = function, guess = guess, tolerance = 0.05)
        return temp

    def stability_protocol(self, fixedpoint, iterations):
        temp_master = [[],[],[]]
        functions = [self.monte_metro, self.monte_uniform, self.monte_vegas]
        for index, func in enumerate(functions):
            start = time.time()
            for i in range(iterations):
                temp_master[index].append(func(guess))
            print('Function #{} done!'.format(index))
        return temp_master


    def custom_log(self, data = [], comments = None):
        log = open('markii_log.txt', 'a+')
        separator = '\n{}\n'.format(''.join('#' for i in range(10)))
        info = separator + 'Datetime = {}\n'.format(data[0]) + 'Time taken (s) = {:.4f}\n'.format(data[1]) + 'Initial parameter guess (atomic untis) = {}\n'.format(data[2]) + 'Optimised parameter (atomic units) = {}\n'.format(data[3]) + 'Optimised function value (atomic units) = {:.4f} with std = {:.4f}\n'.format(data[4], self.std) + 'Iterations = {}, Evaluations/iter = {} with rejection ratio = {:.4f}\n'.format(self.iters, self.evals, self.reject_ratio) + 'Comments = ({})\n'.format(comments)
        print(info)
        log.write(info)
        log.close()


    def __enter__(self):
        return self

    def __exit__(self, e_type, e_val, traceback):
        print('Noble object self-destructing')

@njit
def jit_grad_des(func, guess, tolerance):

    cycle_count = 0
    position0 = guess
    ep = 0.00000001
    step = 0.2
    step_ratio = 0.5
    epsilons = np.identity(len(guess))*ep
    delta1 = 10.
    satisfied = False
    point_collection = []


    while satisfied == False:
        cycle_count+=1
        print('Cycle', cycle_count)
        value1 = func(pfunc, qfunc, sample_iter = 10000, walkers = 40, alpha = position0, dims = 6)[0]

        vector = np.zeros(len(guess))
        for i in range(len(guess)):
            vector[i] = (func(pfunc, qfunc, sample_iter = 10000, walkers = 40, alpha = position0 + epsilons[i], dims = 6)[0] - func(pfunc, qfunc, sample_iter = 10000, walkers = 40, alpha = position0, dims = 6)[0])/ep

        vectornorm = vector/(np.linalg.norm(vector))
        print(vectornorm)
        position0 += -vectornorm * step
        value2 = func(pfunc, qfunc, sample_iter = 10000, walkers = 40, alpha = position0, dims = 6)[0]
        delta1 = value2 - value1

        positionprime = position0 + vectornorm*(step*step_ratio)
        value3 = func(pfunc, qfunc, sample_iter = 10000, walkers = 40, alpha = positionprime, dims = 6)[0]
        delta2 = value3 - value1
        if delta2 < delta1:
            print('shrink!')
            position0 = positionprime
            step = step*step_ratio
            delta1 = delta2
            value2 = value3

        point_collection.append(value2)
        if convergence == 'strict':
            satisfied = abs(delta1) < tolerance
            finalvalue = value2

        elif convergence == 'fuzzy':
            if len(point_collection) >= 5:
                data_set = point_collection[-5:]
                print('std', statistics.pstdev(data_set))
                print('mean', statistics.mean(data_set))
                satisfied = abs(statistics.pstdev(data_set)/statistics.mean(data_set)) < tolerance
                finalvalue = statistics.mean(data_set)

    print('Convergence sucess!')
    return finalvalue, position0



n = Noble(evals = 10000, iters = 10, verbose = False)



evals = [10000, 50000, 100000, 500000, 1000000]
paras = [2,1]
guesses = [np.array([2.0,0.2]),  np.array([0.2])]
trials = ['twopara1', 'onepara1']

for x, para in enumerate(paras):
    data = []
    for i in evals:
        n.evals = i
        psilet_args = {'electrons': 2, 'alphas': para, 'coordsys': 'cartesian'}
        n.ham = psiham.HamLet(trial_expr = trials[x], **psilet_args)
        n.numba_pdf()
        n.numba_expec()
        e_min = n.nead_min(n.monte_uniform, guesses[x])[3][1]
        data.append(np.array([e_min, n.std]))
        print(data)
        print('{} and {}'.format(para,i), 'done')
    np.savetxt('__results__/processed/uniform_test/{}para_{}iterations_10avg.txt'.format(para, i), np.c_[np.array(data)])




# # n.stability_protocol(np.array([1.9, 2.0, 0.2]), 100)
# temp =[]
# for i in range(100):
#     start = time.time()
#     result = n.grad_descent(n.monte_metro, np.array([1.6, 2.0, 0.2]))
#     temp.append(result[0])
#     print(time.time() - start)





#%%%%%%%%%%%%%

# print(np.mean(temp))
# plt.hist(temp, bins = 15)

# iterations = 100
# np.savetxt('__results__/stability_study/mcmc_{}iterations.txt'.format(iterations), np.c_[np.array(temp0)])
# np.savetxt('__results__/stability_study/uniform_{}iterations.txt'.format(iterations), np.c_[np.array(temp1)])
# np.savetxt('__results__/stability_study/vegas.txt_{}iterations.txt'.format(iterations), np.c_[np.array(temp2)])
#
#
# plt.hist(temp0, bins = 10, histtype=u'step')
# plt.hist(temp1, bins = 10, histtype=u'step')
# plt.hist(temp2, bins = 10, histtype=u'step')
# plt.legend(['Metro', 'Uniform', 'Vegas'])
#
#
# plt.show()

#%%%%%%%%%%%
