Skip to content

Analysis

pairplot(samples, points=None, limits=None, subset=None, upper='hist', lower=None, diag='hist', figsize=(10, 10), labels=None, ticks=None, offdiag=None, diag_kwargs=None, upper_kwargs=None, lower_kwargs=None, fig_kwargs=None, fig=None, axes=None, **kwargs)

Plot samples in a 2D grid showing marginals and pairwise marginals.

Each of the diagonal plots can be interpreted as a 1D-marginal of the distribution that the samples were drawn from. Each upper-diagonal plot can be interpreted as a 2D-marginal of the distribution.

Parameters:

Name Type Description Default
samples Union[List[ndarray], List[Tensor], ndarray, Tensor]

Samples used to build the histogram.

required
points Optional[Union[List[ndarray], List[Tensor], ndarray, Tensor]]

List of additional points to scatter.

None
limits Optional[Union[List, Tensor]]

Array containing the plot xlim for each parameter dimension. If None, just use the min and max of the passed samples

None
subset Optional[List[int]]

List containing the dimensions to plot. E.g. subset=[1,3] will plot plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and, if they exist, the 4th, 5th and so on).

None
upper Optional[Union[List[Optional[UpperLiteral]], UpperLiteral]]

Plotting style for upper diagonal, {hist, scatter, contour, kde, None}.

'hist'
lower Optional[Union[List[Optional[LowerLiteral]], LowerLiteral]]

Plotting style for upper diagonal, {hist, scatter, contour, kde, None}.

None
diag Optional[Union[List[Optional[DiagLiteral]], DiagLiteral]]

Plotting style for diagonal, {hist, scatter, kde}.

'hist'
figsize Tuple

Size of the entire figure.

(10, 10)
labels Optional[List[str]]

List of strings specifying the names of the parameters.

None
ticks Optional[Union[List, Tensor]]

Position of the ticks.

None
offdiag Optional[Union[List[Optional[str]], str]]

deprecated, use upper instead.

None
diag_kwargs KwargsType[DiagOptions]

Additional arguments to adjust the diagonal plot, see the source code in KdeDiagOptions, HistDiagOptions or ScatterDiagOptions.

None
upper_kwargs KwargsType[OffDiagOptions]

Additional arguments to adjust the upper diagonal plot, see the source code in KdeOffDiagOptions, HistOffDiagOptions, ScatterOffDiagOptions, ContourOffDiagOptions or PlotOffDiagOptions.

None
lower_kwargs KwargsType[OffDiagOptions]

Additional arguments to adjust the lower diagonal plot, see the source code in KdeOffDiagOptions, HistOffDiagOptions, ScatterOffDiagOptions, ContourOffDiagOptions or PlotOffDiagOptions.

None
fig_kwargs Optional[Union[Dict, FigOptions]]

Additional arguments to adjust the overall figure, see the source code in FigOptions

None
fig Optional[FigureBase]

matplotlib figure to plot on.

None
axes Optional[Axes]

matplotlib axes corresponding to fig.

None
**kwargs Optional[Any]

Additional arguments to adjust the plot (deprecated).

{}

Returns: figure and axis of posterior distribution plot

Source code in sbi/analysis/plot.py
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
def pairplot(
    samples: Union[List[np.ndarray], List[torch.Tensor], np.ndarray, torch.Tensor],
    points: Optional[
        Union[List[np.ndarray], List[torch.Tensor], np.ndarray, torch.Tensor]
    ] = None,
    limits: Optional[Union[List, torch.Tensor]] = None,
    subset: Optional[List[int]] = None,
    upper: Optional[Union[List[Optional[UpperLiteral]], UpperLiteral]] = "hist",
    lower: Optional[Union[List[Optional[LowerLiteral]], LowerLiteral]] = None,
    diag: Optional[Union[List[Optional[DiagLiteral]], DiagLiteral]] = "hist",
    figsize: Tuple = (10, 10),
    labels: Optional[List[str]] = None,
    ticks: Optional[Union[List, torch.Tensor]] = None,
    offdiag: Optional[Union[List[Optional[str]], str]] = None,
    diag_kwargs: KwargsType[DiagOptions] = None,
    upper_kwargs: KwargsType[OffDiagOptions] = None,
    lower_kwargs: KwargsType[OffDiagOptions] = None,
    fig_kwargs: Optional[Union[Dict, FigOptions]] = None,
    fig: Optional[FigureBase] = None,
    axes: Optional[Axes] = None,
    **kwargs: Optional[Any],
) -> Tuple[FigureBase, Axes]:
    """
    Plot samples in a 2D grid showing marginals and pairwise marginals.

    Each of the diagonal plots can be interpreted as a 1D-marginal of the distribution
    that the samples were drawn from. Each upper-diagonal plot can be interpreted as a
    2D-marginal of the distribution.

    Args:
        samples: Samples used to build the histogram.
        points: List of additional points to scatter.
        limits: Array containing the plot xlim for each parameter dimension. If None,
            just use the min and max of the passed samples
        subset: List containing the dimensions to plot. E.g. subset=[1,3] will plot
            plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and,
            if they exist, the 4th, 5th and so on).
        upper: Plotting style for upper diagonal, {hist, scatter, contour, kde,
            None}.
        lower: Plotting style for upper diagonal, {hist, scatter, contour, kde,
            None}.
        diag: Plotting style for diagonal, {hist, scatter, kde}.
        figsize: Size of the entire figure.
        labels: List of strings specifying the names of the parameters.
        ticks: Position of the ticks.
        offdiag: deprecated, use upper instead.
        diag_kwargs: Additional arguments to adjust the diagonal plot,
            see the source code in `KdeDiagOptions`, `HistDiagOptions` or
            `ScatterDiagOptions`.
        upper_kwargs: Additional arguments to adjust the upper diagonal plot,
            see the source code in `KdeOffDiagOptions`, `HistOffDiagOptions`,
            `ScatterOffDiagOptions`, `ContourOffDiagOptions` or `PlotOffDiagOptions`.
        lower_kwargs: Additional arguments to adjust the lower diagonal plot,
            see the source code in `KdeOffDiagOptions`, `HistOffDiagOptions`,
            `ScatterOffDiagOptions`, `ContourOffDiagOptions` or `PlotOffDiagOptions`.
        fig_kwargs: Additional arguments to adjust the overall figure,
            see the source code in `FigOptions`
        fig: matplotlib figure to plot on.
        axes: matplotlib axes corresponding to fig.
        **kwargs: Additional arguments to adjust the plot (deprecated).

    Returns: figure and axis of posterior distribution plot
    """

    upper = _prepare_upper(offdiag, upper)  # type: ignore

    plotting_styles = [
        (diag, DiagLiteral, "diag"),
        (upper, UpperLiteral, "upper"),
        (lower, LowerLiteral, "lower"),
    ]

    for plotting_style, literal, argument_name in plotting_styles:
        _validate_plotting_style(plotting_style, literal, argument_name)

    # Backwards compatibility
    if len(kwargs) > 0:
        fig, axes = _use_deprecated_plot(
            pairplot_dep,
            samples=samples,
            points=points,
            limits=limits,
            subset=subset,
            offdiag=offdiag,
            diag=diag,
            figsize=figsize,
            labels=labels,
            ticks=ticks,
            upper=upper,
            fig=fig,
            axes=axes,
            **kwargs,
        )
        return fig, axes

    samples, dim, limits = prepare_for_plot(samples, limits, points)

    # prepare figure kwargs
    fig_kwargs_filled = _prepare_fig_kwargs(fig_kwargs, samples)

    # Prepare diag
    diag_kwargs_filled, diag_func = _prepare_kwargs(
        plot=diag,  # type: ignore
        samples=samples,
        get_plot_funcs=get_diag_funcs,
        get_default_kwargs=get_default_diag_kwargs,
        plot_kwargs=diag_kwargs,
    )

    # Prepare upper
    upper_kwargs_filled, upper_func = _prepare_kwargs(
        plot=upper,  # type: ignore
        samples=samples,
        get_plot_funcs=get_offdiag_funcs,
        get_default_kwargs=get_default_offdiag_kwargs,
        plot_kwargs=upper_kwargs,
    )

    # Prepare lower
    lower_kwargs_filled, lower_func = _prepare_kwargs(
        plot=lower,  # type: ignore
        samples=samples,
        get_plot_funcs=get_offdiag_funcs,
        get_default_kwargs=get_default_offdiag_kwargs,
        plot_kwargs=lower_kwargs,
    )

    return _arrange_grid(
        diag_func,
        upper_func,
        lower_func,
        diag_kwargs_filled,
        upper_kwargs_filled,
        lower_kwargs_filled,
        samples,
        points,
        limits,
        subset,
        figsize,
        labels,
        ticks,
        fig,
        axes,
        fig_kwargs_filled,
    )

marginal_plot(samples, points=None, limits=None, subset=None, diag='hist', figsize=(10, 2), labels=None, ticks=None, diag_kwargs=None, fig_kwargs=None, fig=None, axes=None, **kwargs)

Plot samples in a row showing 1D marginals of selected dimensions.

Each of the plots can be interpreted as a 1D-marginal of the distribution that the samples were drawn from.

Parameters:

Name Type Description Default
samples Union[List[ndarray], List[Tensor], ndarray, Tensor]

Samples used to build the histogram.

required
points Optional[Union[List[ndarray], List[Tensor], ndarray, Tensor]]

List of additional points to scatter.

None
limits Optional[Union[List, Tensor]]

Array containing the plot xlim for each parameter dimension. If None, just use the min and max of the passed samples

None
subset Optional[List[int]]

List containing the dimensions to plot. E.g. subset=[1,3] will plot plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and, if they exist, the 4th, 5th and so on).

None
diag Optional[Union[List[Optional[str]], str]]

Plotting style for 1D marginals, {hist, kde cond, None}.

'hist'
figsize Optional[Tuple]

Size of the entire figure.

(10, 2)
labels Optional[List[str]]

List of strings specifying the names of the parameters.

None
ticks Optional[Union[List, Tensor]]

Position of the ticks.

None
diag_kwargs KwargsType[DiagOptions]

Additional arguments to adjust the diagonal plot, see the source code in KdeDiagOptions, HistDiagOptions or ScatterDiagOptions.

None
fig_kwargs Optional[Union[Dict, FigOptions]]

Additional arguments to adjust the overall figure, see the source code in FigOptions.

None
fig Optional[FigureBase]

matplotlib figure to plot on.

None
axes Optional[Axes]

matplotlib axes corresponding to fig.

None
**kwargs Optional[Any]

Additional arguments to adjust the plot (deprecated)

{}

Returns: figure and axis of posterior distribution plot

Source code in sbi/analysis/plot.py
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def marginal_plot(
    samples: Union[List[np.ndarray], List[torch.Tensor], np.ndarray, torch.Tensor],
    points: Optional[
        Union[List[np.ndarray], List[torch.Tensor], np.ndarray, torch.Tensor]
    ] = None,
    limits: Optional[Union[List, torch.Tensor]] = None,
    subset: Optional[List[int]] = None,
    diag: Optional[Union[List[Optional[str]], str]] = "hist",
    figsize: Optional[Tuple] = (10, 2),
    labels: Optional[List[str]] = None,
    ticks: Optional[Union[List, torch.Tensor]] = None,
    diag_kwargs: KwargsType[DiagOptions] = None,
    fig_kwargs: Optional[Union[Dict, FigOptions]] = None,
    fig: Optional[FigureBase] = None,
    axes: Optional[Axes] = None,
    **kwargs: Optional[Any],
) -> Tuple[FigureBase, Axes]:
    """
    Plot samples in a row showing 1D marginals of selected dimensions.

    Each of the plots can be interpreted as a 1D-marginal of the distribution
    that the samples were drawn from.

    Args:
        samples: Samples used to build the histogram.
        points: List of additional points to scatter.
        limits: Array containing the plot xlim for each parameter dimension. If None,
            just use the min and max of the passed samples
        subset: List containing the dimensions to plot. E.g. subset=[1,3] will plot
            plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and,
            if they exist, the 4th, 5th and so on).
        diag: Plotting style for 1D marginals, {hist, kde cond, None}.
        figsize: Size of the entire figure.
        labels: List of strings specifying the names of the parameters.
        ticks: Position of the ticks.
        diag_kwargs: Additional arguments to adjust the diagonal plot,
            see the source code in `KdeDiagOptions`, `HistDiagOptions` or
            `ScatterDiagOptions`.
        fig_kwargs: Additional arguments to adjust the overall figure,
            see the source code in `FigOptions`.
        fig: matplotlib figure to plot on.
        axes: matplotlib axes corresponding to fig.
        **kwargs: Additional arguments to adjust the plot (deprecated)
    Returns: figure and axis of posterior distribution plot
    """

    # backwards compatibility
    if len(kwargs) > 0:
        fig, axes = _use_deprecated_plot(
            marginal_plot_dep,
            samples=samples,
            points=points,
            limits=limits,
            subset=subset,
            diag=diag,
            figsize=figsize,
            labels=labels,
            ticks=ticks,
            fig=fig,
            axes=axes,
            **kwargs,
        )
        return fig, axes

    samples, _, limits = prepare_for_plot(samples, limits)

    # prepare kwargs and functions of the subplots
    diag_kwargs_filled, diag_func = _prepare_kwargs(
        plot=diag,
        samples=samples,
        get_plot_funcs=get_diag_funcs,
        get_default_kwargs=get_default_diag_kwargs,
        plot_kwargs=diag_kwargs,
    )

    # prepare fig_kwargs
    fig_kwargs_filled = _prepare_fig_kwargs(fig_kwargs, samples)

    # generate plot
    return _arrange_grid(
        diag_func,
        [None],
        [None],
        diag_kwargs_filled,
        [None],
        [None],
        samples,
        points,
        limits,
        subset,
        figsize,
        labels,
        ticks,
        fig,
        axes,
        fig_kwargs_filled,
    )

conditional_pairplot(density, condition, limits, points=None, subset=None, resolution=50, figsize=(10, 10), labels=None, ticks=None, fig=None, axes=None, **kwargs)

Plot conditional distribution given all other parameters.

The conditionals can be interpreted as slices through the density at a location given by condition.

For example: Say we have a 3D density with parameters \(\theta_0\), \(\theta_1\), \(\theta_2\) and a condition \(c\) passed by the user in the condition argument. For the plot of \(\theta_0\) on the diagonal, this will plot the conditional \(p(\theta_0 | \theta_1=c[1], \theta_2=c[2])\). For the upper diagonal of \(\theta_1\) and \(\theta_2\), it will plot \(p(\theta_1, \theta_2 | \theta_0=c[0])\). All other diagonals and upper-diagonals are built in the corresponding way.

Parameters:

Name Type Description Default
density Any

Probability density with a log_prob() method.

required
condition Tensor

Condition that all but the one/two regarded parameters are fixed to. The condition should be of shape (1, dim_theta), i.e. it could e.g. be a sample from the posterior distribution.

required
limits Union[List, Tensor]

Limits in between which each parameter will be evaluated.

required
points Optional[Union[List[ndarray], List[Tensor], ndarray, Tensor]]

Additional points to scatter.

None
subset Optional[List[int]]

List containing the dimensions to plot. E.g. subset=[1,3] will plot plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and, if they exist, the 4th, 5th and so on)

None
resolution int

Resolution of the grid at which we evaluate the pdf.

50
figsize Tuple

Size of the entire figure.

(10, 10)
labels Optional[List[str]]

List of strings specifying the names of the parameters.

None
ticks Optional[Union[List, Tensor]]

Position of the ticks.

None
points_colors

Colors of the points.

required
fig

matplotlib figure to plot on.

None
axes

matplotlib axes corresponding to fig.

None
**kwargs

Additional arguments to adjust the plot, e.g., samples_colors, points_colors and many more, see the source code in _get_default_opts() in sbi.analysis.plot for details.

{}

Returns: figure and axis of posterior distribution plot

Source code in sbi/analysis/plot.py
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
def conditional_pairplot(
    density: Any,
    condition: torch.Tensor,
    limits: Union[List, torch.Tensor],
    points: Optional[
        Union[List[np.ndarray], List[torch.Tensor], np.ndarray, torch.Tensor]
    ] = None,
    subset: Optional[List[int]] = None,
    resolution: int = 50,
    figsize: Tuple = (10, 10),
    labels: Optional[List[str]] = None,
    ticks: Optional[Union[List, torch.Tensor]] = None,
    fig=None,
    axes=None,
    **kwargs,
):
    r"""
    Plot conditional distribution given all other parameters.

    The conditionals can be interpreted as slices through the `density` at a location
    given by `condition`.

    For example:
    Say we have a 3D density with parameters $\theta_0$, $\theta_1$, $\theta_2$ and
    a condition $c$ passed by the user in the `condition` argument.
    For the plot of $\theta_0$ on the diagonal, this will plot the conditional
    $p(\theta_0 | \theta_1=c[1], \theta_2=c[2])$. For the upper
    diagonal of $\theta_1$ and $\theta_2$, it will plot
    $p(\theta_1, \theta_2 | \theta_0=c[0])$. All other diagonals and upper-diagonals
    are built in the corresponding way.

    Args:
        density: Probability density with a `log_prob()` method.
        condition: Condition that all but the one/two regarded parameters are fixed to.
            The condition should be of shape (1, dim_theta), i.e. it could e.g. be
            a sample from the posterior distribution.
        limits: Limits in between which each parameter will be evaluated.
        points: Additional points to scatter.
        subset: List containing the dimensions to plot. E.g. subset=[1,3] will plot
            plot only the 1st and 3rd dimension but will discard the 0th and 2nd (and,
            if they exist, the 4th, 5th and so on)
        resolution: Resolution of the grid at which we evaluate the `pdf`.
        figsize: Size of the entire figure.
        labels: List of strings specifying the names of the parameters.
        ticks: Position of the ticks.
        points_colors: Colors of the `points`.

        fig: matplotlib figure to plot on.
        axes: matplotlib axes corresponding to fig.
        **kwargs: Additional arguments to adjust the plot, e.g., `samples_colors`,
            `points_colors` and many more, see the source code in `_get_default_opts()`
            in `sbi.analysis.plot` for details.

    Returns: figure and axis of posterior distribution plot
    """
    device = density._device if hasattr(density, "_device") else "cpu"

    # Setting these is required because _pairplot_scaffold will check if opts['diag'] is
    # `None`. This would break if opts has no key 'diag'. Same for 'upper'.
    diag = "cond"
    offdiag = "cond"

    opts = _get_default_opts()
    # update the defaults dictionary by the current values of the variables (passed by
    # the user)
    opts = update(opts, locals())
    opts = update(opts, kwargs)
    opts["lower"] = None

    dim, limits, eps_margins = prepare_for_conditional_plot(condition, opts)
    diag_func = get_conditional_diag_func(opts, limits, eps_margins, resolution)

    def offdiag_func(row, col, **kwargs):
        p_image = (
            eval_conditional_density(
                opts["density"],
                opts["condition"].to(device),
                limits.to(device),
                row,
                col,
                resolution=resolution,
                eps_margins1=eps_margins[row],
                eps_margins2=eps_margins[col],
            )
            .to("cpu")
            .numpy()
        )
        plt.imshow(
            p_image.T,
            origin="lower",
            extent=(
                limits[col, 0].item(),
                limits[col, 1].item(),
                limits[row, 0].item(),
                limits[row, 1].item(),
            ),
            aspect="auto",
        )

    return _arrange_plots(
        diag_func, offdiag_func, dim, limits, points, opts, fig=fig, axes=axes
    )

conditional_corrcoeff(density, limits, condition, subset=None, resolution=50)

Returns the conditional correlation matrix of a distribution.

To compute the conditional distribution, we condition all but two parameters to values from condition, and then compute the Pearson correlation coefficient \(\rho\) between the remaining two parameters under the distribution density. We do so for any pair of parameters specified in subset, thus creating a matrix containing conditional correlations between any pair of parameters.

If condition is a batch of conditions, this function computes the conditional correlation matrix for each one of them and returns the mean.

Parameters:

Name Type Description Default
density Any

Probability density function with .log_prob() function.

required
limits Tensor

Limits within which to evaluate the density.

required
condition Tensor

Values to condition the density on. If a batch of conditions is passed, we compute the conditional correlation matrix for each of them and return the average conditional correlation matrix.

required
subset Optional[List[int]]

Evaluate the conditional distribution only on a subset of dimensions. If None this function uses all dimensions.

None
resolution int

Number of grid points on which the conditional distribution is evaluated. A higher value increases the accuracy of the estimated correlation but also increases the computational cost.

50

Returns: Average conditional correlation matrix of shape either (num_dim, num_dim) or (len(subset), len(subset)) if subset was specified.

Source code in sbi/analysis/conditional_density.py
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def conditional_corrcoeff(
    density: Any,
    limits: Tensor,
    condition: Tensor,
    subset: Optional[List[int]] = None,
    resolution: int = 50,
) -> Tensor:
    r"""Returns the conditional correlation matrix of a distribution.

    To compute the conditional distribution, we condition all but two parameters to
    values from `condition`, and then compute the Pearson correlation
    coefficient $\rho$ between the remaining two parameters under the distribution
    `density`. We do so for any pair of parameters specified in `subset`, thus
    creating a matrix containing conditional correlations between any pair of
    parameters.

    If `condition` is a batch of conditions, this function computes the conditional
    correlation matrix for each one of them and returns the mean.

    Args:
        density: Probability density function with `.log_prob()` function.
        limits: Limits within which to evaluate the `density`.
        condition: Values to condition the `density` on. If a batch of conditions is
            passed, we compute the conditional correlation matrix for each of them and
            return the average conditional correlation matrix.
        subset: Evaluate the conditional distribution only on a subset of dimensions.
            If `None` this function uses all dimensions.
        resolution: Number of grid points on which the conditional distribution is
            evaluated. A higher value increases the accuracy of the estimated
            correlation but also increases the computational cost.

    Returns: Average conditional correlation matrix of shape either `(num_dim, num_dim)`
    or `(len(subset), len(subset))` if `subset` was specified.
    """

    device = density._device if hasattr(density, "_device") else "cpu"

    subset_ = subset if subset is not None else range(condition.shape[1])

    correlation_matrices = []
    for cond in condition:
        correlation_matrices.append(
            torch.stack([
                compute_corrcoeff(
                    eval_conditional_density(
                        density,
                        cond.to(device),
                        limits.to(device),
                        dim1=dim1,
                        dim2=dim2,
                        resolution=resolution,
                    ),
                    limits[[dim1, dim2]].to(device),
                )
                for dim1 in subset_
                for dim2 in subset_
                if dim1 < dim2
            ])
        )

    average_correlations = torch.mean(torch.stack(correlation_matrices), dim=0)

    # `average_correlations` is still a vector containing the upper triangular entries.
    # Below, assemble them into a matrix:
    av_correlation_matrix = torch.zeros((len(subset_), len(subset_)), device=device)
    triu_indices = torch.triu_indices(
        row=len(subset_), col=len(subset_), offset=1, device=device
    )
    av_correlation_matrix[triu_indices[0], triu_indices[1]] = average_correlations

    # Make the matrix symmetric by copying upper diagonal to lower diagonal.
    av_correlation_matrix = torch.triu(av_correlation_matrix) + torch.tril(
        av_correlation_matrix.T
    )

    av_correlation_matrix.fill_diagonal_(1.0)
    return av_correlation_matrix