# Bumps documentation

## bumps.dream.core

# Source code for bumps.dream.core

    """
    DiffeRential Evolution Adaptive Metropolis algorithm

    DREAM runs multiple different chains simultaneously for global exploration,
    and automatically tunes the scale and orientation of the proposal
    distribution using differential evolution.  The algorithm maintains
    detailed balance and ergodicity and works well and efficient for a large
    range of problems, especially in the presence of high-dimensionality and
    multimodality.

    DREAM developed by Jasper A. Vrugt and Cajo ter Braak

    This algorithm has been described in:

       Vrugt, J.A., C.J.F. ter Braak, M.P. Clark, J.M. Hyman, and B.A. Robinson,
          *Treatment of input uncertainty in hydrologic modeling: Doing hydrology
          backward with Markov chain Monte Carlo simulation*,
          Water Resources Research, 44, W00B09, 2008.
          `doi:10.1029/2007WR006720 <http://dx.doi.org/10.1029/2007WR006720>`_

       Vrugt, J.A., C.J.F. ter Braak, C.G.H. Diks, D. Higdon, B.A. Robinson,
           and J.M. Hyman,
           *Accelerating Markov chain Monte Carlo simulation by differential
           evolution with self-adaptive randomized subspace sampling*,
           International Journal of Nonlinear Sciences and Numerical Simulation,
           10(3), 271-288, 2009.

       Vrugt, J.A., C.J.F. ter Braak, H.V. Gupta, and B.A. Robinson,
           *Equifinality of formal (DREAM) and informal (GLUE) Bayesian approaches
           in hydrologic modeling*,
           Stochastic Environmental Research and Risk Assessment,
           1-16, 2009, In Press.
           `doi:10.1007/s00477-008-0274-y
           <http://dx.doi.org/10.1007/s00477-008-0274-y>`_

    For more information please read:

       Ter Braak, C.J.F.,
           *A Markov Chain Monte Carlo version of the genetic algorithm Differential
           Evolution: easy Bayesian computing for real parameter spaces*,
           Stat. Comput., 16, 239 - 249, 2006.
           `doi:10.1007/s11222-006-8769-1
           <http://dx.doi.org/10.1007/s11222-006-8769-1>`_

       Vrugt, J.A., H.V. Gupta, W. Bouten and S. Sorooshian,
           *A Shuffled Complex Evolution Metropolis algorithm for optimization
           and uncertainty assessment of hydrologic model parameters*,
           Water Resour. Res., 39 (8), 1201, 2003.
           `doi:10.1029/2002WR001642 <http://dx.doi.org/10.1029/2002WR001642>`_

       Ter Braak, C.J.F., and J.A. Vrugt,
           *Differential Evolution Markov Chain with snooker updater
           and fewer chains*,
           Statistics and Computing, 2008.
           `doi:10.1007/s11222-008-9104-9
           <http://dx.doi.org/2008.10.1007/s11222-008-9104-9>`_

       Vrugt, J.A., C.J.F. ter Braak, and J.M. Hyman,
           *Differential evolution adaptive Metropolis with snooker update and
           sampling from past states*,
           SIAM journal on Optimization, 2009.

       Vrugt, J.A., C.J.F. ter Braak, and J.M. Hyman,
           *Parallel Markov chain Monte Carlo simulation on distributed computing
           networks using multi-try Metropolis with sampling from past states*,
           SIAM journal on Scientific Computing, 2009.

       G. Schoups, and J.A. Vrugt,
           *A formal likelihood function for Bayesian inference of hydrologic
           models with correlated, heteroscedastic and non-Gaussian errors*,
           Water Resources Research, 2010, In Press.

       G. Schoups, J.A. Vrugt, F. Fenicia, and N.C. van de Giesen,
           *Inaccurate numerical solution of hydrologic models corrupts efficiency
           and robustness of MCMC simulation*,
           Water Resources Research, 2010, In Press.

    Copyright (c) 2008, Los Alamos National Security, LLC
    All rights reserved.

    Copyright 2008. Los Alamos National Security, LLC. This software was produced
    under U.S. Government contract DE-AC52-06NA25396 for Los Alamos National
    Laboratory (LANL), which is operated by Los Alamos National Security, LLC
    for the U.S. Department of Energy. The U.S. Government has rights to use,
    reproduce, and distribute this software.

    NEITHER THE GOVERNMENT NOR LOS ALAMOS NATIONAL SECURITY, LLC MAKES ANY
    WARRANTY, EXPRESS OR IMPLIED, OR ASSUMES ANY LIABILITY FOR THE USE OF
    THIS SOFTWARE.  If software is modified to produce derivative works, such
    modified software should be clearly marked, so as not to confuse it with
    the version available from LANL.

    Additionally, redistribution and use in source and binary forms, with or without
    modification, are permitted provided that the following conditions are met:

    * Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.

    * Redistributions in binary form must reproduce the above copyright notice,
      this list of conditions and the following disclaimer in the documentation
      and/or other materials provided with the distribution.

    * Neither the name of Los Alamos National Security, LLC, Los Alamos National
      Laboratory, LANL the U.S. Government, nor the names of its contributors
      may be used to endorse or promote products derived from this software
      without specific prior written permission.

    THIS SOFTWARE IS PROVIDED BY LOS ALAMOS NATIONAL SECURITY, LLC AND
    CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING,
    BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
    FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL LOS ALAMOS NATIONAL
    SECURITY, LLC OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
    SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
    PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
    OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
    WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
    OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
    ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

    MATLAB code written by Jasper A. Vrugt, Center for NonLinear Studies (CNLS)

    Written by Jasper A. Vrugt: vrugt@lanl.gov

    Version 0.5: June 2008
    Version 1.0: October 2008  Adaption updated and generalized CR implementation

    2010-04-20 Paul Kienzle
    * Convert to python
    Complete changelog on github.com/bumps/bumps
    """

    __all__ = ["Dream"]

    import sys
    import time
    from ctypes import c_double
    from typing import Protocol, List, Sequence, Optional

    import numpy as np
    from numpy.typing import NDArray

    from .state import MCMCDraw
    from .metropolis import metropolis, metropolis_dr, dr_step
    from .crossover import AdaptiveCrossover, LogAdaptiveCrossover, Crossover
    from .diffev import de_step
    from .bounds import make_bounds_handler
    from .compiled import dll
    from .util import rng
    from .convergence import ks_converged

    # Everything should be available in state, but lets be lazy for now
    _LAST_TIME = 0
    """Used by console_monitor to count time until the next update"""

    def console_monitor(state, pop, logp):
        """
        Print progress of fit on the console.

        This is the default update monitor when running DREAM standalone. It is not
        used by the bumps package.
        """
        global _LAST_TIME
        if state.generation == 1:
            print("#gen", "logp(x)", " ".join("<%s>" % par for par in state.labels))
            _LAST_TIME = time.perf_counter()

        current_time = time.perf_counter()
        if current_time >= _LAST_TIME + 1:
            _LAST_TIME = current_time
            print(state.generation, state._best_logp, " ".join("%.15g" % v for v in state._best_x))
            sys.stdout.flush()

    class Model(Protocol):
        labels: List[str]
        """Labels for all the parameters"""
        bounds: Sequence[Sequence[float]]
        """Bounds for each parameter as a pair of sequences of the same length as the labels."""

        def map(self, pop: NDArray) -> NDArray:
            """
            Function which takes an array of [k x n] and returns an array of n where k is
            the number of parameters.  TODO: check if
            """
            ...

    class Dream(object):
        """
        Data structure containing the details of the running DREAM analysis code.
        """

        model: Model = None
        # Sampling parameters
        burn: int = 0
        draws: int = 100000
        thinning: int = 1
        # TODO: change the default outlier test to IQR and control with options
        outlier_test = "none"
        population: NDArray = None
        #: convergence criteria
        alpha: float = 0.01
        # DE parameters
        DE_steps: int = 10
        DE_pairs: int = 3
        DE_eps: float = 0.05
        DE_snooker_rate: float = 0.1
        DE_noise: float = 1e-6
        bounds_style: str = "reflect"  # reflect, clip, fold, randomize, none
        # Crossover parameters
        CR: Crossover = None
        CR_spacing: str = "linear"  # 'log' or 'linear'
        # Delay rejection parameters
        use_delayed_rejection: bool = False
        DR_scale: float = 1  # 1-sigma step size using cov of population
        # Local optimizer best fit injection  The optimizer has
        # the following interface:
        #    x, fx = goalseek(mapper, bounds_handler, pop, fpop)
        # where:
        #    x, fx are the local optimum point and its value
        #    pop is the starting population
        #    fpop is the nllf for each point in pop
        #    mapper is a function which takes pop and returns fpop
        #    bounds_handler takes pop and forces all points into the range
        goalseek_optimizer = None
        goalseek_interval: int = 2e9  # close enough to never
        goalseek_minburn: int = 1000
        state: MCMCDraw = None

        def __init__(self, **kw):
            self.monitor = console_monitor
            for k, v in kw.items():
                if hasattr(self, k):
                    setattr(self, k, v)
                else:
                    raise TypeError("Unknown attribute " + k)

            self._initialized = False

        def sample(self, state: Optional[MCMCDraw] = None, abort_test=lambda: False):
            """
            Pull the requisite number of samples from the distribution
            """
            if not self._initialized:
                self._initialized = True
            self.state = state
            try:
                _run_dream(self, abort_test=abort_test)
            except KeyboardInterrupt:
                pass
            return self.state

    def _run_dream(dream: Dream, abort_test=lambda: False):
        """
        Collect posterior distribution samples using DREAM sampler.
        """

        if dll:
            dll.rand_init(rng.randint(1e9))
        # Step 1: Sample s points in the parameter space
        # [PAK] I moved this out of dream so that the user can use whatever
        # complicated sampling scheme they want.  Unfortunately, this means
        # the user needs to know some complex sampling scheme.
        if dream.population is None:
            raise ValueError("initial population not defined")

        # Remember the problem dimensions
        n_gen, n_chain, n_var = dream.population.shape
        n_pop = n_gen * n_chain

        if dream.CR is None:
            if dream.CR_spacing == "log":
                dream.CR = LogAdaptiveCrossover(n_var)
            else:  # linear
                dream.CR = AdaptiveCrossover(3)

        # Step 2: Calculate posterior density associated with each value in x
        apply_bounds = make_bounds_handler(dream.model.bounds, style=dream.bounds_style)

        # Record initial state
        allocate_state(dream)
        state = dream.state
        # Reset plot portion to 100% during fitting so that the intermediate plots show the
        # full set of samples.
        state.portion = 1.0
        state.labels = dream.model.labels
        previous_draws = state.draws
        if previous_draws:
            x, logp = state._last_gen()
        else:
            # No initial state, so evaluate initial population
            for x in dream.population:
                apply_bounds(x)
                # ********************** MAP *****************************
                logp = dream.model.map(x)
                state._generation(new_draws=n_chain, x=x, logp=logp, accept=n_chain, force_keep=True)
                dream.monitor(state, x, logp)
            # print(f"Initial pop has draws={state.draws} gen={state.generation} pop={dream.population.shape}")

        state._update(CR_weight=dream.CR.weight)

        # Now start drawing samples
        # print "previous draws", previous_draws, "new draws", dream.draws+dream.burn
        last_goalseek = (dream.draws + dream.burn) / n_pop - dream.goalseek_minburn
        next_goalseek = state.generation + dream.goalseek_interval if dream.goalseek_optimizer else 1e100

        if dll:
            xtry = np.empty((n_chain, n_var), "d")
            step_alpha = np.empty(n_chain, "d")
            CR_used = np.empty(n_chain, "d")
        scale = 1.0
        # serial_time = parallel_time = 0.
        # last_time = time.time()

        # Make sure that the pop we are drawing doesn't need to be copied before
        # sending it to the compiled C code.
        pop = state._draw_pop()
        assert pop.ctypes.data == np.ascontiguousarray(pop).ctypes.data

        # frame = 0  # For debugging...
        next_outlier_test = max(state.Ngen, 2 * state.Ngen - 10)
        next_convergence_test = state.Ngen
        final_gen = dream.draws + dream.burn
        # print(f"running dream with {dream.draws} draws and {dream.burn} burn; final gen={final_gen} draws={state.draws} gen={state.generation}")
        while state.draws < final_gen:
            # Age the population using differential evolution
            dream.CR.reset()
            CR = np.ascontiguousarray(np.vstack((dream.CR.CR, dream.CR.weight)).T, "d")
            for gen in range(dream.DE_steps):
                # Define the current locations and associated posterior densities
                xold, logp_old = x, logp
                pop = state._draw_pop()
                # print(pop.ctypes.data, np.ascontiguousarray(pop).ctypes.data)
                # print(pop)
                # print("gen", gen, pop.shape)

                # Generate candidates for each sequence
                if dll is None:
                    xtry, step_alpha, CR_used = de_step(
                        n_chain,
                        pop,
                        CR,
                        max_pairs=dream.DE_pairs,
                        eps=dream.DE_eps,
                        snooker_rate=dream.DE_snooker_rate,
                        noise=dream.DE_noise,
                        scale=scale,
                    )
                else:
                    dll.de_step(
                        n_chain,
                        n_var,
                        len(CR),
                        pop.ctypes,
                        CR.ctypes,
                        dream.DE_pairs,
                        c_double(dream.DE_eps),
                        c_double(dream.DE_snooker_rate),
                        c_double(dream.DE_noise),
                        c_double(scale),
                        xtry.ctypes,
                        step_alpha.ctypes,
                        CR_used.ctypes,
                    )
                # print("try", xtry)

                # PAK: Try a local optimizer every N generations
                if next_goalseek <= state.generation <= last_goalseek:
                    best, logp_best = dream.goalseek_optimizer(dream.model.map, apply_bounds, xold, logp_old)
                    xtry[0] = best
                    # Note: it is slightly inefficient to throw away logp_best,
                    # but it makes the the code cleaner if we do
                    next_goalseek += dream.goalseek_interval

                # Compute the likelihood of the candidates
                apply_bounds(xtry)
                # ********************** MAP *****************************
                # next_time = time.time()
                # serial_time += next_time - last_time
                # last_time = next_time
                logp_try = dream.model.map(xtry)
                # next_time = time.time()
                # parallel_time  += next_time - last_time
                # last_time = next_time
                draws = len(logp_try)

                # Apply the metropolis acceptance/rejection rule
                x, logp, alpha, accept = metropolis(xtry, logp_try, xold, logp_old, step_alpha)

                # Process delayed rejection
                # PAK NOTE: this updates according to the covariance matrix of the
                # current sample, which may be useful on unimodal systems, but
                # doesn't seem to be of any value in general; the DREAM papers
                # found that the acceptance rate did indeed improve with delayed
                # rejection, but the overall performance did not.  Worse, this
                # requires a linear system solution O(nPop^3) which can be near
                # singular for complex posterior distributions.
                if dream.use_delayed_rejection and not accept.all():
                    # Generate alternate candidates using the covariance of xold
                    xdr, r = dr_step(x=xold, scale=dream.DR_scale)

                    # Compute the likelihood of the new candidates
                    reject = ~accept
                    apply_bounds(xdr)
                    # ********************** MAP *****************************
                    logp_xdr = dream.model.map(xdr[reject])
                    draws += len(logp_xdr)

                    # Apply the metropolis delayed rejection rule.
                    x[reject], logp[reject], alpha[reject], accept[reject] = metropolis_dr(
                        xtry[reject],
                        logp_try[reject],
                        x[reject],
                        logp[reject],
                        xold[reject],
                        logp_old[reject],
                        alpha[reject],
                        r,
                    )

                # els = zip(logp_old, logp_try, logp, x[:, -2], x[:, -1], accept)
                # print "pop", "\n ".join((("%12.3e "*(len(el)-1))%el[:-1])
                #                       +("T " if el[-3]<=el[-2] else "  ")
                #                       +("accept" if el[-1] else "")
                #                       for el in els)

                # Update Sequences with the new population.
                state._generation(draws, x, logp, accept)
                # ********************** NOTIFY **************************
                dream.monitor(state, x, logp)
                # print state.generation, ":", state._best_logp

                # Keep track of which CR ratios were successful
                if state.draws <= dream.burn:
                    dream.CR.update(xold, x, CR_used)

                if abort_test():
                    break

            # print("serial&parallel",serial_time,parallel_time)
            # End of differential evolution aging
            # ---------------------------------------------------------------------

            # Save update information
            state._update(CR_weight=dream.CR.weight)

            if abort_test():
                # print("<<< DREAM abort >>>")
                break

            # if state.draws <= 0.1 * dream.draws:
            if state.draws <= dream.burn:
                # Adapt the crossover ratio, but only during burn-in.
                dream.CR.adapt()

            if False:
                # TODO: Suppress scale update until we have a chance to verify that it
                # doesn't skew the resulting statistics.
                _, r = state.acceptance_rate()
                ravg = np.mean(r[-dream.DE_steps :])
                if ravg > 0.4:
                    scale *= 1.01
                elif ravg < 0.2:
                    scale /= 1.01

            if state.generation >= next_convergence_test:
                converged = ks_converged(state, alpha=dream.alpha)
                # print(f"ks convergence for α={dream.alpha} is {converged}")
            else:
                converged = False

            # See whether there are any outlier chains, and remove them
            # Only do this once per frame, and only if there is some time
            # left to adapt the distribution.  Also do it if we believe
            # that we have converged.
            if (converged or state.generation >= next_outlier_test) and state.generation + state.Ngen < final_gen:
                outliers = state.remove_outliers(x, logp, dream.outlier_test)
                next_outlier_test = state.generation + 2 * state.Ngen
                if len(outliers):
                    # TODO: use monitors to report arbitrary information
                    print("step %d trimmed %d outliers from %d chains" % (state.generation, len(outliers), len(logp)))
                    next_convergence_test = state.generation + state.Ngen
                    converged = False

            if converged:
                # print("<<< DREAM converged >>>")
                # _show_logp_frame(dream, state, frame+1)
                break

            # Draw the next frame (for debugging...)
            # next_frame = state.generation // state.Ngen
            # if frame != next_frame: _show_logp_frame(dream, state, next_frame)
            # frame = next_frame

    def _show_logp_frame(dream, state, frame):
        from pylab import clf, savefig
        from . import views

        clf()
        views.plot_logp(state)
        savefig("logp%03d.png" % frame)
        clf()

    def allocate_state(dream):
        """
        Estimate the size of the output
        """
        # Determine problem dimensions from the initial population
        n_pop, n_chain, n_var = dream.population.shape
        steps = dream.DE_steps
        thinning = dream.thinning
        n_cr = len(dream.CR.CR)
        draws = dream.draws

        n_update = int(np.ceil(draws / (steps * n_chain)))
        n_gen = n_update * steps
        n_thin = int(np.ceil(n_gen / thinning))

        # print("new state", n_var, n_chain, n_cr, n_gen, n_thin, n_update, draws, steps)
        if dream.state is not None:
            # print("existing", dream.state.Nvar, dream.state.Npop, dream.state.Ncr)
            dream.state.resize(n_gen, n_thin, n_update, n_var, n_chain, n_cr, thinning)
        else:
            dream.state = MCMCDraw(n_gen, n_thin, n_update, n_var, n_chain, n_cr, thinning)