# Bumps documentation

## bumps.curve

# Source code for bumps.curve

    """
    Build a bumps model from a function and data.

    Example
    -------

    Given a function *sin_model* which computes a sine wave at times *t*::

        from numpy import sin
        def sin_model(t, freq, phase):
            return sin(2*pi*(freq*t + phase))

    and given data *(y,dy)* measured at times *t*, we can define the fit
    problem as follows::

        from bumps.names import *
        M = Curve(sin_model, t, y, dy, freq=20)

    The *freq* and *phase* keywords are optional initial values for the model
    parameters which otherwise default to zero.  The model parameters can be
    accessed as attributes on the model to set fit range::

        M.freq.range(2, 100)
        M.phase.range(0, 1)

    As usual, you can initialize or assign parameter expressions to the the
    parameters if you want to tie parameters together within or between models.

    Note: there is sometimes difficulty getting bumps to recognize the function
    during fits, which can be addressed by putting the definition in a separate
    file on the python path.  With the windows binary distribution of bumps,
    this can be done in the problem definition file with the following code::

        import os
        from bumps.names import *
        sys.path.insert(0, os.getcwd())

    The model function can then be imported from the external module as usual::

        from sin_model import sin_model
    """

    __all__ = ["Curve", "PoissonCurve", "plot_err"]

    import inspect
    import warnings
    from typing import Callable, Literal
    from copy import deepcopy

    import numpy as np
    from numpy import log, pi, sqrt

    from .parameter import Parameter

    def _parse_pars(fn, init=None, skip=0, name=""):
        """
        Extract parameter names from function definition.

        *fn* is the function definition.  This could be declared as
        *fn(p1, p2, p3, ...)* where *p1*, etc. are the fittable parameters.

        *init* is a dictionary of initial values for the parameters,
        overriding any default values.  If called from a constructor with
        **kwargs representing unknown named arguments, use *init=kwargs*.

        *skip* is the number of parameters to skip.  This will be *skip=0*
        for a function which defines the log likelihood directly or one
        that returns a set of residuals. For parameterized curves such as
        *fn(x, p1, p2, ...)* use *skip=1*.  For surfaces with
        *fn(x, y, p1, p2, ...)* use *skip=2*.

        *name* is added to each parameter name to differentiate it from other
        parameters in the same fit.

        A default value in the function definition such as *pk=value* will
        be set as the default value for the parameter.  If the default is
        *pk=None* then the parameter will be non-fittable, and instead set
        through *init*.
        """
        sig = inspect.signature(fn)
        params = sig.parameters.values()
        pnames = [p.name for p in params]

        valid = [p.kind in (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD) for p in params]
        if not all(valid):
            raise TypeError(f"Only positional and keyword arguments allowed for {fn.__name__}")

        # TODO: need "self" handling for passed methods
        # Skip the first argument if it is x or maybe skip x, y.
        pnames = pnames[skip:]

        # Parameters default to zero
        defaults = dict((p, 0) for p in pnames)

        # If the function provides default values, use those.
        for param in list(params)[skip:]:
            if param.default is not inspect.Parameter.empty:
                defaults[param.name] = param.default

        # Non-fittable parameters need to be sent in as None
        state_vars = set(p for p, v in defaults.items() if v is None)

        # Regardless, use any values specified in the constructor, but first
        # check that they exist as function parameters.
        invalid = set(init.keys()) - set(pnames)
        if invalid:
            raise TypeError("Invalid initializers: %s" % ", ".join(sorted(invalid)))
        defaults.update(init)

        # Build parameters out of ranges and initial values
        # maybe:  name=(p+name if name.startswith('_') else name+p)
        pars = dict((p, Parameter.default(defaults[p], name=name + p)) for p in pnames if p not in state_vars)

        state = dict((p, v) for p, v in defaults.items() if p in state_vars)

        # print("pars", pars)
        # print("state", state)
        return pars, state

    def _assign_pars(obj, pars):
        # Make parameters accessible as model attributes
        for k, v in pars.items():
            if hasattr(obj, k):
                raise TypeError("Parameter cannot be named %s" % k)
            setattr(obj, k, v)

    class Curve(object):
        r"""
        Model a measurement with a user defined function.

        The function *fn(x,p1,p2,...)* should return the expected value *y* for
        each point *x* given the parameters *p1*, *p2*, etc.  *dy* is the
        uncertainty for each measured value *y*.  If not specified, it defaults
        to 1. Multi-valued functions, which return multiple *y* values for each
        *x* value, should have *x* as a vector of length *n* and *y*, *dy* as
        arrays of size *[n, k]*.

        Initial values for the parameters can be set as *p=value* arguments to
        *Curve*. If no value is set, then the initial value will be taken from
        the default value given in the definition of *fn*, or set to 0 if the
        parameter is not defined with an initial value.  Arbitrary non-fittable
        data can be passed to the function as parameters, but only if the
        parameter is given a default value of *None* in the function definition,
        and has the initial value set as an argument to *Curve*.  Defining
        *state=dict(key=value, ...)* before *Curve*, and calling *Curve* as
        *Curve(..., \*\*state)* works pretty well.

        *Curve* takes the following special keyword arguments:

        * *name* is added to each parameter name when the parameter is defined.
          The filename for the data is a good choice, since this allows you to keep
          the parameters straight when fitting multiple datasets simultaneously.

        * *plot* is an alternative plotting function. The function should be
          defined as *plot(x,y,dy,fy,\*\*kw)*. The keyword arguments will be
          filled with the values of the parameters used to compute *fy*.  It
          will be easiest to list the parameters you need to make your plot
          as positional arguments after *x,y,dy,fy* in the plot function
          declaration.  For example, *plot(x,y,dy,fy,p3,\*\*kw)* will make the
          value of parameter *p3* available as a variable in your function.  The
          special keyword *view* will be a string containing *linear*, *log*,
          *logx*, or *loglog*.  If only showing the residuals, the string
          will be *residual*.

        * *plot_x* is an array giving the sample points to use when plotting
          the theory function, if different from the *x* values at which the
          function is sampled.  Use this to draw a smooth curve between the
          fitted points.  This value is ignored if you provide your own plot
          function.

        * *labels* are the axis labels for the plot.  This should include
          units in parentheses. If the function is multi-valued then
          use *['x axis', 'y axis', 'line 1', 'line 2', ...]*.

        The data uncertainty is assumed to follow a gaussian distribution.
        If measurements draw from some other uncertainty distribution, then
        subclass Curve and replace nllf with the correct probability given the
        residuals.  See the implementation of :class:`PoissonCurve` for an example.
        """

        def __init__(self, fn, x, y, dy=None, name="", labels=None, plot=None, plot_x=None, **kwargs):
            self.x, self.y = np.asarray(x), np.asarray(y)
            if dy is None:
                self.dy = 1
            else:
                self.dy = np.asarray(dy)
                if (self.dy <= 0).any():
                    raise ValueError("measurement uncertainty must be positive")

            if len(self.x.shape) == 1 and len(self.y.shape) > 1:
                num_curves = self.y.shape[0]
            else:
                num_curves = 1
            self._num_curves = num_curves  # use same value everywhere

            # interpret labels parameter
            if labels is None:
                labels = ["x", "y"]
            elif len(labels) < 2 or len(labels) != num_curves + 2:
                if num_curves > 1:
                    lines = "line1, ..., line%d" % num_curves
                else:
                    lines = "line"
                raise TypeError("labels should be [x, y, %s]" % lines)

            if len(labels) == 2:
                if num_curves > 1:
                    line_labels = ["y%d" % k for k in range(num_curves)]
                else:
                    line_labels = [labels[1]]
                labels = list(labels) + line_labels
            self.labels = labels

            # TODO: self.fn is a duplicate of self._function below. Deprecated?
            self.fn = fn
            self.name = name  # if name else fn.__name__ + " "
            self.plot_x = plot_x

            pars, state = _parse_pars(fn, init=kwargs, skip=1, name=name)

            # Make parameters accessible as model attributes
            _assign_pars(self, pars)
            # _assign_pars(state, self)  # ... and state variables as well

            # Remember the function, parameters, and number of parameters
            # Note: we are remembering the parameter names and not the
            # parameters themselves so that the caller can tie parameters
            # together using model1.par = model2.par.  Otherwise we would
            # need to override __setattr__ to intercept assignment to the
            # parameter attributes and redirect them to the a _pars dictionary.
            # ... and similarly for state if we decide to make them attributes.
            self._function = fn
            self._pnames = list(sorted(pars.keys()))
            self._state = state
            self._plot = plot
            self._cached_theory = None
            self._webview_plots = {}

        def update(self):
            self._cached_theory = None

        def parameters(self):
            return dict((p, getattr(self, p)) for p in self._pnames)

        def numpoints(self):
            return np.prod(self.y.shape)

        def theory(self, x=None):
            # Use cache if x is None, otherwise compute theory with x.
            if x is None:
                if self._cached_theory is None:
                    self._cached_theory = self._compute_theory(self.x)
                return self._cached_theory
            return self._compute_theory(x)

        def _compute_theory(self, x):
            kw = self._fetch_pars()
            return self._function(x, **kw)

        def _fetch_pars(self):
            kw = dict((p, getattr(self, p).value) for p in self._pnames)
            kw.update(self._state)
            return kw

        def simulate_data(self, noise=None):
            theory = self.theory()
            if noise is not None:
                if noise == "data":
                    pass
                elif noise < 0:
                    self.dy = np.full_like(theory, -noise)
                else:
                    self.dy = 0.01 * noise * abs(theory)
            self.y = theory + np.random.randn(*theory.shape) * self.dy

        def residuals(self):
            return (self.theory() - self.y) / self.dy

        def nllf(self):
            r = self.residuals()
            return 0.5 * np.sum(r**2)

        def save(self, basename):
            # TODO: need header line with state vars as json
            # TODO: need to support nD x,y,dy
            if len(self.x.shape) > 1:
                warnings.warn("Save not supported for nD x values")
                return

            theory = self.theory()
            if self._num_curves > 1:
                # Multivalued y, dy for single valued x.
                columns = [self.x]
                headers = ["x"]
                for k, (y, dy, fx) in enumerate(zip(self.y, self.dy, theory)):
                    columns.extend((y, dy, fx))
                    headers.extend(("y[%d]" % (k + 1), "dy[%d]" % (k + 1), "fx[%d]" % (k + 1)))
            else:
                # Single-valued y, dy for single valued x.
                headers = ["x", "y", "dy", "fy"]
                columns = [self.x, self.y, self.dy, theory]
            data = np.vstack(columns)
            outfile = basename + ".dat"
            with open(outfile, "w") as fd:
                fd.write("# " + "\t ".join(headers) + "\n")
                np.savetxt(fd, data.T)

        def plot(self, view=None):
            if self._plot is not None:
                kw = self._fetch_pars()
                self._plot(self.x, self.y, self.dy, self.theory(), view=view, **kw)
                return

            import pylab
            from .plotutil import coordinated_colors

            x = self.x
            if self.plot_x is not None:
                theory_x, theory_y = self.plot_x, self.theory(self.plot_x)
            else:
                theory_x, theory_y = x, self.theory()
            resid = self.residuals()

            if self._num_curves > 1:
                y, dy, theory_y, resid = self.y.T, self.dy.T, theory_y.T, resid.T
            else:
                y, dy, theory_y, resid = (v[:, None] for v in (self.y, self.dy, theory_y, resid))

            colors = tuple(coordinated_colors() for _ in range(self._num_curves))
            labels = self.labels

            # print "kw_plot",kw
            if view == "residual":
                _plot_resids(x, resid, colors, labels=labels, view=view)
            else:
                plot_ratio = 4
                h = pylab.subplot2grid((plot_ratio, 1), (0, 0), rowspan=plot_ratio - 1)
                for tick_label in h.get_xticklabels():
                    tick_label.set_visible(False)
                _plot_fits(data=(x, y, dy), theory=(theory_x, theory_y), colors=colors, labels=labels, view=view)
                # pylab.gca().xaxis.set_visible(False)
                # pylab.gca().spines['bottom'].set_visible(False)
                # pylab.gca().set_xticks([])

                pylab.subplot2grid((plot_ratio, 1), (plot_ratio - 1, 0), sharex=h)
                _plot_resids(x, resid, colors=colors, labels=labels, view=view)

        def register_webview_plot(
            self, plot_title: str, plot_function: Callable, change_with: Literal["parameter", "uncertainty"]
        ):
            # Plot function syntax: f(model, problem, state)
            # change_with = 'parameter' or 'uncertainty'

            self._webview_plots[plot_title] = dict(func=plot_function, change_with=change_with)

        @property
        def webview_plots(self):
            return self._webview_plots

    def _plot_resids(x, resid, colors, labels, view):
        import pylab

        pylab.axhline(y=1, ls="dotted", color="k")
        pylab.axhline(y=0, ls="solid", color="k")
        pylab.axhline(y=-1, ls="dotted", color="k")
        for k, color in enumerate(colors):
            pylab.plot(x, resid[:, k], ".", color=color["base"])
        pylab.gca().locator_params(axis="y", tight=True, nbins=4)
        pylab.xlabel(labels[0])
        pylab.ylabel("(f(x)-y)/dy")
        if view == "logx":
            pylab.xscale("log")
        elif view == "loglog":
            pylab.xscale("log")

    def _plot_fits(data, theory, colors, labels, view):
        import pylab

        x, y, dy = data
        theory_x, theory_y = theory
        for k, color in enumerate(colors):
            pylab.errorbar(x, y[:, k], yerr=dy[:, k], fmt=".", color=color["base"], label="_")
            pylab.plot(theory_x, theory_y[:, k], "-", color=color["dark"], label=labels[k + 2])
        # Note: no xlabel since it is supplied by the residual plot below this plot
        pylab.ylabel(labels[1])
        if len(colors) > 1:
            pylab.legend()
        if view == "log":
            pylab.xscale("linear")
            pylab.yscale("log")
        elif view == "logx":
            pylab.xscale("log")
            pylab.yscale("linear")
        elif view == "logy":
            pylab.xscale("linear")
            pylab.yscale("log")
        elif view == "loglog":
            pylab.xscale("log")
            pylab.yscale("log")
        else:  # view == 'linear'
            pylab.xscale("linear")
            pylab.yscale("linear")

    def plot_resid(x, resid):
        """
        **DEPRECATED**
        """
        import pylab

        pylab.axhline(y=1, ls="dotted", color="k")
        pylab.axhline(y=0, ls="solid", color="k")
        pylab.axhline(y=-1, ls="dotted", color="k")
        pylab.plot(x, resid, ".")
        pylab.gca().locator_params(axis="y", tight=True, nbins=4)
        pylab.ylabel("Residuals")

    def plot_err(x, y, dy, fy, view=None, **kw):
        """
        **DEPRECATED**: subclass Curve and override the plot function.

        Plot data *y* and error *dy* against *x*.

        *view* is one of linear, log, logx or loglog.
        """
        import pylab

        pylab.errorbar(x, y, yerr=dy, fmt=".")
        pylab.plot(x, fy, "-")
        if view == "log":
            pylab.xscale("linear")
            pylab.yscale("log")
        elif view == "logx":
            pylab.xscale("log")
            pylab.yscale("linear")
        elif view == "loglog":
            pylab.xscale("log")
            pylab.yscale("log")
        else:  # view == 'linear'
            pylab.xscale("linear")
            pylab.yscale("linear")

    _LOGFACTORIAL = np.array([log(np.prod(np.arange(1.0, k + 1))) for k in range(21)])

    def logfactorial(n):
        """Compute the log factorial for each element of an array"""
        result = np.empty(n.shape, dtype="double")
        idx = n <= 20
        result[idx] = _LOGFACTORIAL[np.asarray(n[idx], "int32")]
        n = n[~idx]
        result[~idx] = n * log(n) - n + log(n * (1 + 4 * n * (1 + 2 * n))) / 6 + log(pi) / 2
        return result

    class PoissonCurve(Curve):
        r"""
        Model a measurement with Poisson uncertainty.

        The nllf is calculated using Poisson probabilities, but the curve itself
        is displayed using the approximation that $\sigma_y \approx \sqrt(y)$.

        See :class:`Curve` for details.
        """

        def __init__(self, fn, x, y, name="", **fnkw):
            dy = sqrt(y) + (y == 0) if y is not None else None
            Curve.__init__(self, fn, x, y, dy, name=name, **fnkw)
            self._logfacty = logfactorial(y) if y is not None else None
            self._logfactysum = np.sum(self._logfacty)

        ## Assume gaussian residuals for now
        # def residuals(self):
        #    # TODO: provide individual probabilities as residuals
        #    # or perhaps the square roots --- whatever gives a better feel for
        #    # which points are driving the fit
        #    theory = self.theory()
        #    return np.sqrt(self.y * log(theory) - theory - self._logfacty)

        def nllf(self):
            theory = self.theory()
            if (theory <= 0).any():
                return 1e308
            return -sum(self.y * log(theory) - theory) + self._logfactysum

        def simulate_data(self, noise=None):
            theory = self.theory()
            self.y = np.random.poisson(theory)
            self.dy = sqrt(self.y) + (self.y == 0)
            self._logfacty = logfactorial(self.y)
            self._logfactysum = np.sum(self._logfacty)