Chore: extract calculations from interface and clean up

fanpy/interface/fanci/legacy.py

@@ -3,12 +3,11 @@
 Based of the original class from FanCI code.
 """
 
-from fanpy.tools import slater
-
 from fanpy.wfn.base import BaseWavefunction
 from fanpy.eqn.base import BaseSchrodinger
 from fanpy.wfn.composite.product import ProductWavefunction
 from fanpy.interface.fanci.alias import Alias
+from fanpy.interface.fanci.utils import convert_pyci_occs_to_fanpy_sds
 from fanpy.tools.performance import current_memory
 
 from typing import Any, Callable, Dict, List, Sequence, Tuple, Union
@@ -19,9 +18,8 @@
 import math
 
 import pyci
-import cma
 import numpy as np
-from scipy.optimize import OptimizeResult, least_squares, root, minimize
+from scipy.optimize import OptimizeResult, least_squares, root
 
 __all__ = [
     "ProjectedSchrodingerFanCI",
@@ -274,7 +272,7 @@ def __init__(
         self._pspace.setflags(write=False)
         self._sspace.setflags(write=False)
         self._mask_view.setflags(write=False)
-
+    # todo: check if I can move the functions to tools: this is used in multiple places. 
     def make_norm_constraint(self):
         def f(x: np.ndarray) -> float:
             r""" "
@@ -529,7 +527,7 @@ def unfreeze_parameter(self, *params: Sequence[int]) -> None:
 
         # Update nactive
         self._nactive = self._mask.sum()
-
+    # todo: check if test exists for this. PyCI might have it since it is from it. May have to extract
     def compute_objective(self, x: np.ndarray) -> np.ndarray:
         """
         Compute the FanCI objective function.
@@ -580,6 +578,7 @@ def compute_objective(self, x: np.ndarray) -> np.ndarray:
         # Return objective
         return f
 
+    # todo: same as compute_objective
     def compute_jacobian(self, x: np.ndarray) -> np.ndarray:
         """
         Compute the Jacobian of the FanCI objective function.
@@ -652,6 +651,7 @@ def compute_jacobian(self, x: np.ndarray) -> np.ndarray:
         # Return Jacobian
         return jac
 
+    #todo: same as norm constraint
     def make_param_constraint(self, i: int, val: float) -> Tuple[Callable, Callable]:
         """
         Generate parameter constraint functions.
@@ -690,7 +690,8 @@ def dfdx(x: np.ndarray) -> np.ndarray:
             return y
 
         return f, dfdx
-
+
+    #todo: same as norm constraint
     def make_det_constraint(self, i: int, val: float) -> Tuple[Callable, Callable]:
         """
         Generate determinant overlap constraint functions.
@@ -1113,23 +1114,7 @@ def compute_overlap(self, x: np.ndarray, occs_array: Union[np.ndarray, str]) ->
             raise ValueError("invalid `occs_array` argument")
         # FIXME: converting occs_array to slater determinants to be converted back to indices is a waste
         # convert slater determinants
-        sds = []
-        if isinstance(occs_array[0, 0], np.ndarray): # pspace generated with FCI
-            for i, occs in enumerate(occs_array):
-                # FIXME: CHECK IF occs IS BOOLEAN OR INTEGERS
-                # convert occupation vector to sd
-                if occs.dtype == bool:
-                    occs = np.where(occs)[0]
-                sd = slater.create(0, *occs[0])
-                sd = slater.create(sd, *(occs[1] + self.fanpy_wfn.nspatial))
-                sds.append(sd)
-        else: # pspace generated with DOCI
-            for i, occs in enumerate(occs_array):
-                if occs.dtype == bool:
-                    occs = np.where(occs)
-                sd = slater.create(0, *occs)
-                sd = slater.create(sd, *(occs + self.fanpy_wfn.nspatial))
-                sds.append(sd)
+        sds = convert_pyci_occs_to_fanpy_sds(occs_array, nspatial = self.fanpy_wfn.nspatial)
 
         # Feed in parameters into fanpy wavefunction
         for component, indices in self.param_selection.items():
@@ -1183,23 +1168,7 @@ def compute_overlap_deriv(
         # FIXME: converting occs_array to slater determinants to be converted back to indices is
         # a waste
         # convert slater determinants
-        sds = []
-        if isinstance(occs_array[0, 0], np.ndarray): # pspace generated with FCI
-            for i, occs in enumerate(occs_array):
-                # FIXME: CHECK IF occs IS BOOLEAN OR INTEGERS
-                # convert occupation vector to sd
-                if occs.dtype == bool:
-                    occs = np.where(occs)[0]
-                sd = slater.create(0, *occs[0])
-                sd = slater.create(sd, *(occs[1] + self.fanpy_wfn.nspatial))
-                sds.append(sd)
-        else: # pspace generated with DOCI
-            for i, occs in enumerate(occs_array):
-                if occs.dtype == bool:
-                    occs = np.where(occs)
-                sd = slater.create(0, *occs)
-                sd = slater.create(sd, *(occs + self.fanpy_wfn.nspatial))
-                sds.append(sd)
+        sds = convert_pyci_occs_to_fanpy_sds(occs_array, nspatial = self.fanpy_wfn.nspatial)
 
         # Select sds according to selected chunks
         s_chunk, f_chunk = chunk_idx
@@ -1262,13 +1231,13 @@ def compute_overlap_double_deriv(self, x: np.ndarray, occs_array: Union[np.ndarr
         -------
         ovlp_hessian : np.ndarray
             Shape: (N_SD, nparam, nparam)
-        """
-        if occs_array == "P":
+       """
+        if isinstance(occs_array, np.ndarray):
+            pass
+        elif occs_array == "P":
             occs_array = self.pspace
         elif occs_array == "S":
             occs_array = self.sspace
-        elif isinstance(occs_array, np.ndarray):
-            pass
         else:
             raise ValueError("invalid `occs_array` argument")
         from fanpy.wfn.cc.base import BaseCC
@@ -1280,22 +1249,7 @@ def compute_overlap_double_deriv(self, x: np.ndarray, occs_array: Union[np.ndarr
         "Must be a CC-type FanPy wavefunction (BaseCC) with "
         "`get_overlap_double_derivative(sd)` implemented.")
         # Convert occs -> Slater determinants
-        sds = []
-        if isinstance(occs_array[0, 0], np.ndarray):
-            for i, occs in enumerate(occs_array):
-                # FIXME: CHECK IF occs IS BOOLEAN OR INTEGERS
-                # convert occupation vector to sd
-                if occs.dtype == bool:
-                    occs = np.where(occs)[0]
-                sd = slater.create(0, *occs[0])
-                sd = slater.create(sd, *(occs[1] + self.fanpy_wfn.nspatial))
-                sds.append(sd)
-        else:
-            for i, occs in enumerate(occs_array):
-                if occs.dtype == bool:
-                    occs = np.where(occs)
-                sd = slater.create(0, *occs)
-                sds.append(sd)        
+        sds = convert_pyci_occs_to_fanpy_sds(occs_array, nspatial = self.fanpy_wfn.nspatial)      
 
         # Update Fanpy params
         for component, indices in self.param_selection.items():
@@ -1331,40 +1285,17 @@ def compute_objective(self, x: np.ndarray) -> np.ndarray:
             Objective vector.
 
         """
-        if self.objective_type == "projected":
-            output = super().compute_objective(x)
-            self.print_queue["Electronic Energy"] = x[-1]
-            self.print_queue["Cost"] = np.sum(output[: self.nproj] ** 2)
-            self.print_queue["Cost from constraints"] = np.sum(output[self.nproj :] ** 2)
-            if self.step_print:
-                print("(Mid Optimization) Electronic Energy: {}".format(self.print_queue["Electronic Energy"]))
-                print("(Mid Optimization) Cost: {}".format(self.print_queue["Cost"]))
-                if self.constraints:
-                    print(
-                        "(Mid Optimization) Cost from constraints: {}".format(self.print_queue["Cost from constraints"])
-                    )
-        else:
-            # NOTE: ignores energy and constraints
-            # Allocate objective vector
-            output = np.zeros(self.nproj, dtype=pyci.c_double)
-
-            # Compute overlaps of determinants in sspace:
-            #
-            #   c_m
-            #
-            ovlp = self.compute_overlap(x[:-1], "S")
-
-            # Compute objective function:
-            #
-            #   f_n = (\sum_n <\Psi|n> <n|H|\Psi>) / \sum_n <\Psi|n> <n|\Psi>
-            #
-            # Note: we update ovlp in-place here
-            self.ci_op(ovlp, out=output)
-            output = np.sum(output * ovlp[: self.nproj])
-            output /= np.sum(ovlp[: self.nproj] ** 2)
-            self.print_queue["Electronic Energy"] = output
-            if self.step_print:
-                print("(Mid Optimization) Electronic Energy: {}".format(self.print_queue["Electronic Energy"]))
+        output = super().compute_objective(x)
+        self.print_queue["Electronic Energy"] = x[-1]
+        self.print_queue["Cost"] = np.sum(output[: self.nproj] ** 2)
+        self.print_queue["Cost from constraints"] = np.sum(output[self.nproj :] ** 2)
+        if self.step_print:
+            print("(Mid Optimization) Electronic Energy: {}".format(self.print_queue["Electronic Energy"]))
+            print("(Mid Optimization) Cost: {}".format(self.print_queue["Cost"]))
+            if self.constraints:
+                print(
+                    "(Mid Optimization) Cost from constraints: {}".format(self.print_queue["Cost from constraints"])
+                )
 
         if self.step_save:
             self.save_params()
@@ -1388,66 +1319,10 @@ def compute_jacobian(self, x: np.ndarray) -> np.ndarray:
             Jacobian matrix.
 
         """
-        if self.objective_type == "projected":
-            output = super().compute_jacobian(x)
-            self.print_queue["Norm of the Jacobian"] = np.linalg.norm(output)
-            if self.step_print:
-                print("(Mid Optimization) Norm of the Jacobian: {}".format(self.print_queue["Norm of the Jacobian"]))
-        else:
-            # NOTE: ignores energy and constraints
-            # Allocate Jacobian matrix (in transpose memory order)
-            output = np.zeros((self.nproj, self.nactive), order="F", dtype=pyci.c_double)
-            integrals = np.zeros(self.nproj, dtype=pyci.c_double)
-
-            # Compute Jacobian:
-            #
-            #   J_{nk} = d(<n|H|\Psi>)/d(p_k) - E d(<n|\Psi>)/d(p_k) - dE/d(p_k) <n|\Psi>
-            #   J_{nk} = (\sum_n d<\Psi|n> <n|H|\Psi> + <\Psi|n> d<n|H|\Psi>) / \sum_n <\Psi|n>^2 -
-            #            (\sum_n <\Psi|n> <n|H|\Psi>) / (\sum_n <\Psi|n> <n|\Psi>)^2 * (2 \sum_n <\Psi|n>)
-            #   J_{nk} = ((\sum_n d<\Psi|n> <n|H|\Psi> + <\Psi|n> d<n|H|\Psi>) (\sum_n <\Psi|n>^2)
-            #             - (\sum_n <\Psi|n> <n|H|\Psi>) * (2 \sum_n <\Psi|n> d<\Psi|n>))
-            #            / (\sum_n <\Psi|n>^2)^2
-            #   J_{nk} = ((\sum_n d<\Psi|n> <n|H|\Psi> + <\Psi|n> d<n|H|\Psi>) N
-            #             - H * (2 \sum_n <\Psi|n> d<\Psi|n>))
-            #            / N^2
-            #   J_{nk} = (\sum_n N (d<\Psi|n> <n|H|\Psi> + <\Psi|n> d<n|H|\Psi>) - 2 H <\Psi|n> d<\Psi|n>)
-            #            / N^2
-            #
-            # Compute overlap derivatives in sspace:
-            #
-            #   d(c_m)/d(p_k)
-            #
-            overlaps = self.compute_overlap(x[:-1], "S")
-            norm = np.sum(overlaps[: self.nproj] ** 2)
-            self.ci_op(overlaps, out=integrals)
-            energy_integral = np.sum(overlaps[: self.nproj] * integrals)
-
-            d_ovlp = self.compute_overlap_deriv(x[:-1], "S")
-
-            # Iterate over remaining columns of Jacobian and d_ovlp
-            for output_col, d_ovlp_col in zip(output.transpose(), d_ovlp.transpose()):
-                #
-                # Compute each column of the Jacobian:
-                #
-                #   d(<n|H|\Psi>)/d(p_k) = <m|H|n> d(c_m)/d(p_k)
-                #
-                #   E d(<n|\Psi>)/d(p_k) = E \delta_{nk} d(c_n)/d(p_k)
-                #
-                # Note: we update d_ovlp in-place here
-                self.ci_op(d_ovlp_col, out=output_col)
-                output_col *= overlaps[: self.nproj]
-                output_col += d_ovlp_col[: self.nproj] * integrals
-                output_col *= norm
-                output_col -= 2 * energy_integral * overlaps[: self.nproj] * d_ovlp_col[: self.nproj]
-                output_col /= norm**2
-            output = np.sum(output, axis=0)
-            self.print_queue["Norm of the gradient of the energy"] = np.linalg.norm(output)
-            if self.step_print:
-                print(
-                    "(Mid Optimization) Norm of the gradient of the energy: {}".format(
-                        self.print_queue["Norm of the gradient of the energy"]
-                    )
-                )
+        output = super().compute_jacobian(x)
+        self.print_queue["Norm of the Jacobian"] = np.linalg.norm(output)
+        if self.step_print:
+            print("(Mid Optimization) Norm of the Jacobian: {}".format(self.print_queue["Norm of the Jacobian"]))
 
         if self.step_save:
             self.save_params()
@@ -1596,32 +1471,8 @@ def optimize(
             opt_kwargs.setdefault("options", {})
             opt_kwargs["options"].setdefault("xtol", 1.0e-9)
             opt_kwargs.setdefault("callback", self.print)
-        elif mode == "cma":
-            optimizer = cma.fmin
-            opt_kwargs.setdefault("sigma0", 0.01)
-            opt_kwargs.setdefault("options", {})
-            opt_kwargs["options"].setdefault("ftarget", None)
-            opt_kwargs["options"].setdefault("timeout", np.inf)
-            opt_kwargs["options"].setdefault("tolfun", 1e-11)
-            opt_kwargs["options"].setdefault("verb_log", 0)
-            if self.objective_type != "energy":
-                raise ValueError("objective_type must be energy")
-        elif mode == "bfgs":
-            if self.objective_type != "energy":
-                raise ValueError("objective_type must be energy")
-            optimizer = minimize
-            opt_kwargs["method"] = "bfgs"
-            opt_kwargs.setdefault("options", {"gtol": 1e-8})
-            # opt_kwargs["options"]['schrodinger'] = objective
-            opt_kwargs.setdefault("callback", self.print)
-        elif mode == "trustregion":
-            raise NotImplementedError
-        elif mode == "trf":
-            if self.objective_type != "projected":
-                raise ValueError("objective_type must be energy")
-            raise NotImplementedError
         else:
-            raise ValueError("invalid mode parameter")
+            raise ValueError("Invalid mode parameter. Only 'root' and 'lstsq' are supported for the ProjectedSchrodingerFanCI class.")
 
         # Run optimizer
         results = optimizer(*opt_args, **opt_kwargs)