"""Plot correlated data"""
from __future__ import annotations
import itertools
from typing import Any
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.collections import PolyCollection
from numpy.typing import ArrayLike, NDArray
from scipy.optimize import minimize
[docs]
def corlines(
x: NDArray[Any],
y: NDArray[Any],
ycov: NDArray[Any],
*,
corlinestyle: str = ":",
cormarker: str = "_",
ax: None | Any = None,
**kwargs: Any,
) -> Any:
"""Plot data points with error bars and correlation lines.
The correlation lines indicate the correlatio between neighbouring data
points. They are attached to the vertical error bars at a relative height
corresponding to the correlation coefficient between the data points. For
positive correlations, they are attached on the same sides, for negative
correlation at opposing sides.
Parameters
----------
x, y : numpy.ndarray
The data x and y coordinates to be plotted.
ycov : numpy.ndarray
The covariance matrix describing the uncertainties of the y-values. The
error bars will correspond the the square root of the diagonal entries.
corlinestyle : str, default=":"
The Matplotlib linestyle for the correlation lines.
cormarker : str, default="_"
The Matplotlib marker used where the correlation lines attach to the
vertical error bars.
ax : matplotlib.axes.Axes, optional
Axes object to plot onto
**kwargs : dict, optional
All other keyword arguments are passed to :py:meth:`matplotlib.axes.Axes.errorbar`
Returns
-------
matplotlib.container.ErrorbarContainer
The return value of the :py:meth:`matplotlib.axes.Axes.errorbar` method.
Notes
-----
Where the correlation lines attach to the vertical error bars, gives an
indication of how much of the variance in the given data point is "caused"
by the neighbouring data points. Also, if the value of the neighbouring
data point is fixed to plus or minus 1 sigma away from its mean position,
the mean of the given data point is shifted to the position where the
correlation line attaches. Of course, this is a symmetric relationship and
the "fixing" and "causing" can equally be read in the opposite direction.
Examples
--------
.. plot::
:include-source: True
Basic usage:
>>> import numpy as np
>>> from matplotlib import pyplot as plt
>>> from nustattools import plotting as nuplt
>>> rng = np.random.default_rng()
>>> x = np.linspace(0, 10, 5)
>>> u = x[:,np.newaxis] / 4
>>> u[-2] *= -1
>>> cov = np.eye(5) + u@u.T
>>> y = rng.multivariate_normal(np.zeros(5), cov)
>>> nuplt.corlines(x, y, cov, marker="x")
"""
if ax is None:
ax = plt.gca()
# Plot error bars
yerr = np.sqrt(np.diag(ycov))
fmt = kwargs.pop("fmt", " ")
bars = ax.errorbar(x, y, yerr=yerr, fmt=fmt, **kwargs)
color = bars.lines[0].get_color()
zorder = bars.lines[0].zorder
# Get correlations between neighbours
yerr_safe = np.where(yerr > 0, yerr, 1e-12)
ycor = ycov / yerr_safe[:, np.newaxis] / yerr_safe[np.newaxis, :]
ncor = np.diag(ycor, k=1)
# Plot lines
for i, c in enumerate(ncor):
ax.plot(
[x[i], x[i + 1]],
[y[i] + yerr[i] * np.abs(c), y[i + 1] + yerr[i + 1] * c],
color=color,
linestyle=corlinestyle,
marker=cormarker,
zorder=zorder,
)
ax.plot(
[x[i], x[i + 1]],
[y[i] - yerr[i] * np.abs(c), y[i + 1] - yerr[i + 1] * c],
color=color,
linestyle=corlinestyle,
marker=cormarker,
zorder=zorder,
)
return bars
[docs]
def wedgeplot(
x: NDArray[Any],
y: NDArray[Any],
dy: NDArray[Any],
*,
wedgewidth: Any = None,
ax: Any = None,
**kwargs: Any,
) -> Any:
"""Plot vertical wedges at the given data points with the given lengths.
Parameters
----------
x, y, dy : numpy.ndarray
The data x and y coordinates and length of the wedges to be plotted.
wedgewidth : optional
The width of the wedges in axes coordinates. Can be a single number, so
it is equal for all data points; an iterable of numbers so it is
different for each, or an iterable of pairs of numbers, so there is an
asymmetric width for each.
ax : matplotlib.axes.Axes, optional
Axes object to plot onto
**kwargs : dict, optional
All other keyword arguments are passed to :py:class:`matplotlib.collections.PolyCollection`
Returns
-------
matplotlib.collections.PolyCollection
Examples
--------
.. plot::
:include-source: True
Basic usage:
>>> import numpy as np
>>> from matplotlib import pyplot as plt
>>> from nustattools import plotting as nuplt
>>> rng = np.random.default_rng()
>>> x = np.linspace(0, 10, 5)
>>> u = x[:,np.newaxis] / 4
>>> u[-2] *= -1
>>> cov = np.eye(5) + u@u.T
>>> err = np.sqrt(np.diag(cov))
>>> y = rng.multivariate_normal(np.zeros(5), cov)
>>> up = nuplt.wedgeplot(x, y, err, color="C2")
>>> down = nuplt.wedgeplot(x, y, -err, color="C3")
>>> down.set_facecolor("C1")
"""
if ax is None:
ax = plt.gca()
if wedgewidth is None:
# Try to guess a reasonable width from the data
ww = min(np.min(np.diff(x)) * 0.9, (np.max(x) - np.min(x)) / 15) # type: ignore[operator]
wedgewidth = itertools.cycle([ww])
try:
ww_cycle = itertools.cycle(wedgewidth)
except TypeError:
ww_cycle = itertools.cycle([wedgewidth])
# Plot create wedges
paths = []
for xx, yy, dd, w in zip(x, y, dy, ww_cycle):
try:
dxm = w[0]
dxp = w[1]
except (IndexError, TypeError):
dxm = w / 2
dxp = w / 2
points = [
(xx - dxm, yy),
(xx, yy + dd),
(xx + dxp, yy),
]
paths.append(points)
# Make sure the axis is scaled to include everything
ax.update_datalim(points)
col = PolyCollection(paths, **kwargs)
ax.add_collection(col)
ax.autoscale()
return col
[docs]
def pcplot(
x: NDArray[Any],
y: NDArray[Any],
ycov: NDArray[Any],
*,
componentwidth: Any = None,
scaling: float | str = "conditional-mincor",
poshatch: str = "/" * 5,
neghatch: str = "\\" * 2,
drawcorlines: bool = True,
drawconditional: bool = True,
normalize: bool = True,
ax: Any = None,
return_dict: None | dict[Any, Any] = None,
**kwargs: Any,
) -> Any:
"""Plot data points with 1st PCA component and correlation lines.
The contribution of the first principal component is subtracted from the
covariance and the remainder plotted with :py:func:`corlines`. Then the
difference to the full covariance matrix is plotted with the type of infill
indicating the direction of the first principal component.
Parameters
----------
x, y : numpy.ndarray
The data x and y coordinates to be plotted.
ycov : numpy.ndarray
The covariance matrix describing the uncertainties of the y-values. The
error bars will correspond the the square root of the diagonal entries.
componentwidth : optional
The width of the hatched areas indicating the 1st principal component
in axes coordinates. Can be a single number, so it is equal for all
data points; an iterable of numbers so it is different for each, or an
iterable of pairs of numbers, so there is an asymmetric width for each.
scaling: default="conditional-mincor"
Determines how the length of the first principal component is scaled
before removing its contribution from the covariance. If a
:py:class:`float`, the contribution is scaled with that value. At 0.0,
nothing is removed, at 1.0 the component is removed completely and the
remaining covariance's rank will reduce by 1. See `Notes` for an explanation
of the other options.
poshatch: str, optional
The Matplotlib hatch styles for the positive direction of the first
principal component.
neghatch: str, optional
The Matplotlib hatch styles for the negative direction of the first
principal component.
drawcorlines: default=True
Whether to draw correlation lines of the remaining covariance.
drawconditional: default=True
Whether to draw the conditional uncertainty of each data point, i.e.
the allowed variance if all other points are fixed. The filling of the
triangles indicates the direction of the last (smallest) principal
component.
normalize: default=True
If ``True``, the covariance is scaled such that all diagonals are 1,
and the PCA is run on the correlation matrix. If ``False``, the PCA is
run on the covariance matrix directly. In the latter case, different
error scales for different data points will have a strong influence on
the selection of the components.
ax : matplotlib.axes.Axes, optional
Axes object to plot onto
return_dict : dict, optional
Dictionary to store some of the intermediary steps of the covariance
decompositions.
**kwargs : dict, optional
All other keyword arguments are passed to :py:func:`corlines`
Returns
-------
matplotlib.container.ErrorbarContainer
The return value of the :py:func:`corlines` function.
Notes
-----
This plotting style is most useful for data where the first principal
component dominates the covariance of the data and/or there is a single
last/lowest principal component that constrains the variation much more
than the error bars suggest.
The `scaling` argument support a couple of modes to automatically determine
the desired scaling factor:
``"mincor"``
The component will be scaled such that the overall correlation in the
remaining covariance is minimized.
``"second"``
The component will be scaled such that the remaining contribution of
the first principal component is equal to the second principal
component.
``"last"``
The component will be scaled such that its contribution is equal to the
last principal component.
``"conditional"``
The scaling is maximised, while ensuring that the diagonal elements of
the remaining covariance are at least as big as the corresponding
conditional uncertainties of each bin.
``"conditional-mincor"``
The overall correlation in the remaining covariance is minimized under
the same constraints as in the ``"conditional"`` case.
Examples
--------
.. plot::
:include-source: True
Basic usage:
>>> import numpy as np
>>> from matplotlib import pyplot as plt
>>> from nustattools import plotting as nuplt
>>> rng = np.random.default_rng()
>>> x = np.linspace(0, 10, 5)
>>> u = x[:,np.newaxis] / 4
>>> u[-2] *= -1
>>> cov = np.eye(5) + u@u.T
>>> y = rng.multivariate_normal(np.zeros(5), cov)
>>> nuplt.pcplot(x, y, cov, marker="x")
.. plot::
:include-source: True
Compare scalings:
>>> import numpy as np
>>> from matplotlib import pyplot as plt
>>> from nustattools import plotting as nuplt
>>> rng = np.random.default_rng()
>>> x = np.linspace(0, 10, 5)
>>> u = x[:,np.newaxis] / 4
>>> u[-2] *= -1
>>> cov = np.eye(5) + u@u.T
>>> y = rng.multivariate_normal(np.zeros(5), cov)
>>> nuplt.pcplot(x, y, cov, componentwidth=1, scaling="last", label="last")
>>> nuplt.pcplot(x, y, cov, componentwidth=[(0.4,0)], scaling="second", label="second")
>>> nuplt.pcplot(x, y, cov, componentwidth=[(0,0.4)], scaling="mincor", label="mincor")
>>> plt.legend()
.. plot::
:include-source: True
Rank deficient covariance:
>>> import numpy as np
>>> from matplotlib import pyplot as plt
>>> from nustattools import plotting as nuplt
>>> rng = np.random.default_rng()
>>> x = np.linspace(0, 10, 5)
>>> u = x[:,np.newaxis] / 4
>>> u[-2] *= -1
>>> cov = np.eye(5) + u@u.T
>>> # Matrix to project to constant sum of data points
>>> A = np.eye(5) - np.ones((5,5)) * 1/5
>>> cov = A @ cov @ A.T
>>> y = rng.multivariate_normal(np.zeros(5), cov)
>>> nuplt.pcplot(x, y, cov)
"""
if not drawcorlines:
kwargs.update({"corlinestyle": "", "cormarker": ""})
yerr = np.sqrt(np.diag(ycov))
yerr_safe = np.where(yerr > 0, yerr, 1e-12)
if normalize:
ycor = ycov / yerr_safe[:, np.newaxis] / yerr_safe[np.newaxis, :]
yerrscale = yerr
else:
ycor = ycov
yerrscale = 1.0
# Conditional errors, i.e. if all other components are fixed
# Make sure ycov is invertible by inflating the diagonal elements a tiny bit
ycov_diag = np.diag(ycov)
ycov_diag = np.where(ycov_diag == 0, np.max(ycov_diag), ycov_diag)
ycov_safe = ycov + np.diag(ycov_diag) * 1e-12
ycovinv = np.linalg.inv(ycov_safe)
yconderr = 1 / np.sqrt(np.diag(ycovinv))
# Get first and last principal components
q, d, _ = np.linalg.svd(ycor)
w = q[:, -1]
u = q[:, 0]
# Don't remove all of 1st principal component.
# Otherwise the remaining K will be degenerate.
# This also ensures that we do nothing if ycov in uncorrelated.
if not isinstance(scaling, float) and scaling not in (
"second",
"last",
"mincor",
"conditional",
"conditional-mincor",
):
e = f"Unknown scaling: {scaling}"
raise ValueError(e)
s: float = 1.0
if isinstance(scaling, float):
# Scale from 0 to maximum allowed
s = scaling
elif scaling == "second":
# Scale so remaining contribution is same as second PCA component
s = np.sqrt(1 - d[1] / d[0])
elif scaling == "last":
# Scale so remaining contribution is same as last PCA component
s = np.sqrt(1 - d[-1] / d[0])
else:
if "conditional" in scaling:
# Scale so remaining covaraince diagonals are >= the conditional uncertainties
with np.errstate(divide="ignore", invalid="ignore"):
ss = np.sqrt(
np.nanmin((yerr**2 - yconderr**2) / (d[0] * (yerrscale * u) ** 2))
)
s = min(1, ss)
if "mincor" in scaling:
# Scale to minimize total correlation in remaining covariance
def fun(x: ArrayLike) -> Any:
v = u * yerrscale * x * np.sqrt(d[0])
V = v[:, np.newaxis] @ v[np.newaxis, :]
# Ignore degenerate components
L = (ycov - V)[d > 0, :][:, d > 0]
with np.errstate(divide="ignore", invalid="ignore"):
# Ignore divisions by zero when we scale by 1.0
det = np.linalg.det(L)
return np.prod(np.diag(L)) / det
# Start close to scaling to second, non-zero PCA component
# Ensures that we do nothing if everything is already uncorrelated
dl = d[d > 0]
ret = minimize(fun, x0=(1 - np.sqrt(dl[1] / dl[0])), bounds=[(0.0, s)])
s = ret.x
u *= yerrscale * s * np.sqrt(d[0])
U = u[:, np.newaxis] @ u[np.newaxis, :]
K = ycov - U
if np.any(np.diag(K) < 0):
e = "Remaining covariance has negative diagonal elements! Try a less aggressive scaling?"
raise RuntimeError(e)
if ax is None:
ax = plt.gca()
if componentwidth is None:
# Try to guess a reasonable width from the data
cw = min(np.min(np.diff(x)) * 0.9, (np.max(x) - np.min(x)) / 15) # type: ignore[operator]
componentwidth = itertools.cycle([cw])
try:
cw_cycle = itertools.cycle(componentwidth)
except TypeError:
cw_cycle = itertools.cycle([componentwidth])
# Plot error bars with correlation lines
bars = corlines(x, y, K, ax=ax, **kwargs)
color = bars.lines[0].get_color()
zorder = bars.lines[0].zorder
# Plot first principal component
Kerr = np.sqrt(np.diag(K))
xx: list[float] = []
yy: list[float] = []
e_min: list[float] = []
e_max: list[float] = []
fill: list[bool] = []
for i, (xs, ys, cw) in enumerate(zip(x, y, cw_cycle)):
try:
dxm = cw[0]
dxp = cw[1]
except (IndexError, TypeError):
dxm = cw / 2
dxp = cw / 2
su = np.sign(u[i])
su = 1 if su == 0 else su
emin = Kerr[i] * su
emax = yerr[i] * su
# Turn every data point into three so we can use fill_between
# and switch off filling in between points
xx.extend((xs - dxm, xs + dxp, xs + dxp))
yy.extend((ys,) * 3)
e_min.extend((emin,) * 3)
e_max.extend((emax,) * 3)
fill.extend((True, True, False))
xx_arr = np.array(xx)
yy_arr = np.array(yy)
e_min_arr = np.array(e_min)
e_max_arr = np.array(e_max)
fill_arr = np.array(fill)
# Draw first component
ax.fill_between(
xx_arr,
yy_arr + e_min_arr,
yy_arr + e_max_arr,
where=fill_arr,
hatch=poshatch,
facecolor="none",
edgecolor=color,
zorder=zorder,
)
ax.fill_between(
xx_arr,
yy_arr - e_min_arr,
yy_arr - e_max_arr,
where=fill_arr,
hatch=neghatch,
facecolor="none",
edgecolor=color,
zorder=zorder,
)
if drawconditional:
# Draw conditional probabilities and last component
sw = np.sign(w)
sw = np.where(sw == 0, 1, sw)
yb = y + Kerr * sw
yd = -(Kerr - yconderr) * sw
tri_col_pos = wedgeplot(
x, yb, yd, wedgewidth=componentwidth, closed=True, zorder=zorder
)
yb = y - Kerr * sw
yd = (Kerr - yconderr) * sw
tri_col_neg = wedgeplot(
x, yb, yd, wedgewidth=componentwidth, closed=True, zorder=zorder
)
tri_col_pos.set_linewidth(1)
tri_col_pos.set_color(color)
tri_col_pos.set_facecolor("none")
tri_col_neg.set_linewidth(1)
tri_col_neg.set_color(color)
tri_col_neg.set_alpha(0.8)
if return_dict is not None:
return_dict.update(
{
"K": K,
"u": u,
"w": w,
"yconderr": yconderr,
}
)
return bars
__all__ = ["corlines", "pcplot", "wedgeplot"]