"""
Two-dimensional distance-versus-reddening posterior plots.
"""
import numpy as np
from matplotlib import pyplot as plt
from ..utils.math import AVLIM, RVLIM
from .binning import _per_object, bin_pdfs_distred
__all__ = ["dist_vs_red"]
[docs]
def dist_vs_red(
data,
ebv=None,
dist_type="distance_modulus",
lndistprior=None,
coord=None,
avlim=AVLIM,
rvlim=RVLIM,
weights=None,
parallax=None,
parallax_err=None,
Nr=300,
cmap="Blues",
bins=300,
span=None,
smooth=(0.015, 0.005),
plot_kwargs=None,
truths=None,
truth_color="red",
truth_kwargs=None,
rstate=None,
fig=None,
):
"""
Generate a 2-D plot of distance vs reddening.
Parameters
----------
data : 3-tuple or 4-tuple of `~numpy.ndarray`
Arrays of shape ``(Nsamps,)`` or ``(Nobj, Nsamps)`` containing either
saved ``(dists, reds, dreds)`` draws in ``(kpc, mag, -)`` or
``(scales, avs, rvs, covs_sar)`` used to regenerate draws under the
distance and parallax priors. Saved draws are preferred because
regeneration applies only the distance and parallax priors rather than
the dust-map and label-dependent priors used during fitting. When
``(Nobj, Nsamps)`` arrays are passed, the plotted histogram is the
mean of the per-object 2-D PDFs.
ebv : bool, optional
If True, convert :math:`A(V)` to :math:`E(B-V) = A(V) / R(V)`. Default
is ``None`` (plot :math:`A(V)` in mag).
dist_type : str, optional
Distance quantity to plot on the x-axis (``'parallax'``, ``'scale'``,
``'distance'``, or ``'distance_modulus'``). Default is
``'distance_modulus'``.
lndistprior : func, optional
The log-distance prior function used. If not provided,
`~brutus.priors.logp_galactic_structure` is assumed.
coord : 2-tuple or list of 2-tuples, optional
The galactic `(l, b)` coordinates passed to `lndistprior`. A single
`(l, b)` pair is shared across all objects; for multi-object input
a list with one pair per object is also accepted.
avlim : 2-tuple, optional
The Av limits used to truncate regenerated draws. Without `span`,
they also set the A(V) range of the bins and of the smoothing.
Default is `AVLIM`, (0, 20), the `BruteForce.fit` default.
rvlim : 2-tuple, optional
The Rv limits used to truncate results. Default is `RVLIM`, (1, 8).
weights : `~numpy.ndarray` of shape `(Nsamps)` or `(Nobj, Nsamps)`, optional
An optional set of importance weights used to reweight the samples.
parallax : float or `~numpy.ndarray` of shape `(Nobj,)`, optional
The parallax estimate for each source in mas. A float applies to
every object.
parallax_err : float or `~numpy.ndarray` of shape `(Nobj,)`, optional
The parallax error for each source in mas. A float applies to
every object.
Nr : int, optional
The number of Monte Carlo realizations used when sampling using the
provided parallax prior. Default is `300`.
cmap : str, optional
The colormap used when plotting. Default is `'Blues'`.
bins : int or list of ints with length `(ndim,)`, optional
The number of bins to be used in each dimension. Default is `300`.
span : iterable with shape `(2, 2)`, optional
`((av_min, av_max), (d_min, d_max))`, A(V) in mag and distance in
kpc, as in `bin_pdfs_distred`. If not provided, A(V) spans `avlim`
and the distance spans distance modulus 4 to 19.
smooth : float or 2-tuple of floats, optional
The smoothing, as in `bin_pdfs_distred`. Default is `(0.015, 0.005)`.
The distance smoothing never exceeds what the parallax allows.
plot_kwargs : dict, optional
Extra keyword arguments to be used when plotting the smoothed
2-D histograms.
truths : iterable with shape `(ndim,)`, optional
A list of reference values that will be overplotted on the traces and
marginalized 1-D posteriors as solid horizontal/vertical lines.
Individual values can be exempt using `None`. Default is `None`.
truth_color : str or iterable with shape `(ndim,)`, optional
A `~matplotlib`-style color (either a single color or a different
value for each subplot) used when plotting `truths`.
Default is `'red'`.
truth_kwargs : dict, optional
Extra keyword arguments that will be used for plotting the vertical
and horizontal lines with `truths`.
rstate : `~numpy.random.RandomState`, optional
`~numpy.random.RandomState` instance.
fig : (`~matplotlib.figure.Figure`, `~matplotlib.axes.Axes`), optional
The figure and axes to draw on. Default is the current axes,
`~matplotlib.pyplot.gca`.
Returns
-------
hist2d : (counts, xedges, yedges, `~matplotlib.figure.Image`)
Output 2-D histogram.
"""
# Initialize values.
truth_kwargs = dict(truth_kwargs or {})
if plot_kwargs is None:
plot_kwargs = dict()
# Set defaults for truth plotting
truth_kwargs.setdefault("linestyle", "solid")
truth_kwargs.setdefault("linewidth", 2)
truth_kwargs.setdefault("alpha", 0.7)
# Handle single object case - convert to array format expected by bin_pdfs_distred
# bin_pdfs_distred expects (n_objects, n_samples) shape
if len(data[0].shape) == 1:
# Single object case: convert from (n_samples,) to (1, n_samples)
if len(data) == 3: # (dists, reds, dreds)
data = tuple(arr[None, :] for arr in data) # Add object dimension
elif len(data) == 4: # (scales, avs, rvs, covs_sar)
data = (
data[0][None, :],
data[1][None, :],
data[2][None, :],
data[3][None, :],
)
single_object = True
# Convert coord to list format
if coord is not None:
coord = [coord]
# Convert weights to (Nobj, Nsamps) format
if weights is not None:
weights = np.atleast_2d(np.asarray(weights))
else:
# Multi-object case
single_object = False
# A single set of per-sample weights is shared across all objects;
# broadcast it to the (Nobj, Nsamps) shape bin_pdfs_distred expects.
if weights is not None:
weights = np.asarray(weights)
if weights.ndim == 1:
nobj, nsamps = data[0].shape
if weights.shape[0] != nsamps:
raise ValueError(
f"`weights` has shape {weights.shape}, which matches "
f"neither (Nsamps,) = ({nsamps},) nor "
f"(Nobj, Nsamps) = {(nobj, nsamps)}."
)
weights = np.broadcast_to(weights, (nobj, nsamps))
nobj = data[0].shape[0]
parallax = _per_object(parallax, "parallax", nobj)
parallax_err = _per_object(parallax_err, "parallax_err", nobj)
# Use bin_pdfs_distred to do all the heavy lifting for data preparation
binned_vals, xedges, yedges = bin_pdfs_distred(
data,
cdf=False,
ebv=ebv,
dist_type=dist_type,
lndistprior=lndistprior,
coord=coord,
avlim=avlim,
rvlim=rvlim,
weights=weights,
parallaxes=parallax,
parallax_errors=parallax_err,
Nr=Nr,
bins=bins,
span=span,
smooth=smooth,
rstate=rstate,
verbose=False,
)
# For single object, extract the first (and only) object's data
if single_object:
H = binned_vals[0]
else:
# For multiple objects, stack (average) the per-object 2-D PDFs so the
# plot reflects the whole sample, as in dust-mapping sightline stacks.
H = binned_vals.mean(axis=0)
# Set up axis labels
if ebv:
ylabel = r"$E(B-V)$ [mag]"
else:
ylabel = r"$A_v$ [mag]"
if dist_type == "scale":
xlabel = r"$s$"
elif dist_type == "parallax":
xlabel = r"$\pi$ [mas]"
elif dist_type == "distance":
xlabel = r"$d$ [kpc]"
elif dist_type == "distance_modulus":
xlabel = r"$\mu$"
# Determine plot extent
xlims = [xedges[0], xedges[-1]]
ylims = [yedges[0], yedges[-1]]
# Generate the plot
ax = plt.gca() if fig is None else fig[1]
img = ax.imshow(
H.T,
cmap=cmap,
aspect="auto",
interpolation="none",
origin="lower",
extent=[xlims[0], xlims[1], ylims[0], ylims[1]],
**plot_kwargs,
)
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
# Add truth values if provided
if truths is not None:
if truths[0] is not None: # x-axis truth
for t in np.atleast_1d(truths[0]):
ax.axvline(t, color=truth_color, **truth_kwargs)
if truths[1] is not None: # y-axis truth
for t in np.atleast_1d(truths[1]):
ax.axhline(t, color=truth_color, **truth_kwargs)
return H, xedges, yedges, img