17569160869

Committed 09 Sep 2025 01:41AM UTC coverage: 91.465% (-0.1%) from 91.614%

Build # 17569160869

Build Type

push

github

Committed by

morganjwilliams

Commit Message

Add uncertainties, add optional deps for pyproject.toml; WIP demo NB

Run Details

6226 of 6807 relevant lines covered (91.46%)

10.97 hits per line

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

93.94

/pyrolite/plot/density/__init__.py

"""
Kernel desnity estimation plots for geochemical data.
"""

import copy

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.ticker import MaxNLocator

from ...comp.codata import close
from ...util.log import Handle
from ...util.meta import get_additional_params, subkwargs
from ...util.plot.axes import add_colorbar, init_axes
from ...util.plot.density import (
    get_axis_density_methods,
    percentile_contour_values_from_meshz,
    plot_Z_percentiles,
)
from ...util.plot.style import DEFAULT_CONT_COLORMAP
from .grid import DensityGrid
from .ternary import ternary_heatmap

logger = Handle(__name__)


def density(
    arr,
    ax=None,
    logx=False,
    logy=False,
    bins=25,
    mode="density",
    extent=None,
    contours=[],
    percentiles=True,
    relim=True,
    cmap=DEFAULT_CONT_COLORMAP,
    shading="auto",
    vmin=0.0,
    colorbar=False,
    **kwargs,
):
    """
    Creates diagramatic representation of data density and/or frequency for either
    binary diagrams (X-Y) or ternary plots.
    Additional arguments are typically forwarded
    to respective :mod:`matplotlib` functions
    :func:`~matplotlib.pyplot.pcolormesh`,
    :func:`~matplotlib.pyplot.hist2d`,
    :func:`~matplotlib.pyplot.hexbin`,
    :func:`~matplotlib.pyplot.contour`, and
    :func:`~matplotlib.pyplot.contourf` (see Other Parameters, below).

    Parameters
    ----------
    arr : :class:`numpy.ndarray`
        Dataframe from which to draw data.
    ax : :class:`matplotlib.axes.Axes`, `None`
        The subplot to draw on.
    logx : :class:`bool`, `False`
        Whether to use a logspaced *grid* on the x axis. Values strictly >0 required.
    logy : :class:`bool`, `False`
        Whether to use a logspaced *grid* on the y axis. Values strictly >0 required.
    bins : :class:`int`, 20
        Number of bins used in the gridded functions (histograms, KDE evaluation grid).
    mode : :class:`str`, 'density'
        Different modes used here: ['density', 'hexbin', 'hist2d']
    extent : :class:`list`
        Predetermined extent of the grid for which to from the histogram/KDE. In the
        general form (xmin, xmax, ymin, ymax).
    contours : :class:`list`
        Contours to add to the plot, where :code:`mode='density'` is used.
    percentiles :  :class:`bool`, `True`
        Whether contours specified are to be converted to percentiles.
    relim : :class:`bool`, :code:`True`
        Whether to relimit the plot based on xmin, xmax values.
    cmap : :class:`matplotlib.colors.Colormap`
        Colormap for mapping surfaces.
    vmin : :class:`float`, 0.
        Minimum value for colormap.
    shading : :class:`str`, 'auto'
        Shading to apply to pcolormesh.
    colorbar : :class:`bool`, False
        Whether to append a linked colorbar to the generated mappable image.

    {otherparams}

    Returns
    -------
    :class:`matplotlib.axes.Axes`
        Axes on which the densityplot is plotted.


    .. seealso::

        Functions:

            :func:`matplotlib.pyplot.pcolormesh`
            :func:`matplotlib.pyplot.hist2d`
            :func:`matplotlib.pyplot.contourf`

    Notes
    -----
    The default density estimates and derived contours are generated based on
    kernel density estimates. Assumptions around e.g. 95% of points lying within
    a 95% contour won't necessarily be valid for non-normally distributed data
    (instead, this represents the approximate 95% percentile on the kernel
    density estimate). Note that contours are currently only generated; for
    `mode="density"`; future updates may allow the use of a histogram
    basis, which would give results closer to 95% data percentiles.

    Todo
    ----
    * Allow generation of contours from histogram data, rather than just
        the kernel density estimate.
    * Implement an option and filter to 'scatter' points below the minimum threshold
        or maximum percentile contours.
    """
    if (mode == "density") & np.isclose(vmin, 0.0):  # if vmin is not specified
        vmin = 0.02  # 2% max height | 98th percentile

    if arr.shape[-1] == 3:
        projection = "ternary"
    else:
        projection = None

    ax = init_axes(ax=ax, projection=projection, **kwargs)

    pcolor, contour, contourf = get_axis_density_methods(ax)
    background_color = (*ax.patch.get_facecolor()[:-1], 0.0)

    if cmap is not None:
        if isinstance(cmap, str):
            cmap = plt.get_cmap(cmap)
        cmap = copy.copy(cmap)  # without this, it would modify the global cmap
        cmap.set_under((1, 1, 1, 0))

    if mode == "density":
        cbarlabel = "Kernel Density Estimate"
    else:
        cbarlabel = "Frequency"

    valid_rows = np.isfinite(arr).all(axis=-1)

    if (mode in ["hexbin", "hist2d"]) and contours:
        raise NotImplementedError(
            "Contours are not currently implemented for 'hexbin' or 'hist2d' modes."
        )

    if (arr.size > 0) and valid_rows.any():
        # Data can't be plotted if there's any nans, so we can exclude these
        arr = arr[valid_rows]

        if projection is None:  # binary
            x, y = arr.T
            grid = DensityGrid(
                x,
                y,
                bins=bins,
                logx=logx,
                logy=logy,
                extent=extent,
                **subkwargs(kwargs, DensityGrid),
            )
            if mode == "hexbin":
                # extent values are exponents (i.e. 3 -> 10**3)
                mappable = ax.hexbin(
                    x,
                    y,
                    gridsize=bins,
                    cmap=cmap,
                    extent=grid.get_hex_extent(),
                    xscale=["linear", "log"][logx],
                    yscale=["linear", "log"][logy],
                    **subkwargs(kwargs, ax.hexbin),
                )

            elif mode == "hist2d":
                _, _, _, im = ax.hist2d(
                    x,
                    y,
                    bins=[grid.grid_xe, grid.grid_ye],
                    range=grid.get_range(),
                    cmap=cmap,
                    cmin=[0, 1][vmin > 0],
                    **subkwargs(kwargs, ax.hist2d),
                )
                mappable = im

            elif mode == "density":
                zei = grid.kdefrom(
                    arr,
                    xtransform=[lambda x: x, np.log][logx],
                    ytransform=[lambda y: y, np.log][logy],
                    mode="edges",
                    **subkwargs(kwargs, grid.kdefrom),
                )

                if percentiles:  # 98th percentile
                    vmin = percentile_contour_values_from_meshz(zei, [1.0 - vmin])[1][0]
                    logger.debug(
                        "Updating `vmin` to percentile equiv: {:.2f}".format(vmin)
                    )

                if not contours:
                    # pcolormesh using bin edges
                    mappable = pcolor(
                        grid.grid_xei,
                        grid.grid_yei,
                        zei,
                        cmap=cmap,
                        vmin=vmin,
                        shading=shading,
                        **subkwargs(kwargs, pcolor),
                    )
                    mappable.set_edgecolor(background_color)
                    mappable.set_linestyle("None")
                    mappable.set_lw(0.0)
                else:
                    mappable = _add_contours(
                        grid.grid_xei,
                        grid.grid_yei,
                        zi=zei.reshape(grid.grid_xei.shape),
                        ax=ax,
                        contours=contours,
                        percentiles=percentiles,
                        cmap=cmap,
                        vmin=vmin,
                        **kwargs,
                    )
            if relim and (extent is not None):
                ax.axis(extent)
        elif projection == "ternary":  # ternary
            if shading == "auto":
                shading = "flat"  # auto cant' be passed to tripcolor
            # zeros make nans in this case, due to the heatmap calculations
            arr[~(arr > 0).all(axis=1), :] = np.nan
            arr = close(arr)
            if mode == "hexbin":
                raise NotImplementedError
            # density, histogram etc parsed here
            coords, zi, _ = ternary_heatmap(arr, bins=bins, mode=mode)

            if percentiles:  # 98th percentile
                vmin = percentile_contour_values_from_meshz(zi, [1.0 - vmin])[1][0]
                logger.debug("Updating `vmin` to percentile equiv: {:.2f}".format(vmin))

            # remove coords where H==0, as ax.tripcolor can't deal with variable alpha :'(
            fltr = (zi != 0) & (zi >= vmin)
            coords = coords[fltr.flatten(), :]
            zi = zi[fltr]

            if not contours:
                tri_poly_collection = pcolor(
                    *coords.T,
                    zi.flatten(),
                    cmap=cmap,
                    vmin=vmin,
                    shading=shading,
                    **subkwargs(kwargs, pcolor),
                )

                mappable = tri_poly_collection
            else:
                mappable = _add_contours(
                    *coords.T,
                    zi=zi.flatten(),
                    ax=ax,
                    contours=contours,
                    percentiles=percentiles,
                    cmap=cmap,
                    vmin=vmin,
                    **kwargs,
                )
            ax.set_aspect("equal")
        else:
            if not arr.ndim in [0, 1, 2]:
                raise NotImplementedError

        if colorbar:
            cbkwargs = kwargs.copy()
            cbkwargs["label"] = cbarlabel
            add_colorbar(mappable, **cbkwargs)

    return ax


def _add_contours(
    *coords,
    zi=None,
    ax=None,
    contours=[],
    cmap=DEFAULT_CONT_COLORMAP,
    vmin=0.0,
    extent=None,
    **kwargs,
):
    """
    Add density-based contours to a plot.
    """
    # get the contour levels
    percentiles = kwargs.pop("percentiles", True)
    levels = contours or kwargs.get("levels", None)
    pcolor, contour, contourf = get_axis_density_methods(ax)
    if percentiles and not isinstance(levels, int):
        # plot individual percentile contours
        _cs = plot_Z_percentiles(
            *coords,
            zi=zi,
            ax=ax,
            percentiles=levels,
            extent=extent,
            cmap=cmap,
            **kwargs,
        )
        mappable = _cs
    else:
        # plot interval contours
        if levels is None:
            levels = MaxNLocator(nbins=10).tick_values(zi.min(), zi.max())
        elif isinstance(levels, int):
            levels = MaxNLocator(nbins=levels).tick_values(zi.min(), zi.max())
        else:
            raise NotImplementedError
        # filled contours
        mappable = contourf(
            *coords, zi, extent=extent, levels=levels, cmap=cmap, vmin=vmin, **kwargs
        )
        # contours
        contour(
            *coords, zi, extent=extent, levels=levels, cmap=cmap, vmin=vmin, **kwargs
        )
    return mappable


_add_additional_parameters = True

density.__doc__ = density.__doc__.format(
    otherparams=[
        "",
        get_additional_params(
            density,
            plt.pcolormesh,
            plt.hist2d,
            plt.hexbin,
            plt.contour,
            plt.contourf,
            header="Other Parameters",
            indent=4,
            subsections=True,
        ),
    ][_add_additional_parameters]
)

1	"""
2	Kernel desnity estimation plots for geochemical data.
3	"""
4
5	import copy	12✔
6
7	import matplotlib.pyplot as plt	12✔
8	import numpy as np	12✔
9	from matplotlib.ticker import MaxNLocator	12✔
10
11	from ...comp.codata import close	12✔
12	from ...util.log import Handle	12✔
13	from ...util.meta import get_additional_params, subkwargs	12✔
14	from ...util.plot.axes import add_colorbar, init_axes	12✔
15	from ...util.plot.density import (	12✔
16	get_axis_density_methods,
17	percentile_contour_values_from_meshz,
18	plot_Z_percentiles,
19	)
20	from ...util.plot.style import DEFAULT_CONT_COLORMAP	12✔
21	from .grid import DensityGrid	12✔
22	from .ternary import ternary_heatmap	12✔
23
24	logger = Handle(__name__)	12✔
25
26
27	def density(	12✔
28	arr,
29	ax=None,
30	logx=False,
31	logy=False,
32	bins=25,
33	mode="density",
34	extent=None,
35	contours=[],
36	percentiles=True,
37	relim=True,
38	cmap=DEFAULT_CONT_COLORMAP,
39	shading="auto",
40	vmin=0.0,
41	colorbar=False,
42	**kwargs,
43	):
44	"""
45	Creates diagramatic representation of data density and/or frequency for either
46	binary diagrams (X-Y) or ternary plots.
47	Additional arguments are typically forwarded
48	to respective :mod:`matplotlib` functions
49	:func:`~matplotlib.pyplot.pcolormesh`,
50	:func:`~matplotlib.pyplot.hist2d`,
51	:func:`~matplotlib.pyplot.hexbin`,
52	:func:`~matplotlib.pyplot.contour`, and
53	:func:`~matplotlib.pyplot.contourf` (see Other Parameters, below).
54
55	Parameters
56	----------
57	arr : :class:`numpy.ndarray`
58	Dataframe from which to draw data.
59	ax : :class:`matplotlib.axes.Axes`, `None`
60	The subplot to draw on.
61	logx : :class:`bool`, `False`
62	Whether to use a logspaced grid on the x axis. Values strictly >0 required.
63	logy : :class:`bool`, `False`
64	Whether to use a logspaced grid on the y axis. Values strictly >0 required.
65	bins : :class:`int`, 20
66	Number of bins used in the gridded functions (histograms, KDE evaluation grid).
67	mode : :class:`str`, 'density'
68	Different modes used here: ['density', 'hexbin', 'hist2d']
69	extent : :class:`list`
70	Predetermined extent of the grid for which to from the histogram/KDE. In the
71	general form (xmin, xmax, ymin, ymax).
72	contours : :class:`list`
73	Contours to add to the plot, where :code:`mode='density'` is used.
74	percentiles : :class:`bool`, `True`
75	Whether contours specified are to be converted to percentiles.
76	relim : :class:`bool`, :code:`True`
77	Whether to relimit the plot based on xmin, xmax values.
78	cmap : :class:`matplotlib.colors.Colormap`
79	Colormap for mapping surfaces.
80	vmin : :class:`float`, 0.
81	Minimum value for colormap.
82	shading : :class:`str`, 'auto'
83	Shading to apply to pcolormesh.
84	colorbar : :class:`bool`, False
85	Whether to append a linked colorbar to the generated mappable image.
86
87	{otherparams}
88
89	Returns
90	-------
91	:class:`matplotlib.axes.Axes`
92	Axes on which the densityplot is plotted.
93
94
95	.. seealso::
96
97	Functions:
98
99	:func:`matplotlib.pyplot.pcolormesh`
100	:func:`matplotlib.pyplot.hist2d`
101	:func:`matplotlib.pyplot.contourf`
102
103	Notes
104	-----
105	The default density estimates and derived contours are generated based on
106	kernel density estimates. Assumptions around e.g. 95% of points lying within
107	a 95% contour won't necessarily be valid for non-normally distributed data
108	(instead, this represents the approximate 95% percentile on the kernel
109	density estimate). Note that contours are currently only generated; for
110	`mode="density"`; future updates may allow the use of a histogram
111	basis, which would give results closer to 95% data percentiles.
112
113	Todo
114	----
115	* Allow generation of contours from histogram data, rather than just
116	the kernel density estimate.
117	* Implement an option and filter to 'scatter' points below the minimum threshold
118	or maximum percentile contours.
119	"""
120	if (mode == "density") & np.isclose(vmin, 0.0): # if vmin is not specified	12✔
121	vmin = 0.02 # 2% max height \| 98th percentile	12✔
122
123	if arr.shape[-1] == 3:	12✔
124	projection = "ternary"	12✔
125	else:
126	projection = None	12✔
127
128	ax = init_axes(ax=ax, projection=projection, **kwargs)	12✔
129
130	pcolor, contour, contourf = get_axis_density_methods(ax)	12✔
131	background_color = (*ax.patch.get_facecolor()[:-1], 0.0)	12✔
132
133	if cmap is not None:	12✔
134	if isinstance(cmap, str):	12✔
135	cmap = plt.get_cmap(cmap)	12✔
136	cmap = copy.copy(cmap) # without this, it would modify the global cmap	12✔
137	cmap.set_under((1, 1, 1, 0))	12✔
138
139	if mode == "density":	12✔
140	cbarlabel = "Kernel Density Estimate"	12✔
141	else:
142	cbarlabel = "Frequency"	12✔
143
144	valid_rows = np.isfinite(arr).all(axis=-1)	12✔
145
146	if (mode in ["hexbin", "hist2d"]) and contours:	12✔
147	raise NotImplementedError(	×
148	"Contours are not currently implemented for 'hexbin' or 'hist2d' modes."
149	)
150
151	if (arr.size > 0) and valid_rows.any():	12✔
152	# Data can't be plotted if there's any nans, so we can exclude these
153	arr = arr[valid_rows]	12✔
154
155	if projection is None: # binary	12✔
156	x, y = arr.T	12✔
157	grid = DensityGrid(	12✔
158	x,
159	y,
160	bins=bins,
161	logx=logx,
162	logy=logy,
163	extent=extent,
164	**subkwargs(kwargs, DensityGrid),
165	)
166	if mode == "hexbin":	12✔
167	# extent values are exponents (i.e. 3 -> 10**3)
168	mappable = ax.hexbin(	12✔
169	x,
170	y,
171	gridsize=bins,
172	cmap=cmap,
173	extent=grid.get_hex_extent(),
174	xscale=["linear", "log"][logx],
175	yscale=["linear", "log"][logy],
176	**subkwargs(kwargs, ax.hexbin),
177	)
178
179	elif mode == "hist2d":	12✔
180	_, _, _, im = ax.hist2d(	12✔
181	x,
182	y,
183	bins=[grid.grid_xe, grid.grid_ye],
184	range=grid.get_range(),
185	cmap=cmap,
186	cmin=[0, 1][vmin > 0],
187	**subkwargs(kwargs, ax.hist2d),
188	)
189	mappable = im	12✔
190
191	elif mode == "density":	12✔
192	zei = grid.kdefrom(	12✔
193	arr,
194	xtransform=[lambda x: x, np.log][logx],
195	ytransform=[lambda y: y, np.log][logy],
196	mode="edges",
197	**subkwargs(kwargs, grid.kdefrom),
198	)
199
200	if percentiles: # 98th percentile	12✔
201	vmin = percentile_contour_values_from_meshz(zei, [1.0 - vmin])[1][0]	12✔
202	logger.debug(	12✔
203	"Updating `vmin` to percentile equiv: {:.2f}".format(vmin)
204	)
205
206	if not contours:	12✔
207	# pcolormesh using bin edges
208	mappable = pcolor(	12✔
209	grid.grid_xei,
210	grid.grid_yei,
211	zei,
212	cmap=cmap,
213	vmin=vmin,
214	shading=shading,
215	**subkwargs(kwargs, pcolor),
216	)
217	mappable.set_edgecolor(background_color)	12✔
218	mappable.set_linestyle("None")	12✔
219	mappable.set_lw(0.0)	12✔
220	else:
221	mappable = _add_contours(	12✔
222	grid.grid_xei,
223	grid.grid_yei,
224	zi=zei.reshape(grid.grid_xei.shape),
225	ax=ax,
226	contours=contours,
227	percentiles=percentiles,
228	cmap=cmap,
229	vmin=vmin,
230	**kwargs,
231	)
232	if relim and (extent is not None):	12✔
233	ax.axis(extent)	×
234	elif projection == "ternary": # ternary	12✔
235	if shading == "auto":	12✔
236	shading = "flat" # auto cant' be passed to tripcolor	12✔
237	# zeros make nans in this case, due to the heatmap calculations
238	arr[~(arr > 0).all(axis=1), :] = np.nan	12✔
239	arr = close(arr)	12✔
240	if mode == "hexbin":	12✔
241	raise NotImplementedError	12✔
242	# density, histogram etc parsed here
243	coords, zi, _ = ternary_heatmap(arr, bins=bins, mode=mode)	12✔
244
245	if percentiles: # 98th percentile	12✔
246	vmin = percentile_contour_values_from_meshz(zi, [1.0 - vmin])[1][0]	12✔
247	logger.debug("Updating `vmin` to percentile equiv: {:.2f}".format(vmin))	12✔
248
249	# remove coords where H==0, as ax.tripcolor can't deal with variable alpha :'(
250	fltr = (zi != 0) & (zi >= vmin)	12✔
251	coords = coords[fltr.flatten(), :]	12✔
252	zi = zi[fltr]	12✔
253
254	if not contours:	12✔
255	tri_poly_collection = pcolor(	12✔
256	*coords.T,
257	zi.flatten(),
258	cmap=cmap,
259	vmin=vmin,
260	shading=shading,
261	**subkwargs(kwargs, pcolor),
262	)
263
264	mappable = tri_poly_collection	12✔
265	else:
266	mappable = _add_contours(	12✔
267	*coords.T,
268	zi=zi.flatten(),
269	ax=ax,
270	contours=contours,
271	percentiles=percentiles,
272	cmap=cmap,
273	vmin=vmin,
274	**kwargs,
275	)
276	ax.set_aspect("equal")	12✔
277	else:
278	if not arr.ndim in [0, 1, 2]:	×
279	raise NotImplementedError	×
280
281	if colorbar:	12✔
282	cbkwargs = kwargs.copy()	12✔
283	cbkwargs["label"] = cbarlabel	12✔
284	add_colorbar(mappable, **cbkwargs)	12✔
285
286	return ax	12✔
287
288
289	def _add_contours(	12✔
290	*coords,
291	zi=None,
292	ax=None,
293	contours=[],
294	cmap=DEFAULT_CONT_COLORMAP,
295	vmin=0.0,
296	extent=None,
297	**kwargs,
298	):
299	"""
300	Add density-based contours to a plot.
301	"""
302	# get the contour levels
303	percentiles = kwargs.pop("percentiles", True)	12✔
304	levels = contours or kwargs.get("levels", None)	12✔
305	pcolor, contour, contourf = get_axis_density_methods(ax)	12✔
306	if percentiles and not isinstance(levels, int):	12✔
307	# plot individual percentile contours
308	_cs = plot_Z_percentiles(	12✔
309	*coords,
310	zi=zi,
311	ax=ax,
312	percentiles=levels,
313	extent=extent,
314	cmap=cmap,
315	**kwargs,
316	)
317	mappable = _cs	12✔
318	else:
319	# plot interval contours
320	if levels is None:	12✔
321	levels = MaxNLocator(nbins=10).tick_values(zi.min(), zi.max())	×
322	elif isinstance(levels, int):	12✔
323	levels = MaxNLocator(nbins=levels).tick_values(zi.min(), zi.max())	12✔
324	else:
325	raise NotImplementedError	×
326	# filled contours
327	mappable = contourf(	12✔
328	coords, zi, extent=extent, levels=levels, cmap=cmap, vmin=vmin, *kwargs
329	)
330	# contours
331	contour(	12✔
332	coords, zi, extent=extent, levels=levels, cmap=cmap, vmin=vmin, *kwargs
333	)
334	return mappable	12✔
335
336
337	_add_additional_parameters = True	12✔
338
339	density.__doc__ = density.__doc__.format(	12✔
340	otherparams=[
341	"",
342	get_additional_params(
343	density,
344	plt.pcolormesh,
345	plt.hist2d,
346	plt.hexbin,
347	plt.contour,
348	plt.contourf,
349	header="Other Parameters",
350	indent=4,
351	subsections=True,
352	),
353	][_add_additional_parameters]
354	)

morganjwilliams / pyrolite / 17569160869

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