15215176814

Committed 23 May 2025 04:44PM UTC coverage: 87.5% (-0.6%) from 88.129%

Build # 15215176814

Build Type

push

github

Committed by

APN-Pucky

Commit Message

Report coverals

Run Details

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

62.83

/smpl/plot2d.py

import matplotlib.pyplot as plt
import numpy as np
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"],
    "xaxis": [None, "."],
    "yaxis": [None, "."],
    "zaxis": [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 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)
    if "title" in kwargs and kwargs["title"] is not None:
        plt.title(kwargs["title"])


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, xaxis=None, yaxis=None, zaxis=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(zaxis)
    plt.xlabel(xaxis)
    plt.ylabel(yaxis)


def map_vplot(
    tvx,
    tvy,
    tvz,
    xaxis=None,
    yaxis=None,
    zaxis=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
        xaxis, yaxis = yaxis, xaxis
        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.images.append(im)
    ax.set_xlim(xl, xm)
    ax.set_ylim(yl, ym)

    cb = plt.colorbar(im)
    cb.set_label(zaxis)
    plt.xlabel(xaxis)
    plt.ylabel(yaxis)


def scatter_vplot(
    vx,
    vy,
    vz,
    xaxis=None,
    yaxis=None,
    zaxis=None,
    logz=True,
    sort=True,
    fill_missing=True,
    zscale=1.0,
    **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(zaxis)
    plt.xlabel(xaxis)
    plt.ylabel(yaxis)

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

APN-Pucky / smpl / 15215176814

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