mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/optimizers/MicroOptimizer.py

@@ -0,0 +1,169 @@
+import logging
+import os
+from pathlib import Path
+import sys
+
+import numpy as np
+
+from pysisyphus.optimizers.closures import bfgs_multiply
+from pysisyphus.optimizers.hessian_updates import double_damp
+from pysisyphus.optimizers.poly_fit import poly_line_search
+
+
+class MicroOptimizer:
+    def __init__(
+        self,
+        geom,
+        step="lbfgs",
+        line_search=True,
+        max_cycles=100_000_000,
+        max_step=0.2,
+        keep_last=10,
+        rms_force=None,
+        double_damp=True,
+        dump=False,
+        **kwargs,
+    ):
+        self.geometry = geom
+        self.step_funcs = {
+            "sd": self.sd_step,
+            "cg": self.cg_step,
+            "lbfgs": self.lbfgs_step,
+        }
+        self.step_func = self.step_funcs[step]
+        self.line_search = line_search
+        self.max_cycles = max_cycles
+        self.max_step = max_step
+        self.keep_last = keep_last
+        self.double_damp = double_damp
+        self.rms_force = rms_force
+        self.dump = dump
+
+        self.prev_energy = None
+        self.prev_step = None
+        self.prev_forces = None
+        self.coord_diffs = list()
+        self.grad_diffs = list()
+
+        self.trj_fn = "microopt_trj.xyz"
+        if Path(self.trj_fn).exists():
+            os.remove(self.trj_fn)
+
+        self.logger = logging.getLogger("optimizer")
+
+        if kwargs:
+            msg = "Got unsupported keyword arguments: " + ", ".join(
+                [f"{k}={v}" for k, v in kwargs.items()]
+            )
+            self.log(msg)
+        self.is_converged = False
+
+    def log(self, msg):
+        self.logger.debug(msg)
+
+    def run(self):
+        self.geometry.reparametrize()
+
+        for self.cur_cycle in range(self.max_cycles):
+            if self.dump:
+                with open(self.trj_fn, "a") as handle:
+                    handle.write(self.geometry.as_xyz() + "\n")
+
+            results = self.geometry.get_energy_and_forces_at(self.geometry.coords)
+            forces = results["forces"]
+            energy = results["energy"]
+            rms = np.sqrt(np.mean(forces ** 2))
+            print(f"{self.cur_cycle:03d} rms(f)={rms:.6f}")
+            sys.stdout.flush()
+            if self.rms_force and rms <= self.rms_force:
+                print("Converged!")
+                self.is_converged = True
+                break
+
+            self.take_step(energy, forces)
+
+    def take_step(self, energy, forces, return_step=False):
+        self.log(
+            f"Cycle {self.cur_cycle:03d}, energy={energy:.6f} au, "
+            f"norm(forces)={np.linalg.norm(forces):.6f}"
+        )
+
+        ip_gradient, ip_step = None, None
+        if self.line_search and (self.cur_cycle > 0):
+            ip_energy, ip_gradient, ip_step = poly_line_search(
+                energy,
+                self.prev_energy,
+                -forces,
+                -self.prev_forces,
+                self.prev_step,
+            )
+        # Use the interpolated gradient for the step if interpolation succeeded
+        if (ip_gradient is not None) and (ip_step is not None):
+            forces = -ip_gradient
+        # Keep the original gradient when the interpolation failed, but use
+        # zero (interpolated) step.
+        else:
+            ip_step = np.zeros_like(forces)
+
+        # Calculate actual step with potentially interpolated forces
+        if self.cur_cycle == 0:
+            step = self.sd_step(forces)
+        else:
+            step = self.step_func(forces)
+
+        # Form full step. If we did not interpolate or interpolation failed ip_step will be zero.
+        step = step + ip_step
+        # Restrict absolute value of step vector entries
+        step *= min(self.max_step / np.abs(step).max(), 1)
+
+        self.geometry.reparametrize()
+
+        self.prev_step = step
+        self.prev_energy = energy
+        self.prev_forces = forces
+        if return_step:
+            return step
+        else:
+            new_coords = self.geometry.coords + step
+            self.geometry.coords = new_coords
+
+    def optimize(self):
+        if not hasattr(self, "cur_cycle"):
+            self.cur_cycle = 0
+        step = self.take_step(
+            self.geometry.energy, self.geometry.forces, return_step=True
+        )
+        self.cur_cycle += 1
+        return step
+
+    def sd_step(self, forces):
+        step = forces
+        return step
+
+    def cg_step(self, forces):
+        beta = forces.dot(forces) / self.prev_forces.dot(self.prev_forces)
+        step = forces + beta * self.prev_step
+        return step
+
+    def lbfgs_step(self, forces):
+        y = self.prev_forces - forces
+        s = self.prev_step
+
+        if self.double_damp:
+            s, y = double_damp(s, y, s_list=self.coord_diffs, y_list=self.grad_diffs)
+        self.grad_diffs.append(y)
+        self.coord_diffs.append(s)
+
+        # Drop superfluous oldest vectors
+        self.coord_diffs = self.coord_diffs[-self.keep_last :]
+        self.grad_diffs = self.grad_diffs[-self.keep_last :]
+
+        step = bfgs_multiply(
+            self.coord_diffs,
+            self.grad_diffs,
+            forces,
+            gamma_mult=False,
+            logger=self.logger,
+        )
+
+        return step

pysisyphus/optimizers/NCOptimizer.py

@@ -0,0 +1,90 @@
+# [1] https://aip.scitation.org/doi/pdf/10.1063/1.1498468
+#     Bour, 2002
+
+import numpy as np
+
+from pysisyphus.optimizers.HessianOptimizer import HessianOptimizer
+from pysisyphus.helpers_pure import eigval_to_wavenumber 
+
+
+class NCOptimizer(HessianOptimizer):
+
+    def __init__(self, geometry, *args, freeze_modes=None, **kwargs):
+        super().__init__(geometry, **kwargs)
+
+        assert not self.is_cos and self.geometry.coord_type == "cart", \
+            "NCOptimizer can't be used with ChainOfStates-methods and " \
+            "coordinate systems beside cartesians ('coord_type: cart')."
+
+        self.freeze_modes = freeze_modes
+        
+        self.sqrt_m = np.sqrt(self.geometry.masses_rep)
+
+    def optimize(self):
+        grad = self.geometry.gradient
+        self.forces.append(-grad.copy())
+        self.energies.append(self.geometry.energy)
+
+        if self.cur_cycle > 0:
+            self.update_trust_radius()
+            self.update_hessian()
+
+        # eigvals_org, eigvecs_org = np.linalg.eigh(self.H)
+        # grad_trans = eigvals_org.T.dot(grad)
+
+        # Mass-weighted hessian
+        H_mw = self.geometry.mass_weigh_hessian(self.H)
+        # Project out translation and rotation
+        if self.geometry.coords.size > 3:
+            H_mw = self.geometry.eckart_projection(H_mw)
+        eigvals, eigvecs = np.linalg.eigh(H_mw)
+        # Drop translational/rotational modes as they will have
+        # small (zero) eigenvalues. Zero the corresponding gradient entries.
+        eigvals, eigvecs, small_inds = self.filter_small_eigvals(eigvals, eigvecs, mask=True)
+
+        wavenumbers = eigval_to_wavenumber(eigvals)
+
+        freeze_inds = np.zeros_like(wavenumbers, dtype=bool)
+        if self.freeze_modes:
+            freeze_inds = wavenumbers < self.freeze_modes
+            eigvals = eigvals[~freeze_inds]
+            eigvecs = eigvecs[:,~freeze_inds]
+            self.log(f"{np.sum(freeze_inds)} normal modes will be frozen.")
+
+        frozen_str = ["(frozen)" if frozen else "" for frozen in freeze_inds]
+        wavenumber_str = "\n".join([
+                            f"\t{i:> 3d}: {wn:> 8.2f} cm⁻¹ {frz}"
+                            for i, (wn, frz) in enumerate(zip(wavenumbers, frozen_str))
+        ])
+        self.log("Frequencies:\n" + wavenumber_str)
+
+        # Transform gradient to eigensystem of Hessian. We also have to use
+        # the mass-weighted gradient.
+        grad_q = eigvecs.T.dot(self.geometry.mw_gradient)
+
+        mw_H_aug = self.get_augmented_hessian(eigvals, grad_q)
+        # Discard eigenvector for now
+        mw_step, eigval, nu, _ = self.solve_rfo(mw_H_aug, "min")
+        # Transform back to original basis. Right now the step is still in
+        # mass-weighted coordinates.
+        mw_step = eigvecs @ mw_step
+        # Un-massweigh step
+        step = mw_step / self.sqrt_m
+
+        step_norm = np.linalg.norm(step)
+
+        if step_norm > self.trust_radius:
+            step = step / step_norm * self.trust_radius
+
+        if self.freeze_modes:
+            # With frozen modes we only want to consider the gradient contributions
+            # from non-frozen modes. So here we transform back the gradient and
+            # un-weigh it.
+            grad = eigvecs.dot(grad_q) * self.sqrt_m
+            self.modified_forces.append(-grad)
+
+        quadratic_prediction = step @ grad + 0.5 * step @ self.H @ step
+        rfo_prediction = quadratic_prediction / (1 + step @ step)
+        self.predicted_energy_changes.append(rfo_prediction)
+
+        return step