add adamax; add --xtol and --ftol

Author: terhorst

smcpp/analysis.py

@@ -37,8 +37,6 @@ def __init__(self, files, args):
             args.polarization_error = 0.5
         if args.polarization_error > 0.:
             logger.info("Polarization error p=%f", args.polarization_error)
-        if args.factr:
-            args.solver_args['factr'] = args.factr
         # Data-related stuff
         self._load_data(files)
         self._validate_data()
@@ -109,7 +107,7 @@ def _perform_thinning(self, thinning):
         elif np.any(ns > 0):
             logger.warn("Not thinning yet undistinguished lineages are present")
 
-    def _normalize_data(self, length_cutoff, no_filter):
+    def _normalize_data(self, length_cutoff, filter):
         ## break up long spans
         self._contigs, attrs = estimation_tools.break_long_spans(self._contigs, length_cutoff)
         if not attrs:
@@ -128,7 +126,7 @@ def _normalize_data(self, length_cutoff, no_filter):
             var = np.average((het - avg) ** 2, weights=w) * (n / (n - 1.))
             sd = np.sqrt(var)
             logger.debug("Average/sd het:%f(%f)", avg, sd)
-            if not no_filter:
+            if filter:
                 logger.debug("Keeping contigs within +-3 s.d. of mean")
         logger.debug("Average heterozygosity (derived / total bases) by data set (* = dropped)")
         ci = 0
@@ -139,7 +137,7 @@ def _normalize_data(self, length_cutoff, no_filter):
             for attr in attrs[key]:
                 het = attr[-1]
                 mytpl = tpl
-                if no_filter or abs(het - avg) <= 3 * sd:
+                if not filter or abs(het - avg) <= 3 * sd:
                     new_contigs.append(self._contigs[ci])
                 else:
                     mytpl += " *"
@@ -282,7 +280,7 @@ def __init__(self, files, args):
                                                        self.rescale(args.tK), 
                                                        knot_spans, args.offset))
         # Perform initial filtering for weird contigs
-        self._normalize_data(args.length_cutoff, args.no_filter)
+        self._normalize_data(args.length_cutoff, args.filter)
 
         # Initialize members
         self._init_parameters(args.theta, args.rho)
@@ -295,7 +293,8 @@ def __init__(self, files, args):
             self._init_optimizer(args, files, args.outdir,
                     1,  # set block-size to knots
                     "L-BFGS-B",  # TNC tends to overfit for initial pass
-                    args.tolerance, learn_rho=False)
+                    args.xtol, args.ftol,
+                    learn_rho=False)
             self._optimizer.run(1)
 
         # Thin the data
@@ -305,7 +304,7 @@ def __init__(self, files, args):
         self._init_hidden_states(args.prior_model, args.M)
         self._init_inference_manager(args.polarization_error)
         self._init_optimizer(args, files, args.outdir, args.blocks,
-                args.algorithm, args.tolerance, learn_rho=True)
+                args.algorithm, args.xtol, args.ftol, learn_rho=True)
 
     def _init_parameters(self, theta=None, rho=None):
         ## Set theta and rho to their default parameters
@@ -363,16 +362,16 @@ def _init_model(self, pieces, N0, t1, tK, spline_class):
                 mods[-1][-1] = y0
             self._model = SMCTwoPopulationModel(mods[0], mods[1], split)
 
-    def _init_optimizer(self, args, files, outdir, blocks, algorithm, tolerance, learn_rho):
+    def _init_optimizer(self, args, files, outdir, blocks, algorithm, xtol, ftol, learn_rho):
         if self.npop == 1:
             self._optimizer = optimizer.SMCPPOptimizer(
-                self, algorithm, tolerance, blocks, args.solver_args)
+                self, algorithm, xtol, ftol, blocks, args.solver_args)
             # Also optimize knots in 1 pop case. Not yet implemented
             # for two pop case.
             # self._optimizer.register(optimizer.KnotOptimizer())
         elif self.npop == 2:
             self._optimizer = optimizer.TwoPopulationOptimizer(
-                self, algorithm, tolerance, blocks, args.solver_args)
+                self, algorithm, xtol, ftol, blocks, args.solver_args)
             smax = np.sum(self._model.distinguished_model.s)
             self._optimizer.register(
                 optimizer.ParameterOptimizer("split", (0., smax), "model"))
@@ -397,26 +396,26 @@ def __init__(self, files, args):
 
         self._hidden_states = np.array([0., np.inf])
         self._init_inference_manager(False)
-        self._init_optimizer(args, files, args.outdir, args.algorithm, args.tolerance, args.blocks, False)
+        self._init_optimizer(args, files, args.outdir, args.algorithm, args.xtol, args.ftol, args.blocks, False)
         # Hack to only estimate split time.
         self._optimizer.run(1)
 
         # After inferring initial split time, thin
         self._perform_thinning(args.thinning)
-        self._normalize_data(args.length_cutoff, args.no_filter)
+        self._normalize_data(args.length_cutoff, args.filter)
 
         self._init_hidden_states(args.pop1, args.M)
         self._init_inference_manager(False)
-        self._init_optimizer(args, files, args.outdir, args.algorithm, args.tolerance, args.blocks)
+        self._init_optimizer(args, files, args.outdir, args.algorithm, args.xtol, args.ftol, args.blocks)
 
     def _validate_data(self):
         BaseAnalysis._validate_data(self)
         if not any(c.npop == 2 for c in self._contigs):
             logger.error("Data contains no joint frequency spectrum information. Split estimation is impossible.")
             sys.exit(1)
 
-    def _init_optimizer(self, args, files, outdir, algorithm, tolerance, blocks, save=True):
-        self._optimizer = optimizer.TwoPopulationOptimizer(self, algorithm, tolerance, blocks, args.solver_args)
+    def _init_optimizer(self, args, files, outdir, algorithm, xtol, ftol, blocks, save=True):
+        self._optimizer = optimizer.TwoPopulationOptimizer(self, algorithm, xtol, ftol, blocks, args.solver_args)
         smax = np.sum(self._model.distinguished_model.s)
         self._optimizer.register(optimizer.ParameterOptimizer("split", (0., smax), "model"))
         if save:

smcpp/commands/command.py

@@ -47,9 +47,7 @@ def add_common_estimation_args(parser):
                       help="omit sequences < cutoff. default: 10000", default=10000, type=int)
     data.add_argument('--thinning', help="only emit full SFS every <k>th site. default: 500 * n.",
                       default=None, type=int, metavar="k")
-    data.add_argument('--no-filter', help="do not drop contigs with extreme heterozygosity. "
-                                          "(not recommended unless data set is small)",
-                      action="store_true", default=False)
+    data.add_argument('--filter', help=argparse.SUPPRESS, action="store_true", default=False)
 
     optimizer = parser.add_argument_group("Optimization parameters")
     optimizer.add_argument(
@@ -60,13 +58,16 @@ def add_common_estimation_args(parser):
                            default="L-BFGS-B", help=argparse.SUPPRESS)
     optimizer.add_argument('--blocks', type=int, 
             help="number of coordinate ascent blocks. default: min(4, K)")
-    optimizer.add_argument('--factr', type=float,
-                           default=1e-9, help=argparse.SUPPRESS)
+    optimizer.add_argument("--ftol", type=float, default=1e-3,
+                           help="stopping criterion for relative improvement in loglik "
+                           "in EM algorithm. algorithm will terminate when "
+                           "|loglik' - loglik| / loglik < ftol")
+    optimizer.add_argument('--xtol', type=float,
+                           default=.001,
+                           help=r"x tolerance for optimizer. "
+                           "optimizer will stop when |x' - x|_\infty < xtol")
     optimizer.add_argument('--regularization-penalty',
                            type=float, help="regularization penalty", default=1.)
-    optimizer.add_argument("--tolerance", type=float, default=1e-4,
-                           help="stopping criterion for relative improvement in loglik "
-                           "in EM algorithm")
     optimizer.add_argument('--Nmin', type=float,
                            help="Lower bound on effective population size (in units of N0)",
                            default=.01)

smcpp/model.py

@@ -92,7 +92,7 @@ def K(self):
         return len(self.knots)
 
     def randomize(self):
-        self[:] += np.random.normal(0., .01, size=len(self[:]))
+        self[:] += np.random.normal(0., .0001, size=len(self[:]))
 
     @property
     def knots(self):

smcpp/optimizer.py

@@ -21,6 +21,32 @@
 
 logger = logging.getLogger(__name__)
 
+def AdaMax(f, x0, args, jac, bounds, alpha=0.1, b1=0.9, b2=0.999, eps=1e-3, **kwargs):
+    assert jac
+    bounds = np.array(bounds)
+    def _f(x0):
+        return tuple(q(x0, *args) for q in (f, jac))
+    obj, grad = _f(x0)
+    theta = x0.copy()
+    t = 0
+    mt = 0
+    ut = 0
+    while True:
+        t += 1
+        ft, gt = _f(theta)
+        mt = b1 * mt + (1. - b1) * gt
+        ut = np.maximum(b2 * ut, abs(gt))
+        delta = -(alpha / (1. - b1 ** t)) * mt / ut
+        if np.linalg.norm(delta) < eps:
+            break
+        theta = box_constrain(theta + delta, bounds)
+        if 'callback' in kwargs:
+            kwargs['callback'](theta)
+    return scipy.optimize.OptimizeResult({'x': theta, 'fun': ft})
+
+class ConvergedException(Exception):
+    "Thrown when optimizer reaches stopping criterion."
+    pass
 
 class EMTerminationException(Exception):
     "Thrown when EM algorithm reaches stopping criterion."
@@ -31,11 +57,12 @@ class AbstractOptimizer(Observable):
     '''
     Abstract representation of the execution flow of the optimizer.
     '''
-    def __init__(self, analysis, algorithm, tolerance, blocks, solver_args={}):
+    def __init__(self, analysis, algorithm, ftol, xtol, blocks, solver_args={}):
         Observable.__init__(self)
         self._analysis = analysis
         self._algorithm = algorithm
-        self._tolerance = tolerance
+        self._ftol = ftol
+        self._xtol = xtol
         self._blocks = blocks
         self._solver_args = solver_args
 
@@ -66,6 +93,7 @@ def _f(self, x, analysis, coords, k=None):
         return ret
 
     def _minimize(self, x0, coords, bounds):
+        self._xk = None
         if os.environ.get("SMCPP_GRADIENT_CHECK", False):
             print("\n\ngradient check")
             y, dy = self._f(x0, self._analysis, coords)
@@ -74,12 +102,23 @@ def _minimize(self, x0, coords, bounds):
                 y1, _ = self._f(x0, self._analysis, coords)
                 print("***grad", i, y1, (y1 - y) * 1e8, dy[i])
                 x0[i] -= 1e-8
-        return minimize_proxy(self._f, x0, 
-                              jac=True,
-                              args=(self._analysis, coords),
-                              bounds=bounds,
-                              options=self._solver_args,
-                              method=self._algorithm)
+        try:
+            if self._algorithm == "AdaMax":
+                alg = AdaMax
+            else:
+                alg = self._algorithm
+            res = scipy.optimize.minimize(self._f, x0,
+                    jac=True,
+                    args=(self._analysis, coords),
+                    bounds=bounds,
+                    options=self._solver_args,
+                    callback=self._callback,
+                    method=alg)
+            return res
+        except ConvergedException:
+            logger.debug("Converged: |xk - xk_1| < %g", self._xtol)
+            return scipy.optimize.OptimizeResult(
+                {'x': self._xk, 'fun': self._f(self._xk, self._analysis, coords)[0]})
 
     def run(self, niter):
         self.update_observers('begin')
@@ -108,6 +147,15 @@ def run(self, niter):
         # Conclude the optimization and perform any necessary callbacks.
         self.update_observers('optimization finished')
 
+    def _callback(self, xk):
+        if self._xk is None:
+            self._xk = xk
+            return
+        delta = max(abs(xk - self._xk))
+        self._xk = xk
+        if delta < self._xtol:
+            raise ConvergedException()
+
     def update_observers(self, *args, **kwargs):
         kwargs.update({
             'optimizer': self,
@@ -158,7 +206,7 @@ def update(self, message, *args, **kwargs):
             improvement = (self._old_loglik - ll) / self._old_loglik
             logger.info("New loglik: %f\t(old: %f [%f%%])",
                     ll, self._old_loglik, 100. * improvement)
-            tol = kwargs['optimizer']._tolerance
+            tol = kwargs['optimizer']._ftol
             if improvement < 0:
                 logger.warn("Loglik decreased")
             elif improvement < tol:
@@ -362,8 +410,8 @@ def write(x):
 class SMCPPOptimizer(AbstractOptimizer):
     'Model fitting for one population.'
 
-    def __init__(self, analysis, algorithm, tolerance, blocks, solver_args):
-        AbstractOptimizer.__init__(self, analysis, algorithm, tolerance, blocks, solver_args)
+    def __init__(self, analysis, algorithm, xtol, ftol, blocks, solver_args):
+        AbstractOptimizer.__init__(self, analysis, algorithm, xtol, ftol, blocks, solver_args)
         observers = [
             HiddenStateOccupancyPrinter(),
             ProgressPrinter(),
@@ -405,32 +453,5 @@ def _coordinates(self):
     def _bounds(self, coords):
         return SMCPPOptimizer._bounds(self, coords[1])
 
-AdaMaxResult = namedtuple('AdaMaxResult', 'x fun')
-
 def box_constrain(x, bounds):
     return np.maximum(np.minimum(x, bounds[:, 1]), bounds[:, 0])
-
-def AdaMax(f, x0, jac, args, bounds, alpha=0.0002, b1=0.9, b2=0.999, eps=1e-3, **kwargs):
-    assert jac == True
-    bounds = np.array(bounds)
-    obj, grad = f(x0, *args)
-    m0 = 0
-    u0 = 0
-    theta = x0.copy()
-    t = 0
-    mt = 0
-    while True:
-        t += 1
-        ft, gt = f(theta, *args)
-        mt = b1 * mt + (1. - b1) * gt
-        ut = np.maximum(b2, np.abs(gt))
-        delta = -(alpha / (1. - b1 ** t)) * mt / ut
-        if np.linalg.norm(delta) < eps:
-            break
-        theta = box_constrain(theta + delta, bounds)
-    return AdaMaxResult(x=theta, fun=ft)
-
-def minimize_proxy(f, x0, *args, **kwargs):
-    if kwargs['method'] == "AdaMax":
-        return AdaMax(f, x0, *args, **kwargs)
-    return scipy.optimize.minimize(f, x0, *args, **kwargs)