Overhauls Model class used for model/gradient manipulation

seisflows/examples/ex2_hh_w_pyatoa.py

@@ -83,8 +83,8 @@ def __init__(self, ntask=None, niter=None, nsta=None, nproc=None,
         }
 
         # Adjust the existing parameter list
-        self._parameters["smooth_h"] = 5000.
-        self._parameters["smooth_v"] = 5000.
+        self._parameters["smooth_h"] = 20000.
+        self._parameters["smooth_v"] = 20000.
 
         # Pyaflowa preprocessing parameters
         self._parameters["min_period"] = 10.  # filter bounds define windows

seisflows/optimize/NLCG.py

@@ -7,6 +7,8 @@
 from seisflows import logger
 from seisflows.optimize.gradient import Gradient
 from seisflows.tools.math import dot
+from seisflows.tools.specfem_model import Model
+from seisflows.tools import msg
 from seisflows.plugins import line_search as line_search_dir
 
 
@@ -33,11 +35,15 @@ class NLCG(Gradient):
     """
     __doc__ = Gradient.__doc__ + __doc__
 
-    def __init__(self, nlcg_max=np.inf, nlcg_thresh=np.inf,
+    def __init__(self, nlcg_max=np.inf, nlcg_thresh=np.inf, 
                  calc_beta="pollak_ribere", **kwargs):
         """NLCG-specific input parameters"""
         super().__init__(**kwargs)
 
+        # Be honest!
+        logger.warning(msg.mjr("THIS OPTIMIZATION MODULE IS NOT WELL TESTED "
+                               "USE AT YOUR OWN RISK"))
+
         # Overwrite user-chosen line search. L-BFGS requires 'Backtrack'ing LS
         if self.line_search_method.title() != "Bracket":
             logger.warning(f"NLCG optimization requires 'bracket'ing line "
@@ -53,7 +59,7 @@ def __init__(self, nlcg_max=np.inf, nlcg_thresh=np.inf,
         self.NLCG_thresh = nlcg_thresh
 
         # Check paramter validity
-        _acceptable_calc_beta = ["pollak_ribere", "fletcher_reeves"]
+        _acceptable_calc_beta = ["fletcher_reeves"]  # "pollak_ribere" 
         assert(calc_beta in _acceptable_calc_beta), (
             f"unacceptable `calc_beta`, must be in {_acceptable_calc_beta}"
         )
@@ -72,7 +78,7 @@ def checkpoint(self):
         checkpoint_dict = np.load(self.path._checkpoint)
         checkpoint_dict["NLCG_iter"] = self._NLCG_iter
 
-        np.savez(file=self.path._checkpoint, **dict_out)  # NOQA
+        np.savez(file=self.path._checkpoint, **checkpoint_dict)  # NOQA
 
     def compute_direction(self):
         """
@@ -97,54 +103,57 @@ def compute_direction(self):
         self._NLCG_iter += 1
 
         # Load the current gradient direction
-        g_new = self.load_vector("g_new")
-        p_new = g_new.copy()
+        g_new = Model(self.path._g_new)
+        g_old = Model(self.path._g_old)
+        p_old = Model(self.path._p_old)
 
         # CASE 1: If first iteration, search direction is the current gradient
         if self._NLCG_iter == 1:
             logger.info("first NLCG iteration, setting search direction "
                         "as inverse gradient")
-            p_new.update(vector=-1 * g_new.vector)
-            restarted = False
+            # p_new = -g
+            return g_new.apply(actions=["*"], values=[-1], 
+                               export_to=self.path._p_new)
         # CASE 2: Force restart if the iterations have surpassed the maximum
         # number of allowable iter
         elif self._NLCG_iter > self.NLCG_max:
-            logger.info("restarting NLCG due to periodic restart "
-                        "condition. setting search direction as inverse "
-                        "gradient")
+            logger.info("restarting NLCG due to periodic restart condition. "
+                        "setting search direction as inverse gradient")
             self.restart()
-            p_new.update(vector=-1 * g_new.vector)
-            restarted = True
-        # Normal NLCG direction compuitation
+            # p_new = -g
+            return g_new.apply(actions=["*"], values=[-1], 
+                               export_to=self.path._p_new)
+        # Normal NLCG direction computation
         else:
-            # Compute search direction
-            g_old = self.load_vector("g_old")
-            p_old = self.load_vector("p_old")
-            beta = self._calc_beta(g_new.vector, g_old.vector)
-
-            # Apply preconditioner and calc. scale factor for search dir. (beta)
-            _p_new_vec = (
-                    -1 * self._precondition(g_new.vector) + beta * p_old.vector
-            )
-            p_new.update(vector=_p_new_vec)
-
-            # Check restart conditions, return search direction and statusa
-            if check_conjugacy(g_new.vector, g_old.vector) > self.NLCG_thresh:
+            if check_conjugacy(g_new, g_old) > self.NLCG_thresh:
                 logger.info("restarting NLCG due to loss of conjugacy")
                 self.restart()
-                p_new.update(vector=-1 * g_new.vector)
-                restarted = True
-            elif check_descent(p_new.vector, g_new.vector) > 0.:
+                # p_new = -g
+                return g_new.apply(actions=["*"], values=[-1], 
+                                   export_to=self.path._p_new)
+
+            # Compute beta, the scaling factor for the search direction. User
+            # has chosen a specific underlying function to do this
+            beta = self._calc_beta()
+
+            # Optional step, apply preconditioner in place
+            if self.preconditioner:
+                g_new = self.precondition(q=g_new)
+
+            # p_new = -1 * g + B * p_old (couldn't figure out how to do this in
+            # one action so we apply three)
+            p_new = g_new.apply(actions=["*"], values=[-1])
+            p_old = p_old.apply(actions=["*"], values=[beta], 
+                                export_to=self.path._scratch)
+            p_new = p_new.apply(actions=["+"], values=[p_old])
+
+            # Check restart conditions, return search direction and statusa
+            if check_descent(p_new.get("vector"), g_new.get("vector")) > 0.:
                 logger.info("restarting NLCG, not a descent direction")
                 self.restart()
-                p_new.update(vector=-1 * g_new.vector)
-                restarted = True
-            else:
-                # p_new = p_new
-                restarted = False
-
-        # Save values to disk and memory
-        self._restarted = restarted
+                # p_new = -g
+                return g_new.apply(actions=["*"], values=[-1], 
+                                   export_to=self.path._p_new)
 
         return p_new
 
@@ -153,65 +162,76 @@ def restart(self):
         Overwrite the Base restart class and include a restart of the NLCG
         """
         logger.info("restarting NLCG optimization algorithm")
-
-        g_new = self.load_vector("g_new")
-        p_new = g_new.copy()
-        p_new.update(vector=-1 * g_new.vector)
-        self.save_vector("p_new", p_new)
-
         self._line_search.clear_search_history()
-        self._restarted = 1
+        self._restarted = True
         self._NLCG_iter = 1
 
-    def _fletcher_reeves(self, g_new, g_old):
+    def _fletcher_reeves(self):
         """
-        One method for calculating beta in the NLCG Algorithm from
+        Calculating scale factor beta in the NLCG Algorithm following:
         Fletcher & Reeves, 1964
 
-        :type g_new: np.array
-        :param g_new: new search direction
-        :type g_old: np.array
-        :param g_old: old search direction
         :rtype: float
         :return: beta, the scale factor to apply to the old search direction to
             determine the new search direction
         """
-        num = dot(self._precondition(g_new), g_new)
-        den = dot(g_old, g_old)
-        beta = num / den
-        return beta
+        # Use current and old gradients
+        g_new = Model(self.path._g_new)
+        g_old = Model(self.path._g_old)
 
-    def _pollak_ribere(self, g_new, g_old):
+        if self.preconditioner:
+            g_new_prcnd = self.precondition(g_new, export_to=self.path._scratch)
+        else:
+            g_new_prcnd = g_new
+
+        return g_new_prcnd.dot(g_new) / g_old.dot(g_old)
+
+    def _pollak_ribere(self):
         """
-        One method for calculating beta in the NLCG Algorithm from
+        Calculating scale factor beta in the NLCG Algorithm following:
         Polak & Ribiere, 1969
 
-        :type g_new: np.array
-        :param g_new: new search direction
-        :type g_old: np.array
-        :param g_old: old search direction
         :rtype: float
         :return: beta, the scale factor to apply to the old search direction to
             determine the new search direction
         """
-        num = dot(self._precondition(g_new), g_new - g_old)
-        den = dot(g_old, g_old)
-        beta = num / den
-        return beta
+        # BC Did not have time to finish updating this, need to figure out how
+        # to store two scratch variables at the same time. See comment below
+        raise NotImplementedError("This function is currently not working, if "
+                                  "you would like to use it, please open a " 
+                                  "GitHub issue")
+
+        # Use current and old gradients
+        g_new = Model(self.path._g_new)
+        g_old = Model(self.path._g_old)
+        p_old = Model(self.path._p_old)
+
+        if self.preconditioner:
+            g_new_prcnd = self.precondition(g_new, export_to=self.path._scratch)
+        else:
+            g_new_prcnd = g_new
 
+        # This wont work because both `dg` and `g_new_prcnd` need to be stored
+        # in scratch. I guess I could make another scratch directory or store it
+        # somewhere else but I feel like there is a cleaner way
+        dg = g_new.apply(action=["-"], values=g_old, 
+                         export_to=self.path._scratch)
+
+        return g_new_prcndt.dot(dg) / g_old.dot(g_old)
+
 
 def check_conjugacy(g_new, g_old):
     """
     Check for conjugacy between two vectors
 
-    :type g_new: np.array
+    :type g_new: seisflows.tools.specfem_model.Model
     :param g_new: new search direction
-    :type g_old: np.array
+    :type g_old: seisflows.tools.specfem_model.Model
     :param g_old: old search direction
     :rtype: float
     :return: an element that proves conjugacy
     """
-    return abs(dot(g_new, g_old) / dot(g_new, g_new))
+    return abs(g_new.dot(g_old) / g_new.dot(g_new))
 
 
 def check_descent(p_new, g_new):
@@ -226,7 +246,7 @@ def check_descent(p_new, g_new):
     :return: the angle between search direction and gradient direction, should
         be negative to ensure descent
     """
-    return dot(p_new, g_new) / dot(g_new, g_new)
+    return p_new.dot(g_new) / g_new.dot(g_new)