<gh_stars>1-10
import time
import numpy

import matplotlib.pyplot as plt
from pyscf import gto, scf
import scipy
from scipy.optimize import minimize

import jax.numpy as jnp
from jax import grad, jit, random

from jax.config import config
config.update("jax_enable_x64", True)

import adscf

key = random.PRNGKey(0)

x = []
y = []

x_aug = []
y_aug = []

x_scf = []
y_scf = []

for i in range(5, 31):
    R = 0.1 * i
    print(f"interatomic distance: {R:.2f}")

    mol = gto.Mole()
    mol.charge = 0
    mol.spin = 0

    mol.build(atom = f'H 0.0 0.0 0.0; F 0.0 0.0 {R:.2f}',
              basis ='STO-3G', unit='Angstrom')

    calcEnergy, gradEnergy = adscf.calcEnergy_create(mol)

    start = time.time()

    # RHF energy calculation by PySCF
    mf = scf.RHF(mol)
    mf.scf()
    elapsed_time = time.time() - start
    print ("SCF: {:.3f} ms".format(elapsed_time * 1000))
    e_scf = scf.hf.energy_tot(mf)
    x_scf.append(R)
    y_scf.append(e_scf)

    # Curvilinear search using Cayley transformation
    start = time.time()

    # parameters
    tau = 1.0
    tau_m = 1e-10
    tau_M = 1e10
    rho = 1e-4
    delta = 0.1
    eta = 0.5
    epsilon = 1e-6
    max_iter = 5000

    # 1. initialize X0
    S = mol.intor_symmetric('int1e_ovlp')  # overlap matrix
    S64 = numpy.asarray(S, dtype=numpy.float64)
    X_np = scipy.linalg.inv(scipy.linalg.sqrtm(S64))
    X = jnp.asarray(X_np)
    
    # 2. set C=f(X0) and Q0=1
    C = calcEnergy(X)
    Q = 1.0

    # 3. calculate G0 and A0
    G = gradEnergy(X)
    A = G @ X.T @ S - S @ X @ G.T

    # function to calculate Y(tau)
    I = jnp.identity(len(S))
    def Y_tau(tau, X, A):
        return jnp.linalg.inv(I + 0.5 * tau * A @ S) @ (I - 0.5 * tau * A @ S) @ X
    
    # main loop
    for k in range(max_iter):
        Y = Y_tau(tau, X, A)
        A_norm = jnp.linalg.norm(A, "fro")
        X_old, Q_old, G_old = X, Q, G
    
        # 5
        while calcEnergy(Y) > C - rho * tau * A_norm**2.0:
            tau *= delta    # 6
            Y = Y_tau(tau, X, A)
    
        # 8
        X_new = Y
        Q_new = eta * Q + 1.0
        C = (eta * Q * C + calcEnergy(X_new)) / Q_new

        # 9
        G_new = gradEnergy(X_new)
        A_new = G_new @ X_new.T @ S - S @ X_new @ G_new.T

        # 10
        Sk = X_new - X
        Yk = G_new - G
        if k % 2 == 0:
            tau_k = jnp.trace(Sk.T @ Sk) / abs(jnp.trace(Sk.T @ Yk))
        else:
            tau_k = abs(jnp.trace(Sk.T @ Yk)) / jnp.trace(Yk.T @ Yk)
        tau = max(min(tau_k, tau_M), tau_m)

        # Update variables for next iteration
        X, Q, G, A = X_new, Q_new, G_new, A_new

        # Check loop condition (4)
        cond = jnp.linalg.norm(A @ X)
        if cond < epsilon:
            break

    elapsed_time = time.time() - start
    print ("Curvilinear search: {:.3f} ms".format(elapsed_time*1000))
    e = calcEnergy(X)+mol.energy_nuc()
    print(f"total energy = {e}")
    x.append(R)
    y.append(e)

    # augmented Lagrangian
    @jit
    def orthogonality(x):
        C = jnp.reshape(x, [len(S), len(S)])
        return jnp.linalg.norm(C.transpose()@S@C - jnp.identity(len(S)))

    start = time.time()
    x0 = random.uniform(key, (S.size,))

    # 1
    mu = 1.0
    lam = 0.0

    constraint = orthogonality(x0)

    # 2
    while constraint > 1e-6:
        def target(x):
            h = orthogonality(x)
            return calcEnergy(x) + mu * h ** 2.0 + lam * h

        # 3
        res = minimize(jit(target), x0, jac=jit(grad(jit(target))), method="BFGS", options={'maxiter': 100})
        x0 = res.x
        constraint = orthogonality(x0)
        # 4
        lam += 2.0 * mu * constraint
        # 5
        mu *= 2.0
    elapsed_time = time.time() - start
    print ("Augmented: {:.3f} s".format(elapsed_time*1000))
    energy = res.fun+mol.energy_nuc()
    print(f"calculated energy = {energy}")
    x_aug.append(R)
    y_aug.append(energy)

p2 = plt.plot(x_scf, y_scf, marker="o")
p1 = plt.plot(x_aug, y_aug, marker="*")
p0 = plt.plot(x, y, marker="x")
plt.xlabel("interatomic distance (Å)", fontsize=16)
plt.ylabel("total energy (Eh)", fontsize=16)
plt.legend((p0[0], p1[0], p2[0]),
           ("Curvilinear search", "Augmented Lagrangian", "PySCF"),
           loc='upper right')
plt.savefig("result.png", dpi=300)

import time
import numpy

import matplotlib.pyplot as plt
from pyscf import gto, scf
import scipy
from scipy.optimize import minimize

import jax.numpy as jnp
from jax import grad, jit, random

from jax.config import config
config.update("jax_enable_x64", True)

import adscf

key = random.PRNGKey(0)

x = []
y = []

x_aug = []
y_aug = []

x_scf = []
y_scf = []

for i in range(5, 31):
    R = 0.1 * i
    print(f"interatomic distance: {R:.2f}")

    mol = gto.Mole()
    mol.charge = 0
    mol.spin = 0

    mol.build(atom = f'H 0.0 0.0 0.0; F 0.0 0.0 {R:.2f}',
              basis ='STO-3G', unit='Angstrom')

    calcEnergy, gradEnergy = adscf.calcEnergy_create(mol)

    start = time.time()

    # RHF energy calculation by PySCF
    mf = scf.RHF(mol)
    mf.scf()
    elapsed_time = time.time() - start
    print ("SCF: {:.3f} ms".format(elapsed_time * 1000))
    e_scf = scf.hf.energy_tot(mf)
    x_scf.append(R)
    y_scf.append(e_scf)

    # Curvilinear search using Cayley transformation
    start = time.time()

    # parameters
    tau = 1.0
    tau_m = 1e-10
    tau_M = 1e10
    rho = 1e-4
    delta = 0.1
    eta = 0.5
    epsilon = 1e-6
    max_iter = 5000

    # 1. initialize X0
    S = mol.intor_symmetric('int1e_ovlp')  # overlap matrix
    S64 = numpy.asarray(S, dtype=numpy.float64)
    X_np = scipy.linalg.inv(scipy.linalg.sqrtm(S64))
    X = jnp.asarray(X_np)
    
    # 2. set C=f(X0) and Q0=1
    C = calcEnergy(X)
    Q = 1.0

    # 3. calculate G0 and A0
    G = gradEnergy(X)
    A = G @ X.T @ S - S @ X @ G.T

    # function to calculate Y(tau)
    I = jnp.identity(len(S))
    def Y_tau(tau, X, A):
        return jnp.linalg.inv(I + 0.5 * tau * A @ S) @ (I - 0.5 * tau * A @ S) @ X
    
    # main loop
    for k in range(max_iter):
        Y = Y_tau(tau, X, A)
        A_norm = jnp.linalg.norm(A, "fro")
        X_old, Q_old, G_old = X, Q, G
    
        # 5
        while calcEnergy(Y) > C - rho * tau * A_norm**2.0:
            tau *= delta    # 6
            Y = Y_tau(tau, X, A)
    
        # 8
        X_new = Y
        Q_new = eta * Q + 1.0
        C = (eta * Q * C + calcEnergy(X_new)) / Q_new

        # 9
        G_new = gradEnergy(X_new)
        A_new = G_new @ X_new.T @ S - S @ X_new @ G_new.T

        # 10
        Sk = X_new - X
        Yk = G_new - G
        if k % 2 == 0:
            tau_k = jnp.trace(Sk.T @ Sk) / abs(jnp.trace(Sk.T @ Yk))
        else:
            tau_k = abs(jnp.trace(Sk.T @ Yk)) / jnp.trace(Yk.T @ Yk)
        tau = max(min(tau_k, tau_M), tau_m)

        # Update variables for next iteration
        X, Q, G, A = X_new, Q_new, G_new, A_new

        # Check loop condition (4)
        cond = jnp.linalg.norm(A @ X)
        if cond < epsilon:
            break

    elapsed_time = time.time() - start
    print ("Curvilinear search: {:.3f} ms".format(elapsed_time*1000))
    e = calcEnergy(X)+mol.energy_nuc()
    print(f"total energy = {e}")
    x.append(R)
    y.append(e)

    # augmented Lagrangian
    @jit
    def orthogonality(x):
        C = jnp.reshape(x, [len(S), len(S)])
        return jnp.linalg.norm(C.transpose()@S@C - jnp.identity(len(S)))

    start = time.time()
    x0 = random.uniform(key, (S.size,))

    # 1
    mu = 1.0
    lam = 0.0

    constraint = orthogonality(x0)

    # 2
    while constraint > 1e-6:
        def target(x):
            h = orthogonality(x)
            return calcEnergy(x) + mu * h ** 2.0 + lam * h

        # 3
        res = minimize(jit(target), x0, jac=jit(grad(jit(target))), method="BFGS", options={'maxiter': 100})
        x0 = res.x
        constraint = orthogonality(x0)
        # 4
        lam += 2.0 * mu * constraint
        # 5
        mu *= 2.0
    elapsed_time = time.time() - start
    print ("Augmented: {:.3f} s".format(elapsed_time*1000))
    energy = res.fun+mol.energy_nuc()
    print(f"calculated energy = {energy}")
    x_aug.append(R)
    y_aug.append(energy)

p2 = plt.plot(x_scf, y_scf, marker="o")
p1 = plt.plot(x_aug, y_aug, marker="*")
p0 = plt.plot(x, y, marker="x")
plt.xlabel("interatomic distance (Å)", fontsize=16)
plt.ylabel("total energy (Eh)", fontsize=16)
plt.legend((p0[0], p1[0], p2[0]),
           ("Curvilinear search", "Augmented Lagrangian", "PySCF"),
           loc='upper right')
plt.savefig("result.png", dpi=300)
# RHF energy calculation by PySCF # Curvilinear search using Cayley transformation # parameters # 1. initialize X0 # overlap matrix # 2. set C=f(X0) and Q0=1 # 3. calculate G0 and A0 # function to calculate Y(tau) # main loop # 5 # 6 # 8 # 9 # 10 # Update variables for next iteration # Check loop condition (4) # augmented Lagrangian # 1 # 2 # 3 # 4 # 5