15236750782

Committed 25 May 2025 10:07AM UTC coverage: 87.186% (-0.3%) from 87.5%

Build # 15236750782

Build Type

push

github

Committed by

web-flow

Commit Message

Update README.md

Run Details

973 of 1116 relevant lines covered (87.19%)

4.36 hits per line

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

60.5

/smpl/plot2d.py

import warnings

import matplotlib.pyplot as plt
import numpy as np
import smplr
from matplotlib import colors
from matplotlib.image import NonUniformImage

from smpl import doc, util
from smpl import plot as splot

default = {
    "title": [None, "Plot title"],
    "xlabel": [None, "."],
    "ylabel": [None, "."],
    "zlabel": [None, "."],
    "logz": [True, "Colorbar in logarithmic scale."],
    "style": [
        "pcolormesh",
        "Plot via an image ('image') or scatter ('scatter') or mesh ('pcolormesh').",
    ],
    "interpolation": [
        "nearest",
        "Only 'nearest' or 'bilinear' for nonuniformimage. Check https://matplotlib.org/stable/gallery/images_contours_and_fields/interpolation_methods.html#interpolations-for-imshow",
    ],
    "cmap": [
        "viridis",
        "Good default color map for missing datapoints since it does not include white.",
    ],
    # 'zscale' : [None,"Rescale z values."],
}

# @doc.insert_str("\tDefault kwargs\n\n\t")


@doc.append_str(doc.array_table(default, init=False))
@doc.append_str(
    doc.array_table({"plot2d_kwargs": ["default", "description"]}, bottom=False)
)
def plot2d_kwargs(kwargs):
    """Set default plot2d_kwargs if not set."""
    for k, v in default.items():
        if k not in kwargs:
            kwargs[k] = v[0]
    return kwargs


def plot2d(datax, datay, dataz, **kwargs):
    """
    Creates a 2D-Plot.

    Parameters
    ----------
    **kwargs : optional
        see :func:`plot2d_kwargs`.
    """
    kwargs = plot2d_kwargs(kwargs)
    if "xaxis" in kwargs and not kwargs["xlabel"]:
        warnings.warn("xaxis is deprecated. Use xlabel instead.", DeprecationWarning, 2)
        kwargs["xlabel"] = kwargs["xaxis"]
    if "yaxis" in kwargs and not kwargs["ylabel"]:
        warnings.warn("yaxis is deprecated. Use ylabel instead.", DeprecationWarning, 2)
        kwargs["ylabel"] = kwargs["yaxis"]
    if "zaxis" in kwargs and not kwargs["zlabel"]:
        warnings.warn("zaxis is deprecated. Use zlabel instead.", DeprecationWarning, 2)
        kwargs["zlabel"] = kwargs["zaxis"]

    if util.has("axes", kwargs) and kwargs["axes"] is not None:
        plt.sca(kwargs["axes"])
    if kwargs["style"] == "pcolormesh":
        pcolormesh_vplot(datax, datay, dataz, **kwargs)
    elif kwargs["style"] == "image":
        map_vplot(datax, datay, dataz, **kwargs)
    elif kwargs["style"] == "scatter":
        scatter_vplot(datax, datay, dataz, **kwargs)


def sort_xyz(x, y, z):
    p1 = x.argsort(kind="stable")
    x = np.copy(x[p1])
    y = np.copy(y[p1])
    z = np.copy(z[p1])
    p2 = y.argsort(kind="stable")
    x = x[p2]
    y = y[p2]
    z = z[p2]
    return x, y, z


def pcolormesh_vplot(
    tvx,
    tvy,
    tvz,
    xlabel=None,
    ylabel=None,
    zlabel=None,
    logz=True,
    zscale=1.0,
    **kwargs,
):
    """
    Advantage over matplotlibs pcolor(mesh) is that does not require a meshgrid. Instead it uses the data points directly in three lists.
    """
    vx = np.copy(tvx)
    vy = np.copy(tvy)
    vz = np.copy(tvz)
    assert vx.shape == vy.shape == vz.shape

    if len(vz.shape) < 2:
        mesh = np.meshgrid(np.unique(vx), np.unique(vy))
        X, Y = mesh
        # set Z to values of vz on the meshgrid
        Z = np.empty(mesh[0].shape)
        Z[:] = np.nan
        for i, _ in enumerate(vx):
            Z[(mesh[0] == vx[i]) & (mesh[1] == vy[i])] = splot.unv(vz[i])
        Z[:] *= zscale
    else:
        X = vx
        Y = vy
        Z = vz * zscale

    plt.pcolormesh(
        X, Y, Z, norm=colors.LogNorm() if logz else None, cmap=kwargs["cmap"]
    )

    # ax.set_xlim(xl, xm)
    # ax.set_ylim(yl, ym)

    cb = plt.colorbar()
    cb.set_label(zlabel)
    smplr.style_plot2d(xlabel=xlabel, ylabel=ylabel, **kwargs)


def map_vplot(
    tvx,
    tvy,
    tvz,
    xlabel=None,
    ylabel=None,
    zlabel=None,
    logz=True,
    sort=True,
    fill_missing=True,
    zscale=1.0,
    **kwargs,
):
    """ """
    vx = np.copy(tvx)
    vy = np.copy(tvy)
    vz = np.copy(tvz)
    if fill_missing:
        # TODO speed up
        for x in vx:
            for y in vy:
                ex = np.any(np.logical_and((vx == x), (vy == y)))
                if not ex:
                    vx = np.append(vx, x)
                    vy = np.append(vy, y)
                    vz = np.append(vz, 0)
    if sort:
        vx, vy, vz = sort_xyz(vx, vy, vz)

    s = 1
    while vy[s] == vy[s - 1]:
        s = s + 1
    if s == 1:
        # print("flipped x y ")
        while vx[s] == vx[s - 1]:
            s = s + 1
        if s == 1:
            print("error too small map")
            return
        # x, y = y, x
        xlabel, ylabel = ylabel, xlabel
        vx, vy = vy, vx

    grid = splot.unv(vz).reshape((int(np.rint(np.size(vx) / s)), s)) * zscale

    _, ax = plt.subplots(nrows=1, ncols=1, constrained_layout=True)
    im = None
    xl = vx.min() + (vx.min() / 2) - vx[vx != vx.min()].min() / 2
    xm = vx.max() + (vx.max() / 2) - vx[vx != vx.max()].max() / 2
    yl = vy.min() + (vy.min() / 2) - vy[vy != vy.min()].min() / 2
    ym = vy.max() + (vy.max() / 2) - vy[vy != vy.max()].max() / 2
    im = NonUniformImage(
        ax,
        origin="lower",
        cmap=kwargs["cmap"],
        interpolation=kwargs["interpolation"],
        extent=(xl, xm, yl, ym),
        norm=colors.LogNorm() if logz else None,
    )

    im.set_data(np.unique(vx), np.unique(vy), grid)
    ax.add_image(im)
    # ax.images.append(im)
    ax.set_xlim(xl, xm)
    ax.set_ylim(yl, ym)

    cb = plt.colorbar(im)
    cb.set_label(zlabel)
    smplr.style_plot2d(xlabel=xlabel, ylabel=ylabel, **kwargs)


def scatter_vplot(
    vx,
    vy,
    vz,
    xlabel=None,
    ylabel=None,
    zlabel=None,
    logz=True,
    sort=True,
    **kwargs,
):
    if sort:
        vx, vy, vz = sort_xyz(vx, vy, vz)

    _, ax = plt.subplots(nrows=1, ncols=1, constrained_layout=True)
    xl = vx.min() + (vx.min() / 2) - vx[vx != vx.min()].min() / 2
    xm = vx.max() + (vx.max() / 2) - vx[vx != vx.max()].max() / 2
    yl = vy.min() + (vy.min() / 2) - vy[vy != vy.min()].min() / 2
    ym = vy.max() + (vy.max() / 2) - vy[vy != vy.max()].max() / 2

    s = plt.scatter(
        np.concatenate((vx, vx, vx)),
        np.concatenate((vy, vy, vy)),
        c=np.concatenate(
            (
                splot.unv(vz) + splot.usd(vz),
                splot.unv(vz) - splot.usd(vz),
                splot.unv(vz),
            )
        ),
        s=np.concatenate(
            (
                [(3 * plt.rcParams["lines.markersize"]) ** 2 for i in range(len(vx))],
                [(2 * plt.rcParams["lines.markersize"]) ** 2 for i in range(len(vx))],
                [(plt.rcParams["lines.markersize"]) ** 2 for i in range(len(vx))],
            )
        ),
        norm=colors.LogNorm() if logz else None,
        cmap=kwargs["cmap"],
    )

    ax.set_xlim(xl, xm)
    ax.set_ylim(yl, ym)

    cb = plt.colorbar(s)
    cb.set_label(zlabel)
    smplr.style_plot2d(xlabel=xlabel, ylabel=ylabel, **kwargs)

1	import warnings	5✔
2
3	import matplotlib.pyplot as plt	5✔
4	import numpy as np	5✔
5	import smplr	5✔
6	from matplotlib import colors	5✔
7	from matplotlib.image import NonUniformImage	5✔
8
9	from smpl import doc, util	5✔
10	from smpl import plot as splot	5✔
11
12	default = {	5✔
13	"title": [None, "Plot title"],
14	"xlabel": [None, "."],
15	"ylabel": [None, "."],
16	"zlabel": [None, "."],
17	"logz": [True, "Colorbar in logarithmic scale."],
18	"style": [
19	"pcolormesh",
20	"Plot via an image ('image') or scatter ('scatter') or mesh ('pcolormesh').",
21	],
22	"interpolation": [
23	"nearest",
24	"Only 'nearest' or 'bilinear' for nonuniformimage. Check https://matplotlib.org/stable/gallery/images_contours_and_fields/interpolation_methods.html#interpolations-for-imshow",
25	],
26	"cmap": [
27	"viridis",
28	"Good default color map for missing datapoints since it does not include white.",
29	],
30	# 'zscale' : [None,"Rescale z values."],
31	}
32
33	# @doc.insert_str("\tDefault kwargs\n\n\t")
34
35
36	@doc.append_str(doc.array_table(default, init=False))	5✔
37	@doc.append_str(	5✔
38	doc.array_table({"plot2d_kwargs": ["default", "description"]}, bottom=False)
39	)
40	def plot2d_kwargs(kwargs):	5✔
41	"""Set default plot2d_kwargs if not set."""
42	for k, v in default.items():	5✔
43	if k not in kwargs:	5✔
44	kwargs[k] = v[0]	5✔
45	return kwargs	5✔
46
47
48	def plot2d(datax, datay, dataz, **kwargs):	5✔
49	"""
50	Creates a 2D-Plot.
51
52	Parameters
53	----------
54	**kwargs : optional
55	see :func:`plot2d_kwargs`.
56	"""
57	kwargs = plot2d_kwargs(kwargs)	5✔
58	if "xaxis" in kwargs and not kwargs["xlabel"]:	5✔
59	warnings.warn("xaxis is deprecated. Use xlabel instead.", DeprecationWarning, 2)	×
60	kwargs["xlabel"] = kwargs["xaxis"]	×
61	if "yaxis" in kwargs and not kwargs["ylabel"]:	5✔
62	warnings.warn("yaxis is deprecated. Use ylabel instead.", DeprecationWarning, 2)	×
63	kwargs["ylabel"] = kwargs["yaxis"]	×
64	if "zaxis" in kwargs and not kwargs["zlabel"]:	5✔
65	warnings.warn("zaxis is deprecated. Use zlabel instead.", DeprecationWarning, 2)	×
66	kwargs["zlabel"] = kwargs["zaxis"]	×
67
68	if util.has("axes", kwargs) and kwargs["axes"] is not None:	5✔
69	plt.sca(kwargs["axes"])	×
70	if kwargs["style"] == "pcolormesh":	5✔
71	pcolormesh_vplot(datax, datay, dataz, **kwargs)	5✔
72	elif kwargs["style"] == "image":	5✔
73	map_vplot(datax, datay, dataz, **kwargs)	×
74	elif kwargs["style"] == "scatter":	5✔
75	scatter_vplot(datax, datay, dataz, **kwargs)	5✔
76
77
78	def sort_xyz(x, y, z):	5✔
79	p1 = x.argsort(kind="stable")	5✔
80	x = np.copy(x[p1])	5✔
81	y = np.copy(y[p1])	5✔
82	z = np.copy(z[p1])	5✔
83	p2 = y.argsort(kind="stable")	5✔
84	x = x[p2]	5✔
85	y = y[p2]	5✔
86	z = z[p2]	5✔
87	return x, y, z	5✔
88
89
90	def pcolormesh_vplot(	5✔
91	tvx,
92	tvy,
93	tvz,
94	xlabel=None,
95	ylabel=None,
96	zlabel=None,
97	logz=True,
98	zscale=1.0,
99	**kwargs,
100	):
101	"""
102	Advantage over matplotlibs pcolor(mesh) is that does not require a meshgrid. Instead it uses the data points directly in three lists.
103	"""
104	vx = np.copy(tvx)	5✔
105	vy = np.copy(tvy)	5✔
106	vz = np.copy(tvz)	5✔
107	assert vx.shape == vy.shape == vz.shape	5✔
108
109	if len(vz.shape) < 2:	5✔
110	mesh = np.meshgrid(np.unique(vx), np.unique(vy))	5✔
111	X, Y = mesh	5✔
112	# set Z to values of vz on the meshgrid
113	Z = np.empty(mesh[0].shape)	5✔
114	Z[:] = np.nan	5✔
115	for i, _ in enumerate(vx):	5✔
116	Z[(mesh[0] == vx[i]) & (mesh[1] == vy[i])] = splot.unv(vz[i])	5✔
117	Z[:] *= zscale	5✔
118	else:
119	X = vx	5✔
120	Y = vy	5✔
121	Z = vz * zscale	5✔
122
123	plt.pcolormesh(	5✔
124	X, Y, Z, norm=colors.LogNorm() if logz else None, cmap=kwargs["cmap"]
125	)
126
127	# ax.set_xlim(xl, xm)
128	# ax.set_ylim(yl, ym)
129
130	cb = plt.colorbar()	5✔
131	cb.set_label(zlabel)	5✔
132	smplr.style_plot2d(xlabel=xlabel, ylabel=ylabel, **kwargs)	5✔
133
134
135	def map_vplot(	5✔
136	tvx,
137	tvy,
138	tvz,
139	xlabel=None,
140	ylabel=None,
141	zlabel=None,
142	logz=True,
143	sort=True,
144	fill_missing=True,
145	zscale=1.0,
146	**kwargs,
147	):
148	""" """
149	vx = np.copy(tvx)	×
150	vy = np.copy(tvy)	×
151	vz = np.copy(tvz)	×
152	if fill_missing:	×
153	# TODO speed up
154	for x in vx:	×
155	for y in vy:	×
156	ex = np.any(np.logical_and((vx == x), (vy == y)))	×
157	if not ex:	×
158	vx = np.append(vx, x)	×
159	vy = np.append(vy, y)	×
160	vz = np.append(vz, 0)	×
161	if sort:	×
162	vx, vy, vz = sort_xyz(vx, vy, vz)	×
163
164	s = 1	×
165	while vy[s] == vy[s - 1]:	×
166	s = s + 1	×
167	if s == 1:	×
168	# print("flipped x y ")
169	while vx[s] == vx[s - 1]:	×
170	s = s + 1	×
171	if s == 1:	×
172	print("error too small map")	×
173	return	×
174	# x, y = y, x
175	xlabel, ylabel = ylabel, xlabel	×
176	vx, vy = vy, vx	×
177
178	grid = splot.unv(vz).reshape((int(np.rint(np.size(vx) / s)), s)) * zscale	×
179
180	_, ax = plt.subplots(nrows=1, ncols=1, constrained_layout=True)	×
181	im = None	×
182	xl = vx.min() + (vx.min() / 2) - vx[vx != vx.min()].min() / 2	×
183	xm = vx.max() + (vx.max() / 2) - vx[vx != vx.max()].max() / 2	×
184	yl = vy.min() + (vy.min() / 2) - vy[vy != vy.min()].min() / 2	×
185	ym = vy.max() + (vy.max() / 2) - vy[vy != vy.max()].max() / 2	×
186	im = NonUniformImage(	×
187	ax,
188	origin="lower",
189	cmap=kwargs["cmap"],
190	interpolation=kwargs["interpolation"],
191	extent=(xl, xm, yl, ym),
192	norm=colors.LogNorm() if logz else None,
193	)
194
195	im.set_data(np.unique(vx), np.unique(vy), grid)	×
196	ax.add_image(im)	×
197	# ax.images.append(im)
198	ax.set_xlim(xl, xm)	×
199	ax.set_ylim(yl, ym)	×
200
201	cb = plt.colorbar(im)	×
202	cb.set_label(zlabel)	×
203	smplr.style_plot2d(xlabel=xlabel, ylabel=ylabel, **kwargs)	×
204
205
206	def scatter_vplot(	5✔
207	vx,
208	vy,
209	vz,
210	xlabel=None,
211	ylabel=None,
212	zlabel=None,
213	logz=True,
214	sort=True,
215	**kwargs,
216	):
217	if sort:	5✔
218	vx, vy, vz = sort_xyz(vx, vy, vz)	5✔
219
220	_, ax = plt.subplots(nrows=1, ncols=1, constrained_layout=True)	5✔
221	xl = vx.min() + (vx.min() / 2) - vx[vx != vx.min()].min() / 2	5✔
222	xm = vx.max() + (vx.max() / 2) - vx[vx != vx.max()].max() / 2	5✔
223	yl = vy.min() + (vy.min() / 2) - vy[vy != vy.min()].min() / 2	5✔
224	ym = vy.max() + (vy.max() / 2) - vy[vy != vy.max()].max() / 2	5✔
225
226	s = plt.scatter(	5✔
227	np.concatenate((vx, vx, vx)),
228	np.concatenate((vy, vy, vy)),
229	c=np.concatenate(
230	(
231	splot.unv(vz) + splot.usd(vz),
232	splot.unv(vz) - splot.usd(vz),
233	splot.unv(vz),
234	)
235	),
236	s=np.concatenate(
237	(
238	[(3 * plt.rcParams["lines.markersize"]) ** 2 for i in range(len(vx))],
239	[(2 * plt.rcParams["lines.markersize"]) ** 2 for i in range(len(vx))],
240	[(plt.rcParams["lines.markersize"]) ** 2 for i in range(len(vx))],
241	)
242	),
243	norm=colors.LogNorm() if logz else None,
244	cmap=kwargs["cmap"],
245	)
246
247	ax.set_xlim(xl, xm)	5✔
248	ax.set_ylim(yl, ym)	5✔
249
250	cb = plt.colorbar(s)	5✔
251	cb.set_label(zlabel)	5✔
252	smplr.style_plot2d(xlabel=xlabel, ylabel=ylabel, **kwargs)	5✔

APN-Pucky / smpl / 15236750782

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

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