Source code for autofit.non_linear.search.mle.bfgs.search

from __future__ import annotations

from typing import Optional, TYPE_CHECKING

from autofit import exc

from autofit.mapper.prior_model.abstract import AbstractPriorModel
from autofit.non_linear.search.mle.abstract_mle import AbstractMLE
from autofit.non_linear.analysis import Analysis
from autofit.non_linear.fitness import Fitness
from autofit.non_linear.clipper import (
    AbstractClipper,
    ClipperNone,
    ClipperPriorBoxJoint,
)
from autofit.non_linear.initializer import AbstractInitializer
from autofit.non_linear.samples.sample import Sample
from autofit.non_linear.samples.samples import Samples

import numpy as np

if TYPE_CHECKING:
    from autofit.database.sqlalchemy_ import sa


class AbstractBFGS(AbstractMLE):

    # Same contract as ``AbstractMultiStartGradient``: bounds handed to scipy
    # change where the fit converges, so the clipper forks the identifier here
    # too, re-keying existing (L)BFGS output directories (PyAutoFit#1493).
    # ``Drawer`` deliberately does NOT declare this: it inherits the attribute
    # from ``AbstractMLE`` but never consumes it, and a setting that cannot
    # affect the result must not re-key stored results.
    __identifier_fields__ = ("clipper",)

    method = None

    def __init__(
        self,
        name: Optional[str] = None,
        path_prefix: Optional[str] = None,
        unique_tag: Optional[str] = None,
        tol: Optional[float] = None,
        disp: bool = False,
        eps: float = 1.0e-08,
        ftol: float = 2.220446049250313e-09,
        gtol: float = 1.0e-05,
        iprint: float = -1.0,
        maxcor: int = 10,
        maxfun: int = 15000,
        maxiter: int = 15000,
        maxls: int = 20,
        initializer: Optional[AbstractInitializer] = None,
        clipper: Optional[AbstractClipper] = None,
        iterations_per_full_update: int = None,
        iterations_per_quick_update: int = None,
        silence: bool = False,
        session: Optional[sa.orm.Session] = None,
        **kwargs
    ):
        """
        Abstract wrapper for the BFGS and L-BFGS scipy non-linear searches.

        Parameters
        ----------
        tol
            Tolerance for termination.
        disp
            Set to True to print convergence messages.
        maxiter
            Maximum number of iterations.
        maxfun
            Maximum number of function evaluations.
        clipper
            Enforcement of prior support (see :mod:`autofit.non_linear.clipper`).
            Unlike the multi-start gradient searches, which project the parameters
            themselves after every update, this search declares the box to scipy
            and lets scipy enforce it — ``L-BFGS-B`` supports box bounds natively.

            Default ``ClipperNone``, under which no ``bounds=`` is passed at all
            and behaviour is unchanged. Only bound-supporting methods accept a
            real clipper; see ``_bounds_from``.
        """

        super().__init__(
            name=name,
            path_prefix=path_prefix,
            unique_tag=unique_tag,
            initializer=initializer,
            clipper=clipper,
            iterations_per_quick_update=iterations_per_quick_update,
            iterations_per_full_update=iterations_per_full_update,
            silence=silence,
            session=session,
            **kwargs
        )

        self.tol = tol
        self.disp = disp
        self.eps = eps
        self.ftol = ftol
        self.gtol = gtol
        self.iprint = iprint
        self.maxcor = maxcor
        self.maxfun = maxfun
        self.maxiter = maxiter
        self.maxls = maxls

        self.logger.debug(f"Creating {self.method} Search")

    # The scipy methods that accept box bounds. Plain ``BFGS`` is deliberately
    # absent: scipy does not reject bounds it cannot use, it *ignores* them behind
    # a ``UserWarning`` and returns the unconstrained optimum. A user who asked for
    # prior-support enforcement and received an unbounded fit plus a log line is
    # exactly the silent-wrong-answer this class is meant to prevent, so
    # ``_bounds_from`` raises instead.
    _BOUND_SUPPORTING_METHODS = ("L-BFGS-B", "TNC", "SLSQP")

    def _bounds_from(self, model):
        """
        The ``scipy.optimize.Bounds`` for ``model``, or ``None`` when no clipper is
        configured (in which case no ``bounds=`` is passed at all and behaviour is
        unchanged).

        The conversion is not incidental. ``AbstractClipper.bounds_from_model``
        returns ``(lower, upper)`` as two arrays, which is the shape ``project``
        broadcasts against — but ``optimize.minimize`` reads a ``(lower, upper)``
        tuple as a *sequence of ``(min, max)`` pairs*. For a two-parameter model
        that is a valid-looking pair sequence, so scipy pins each parameter to a
        constant and returns a wrong fit with no error and no warning; at every
        other dimensionality it raises. Building an explicit ``Bounds`` is what
        makes the intent unambiguous.

        A :class:`~autofit.non_linear.clipper.ClipperPriorBoxJoint` on a model that
        actually declares a ball is rejected rather than degraded to its box. A
        ball is not expressible as a ``scipy.optimize.Bounds`` at all -- no
        per-coordinate interval can exclude the corners of a square -- so silently
        handing scipy the box alone would give the caller an
        unconstrained-in-the-corners fit from a clipper they chose precisely to
        constrain them, which is the same class of silent wrong answer the
        ``_BOUND_SUPPORTING_METHODS`` check below exists to prevent.

        The refusal is keyed on the *model*, not on the clipper's type, because
        the joint clipper is a strict no-op on a model whose classes declare no
        geometry (see :class:`ClipperPriorBoxJoint`); refusing that case too would
        stop the clipper being configured once for a whole pipeline.
        """
        # Imported lazily, as everywhere else in autofit -- no module in the
        # package pulls scipy in at import time.
        from scipy import optimize

        if isinstance(self.clipper, ClipperNone):
            return None

        if (
            isinstance(self.clipper, ClipperPriorBoxJoint)
            and model.ball_constraint_index_pairs()
        ):
            raise exc.SearchException(
                f"A {type(self.clipper).__name__} was passed to "
                f"{type(self).__name__}, which enforces its bounds through "
                "scipy, together with a model declaring a ball constraint. A "
                "ball cannot be expressed as a `scipy.optimize.Bounds` -- only "
                "the box could be passed on, silently dropping the ball this "
                "clipper exists for. Use ClipperPriorBox here, or a search that "
                "projects its own steps (e.g. af.MultiStartAdam)."
            )

        if self.method not in self._BOUND_SUPPORTING_METHODS:
            raise exc.SearchException(
                f"A {type(self.clipper).__name__} was passed to a search using "
                f"method '{self.method}', which does not support box bounds. "
                f"SciPy would silently ignore the bounds and return an "
                f"unconstrained fit. Use one of "
                f"{', '.join(self._BOUND_SUPPORTING_METHODS)} (e.g. af.LBFGS) or "
                f"leave the clipper as the default ClipperNone."
            )

        lower, upper = self.clipper.bounds_from_model(model)
        return optimize.Bounds(lower, upper)

    @property
    def options(self):
        return {
            "disp": self.disp,
            "eps": self.eps,
            "ftol": self.ftol,
            "gtol": self.gtol,
            "iprint": self.iprint,
            "maxcor": self.maxcor,
            "maxfun": self.maxfun,
            "maxiter": self.maxiter,
            "maxls": self.maxls,
        }

    def _fit(
        self,
        model: AbstractPriorModel,
        analysis: Analysis,
    ):
        """
        Fit a model using the scipy L-BFGS method and the Analysis class which contains the data and returns the log
        likelihood from instances of the model, which the `NonLinearSearch` seeks to maximize.

        Parameters
        ----------
        model
            The model which generates instances for different points in parameter space.
        analysis
            Contains the data and the log likelihood function which fits an instance of the model to the data,
            returning the log likelihood the `NonLinearSearch` maximizes.

        Returns
        -------
        A result object comprising the Samples object that inclues the maximum log likelihood instance and full
        chains used by the fit.
        """
        from scipy import optimize

        # Resolved once, before any stepping: an unsupported method should fail
        # immediately rather than after the first chunk of iterations.
        bounds = self._bounds_from(model=model)

        fitness = Fitness(
            model=model,
            analysis=analysis,
            paths=self.paths,
            fom_is_log_likelihood=False,
            resample_figure_of_merit=-np.inf,
            convert_to_chi_squared=True,
            store_history=self.should_plot_start_point,
            iterations_per_quick_update=self.iterations_per_quick_update,
            background_quick_update=self.quick_update_background,
            live_visual_update=self.live_visual_update,
        )

        try:
            search_internal_dict = self.paths.load_search_internal()

            x0 = search_internal_dict["x0"]
            total_iterations = search_internal_dict["total_iterations"]

            self.logger.info(
                "Resuming LBFGS non-linear search (previous samples found)."
            )

        except (FileNotFoundError, TypeError):

            (
                unit_parameter_lists,
                parameter_lists,
                log_posterior_list,
            ) = self.initializer.samples_from_model(
                total_points=1,
                model=model,
                fitness=fitness,
                paths=self.paths,
                n_cores=self.number_of_cores,
            )

            x0 = np.asarray(parameter_lists[0])

            total_iterations = 0

            self.logger.info(
               f"Starting new {self.method} non-linear search (no previous samples found)."
            )

            self.plot_start_point(
                parameter_vector=x0,
                model=model,
                analysis=analysis,
            )

        while total_iterations < self.maxiter:

            iterations_remaining = self.maxiter - total_iterations
            # SciPy tolerates an integral float ``maxiter``, but a fractional one
            # would quietly acquire ceiling semantics instead of being rejected
            # (PyAutoFit#1422).
            iterations = self._steps_until_full_update(iterations_remaining)

            if iterations > 0:
                options = dict(self.options)
                options["maxiter"] = iterations

                # ``bounds`` is ``None`` under the default ``ClipperNone``, and
                # ``minimize(bounds=None)`` is the same call as omitting it, so the
                # default path is unchanged.
                #
                # Both branches are exposed to the prior-support problem:
                # ``UniformPrior.log_prior_from_value`` returns ``-inf`` outside
                # its bounds on the NumPy and JAX paths alike (PyAutoFit#1489
                # restored the NumPy-side wall), so an unclipped step out of the
                # box makes the objective non-finite on either branch. Bounds are
                # passed on both — they are correct on both, and having them
                # diverge by branch would be a trap of its own.
                if analysis._use_jax:

                    search_internal = optimize.minimize(
                        fun=fitness._jit,
                        x0=x0,
                        method=self.method,
                        options=options,
                        tol=self.tol,
                        bounds=bounds,
                    )
                else:

                    search_internal = optimize.minimize(
                        fun=fitness.__call__,
                        x0=x0,
                        method=self.method,
                        options=options,
                        tol=self.tol,
                        bounds=bounds,
                    )

                total_iterations += search_internal.nit

                search_internal.log_posterior_list = -0.5 * fitness(
                    parameters=search_internal.x
                )

                if self.should_plot_start_point:

                    search_internal.parameters_history_list = fitness.parameters_history_list
                    search_internal.log_likelihood_history_list = fitness.log_likelihood_history_list

                self.paths.save_search_internal(
                    obj=search_internal,
                )

                x0 = search_internal.x

                if search_internal.nit < iterations:
                    return search_internal, fitness

                self.perform_update(
                    model=model,
                    analysis=analysis,
                    during_analysis=True,
                    fitness=fitness,
                    search_internal=search_internal,
                )

        self.logger.info(f"{self.method} sampling complete.")

        return search_internal, fitness

    def samples_via_internal_from(
        self, model: AbstractPriorModel, search_internal=None
    ):
        """
        Returns a `Samples` object from the LBFGS internal results.

        The samples contain all information on the parameter space sampling (e.g. the parameters,
        log likelihoods, etc.).

        The internal search results are converted from the native format used by the search to lists of values
        (e.g. `parameter_lists`, `log_likelihood_list`).

        Parameters
        ----------
        model
            Maps input vectors of unit parameter values to physical values and model instances via priors.
        """

        if search_internal is None:
            search_internal = self.paths.load_search_internal()

        x0 = search_internal.x
        total_iterations = search_internal.nit

        if self.should_plot_start_point:

            parameter_lists = search_internal.parameters_history_list
            log_prior_list = model.log_prior_list_from(parameter_lists=parameter_lists)
            log_likelihood_list = search_internal.log_likelihood_history_list

        else:

            parameter_lists = [list(x0)]
            log_prior_list = model.log_prior_list_from(parameter_lists=parameter_lists)
            log_posterior_list = np.array([search_internal.log_posterior_list])
            log_likelihood_list = [
                lp - prior for lp, prior in zip(log_posterior_list, log_prior_list)
            ]

        weight_list = len(log_likelihood_list) * [1.0]

        sample_list = Sample.from_lists(
            model=model,
            parameter_lists=parameter_lists,
            log_likelihood_list=log_likelihood_list,
            log_prior_list=log_prior_list,
            weight_list=weight_list,
        )

        samples_info = {
            "total_iterations": total_iterations,
            # Prior-support enforcement (PyAutoFit#1477). The name only, with
            # deliberately NO ``n_clipped_lane_steps`` alongside it: this search
            # is declarative, handing ``optimize.Bounds`` to scipy and letting
            # scipy enforce, so ``Clipper.project`` is never called, no mask is
            # produced and there is nothing to count. Writing a ``0`` here would
            # read as "the clipper never fired" when it means "this search
            # cannot know"; the summary renders the absent key as "not measured"
            # instead.
            "clipper": type(self.clipper).__name__,
            "time": self.timer.time if self.timer else None,
        }

        return Samples(
            model=model,
            sample_list=sample_list,
            samples_info=samples_info,
        )


[docs] class BFGS(AbstractBFGS): """ The BFGS non-linear search, which wraps the scipy Broyden-Fletcher-Goldfarb-Shanno (BFGS) algorithm. For a full description of the scipy BFGS method, checkout its documentation: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-bfgs.html#optimize-minimize-bfgs """ method = "BFGS"
[docs] class LBFGS(AbstractBFGS): """ The L-BFGS non-linear search, which wraps the scipy Limited-memory Broyden-Fletcher-Goldfarb-Shanno (L-BFGS) algorithm. For a full description of the scipy L-BFGS method, checkout its documentation: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-lbfgsb.html """ method = "L-BFGS-B" # SciPy 1.15 deprecated the L-BFGS-B ``disp`` / ``iprint`` options (removal # slated for 1.18) — the solver no longer emits its Fortran-side verbose # output, so passing them buys nothing but a DeprecationWarning per # ``minimize`` call. The constructor still accepts both for API stability; # they simply never reach scipy for this method. @property def options(self): options = dict(super().options) del options["disp"] del options["iprint"] return options