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.
[docs] @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}")
[docs] def update(self): self._cached_theory = None
[docs] def parameters(self): return self.pars
[docs] def numpoints(self): return np.prod(self.y.shape)
[docs] 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
[docs] 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
[docs] def residuals(self): return (self.theory() - self.y) / self.dy
[docs] def nllf(self): r = self.residuals() return 0.5 * np.sum(r**2)
[docs] 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)
[docs] 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 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)
[docs] 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")
[docs] 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
[docs] 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)
[docs] def nllf(self): theory = self.theory() if (theory <= 0).any(): return 1e308 return -sum(self.y * log(theory) - theory) + self._logfactysum
[docs] 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)