Generating plots via plots attribute

example/lsqfit_basic_example.py

@@ -1,13 +1,15 @@
+"""Example application of lsqfit-gui on a "basic fit" from the lsqfit documentation.
 
-"""Example application of lsqfit-gui on a "basic fit" from the lsqfit documentation:
- https://lsqfit.readthedocs.io/en/latest/overview.html#basic-fits """
+See also https://lsqfit.readthedocs.io/en/latest/overview.html#basic-fits
+"""
 
 import lsqfitgui
+from lsqfitgui.plot.uncertainty import wrap_plot_gvar, plot_gvar
 import lsqfit
 import numpy as np
 import gvar as gv
 
-'''
+"""
 # An even more basic example; comment-out the other main() instead
 def main():
     x, y = make_data()              # collect fit data
@@ -16,62 +18,111 @@ def main():
     fit = lsqfit.nonlinear_fit(data=(x, y), fcn=fcn, prior=prior, p0=p0)
 
     lsqfitgui.run_server(fit)
-'''
+"""
+
 
 def main():
+    x, y = make_data()
+    y_eff_mass = np.array([])
+    dx = np.roll(x, -1) - x
+    y_eff_mass = np.log(y / np.roll(y, -1)) / dx
+    y_eff_wf = np.exp(y_eff_mass * x) * y
+
+    lsqfitgui.run_server(
+        fit_setup_function=generate_fit,
+        fit_setup_kwargs={"n_exp": 4},
+        meta_config=[
+            {"name": "n_exp", "type": "number", "min": 1, "max": 10, "step": 1}
+        ],
+        additional_plots=[
+            {
+                "name": "Effective mass",
+                "fcn": eff_mass,
+                "kwargs": {},  # passed to eff_mass function directly
+                "static_plot_gvar": {  # calls plot_gvar on the same figure
+                    "x": x,
+                    "y": y_eff_mass,
+                    "kind": "errorbars",
+                    "scatter_kwargs": {"name": "Data"},
+                },
+            },
+            {
+                "name": "Effective WF",
+                # "fcn": eff_wf, # both would do the same thing
+                "fcn": plot_eff_wf,
+                "kwargs": {},  # passed to eff_mass function directly
+                "static_plot_gvar": {  # calls plot_gvar on the same figure
+                    "x": x,
+                    "y": y_eff_wf,
+                    "kind": "errorbars",
+                    "scatter_kwargs": {"name": "Data"},
+                },
+            },
+        ],
+    )
 
-   x , y = make_data()
-   y_eff_mass = np.array([])
-   dx = np.roll(x, -1) - x
-   y_eff_mass = np.log(y/np.roll(y, -1))/dx
-   y_eff_wf = np.exp(y_eff_mass *x)*y
-
-   lsqfitgui.run_server(
-      fit_setup_function=generate_fit,
-      fit_setup_kwargs={"n_exp": 4},
-      meta_config=[
-         {"name": "n_exp", "type": "number", "min": 1, "max": 10, "step": 1}
-      ], 
-      plot_fcns={'eff_mass' : eff_mass}, #, 'eff_wf' :eff_wf},
-      transformed_y={'eff_mass' : y_eff_mass, 'eff_wf' : y_eff_wf}
-   )
 
 def generate_fit(**meta):
-   """Generate a fit for specified meta information."""
-   x, y = make_data()              # collect fit data
-   p0 = None                       # make larger fits go faster (opt.)
-   prior = make_prior(nexp=meta["n_exp"])
-   fit = lsqfit.nonlinear_fit(data=(x, y), fcn=fcn, prior=prior, p0=p0)
-   fit.meta = meta
-   return fit
-
-def fcn(x, p):                     # function used to fit x, y data
-   a = p['a']                      # array of a[i]s
-   E = p['E']                      # array of E[i]s
-   return np.array(sum(ai * np.exp(-Ei * x) for ai, Ei in zip(a, E)))
-
-#@lsqfitgui.plot
+    """Generate a fit for specified meta information."""
+    x, y = make_data()  # collect fit data
+    p0 = None  # make larger fits go faster (opt.)
+    prior = make_prior(nexp=meta["n_exp"])
+    fit = lsqfit.nonlinear_fit(data=(x, y), fcn=fcn, prior=prior, p0=p0)
+    fit.meta = meta
+    return fit
+
+
+def fcn(x, p):
+    """Multi exponential fit function.
+
+    `sum([a_i * exp(-E_i * x)])`
+
+    Arguemnts of prior p:
+        a: array of a[i]s
+        E: array of E[i]s
+    """
+    return np.array(sum(ai * np.exp(-Ei * x) for ai, Ei in zip(p["a"], p["E"])))
+
+
+@wrap_plot_gvar(kind="band", scatter_kwargs={"opacity": 0.5, "name": "Fit"})
 def eff_mass(x, p):
-   return np.log(fcn(x, p)/fcn(x+1, p))
+    return np.log(fcn(x, p) / fcn(x + 1, p))
 
-#@lsqfitgui.plot
+
+@wrap_plot_gvar(kind="band", scatter_kwargs={"opacity": 0.5, "name": "Fit"})
 def eff_wf(x, p):
-   return np.exp(np.log(fcn(x, p)/fcn(x+1, p))*x) *fcn(x, p) 
+    return np.exp(np.log(fcn(x, p) / fcn(x + 1, p)) * x) * fcn(x, p)
+
+
+# The above call does effectively the below code
+
 
-def make_prior(nexp):              # make priors for fit parameters
-   prior = gv.BufferDict()         # any dictionary works
-   prior['a'] = [gv.gvar(0.5, 0.4) for i in range(nexp)]
-   prior['E'] = [gv.gvar(i+1, 0.4) for i in range(nexp)]
-   return prior
+def plot_eff_wf(fit):
+    ff = fcn(fit.x, fit.p)
+    return plot_gvar(
+        fit.x,
+        np.exp(np.log(ff / fcn(fit.x + 1, fit.p)) * fit.x) * ff,
+        kind="band",
+        scatter_kwargs={"opacity": 0.5, "name": "Fit"},
+    )
 
-def make_data():                   # assemble fit data
-   x = np.array([  5.,   6.,   7.,   8.,   9.,  10.,  12.,  14.])
-   ymean = np.array(
+
+def make_prior(nexp):  # make priors for fit parameters
+    prior = gv.BufferDict()  # any dictionary works
+    prior["a"] = [gv.gvar(0.5, 0.4) for i in range(nexp)]
+    prior["E"] = [gv.gvar(i + 1, 0.4) for i in range(nexp)]
+    return prior
+
+
+def make_data():  # assemble fit data
+    # fmt: off
+    x = np.array([  5.,   6.,   7.,   8.,   9.,  10.,  12.,  14.])
+    ymean = np.array(
        [  4.5022829417e-03,   1.8170543788e-03,   7.3618847843e-04,
           2.9872730036e-04,   1.2128831367e-04,   4.9256559129e-05,
           8.1263644483e-06,   1.3415253536e-06]
        )
-   ycov = np.array(
+    ycov = np.array(
        [[ 2.1537808808e-09,   8.8161794696e-10,   3.6237356558e-10,
           1.4921344875e-10,   6.1492842463e-11,   2.5353714617e-11,
           4.3137593878e-12,   7.3465498888e-13],
@@ -97,7 +148,9 @@ def make_data():                   # assemble fit data
           5.5986403288e-14,   2.3914810594e-14,   1.0244738582e-14,
           1.8976295583e-15,   3.5672355835e-16]]
        )
-   return x, gv.gvar(ymean, ycov)
+    # fmt: on
+    return x, gv.gvar(ymean, ycov)
+
 
-if __name__ == '__main__':
-    main()
+if __name__ == "__main__":
+    main()

example/lsqfit_stability_example.py

@@ -0,0 +1,168 @@
+"""Example application of lsqfit-gui on a "basic fit" from the lsqfit documentation.
+
+See also https://lsqfit.readthedocs.io/en/latest/overview.html#basic-fits
+"""
+from itertools import product
+
+import numpy as np
+import pandas as pd
+
+import gvar as gv
+import lsqfit
+
+import plotly.express as px
+
+import lsqfitgui
+
+
+T_MIN_RANGE = range(0, 4)
+T_MAX_RANGE = range(8, 9)
+N_EXP_RANGE = range(2, 5)
+RANGE_PRODUCT = list(product(T_MIN_RANGE, T_MAX_RANGE, N_EXP_RANGE))
+KEYS = ("t_min", "t_max", "n_exp")
+
+
+def k2s(**kwargs) -> str:
+    """Maps multiple arguments to a string.
+
+    Needed since gui labels need to be strings.
+    """
+    return ", ".join([f"{key}={value}" for key, value in kwargs.items()])
+
+
+def generate_fit():
+    """Generate a fit for specified meta information."""
+    xx = {}
+    yy = gv.BufferDict()
+    prior = gv.BufferDict()
+    p0 = None
+    for t_min, t_max, n_exp in RANGE_PRODUCT:
+        kwargs = dict(zip(KEYS, (t_min, t_max, n_exp)))
+
+        xx[k2s(**kwargs)], yy[k2s(**kwargs)] = make_data(t_min=t_min, t_max=t_max)
+        for key, val in make_prior(nexp=n_exp).items():
+            prior[k2s(**kwargs, key=key)] = val
+
+    fit = lsqfit.nonlinear_fit(data=(xx, yy), fcn=fcn, prior=prior, p0=p0,)
+    return fit
+
+
+def fcn(x, p, individual_priors: bool = False):
+    """Multi exponential fit function.
+
+    `sum([a_i * exp(-E_i * x)])`
+
+    Arguemnts of prior p:
+        a: array of a[i]s
+        E: array of E[i]s
+    """
+    out = {}
+    for args in RANGE_PRODUCT:
+
+        kwargs = dict(zip(KEYS, args))
+        a = p[k2s(**kwargs, key="a")]
+        E = p[k2s(**kwargs, key="E")]
+        xx = x[k2s(**kwargs)]
+
+        out[k2s(**kwargs)] = np.array(
+            sum([ai * np.exp(-Ei * xx) for ai, Ei in zip(a, E)])
+        )
+    return out
+
+
+def make_prior(nexp):  # make priors for fit parameters
+    prior = gv.BufferDict()  # any dictionary works
+    prior["a"] = [gv.gvar(0.5, 0.4) for i in range(nexp)]
+    prior["E"] = [gv.gvar(i + 1, 0.4) for i in range(nexp)]
+    return prior
+
+
+def make_data(t_min=0, t_max=8):  # assemble fit data
+
+    if t_min < 0 or t_max > 8:
+        raise ValueError("t rannge not valid")
+    # fmt: off
+    x = np.array([  5.,   6.,   7.,   8.,   9.,  10.,  12.,  14.])
+    ymean = np.array(
+       [  4.5022829417e-03,   1.8170543788e-03,   7.3618847843e-04,
+          2.9872730036e-04,   1.2128831367e-04,   4.9256559129e-05,
+          8.1263644483e-06,   1.3415253536e-06]
+       )
+    ycov = np.array(
+       [[ 2.1537808808e-09,   8.8161794696e-10,   3.6237356558e-10,
+          1.4921344875e-10,   6.1492842463e-11,   2.5353714617e-11,
+          4.3137593878e-12,   7.3465498888e-13],
+       [  8.8161794696e-10,   3.6193461816e-10,   1.4921610813e-10,
+          6.1633547703e-11,   2.5481570082e-11,   1.0540958082e-11,
+          1.8059692534e-12,   3.0985581496e-13],
+       [  3.6237356558e-10,   1.4921610813e-10,   6.1710468826e-11,
+          2.5572230776e-11,   1.0608148954e-11,   4.4036448945e-12,
+          7.6008881270e-13,   1.3146405310e-13],
+       [  1.4921344875e-10,   6.1633547703e-11,   2.5572230776e-11,
+          1.0632830128e-11,   4.4264622187e-12,   1.8443245513e-12,
+          3.2087725578e-13,   5.5986403288e-14],
+       [  6.1492842463e-11,   2.5481570082e-11,   1.0608148954e-11,
+          4.4264622187e-12,   1.8496194125e-12,   7.7369196122e-13,
+          1.3576009069e-13,   2.3914810594e-14],
+       [  2.5353714617e-11,   1.0540958082e-11,   4.4036448945e-12,
+          1.8443245513e-12,   7.7369196122e-13,   3.2498644263e-13,
+          5.7551104112e-14,   1.0244738582e-14],
+       [  4.3137593878e-12,   1.8059692534e-12,   7.6008881270e-13,
+          3.2087725578e-13,   1.3576009069e-13,   5.7551104112e-14,
+          1.0403917951e-14,   1.8976295583e-15],
+       [  7.3465498888e-13,   3.0985581496e-13,   1.3146405310e-13,
+          5.5986403288e-14,   2.3914810594e-14,   1.0244738582e-14,
+          1.8976295583e-15,   3.5672355835e-16]]
+       )
+    # fmt: on
+    return x[t_min:t_max], gv.gvar(ymean, ycov)[t_min:t_max]
+
+
+def plot_stability(fit, **kwargs):
+    """Creates a stability plot."""
+    delta_n_exp = N_EXP_RANGE[-1] - N_EXP_RANGE[0]
+
+    data = []
+    e0s = []
+    for t_min, t_max, n_exp in RANGE_PRODUCT:
+        kwargs = dict(zip(KEYS, (t_min, t_max, n_exp)))
+        e0 = fit.p[k2s(**kwargs, key="E")][0]
+        data.append(
+            {
+                "t_min": t_min  # slightly shift x values for differnt markers
+                + (n_exp - N_EXP_RANGE[0] - delta_n_exp / 2) / delta_n_exp * 0.1,
+                "t_max": str(t_max),  # This makes plotly use discrete colors
+                "n_exp": str(n_exp),
+                "E0 mean": e0.mean,
+                "E0 sdev": e0.sdev,
+            }
+        )
+        e0s.append(e0)
+    df = pd.DataFrame(data)
+
+    fig = px.scatter(
+        df,
+        x="t_min",
+        y="E0 mean",
+        error_y="E0 sdev",
+        color="n_exp",
+        symbol="n_exp",
+        facet_row="t_max",
+    )
+    mean = gv.mean(e0s).mean()
+    sdev = np.sqrt(gv.mean(e0s).std(ddof=1) ** 2 + gv.sdev(e0s).mean() ** 2)
+    fig.add_hrect(mean - sdev, mean + sdev, opacity=0.5, fillcolor="gray")
+    fig.add_hline(mean, line_color="gray")
+    return fig
+
+
+def main():
+    """Run the stability plot server."""
+    lsqfitgui.run_server(
+        fit=generate_fit(),
+        additional_plots=[{"name": "Stability plot", "fcn": plot_stability}],
+    )
+
+
+if __name__ == "__main__":
+    main()