New corrections

fanpy/eqn/energy_oneside.py

@@ -104,6 +104,7 @@ def __init__(
         step_save=True,
         tmpfile="",
         refwfn=None,
+        mpi_comm=None,
     ):
         """Initialize the objective instance.
 
@@ -144,6 +145,9 @@ def __init__(
             space.
             By default, the given wavefunction is used as the reference by using a complete
             projection space.
+        mpi_comm : mpi4py.MPI.Comm
+            MPI communicator used to split projection-space sums across ranks.
+            Default disables MPI parallelism.
 
         Raises
         ------
@@ -163,6 +167,7 @@ def __init__(
             step_print=step_print,
             step_save=step_save,
             tmpfile=tmpfile,
+            mpi_comm=mpi_comm,
         )
         self.assign_refwfn(refwfn)
 
@@ -218,6 +223,10 @@ def assign_refwfn(self, refwfn=None):  # pylint: disable=R0912
         if slater.is_sd_compatible(refwfn):
             refwfn = [refwfn]
 
+        if hasattr(refwfn, "iter_chunks"):
+            self.refwfn = refwfn
+            return
+
         if __debug__:
             if isinstance(refwfn, (list, tuple)):
                 for sd in refwfn:  # pylint: disable=C0103
@@ -251,7 +260,7 @@ def assign_refwfn(self, refwfn=None):  # pylint: disable=R0912
                         "spin orbitals as the given wavefunction."
                     )
             else:
-                raise TypeError("Projection space must be given as a list or a tuple.")
+                raise TypeError("Projection space must be given as a list, tuple, or chunked projection space.")
 
         self.refwfn = refwfn
 
@@ -284,14 +293,15 @@ def objective(self, params, assign=True, normalize=False, save=True):
 
         energy = self.get_energy_one_proj(self.refwfn)
 
-        if self.step_print:
+        is_mpi_root = self.mpi_comm is None or self.mpi_comm.Get_rank() == 0
+        if self.step_print and is_mpi_root:
             print("(Mid Optimization) Electronic energy: {}".format(energy))
-        else:
+        elif is_mpi_root:
             self.print_queue["Electronic energy"] = energy
 
         return energy
 
-    def gradient(self, params, assign=True, normalize=True, save=True):
+    def gradient(self, params, assign=True, normalize=False, save=True):
         """Return the gradient of the energy integrated against the reference wavefunction.
 
         See `BaseSchrodinger.get_energy_one_proj` for details.
@@ -315,15 +325,17 @@ def gradient(self, params, assign=True, normalize=True, save=True):
         if normalize and hasattr(self.wfn, "normalize"):
             self.wfn.normalize(self.refwfn)
         # Save params
-        if save:
+        if save and self.step_save:
             self.save_params()
 
         grad = self.get_energy_one_proj(self.refwfn, True)
 
         grad_norm = np.linalg.norm(grad)
-        if self.step_print:
+        is_mpi_root = self.mpi_comm is None or self.mpi_comm.Get_rank() == 0
+        if self.step_print and is_mpi_root:
             print("(Mid Optimization) Norm of the gradient of the energy: {}".format(grad_norm))
-        else:
+        elif is_mpi_root:
             self.print_queue["Norm of the gradient of the energy"] = grad_norm
 
         return grad
+

fanpy/solver/equation.py

@@ -175,9 +175,9 @@ def update_iteration(*args):  # pylint: disable=W0613
         """Clean up at the end of each iteration."""
         # update hamiltonian
         if objective.indices_component_params[ham].size > 0:
-            ham.update_prev_prams = True
+            ham.update_prev_params = True
             ham.assign_params(ham.params)
-            ham.update_prev_prams = False
+            ham.update_prev_params = False
         # save parameters
         objective.save_params()
         # print
@@ -194,3 +194,122 @@ def update_iteration(*args):  # pylint: disable=W0613
         output["energy"] = output["function"]
 
     return output
+
+
+def basinhopping(objective, use_gradient=True, **kwargs):
+    """Solve an equation using `scipy.optimize.basinhopping`.
+
+    Basin-hopping performs repeated random perturbations followed by local minimizations. It is useful
+    when a local optimizer such as BFGS is getting trapped in one basin.
+
+    Parameters
+    ----------
+    objective : BaseSchrodinger
+        Instance that contains the function that will be optimized.
+    use_gradient : bool
+        Option to pass the objective gradient to the local minimizer.
+        Default is True.
+    kwargs : dict
+        Keyword arguments to `scipy.optimize.basinhopping`. Local minimizer options should be given
+        in `minimizer_kwargs`. By default, the local minimizer is BFGS with the objective gradient
+        when `use_gradient` is True, and Powell otherwise.
+
+    Returns
+    -------
+    Dictionary with the following keys and values:
+    success : bool
+        True if optimization succeeded.
+    params : np.ndarray
+        Parameters at the end of the optimization.
+    energy : float
+        Energy after optimization.
+        Only available for objectives that are EnergyOneSideProjection, EnergyTwoSideProjection, and
+        LeastSquaresEquations instances.
+    function : float
+        Final objective value.
+    message : str
+        Message returned by the optimizer.
+    internal : OptimizeResult
+        Returned value of `scipy.optimize.basinhopping`.
+
+    Raises
+    ------
+    TypeError
+        If objective is not BaseSchrodinger instance.
+    ValueError
+        If objective has more than one equation.
+
+    """
+    import scipy.optimize  # pylint: disable=C0415
+
+    if not isinstance(objective, BaseSchrodinger):
+        raise TypeError("Objective must be a BaseSchrodinger instance.")
+    if objective.num_eqns != 1:
+        raise ValueError("Objective must contain only one equation.")
+
+    minimizer_kwargs = kwargs.setdefault("minimizer_kwargs", {})
+    if "args" in kwargs:
+        minimizer_kwargs.setdefault("args", kwargs.pop("args"))
+    if use_gradient:
+        minimizer_kwargs.setdefault("method", "BFGS")
+        minimizer_kwargs.setdefault("jac", objective.gradient)
+        minimizer_kwargs.setdefault("options", {})
+        minimizer_kwargs["options"].setdefault("gtol", 1e-8)
+        minimizer_kwargs["options"].setdefault("ftol", 1e-12)
+        minimizer_kwargs["options"].setdefault("maxiter", 500)
+    else:
+        minimizer_kwargs.setdefault("method", "Powell")
+        minimizer_kwargs.setdefault("options", {})
+        minimizer_kwargs["options"].setdefault("xtol", 1e-9)
+        minimizer_kwargs["options"].setdefault("ftol", 1e-9)
+
+    ham = objective.ham
+    ham.update_prev_params = False
+    user_callback = kwargs.get("callback", None)
+    user_minimizer_callback = minimizer_kwargs.get("callback", None)
+
+    def cleanup():
+        """Clean up during local minimization and after each basin minimum."""
+        if objective.indices_component_params[ham].size > 0:
+            ham.update_prev_params = True
+            ham.assign_params(ham.params)
+            ham.update_prev_params = False
+        objective.save_params()
+        for key, value in objective.print_queue.items():
+            print("(Mid Optimization) {}: {}".format(key, value))
+
+    def update_local_minimize(*args, **callback_kwargs):
+        """Clean up at the end of a local minimizer iteration."""
+        cleanup()
+        if user_minimizer_callback is not None:
+            return user_minimizer_callback(*args, **callback_kwargs)
+        return None
+
+    def update_basinhopping(x, f, accept):  # pylint: disable=C0103
+        """Clean up at the end of a basin-hopping step."""
+        cleanup()
+        if user_callback is not None:
+            return user_callback(x, f, accept)
+        return None
+
+    minimizer_kwargs["callback"] = update_local_minimize
+    kwargs["callback"] = update_basinhopping
+
+    results = scipy.optimize.basinhopping(objective.objective, objective.active_params, **kwargs)
+
+    output = {}
+    output["success"] = results.success
+    output["params"] = results.x
+    output["function"] = results.fun
+    output["message"] = results.message
+    output["internal"] = results
+    if isinstance(objective, LeastSquaresEquations):
+        output["energy"] = objective.energy.params
+    elif isinstance(objective, (EnergyOneSideProjection, EnergyTwoSideProjection)):  # pragma: no branch
+        output["energy"] = results.fun
+
+    objective.assign_params(results.x)
+    objective.save_params()
+
+    return output
+

fanpy/wfn/cc/base.py

@@ -655,7 +655,7 @@ def _olp(sd1):
             """Cached _olp method without caching the instance."""
             return self._olp(sd1)
 
-        @functools.lru_cache(maxsize=0, typed=False)
+        @functools.lru_cache(maxsize=memory, typed=False)
         def _olp_deriv(sd1):
             """Cached _olp_deriv method without caching the instance."""
             return self._olp_deriv(sd1)