# 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, Union, Optional, Any, Dict, List, Callable
    from dataclasses import dataclass

    import numpy as np
    from numpy import log, pi, sqrt
    from numpy.typing import ArrayLike

    from .util import NDArray
    from .parameter import Parameter, ValueProtocol

    # import sys; sys.setrecursionlimit(60)  # For debugging getattr/setattr

    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

    # TODO: separate the dataclass from the builder function
    # As currently implemented we need artificial parameters
    # "pars" and "state" to allow deserialization of the dataclass
    # but these are parameters that the user should never touch.
    # Instead, define a cleaner Curve object which can be dumped
    # and loaded independent of the calling parameters.

    @dataclass(init=False, eq=False)
    class Curve:
        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 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.
        """

        # TODO: Add resolution and uncertainty in x. Resolution will require oversampling.
        # TODO: Move data into its own class (pre 1.0?)
        # The data class can deal with x, y, dy, resolution, jitter, plotter, plot_x, labels
        fn: Callable
        x: NDArray
        y: NDArray
        dy: NDArray
        """Data uncertainty"""
        name: str
        """Name of the model"""
        pars: Dict[str, Parameter] = None  # Needs to be None initially for getattr/setattr to work
        """Fittable parameters to the model"""
        state: Dict[str, Any]
        """Nonfittable parameters set during initialization. Values should be serializable."""
        labels: List[str]
        plot_x: Optional[NDArray]
        plotter: Optional[Callable]

        def __init__(
            self,
            fn: Callable,
            x: ArrayLike,
            y: ArrayLike,
            dy: Optional[ArrayLike] = None,
            name: Optional[str] = "",
            labels: Optional[List[str]] = None,
            plotter: Optional[Callable[..., None]] = None,
            plot_x: Optional[NDArray] = None,
            plot: Optional[Callable] = None,
            pars: Optional[Dict[str, Parameter]] = None,
            state: Optional[Dict[str, Any]] = None,
            **kwargs,
        ):
            self.x, self.y = np.asarray(x), np.asarray(y)
            if dy is None:
                self.dy = np.ones_like(y)
            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

            self.plot_x = plot_x
            # TODO: drop plot= option; likely nobody is using it, and this is a 1.0 release
            if plot is not None:
                warnings.warn("***DEPRECATED*** use Curve(..., plotter=...) instead of plot=...")
            self.plotter = plotter

            # _assign_pars(state, self)  # ... and state variables as well
            self._cached_theory = None
            self._webview_plots = {}

            self.fn = fn
            self.name = name  # if name else fn.__name__ + " "
            if pars is None:
                pars, state = _parse_pars(self.fn, init=kwargs, skip=1, name=name)
            # else: print("restoring", {k: float(v) for k, v in pars.items()}, state)
            # print("parsed", self.fn, pars, state)

            # 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.state = state

            # Need to do collision test before making parameters available as attributes
            # otherwise every parameter is a collision. Need to assign the pars object
            # last so that we know all the variables that could collide.
            collisions = [key for key in pars.keys() if hasattr(self, key)]
            if collisions:
                raise ValueError(f"parameter names shadowed by Curve object: {', '.join(collisions)}")
            self.pars = pars

        # Allow dot access to members of the parameter dictionary. Existing attributes
        # of the object take precedence.
        def __setattr__(self, key, value):
            # print(f"setting {key}")
            if self.pars and key in self.pars and key not in self.__dict__:
                if not isinstance(value, ValueProtocol):
                    raise TypeError("Can only assign parameter or expression to a parameter slot")
                self.__dict__["pars"][key] = value
            else:
                super().__setattr__(key, value)

        def __getattr__(self, key):
            # print(f"getting {key}")
            if self.pars and key in self.pars:
                return self.pars[key]
            raise AttributeError(f"{type(self)!r} has no attribute {key!r}")

        def update(self):
            self._cached_theory = None

        def parameters(self):
            return self.pars

        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.fn(x, **kw)

        def _fetch_pars(self):
            kw = {k: float(v) for k, v in self.pars.items()}
            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.plotter is not None:
                kw = self._fetch_pars()
                self.plotter(self.x, self.y, self.dy, self.theory(), view=view, **kw)
                return

            import matplotlib.pyplot as plt
            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))
            # clear figure after calling coordinated_colors, which can create an axis object
            plt.clf()
            labels = self.labels

            # print "kw_plot",kw
            if view == "residual":
                _plot_resids(x, resid, colors, labels=labels, view=view)
            else:
                plot_ratio = 4
                h = plt.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)
                # plt.gca().xaxis.set_visible(False)
                # plt.gca().spines['bottom'].set_visible(False)
                # plt.gca().set_xticks([])

                plt.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 matplotlib.pyplot as plt

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

    def _plot_fits(data, theory, colors, labels, view):
        import matplotlib.pyplot as plt

        x, y, dy = data
        theory_x, theory_y = theory
        for k, color in enumerate(colors):
            plt.errorbar(x, y[:, k], yerr=dy[:, k], fmt=".", color=color["base"], label="_")
            plt.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
        plt.ylabel(labels[1])
        if len(colors) > 1:
            plt.legend()
        if view == "log":
            plt.xscale("linear")
            plt.yscale("log")
        elif view == "logx":
            plt.xscale("log")
            plt.yscale("linear")
        elif view == "logy":
            plt.xscale("linear")
            plt.yscale("log")
        elif view == "loglog":
            plt.xscale("log")
            plt.yscale("log")
        else:  # view == 'linear'
            plt.xscale("linear")
            plt.yscale("linear")

    def plot_resid(x, resid):
        """
        **DEPRECATED**
        """
        import matplotlib.pyplot as plt

        plt.axhline(y=1, ls="dotted", color="k")
        plt.axhline(y=0, ls="solid", color="k")
        plt.axhline(y=-1, ls="dotted", color="k")
        plt.plot(x, resid, ".")
        plt.gca().locator_params(axis="y", tight=True, nbins=4)
        plt.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 matplotlib.pyplot as plt

        plt.errorbar(x, y, yerr=dy, fmt=".")
        plt.plot(x, fy, "-")
        if view == "log":
            plt.xscale("linear")
            plt.yscale("log")
        elif view == "logx":
            plt.xscale("log")
            plt.yscale("linear")
        elif view == "loglog":
            plt.xscale("log")
            plt.yscale("log")
        else:  # view == 'linear'
            plt.xscale("linear")
            plt.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

    @dataclass(init=False, eq=False)
    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)$.

        Note that the *dy* argument is ignored. It is only present for deserialization
        of the Curve base class.

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

        def __init__(self, fn, x, y, dy=None, name="", **fnkw):
            # dy ignored. It is needed for dataclass deserialization
            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)