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[str]], str]]

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

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

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

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

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 Optional[Union[List[Optional[Dict]], Dict]]

Additional arguments to adjust the diagonal plot, see the source code in _get_default_diag_kwarg()

None
upper_kwargs Optional[Union[List[Optional[Dict]], Dict]]

Additional arguments to adjust the upper diagonal plot, see the source code in _get_default_offdiag_kwarg()

None
lower_kwargs Optional[Union[List[Optional[Dict]], Dict]]

Additional arguments to adjust the lower diagonal plot, see the source code in _get_default_offdiag_kwarg()

None
fig_kwargs Optional[Dict]

Additional arguments to adjust the overall figure, see the source code in _get_default_fig_kwargs()

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
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
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[str]], str]] = "hist",
    lower: Optional[Union[List[Optional[str]], str]] = None,
    diag: Optional[Union[List[Optional[str]], str]] = "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: Optional[Union[List[Optional[Dict]], Dict]] = None,
    upper_kwargs: Optional[Union[List[Optional[Dict]], Dict]] = None,
    lower_kwargs: Optional[Union[List[Optional[Dict]], Dict]] = None,
    fig_kwargs: Optional[Dict] = 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 `_get_default_diag_kwarg()`
        upper_kwargs: Additional arguments to adjust the upper diagonal plot,
            see the source code in `_get_default_offdiag_kwarg()`
        lower_kwargs: Additional arguments to adjust the lower diagonal plot,
            see the source code in `_get_default_offdiag_kwarg()`
        fig_kwargs: Additional arguments to adjust the overall figure,
            see the source code in `_get_default_fig_kwargs()`
        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:
        warn(
            f"you passed deprecated arguments **kwargs: {[key for key in kwargs]}, use "
            "fig_kwargs instead. We continue calling the deprecated pairplot function",
            DeprecationWarning,
            stacklevel=2,
        )
        fig, axes = pairplot_dep(
            samples,
            points,
            limits,
            subset,
            offdiag,
            diag,
            figsize,
            labels,
            ticks,
            upper,
            fig,
            axes,
            **kwargs,
        )
        return fig, axes

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

    # prepate figure kwargs
    fig_kwargs_filled = _get_default_fig_kwargs()
    # update the defaults dictionary with user provided values
    fig_kwargs_filled = _update(fig_kwargs_filled, fig_kwargs)

    # checks.
    if fig_kwargs_filled["legend"]:
        assert len(fig_kwargs_filled["samples_labels"]) >= len(
            samples
        ), "Provide at least as many labels as samples."
    if offdiag is not None:
        warn("offdiag is deprecated, use upper or lower instead.", stacklevel=2)
        upper = offdiag

    # Prepare diag
    diag_list = to_list_string(diag, len(samples))
    diag_kwargs_list = to_list_kwargs(diag_kwargs, len(samples))
    diag_func = get_diag_funcs(diag_list)
    diag_kwargs_filled = []
    for i, (diag_i, diag_kwargs_i) in enumerate(zip(diag_list, diag_kwargs_list)):
        diag_kwarg_filled_i = _get_default_diag_kwargs(diag_i, i)
        # update the defaults dictionary with user provided values
        diag_kwarg_filled_i = _update(diag_kwarg_filled_i, diag_kwargs_i)
        diag_kwargs_filled.append(diag_kwarg_filled_i)

    # Prepare upper
    upper_list = to_list_string(upper, len(samples))
    upper_kwargs_list = to_list_kwargs(upper_kwargs, len(samples))
    upper_func = get_offdiag_funcs(upper_list)
    upper_kwargs_filled = []
    for i, (upper_i, upper_kwargs_i) in enumerate(zip(upper_list, upper_kwargs_list)):
        upper_kwarg_filled_i = _get_default_offdiag_kwargs(upper_i, i)
        # update the defaults dictionary with user provided values
        upper_kwarg_filled_i = _update(upper_kwarg_filled_i, upper_kwargs_i)
        upper_kwargs_filled.append(upper_kwarg_filled_i)

    # Prepare lower
    lower_list = to_list_string(lower, len(samples))
    lower_kwargs_list = to_list_kwargs(lower_kwargs, len(samples))
    lower_func = get_offdiag_funcs(lower_list)
    lower_kwargs_filled = []
    for i, (lower_i, lower_kwargs_i) in enumerate(zip(lower_list, lower_kwargs_list)):
        lower_kwarg_filled_i = _get_default_offdiag_kwargs(lower_i, i)
        # update the defaults dictionary with user provided values
        lower_kwarg_filled_i = _update(lower_kwarg_filled_i, lower_kwargs_i)
        lower_kwargs_filled.append(lower_kwarg_filled_i)

    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 Optional[Union[List[Optional[Dict]], Dict]]

Additional arguments to adjust the diagonal plot, see the source code in _get_default_diag_kwarg()

None
fig_kwargs Optional[Dict]

Additional arguments to adjust the overall figure, see the source code in _get_default_fig_kwargs()

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
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
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: Optional[Union[List[Optional[Dict]], Dict]] = None,
    fig_kwargs: Optional[Dict] = 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 `_get_default_diag_kwarg()`
        fig_kwargs: Additional arguments to adjust the overall figure,
            see the source code in `_get_default_fig_kwargs()`
        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:
        warn(
            "**kwargs are deprecated, use fig_kwargs instead. "
            "calling the to be deprecated marginal_plot function",
            DeprecationWarning,
            stacklevel=2,
        )
        fig, axes = marginal_plot_dep(
            samples,
            points,
            limits,
            subset,
            diag,
            figsize,
            labels,
            ticks,
            fig,
            axes,
            **kwargs,
        )
        return fig, axes

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

    # prepare kwargs and functions of the subplots
    diag_list = to_list_string(diag, len(samples))
    diag_kwargs_list = to_list_kwargs(diag_kwargs, len(samples))
    diag_func = get_diag_funcs(diag_list)
    diag_kwargs_filled = []
    for i, (diag_i, diag_kwargs_i) in enumerate(zip(diag_list, diag_kwargs_list)):
        diag_kwarg_filled_i = _get_default_diag_kwargs(diag_i, i)
        diag_kwarg_filled_i = _update(diag_kwarg_filled_i, diag_kwargs_i)
        diag_kwargs_filled.append(diag_kwarg_filled_i)

    # prepare fig_kwargs
    fig_kwargs_filled = _get_default_fig_kwargs()
    fig_kwargs_filled = _update(fig_kwargs_filled, fig_kwargs)

    # 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
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
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
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