# Bumps documentation

## bumps.dream.state

# Source code for bumps.dream.state

    """
    Sampling history for MCMC.

    MCMC keeps track of a number of things during sampling.

    The results may be queried as follows::

        draws, generation, thinning, total_generations
        sample(condition) returns draws, points, logp
        gen_logp()        returns genid, logp
        acceptance_rate() returns draws, AR
        traces()          returns genid, chains
        CR_weight()       returns draws, CR_weight
        best()            returns best_x, best_logp
        outliers()        returns outliers
        show()/save(file)/load(file)

    Data is stored in circular arrays, which keeps the last N generations and
    throws the rest away.

    draws is the total number of draws from the sampler.

    generation is the total number of generations. Due to tests for partially
    full circular buffers throughout the code (state.generation < state.Ngen)
    we are resetting generation to the size of the stored history on resume,
    and setting the generation_offset to the start of the history. If you need the
    number of generations across resume then use total_generations.

    thinning is the number of generations per stored sample.

    draws[i] is the number of draws including those required to produce the
    information in the corresponding return vector.  Note that draw numbers
    need not be linearly spaced, since techniques like delayed rejection
    will result in a varying number of samples per generation.

    logp[i] is the set of log likelihoods, one for each member of the population.
    The logp() method returns the complete set, and the sample() method returns
    a thinned set, with on element of logp[i] for each vector point[i, :].

    AR[i] is the acceptance rate at generation i, showing the proportion of
    proposed points which are accepted into the population.

    chains[i, :, :] is the set of points in the differential evolution population
    at thinned generation i.  Ideally, the thinning rate of the MCMC process
    is chosen so that thinned generations i and i+1 are independent samples
    from the posterior distribution, though there is a chance that this may
    not be the case, and indeed, some points in generation i+1 may be identical
    to those in generation i.  Actual generation number is i*thinning.

    points[i, :] is the ith point in a returned sample.  The i is just a place
    holder; there is no inherent ordering to the sample once they have been
    extracted from the chains.  Note that the sample may be from a marginal
    distribution.

    R[i] is the Gelman R statistic measuring convergence of the Markov chain.

    CR_weight[i] is the set of weights used for selecting between the crossover
    ratios available to the candidate generation process of differential
    evolution.  These will be fixed early in the sampling, even when adaptive
    differential evolution is selected.

    outliers[i] is a vector containing the thinned generation number at which
    an outlier chain was removed, the id of the chain that was removed and
    the id of the chain that replaced it.  We leave it to the reader to decide
    if the cloned samples, point[:generation, :, removed_id], should be included
    in further analysis.

    best_logp is the highest log likelihood observed during the analysis and
    best_x is the corresponding point at which it was observed.

    generation is the last generation number
    """

    # TODO: state should be collected in files as we go

    __all__ = ["MCMCDraw", "Draw", "dream_load", "h5load", "h5dump"]

    import os.path
    import re
    import gzip
    from typing import List, Dict, TextIO, Tuple, Union, Optional, Callable
    from pathlib import Path
    import warnings
    from fnmatch import fnmatch

    import numpy as np
    from numpy import empty, sum, asarray, inf, argmax, hstack, dstack
    from numpy import loadtxt, savetxt, reshape
    from numpy.typing import NDArray

    from .convergence import burn_point
    from .outliers import identify_outliers
    from .util import rng
    from .gelman import gelman

    # Don't load hdf5 if you don't need it
    from typing import TYPE_CHECKING

    if TYPE_CHECKING:
        from h5py import Group

    EXT = ".mc.gz"
    CREATE = gzip.open
    # EXT = ".mc"
    # CREATE = open

    SelectionType = Optional[Dict[Union[int, str], Tuple[float, float]]]

    UNCERTAINTY_DTYPE = "d"
    MAX_LABEL_LENGTH = 1024
    LABEL_DTYPE = f"|S{MAX_LABEL_LENGTH}"
    H5_COMPRESSION = 5

    # TODO: Unused code. Using webview.server.cli.reload_export instead

    def dream_load(store):
        """
        Load the saved DREAM state.

        *store* is the path to the stored state, either as an HDF5 history file
        or as a bumps export directory. If the directory contains multiple exports
        then use the path to the .par file within the directory.

        See also: h5load, load_state
        """
        from pathlib import Path
        import h5py

        store = Path(store).expanduser().resolve()
        modelname = None

        if not store.exists():
            raise RuntimeError(f"Path {store} does not exist.")

        # First try the HDF reader since the user can save the session to any file.
        if store.is_file():
            h5 = None
            try:
                with h5py.File(store) as h5:
                    if "fitting/fit_state" in h5:
                        return h5load(h5["fitting/fit_state"])
                    else:
                        raise RuntimeError(f"fitting/fit_state not found in {store}")
            except Exception:
                if h5 is not None:
                    raise
            # Fall through to the non-hdf5 loader

        if store.is_dir():
            # It's a directory: look for model name in store/*.par
            pars = list(store.glob("*.par"))
            if not pars:
                raise RuntimeError("No .par file found in %s" % store)
            if len(pars) > 1:
                print(pars)
                raise RuntimeError("Multiple models found in %s: %s" % (store, ", ".join(p.name for p in pars)))
            modelname = pars[0].stem
        else:
            # It's a filename: parse store=path/model.par
            if store.suffix != ".par":
                raise RuntimeError(f"Expected model.par file, got {store}")
            modelname = store.stem
            store = store.parent

        # Reload the MCMC data
        basename = str(store / modelname)
        state = load_state(basename)

        return state

    def _h5_write_field(group: "Group", field: str, data: Union[NDArray, str]):
        import h5py

        # print(f"h5 writing {field} in {group} with {data}")
        if isinstance(data, str):
            dtype = h5py.string_dtype(encoding="utf-8", length=len(data))
            return group.create_dataset(field, data=data, dtype=dtype)
        elif isinstance(data, (int, np.integer)):
            return group.create_dataset(field, data=data, dtype=np.int64)
        elif isinstance(data, (float, np.floating)):
            return group.create_dataset(field, data=data, dtype=np.double)
        elif np.isscalar(data):
            return group.create_dataset(field, data=data, dtype=data.dtype)
        else:
            return group.create_dataset(field, data=data, dtype=data.dtype, compression=H5_COMPRESSION)

    def _h5_read_field(group: "Group", field: str):
        raw_data = group[field][()]
        size = raw_data.size
        if isinstance(raw_data, np.integer):
            return int(raw_data)
        elif isinstance(raw_data, np.floating):
            return float(raw_data)
        elif size is not None and size > 0:
            return raw_data
        else:
            return None

    def h5dump(group: "Group", state: "MCMCDraw"):
        class Fields:
            total_generations = state.total_generations
            thinning = state.thinning
            gen_draws, gen_logp = state.logp(full=True)
            _, AR = state.acceptance_rate()
            thin_draws, thin_point, thin_logp = state.chains()
            update_draws, update_CR_weight = state.CR_weight()
            labels = np.array(state.labels, dtype=LABEL_DTYPE)
            # These are marked outliers. They should still be marked when reloading the
            # state, but they need to be cleared when updates are given. Make sure the
            # request for the population from dream includes the outliers.
            good_chains = None if isinstance(state._good_chains, slice) else state._good_chains
            # In case the MCMC chains are stored as 32-bit, save the current population
            # as 64-bit so that resume is consistent.
            current_population, current_logp = state._last_gen()
            best_x, best_logp = state.best()
            best_gen = state._best_gen
            portion = state.portion

        # print(f"wrote {Fields.total_generations} generations")
        # print(f"{state._gen_offset=} {state.Ngen=}")
        # print("wrote", Fields.__dict__)
        for field, value in Fields.__dict__.items():
            if field[0] != "_" and value is not None:
                try:
                    _h5_write_field(group, field, value)
                except Exception as exc:
                    exc.args = (f"{exc.args[0]} while writing {field}",) + exc.args[1:]
                    raise

    def h5load(group: "Group"):
        # print(f"dream.state.h5load {group}")
        class Fields: ...

        for field in group:
            setattr(Fields, field, _h5_read_field(group, field))
        # print("read", Fields.__dict__)
        # Guess dimensions
        Ngen = Fields.gen_draws.shape[0]
        Nthin, Npop, Nvar = Fields.thin_point.shape
        Nupdate, Ncr = Fields.update_CR_weight.shape
        # Nthin -= skip
        good_chains = getattr(Fields, "good_chains", None)
        total_generations = getattr(Fields, "total_generations", Ngen)
        thinning = getattr(Fields, "thinning", 1)
        current_population = getattr(Fields, "current_population", None)
        if current_population is None:
            current_population = Fields.thin_point[-1].astype("d")
        current_logp = getattr(Fields, "current_logp", None)
        if current_logp is None:
            current_logp = Fields.thin_logp[-1].astype("d")
        best_x = getattr(Fields, "best_x", None)
        best_logp = getattr(Fields, "best_logp", 0.0)
        best_gen = getattr(Fields, "best_gen", 0)
        portion = getattr(Fields, "portion", 1.0)

        # Create empty draw and fill it with loaded data
        state = MCMCDraw(0, 0, 0, 0, 0, 0, thinning)
        state.portion = portion
        state.draws = Ngen * Npop
        state.labels = [label.decode() for label in Fields.labels]
        state.generation = Ngen
        # print(f"loading {total_generations=} {Ngen=}")
        state._gen_offset = total_generations - Ngen
        state._gen_index = 0
        state._gen_draws = Fields.gen_draws.astype(np.int64)
        state._gen_acceptance_rate = Fields.AR
        state._gen_logp = Fields.gen_logp
        state._thin_count = Nthin
        state._thin_index = 0
        state._thin_draws = Fields.thin_draws.astype(np.int64)
        state._thin_logp = Fields.thin_logp
        state._thin_point = Fields.thin_point
        state._gen_current = current_population
        state._gen_current_logp = current_logp
        state._update_count = Nupdate
        state._update_index = 0
        state._update_draws = Fields.update_draws.astype(np.int64)
        state._update_CR_weight = Fields.update_CR_weight
        state._outliers = []

        if best_x is not None:
            state._best_x = best_x
            state._best_logp = best_logp
            state._best_gen = best_gen
        else:
            bestidx = np.unravel_index(np.argmax(Fields.thin_logp), Fields.thin_logp.shape)
            state._best_logp = Fields.thin_logp[bestidx]
            state._best_x = Fields.thin_point[bestidx]
            # We are not multiplying the following by thinning because thinning and best_gen
            # were added at the same time, so the reloaded thinning=1. Even if it were not
            # one this could still be wrong since thinning may have changed on resume.
            state._best_gen = bestidx[0]

        state._good_chains = slice(None, None) if good_chains is None else good_chains.astype(int)

        return state

    def save_state(state: "MCMCDraw", filename: str):
        # import sys; trace = sys.stdout
        # trace = open(filename+"-trace.mc", "wt")
        # trace.write("starting trace\n")

        # Build 2-D data structures
        # trace.write("extracting draws, logp\n")
        draws, logp = state.logp(full=True)
        # trace.write("extracting acceptance rate\n")
        _, AR = state.acceptance_rate(portion=1)
        # trace.write("building chain from draws, AR and logp\n")
        chain = hstack((draws[:, None], AR[:, None], logp))

        # trace.write("extracting point, logp\n")
        _, point, logp = state.chains()
        Nthin, Npop, Nvar = point.shape
        # trace.write("shape is %d,%d,%d\n" % (Nthin, Npop, Nvar))
        # trace.write("adding logp to point\n")
        point = dstack((logp[:, :, None], point))
        # trace.write("collapsing to draws x point\n")
        point = reshape(point, (point.shape[0] * point.shape[1], point.shape[2]))

        # trace.write("extracting CR_weight\n")
        draws, CR_weight = state.CR_weight()
        Nupdate, Ncr = CR_weight.shape
        # trace.write("building stats\n")
        stats = hstack((draws[:, None], CR_weight))

        # TODO: missing _outliers from save_state

        # Write convergence info
        # trace.write("writing chain\n")
        fid = CREATE(filename + "-chain" + EXT, "wb")
        fid.write(f"# draws acceptance_rate {Npop}*logp\n".encode())
        savetxt(fid, chain)
        fid.close()

        # Write point info
        # trace.write("writing point\n")
        fid = CREATE(filename + "-point" + EXT, "wb")
        fid.write(f"# logp point (Nthin x Npop x Nvar = [{Nthin},{Npop},{Nvar}])\n".encode())
        savetxt(fid, point)
        fid.close()

        # Write stats
        # trace.write("writing stats\n")
        fid = CREATE(filename + "-stats" + EXT, "wb")
        fid.write(f"# draws {Ncr}*CR_weight\n".encode())
        savetxt(fid, stats)
        fid.close()
        # trace.write("done state save\n")
        # trace.close()

    IND_PAT = re.compile("-1#IND")
    INF_PAT = re.compile("1#INF")

    def loadtxt_with_fallback(file: TextIO, report=0):
        """
        Try to load the file with numpy.loadtxt, but if it fails then
        fall back to loadtxt_MSVC, which is adapted for windows non-finite numbers.
        """

        # record current position in file so that we can reset it if loadtxt fails
        pos = file.tell()
        try:
            return loadtxt(file)
        except Exception:
            file.seek(pos)
            return loadtxt_MSVC(file, report=report)

    def loadtxt_MSVC(fh: TextIO, report=0):
        """
        Like numpy loadtxt, but adapted for windows non-finite numbers.
        """

        res = []
        section = 0
        lineno = 0
        for line in fh:
            lineno += 1
            if report and lineno % report == 0:
                print("read", section * report)
                section += 1
            IND_PAT.sub("nan", line)
            INF_PAT.sub("inf", line)
            line = line.split("#")[0].strip()
            values = line.split()
            if len(values) > 0:
                try:
                    res.append([float(v) for v in values])
                except ValueError:
                    print("Parse error:", values)
        return asarray(res)

    def path_contains_saved_state(filename: str) -> bool:
        chain_file = filename + "-chain" + EXT
        return os.path.exists(chain_file)

    def openmc(filename: str) -> TextIO:
        # If filename ends in .mc.gz, also check for a .mc file.
        # If filename ends in .mc, also check for a .mc.gz file.
        if filename.endswith(".gz"):
            if os.path.exists(filename):
                # print("opening with gzip")
                fh = gzip.open(filename, "rt")
            elif os.path.exists(filename[:-3]):
                fh = open(filename[:-3], "rt")
            else:
                raise RuntimeError("file %s does not exist" % filename)
        else:
            if os.path.exists(filename):
                fh = open(filename, "rt")
            elif os.path.exists(filename + ".gz"):
                # print("opening with gzip")
                fh = gzip.open(filename + ".gz", "rt")
            else:
                raise RuntimeError("file %s does not exist" % filename)
        return fh

    def load_state(filename: str, skip=0, report=0, derived_vars=0):
        """
        *filename* is the path to the saved MCMC state up to the final -chain.mc, etc.
        Any extension will be removed before using.

        *derived_vars* is the number of columns added to each point, derived
        from other columns in that point. The newer set_derived_vars interface generates
        the derived variables on demand rather than storing them in the state object and
        so it will be zero always.
        """

        # Trim extension from filename
        filename = str(Path(filename).with_suffix(""))

        # Read chain file
        with openmc(filename + "-chain" + EXT) as fid:
            chain = loadtxt_with_fallback(fid)

        # Read point file
        with openmc(filename + "-point" + EXT) as fid:
            line = fid.readline()
            point_dims = line[line.find("[") + 1 : line.find("]")]
            Nthin, Npop, Nvar = eval(point_dims)
            for _ in range(skip * Npop):
                fid.readline()
            point = loadtxt_with_fallback(fid, report=report * Npop)

        # Read stats file
        with openmc(filename + "-stats" + EXT) as fd:
            stats_header = fd.readline()
            stats = loadtxt_with_fallback(fd)

        # Determine number of R-stat stored in the stats file
        if "R-stat" in stats_header:
            # Old header looks like:
            #     # draws {Nvar}*R-stat {Ncr}*CR_weight
            # however, number of R-stat stored in stats file is the number of
            # variables stored each generation, not including the derived variables
            # calculated after the MCMC has completed.
            num_r = int(stats_header.split("*")[0].split()[-1]) - derived_vars
        else:
            num_r = 0

        # Guess dimensions
        Ngen = chain.shape[0]
        thinning = 1
        Nthin -= skip
        Nupdate = stats.shape[0]

        # Create empty draw and fill it with loaded data
        state = MCMCDraw(0, 0, 0, 0, 0, 0, thinning)
        # print("gen, var, pop", Ngen, Nvar, Npop)
        state.portion = 1.0
        state.draws = Ngen * Npop
        state.generation = Ngen
        state._gen_offset = 0
        state._gen_index = 0
        state._gen_draws = chain[:, 0].astype(np.int64)
        state._gen_acceptance_rate = chain[:, 1]
        state._gen_logp = chain[:, 2:]
        state._thin_count = Ngen // thinning
        state._thin_index = 0
        state._thin_draws = state._gen_draws[(skip + 1) * thinning - 1 :: thinning]
        state._thin_logp = point[:, 0].reshape((Nthin, Npop))
        state._thin_point = reshape(point[:, 1 : Nvar + 1 - derived_vars], (Nthin, Npop, -1))
        state._gen_current = state._thin_point[-1].copy()
        state._gen_current_logp = state._thin_logp[-1].copy()
        state._update_count = Nupdate
        state._update_index = 0
        state._update_draws = stats[:, 0].astype(np.int64)
        state._update_CR_weight = stats[:, 1 + num_r :]
        state._outliers = []

        bestidx = np.argmax(point[:, 0])
        state._best_logp = point[bestidx, 0]
        state._best_x = point[bestidx, 1 : Nvar + 1 - derived_vars]
        state._best_gen = 0

        # Clean up _thin dimensions.
        # It seems that some 0.x versions of bumps are yielding thin draws one shorter
        # than logp and points, so throw away the first logp and points.
        if state._thin_draws.shape[0] == state._thin_logp.shape[0] - 1:
            # print(f"shapes differ {state._gen_draws.shape=}, {state._thin_draws.shape=}, {state._thin_logp.shape=}")
            state._thin_logp = state._thin_logp[1:]
            state._thin_point = state._thin_point[1:]

        # Load labels from the .par file if it exists.
        parfile = Path(f"{filename}.par")
        if parfile.exists():
            try:
                text = parfile.read_text().strip()
                lines = text.split("\n")
                labels = [line.rsplit(None, 1)[0].strip() for line in lines]
                state.labels = labels
            except Exception as exc:
                print(f"Couldn't load labels from {parfile}: {exc}")

        return state

    # TODO: Unused code. Using webview.server.cli.reload_export instead
    def reload_state_from_export(parfile):
        """
        For the simplified fit with an export path we can reload
        the bumps dream state from the exported MCMC files.

        *parfile* is the name of the model.par file saved in the
        export directory.

        Use *state.show()* for the returned state to produce the plots.
        """
        assert parfile.endswith(".par"), "need model.par file"
        path = str(parfile)[:-4]
        state = load_state(path)
        state.mark_outliers()
        state.portion = state.trim_portion()
        return state

    class MCMCDraw(object):
        """
        Initialization parameters:

        *Ngen* is the number of generations to store. These are retrievable through
        *gen_logp()* and *acceptance_rate()* methods. Note that this is not the same
        as the length of the saved chains.

        *Nthin* is the number of thinned generations to store. These are retrievable
        through the *draw()* and *traces()* methods.

        *Nupdate* is the number of crossover ratio sets to save. The DREAM algorithm
        runs in batches of steps, with various checks such as convergence and
        crossover ratio updates after each batch. The results are retrievable
        through the *CR_weight()* method.

        *Nvar* is the number of parameters stored in each point.

        *Npop* is the number of running chains. Note that unlike the --pop command
        line option, this is not the number of chains per parameter.

        *Ncr* is the number of crossover ratios. DREAM maintains a fixed set of
        crossover ratios, choosing amongst them at each DE step. Depending on the
        success rate of the steps, it will adapt the weights after every batch of
        updates.

        *thinning* is the number of update generations to skip between saves to the
        MC chains. Additional thinning can be done when drawing samples from the
        saved points.

        Attributes and properties:

        *Ngen*, *Nthin*, *Nupdate*, *Nvar*, *Npop*, *Ncr* gives the size of the
        buffers stored in state. If the buffers are not yet full, the returned sizes
        for the *traces()*, *gen_logp()*, *acceptance_rate()*, *CR_weight()* and
        *draw()* methods will be smaller.

        *generation* is the number of generations in the current fit run. This is
        reset to zero each time the fit is resumed.

        *draws* is the number of draws in the current fit run. This is reset to zero
        each time the fit is resumed. Use *total_generations* time *Npop* if you
        want the number of draws including all previous runs.

        *total_generations* is the total number of generations for the fit across
        all fit runs.

        *title* a title for the plot

        *labels* a list of names for each parameter

        *thinning* amount of thinning between the gen_logp buffer and the traces
        buffer. DREAM stores every nth generation starting at generation n
        (1-origin). It is possible to change thinning on resume, which will change
        the short range correlation in the traces, however these changes are not
        recorded. Generation number on the trace plot will be incorrect until DREAM
        burns through the entire buffer.

        *portion* gives the fraction of each chain to use for statistical plots.  A
        value can be automatically determined by calling *state.portion =
        state.trim_portion()*, or supplied by the user to various functions that
        work with the chains.  The value should be between 0.0 and 1.0, where 1.0
        means the entire chain is used for statistical plots.
        """

        _derived_fn: Callable[[NDArray], NDArray] = None
        _derived_labels: List[str] = None
        _labels = None
        _integer_vars = None  # boolean array of integer variables, or None
        title = None

        def __init__(
            self,
            Ngen: int,
            Nthin: int,
            Nupdate: int,
            Nvar: int,
            Npop: int,
            Ncr: int,
            thinning: int,
        ):
            # self.generation and self.draws are used to control the number of iterations in
            # the dream loop, and must therefore be set to the current size of the state on
            # resume rather than the cumulative across all runs. To retrieve the totals use
            # generations=self.total_generations and draws=self.total_generations*self.Npop.

            # Total number of draws so far
            self.draws = 0
            self.portion = 1.0

            # Maximum observed likelihood
            # Note: with thinning and burn the best may not be in the set of samples
            self._best_x = None
            self._best_logp = -inf
            self._best_gen = 0

            # TODO: Scrap the per-generation data. It takes memory and disk, but doesn't add value.
            # Per generation iteration
            self.generation = 0
            self._gen_offset = 0
            self._gen_index = 0
            self._gen_draws = empty(Ngen, np.int64)
            self._gen_logp = empty((Ngen, Npop), dtype=UNCERTAINTY_DTYPE)
            self._gen_acceptance_rate = empty(Ngen, dtype=UNCERTAINTY_DTYPE)

            # TODO: should current generation be single precision if the model is single precision?
            # The current generation, stored as double precision so that nothing
            # is lost on resume.
            self._gen_current = empty((Npop, Nvar), dtype="d")
            self._gen_current_logp = empty((Npop,), dtype="d")

            # Per thinned generation iteration
            self.thinning = thinning
            self._thin_index = 0
            self._thin_count = 0
            self._thin_timer = 0
            self._thin_draws = empty(Nthin, np.int64)
            self._thin_point = empty((Nthin, Npop, Nvar), dtype=UNCERTAINTY_DTYPE)
            self._thin_logp = empty((Nthin, Npop), dtype=UNCERTAINTY_DTYPE)

            # Per update iteration
            self._update_index = 0
            self._update_count = 0
            self._update_draws = empty(Nupdate, np.int64)
            self._update_CR_weight = empty((Nupdate, Ncr), dtype=UNCERTAINTY_DTYPE)

            self._outliers = []

            # Query functions will not return outlier chains; initially, all
            # chains are marked as good.  Call mark_outliers to remove
            # outlier chains from the set.
            self._good_chains = slice(None, None)

        @property
        def Ngen(self):
            return self._gen_draws.shape[0]

        @property
        def total_generations(self):
            return self._gen_offset + self.generation

        @property
        def Nsamples(self):
            return self._gen_logp.size

        @property
        def Nthin(self):
            return self._thin_draws.shape[0]

        @property
        def Nupdate(self):
            return self._update_draws.shape[0]

        @property
        def Nvar(self):
            """Number of parameters in the fit"""
            return self._thin_point.shape[2]

        @property
        def Npop(self):
            return self._gen_logp.shape[1]

        @property
        def Ncr(self):
            return self._update_CR_weight.shape[1]

        def resize(self, Ngen: int, Nthin: int, Nupdate: int, Nvar: int, Npop: int, Ncr: int, thinning: int):
            if self.Nvar != Nvar or self.Npop != Npop or self.Ncr != Ncr:
                raise ValueError("Cannot change Nvar, Npop or Ncr on resize")

            # print("resize")
            # print(self.generation, self.Ngen, Ngen)
            # print(self._update_count, self.Nupdate, Nupdate)
            # print(self._thin_count, self.Nthin, Nthin)

            # When resuming from live state, make sure we unroll the state before resizing.
            # This is probably a no-op since plotting will already have forced an unroll.
            self._unroll()

            self.thinning = thinning

            def buf_resize(v, new_size, position):
                """
                Resize a circular buffer to *new_size*.

                *position* is the number of entries that have been added to the buffer. This
                may be less than the existing size if the buffer is not yet filled with one
                complete cycle. Assume the buffer has been unrolled so that the next location
                index is zero. Don't assume that the next location index is *position%size*.

                When extending, we add space to the end of the buffer and set the next index
                after the previous end of the buffer. When contracting, we keep the last N entries
                in the buffer. If there are fewer than N entries, keep all of them and include
                enough empty entries to make the buffer size N.

                Need to set the next index to min(position, old_size, new_size)%new_size
                """
                if v.shape[0] < new_size:  # expand
                    v = np.resize(v, (new_size, *v.shape[1:]))
                elif position < new_size:  # contract when not enough
                    v = v[:new_size]
                else:  # contract when have enough
                    v = v[-new_size:]
                return v

            self._gen_offset += max(self.generation - Ngen, 0)
            self.generation = min(self.generation, Ngen, self.Ngen)
            self.draws = self.generation * self.Npop
            self._gen_index = 0
            # Reset sample size to current size on resume.
            self.generation = min(self.generation, self.Ngen, Ngen)
            self.draws = self.generation * self.Npop
            self._gen_index = self.generation % Ngen
            self._gen_draws = buf_resize(self._gen_draws, Ngen, self.generation)
            self._gen_logp = buf_resize(self._gen_logp, Ngen, self.generation)
            self._gen_acceptance_rate = buf_resize(self._gen_acceptance_rate, Ngen, self.generation)

            self._thin_count = min(self._thin_count, self.Nthin, Nthin)
            self._thin_index = self._thin_count % Nthin
            self._thin_draws = buf_resize(self._thin_draws, Nthin, self._thin_count)
            self._thin_point = buf_resize(self._thin_point, Nthin, self._thin_count)
            self._thin_logp = buf_resize(self._thin_logp, Nthin, self._thin_count)

            self._update_count = min(self._update_count, self.Nupdate, Nupdate)
            self._update_index = self._update_count % Nupdate
            self._update_draws = buf_resize(self._update_draws, Nupdate, self._update_count)
            self._update_CR_weight = buf_resize(self._update_CR_weight, Nupdate, self._update_count)

        def save(self, filename: Union[Path, str]):
            save_state(self, filename)

        def trim_portion(self):
            """
            Estimate the point at which the Markov chains have reached equilibrium,
            returning it as a portion of the currently stored length.

            If not converged, then trim the first half of the samples.
            """
            index = burn_point(self)
            saved_gens = min(self.Ngen, self.generation)
            portion = 1 - (index / saved_gens) if index >= 0 else 0.5
            return portion

        def trim_index(self, portion: Optional[float] = None, generation: Optional[int] = None):
            """
            Returns the generation index corresponding to the trim portion. The returned
            generation is relative to the start of the sampling, even if resume has been
            called multiple times to extend the burn or the sample size.

            The optional *generation* parameter is needed in case we have a state object
            that is out of date with respect to the number of generations seen so far.
            This can happen when the state is transferred to the user interface thread
            much less frequently than the step monitor.
            """
            # Note: self.total_generations = self.generation + self._gen_offset
            # If generation is provided, it represents self.total_generations, so
            # we can reconstruct the number of saved generations we would see if the
            # state object were up to date.
            if generation is None:
                total_gens = self.total_generations
                saved_gens = min(self.Ngen, self.generation)
            else:
                total_gens = generation
                saved_gens = min(self.Ngen, generation - self._gen_offset)
            portion = self.portion if portion is None else portion
            return total_gens - int(portion * saved_gens)

        def show(self, portion: Optional[float] = None, figfile: Union[str, Path, None] = None):
            from .views import plot_all

            plot_all(self, portion=portion, figfile=figfile)

        # TODO: _last_gen and _draw_pop do similar things. Can they be merged?
        def _last_gen(self):
            """
            Returns x, logp for most recently saved generation. This may be several
            generations ago if thinning is enabled.
            """
            return self._gen_current, self._gen_current_logp

        def _generation(self, new_draws: int, x: NDArray, logp: NDArray, accept: NDArray, force_keep: bool = False):
            """
            Called from dream.py after each generation is completed with
            a set of accepted points and their values.
            """
            # Keep track of the total number of draws
            # Note: this is first so that we tag the record with the number of
            # draws taken so far, including the current draw.
            self.draws += new_draws
            self.generation += 1

            # Record if this is the best so far
            maxid = argmax(logp)
            if logp[maxid] > self._best_logp:
                self._best_logp = logp[maxid]
                self._best_x = x[maxid, :] + 0  # Force a copy
                self._best_gen = self.total_generations
                # print("new best", logp[maxid], self.generation)

            # Record acceptance rate and cost
            i = self._gen_index
            # print("generation", i, self.draws, "\n x", x, "\n logp", logp, "\n accept", accept)
            self._gen_draws[i] = self.draws
            self._gen_acceptance_rate[i] = 100 * sum(accept) / new_draws
            self._gen_logp[i] = logp
            i = i + 1
            if i == len(self._gen_draws):
                i = 0
            self._gen_index = i

            # Keep every nth iteration
            self._thin_timer += 1
            if self._thin_timer == self.thinning or force_keep:
                self._thin_timer = 0
                self._thin_count += 1
                i = self._thin_index
                self._thin_draws[i] = self.draws
                self._thin_point[i] = x
                self._thin_logp[i] = logp
                i = i + 1
                if i == len(self._thin_draws):
                    i = 0
                self._thin_index = i
            self._gen_current[:] = x
            self._gen_current_logp[:] = logp

        def _update(self, CR_weight: NDArray):
            """
            Called from dream.py when a series of DE steps is completed and
            summary statistics/adaptations are ready to be stored.
            """
            self._update_count += 1
            i = self._update_index
            # print("update", i, self.draws, "\n CR weight", CR_weight)
            self._update_draws[i] = self.draws
            self._update_CR_weight[i] = CR_weight
            i = i + 1
            if i == len(self._update_draws):
                i = 0
            self._update_index = i

        @property
        def labels(self):
            """
            Labels associated with each parameter in a point. This doesn't
            include the labels for derived quantities. For these you will
            need to query state.draw().labels.
            """
            if self._labels is None:
                result = ["P%d" % i for i in range(self._thin_point.shape[2])]
            else:
                result = self._labels
            return result

        @labels.setter
        def labels(self, v: List[str]):
            self._labels = v

        def _draw_pop(self):
            """
            Return the current population.
            """
            return self._gen_current

        def _unroll(self):
            """
            Unroll the circular queue so that data access can be done inplace.

            Call this when done stepping, and before plotting.  Calls to
            logp, sample, etc. assume the data is already unrolled.
            """
            if self.generation > self._gen_index > 0:
                self._gen_draws[:] = np.roll(self._gen_draws, -self._gen_index, axis=0)
                self._gen_logp[:] = np.roll(self._gen_logp, -self._gen_index, axis=0)
                self._gen_acceptance_rate[:] = np.roll(self._gen_acceptance_rate, -self._gen_index, axis=0)
                self._gen_index = 0

            if self._thin_count > self._thin_index > 0:
                self._thin_draws[:] = np.roll(self._thin_draws, -self._thin_index, axis=0)
                self._thin_point[:] = np.roll(self._thin_point, -self._thin_index, axis=0)
                self._thin_logp[:] = np.roll(self._thin_logp, -self._thin_index, axis=0)
                self._thin_index = 0

            if self._update_count > self._update_index > 0:
                self._update_draws[:] = np.roll(self._update_draws, -self._update_index, axis=0)
                self._update_CR_weight[:] = np.roll(self._update_CR_weight, -self._update_index, axis=0)
                self._update_index = 0

        def remove_outliers(self, x: NDArray, logp: NDArray, test: str = "IQR"):
            """
            Replace outlier chains with clones of good ones.  This should happen
            early in the sampling processes so the clones have an opportunity
            to evolve their own identity.  Only the head of the chain is modified.

            *state* contains the chains, with log likelihood for each point.

            *x*, *logp* are the current population and the corresponding
            log likelihoods; these are updated with cloned chain values.

            *test* is the name of the test to use (one of IQR, Grubbs, Mahal
            or none). See :func:`.outliers.identify_outliers` for details.

            Updates *state*, *x* and *logp* to reflect the changes.

            Returns a list of the outliers that were removed.
            """
            # Grab the last part of the chain histories
            _, chains = self.logp()
            chain_len, Nchains = chains.shape
            outliers = identify_outliers(test, chains, x)
            # if len(outliers): print("old llf", logp[outliers])

            # Loop over each outlier chain, replacing each with another
            for old in outliers:
                # Draw another chain at random, with replacement
                # TODO: consider using relative likelihood as a weight factor
                while True:
                    new = rng.randint(Nchains)
                    if new not in outliers:
                        break
                # Update the saved state and current population
                self._replace_outlier(old=old, new=new)
                x[old, :] = x[new, :]
                logp[old] = logp[new]

            # if len(outliers): print("new llf", logp[outliers])
            return outliers

        def _replace_outlier(self, old: int, new: int):
            """
            Called from outliers.py when a chain is replaced by the
            clone of another.
            """
            self._outliers.append((self._thin_index, old, new))

            # 2017-10-06 [PAK] only replace the head, not the full chain
            index = self._gen_index
            self._gen_current[old] = self._gen_current[new]
            self._gen_logp[index, old] = self._gen_logp[index, new]
            self._thin_logp[index, old] = self._thin_logp[index, new]
            self._thin_point[index, old, :] = self._thin_point[index, new, :]

        def mark_outliers(self, test: str = "IQR", portion: Optional[float] = None):
            """
            Mark some chains as outliers but don't remove them.  This can happen
            after drawing is complete, so that chains that did not converge are
            not included in the statistics.

            *test* is 'IQR', 'mahal', 'grubbs', or 'none'.

            *portion* indicates what portion of the samples should be included
            in the outlier test.  If None, then the stored portion is used.
            """
            _, chains, logp = self.chains()
            if test == "none":
                self._good_chains = slice(None, None)
            else:
                thin_saved_gens = chains.shape[0]
                portion = self.portion if portion is None else portion
                start = int(thin_saved_gens * (1 - portion))
                outliers = identify_outliers(test, logp[start:], chains[-1])
                # print("outliers", outliers, logp.shape, chains.shape)
                if len(outliers) > 0:
                    self._good_chains = np.array([i for i in range(logp.shape[1]) if i not in outliers])
                else:
                    self._good_chains = slice(None, None)

                # print(self._good_chains)

        def logp(self, full: bool = False):
            """
            Return the cumulative number of draws and the log likelihoods for each
            chain.

            Note that draw[i] represents the total number of samples taken,
            including those for the samples in logp[i].

            If full is True, then return all chains, not just good chains.

            ** Deprecated ** Use state.gen_logp() instead.
            """
            # self._unroll()
            # draws, logp = self._gen_draws, self._gen_logp
            # if self.generation == self._gen_index:
            #    draws, logp = [v[:self.generation] for v in (draws, logp)]

            # Don't do a full unroll here
            if self.generation == self._gen_index:
                draws = self._gen_draws[: self._gen_index]
                logp = self._gen_logp[: self._gen_index]
            elif self._gen_index > 0:
                draws = np.roll(self._gen_draws, -self._gen_index, axis=0)
                logp = np.roll(self._gen_logp, -self._gen_index, axis=0)
            else:
                draws = self._gen_draws
                logp = self._gen_logp

            return draws, (logp if full else logp[:, self._good_chains])

        def gen_logp(self, portion: Optional[float] = None, outliers: bool = False):
            """
            Return the iteration number and the log likelihood for each point in
            the individual sequences in that iteration.

            For example, to plot the convergence of each sequence::

                genid, logp = state.gen_logp()
                plot(genid, logp)

            *portion* is the amount to trim from the head of the chain, or None
            to use the stored burn point.

            If *outliers* is True, then return all chains, not just good chains.
            """
            # Don't do a full unroll here
            if self.generation == self._gen_index:
                logp = self._gen_logp[: self._gen_index]
            elif self._gen_index > 0:
                logp = np.roll(self._gen_logp, -self._gen_index, axis=0)
            else:
                logp = self._gen_logp

            stop, step = self.total_generations, 1  # No thinning on _gen_logp, _gen_draws
            start = stop - (len(logp) - 1) * step
            genid = np.arange(start, stop + 1, step)  # 1-origin, end-inclusive
            assert len(logp) == len(genid)

            portion = self.portion if portion is None else portion
            start = int((1 - portion) * len(logp))
            gen_index = slice(start, None)
            chain_index = slice(None) if outliers else self._good_chains
            return genid[gen_index], logp[gen_index, chain_index]

        def logp_slice(self, n: int):
            """
            Return a slice of the logp chains, either the first n if n > 0
            or the last n if n < 0.  Avoids unrolling the circular buffer if
            possible.
            """
            if n < 0:  # tail
                if self._gen_index >= -n:
                    return self._gen_logp[self._gen_index + n : self._gen_index]
                elif self._gen_index == 0:
                    return self._gen_logp[n:]
                else:  # unroll across boundary
                    return np.vstack((self._gen_logp[n + self._gen_index :], self._gen_logp[: self._gen_index]))
            else:  # head
                if self.generation < self.Ngen:
                    return self._gen_logp[:n]
                elif self._gen_index + n <= self.Ngen:
                    return self._gen_logp[self._gen_index : self._gen_index + n]
                else:
                    return np.vstack((self._gen_logp[self._gen_index :], self._gen_logp[-n + self._gen_index :]))

        def min_slice(self, n: int):
            """
            Return the minimum logp for n slices, from the head if positive
            or the tail if negative.

            This is a specialized function so it can be fast.  Convergence
            can be quickly rejected if the min in a short head is smaller
            than the min in a long tail.  Unfortunately, if the data is
            wrapped, then the max function will cost extra.
            """
            # Copy the logic of slice
            if n < 0:  # tail
                if self._gen_index >= -n:
                    return np.min(self._gen_logp[self._gen_index + n : self._gen_index])
                elif self._gen_index == 0:
                    return np.min(self._gen_logp[n:])
                else:  # max across boundary
                    return min(np.min(self._gen_logp[n + self._gen_index :]), np.min(self._gen_logp[: self._gen_index]))
            else:  # head
                if self.generation < self.Ngen:
                    return np.min(self._gen_logp[:n])
                elif self._gen_index + n <= self.Ngen:
                    return np.min(self._gen_logp[self._gen_index : self._gen_index + n])
                else:
                    return min(np.min(self._gen_logp[self._gen_index :]), np.min(self._gen_logp[-n + self._gen_index :]))

        def acceptance_rate(self, portion: Optional[float] = None):
            """
            Return the iteration number and the acceptance rate for that iteration.

            For example, to plot the acceptance rate over time::

                genid, AR = state.acceptance_rate()
                plot(genid, AR)

            V1.0.2: Now returns (genid, AR) rather than (num_draws, AR)
            """
            AR = self._gen_acceptance_rate
            if self.generation == self._gen_index:
                AR = AR[: self.generation]
            elif self._gen_index > 0:
                AR = np.roll(AR, -self._gen_index, axis=0)

            # Generation numbers for the buffered acceptance rates. The acceptance
            # rate is stored for each generation regardless of thinning, so this is
            # a simple increment.
            last, step = self.total_generations, 1
            first = last - (len(AR) - 1) * step
            generation = np.arange(first, last + 1, step)  # 1-origin, end-inclusive

            portion = self.portion if portion is None else portion
            index = int((1 - portion) * len(AR))
            return generation[index:], AR[index:]

        def chains(self):
            """
            Returns the observed Markov chains and the corresponding likelihoods.

            The return value is a tuple (*draws*, *chains*, *logp*).

            *draws* is the number of samples taken up to and including the samples
            for the current generation.

            *chains* is a three dimensional array of generations X chains X vars
            giving the set of points observed for each chain in every generation.
            Only the thinned samples are returned.

            *logp* is a two dimensional array of generation X population giving
            the log likelihood of observing the set of variable values given in
            chains.
            """
            self._unroll()
            retval = self._thin_draws, self._thin_point, self._thin_logp
            if self._thin_count == self._thin_index:
                retval = [v[: self._thin_count] for v in retval]
            return retval

        def traces(self, portion: Optional[float] = None, thin: int = 1, outliers: bool = False):
            """
            Grab traces for trace plot.

            *portion* gives the starting point of the trace, skipping over the initial frames
            that may not have converged.

            *thin* is the amount of additional thinning to apply on the traces.

            *outliers* (default False) is True if bad chains should be included in the trace.

            Returns *generation* [Ngen] and *chains* [Ngen x Nchain x Nvar]
            """
            portion = self.portion if portion is None else portion

            # TODO: Chains code is duplicated in Draw
            # Draw samples from the state according to portion, thinning and outliers.
            # If we also want derived variables and value range limits then we should
            # switch to the code in Draw.
            _, chains, _ = self.chains()
            start = int((1 - portion) * len(chains))
            gen_index = slice(start, None, thin)
            chain_index = slice(None) if outliers else self._good_chains
            chains = chains[gen_index, chain_index, :]

            num_steps = chains.shape[0]
            last, step = self.total_generations, self.thinning * thin
            first = last - (num_steps - 1) * step
            generation = np.arange(first + 1, last + 2, step)  # 1-origin inclusive
            return generation, chains

        def gelman(self, portion: Optional[float] = None):
            """
            Compute the Gelman and Rubin R-statistic for the Markov chains.
            """
            # Note: Not allowing a portion parameter here because that would
            # require unrolling the chains or otherwise making the logic in
            # gelman() too complicated. Keep it simple so that we can show
            # a simple R statistic over time without thrashing memory.
            if self.generation < self.Ngen:
                return gelman(self._thin_point[: self.generation], portion=1.0)
            else:
                return gelman(self._thin_point, portion=1.0)

        def CR_weight(self):
            """
            Return the crossover ratio weights to be used in the next generation.

            For example, to see if the adaptive CR is stable use::

                draw, weight = state.CR_weight()
                plot(draw, weight)

            See :mod:`.crossover` for details.
            """
            self._unroll()
            retval = self._update_draws, self._update_CR_weight
            if self._update_count == self._update_index:
                retval = [v[: self._update_count] for v in retval]
            return retval

        def outliers(self):
            """
            Return a list of outlier removal operations.

            Each outlier operation is a tuple giving the thinned generation
            in which it occurred, the old chain id and the new chain id.

            The chains themselves have already been updated to reflect the
            removal.

            Curiously, it is possible for the maximum likelihood seen so far
            to be removed by this operation.
            """
            return asarray(self._outliers, "i")

        def best(self):
            """
            Return the best point seen and its log likelihood.
            """
            return self._best_x, self._best_logp

        def stable_best(self):
            """
            Return True if the at least one full cycle of the circular
            buffer has passed since the best logp was first observed.
            """
            # print(f"stable_best {self._best_gen} + {self.Ngen} <= {self.total_generations}")
            return self._best_gen + self.Ngen <= self.total_generations

        def keep_best(self):
            """
            Place the best point at the end of the last good chain.

            Good chains are defined by mark_outliers.

            Because the Markov chain is designed to wander the parameter
            space, the best individual seen during the random walk may have
            been observed during the burn-in period, and may no longer be
            present in the chain.  If this is the case, replace the final
            point with the best, otherwise swap the positions of the final
            and the best.
            """

            # Get state as a 1D array
            _, chains, logp = self.chains()
            Ngen, Npop, Nvar = chains.shape
            points = reshape(chains, (Ngen * Npop, Nvar))
            logp = reshape(logp, Ngen * Npop)

            # Set the final position to the end of the last good chain.  If
            # mark_outliers has not been called, then _good_chains will
            # just be slice(None, None)
            if isinstance(self._good_chains, slice):
                final = -1
            else:
                final = self._good_chains[-1] - Npop

            # Find the location of the best point if it exists and swap with
            # the final position, or overwrite the final position if it does
            # not exist.
            idx = np.where(logp == self._best_logp)[0]
            if len(idx) == 0:
                logp[final] = self._best_logp
                points[final, :] = self._best_x
            else:
                # If there are multiple best values, arbitrarily choose one of them.
                # The chances are very high that it is the same point that was kept
                # after a proposed step was rejected.
                idx = idx[0]
                logp[final], logp[idx] = logp[idx], logp[final]
                points[final, :], points[idx, :] = points[idx, :], points[final, :]

        def sample(self, **kw):
            """
            Return a sample from the posterior distribution.

            **Deprecated** use :meth:`draw` instead.
            """
            drawn = self.draw(**kw)
            return drawn.points, drawn.logp

        def entropy(
            self,
            vars: Optional[List[int]] = None,
            portion: Optional[float] = None,
            selection: SelectionType = None,
            n_est: int = 10000,
            thin: Optional[int] = None,
            method: Optional[str] = None,
        ):
            r"""
            Return entropy estimate and uncertainty from an MCMC draw.

            *portion* is the portion of each chain to use (uses self.portion if None).

            *vars* is the set of variables to marginalize over.  It is None for
            the visible variables, or a list of variables.

            *vars* is the list of variables to use for marginalization.

            *selection* sets the range each parameter in the returned distribution,
            using {variable: (low, high)}. Missing variables use the full range.

            *n_est* is the number of points to use from the draw when estimating
            the entropy (default=10000).

            *thin* is the amount of thinning to use when selecting points from the
            draw.

            *method* determines which entropy calculation to use:

            * gmm: fit sample to a gaussian mixture model (GMM) with $5 \sqrt{d}$
              components where $d$ is the number fitted parameters and estimate
              entropy by sampling from the GMM.

            * llf: estimates likelihood scale factor from ratio of density
              estimate to model likelihood, then computes Monte Carlo entropy
              from sample; this does not work for marginal likelihood estimates.
              DOI:10.1109/CCA.2010.5611198

            * mvn: fit sample to a multi-variate Gaussian and return the entropy
              of the best fit gaussian; uses bootstrap to estimate uncertainty.

            * wnn: estimate entropy from nearest-neighbor distances in sample.
              DOI:10.1214/18-AOS1688
            """
            from . import entropy

            # Get the sample from the state.
            # set default thinning to max((steps * samples/step) // n_est, 1)
            if thin is None:
                Nsteps = min(self.Nthin, self._thin_count)
                thin = max(Nsteps * self.Npop // n_est, 1)
                # print("thin", thin, Nsteps, self.Npop, self.Nthin, self._thin_count)
            drawn = self.draw(portion=portion, vars=vars, selection=selection, thin=thin)

            # TODO: don't print within a library function!
            M = entropy.MVNEntropy(drawn.points)
            print("Entropy from MVN: %s" % str(M))

            if method is None:
                # TODO: change default to gmm
                method = "llf"

            if method == "llf":
                S, Serr = entropy.entropy(drawn.points, drawn.logp, N_entropy=n_est)
                # print("Entropy from llf (Kramer): %s"%str(S))
            elif method == "gmm":
                # Try pure gmm ... pretty good
                S, Serr = entropy.gmm_entropy(drawn.points, n_est=n_est)
                # print("Entropy from gmm: %g +/- %g"% (S, Serr))
            elif method == "wnn":
                # Try pure wnn ... no good
                S, Serr = entropy.wnn_entropy(drawn.points, n_est=n_est)
                # print("Entropy from wnn: %s"%str(S))
            elif method == "mvn":
                S, Serr = entropy.mvn_entropy_bootstrap(drawn.points)
                # print("Entropy from mvn: %s"%str(S))
            else:
                raise ValueError("unknown method %r" % method)

            # Always return entropy estimate from draw, even if it is normal
            return S, Serr

        def show_labels(self):
            """
            List of parameters in the state, with the parameter number for each. Use this
            to help with the inputs to state.draw().
            """
            print("\n".join(f"p[{k}] = {v}" for k, v in enumerate(self.labels)))

        def draw(
            self,
            portion: Optional[float] = None,
            vars: Optional[List[int]] = None,
            exclude: Optional[List[int]] = None,
            selection: SelectionType = None,
            thin: int = 1,
            outliers: bool = False,
            derived: Optional[Callable] = None,
        ):
            """
            Return a sample from the posterior distribution.

            *portion* is the portion of each chain to use

            *vars* is the list of vars that are active. Variables can be
            0-origin integers or strings matching a variable label. Match uses
            standard glob rules, with `*` matching any sequence, `?` matching any
            character, `[abc]` matching one of `abc`, and `[!abc]` not matching
            one of `abc`.

            *selection* sets the range each parameter in the returned distribution,
            using {var: (low, high)}. If *var* matches multiple labels, then use the
            first restriction only. Missing variables use the full range.

            *thin* takes every nth item.

            *outliers* is True if outlier chains should be included (default False).

            To plot the distribution for parameter p1::

                draw = state.draw()
                hist(draw.points[:, 0])

            To plot the interdependence of p1 and p2::

                draw = state.sample()
                plot(draw.points[:, 0], draw.points[:, 1], '.')
            """
            # Fill in defaults
            vars = vars if vars is not None else getattr(self, "_shown", None)
            portion = self.portion if portion is None else portion
            return Draw(
                state=self,
                portion=portion,
                vars=vars,
                exclude=exclude,
                selection=selection,
                thin=thin,
                outliers=outliers,
                derived=derived,
            )

        # TODO: Move processing of visible/integer/derived out of state

        def set_visible_vars(self, labels: List[str]):
            self._shown = [self.labels.index(v) for v in labels]

            # print("\n".join(str(pair) for pair in enumerate(self.labels)))
            # print(labels)
            # print(self._shown)

        def set_integer_vars(self, labels: List[str]):
            """
            Indicate tha variables should be considered integer variables when
            computing statistics.
            """
            self._integer_vars = np.array([var in labels for var in self.labels])

        def set_derived_vars(self, fn: Callable[[NDArray], NDArray], labels: List[str]):
            """
            Define derived variables from the sample. When calling draw() it will add
            columns for the derived variables to each sample.

            *fn* is a function taking points p[:, k] for k in 0 ... samples and
            returning a set of derived variables pj[k] for each sample k.  The
            variables can be returned as any kind of sequence including an
            array or a tuple with one entry per variable.  The caller uses
            asarray to convert the returned variables into a vars X samples array.
            For convenience, a single variable can be returned by itself.

            *labels* are the labels to use for the derived variables.

            The following example adds the new variable x+y = P[0] + P[1]::

                state.derive_vars(lambda p: p[0]+p[1], labels=["x+y"])
            """
            self._derived_fn = fn
            self._derived_labels = labels

        def derive_vars(self, fn: Callable[[NDArray], NDArray], labels: Optional[List[str]] = None):
            """
            *** DEPRECATED ***

            Like set_derived_vars but operating in place, modifying the points in the history.
            """
            warnings.warn("Use state.set_derived_vars() instead of state.derive_vars()", DeprecationWarning, stacklevel=2)

            # Grab all samples as a set of points
            _, chains, _ = self.chains()
            Ngen, Npop, Nvar = chains.shape
            points = reshape(chains, (Ngen * Npop, Nvar))

            # Compute new variables from the points
            newvars = asarray(fn(points.T)).T
            Nnew = newvars.shape[1] if len(newvars.shape) == 2 else 1
            newvars.reshape((Ngen, Npop, Nnew))

            # Extend new variables to be the same length as the stored selection
            Nthin = self._thin_point.shape[0]
            newvars = np.resize(newvars, (Nthin, Npop, Nnew))

            # Add labels for the new variables, if available. This must
            # occur before updating self._thin_point, otherwise a default
            # label may be generated for the new column.
            if labels is not None:
                self.labels = self.labels + labels
            elif self._labels is not None:
                labels = ["P%d" % i for i in range(Nvar, Nvar + Nnew)]
                self.labels = self.labels + labels
            else:  # no labels specified, old or new
                pass

            # Add new variables to the points
            self._thin_point = dstack((self._thin_point, newvars))

    class Draw:
        """
        A set of sample points from an MCMC draw with restrictions applied.

        **Inputs**

        *state* is the :class:`MCMCDraw` containing all points.

        *portion* is the fraction of points to use for the sample, starting
        from the ends of the Markov chains.

        *vars* is the list of vars that are active. Variables can be
        0-origin integers or strings matching a variable label. Match uses
        standard glob rules, with `*` matching any sequence, `?` matching any
        character, `[abc]` matching one of `abc`, and `[!abc]` not matching
        one of `abc`.

        *exclude* all labels that match a pattern in the list. Matching rules
        are the same as for *vars*. If a variable matches both *vars* and
        *exclude* then it is included.

        *selection* restricts the points returned to a specific region of the
        parameter space using {var: (low, high)}. The var can be a 0-origin
        integer or label as for *vars*. It can also be *logp* if restricting by
        log likelihood. If a label matches multiple patterns then choose the
        first restriction (this may change to all that match in future).

        *thin* is the amount of thinning to apply, using every nth point in
        the chain. This reduces autocorrelation in the chains and reduces
        artificial spikes in the marginal distributions. (The distributions
        will still be spiky, but the spikiness will be appropriate for the
        number of points in the bins.)

        **Attributes**

        *points[n,k]* are the set of sampled points, where n is the number of
        points and k is the number of parameters. The earlier points come
        from earlier generations, but this is an accident of the implementation
        and should not be relied on. n is usually the number of chains times the
        number of saved generations, but this will change depending on *portion*,
        *selection*, *thin* and *outliers* parameters. k is usually the number
        of fitted parameters, but the derived and active parameter lists will
        adjust this.

        *Nvar* is the number of dimensions per point in the draw after
        including derived parameters and excluding nuisance parameters.

        *weights* are the weights associated with each point. Since the chains
        are run at a single temperature the weights will all be 1.

        *logp[n]* is the negative log likelihuud for each sampled point.

        *labels* are the labels for the variables remaining after applying derived
        variables and restricting to the requested *vars* list.

        *integers* is an array of flags, one per parameter* indicating whether the
        parameter is float (0) or int (1=floor).  The values in the chain are
        managed as floats, and it is the responsibility of the likelihood function
        to transform them to integers. Models which use round, trunc or ceil rather
        than float are not yet supported.[1]

        *state* is the original MCMCDraw.

        [1] Note that discrete parameter handling is likely to change in the future.
        The derivative based optimizers in particular will need to know that the
        minimum step size for discrete parameters is one when computing the partial
        derivative. If bumps converts the parameters to integers before calling the
        likelihood function then we will only need a discrete flag since we will
        always use the *floor()* function.
        """

        state: MCMCDraw
        vars: Optional[List[Union[int, str]]]
        portion: float
        selection: SelectionType
        thin: int
        Nvar: int

        title: Optional[str]
        labels: List[str]
        integers: NDArray  # boolean
        logp: NDArray
        points: NDArray
        weights: NDArray

        # # MC chains, with mask for included points.
        # generation: NDArray
        # chains: NDArray
        # chains_logp: NDArray
        # mask: Optional[NDArray]

        def __init__(
            self,
            state: MCMCDraw,
            portion: float = 1.0,
            vars: Optional[List[Union[int, str]]] = None,
            exclude: Optional[List[Union[int, str]]] = None,
            selection: SelectionType = None,
            thin: int = 1,
            outliers: bool = False,
            derived: Optional[Callable] = None,
        ):
            self.state = state
            self.portion = portion
            self.thin = thin
            self.selection = selection
            self.labels = state.labels[:]  # copy of labels
            self.integers = (
                state._integer_vars if state._integer_vars is not None else np.zeros(len(self.labels), dtype=bool)
            )
            self.title = state.title

            # TODO: Make a copy?
            # Draw samples from the state according to portion and thinning.
            _, chains, logp = state.chains()
            start = int((1 - portion) * len(chains))
            gen_index = slice(start, None, thin)
            chain_index = slice(None) if outliers else state._good_chains
            chains = chains[gen_index, chain_index, :]
            chains_logp = logp[gen_index, chain_index]
            Ngen, Npop, Nvar = chains.shape

            # Derived variables are created during the draw so that existing data isn't altered.
            # This allows resume to work without extra effort. The code is also much simpler.
            if derived is None:
                derived = state._derived_fn
            if derived is not None:
                newvars = derived(chains.reshape(-1, Nvar).T)
                # If derived returns a dict, then use its keys as parameter labels
                if isinstance(newvars, dict):
                    newlabels, newvars = zip(*newvars.items())
                else:
                    # Backward compatibility: old interface had derived returning a list
                    newlabels = state._derived_labels
                    if not newlabels:
                        newlabels = [f"F{k+1}" for k in range(len(derived))]
                newvars = asarray(newvars).T
                newvars = newvars.reshape(Ngen, Npop, -1)
                chains = np.dstack((chains, newvars))
                Nvar = chains.shape[2]  # Update the number of variables
                self.integers = np.append(self.integers, [False] * newvars.shape[2])
                self.labels.extend(newlabels)

            # Select points within range limits given by selection, if any.
            # This happens before variable subsetting so that numerical indices
            # are independent of the displayed parameters. It happens after
            # calculating derived variables so that we can apply selection to derived values.
            mask = True
            selection = self._selection_pattern_to_index(selection)
            for var, limits in selection.items():
                if var == "logp":
                    mask = mask & (chains_logp >= limits[0]) & (chains_logp <= limits[1])
                else:
                    mask = mask & (chains[..., var] >= limits[0]) & (chains[..., var] <= limits[1])

            # TODO: with pattern matching we've lost the ability to reorder parameters
            # Restrict draw to a subset of the variables. Labels matching include take
            # precedence over those matching exclude.
            include = self._var_pattern_to_index(vars)
            exclude = self._var_pattern_to_index(exclude)
            exclude = [k for k in exclude if k not in include]
            include = [k for k in include if k not in exclude]
            # if vars and not include:
            #     # Note: could be empty because everything is excluded
            #    logging.warning(f"No parameters match {vars} in {self.labels}")
            if include:  # if nothing matches then match everything
                chains = chains[:, :, include]
                # Mask doesn't include the final dimension, so doesn't need to be trimmed.
                Nvar = chains.shape[-1]
                self.labels = [self.labels[v] for v in include]
                self.integers = self.integers[include]
                self.vars = include  # indices of selected items
            else:
                self.vars = None

            # Convert to a vector of selected points.
            if mask is True:
                # All points selected, so simply adjust the view
                self.points = chains.reshape(-1, Nvar)
                self.logp = chains_logp.flatten()
            else:
                # Only some points selected, so grab those under the mask
                self.points = chains[mask, :]
                self.logp = chains_logp[mask]

            # # Keep the thinned chains and the mask for trace plots
            # self.chains = chains
            # self.chains_logp = chains_logp
            # self.mask = None if mask is True else mask
            # stop, step = state.total_generations, thin*state.thinning
            # self.generation = np.arange(stop - (Ngen-1)*step + 1, stop + 2, step) # 1-origin inclusive

            self.weights = None
            self.Nvar = Nvar

            # Cache for the sorted argument indices.
            self._argsort_indices = {}

        def get_argsort_indices(self, var: int):
            """
            Sort var by value. Unlike argsort this caches the results for reuse.
            """
            if var not in self._argsort_indices:
                self._argsort_indices[var] = np.argsort(self.points[:, var].flatten())
            return self._argsort_indices[var]

        def _selection_pattern_to_index(self, selection: dict[int | str, tuple[float, float]] | None):
            """
            For the selection dictionary {pattern: data}, return a dictionary of {index: data}
            where data comes from the first pattern to matched by the label. Patterns are the
            standard unix glob patterns, or integers for specific numbered variables
            (see :meth:`State.show_labels`).
            """
            if selection is None:
                return {}
            result = {}
            if data := selection.get("logp", None):
                result["logp"] = data
            for index, label in enumerate(self.labels):
                for var, data in selection.items():
                    if index == var or (isinstance(var, str) and fnmatch(label, var)):
                        result[index] = data
                        break
            return result

        def _var_pattern_to_index(self, vars: list[int | str] | None):
            """
            Convert the include/exclude list of var patterns into a list of integer
            parameter indices. Patterns are the standard unix glob patterns, or integers
            for specific numbered variables (see :meth:`State.show_labels`)
            """
            if vars is None:
                return []

            def var_match(var, label, index):
                return index == var or (isinstance(var, str) and fnmatch(label, var))

            return [k for k, p in enumerate(self.labels) if any(var_match(v, p, k) for v in vars)]

    def test():
        from numpy.linalg import norm
        from numpy.random import rand
        from numpy import arange

        # Make some fake data
        Nupdate, Nstep = 3, 5
        Ngen = Nupdate * Nstep
        Nvar, Npop, Ncr = 3, 6, 2
        xin = rand(Ngen, Npop, Nvar)
        xin[..., 2] *= 9  # make one parameter range large enough for integer tests
        xin += 0.1  # force above 0.1 so that arrays print nicer
        pin = rand(Ngen, Npop)
        accept = rand(Ngen, Npop) < 0.8
        CRin = rand(Nupdate, Ncr)
        # thinning = 2
        # Nthin = int(Ngen/thinning)

        # Put it into a state
        thinning = 2
        Nthin = int(Ngen / thinning)
        thinned = xin[thinning - 1 :: thinning, :, :]
        state = MCMCDraw(Ngen=Ngen, Nthin=Nthin, Nupdate=Nupdate, Nvar=Nvar, Npop=Npop, Ncr=Ncr, thinning=thinning)
        for i in range(Nupdate):
            state._update(CR_weight=CRin[i])
            for j in range(Nstep):
                gen = i * Nstep + j
                state._generation(new_draws=Npop, x=xin[gen], logp=pin[gen], accept=accept[gen])

        # Check that it got there
        assert (thinned == state.chains()[1]).all()
        draws, logp = state.logp()
        assert norm(draws - Npop * arange(1, Ngen + 1)) == 0
        # Data is generated as double but may be stored as single,
        # so there could be a loss of precision.
        assert norm(logp - pin.astype(UNCERTAINTY_DTYPE)) == 0
        genid, logp = state.gen_logp()
        assert norm(genid - arange(1, Ngen + 1)) == 0
        assert norm(logp - pin) == 0
        genid, AR = state.acceptance_rate()
        assert norm(genid - arange(1, Ngen + 1)) == 0
        assert norm(AR - (100 * sum(accept, axis=1) / Npop).astype(UNCERTAINTY_DTYPE)) == 0
        draws, logp = state.sample()
        # assert norm(draws - thinning*Npop*arange(1, Nthin+1)) == 0
        # assert norm(sample - xin[thinning-1::thinning]) == 0
        # assert norm(logp - pin[thinning-1::thinning]) == 0
        draws, CR = state.CR_weight()
        assert norm(draws - Npop * Nstep * arange(Nupdate)) == 0
        assert norm(CR - CRin.astype(UNCERTAINTY_DTYPE)) == 0
        x, p = state.best()
        bestid = argmax(pin)
        i, j = bestid // Npop, bestid % Npop
        assert pin[i, j] == p
        assert norm(xin[i, j, :] - x) == 0

        draw = state.draw(selection={"P1": (0.4, 1.1)}, portion=1, outliers=True)
        assert ((draw.points[:, 1] >= 0.4) & (draw.points[:, 1] <= 1.1)).all()
        # print(thinned[..., 1])
        # print(thinned[..., 1] >= 0.4, thinned.shape)
        # print(len(draw.points[:, 1]), np.sum(thinned[..., 1] >= 0.4))
        # print("thinned", thinned[thinned[..., 1] >= 0.4][:, 1])
        # print("points", draw.points[:, 1])
        mask = (thinned[..., 1] >= 0.4) & (thinned[..., 1] <= 1.1)
        assert len(draw.points[:, 1]) == np.sum(mask)

        if 0:  # test of deprecated interface
            # Test derived variables (inplace update! ick!)
            state.derive_vars(lambda p: p[0] + p[1], labels=["inplace x+y"])
            assert state.traces()[1].shape[2] == Nvar + 1
            draw = state.draw()
            assert draw.labels[-1] == "inplace x+y"
            assert (draw.points[:, -1] == draw.points[:, 0] + draw.points[:, 1]).all()
            assert draw.points.shape[1] == Nvar + 1
            assert len(draw.labels) == Nvar + 1

        # Test derived variables (ephemeral)
        state.set_derived_vars(lambda p: p[0] + p[1], labels=["x+y"])
        draw = state.draw()
        assert draw.labels[-1] == "x+y"
        assert (draw.points[:, -1] == draw.points[:, 0] + draw.points[:, 1]).all()
        assert draw.points.shape[1] == Nvar + 1
        assert len(draw.labels) == Nvar + 1

        from .stats import var_stats, format_vars

        vstats = var_stats(state.draw(vars=[2, "x+y"]))
        print(format_vars(vstats))

        # Test integer vars
        state.set_integer_vars(labels=["P2"])
        draw = state.draw()
        vstats_int = var_stats(draw)
        # print("integers", draw.integers)
        # print(format_vars(vstats_int))
        assert vstats_int[2].mean == np.mean(np.floor(draw.points[:, 2]))

        draw = state.draw(selection={"P1": (0.4, 1.1)})
        # print(format_vars(var_stats(draw)))
        assert (draw.points[:, 1] >= 0.4).all()
        # print(thinned[..., 1])
        # print(thinned[..., 1] >= 0.4, thinned.shape)
        # print(len(draw.points[:, 1]), np.sum(thinned[..., 1] >= 0.4))
        # print("thinned", thinned[thinned[..., 1] >= 0.4][:, 1])
        # print("points", draw.points[:, 1])
        assert len(draw.points[:, 1]) == np.sum(thinned[..., 1] >= 0.4)

        # Note: check mods to the state data structures last
        # Check that outlier updates properly
        state._replace_outlier(1, 2)
        outliers = state.outliers()
        draws, logp = state.sample()
        assert norm(outliers - asarray([[state._thin_index, 1, 2]])) == 0
        # assert norm(sample[:, 1, :] - xin[thinning-1::thinning, 2, :]) == 0
        # assert norm(sample[:, 2, :] - xin[thinning-1::thinning, 2, :]) == 0
        # assert norm(logp[:, 1] - pin[thinning-1::thinning, 2]) == 0
        # assert norm(logp[:, 2] - pin[thinning-1::thinning, 2]) == 0

    if __name__ == "__main__":
        test()