<reponame>yannra/netket<filename>scripts/hydrogen_chain_supervised_learning.py
import numpy as np
import netket as nk
import sys
import scipy.optimize as spo
import netket.custom.utils as utls

from netket.utils import (
    MPI_comm as _MPI_comm,
    n_nodes as _n_nodes,
    node_number as _rank
)

from netket.stats import (
    statistics as _statistics,
    mean as _mean,
    sum_inplace as _sum_inplace,
)

import netket.custom.fermionic_hilbert as fermhi

import gpw.utils as utls

import numpy_ml as ml

import netket.custom.hubbard as hub

from netket.custom.fermionic_hopping_sampler import FermionicHoppingSampler,FermionicHoppingKernel

from netket.custom.ab_initio_ham import AbInitio

import numpy as np
import netket as nk
import sys
import scipy.optimize as spo
import netket.custom.utils as utls

from netket.utils import (
    MPI_comm as _MPI_comm,
    n_nodes as _n_nodes,
    node_number as _rank
)

import mpi4py.MPI as mpi

from netket.stats import (
    statistics as _statistics,
    mean as _mean,
    sum_inplace as _sum_inplace,
)

import netket.custom.fermionic_hilbert as fermhi

import gpw.utils as utls

import numpy_ml as ml

import netket.custom.hubbard as hub

from netket.custom.fermionic_hopping_sampler import FermionicHoppingSampler,FermionicHoppingKernel

from netket.custom.ab_initio_ham import AbInitio

from pyscf import scf, gto, ao2mo, fci

import scipy.optimize as spo


from pyscf import scf, gto, ao2mo, fci, lo
from pyscf.tools import ring

N = int(sys.argv[1])
dist = float(sys.argv[2])
basis_rep = int(sys.argv[3]) # 0: local basis, 1: canonical basis
loss_fun = int(sys.argv[4]) # 0: overlap, 1: squared error

mol = gto.Mole()

mol.build(
    atom = [('H', x) for x in ring.make(12, dist)],
    basis = 'sto-3g',
    symmetry = True,
    unit="Bohr"
)

nelec = mol.nelectron
print('Number of electrons: ',nelec)

myhf = scf.RHF(mol)
ehf = myhf.scf()
norb = myhf.mo_coeff.shape[1]
print('Number of molecular orbitals: ',norb)

# Get hamiltonian elements
# 1-electron 'core' hamiltonian terms, transformed into MO basis
h1 = np.linalg.multi_dot((myhf.mo_coeff.T, myhf.get_hcore(), myhf.mo_coeff))

# Get 2-electron electron repulsion integrals, transformed into MO basis
eri = ao2mo.incore.general(myhf._eri, (myhf.mo_coeff,)*4, compact=False)

# Previous representation exploited permutational symmetry in storage. Change this to a 4D array.
# Integrals now stored as h2[p,q,r,s] = (pq|rs) = <pr|qs>. Note 8-fold permutational symmetry.
h2 = ao2mo.restore(1, eri, norb)
print('Shape of eri array: ',h2.shape)

if basis_rep == 1:
    # Transform to a local orbital basis
    loc_coeff = lo.orth_ao(mol, 'meta_lowdin')
    ovlp = myhf.get_ovlp()
    # Check that we still have an orthonormal basis, i.e. C^T S C should be the identity
    assert(np.allclose(np.linalg.multi_dot((loc_coeff.T, ovlp, loc_coeff)),np.eye(norb)))
    # Find the hamiltonian in the local basis
    hij_local = np.linalg.multi_dot((loc_coeff.T, myhf.get_hcore(), loc_coeff))
    hijkl_local = ao2mo.restore(1, ao2mo.kernel(mol, loc_coeff), norb)
    h1 = hij_local
    h2 = hijkl_local


# Run FCI
cisolver = fci.direct_spin1.FCISolver(mol)
e, c = cisolver.kernel(h1, h2, norb, nelec)

print('Size of FCI wave function: ',c.shape)
print('E(FCI) = %.12f' % ((e + mol.energy_nuc())/12)) # Add nuclear repulsion energy


nuc_en = mol.energy_nuc()

all_configs = []
amps = []

for i in range(c.shape[0]):
    for j in range(c.shape[1]):
        alpha = np.array(utls.bitfield(fci.cistring.addr2str(norb, nelec//2, i), norb))
        beta = np.array(utls.bitfield(fci.cistring.addr2str(norb, nelec//2, j), norb))
        all_configs.append(alpha + 2*beta)
        amps.append(c[i,j])

all_configs = np.array(all_configs)
amps = np.array(amps)

hi = fermhi.Fermions(norb, n_elec=(nelec//2,nelec//2))

ha = AbInitio(hi, h1, h2)

ma = nk.machine.qgps.QGPSFermionicProdSym(hi, n_bond=N, automorphisms=None, spin_flip_sym=False, dtype=complex)
ma._exp_kern_representation = False
ma._fast_update = False
ma.init_random_parameters(sigma=0.1, start_from_uniform=True)

rng_par = np.random.default_rng(1234)
ma._epsilon = rng_par.normal(scale=0.1, size=ma._epsilon.shape) + 1.j*rng_par.normal(scale=0.5, size=ma._epsilon.shape)
ma._opt_params = ma._epsilon[ma._der_ids >= 0].copy()


amps = amps/np.exp(np.mean(np.log(amps + 0.j)))

indices = np.array_split(np.arange(len(amps)), mpi.COMM_WORLD.size)[_rank]
basis = all_configs[indices]
amplitudes = amps[indices]
weightings = np.ones(len(amplitudes))

learning = nk.custom.SupervisedLearning(ma)

def objective_fun(parameters):
    par = parameters.reshape((2,-1))
    ma.parameters = par[0,:] + 1.j*par[1,:]
    if loss_fun == 0:
        return -learning.overlap(basis, amplitudes, weightings)
    else:
        return learning.mean_squared_error(basis, amplitudes, weightings)


def grad(parameters):
    par = parameters.reshape((2,-1))
    ma.parameters = par[0,:] + 1.j*par[1,:]
    if loss_fun == 0:
        return -learning.overlap_der(basis, amplitudes, weightings)
    else:
        return learning.mean_squared_error_der(basis, amplitudes, weightings)


if mpi.COMM_WORLD.Get_rank() == 0:
    with open("out.txt", "w") as fl:
        fl.write("")

count = 0
steps =100
def callback(parameters):
    global count
    if count % steps == 0:
        val = objective_fun(parameters)
        if loss_fun == 0:
            val = -val
        c_predicted_full = np.zeros(c.size, dtype = complex)
        c_predicted_full.fill(0.0)
        c_predicted_full[indices] = np.exp(ma.log_val(basis))
        c_predicted_full = _sum_inplace(c_predicted_full)
        if _rank == 0:
            c_predicted_full /= np.sqrt(np.sum(abs(c_predicted_full)**2))
            c_predicted = c_predicted_full.reshape(c.shape)
            en = cisolver.energy(h1, h2, c_predicted.real, norb, nelec) + cisolver.energy(h1, h2, c_predicted.imag, norb, nelec)
            print(val, en)
            with open("out.txt", "a") as fl:
                fl.write("{}  {}\n".format(en, val))
    count += 1

par = np.array([ma.parameters.real, ma.parameters.imag]).flatten()

opt = ml.neural_nets.optimizers.Adam()

no_improvement_count = 0
best_loss = objective_fun(par)

while no_improvement_count < 100:
    gradient = grad(par)
    par = opt.update(par, gradient, "pars")
    callback(par)
    loss = objective_fun(par)
    if loss < best_loss:
        best_loss = loss
        no_improvement_count = 0
    else:
        no_improvement_count += 1

if _rank == 0:
    print("Starting min with linesearch", flush=True)

steps = 1
minimization = spo.minimize(objective_fun, par, jac=grad, callback=callback, method="CG")

if _rank == 0:
    print(minimization)

count = 0
callback(minimization.x)
import numpy as np
import netket as nk
import sys
import scipy.optimize as spo
import netket.custom.utils as utls

from netket.utils import (
    MPI_comm as _MPI_comm,
    n_nodes as _n_nodes,
    node_number as _rank
)

from netket.stats import (
    statistics as _statistics,
    mean as _mean,
    sum_inplace as _sum_inplace,
)

import netket.custom.fermionic_hilbert as fermhi

import gpw.utils as utls

import numpy_ml as ml

import netket.custom.hubbard as hub

from netket.custom.fermionic_hopping_sampler import FermionicHoppingSampler,FermionicHoppingKernel

from netket.custom.ab_initio_ham import AbInitio

import numpy as np
import netket as nk
import sys
import scipy.optimize as spo
import netket.custom.utils as utls

from netket.utils import (
    MPI_comm as _MPI_comm,
    n_nodes as _n_nodes,
    node_number as _rank
)

import mpi4py.MPI as mpi

from netket.stats import (
    statistics as _statistics,
    mean as _mean,
    sum_inplace as _sum_inplace,
)

import netket.custom.fermionic_hilbert as fermhi

import gpw.utils as utls

import numpy_ml as ml

import netket.custom.hubbard as hub

from netket.custom.fermionic_hopping_sampler import FermionicHoppingSampler,FermionicHoppingKernel

from netket.custom.ab_initio_ham import AbInitio

from pyscf import scf, gto, ao2mo, fci

import scipy.optimize as spo


from pyscf import scf, gto, ao2mo, fci, lo
from pyscf.tools import ring

N = int(sys.argv[1])
dist = float(sys.argv[2])
basis_rep = int(sys.argv[3]) # 0: local basis, 1: canonical basis
loss_fun = int(sys.argv[4]) # 0: overlap, 1: squared error

mol = gto.Mole()

mol.build(
    atom = [('H', x) for x in ring.make(12, dist)],
    basis = 'sto-3g',
    symmetry = True,
    unit="Bohr"
)

nelec = mol.nelectron
print('Number of electrons: ',nelec)

myhf = scf.RHF(mol)
ehf = myhf.scf()
norb = myhf.mo_coeff.shape[1]
print('Number of molecular orbitals: ',norb)

# Get hamiltonian elements
# 1-electron 'core' hamiltonian terms, transformed into MO basis
h1 = np.linalg.multi_dot((myhf.mo_coeff.T, myhf.get_hcore(), myhf.mo_coeff))

# Get 2-electron electron repulsion integrals, transformed into MO basis
eri = ao2mo.incore.general(myhf._eri, (myhf.mo_coeff,)*4, compact=False)

# Previous representation exploited permutational symmetry in storage. Change this to a 4D array.
# Integrals now stored as h2[p,q,r,s] = (pq|rs) = <pr|qs>. Note 8-fold permutational symmetry.
h2 = ao2mo.restore(1, eri, norb)
print('Shape of eri array: ',h2.shape)

if basis_rep == 1:
    # Transform to a local orbital basis
    loc_coeff = lo.orth_ao(mol, 'meta_lowdin')
    ovlp = myhf.get_ovlp()
    # Check that we still have an orthonormal basis, i.e. C^T S C should be the identity
    assert(np.allclose(np.linalg.multi_dot((loc_coeff.T, ovlp, loc_coeff)),np.eye(norb)))
    # Find the hamiltonian in the local basis
    hij_local = np.linalg.multi_dot((loc_coeff.T, myhf.get_hcore(), loc_coeff))
    hijkl_local = ao2mo.restore(1, ao2mo.kernel(mol, loc_coeff), norb)
    h1 = hij_local
    h2 = hijkl_local


# Run FCI
cisolver = fci.direct_spin1.FCISolver(mol)
e, c = cisolver.kernel(h1, h2, norb, nelec)

print('Size of FCI wave function: ',c.shape)
print('E(FCI) = %.12f' % ((e + mol.energy_nuc())/12)) # Add nuclear repulsion energy


nuc_en = mol.energy_nuc()

all_configs = []
amps = []

for i in range(c.shape[0]):
    for j in range(c.shape[1]):
        alpha = np.array(utls.bitfield(fci.cistring.addr2str(norb, nelec//2, i), norb))
        beta = np.array(utls.bitfield(fci.cistring.addr2str(norb, nelec//2, j), norb))
        all_configs.append(alpha + 2*beta)
        amps.append(c[i,j])

all_configs = np.array(all_configs)
amps = np.array(amps)

hi = fermhi.Fermions(norb, n_elec=(nelec//2,nelec//2))

ha = AbInitio(hi, h1, h2)

ma = nk.machine.qgps.QGPSFermionicProdSym(hi, n_bond=N, automorphisms=None, spin_flip_sym=False, dtype=complex)
ma._exp_kern_representation = False
ma._fast_update = False
ma.init_random_parameters(sigma=0.1, start_from_uniform=True)

rng_par = np.random.default_rng(1234)
ma._epsilon = rng_par.normal(scale=0.1, size=ma._epsilon.shape) + 1.j*rng_par.normal(scale=0.5, size=ma._epsilon.shape)
ma._opt_params = ma._epsilon[ma._der_ids >= 0].copy()


amps = amps/np.exp(np.mean(np.log(amps + 0.j)))

indices = np.array_split(np.arange(len(amps)), mpi.COMM_WORLD.size)[_rank]
basis = all_configs[indices]
amplitudes = amps[indices]
weightings = np.ones(len(amplitudes))

learning = nk.custom.SupervisedLearning(ma)

def objective_fun(parameters):
    par = parameters.reshape((2,-1))
    ma.parameters = par[0,:] + 1.j*par[1,:]
    if loss_fun == 0:
        return -learning.overlap(basis, amplitudes, weightings)
    else:
        return learning.mean_squared_error(basis, amplitudes, weightings)


def grad(parameters):
    par = parameters.reshape((2,-1))
    ma.parameters = par[0,:] + 1.j*par[1,:]
    if loss_fun == 0:
        return -learning.overlap_der(basis, amplitudes, weightings)
    else:
        return learning.mean_squared_error_der(basis, amplitudes, weightings)


if mpi.COMM_WORLD.Get_rank() == 0:
    with open("out.txt", "w") as fl:
        fl.write("")

count = 0
steps =100
def callback(parameters):
    global count
    if count % steps == 0:
        val = objective_fun(parameters)
        if loss_fun == 0:
            val = -val
        c_predicted_full = np.zeros(c.size, dtype = complex)
        c_predicted_full.fill(0.0)
        c_predicted_full[indices] = np.exp(ma.log_val(basis))
        c_predicted_full = _sum_inplace(c_predicted_full)
        if _rank == 0:
            c_predicted_full /= np.sqrt(np.sum(abs(c_predicted_full)**2))
            c_predicted = c_predicted_full.reshape(c.shape)
            en = cisolver.energy(h1, h2, c_predicted.real, norb, nelec) + cisolver.energy(h1, h2, c_predicted.imag, norb, nelec)
            print(val, en)
            with open("out.txt", "a") as fl:
                fl.write("{}  {}\n".format(en, val))
    count += 1

par = np.array([ma.parameters.real, ma.parameters.imag]).flatten()

opt = ml.neural_nets.optimizers.Adam()

no_improvement_count = 0
best_loss = objective_fun(par)

while no_improvement_count < 100:
    gradient = grad(par)
    par = opt.update(par, gradient, "pars")
    callback(par)
    loss = objective_fun(par)
    if loss < best_loss:
        best_loss = loss
        no_improvement_count = 0
    else:
        no_improvement_count += 1

if _rank == 0:
    print("Starting min with linesearch", flush=True)

steps = 1
minimization = spo.minimize(objective_fun, par, jac=grad, callback=callback, method="CG")

if _rank == 0:
    print(minimization)

count = 0
callback(minimization.x)
# 0: local basis, 1: canonical basis # 0: overlap, 1: squared error # Get hamiltonian elements # 1-electron 'core' hamiltonian terms, transformed into MO basis # Get 2-electron electron repulsion integrals, transformed into MO basis # Previous representation exploited permutational symmetry in storage. Change this to a 4D array. # Integrals now stored as h2[p,q,r,s] = (pq|rs) = <pr|qs>. Note 8-fold permutational symmetry. # Transform to a local orbital basis # Check that we still have an orthonormal basis, i.e. C^T S C should be the identity # Find the hamiltonian in the local basis # Run FCI # Add nuclear repulsion energy