From 8934166f88ee0029a4b6ce9e339d96cb089f7ce0 Mon Sep 17 00:00:00 2001
From: Brian Zhao <zijun.zhao@emory.edu>
Date: Tue, 26 Nov 2024 15:38:24 -0500
Subject: [PATCH] Add test for untested helper functions

---
 forte/utils/helpers.py               | 41 ++--------------------------
 tests/pytest/helpers/test_helpers.py | 34 +++++++++++++++++++++++
 2 files changed, 36 insertions(+), 39 deletions(-)
 create mode 100644 tests/pytest/helpers/test_helpers.py

diff --git a/forte/utils/helpers.py b/forte/utils/helpers.py
index e1717bfd2..daf310156 100644
--- a/forte/utils/helpers.py
+++ b/forte/utils/helpers.py
@@ -98,45 +98,11 @@ def psi4_casscf(geom, basis, reference, restricted_docc, active, options={}) ->
     # pipe output to the file output.dat
     psi4.core.set_output_file("output.dat", True)
 
-    # psi4.core.clean()
-
     # run scf and return the energy and a wavefunction object (will work only if pass return_wfn=True)
     E_scf, wfn = psi4.energy("casscf", molecule=mol, return_wfn=True)
     return (E_scf, wfn)
 
 
-def psi4_casscf(geom, basis, mo_spaces):
-    """
-    Run a Psi4 SCF.
-    :param geom: a string for molecular geometry
-    :param basis: a string for basis set
-    :param reference: a string for the type of reference
-    :return: a tuple of (scf energy, psi4 Wavefunction)
-    """
-    psi4.core.clean()
-    mol = psi4.geometry(geom)
-
-    psi4.set_options(
-        {
-            "basis": basis,
-            "scf_type": "pk",
-            "e_convergence": 1e-13,
-            "d_convergence": 1e-6,
-            "restricted_docc": mo_spaces["RESTRICTED_DOCC"],
-            "active": mo_spaces["ACTIVE"],
-            "mcscf_maxiter": 100,
-            "mcscf_e_convergence": 1.0e-11,
-            "mcscf_r_convergence": 1.0e-6,
-            "mcscf_diis_start": 20,
-        }
-    )
-    psi4.core.set_output_file("output.dat", False)
-
-    Escf, wfn = psi4.energy("casscf", return_wfn=True)
-    psi4.core.clean()
-    return Escf, wfn
-
-
 def psi4_cubeprop(wfn, path=".", orbs=[], nocc=0, nvir=0, density=False, frontier_orbitals=False, load=False):
     """
     Run a psi4 cubeprop computation to generate cube files from a given Wavefunction object
@@ -184,7 +150,7 @@ def psi4_cubeprop(wfn, path=".", orbs=[], nocc=0, nvir=0, density=False, frontie
 
 
 def prepare_forte_objects(
-    wfn, mo_spaces=None, active_space="ACTIVE", core_spaces=["RESTRICTED_DOCC"], localize=False, localize_spaces=[]
+    wfn, mo_spaces, active_space="ACTIVE", core_spaces=["RESTRICTED_DOCC"], localize=False, localize_spaces=[]
 ):
     """Take a psi4 wavefunction object and prepare the ForteIntegrals, SCFInfo, and MOSpaceInfo objects
 
@@ -235,10 +201,7 @@ def prepare_forte_objects(
     point_group = wfn.molecule().point_group().symbol()
 
     # create a MOSpaceInfo object
-    if mo_spaces is None:
-        mo_space_info = forte.make_mo_space_info(nmopi, point_group, options)
-    else:
-        mo_space_info = forte.make_mo_space_info_from_map(nmopi, point_group, mo_spaces)
+    mo_space_info = forte.make_mo_space_info_from_map(nmopi, point_group, mo_spaces)
 
     # These variables are needed in make_state_weights_map
     nel = wfn.nalpha() + wfn.nbeta()
diff --git a/tests/pytest/helpers/test_helpers.py b/tests/pytest/helpers/test_helpers.py
new file mode 100644
index 000000000..cd0ad5676
--- /dev/null
+++ b/tests/pytest/helpers/test_helpers.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+
+def test_helpers():
+    import math
+    import pytest
+    import psi4
+    import forte
+    import forte.utils
+
+    psi4.core.clean()
+
+    geom = """
+    O  0.0000  0.0000  0.1173
+    H  0.0000  0.7572 -0.4692
+    H  0.0000 -0.7572 -0.4692    
+    symmetry c2v
+    """
+    (E_scf, wfn_scf) = forte.utils.psi4_scf(geom,basis='6-31g',reference='rhf')
+    assert math.isclose(E_scf, -75.98397447271522)
+
+    (E_casscf, wfn_casscf) = forte.utils.psi4_casscf(geom,basis='6-31g',reference='rhf', restricted_docc=[2,0,0,1], active=[2,0,1,1])
+    assert math.isclose(E_casscf, -75.9998515885993)
+
+    res = forte.utils.prepare_forte_objects(wfn_casscf, {"RESTRICTED_DOCC": [2,0,0,1], "ACTIVE": [2,0,1,1]})
+    mo_space_info = res["mo_space_info"]
+    # mo_space_info.corr_absolute_mo("RESTRICTED_DOCC") == [0, 1, 9]
+    assert mo_space_info.corr_absolute_mo("RESTRICTED_DOCC")[2] == 9
+    
+    psi4.core.clean()
+
+if __name__ == "__main__":
+    test_helpers()