pyGSTi 0.9.12__cp38-cp38-win32.whl → 0.9.13__cp38-cp38-win32.whl

pygsti/optimize/simplerlm.py

@@ -0,0 +1,841 @@
+"""
+Custom implementation of the Levenberg-Marquardt Algorithm (but simpler than customlm.py)
+"""
+#***************************************************************************************************
+# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
+# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
+# in this software.
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+# in compliance with the License.  You may obtain a copy of the License at
+# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
+#***************************************************************************************************
+
+import os as _os
+import signal as _signal
+import time as _time
+
+import numpy as _np
+import scipy as _scipy
+
+from pygsti.optimize import arraysinterface as _ari
+from pygsti.optimize.customsolve import custom_solve as _custom_solve
+from pygsti.baseobjs.verbosityprinter import VerbosityPrinter as _VerbosityPrinter
+from pygsti.baseobjs.resourceallocation import ResourceAllocation as _ResourceAllocation
+from pygsti.baseobjs.nicelyserializable import NicelySerializable as _NicelySerializable
+from pygsti.objectivefns.objectivefns import Chi2Function, TimeIndependentMDCObjectiveFunction
+from typing import Callable
+
+#Make sure SIGINT will generate a KeyboardInterrupt (even if we're launched in the background)
+#This may be problematic for multithreaded parallelism above pyGSTi, e.g. Dask,
+#so this can be turned off by setting the PYGSTI_NO_CUSTOMLM_SIGINT environment variable
+if 'PYGSTI_NO_CUSTOMLM_SIGINT' not in _os.environ:
+    _signal.signal(_signal.SIGINT, _signal.default_int_handler)
+
+#constants
+_MACH_PRECISION = 1e-12
+
+
+class OptimizerResult(object):
+    """
+    The result from an optimization.
+
+    Parameters
+    ----------
+    objective_func : ObjectiveFunction
+        The objective function that was optimized.
+
+    opt_x : numpy.ndarray
+        The optimal argument (x) value.  Often a vector of parameters.
+
+    opt_f : numpy.ndarray
+        the optimal objective function (f) value.  Often this is the least-squares
+        vector of objective function values.
+
+    opt_jtj : numpy.ndarray, optional
+        the optimial `dot(transpose(J),J)` value, where `J`
+        is the Jacobian matrix.  This may be useful for computing
+        approximate error bars.
+
+    opt_unpenalized_f : numpy.ndarray, optional
+        the optimal objective function (f) value with any
+        penalty terms removed.
+
+    chi2_k_distributed_qty : float, optional
+        a value that is supposed to be chi2_k distributed.
+
+    optimizer_specific_qtys : dict, optional
+        a dictionary of additional optimization parameters.
+    """
+    def __init__(self, objective_func, opt_x, opt_f=None, opt_jtj=None,
+                 opt_unpenalized_f=None, chi2_k_distributed_qty=None,
+                 optimizer_specific_qtys=None):
+        self.objective_func = objective_func
+        self.x = opt_x
+        self.f = opt_f
+        self.jtj = opt_jtj  # jacobian.T * jacobian
+        self.f_no_penalties = opt_unpenalized_f
+        self.optimizer_specific_qtys = optimizer_specific_qtys
+        self.chi2_k_distributed_qty = chi2_k_distributed_qty
+
+
+class Optimizer(_NicelySerializable):
+    """
+    An optimizer.  Optimizes an objective function.
+    """
+
+    @classmethod
+    def cast(cls, obj):
+        """
+        Cast `obj` to a :class:`Optimizer`.
+
+        If `obj` is already an `Optimizer` it is just returned,
+        otherwise this function tries to create a new object
+        using `obj` as a dictionary of constructor arguments.
+
+        Parameters
+        ----------
+        obj : Optimizer or dict
+            The object to cast.
+
+        Returns
+        -------
+        Optimizer
+        """
+        if isinstance(obj, cls):
+            return obj
+        else:
+            return cls(**obj) if obj else cls()
+
+    def __init__(self):
+        super().__init__()
+
+
+class SimplerLMOptimizer(Optimizer):
+    """
+    A Levenberg-Marquardt optimizer customized for GST-like problems.
+
+    Parameters
+    ----------
+    maxiter : int, optional
+        The maximum number of (outer) interations.
+
+    maxfev : int, optional
+        The maximum function evaluations.
+
+    tol : float or dict, optional
+        The tolerance, specified as a single float or as a dict
+        with keys `{'relx', 'relf', 'jac', 'maxdx'}`.  A single
+        float sets the `'relf'` and `'jac'` elemments and leaves
+        the others at their default values.
+
+    fditer : int optional
+        Internally compute the Jacobian using a finite-difference method
+        for the first `fditer` iterations.  This is useful when the initial
+        point lies at a special or singular point where the analytic Jacobian
+        is misleading.
+
+    first_fditer : int, optional
+        Number of finite-difference iterations applied to the first
+        stage of the optimization (only).  Unused.
+
+    init_munu : tuple, optional
+        If not None, a (mu, nu) tuple of 2 floats giving the initial values
+        for mu and nu.
+
+    oob_check_interval : int, optional
+        Every `oob_check_interval` outer iterations, the objective function
+        (`obj_fn`) is called with a second argument 'oob_check', set to True.
+        In this case, `obj_fn` can raise a ValueError exception to indicate
+        that it is Out Of Bounds.  If `oob_check_interval` is 0 then this
+        check is never performed; if 1 then it is always performed.
+
+    oob_action : {"reject","stop"}
+        What to do when the objective function indicates (by raising a ValueError
+        as described above).  `"reject"` means the step is rejected but the
+        optimization proceeds; `"stop"` means the optimization stops and returns
+        as converged at the last known-in-bounds point.
+
+    oob_check_mode : int, optional
+        An advanced option, expert use only.  If 0 then the optimization is
+        halted as soon as an *attempt* is made to evaluate the function out of bounds.
+        If 1 then the optimization is halted only when a would-be *accepted* step
+        is out of bounds.
+
+    serial_solve_proc_threshold : int, optional
+        When there are fewer than this many processors, the optimizer will solve linear
+        systems serially, using SciPy on a single processor, rather than using a parallelized
+        Gaussian Elimination (with partial pivoting) algorithm coded in Python. Since SciPy's
+        implementation is more efficient, it's not worth using the parallel version until there
+        are many processors to spread the work among.
+
+    lsvec_mode : {'normal', 'percircuit'}
+        Whether the terms used in the least-squares optimization are the "elements" as computed
+        by the objective function's `.terms()` and `.lsvec()` methods (`'normal'` mode) or the
+        "per-circuit quantities" computed by the objective function's `.percircuit()` and
+        `.lsvec_percircuit()` methods (`'percircuit'` mode).
+    """
+
+    @classmethod
+    def cast(cls, obj):
+        if isinstance(obj, cls):
+            return obj
+        if obj:
+            try:
+                return cls(**obj)
+            except:
+                from pygsti.optimize.customlm import CustomLMOptimizer
+                return CustomLMOptimizer(**obj)
+        return cls()
+
+    def __init__(self, maxiter=100, maxfev=100, tol=1e-6, fditer=0, first_fditer=0, init_munu="auto", oob_check_interval=0,
+                 oob_action="reject", oob_check_mode=0, serial_solve_proc_threshold=100, lsvec_mode="normal"):
+
+        super().__init__()
+        if isinstance(tol, float): tol = {'relx': 1e-8, 'relf': tol, 'f': 1.0, 'jac': tol, 'maxdx': 1.0}
+        self.maxiter = maxiter
+        self.maxfev = maxfev
+        self.tol = tol
+        self.fditer = fditer
+        self.first_fditer = first_fditer
+        self.init_munu = init_munu
+        self.oob_check_interval = oob_check_interval
+        self.oob_action = oob_action
+        self.oob_check_mode = oob_check_mode
+        self.array_types = 3 * ('p',) + ('e', 'ep')  # see simplish_leastsq fn "-type"s  -need to add 'jtj' type
+        self.called_objective_methods = ('lsvec', 'dlsvec')  # the objective function methods we use (for mem estimate)
+        self.serial_solve_proc_threshold = serial_solve_proc_threshold
+        self.lsvec_mode = lsvec_mode
+
+    def _to_nice_serialization(self):
+        state = super()._to_nice_serialization()
+        state.update({
+            'maximum_iterations': self.maxiter,
+            'maximum_function_evaluations': self.maxfev,
+            'tolerance': self.tol,
+            'number_of_finite_difference_iterations': self.fditer,
+            'number_of_first_stage_finite_difference_iterations': self.first_fditer,
+            'initial_mu_and_nu': self.init_munu,
+            'out_of_bounds_check_interval': self.oob_check_interval,
+            'out_of_bounds_action': self.oob_action,
+            'out_of_bounds_check_mode': self.oob_check_mode,
+            'array_types': self.array_types,
+            'called_objective_function_methods': self.called_objective_methods,
+            'serial_solve_number_of_processors_threshold': self.serial_solve_proc_threshold,
+            'lsvec_mode': self.lsvec_mode
+        })
+        return state
+
+    @classmethod
+    def _from_nice_serialization(cls, state):
+        return cls(maxiter=state['maximum_iterations'],
+                   maxfev=state['maximum_function_evaluations'],
+                   tol=state['tolerance'],
+                   fditer=state['number_of_finite_difference_iterations'],
+                   first_fditer=state['number_of_first_stage_finite_difference_iterations'],
+                   init_munu=state['initial_mu_and_nu'],
+                   oob_check_interval=state['out_of_bounds_check_interval'],
+                   oob_action=state['out_of_bounds_action'],
+                   oob_check_mode=state['out_of_bounds_check_mode'],
+                   serial_solve_proc_threshold=state['serial_solve_number_of_processors_threshold'],
+                   lsvec_mode=state.get('lsvec_mode', 'normal'))
+
+    def run(self, objective: TimeIndependentMDCObjectiveFunction, profiler, printer):
+
+        """
+        Perform the optimization.
+
+        Parameters
+        ----------
+        objective : ObjectiveFunction
+            The objective function to optimize.
+
+        profiler : Profiler
+            A profiler to track resource usage.
+
+        printer : VerbosityPrinter
+            printer to use for sending output to stdout.
+        """
+        nExtra = objective.ex  # number of additional "extra" elements
+
+        if self.lsvec_mode == 'normal':
+            objective_func = objective.lsvec
+            jacobian = objective.dlsvec
+            nEls = objective.layout.num_elements + nExtra  # 'e' for array types
+        elif self.lsvec_mode == 'percircuit':
+            objective_func = objective.lsvec_percircuit
+            jacobian = objective.dlsvec_percircuit
+            nEls = objective.layout.num_circuits + nExtra  # 'e' for array types
+        else:
+            raise ValueError("Invalid `lsvec_mode`: %s" % str(self.lsvec_mode))
+
+        x0 = objective.model.to_vector()
+        x_limits = objective.model.parameter_bounds
+        # x_limits should be a (num_params, 2)-shaped array, holding on each row the (min, max) values for the
+        #  corresponding parameter (element of the "x" vector) or `None`.  If `None`, then no limits are imposed.
+
+        # Check memory limit can handle what simplish_leastsq will "allocate"
+        nP = len(x0)  # 'p' for array types
+        objective.resource_alloc.check_can_allocate_memory(3 * nP + nEls + nEls * nP + nP * nP)  # see array_types above
+
+        from ..layouts.distlayout import DistributableCOPALayout as _DL
+        if isinstance(objective.layout, _DL):
+            ari = _ari.DistributedArraysInterface(objective.layout, self.lsvec_mode, nExtra)
+        else:
+            ari = _ari.UndistributedArraysInterface(nEls, nP)
+
+        opt_x, converged, msg, mu, nu, norm_f, f = simplish_leastsq(
+            objective_func, jacobian, x0,
+            max_iter=self.maxiter,
+            num_fd_iters=self.fditer,
+            f_norm2_tol=self.tol.get('f', 1.0),
+            jac_norm_tol=self.tol.get('jac', 1e-6),
+            rel_ftol=self.tol.get('relf', 1e-6),
+            rel_xtol=self.tol.get('relx', 1e-8),
+            max_dx_scale=self.tol.get('maxdx', 1.0),
+            init_munu=self.init_munu,
+            oob_check_interval=self.oob_check_interval,
+            oob_action=self.oob_action,
+            oob_check_mode=self.oob_check_mode,
+            resource_alloc=objective.resource_alloc,
+            arrays_interface=ari,
+            serial_solve_proc_threshold=self.serial_solve_proc_threshold,
+            x_limits=x_limits,
+            verbosity=printer - 1, profiler=profiler)
+
+        printer.log("Least squares message = %s" % msg, 2)
+        assert(converged), "Failed to converge: %s" % msg
+        current_v = objective.model.to_vector()
+        if not _np.allclose(current_v, opt_x):  # ensure the last model evaluation was at opt_x
+            objective_func(opt_x)
+            #objective.model.from_vector(opt_x)  # performed within line above
+
+        #DEBUG CHECK SYNC between procs (especially for shared mem) - could REMOVE
+        # if objective.resource_alloc.comm is not None:
+        #     comm = objective.resource_alloc.comm
+        #     v_cmp = comm.bcast(objective.model.to_vector() if (comm.Get_rank() == 0) else None, root=0)
+        #     v_matches_x = _np.allclose(objective.model.to_vector(), opt_x)
+        #     same_as_root = _np.isclose(_np.linalg.norm(objective.model.to_vector() - v_cmp), 0.0)
+        #     if not (v_matches_x and same_as_root):
+        #         raise ValueError("Rank %d CUSTOMLM ERROR: END model vector-matches-x=%s and vector-is-same-as-root=%s"
+        #                          % (comm.rank, str(v_matches_x), str(same_as_root)))
+        #     comm.barrier()  # if we get past here, then *all* processors are OK
+        #     if comm.rank == 0:
+        #         print("OK - model vector == best_x and all vectors agree w/root proc's")
+
+        unpenalized_f = f[0:-objective.ex] if (objective.ex > 0) else f
+        unpenalized_normf = sum(unpenalized_f**2)  # objective function without penalty factors
+        chi2k_qty = objective.chi2k_distributed_qty(norm_f)
+        optimizer_specific_qtys = {'msg': msg, 'mu': mu, 'nu': nu, 'fvec': f}
+        return OptimizerResult(objective, opt_x, norm_f, None, unpenalized_normf, chi2k_qty, optimizer_specific_qtys)
+
+
+
+def damp_coeff_update(mu, nu, half_max_nu, reject_msg, printer):
+    ############################################################################################
+    #
+    #   if this point is reached, either the linear solve failed
+    #   or the error did not reduce.  In either case, reject increment.
+    #
+    ############################################################################################
+    mu *= nu
+    if nu > half_max_nu:  # watch for nu getting too large (&overflow)
+        msg = "Stopping after nu overflow!"
+    else:
+        msg = ""
+    nu = 2 * nu
+    printer.log("      Rejected%s!  mu => mu*nu = %g, nu => 2*nu = %g" % (reject_msg, mu, nu), 2)
+    return mu, nu, msg
+
+
+def jac_guarded(k: int, num_fd_iters: int, obj_fn: Callable, jac_fn: Callable, f, ari, global_x, fdJac_work):
+    if k >= num_fd_iters:
+        Jac = jac_fn(global_x)  # 'EP'-type, but doesn't actually allocate any more mem (!)
+    else:
+        # Note: x holds only number of "fine"-division params - need to use global_x, and
+        # Jac only holds a subset of the derivative and element columns and rows, respectively.
+        f_fixed = f.copy()  # a static part of the distributed `f` resturned by obj_fn - MUST copy this.
+
+        pslice = ari.jac_param_slice(only_if_leader=True)
+        eps = 1e-7
+        #Don't do this: for ii, i in enumerate(range(pslice.start, pslice.stop)): (must keep procs in sync)
+        for i in range(len(global_x)):
+            x_plus_dx = global_x.copy()
+            x_plus_dx[i] += eps
+            fd = (obj_fn(x_plus_dx) - f_fixed) / eps
+            if pslice.start <= i < pslice.stop:
+                fdJac_work[:, i - pslice.start] = fd
+            #if comm is not None: comm.barrier()  # overkill for shared memory leader host barrier
+        Jac = fdJac_work
+    return Jac
+
+
+
+def simplish_leastsq(
+    obj_fn, jac_fn, x0, f_norm2_tol=1e-6, jac_norm_tol=1e-6,
+    rel_ftol=1e-6, rel_xtol=1e-6, max_iter=100, num_fd_iters=0, max_dx_scale=1.0,
+    init_munu="auto", oob_check_interval=0, oob_action="reject", oob_check_mode=0,
+    resource_alloc=None, arrays_interface=None, serial_solve_proc_threshold=100,
+    x_limits=None, verbosity=0, profiler=None
+    ):
+    """
+    An implementation of the Levenberg-Marquardt least-squares optimization algorithm customized for use within pyGSTi.
+
+    This general purpose routine mimic to a large extent the interface used by
+    `scipy.optimize.leastsq`, though it implements a newer (and more robust) version
+    of the algorithm.
+
+    Parameters
+    ----------
+    obj_fn : function
+        The objective function.  Must accept and return 1D numpy ndarrays of
+        length N and M respectively.  Same form as scipy.optimize.leastsq.
+
+    jac_fn : function
+        The jacobian function (not optional!).  Accepts a 1D array of length N
+        and returns an array of shape (M,N).
+
+    x0 : numpy.ndarray
+        Initial evaluation point.
+
+    f_norm2_tol : float, optional
+        Tolerace for `F^2` where `F = `norm( sum(obj_fn(x)**2) )` is the
+        least-squares residual.  If `F**2 < f_norm2_tol`, then mark converged.
+
+    jac_norm_tol : float, optional
+        Tolerance for jacobian norm, namely if `infn(dot(J.T,f)) < jac_norm_tol`
+        then mark converged, where `infn` is the infinity-norm and
+        `f = obj_fn(x)`.
+
+    rel_ftol : float, optional
+        Tolerance on the relative reduction in `F^2`, that is, if
+        `d(F^2)/F^2 < rel_ftol` then mark converged.
+
+    rel_xtol : float, optional
+        Tolerance on the relative value of `|x|`, so that if
+        `d(|x|)/|x| < rel_xtol` then mark converged.
+
+    max_iter : int, optional
+        The maximum number of (outer) interations.
+
+    num_fd_iters : int optional
+        Internally compute the Jacobian using a finite-difference method
+        for the first `num_fd_iters` iterations.  This is useful when `x0`
+        lies at a special or singular point where the analytic Jacobian is
+        misleading.
+
+    max_dx_scale : float, optional
+        If not None, impose a limit on the magnitude of the step, so that
+        `|dx|^2 < max_dx_scale^2 * len(dx)` (so elements of `dx` should be,
+        roughly, less than `max_dx_scale`).
+
+    init_munu : tuple, optional
+        If not None, a (mu, nu) tuple of 2 floats giving the initial values
+        for mu and nu.
+
+    oob_check_interval : int, optional
+        Every `oob_check_interval` outer iterations, the objective function
+        (`obj_fn`) is called with a second argument 'oob_check', set to True.
+        In this case, `obj_fn` can raise a ValueError exception to indicate
+        that it is Out Of Bounds.  If `oob_check_interval` is 0 then this
+        check is never performed; if 1 then it is always performed.
+
+    oob_action : {"reject","stop"}
+        What to do when the objective function indicates (by raising a ValueError
+        as described above).  `"reject"` means the step is rejected but the
+        optimization proceeds; `"stop"` means the optimization stops and returns
+        as converged at the last known-in-bounds point.
+
+    oob_check_mode : int, optional
+        An advanced option, expert use only.  If 0 then the optimization is
+        halted as soon as an *attempt* is made to evaluate the function out of bounds.
+        If 1 then the optimization is halted only when a would-be *accepted* step
+        is out of bounds.
+
+    resource_alloc : ResourceAllocation, optional
+        When not None, an resource allocation object used for distributing the computation
+        across multiple processors.
+
+    arrays_interface : ArraysInterface
+        An object that provides an interface for creating and manipulating data arrays.
+
+    serial_solve_proc_threshold : int optional
+        When there are fewer than this many processors, the optimizer will solve linear
+        systems serially, using SciPy on a single processor, rather than using a parallelized
+        Gaussian Elimination (with partial pivoting) algorithm coded in Python. Since SciPy's
+        implementation is more efficient, it's not worth using the parallel version until there
+        are many processors to spread the work among.
+
+    x_limits : numpy.ndarray, optional
+        A (num_params, 2)-shaped array, holding on each row the (min, max) values for the corresponding
+        parameter (element of the "x" vector).  If `None`, then no limits are imposed.
+
+    verbosity : int, optional
+        Amount of detail to print to stdout.
+
+    profiler : Profiler, optional
+        A profiler object used for to track timing and memory usage.
+
+    Returns
+    -------
+    x : numpy.ndarray
+        The optimal solution.
+    converged : bool
+        Whether the solution converged.
+    msg : str
+        A message indicating why the solution converged (or didn't).
+    """
+    resource_alloc = _ResourceAllocation.cast(resource_alloc)
+    comm = resource_alloc.comm
+    printer = _VerbosityPrinter.create_printer(verbosity, comm)
+    ari = arrays_interface  # shorthand
+
+    msg = ""
+    converged = False
+    half_max_nu = 2**62  # what should this be??
+    tau = 1e-3
+
+    #Allocate potentially shared memory used in loop
+    JTJ = ari.allocate_jtj()
+    minus_JTf = ari.allocate_jtf()
+    x = ari.allocate_jtf()
+    best_x = ari.allocate_jtf()
+    dx = ari.allocate_jtf()
+    new_x = ari.allocate_jtf()
+    optional_jtj_buff = ari.allocate_jtj_shared_mem_buf()
+    fdJac = ari.allocate_jac() if num_fd_iters > 0 else None
+
+    global_x = x0.copy()
+    ari.allscatter_x(global_x, x)
+    global_new_x = global_x.copy()
+    best_x[:] = x[:] 
+    # ^ like x.copy() -the x-value corresponding to min_norm_f ('P'-type)
+
+    if x_limits is not None:
+        x_lower_limits = ari.allocate_jtf()
+        x_upper_limits = ari.allocate_jtf()
+        ari.allscatter_x(x_limits[:, 0], x_lower_limits)
+        ari.allscatter_x(x_limits[:, 1], x_upper_limits)
+    max_norm_dx = (max_dx_scale**2) * len(global_x) if max_dx_scale else None
+    # ^ don't let any component change by more than ~max_dx_scale
+
+
+    f = obj_fn(global_x)  # 'E'-type array
+    norm_f = ari.norm2_f(f)
+    if not _np.isfinite(norm_f):
+        msg = "Infinite norm of objective function at initial point!"
+
+    if len(global_x) == 0:  # a model with 0 parameters - nothing to optimize
+        msg = "No parameters to optimize"
+        converged = True
+
+    mu, nu =  (1, 2) if init_munu == 'auto' else init_munu
+    # ^ We have to set some *some* values in case we exit at the start of the first
+    #   iteration. mu will almost certainly be overwritten before being read.
+    min_norm_f = 1e100  # sentinel
+    best_x_state = (mu, nu, norm_f, f.copy())
+    # ^ here and elsewhere, need f.copy() b/c f is objfn mem
+
+    try:
+
+        for k in range(max_iter):  # outer loop
+            # assume global_x, x, f, fnorm hold valid values
+
+            if len(msg) > 0:
+                break  # exit outer loop if an exit-message has been set
+
+            if norm_f < f_norm2_tol:
+                if oob_check_interval <= 1:
+                    msg = "Sum of squares is at most %g" % f_norm2_tol
+                    converged = True
+                    break
+                else:
+                    printer.log(("** Converged with out-of-bounds with check interval=%d, reverting to last know in-bounds point and setting interval=1 **") % oob_check_interval, 2)
+                    oob_check_interval = 1
+                    x[:] = best_x[:]
+                    mu, nu, norm_f, f[:] = best_x_state
+                    continue
+
+            if profiler: profiler.memory_check("simplish_leastsq: begin outer iter")
+
+            Jac = jac_guarded(k, num_fd_iters, obj_fn, jac_fn, f, ari, global_x, fdJac)
+
+            if profiler:
+                jac_gb = Jac.nbytes/(1024.0**3) if hasattr(Jac, 'nbytes') else _np.NaN
+                vals = ((f.size, global_x.size), jac_gb)
+                profiler.memory_check("simplish_leastsq: after jacobian: shape=%s, GB=%.2f" % vals)
+            
+            Jnorm = _np.sqrt(ari.norm2_jac(Jac))
+            xnorm = _np.sqrt(ari.norm2_x(x))
+            printer.log("--- Outer Iter %d: norm_f = %g, mu=%g, |x|=%g, |J|=%g" % (k, norm_f, mu, xnorm, Jnorm))
+
+            tm = _time.time()
+
+            # Riley note: fill_JTJ is the first place where we try to access J as a dense matrix.
+            ari.fill_jtj(Jac, JTJ, optional_jtj_buff)
+            ari.fill_jtf(Jac, f, minus_JTf)  # 'P'-type
+            minus_JTf *= -1
+
+            if profiler: profiler.add_time("simplish_leastsq: dotprods", tm)
+
+            norm_JTf = ari.infnorm_x(minus_JTf)
+            norm_x = ari.norm2_x(x)
+            pre_reg_data = ari.jtj_pre_regularization_data(JTJ)
+
+            if norm_JTf < jac_norm_tol:
+                if oob_check_interval <= 1:
+                    msg = "norm(jacobian) is at most %g" % jac_norm_tol
+                    converged = True
+                    break
+                else:
+                    printer.log(("** Converged with out-of-bounds with check interval=%d, reverting to last know in-bounds point and setting interval=1 **") % oob_check_interval, 2)
+                    oob_check_interval = 1
+                    x[:] = best_x[:]
+                    mu, nu, norm_f, f[:] = best_x_state
+                    continue
+
+            if k == 0:
+                max_jtj_diag = ari.jtj_max_diagonal_element(JTJ)
+                mu, nu = (tau * max_jtj_diag, 2) if init_munu == 'auto' else init_munu
+                best_x_state = (mu, nu, norm_f, f.copy())
+
+            #determing increment using adaptive damping
+            while True:  # inner loop
+
+                if profiler: profiler.memory_check("simplish_leastsq: begin inner iter")
+
+                # ok if assume fine-param-proc.size == 1 (otherwise need to sync setting local JTJ)
+                ari.jtj_update_regularization(JTJ, pre_reg_data, mu)
+
+                #assert(_np.isfinite(JTJ).all()), "Non-finite JTJ (inner)!" # NaNs tracking
+                #assert(_np.isfinite(minus_JTf).all()), "Non-finite minus_JTf (inner)!" # NaNs tracking
+
+                try:
+                    if profiler: profiler.memory_check("simplish_leastsq: before linsolve")
+                    tm = _time.time()
+                    _custom_solve(JTJ, minus_JTf, dx, ari, resource_alloc, serial_solve_proc_threshold)
+                    if profiler: profiler.add_time("simplish_leastsq: linsolve", tm)
+                except _scipy.linalg.LinAlgError:
+                    reject_msg = " (LinSolve Failure)"
+                    mu, nu, msg = damp_coeff_update(mu, nu, half_max_nu, reject_msg, printer)
+                    if len(msg) == 0:
+                        continue
+                    else:
+                        break
+
+                reject_msg = ""
+                if profiler: profiler.memory_check("simplish_leastsq: after linsolve")
+
+                new_x[:] = x + dx
+                norm_dx = ari.norm2_x(dx)
+
+                #ensure dx isn't too large - don't let any component change by more than ~max_dx_scale
+                if max_norm_dx and norm_dx > max_norm_dx:
+                    dx *= _np.sqrt(max_norm_dx / norm_dx)
+                    new_x[:] = x + dx
+                    norm_dx = ari.norm2_x(dx)
+
+                #apply x limits (bounds)
+                if x_limits is not None:
+                    # Approach 1: project x into valid space by simply clipping out-of-bounds values
+                    for i, (x_el, lower, upper) in enumerate(zip(x, x_lower_limits, x_upper_limits)):
+                        if new_x[i] < lower:
+                            new_x[i] = lower
+                            dx[i] = lower - x_el
+                        elif new_x[i] > upper:
+                            new_x[i] = upper
+                            dx[i] = upper - x_el
+                    norm_dx = ari.norm2_x(dx)
+
+                printer.log("  - Inner Loop: mu=%g, norm_dx=%g" % (mu, norm_dx), 2)
+
+                if norm_dx < (rel_xtol**2) * norm_x:
+                    if oob_check_interval <= 1:
+                        msg = "Relative change, |dx|/|x|, is at most %g" % rel_xtol
+                        converged = True
+                        break
+                    else:
+                        printer.log(("** Converged with out-of-bounds with check interval=%d, reverting to last know in-bounds point and setting interval=1 **") % oob_check_interval, 2)
+                        oob_check_interval = 1
+                        x[:] = best_x[:]
+                        mu, nu, norm_f, f[:] = best_x_state
+                        break
+                elif (norm_x + rel_xtol) < norm_dx * (_MACH_PRECISION**2):
+                    msg = "(near-)singular linear system"
+                    break
+
+                if oob_check_mode == 0 and oob_check_interval > 0:
+                    if k % oob_check_interval == 0:
+                        #Check to see if objective function is out of bounds
+
+                        in_bounds = []
+                        ari.allgather_x(new_x, global_new_x)
+                        try:
+                            new_f = obj_fn(global_new_x, oob_check=True)
+                        except ValueError:  # Use this to mean - "not allowed, but don't stop"
+                            in_bounds.append(False)
+                        else:
+                            in_bounds.append(True)
+
+                        if any(in_bounds):  # In adaptive mode, proceed if *any* cases are in-bounds
+                            new_x_is_known_inbounds = True
+                        else:
+                            MIN_STOP_ITER = 1  # the minimum iteration where an OOB objective stops the optimization
+                            if oob_action == "reject" or k < MIN_STOP_ITER:
+                                reject_msg = " (out-of-bounds)"
+                                mu, nu, msg = damp_coeff_update(mu, nu, half_max_nu, reject_msg, printer)
+                                if len(msg) == 0:
+                                    continue
+                                else:
+                                    break
+                            elif oob_action == "stop":
+                                if oob_check_interval == 1:
+                                    msg = "Objective function out-of-bounds! STOP"
+                                    converged = True
+                                    break
+                                else:  # reset to last know in-bounds point and not do oob check every step
+                                    printer.log(("** Hit out-of-bounds with check interval=%d, reverting to last know in-bounds point and setting interval=1 **") % oob_check_interval, 2)
+                                    oob_check_interval = 1
+                                    x[:] = best_x[:]
+                                    mu, nu, norm_f, f[:] = best_x_state
+                                    break  # restart next outer loop
+                            else:
+                                raise ValueError("Invalid `oob_action`: '%s'" % oob_action)
+                    else:  # don't check this time
+                        ari.allgather_x(new_x, global_new_x)
+                        new_f = obj_fn(global_new_x, oob_check=False)
+                        new_x_is_known_inbounds = False
+                else:
+                    #Just evaluate objective function normally; never check for in-bounds condition
+                    ari.allgather_x(new_x, global_new_x)
+                    new_f = obj_fn(global_new_x)
+                    new_x_is_known_inbounds = oob_check_interval == 0
+                    # ^ assume in bounds if we have no out-of-bounds checks.
+
+                norm_new_f = ari.norm2_f(new_f)
+                if not _np.isfinite(norm_new_f):  # avoid infinite loop...
+                    msg = "Infinite norm of objective function!"
+                    break
+
+                # dL = expected decrease in ||F||^2 from linear model
+                dL = ari.dot_x(dx, mu * dx + minus_JTf)
+                dF = norm_f - norm_new_f      # actual decrease in ||F||^2
+
+                printer.log("      (cont): norm_new_f=%g, dL=%g, dF=%g, reldL=%g, reldF=%g" % (norm_new_f, dL, dF, dL / norm_f, dF / norm_f), 2)
+
+                if dL / norm_f < rel_ftol and dF >= 0 and dF / norm_f < rel_ftol and dF / dL < 2.0:
+                    if oob_check_interval <= 1:  # (if 0 then no oob checking is done)
+                        msg = "Both actual and predicted relative reductions in the sum of squares are at most %g" % rel_ftol
+                        converged = True
+                        break
+                    else:
+                        printer.log(("** Converged with out-of-bounds with check interval=%d, reverting to last know in-bounds point and setting interval=1 **") % oob_check_interval, 2)
+                        oob_check_interval = 1
+                        x[:] = best_x[:]
+                        mu, nu, norm_f, f[:] = best_x_state
+                        break
+
+                if (dL <= 0 or dF <= 0):
+                    reject_msg = " (out-of-bounds)"
+                    mu, nu, msg = damp_coeff_update(mu, nu, half_max_nu, reject_msg, printer)
+                    if len(msg) == 0:
+                        continue
+                    else:
+                        break
+
+                #Check whether an otherwise acceptable solution is in-bounds
+                if oob_check_mode == 1 and oob_check_interval > 0 and k % oob_check_interval == 0:
+                    #Check to see if objective function is out of bounds
+                    try:
+                        obj_fn(global_new_x, oob_check=True)  # don't actually need return val (== new_f)
+                        new_x_is_known_inbounds = True
+                    except ValueError:  # Use this to mean - "not allowed, but don't stop"
+                        MIN_STOP_ITER = 1  # the minimum iteration where an OOB objective can stops the opt.
+                        if oob_action == "reject" or k < MIN_STOP_ITER:
+                            reject_msg = " (out-of-bounds)"
+                            mu, nu, msg = damp_coeff_update(mu, nu, half_max_nu, reject_msg, printer)
+                            if len(msg) == 0:
+                                continue
+                            else:
+                                break
+                        elif oob_action == "stop":
+                            if oob_check_interval == 1:
+                                msg = "Objective function out-of-bounds! STOP"
+                                converged = True
+                                break
+                            else:  # reset to last know in-bounds point and not do oob check every step
+                                printer.log(("** Hit out-of-bounds with check interval=%d, reverting to last know in-bounds point and setting interval=1 **") % oob_check_interval, 2)
+                                oob_check_interval = 1
+                                x[:] = best_x[:]
+                                mu, nu, norm_f, f[:] = best_x_state
+                                break  # restart next outer loop
+                        else:
+                            raise ValueError("Invalid `oob_action`: '%s'" % oob_action)
+
+                # reduction in error: increment accepted!
+                #       ^ Note: if we ever reach this line, then we know that we'll be breaking from the loop.
+                t = 1.0 - (2 * dF / dL - 1.0)**3  # dF/dL == gain ratio
+                # always reduce mu for accepted step when |dx| is small
+                mu_factor = max(t, 1.0 / 3.0) if norm_dx > 1e-8 else 0.3
+                mu *= mu_factor
+                nu = 2
+                x[:] = new_x[:]
+                f[:] = new_f[:]
+                norm_f = norm_new_f
+                global_x[:] = global_new_x[:]
+                printer.log("      Accepted%s! gain ratio=%g  mu * %g => %g" % ("", dF / dL, mu_factor, mu), 2)
+                if new_x_is_known_inbounds and norm_f < min_norm_f:
+                    min_norm_f = norm_f
+                    best_x[:] = x[:]
+                    best_x_state = (mu, nu, norm_f, f.copy())
+
+                #assert(_np.isfinite(x).all()), "Non-finite x!" # NaNs tracking
+                #assert(_np.isfinite(f).all()), "Non-finite f!" # NaNs tracking
+
+                break 
+                # ^ exit inner loop normally ...
+            # end of inner loop
+        # end of outer loop
+        else:
+            #if no break stmt hit, then we've exceeded max_iter
+            msg = "Maximum iterations (%d) exceeded" % max_iter
+            converged = True  # call result "converged" even in this case, but issue warning:
+            printer.warning("Treating result as *converged* after maximum iterations (%d) were exceeded." % max_iter)
+
+    except KeyboardInterrupt:
+        if comm is not None:
+            # ensure all procs agree on what best_x is (in case the interrupt occurred around x being updated)
+            comm.Bcast(best_x, root=0)
+            printer.log("Rank %d caught keyboard interrupt!  Returning the current solution as being *converged*."
+                        % comm.Get_rank())
+        else:
+            printer.log("Caught keyboard interrupt!  Returning the current solution as being *converged*.")
+        msg = "Keyboard interrupt!"
+        converged = True
+
+    if comm is not None:
+        comm.barrier()  # Just to be safe, so procs stay synchronized and we don't free anything too soon
+
+    ari.deallocate_jtj(JTJ)
+    ari.deallocate_jtf(minus_JTf)
+    ari.deallocate_jtf(x)
+    ari.deallocate_jtj_shared_mem_buf(optional_jtj_buff)
+
+    if x_limits is not None:
+        ari.deallocate_jtf(x_lower_limits)
+        ari.deallocate_jtf(x_upper_limits)
+
+    ari.deallocate_jtf(dx)
+    ari.deallocate_jtf(new_x)
+
+    if fdJac is not None:
+        ari.deallocate_jac(fdJac)
+
+    ari.allgather_x(best_x, global_x)
+    ari.deallocate_jtf(best_x)
+
+    mu, nu, norm_f, f[:] = best_x_state
+
+    global_f = _np.empty(ari.global_num_elements(), 'd')
+    ari.allgather_f(f, global_f)
+
+    return global_x, converged, msg, mu, nu, norm_f, global_f