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

APN-Pucky / smpl / 15215464499

23 May 2025 05:02PM UTC coverage: 87.5%. Remained the same
15215464499

push

github

APN-Pucky
renenabel codacy

987 of 1128 relevant lines covered (87.5%)

4.37 hits per line

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

92.93
/smpl/plot.py
1
"""Simplified plotting."""
2

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

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

16

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

34

35
unv = unp.nominal_values
5✔
36
usd = unp.std_devs
5✔
37

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

218

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

234

235
# TODO optimize to minimize number of calls to function
236
# @append_doc(default_kwargs)
237

238

239
def fit(func, *adata, **kwargs):
5✔
240
    """
241
    Fit and plot function to datax and datay.
242

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

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

261
    Examples
262
    --------
263

264
    .. plot::
265
        :include-source:
266

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

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

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

305
    kwargs = plot_kwargs(kwargs)
5✔
306

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

324

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

352

353
# @append_doc(default_kwargs)
354

355

356
def data(*data, function=None, **kwargs):
5✔
357
    """
358
    Plot datay against datax via :func:`fit`
359

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

380

381
# @append_doc(default_kwargs)
382

383

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

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

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

399

400

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

407

408
def _function(
5✔
409
    func,
410
    xfit,
411
    fmt="-",
412
    label=None,
413
    function_color=None,
414
    sigmas=0.0,
415
    alpha=0.4,
416
    **kwargs,
417
):
418
    # kargs = {}
419
    # if util.has("fmt", kwargs):
420
    #    kargs["fmt"] = kwargs["fmt"]
421
    # if util.has("label", kwargs) and kwargs["label"] != "":
422
    #    kargs["label"] = kwargs["label"]
423
    # if util.has("function_color", kwargs) and kwargs["function_color"] != "":
424
    #    kargs["color"] = kwargs["function_color"]
425
    # if util.has("sigmas", kwargs) and kwargs["sigmas"] != "":
426
    #    kargs["sigmas"] = kwargs["sigmas"]
427
    # if util.has("alpha", kwargs) and kwargs["alpha"] != "":
428
    #    kargs["alpha"] = kwargs["alpha"]
429
    # __function(func, xfit,  **kargs)
430

431
    if label == "":
5✔
432
        label = None
×
433
    if function_color == "":
5✔
434
        function_color = None
×
435
    if sigmas == "":
5✔
436
        sigmas = 0.0
×
437
    if alpha == "":
5✔
438
        alpha = 0.4
×
439
    __function(
5✔
440
        func,
441
        xfit,
442
        fmt=fmt,
443
        label=label,
444
        color=function_color,
445
        sigmas=sigmas,
446
        alpha=alpha,
447
        **kwargs,
448
    )
449

450

451
def plt_plt(x, y, fmt, color, label, linestyle, **kwargs):
5✔
452
    if linestyle is None and fmt is not None:
5✔
453
        return plt.plot(x, y, fmt, label=label, color=color, **kwargs)
5✔
454
    if linestyle is not None and fmt is None:
5✔
455
        return plt.plot(x, y, label=label, color=color, linestyle=linestyle, **kwargs)
×
456
    if linestyle is None and fmt is None:
5✔
457
        return plt.plot(x, y, label=label, color=color, **kwargs)
5✔
458

459

460
def __function(
5✔
461
    gfunc,
462
    xlinspace,
463
    fmt="-",
464
    label=None,
465
    color=None,
466
    hatch=None,
467
    sigmas=0.0,
468
    linestyle=None,
469
    alpha=0.4,
470
    **kwargs,
471
):
472
    # filter unused bad kwargs here to avoid passing them down
473
    # TODO it would be better to not pass them down in the first place
474
    for key in [
5✔
475
        "xaxis",
476
        "yaxis",
477
        "xvar",
478
        "xmin",
479
        "xmax",
480
        "xlabel",
481
        "ylabel",
482
        "bins",
483
        "binunc",
484
        "bbox_to_anchor",
485
        "tight",
486
        "residue",
487
        "lpos",
488
        "interpolate",
489
        "params",
490
        "also_fit",
491
        "init",
492
        "frange",
493
        "epsfcn",
494
        "units",
495
        "fselector",
496
        "maxfev",
497
        "sortbyx",
498
        "xerror",
499
        "yerror",
500
        "fixed_params",
501
        "autotqdm",
502
        "fitter",
503
        "title",
504
        "where",
505
        "save",
506
        "prange",
507
        "ss",
508
        "also_data",
509
        "auto_fit",
510
        "data_sigmas",
511
        "logy",
512
        "logx",
513
        "data_color",
514
        "fit_color",
515
        "fit_fmt",
516
        "show",
517
        "size",
518
        "number_format",
519
        "selector",
520
        "fitinline",
521
        "grid",
522
        "hist",
523
        "stairs",
524
        "capsize",
525
        "axes",
526
        "xspace",
527
        "extrapolate",
528
        "extrapolate_min",
529
        "extrapolate_max",
530
        "extrapolate_fmt",
531
        "extrapolate_hatch",
532
        "function_color",
533
        "residue_err",
534
        "interpolate_fmt",
535
        "interpolate_label",
536
        "ncol",
537
        "steps",
538
        "interpolator",
539
        "next_color",
540
    ]:
541
        kwargs.pop(key, None)
5✔
542
    func = gfunc
5✔
543
    x = xlinspace
5✔
544
    l = label
5✔
545

546
    if isinstance(func(x)[0], uncertainties.UFloat):
5✔
547
        if sigmas > 0:
5✔
548
            (ll,) = plt_plt(
5✔
549
                x,
550
                unv(func(x)),
551
                fmt,
552
                label=None,
553
                color=color,
554
                linestyle=linestyle,
555
                **kwargs,
556
            )
557
            y = func(x)
5✔
558
            plt.fill_between(
5✔
559
                x,
560
                unv(y) - sigmas * usd(y),
561
                unv(y) + sigmas * usd(y),
562
                alpha=alpha,
563
                label=l,
564
                color=ll.get_color(),
565
                hatch=hatch,
566
                **kwargs,
567
            )
568
        else:
569
            (ll,) = plt_plt(
5✔
570
                x,
571
                unv(func(x)),
572
                fmt,
573
                label=l,
574
                color=color,
575
                linestyle=linestyle,
576
                **kwargs,
577
            )
578
    else:
579
        (ll,) = plt_plt(
5✔
580
            x, func(x), fmt, label=l, color=color, linestyle=linestyle, **kwargs
581
        )
582
    return ll
5✔
583

584

585
def function(func, *args, **kwargs):
5✔
586
    """
587
    Plot function ``func`` between ``xmin`` and ``xmax``
588

589
    Parameters
590
    ----------
591
    func : function
592
        Function to be plotted between ``xmin`` and ``xmax``, only taking `array_like` ``x`` as parameter
593
    *args : optional
594
        arguments for ``func``
595
    **kwargs : optional
596
        see :func:`plot_kwargs`.
597
    """
598
    if not util.has("xmin", kwargs) or not util.has("xmax", kwargs):
5✔
599
        kwargs["xmin"], kwargs["xmax"] = stat.get_interesting_domain(func)
5✔
600
        # raise Exception("xmin or xmax missing.")
601

602
    # if not util.has('lpos', kwargs) and not util.has('label', kwargs):
603
    #    kwargs['lpos'] = -1
604
    if not util.has("fmt", kwargs):
5✔
605
        kwargs["fmt"] = "-"
5✔
606

607
    if "label" not in kwargs:
5✔
608
        kwargs = plot_kwargs(kwargs)
5✔
609
        kwargs["label"] = get_fnc_legend(func, args, **kwargs)
5✔
610
    else:
611
        kwargs = plot_kwargs(kwargs)
5✔
612

613
    xlin = kwargs["xspace"](kwargs["xmin"], kwargs["xmax"], kwargs["steps"])
5✔
614
    init_plot(kwargs)
5✔
615

616
    # kwargs['lpos'] = 0
617
    # _plot(xfit, func(xfit, *args), **kwargs)
618
    _function(wrap.get_lambda_argd(func, kwargs["xvar"], *args), xlin, **kwargs)
5✔
619
    if kwargs["ss"]:
5✔
620
        save_plot(**kwargs)
5✔
621

622

623
# xaxis="",yaxis="",fit_color=None,save = None,residue_err=True,show=False):
624
def plt_residue(datax, datay, gfunction, rfit, fig, **kwargs):
5✔
625
    function = wrap.get_lambda(gfunction, kwargs["xvar"])
5✔
626
    fig.add_axes((0.1, 0.1, 0.8, 0.2))
5✔
627
    kwargs["yaxis"] = "$\\Delta$" + kwargs["yaxis"]
5✔
628
    kwargs["data_color"] = kwargs["fit_color"]
5✔
629

630
    if kwargs["residue_err"]:
5✔
631
        plt_data(datax, datay - function(datax, *rfit), **kwargs)
5✔
632
    else:
633
        plt_data(unv(datax), unv(datay - function(datax, *rfit)), **kwargs)
5✔
634
    kwargs["lpos"] = -1
5✔
635
    save_plot(**kwargs)
5✔
636

637

638
def data_split(datax, datay, **kwargs):
5✔
639
    return ffit.data_split(datax, datay, **kwargs)
5✔
640

641

642
def _fit(datax, datay, function, **kwargs):
5✔
643
    """
644
    Returns a fit like :func:`fit` but does no plotting.
645
    """
646
    return ffit.fit(datax, datay, function, **kwargs)
5✔
647

648

649
def plt_data(datax, datay, **kwargs):
5✔
650
    """
651
    Plot datay vs datax
652
    """
653
    x, y, xerr, yerr = data_split(datax, datay, **kwargs)
5✔
654
    if xerr is not None:
5✔
655
        xerr = xerr * kwargs["data_sigmas"]
5✔
656
    if yerr is not None:
5✔
657
        yerr = yerr * kwargs["data_sigmas"]
5✔
658

659
    ll = None
5✔
660
    if xerr is None and yerr is None:
5✔
661
        if kwargs["fmt"] is None:
5✔
662
            if kwargs["linestyle"] is None:
×
663
                (ll,) = plt.plot(
×
664
                    x, y, label=kwargs["label"], color=kwargs["data_color"]
665
                )
666
            else:
667
                (ll,) = plt.plot(
×
668
                    x,
669
                    y,
670
                    label=kwargs["label"],
671
                    color=kwargs["data_color"],
672
                    linestyle=kwargs["linestyle"],
673
                )
674
        elif kwargs["fmt"] == "step":
5✔
675
            (ll,) = plt.step(
×
676
                x,
677
                y,
678
                where=kwargs["where"],
679
                label=kwargs["label"],
680
                color=kwargs["data_color"],
681
            )
682
        elif kwargs["fmt"] == "hist":
5✔
683
            (ll,) = plt.step(
5✔
684
                x,
685
                y,
686
                where=kwargs["where"],
687
                label=kwargs["label"],
688
                color=kwargs["data_color"],
689
            )
690
            plt.fill_between(x, y, step="mid", color=ll.get_color())
5✔
691
        else:
692
            (ll,) = plt.plot(
5✔
693
                x, y, kwargs["fmt"], label=kwargs["label"], color=kwargs["data_color"]
694
            )
695
    elif kwargs["fmt"] is None:
5✔
696
        if kwargs["linestyle"] is None:
5✔
697
            (
5✔
698
                ll,
699
                _,
700
                _,
701
            ) = plt.errorbar(
702
                x,
703
                y,
704
                yerr=yerr,
705
                xerr=xerr,
706
                fmt=" ",
707
                capsize=kwargs["capsize"],
708
                label=kwargs["label"],
709
                color=kwargs["data_color"],
710
            )
711
        else:
712
            (
1✔
713
                ll,
714
                _,
715
                _,
716
            ) = plt.errorbar(
717
                x,
718
                y,
719
                yerr=yerr,
720
                xerr=xerr,
721
                fmt=" ",
722
                capsize=kwargs["capsize"],
723
                label=kwargs["label"],
724
                color=kwargs["data_color"],
725
                linestyle=kwargs["linestyle"],
726
            )
727
    elif kwargs["fmt"] == "step":
5✔
728
        (ll,) = plt.step(x, y, where=kwargs["where"], color=kwargs["data_color"])
5✔
729
        if xerr is not None:
5✔
730
            for ix, xv in enumerate(x):
5✔
731
                dx = xerr[ix]
5✔
732
                tx = [xv - dx, xv + dx]
5✔
733
                plt.fill_between(
5✔
734
                    tx,
735
                    y[ix] - yerr[ix],
736
                    y[ix] + yerr[ix],
737
                    label=kwargs["label"] if ix == 1 else None,
738
                    alpha=kwargs["alpha"],
739
                    step="pre",
740
                    color=ll.get_color(),
741
                )
742
        else:
743
            plt.fill_between(
5✔
744
                x,
745
                y - yerr,
746
                y + yerr,
747
                label=kwargs["label"],
748
                alpha=kwargs["alpha"],
749
                step="mid",
750
                color=ll.get_color(),
751
            )
752
    elif kwargs["fmt"] == "hist":
5✔
753
        (
5✔
754
            ll,
755
            _,
756
            _,
757
        ) = plt.errorbar(
758
            x,
759
            y,
760
            yerr=yerr,
761
            xerr=xerr,
762
            fmt=" ",
763
            capsize=kwargs["capsize"],
764
            color="black",
765
        )
766
        plt.fill_between(x, y, step="mid", label=kwargs["label"], color=ll.get_color())
5✔
767
    else:
768
        (
5✔
769
            ll,
770
            _,
771
            _,
772
        ) = plt.errorbar(
773
            x,
774
            y,
775
            yerr=yerr,
776
            xerr=xerr,
777
            fmt=kwargs["fmt"],
778
            capsize=kwargs["capsize"],
779
            label=kwargs["label"],
780
            color=kwargs["data_color"],
781
        )
782
    return ll
5✔
783

784

785
def get_fnc_legend(function, rfit, **kwargs):
5✔
786
    l = wrap.get_latex(function)
5✔
787

788
    vnames = wrap.get_varnames(function, kwargs["xvar"])
5✔
789
    for i in range(1, len(vnames)):
5✔
790
        l = l + ("\n" if not kwargs["fitinline"] or i == 1 else " ")
5✔
791
        l = l + "$" + sympy.latex(sympy.symbols(str(vnames[i]))) + "$="
5✔
792
        if kwargs["units"] is not None and usd(rfit[i - 1]) > 0:
5✔
793
            l = l + "("
5✔
794
        if "number_format" in kwargs:
5✔
795
            l = l + kwargs["number_format"].format(rfit[i - 1])
5✔
796
        else:
797
            l = l + "%s" % (rfit[i - 1])
×
798

799
        if kwargs["units"] is not None and usd(rfit[i - 1]) > 0:
5✔
800
            l = l + ")"
5✔
801
        if kwargs["units"] is not None:
5✔
802
            l = l + " " + kwargs["units"][i - 1]
5✔
803
    return l
5✔
804

805

806
def plt_fit_or_interpolate(
5✔
807
    datax, datay, fitted, l=None, c=None, f=None, ls=None, **kwargs
808
):
809
    # just filter these kwargs out, so they dont get passed down and are replaced by above args
810
    # TODO why not pass label=XXX directly to this?
811
    #      -> probably since there are cases where both e.g. color and replace color are needed
812
    for key in ["color", "label", "fmt", "linestyle", "hatch"]:
5✔
813
        kwargs.pop(key, None)
5✔
814
    if kwargs["prange"] is None:
5✔
815
        x, _, _, _ = ffit.fit_split(datax, datay, **kwargs)
5✔
816
        xfit = kwargs["xspace"](np.min(unv(x)), np.max(unv(x)), kwargs["steps"])
5✔
817
    else:
818
        xfit = kwargs["xspace"](
×
819
            kwargs["prange"][0], kwargs["prange"][1], kwargs["steps"]
820
        )
821
    ll = __function(
5✔
822
        fitted,
823
        xfit,
824
        kwargs["fit_fmt"] if f is not None and ls is None else f,
825
        label=l,
826
        color=kwargs["fit_color"] if c is None else c,
827
        linestyle=ls,
828
        **kwargs,
829
    )
830

831
    if (
5✔
832
        (
833
            (kwargs["frange"] is not None or kwargs["fselector"] is not None)
834
            and util.true("extrapolate", kwargs)
835
        )
836
        or util.has("extrapolate_max", kwargs)
837
        or util.has("extrapolate_min", kwargs)
838
    ):
839
        xxfit = kwargs["xspace"](
5✔
840
            util.get("extrapolate_min", kwargs, np.min(unv(datax))),
841
            util.get("extrapolate_max", kwargs, np.max(unv(datax))),
842
            kwargs["steps"],
843
        )
844
        for pmin, pmax in [
5✔
845
            (np.min(xxfit), np.min(xfit)),
846
            (np.max(xfit), np.max(xxfit)),
847
        ]:
848
            __function(
5✔
849
                fitted,
850
                kwargs["xspace"](pmin, pmax, kwargs["steps"]),
851
                util.get("extrapolate_fmt", kwargs, "--"),
852
                color=ll.get_color(),
853
                hatch=util.get("extrapolate_hatch", kwargs, r"||"),
854
                **kwargs,
855
            )
856
    return ll.get_color(), xfit, fitted(xfit)
5✔
857

858

859
def plt_interpolate(datax, datay, icolor=None, **kwargs):
5✔
860
    """
861
    Interpolate and Plot that Interpolation.
862
    """
863
    inter = interpolate.interpolate(datax, datay, **kwargs)
5✔
864
    kargs = {}
5✔
865
    if isinstance(kwargs["interpolate_fmt"], tuple):
5✔
866
        kargs["ls"] = kwargs["interpolate_fmt"]
×
867
    else:
868
        kargs["f"] = kwargs["interpolate_fmt"]
5✔
869
    if kwargs["interpolate_label"] is not None:
5✔
870
        kargs["l"] = kwargs["interpolate_label"]
5✔
871
    # l = None so that no label
872
    return (
5✔
873
        inter,
874
        *plt_fit_or_interpolate(datax, datay, inter, c=icolor, **kargs, **kwargs),
875
    )
876

877

878
def plt_fit(datax, datay, gfunction, **kwargs):
5✔
879
    """
880
    Fit and Plot that Fit.
881
    """
882
    func = wrap.get_lambda(gfunction, kwargs["xvar"])
5✔
883
    rfit = _fit(datax, datay, gfunction, **kwargs)
5✔
884

885
    def fitted(x):
5✔
886
        return func(x, *rfit)
5✔
887

888
    vnames = wrap.get_varnames(gfunction, kwargs["xvar"])
5✔
889
    for v in vnames[1:]:  # remove fixed parameters from kwargs
5✔
890
        kwargs.pop(v, None)
5✔
891

892
    l = get_fnc_legend(gfunction, rfit, **kwargs)
5✔
893
    return (rfit, *plt_fit_or_interpolate(datax, datay, fitted, l, **kwargs))
5✔
894

895

896
def init_plot(kwargs):
5✔
897
    fig = None
5✔
898
    if util.has("axes", kwargs) and kwargs["axes"] is not None:
5✔
899
        plt.sca(kwargs["axes"])
5✔
900
        fig = kwargs["axes"].get_figure()
5✔
901
    if kwargs["init"] or util.true("residue", kwargs):
5✔
902
        if kwargs["size"] is None:
5✔
903
            fig = plt.figure()
5✔
904
        else:
905
            fig = plt.figure(figsize=kwargs["size"])
×
906
        if kwargs["residue"]:
5✔
907
            fig.add_axes((0.1, 0.3, 0.8, 0.6))
5✔
908
    if util.has("xlabel", kwargs) and kwargs["xlabel"] != "":
5✔
909
        plt.xlabel(kwargs["xlabel"])
5✔
910
    if util.has("ylabel", kwargs) and kwargs["ylabel"] != "":
5✔
911
        plt.ylabel(kwargs["ylabel"])
5✔
912
    if util.has("xaxis", kwargs) and kwargs["xaxis"] != "":
5✔
913
        plt.xlabel(kwargs["xaxis"])
5✔
914
    if util.has("yaxis", kwargs) and kwargs["yaxis"] != "":
5✔
915
        plt.ylabel(kwargs["yaxis"])
5✔
916
    if util.has("next_color", kwargs) and not kwargs["next_color"]:
5✔
917
        lines = plt.gca()._get_lines
5✔
918
        tmp_color = lines._cycler_items[lines._idx]["color"]
5✔
919

920
        if kwargs["data_color"] is None:
5✔
921
            kwargs["data_color"] = tmp_color
5✔
922
        if kwargs["fit_color"] is None:
5✔
923
            kwargs["fit_color"] = tmp_color
5✔
924
        if kwargs["function_color"] is None:
5✔
925
            kwargs["function_color"] = tmp_color
5✔
926
    return fig
5✔
927

928

929
def save_plot(**kwargs):
5✔
930
    """
931
    save plot
932
    """
933
    if "title" in kwargs and kwargs["title"] is not None:
5✔
934
        plt.title(kwargs["title"])
×
935
    if kwargs.get("logy"):
5✔
936
        plt.gca().set_yscale("log")
5✔
937
    if kwargs.get("logx"):
5✔
938
        plt.gca().set_xscale("log")
5✔
939
    if kwargs.get("tight"):
5✔
940
        plt.tight_layout()
5✔
941
    if "lpos" in kwargs and kwargs["lpos"] >= 0:
5✔
942
        if util.has("bbox_to_anchor", kwargs):
5✔
943
            if util.has("ncol", kwargs):
5✔
944
                plt.legend(
5✔
945
                    loc=kwargs["lpos"],
946
                    bbox_to_anchor=kwargs["bbox_to_anchor"],
947
                    ncol=kwargs["ncol"],
948
                    borderaxespad=0,
949
                )
950
            else:
951
                plt.legend(loc=kwargs["lpos"], bbox_to_anchor=kwargs["bbox_to_anchor"])
5✔
952
        else:
953
            plt.legend(loc=kwargs["lpos"])
5✔
954
    # plt.gca().set_xlim([kwargs['xmin'],kwargs['xmax']])
955
    # plt.gca().set_ylim([kwargs['ymin'],kwargs['ymax']])
956
    if "save" in kwargs and kwargs["save"] is not None:
5✔
957
        io.mkdirs(kwargs["save"])
×
958
        plt.savefig(kwargs["save"] + ".pdf")
×
959
    plt.grid(kwargs["grid"])
5✔
960
    if kwargs.get("show"):
5✔
961
        show(**kwargs)
5✔
962

963

964
def show(**kwargs):
5✔
965
    kwargs = plot_kwargs(kwargs)
5✔
966

967
    plt.grid(kwargs["grid"])
5✔
968
    plt.show()
5✔
969

970

971
from smpl.plot2d import plot2d as _plot2d
5✔
972
from smpl.plot2d import plot2d_kwargs as _plot2d_kwargs
5✔
973

974

975
def plot2d(*args, **kwargs):
5✔
976
    _plot2d(*args, **kwargs)
5✔
977

978

979
def plot2d_kwargs(*args, **kwargs):
5✔
980
    _plot2d_kwargs(*args, **kwargs)
×
981

982

983
if __name__ == "__main__":
5✔
984
    import doctest
×
985

986
    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