• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

APN-Pucky / smpl / 15207604833

23 May 2025 09:59AM UTC coverage: 88.129% (-1.7%) from 89.784%
15207604833

push

github

web-flow
copy deprecated files from scipy (#339)

* copy deprecated files from scipy

* fix: tests

* fix: relaxe test

24 of 46 new or added lines in 1 file covered. (52.17%)

1 existing line in 1 file now uncovered.

980 of 1112 relevant lines covered (88.13%)

0.88 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

92.74
/smpl/plot.py
1
"""Simplified plotting."""
2
import itertools
1✔
3

4
import matplotlib
1✔
5
import matplotlib.pylab as pylab
1✔
6
import matplotlib.pyplot as plt
1✔
7
import numpy as np
1✔
8
import sympy
1✔
9
import uncertainties
1✔
10
import uncertainties.unumpy as unp
1✔
11
from matplotlib.pyplot import *
1✔
12

13
# local imports
14
from smpl import doc
1✔
15
from smpl import fit as ffit
1✔
16
from smpl import interpolate, io, stat, util, wrap
1✔
17

18

19
def set_plot_style():
1✔
20
    # fig_size = (8, 6)
21
    # fig_legendsize = 14
22
    # fig_labelsize = 12 # ‘xx-small’, ‘x-small’, ‘small’, ‘medium’, ‘large’, ‘x-large’, ‘xx-large’.
23
    params = {
×
24
        "legend.fontsize": "x-large",
25
        "figure.figsize": (8, 6),
26
        "axes.labelsize": "x-large",
27
        "axes.titlesize": "x-large",
28
        "xtick.labelsize": "x-large",
29
        "ytick.labelsize": "x-large",
30
    }
31
    pylab.rcParams.update(params)
×
32
    matplotlib.rcParams.update(params)
×
33
    # matplotlib.rcParams.update({'font.size': fig_labelsize})
34
    # colors = dict(mcolors.BASE_COLORS, **mcolors.CSS4_COLORS)
35

36

37
unv = unp.nominal_values
1✔
38
usd = unp.std_devs
1✔
39

40
default = {
1✔
41
    #    'params'        :[None      ,"Initial fit parameters",
42
    "title": [None, "Plot title"],
43
    "xlabel": [
44
        "",
45
        "X axis label",
46
    ],
47
    "ylabel": [
48
        "",
49
        "Y axis label",
50
    ],
51
    "label": [
52
        None,
53
        "Legend name of plotted ``data``",
54
    ],
55
    "fmt": [
56
        ".",
57
        "Format for plotting fit function (could be a line style, marker or a combination of both or 'step')",
58
    ],
59
    "where": [
60
        "mid",
61
        "Where to place the ticks. Possible values are 'pre', 'post', 'mid', 'edge'. Might be incompatible with uncertainties.",
62
    ],
63
    "units": [
64
        None,
65
        "Units of the fit parameters as strings. Displayed in the Legend",
66
    ],
67
    "save": [
68
        None,
69
        " File to save the plot",
70
    ],
71
    "lpos": [
72
        0,
73
        "Legend position",
74
    ],
75
    "tight": [
76
        True,
77
        "tight_layout",
78
    ],
79
    #          'frange'        :[None      ,"Limit the fit to given range. First integer is the lowest and second the highest index.",],
80
    "prange": [
81
        None,
82
        "Limit the plot of the fit to given range",
83
    ],
84
    "sigmas": [
85
        0,
86
        "Color the array of given ``sigma`` times uncertainty. Only works if the fit function is coded with ``unp``",
87
    ],
88
    "data_sigmas": [
89
        1,
90
        "Color the array of given ``sigma`` times uncertainty. Only works if the data has uncertainties",
91
    ],
92
    "init": [False, "Initialize a new plot"],
93
    "ss": [
94
        True,
95
        "save, add legends and grid to the plot",
96
    ],
97
    "also_data": [True, " also plot the data"],
98
    "also_fit": [
99
        True,
100
        "also plot the fit",
101
    ],
102
    "auto_fit": [
103
        False,
104
        "automatically fit",
105
    ],
106
    "logy": [
107
        False,
108
        "logarithmic x axis",
109
    ],
110
    "logx": [
111
        False,
112
        "logarithmic y axis",
113
    ],
114
    "function_color": [
115
        None,
116
        "Color of the function plot",
117
    ],
118
    "data_color": [
119
        None,
120
        "Color of the data plot",
121
    ],
122
    "fit_color": [
123
        None,
124
        "Color of the fit plot",
125
    ],
126
    "fit_fmt": [
127
        "-",
128
        "Format of the fit plot",
129
    ],
130
    "residue": [
131
        False,
132
        "Display difference between fit and data in a second plot",
133
    ],
134
    "residue_err": [
135
        True,
136
        "Differences between fit and data will have errorbars",
137
    ],
138
    "show": [
139
        False,
140
        "Call plt.show()",
141
    ],
142
    "size": [
143
        None,
144
        "Size of the plot as a tuple (x,y). Only has an effect if ``init`` is True",
145
    ],
146
    "number_format": [
147
        io.gf(4),
148
        "Format to display numbers.",
149
    ],
150
    # ,          'selector'      :[ None     ,"Function that takes ``x`` and ``y`` as parameters and returns an array mask in order to limit the data points for fitting. Alternatively a mask for selecting elements from datax and datay.",],
151
    # ,          'fixed_params'  :[ True     ,"Enable fixing parameters by choosing the same-named variables from ``kwargs``.",],
152
    # ,          'sortbyx'       :[ True     , "Enable sorting the x and y data so that x is sorted.",],
153
    "interpolate": [False, "Enable interpolation of the data."],
154
    "interpolate_fmt": [
155
        "-",
156
        "Either format string or linestyle tuple.",
157
    ],
158
    "interpolate_label": [
159
        None,
160
        "Label for the interpolation.",
161
    ],
162
    "extrapolate": [
163
        True,
164
        "Enable extrapolation of whole data if fit range is limited by ``frange`` or ``fselector``.",
165
    ],
166
    "extrapolate_min": [
167
        None,
168
        "Lower extrapolation bound",
169
    ],
170
    "extrapolate_max": [
171
        None,
172
        "Higher extrapolation bound",
173
    ],
174
    "extrapolate_fmt": [
175
        "--",
176
        "Format of the extrapolation line",
177
    ],
178
    "extrapolate_hatch": [
179
        r"||",
180
        "Extrapolation shape/hatch for filled area in case of ``sigmas``>0. See https://matplotlib.org/stable/gallery/shapes_and_collections/hatch_style_reference.html",
181
    ],
182
    "bbox_to_anchor": [
183
        None,
184
        "Position in a tuple (x,y),Shift position of the legend out of the main pane. ",
185
    ],
186
    "ncol": [
187
        None,
188
        "Columns in the legend if used with ``bbox_to_anchor``.",
189
    ],
190
    "steps": [
191
        1000,
192
        "resolution of the plotted function",
193
    ],
194
    "fitinline": [
195
        False,
196
        "No newlines for each fit parameter",
197
    ],
198
    "grid": [
199
        True,
200
        "Enable grid for the plot",
201
    ],
202
    "hist": [
203
        False,
204
        "Enable histogram plot",
205
    ],
206
    "stairs": [
207
        False,
208
        "Enable stair plot",
209
    ],
210
    "capsize": [5, "size of cap on error bar plot"],
211
    "axes": [None, "set current axis"],
212
    "linestyle": [None, "linestyle, only active if `fmt`=None"],
213
    "xspace": [
214
        np.linspace,
215
        "xspace gets called with xspace(xmin,xmax,steps) in :func:`function` to get the points of the function that will be drawn.",
216
    ],
217
    "alpha": [0.2, "alpha value for the fill_between plot"],
218
}
219

220

221
@doc.append_doc(ffit.fit_kwargs)
1✔
222
@doc.append_str("\t")
1✔
223
@doc.append_str(doc.array_table(default, init=False))
1✔
224
@doc.append_str(
1✔
225
    doc.array_table({"plot_kwargs        ": ["default", "description"]}, bottom=False)
226
)
227
@doc.append_str("\n\n")
1✔
228
def plot_kwargs(kwargs):
1✔
229
    """Set default :func:`plot_kwargs` if not set."""
230
    kwargs = ffit.fit_kwargs(kwargs)
1✔
231
    for k, v in default.items():
1✔
232
        if not k in kwargs:
1✔
233
            kwargs[k] = v[0]
1✔
234
    return kwargs
1✔
235

236

237
# TODO optimize to minimize number of calls to function
238
# @append_doc(default_kwargs)
239

240

241
def fit(func, *adata, **kwargs):
1✔
242
    """
243
    Fit and plot function to datax and datay.
244

245
    Parameters
246
    ----------
247
    datax : array_like
248
        X data either as ``unp.uarray`` or ``np.array`` or ``list``
249
    datay : array_like
250
        Y data either as ``unp.uarray`` or ``np.array`` or ``list``
251
    function : func
252
        Fit function with parameters: ``x``, ``params``
253
    **kwargs : optional
254
        see :func:`plot_kwargs`.
255
    Fit parameters can be fixed via ``kwargs`` eg. ``a=5``.
256

257
    Returns
258
    -------
259
    array_like
260
        Optimized fit parameters of ``function`` to ``datax`` and ``datay``.
261
        If ``datay`` is complex, both the real and imaginary part are returned.
262

263
    Examples
264
    --------
265

266
    .. plot::
267
        :include-source:
268

269
        >>> from smpl import functions as f
270
        >>> from smpl import plot
271
        >>> param = plot.fit([0,1,2],[0,1,2],f.line)
272
        >>> plot.unv(param).round()[0]
273
        1.0
274

275
    """
276
    function = func
1✔
277
    if "function" in kwargs:
1✔
278
        function = kwargs["function"]
1✔
279
        del kwargs["function"]
1✔
280
        adata = [func, *adata]
1✔
281
    # Fix parameter order if necessary
282
    elif isinstance(function, (list, tuple, np.ndarray)):
1✔
283
        adata = [adata[-1], function, *adata[:-1]]
1✔
284
        function = adata[0]
1✔
285
        adata = adata[1:]
1✔
286
    if util.true("bins", kwargs):
1✔
287
        # yvalue will be overwritten
288
        ndata = [*adata, *adata]
1✔
289
        for i, o in enumerate(adata):
1✔
290
            ndata[2 * i] = o
1✔
291
            ndata[2 * i + 1] = o * 0
1✔
292
        adata = ndata
1✔
293

294
    assert len(adata) % 2 == 0, "data must be pairs of x and y data"
1✔
295
    if len(adata) == 2:
1✔
296
        datax, datay = adata
1✔
297
    else:
298
        rs = []
1✔
299
        for i in range(0, len(adata), 2):
1✔
300
            datax, datay = adata[i], adata[i + 1]
1✔
301
            if util.true("bins", kwargs):
1✔
302
                rs.append(fit(function, datax, **kwargs))
1✔
303
            else:
304
                rs.append(fit(function, datax, datay, **kwargs))
1✔
305
        return rs
1✔
306

307
    kwargs = plot_kwargs(kwargs)
1✔
308

309
    if np.any(np.iscomplex(datay)):
1✔
310
        label = util.get("label", kwargs, "")
1✔
311
        kwargs["label"] = label + "(real)"
1✔
312
        r = fit(datax, datay.real, function=function, **kwargs)
1✔
313
        kwargs["label"] = label + "(imag)"
1✔
314
        i = fit(datax, datay.imag, function=function, **kwargs)
1✔
315
        return r, i
1✔
316
    if kwargs["auto_fit"]:
1✔
317
        best_f, best_ff, lambda_f = ffit.auto(datax, datay, function, **kwargs)
1✔
318
        if best_f is not None:
1✔
319
            del kwargs["auto_fit"]
1✔
320
            fit(datax, datay, best_f, **kwargs)
1✔
321
        return best_f, best_ff, lambda_f
1✔
322
    if kwargs["also_fit"] == False and kwargs["label"] is None and kwargs["lpos"] == 0:
1✔
323
        kwargs["lpos"] = -1
1✔
324
    return _fit_impl(datax, datay, function, **kwargs)
1✔
325

326

327
def _fit_impl(datax, datay, function, **kwargs):
1✔
328
    x = None
1✔
329
    y = None
1✔
330
    rfit = None
1✔
331
    ifit = None
1✔
332
    fig = None
1✔
333
    fig = init_plot(kwargs)
1✔
334
    ll = None
1✔
335
    if kwargs["also_data"]:
1✔
336
        ll = plt_data(datax, datay, **kwargs).get_color()
1✔
337
    if kwargs["interpolate"]:
1✔
338
        ifit, _, x, y = plt_interpolate(datax, datay, icolor=ll, **kwargs)
1✔
339
    if kwargs["also_fit"]:
1✔
340
        assert function is not None, "function must be given"
1✔
341
        rfit, kwargs["fit_color"], _, _ = plt_fit(datax, datay, function, **kwargs)
1✔
342
    if kwargs["ss"]:
1✔
343
        kwargs["oldshow"] = kwargs["show"]
1✔
344
        kwargs["show"] = kwargs["show"] and not kwargs["residue"]
1✔
345
        save_plot(**kwargs)
1✔
346
        kwargs["show"] = kwargs["oldshow"]
1✔
347
    if kwargs["residue"] and fig is not None:
1✔
348
        plt_residue(datax, datay, function, rfit, fig, **kwargs)
1✔
349
    if not kwargs["also_fit"] and kwargs["interpolate"]:
1✔
350
        return (ifit, x, y)
1✔
351
        # return ifit
352
    return rfit
1✔
353

354

355
# @append_doc(default_kwargs)
356

357

358
def data(*data, function=None, **kwargs):
1✔
359
    """
360
    Plot datay against datax via :func:`fit`
361

362
    Parameters
363
    ----------
364
    datax : array_like
365
        X data either as ``unp.uarray`` or ``np.array`` or ``list``
366
    datay : array_like
367
        Y data either as ``unp.uarray`` or ``np.array`` or ``list``
368
    function : func,optional
369
        Fit function with parameters: ``x``, ``params``
370
    **kwargs : optional
371
        see :func:`plot_kwargs`.
372
    Returns
373
    -------
374
    array_like
375
        Optimized fit parameters of ``function`` to ``datax`` and ``datay``
376
    """
377
    if "also_fit" not in kwargs:
1✔
378
        kwargs["also_fit"] = False
1✔
379
    kwargs = plot_kwargs(kwargs)
1✔
380
    return fit(function=function, *data, **kwargs)
1✔
381

382

383
# @append_doc(default_kwargs)
384

385

386
def auto(*adata, funcs=None, **kwargs):
1✔
387
    """
388
    Automatically loop over functions and fit the best one.
389

390
    Parameters
391
    ----------
392
    funcs : function array
393
        functions to consider as fit. Default all ``smpl.functions``.
394
    **kwargs : optional
395
        see :func:`plot_kwargs`.
396

397
    Returns
398
    -------
399
    The best fit function and it's parameters. Also a lambda function where the parameters are already applied.
400

401

402

403
    """
404
    if "auto_fit" not in kwargs:
1✔
405
        kwargs["auto_fit"] = True
1✔
406
    kwargs = plot_kwargs(kwargs)
1✔
407
    return fit(function=funcs, *adata, **kwargs)
1✔
408

409

410
def _function(func, xfit,fmt="-",label=None,function_color=None,sigmas=0.0,alpha=0.4, **kwargs):
1✔
411

412
    #kargs = {}
413
    #if util.has("fmt", kwargs):
414
    #    kargs["fmt"] = kwargs["fmt"]
415
    #if util.has("label", kwargs) and kwargs["label"] != "":
416
    #    kargs["label"] = kwargs["label"]
417
    #if util.has("function_color", kwargs) and kwargs["function_color"] != "":
418
    #    kargs["color"] = kwargs["function_color"]
419
    #if util.has("sigmas", kwargs) and kwargs["sigmas"] != "":
420
    #    kargs["sigmas"] = kwargs["sigmas"]
421
    #if util.has("alpha", kwargs) and kwargs["alpha"] != "":
422
    #    kargs["alpha"] = kwargs["alpha"]
423
    #__function(func, xfit,  **kargs)
424

425
    if label == "":
1✔
426
        label = None
×
427
    if function_color == "":
1✔
428
        function_color = None
×
429
    if sigmas == "":
1✔
430
        sigmas = 0.0
×
431
    if alpha == "":
1✔
432
        alpha = 0.4
×
433
    __function(func, xfit, fmt=fmt, label=label,color=function_color,sigmas=sigmas, alpha=alpha, **kwargs)
1✔
434

435

436
def plt_plt(x, y, fmt, color, label, linestyle,**kwargs):
1✔
437
    if linestyle is None and fmt is not None:
1✔
438
        return plt.plot(x, y, fmt, label=label, color=color, **kwargs)
1✔
439
    elif linestyle is not None and fmt is None:
1✔
440
        return plt.plot(x, y, label=label, color=color, linestyle=linestyle,**kwargs)
×
441
    elif linestyle is None and fmt is None:
1✔
442
        return plt.plot(x, y, label=label, color=color,**kwargs)
1✔
443

444

445
def __function(
1✔
446
    gfunc,
447
    xlinspace,
448
    fmt="-",
449
    label=None,
450
    color=None,
451
    hatch=None,
452
    sigmas=0.0,
453
    linestyle=None,
454
    alpha=0.4,
455
    **kwargs
456
):
457
    # filter unused bad kwargs here to avoid passing them down
458
    # TODO it would be better to not pass them down in the first place
459
    for key in ["xaxis","yaxis","xvar","xmin","xmax","xlabel","ylabel","bins","binunc","bbox_to_anchor","tight","residue","lpos","interpolate","params","also_fit","init","frange","epsfcn","units","fselector","maxfev","sortbyx","xerror","yerror","fixed_params","autotqdm","fitter",
1✔
460
                "title","where","save", "prange","ss","also_data", "auto_fit","data_sigmas", "logy", "logx", "data_color", "fit_color", "fit_fmt", "show", "size", "number_format", "selector", "fitinline", "grid", "hist", "stairs", "capsize", "axes",  "xspace",  "extrapolate", "extrapolate_min", "extrapolate_max", "extrapolate_fmt", "extrapolate_hatch",
461
                "function_color", "residue_err", "interpolate_fmt", "interpolate_label", "ncol", "steps","interpolator",
462
                "next_color",
463
                ]:
464
        if key in kwargs:
1✔
465
            del kwargs[key]
1✔
466
    func = gfunc
1✔
467
    x = xlinspace
1✔
468
    l = label
1✔
469

470
    if isinstance(func(x)[0], uncertainties.UFloat):
1✔
471
        if sigmas > 0:
1✔
472
            (ll,) = plt_plt(
1✔
473
                x, unv(func(x)), fmt, label=None, color=color, linestyle=linestyle,**kwargs
474
            )
475
            y = func(x)
1✔
476
            plt.fill_between(
1✔
477
                x,
478
                unv(y) - sigmas * usd(y),
479
                unv(y) + sigmas * usd(y),
480
                alpha=alpha,
481
                label=l,
482
                color=ll.get_color(),
483
                hatch=hatch,
484
                **kwargs
485
            )
486
        else:
487
            (ll,) = plt_plt(
1✔
488
                x, unv(func(x)), fmt, label=l, color=color, linestyle=linestyle,**kwargs
489
            )
490
    else:
491
        (ll,) = plt_plt(x, func(x), fmt, label=l, color=color, linestyle=linestyle,**kwargs)
1✔
492
    return ll
1✔
493

494

495
def function(func, *args, **kwargs):
1✔
496
    """
497
    Plot function ``func`` between ``xmin`` and ``xmax``
498

499
    Parameters
500
    ----------
501
    func : function
502
        Function to be plotted between ``xmin`` and ``xmax``, only taking `array_like` ``x`` as parameter
503
    *args : optional
504
        arguments for ``func``
505
    **kwargs : optional
506
        see :func:`plot_kwargs`.
507
    """
508
    if not util.has("xmin", kwargs) or not util.has("xmax", kwargs):
1✔
509
        kwargs["xmin"], kwargs["xmax"] = stat.get_interesting_domain(func)
1✔
510
        # raise Exception("xmin or xmax missing.")
511

512
    # if not util.has('lpos', kwargs) and not util.has('label', kwargs):
513
    #    kwargs['lpos'] = -1
514
    if not util.has("fmt", kwargs):
1✔
515
        kwargs["fmt"] = "-"
1✔
516

517
    if "label" not in kwargs:
1✔
518
        kwargs = plot_kwargs(kwargs)
1✔
519
        kwargs["label"] = get_fnc_legend(func, args, **kwargs)
1✔
520
    else:
521
        kwargs = plot_kwargs(kwargs)
1✔
522

523
    xlin = kwargs["xspace"](kwargs["xmin"], kwargs["xmax"], kwargs["steps"])
1✔
524
    init_plot(kwargs)
1✔
525

526
    # kwargs['lpos'] = 0
527
    # _plot(xfit, func(xfit, *args), **kwargs)
528
    _function(wrap.get_lambda_argd(func, kwargs["xvar"], *args), xlin, **kwargs)
1✔
529
    if kwargs["ss"]:
1✔
530
        save_plot(**kwargs)
1✔
531

532

533
# xaxis="",yaxis="",fit_color=None,save = None,residue_err=True,show=False):
534
def plt_residue(datax, datay, gfunction, rfit, fig, **kwargs):
1✔
535
    function = wrap.get_lambda(gfunction, kwargs["xvar"])
1✔
536
    fig.add_axes((0.1, 0.1, 0.8, 0.2))
1✔
537
    kwargs["yaxis"] = "$\\Delta$" + kwargs["yaxis"]
1✔
538
    kwargs["data_color"] = kwargs["fit_color"]
1✔
539

540
    if kwargs["residue_err"]:
1✔
541
        plt_data(datax, datay - function(datax, *rfit), **kwargs)
1✔
542
    else:
543
        plt_data(unv(datax), unv(datay - function(datax, *rfit)), **kwargs)
1✔
544
    kwargs["lpos"] = -1
1✔
545
    save_plot(**kwargs)
1✔
546

547

548
def data_split(datax, datay, **kwargs):
1✔
549
    return ffit.data_split(datax, datay, **kwargs)
1✔
550

551

552
def _fit(datax, datay, function, **kwargs):
1✔
553
    """
554
    Returns a fit like :func:`fit` but does no plotting.
555
    """
556
    return ffit.fit(datax, datay, function, **kwargs)
1✔
557

558

559
def plt_data(datax, datay, **kwargs):
1✔
560
    """
561
    Plot datay vs datax
562
    """
563
    x, y, xerr, yerr = data_split(datax, datay, **kwargs)
1✔
564
    if xerr is not None:
1✔
565
        xerr = xerr * kwargs["data_sigmas"]
1✔
566
    if yerr is not None:
1✔
567
        yerr = yerr * kwargs["data_sigmas"]
1✔
568

569
    ll = None
1✔
570
    if xerr is None and yerr is None:
1✔
571
        if kwargs["fmt"] is None:
1✔
572
            if kwargs["linestyle"] is None:
×
573
                (ll,) = plt.plot(
×
574
                    x, y, label=kwargs["label"], color=kwargs["data_color"]
575
                )
576
            else:
577
                (ll,) = plt.plot(
×
578
                    x,
579
                    y,
580
                    label=kwargs["label"],
581
                    color=kwargs["data_color"],
582
                    linestyle=kwargs["linestyle"],
583
                )
584
        elif kwargs["fmt"] == "step":
1✔
585
            (ll,) = plt.step(
×
586
                x,
587
                y,
588
                where=kwargs["where"],
589
                label=kwargs["label"],
590
                color=kwargs["data_color"],
591
            )
592
        elif kwargs["fmt"] == "hist":
1✔
593
            (ll,) = plt.step(
1✔
594
                x,
595
                y,
596
                where=kwargs["where"],
597
                label=kwargs["label"],
598
                color=kwargs["data_color"],
599
            )
600
            plt.fill_between(x, y, step="mid", color=ll.get_color())
1✔
601
        else:
602
            (ll,) = plt.plot(
1✔
603
                x, y, kwargs["fmt"], label=kwargs["label"], color=kwargs["data_color"]
604
            )
605
    else:
606
        if kwargs["fmt"] is None:
1✔
607
            if kwargs["linestyle"] is None:
1✔
608
                (
1✔
609
                    ll,
610
                    _,
611
                    _,
612
                ) = plt.errorbar(
613
                    x,
614
                    y,
615
                    yerr=yerr,
616
                    xerr=xerr,
617
                    fmt=" ",
618
                    capsize=kwargs["capsize"],
619
                    label=kwargs["label"],
620
                    color=kwargs["data_color"],
621
                )
622
            else:
UNCOV
623
                (
×
624
                    ll,
625
                    _,
626
                    _,
627
                ) = plt.errorbar(
628
                    x,
629
                    y,
630
                    yerr=yerr,
631
                    xerr=xerr,
632
                    fmt=" ",
633
                    capsize=kwargs["capsize"],
634
                    label=kwargs["label"],
635
                    color=kwargs["data_color"],
636
                    linestyle=kwargs["linestyle"],
637
                )
638
        elif kwargs["fmt"] == "step":
1✔
639
            (ll,) = plt.step(x, y, where=kwargs["where"], color=kwargs["data_color"])
1✔
640
            if xerr is not None:
1✔
641
                for ix, xv in enumerate(x):
1✔
642
                    dx = xerr[ix]
1✔
643
                    tx = [xv - dx, xv + dx]
1✔
644
                    plt.fill_between(
1✔
645
                        tx,
646
                        y[ix] - yerr[ix],
647
                        y[ix] + yerr[ix],
648
                        label=kwargs["label"] if ix == 1 else None,
649
                        alpha=kwargs["alpha"],
650
                        step="pre",
651
                        color=ll.get_color(),
652
                    )
653
            else:
654
                plt.fill_between(
1✔
655
                    x,
656
                    y - yerr,
657
                    y + yerr,
658
                    label=kwargs["label"],
659
                    alpha=kwargs["alpha"],
660
                    step="mid",
661
                    color=ll.get_color(),
662
                )
663
        elif kwargs["fmt"] == "hist":
1✔
664
            (
1✔
665
                ll,
666
                _,
667
                _,
668
            ) = plt.errorbar(
669
                x,
670
                y,
671
                yerr=yerr,
672
                xerr=xerr,
673
                fmt=" ",
674
                capsize=kwargs["capsize"],
675
                color="black",
676
            )
677
            plt.fill_between(
1✔
678
                x, y, step="mid", label=kwargs["label"], color=ll.get_color()
679
            )
680
        else:
681
            (
1✔
682
                ll,
683
                _,
684
                _,
685
            ) = plt.errorbar(
686
                x,
687
                y,
688
                yerr=yerr,
689
                xerr=xerr,
690
                fmt=kwargs["fmt"],
691
                capsize=kwargs["capsize"],
692
                label=kwargs["label"],
693
                color=kwargs["data_color"],
694
            )
695
    return ll
1✔
696

697

698
def get_fnc_legend(function, rfit, **kwargs):
1✔
699
    l = wrap.get_latex(function)
1✔
700

701
    vnames = wrap.get_varnames(function, kwargs["xvar"])
1✔
702
    for i in range(1, len(vnames)):
1✔
703
        l = l + ("\n" if not kwargs["fitinline"] or i == 1 else " ")
1✔
704
        l = l + "$" + sympy.latex(sympy.symbols(str(vnames[i]))) + "$="
1✔
705
        if kwargs["units"] is not None and usd(rfit[i - 1]) > 0:
1✔
706
            l = l + "("
1✔
707
        if "number_format" in kwargs:
1✔
708
            l = l + kwargs["number_format"].format(rfit[i - 1])
1✔
709
        else:
710
            l = l + "%s" % (rfit[i - 1])
×
711

712
        if kwargs["units"] is not None and usd(rfit[i - 1]) > 0:
1✔
713
            l = l + ")"
1✔
714
        if kwargs["units"] is not None:
1✔
715
            l = l + " " + kwargs["units"][i - 1]
1✔
716
    return l
1✔
717

718

719
def plt_fit_or_interpolate(
1✔
720
    datax, datay, fitted, l=None, c=None, f=None, ls=None, **kwargs
721
):
722
    # just filter these kwargs out, so they dont get passed down and are replaced by above args
723
    # TODO why not pass label=XXX directly to this?
724
    #      -> probably since there are cases where both e.g. color and replace color are needed
725
    for key in ['color', 'label', 'fmt', 'linestyle', 'hatch']:
1✔
726
        if key in kwargs:
1✔
727
            del kwargs[key]
1✔
728
    if kwargs["prange"] is None:
1✔
729
        x, _, _, _ = ffit.fit_split(datax, datay, **kwargs)
1✔
730
        xfit = kwargs["xspace"](np.min(unv(x)), np.max(unv(x)), kwargs["steps"])
1✔
731
    else:
732
        xfit = kwargs["xspace"](
×
733
            kwargs["prange"][0], kwargs["prange"][1], kwargs["steps"]
734
        )
735
    ll = __function(
1✔
736
        fitted,
737
        xfit,
738
        kwargs["fit_fmt"] if f is not None and ls is None else f,
739
        label=l,
740
        color=kwargs["fit_color"] if c is None else c,
741
        linestyle=ls,
742
        **kwargs
743
    )
744

745
    if (
1✔
746
        (kwargs["frange"] is not None or kwargs["fselector"] is not None)
747
        and util.true("extrapolate", kwargs)
748
        or util.has("extrapolate_max", kwargs)
749
        or util.has("extrapolate_min", kwargs)
750
    ):
751
        xxfit = kwargs["xspace"](
1✔
752
            util.get("extrapolate_min", kwargs, np.min(unv(datax))),
753
            util.get("extrapolate_max", kwargs, np.max(unv(datax))),
754
            kwargs["steps"],
755
        )
756
        for pmin, pmax in [
1✔
757
            (np.min(xxfit), np.min(xfit)),
758
            (np.max(xfit), np.max(xxfit)),
759
        ]:
760
            __function(
1✔
761
                fitted,
762
                kwargs["xspace"](pmin, pmax, kwargs["steps"]),
763
                util.get("extrapolate_fmt", kwargs, "--"),
764
                color=ll.get_color(),
765
                hatch=util.get("extrapolate_hatch", kwargs, r"||"),
766
                **kwargs
767
            )
768
    return ll.get_color(), xfit, fitted(xfit)
1✔
769

770

771
def plt_interpolate(datax, datay, icolor=None, **kwargs):
1✔
772
    """
773
    Interpolate and Plot that Interpolation.
774
    """
775
    inter = interpolate.interpolate(datax, datay, **kwargs)
1✔
776
    kargs = {}
1✔
777
    if isinstance(kwargs["interpolate_fmt"], tuple):
1✔
778
        kargs["ls"] = kwargs["interpolate_fmt"]
×
779
    else:
780
        kargs["f"] = kwargs["interpolate_fmt"]
1✔
781
    if kwargs["interpolate_label"] is not None:
1✔
782
        kargs["l"] = kwargs["interpolate_label"]
1✔
783
    # l = None so that no label
784
    return (
1✔
785
        inter,
786
        *plt_fit_or_interpolate(datax, datay, inter, c=icolor, **kargs, **kwargs),
787
    )
788

789
def plt_fit(datax, datay, gfunction, **kwargs):
1✔
790
    """
791
    Fit and Plot that Fit.
792
    """
793
    func = wrap.get_lambda(gfunction, kwargs["xvar"])
1✔
794
    rfit = _fit(datax, datay, gfunction, **kwargs)
1✔
795

796
    def fitted(x):
1✔
797
        return func(x, *rfit)
1✔
798

799
    vnames = wrap.get_varnames(gfunction, kwargs["xvar"])
1✔
800
    for v in vnames[1:]: # remove fixed parameters from kwargs
1✔
801
        if v in kwargs:
1✔
802
            del kwargs[v]
1✔
803

804
    l = get_fnc_legend(gfunction, rfit, **kwargs)
1✔
805
    return (rfit, *plt_fit_or_interpolate(datax, datay, fitted, l, **kwargs))
1✔
806

807

808
def init_plot(kwargs):
1✔
809
    fig = None
1✔
810
    if util.has("axes", kwargs) and kwargs["axes"] is not None:
1✔
811
        plt.sca(kwargs["axes"])
1✔
812
        fig = kwargs["axes"].get_figure()
1✔
813
    if kwargs["init"] or util.true("residue", kwargs):
1✔
814
        if kwargs["size"] is None:
1✔
815
            fig = plt.figure()
1✔
816
        else:
817
            fig = plt.figure(figsize=kwargs["size"])
×
818
        if kwargs["residue"]:
1✔
819
            fig.add_axes((0.1, 0.3, 0.8, 0.6))
1✔
820
    if util.has("xlabel", kwargs) and kwargs["xlabel"] != "":
1✔
821
        plt.xlabel(kwargs["xlabel"])
1✔
822
    if util.has("ylabel", kwargs) and kwargs["ylabel"] != "":
1✔
823
        plt.ylabel(kwargs["ylabel"])
1✔
824
    if util.has("xaxis", kwargs) and kwargs["xaxis"] != "":
1✔
825
        plt.xlabel(kwargs["xaxis"])
1✔
826
    if util.has("yaxis", kwargs) and kwargs["yaxis"] != "":
1✔
827
        plt.ylabel(kwargs["yaxis"])
1✔
828
    if util.has("next_color", kwargs) and not kwargs["next_color"]:
1✔
829
        it1, it2 = itertools.tee(iter(plt.gca()._get_lines.prop_cycler))
1✔
830
        plt.gca()._get_lines.prop_cycler = it2
1✔
831
        tmp_color = next(it1)["color"]
1✔
832
        if kwargs["data_color"] is None:
1✔
833
            kwargs["data_color"] = tmp_color
1✔
834
        if kwargs["fit_color"] is None:
1✔
835
            kwargs["fit_color"] = tmp_color
1✔
836
        if kwargs["function_color"] is None:
1✔
837
            kwargs["function_color"] = tmp_color
1✔
838
    return fig
1✔
839

840

841
def save_plot(**kwargs):
1✔
842
    """
843
    save plot
844
    """
845
    if "title" in kwargs and kwargs["title"] is not None:
1✔
846
        plt.title(kwargs["title"])
×
847
    if "logy" in kwargs and kwargs["logy"]:
1✔
848
        plt.gca().set_yscale("log")
1✔
849
    if "logx" in kwargs and kwargs["logx"]:
1✔
850
        plt.gca().set_xscale("log")
1✔
851
    if "tight" in kwargs and kwargs["tight"]:
1✔
852
        plt.tight_layout()
1✔
853
    if "lpos" in kwargs and kwargs["lpos"] >= 0:
1✔
854
        if util.has("bbox_to_anchor", kwargs):
1✔
855
            if util.has("ncol", kwargs):
1✔
856
                plt.legend(
1✔
857
                    loc=kwargs["lpos"],
858
                    bbox_to_anchor=kwargs["bbox_to_anchor"],
859
                    ncol=kwargs["ncol"],
860
                    borderaxespad=0,
861
                )
862
            else:
863
                plt.legend(loc=kwargs["lpos"], bbox_to_anchor=kwargs["bbox_to_anchor"])
1✔
864
        else:
865
            plt.legend(loc=kwargs["lpos"])
1✔
866
    # plt.gca().set_xlim([kwargs['xmin'],kwargs['xmax']])
867
    # plt.gca().set_ylim([kwargs['ymin'],kwargs['ymax']])
868
    if "save" in kwargs and not kwargs["save"] is None:
1✔
869
        io.mkdirs(kwargs["save"])
×
870
        plt.savefig(kwargs["save"] + ".pdf")
×
871
    plt.grid(kwargs["grid"])
1✔
872
    if "show" in kwargs and kwargs["show"]:
1✔
873
        show(**kwargs)
1✔
874

875

876
def show(**kwargs):
1✔
877
    kwargs = plot_kwargs(kwargs)
1✔
878

879
    plt.grid(kwargs["grid"])
1✔
880
    plt.show()
1✔
881

882

883
from smpl.plot2d import plot2d as _plot2d
1✔
884
from smpl.plot2d import plot2d_kwargs as _plot2d_kwargs
1✔
885

886

887
def plot2d(*args, **kwargs):
1✔
888
    _plot2d(*args, **kwargs)
1✔
889

890

891
def plot2d_kwargs(*args, **kwargs):
1✔
892
    _plot2d_kwargs(*args, **kwargs)
×
893

894

895
if __name__ == "__main__":
1✔
896
    import doctest
×
897

898
    doctest.testmod()
×
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc