"""
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)