Pair plots with specific chains, coords, and divergences

Questions version agnostic
1

Hi everyone, is there a way to generate a pair plot where (1) the divergences are highlighted for a specific chain and a specific coordinate? I can generate plots with two of those features. But not all three :frowning:

Here’s what I want to work:

  az.plot_pair(trace,
      divergences=True,
      coords={
          'channels':3,
          'chain':0
      },
      var_names=[
          "retention",
          "saturation_point",
          "costumer_aquisition"]);

Which returns the error message: IndexError: boolean index did not match indexed array along dimension 0; dimension is 2000 but corresponding boolean dimension is 8000.

It’s possible to plot the divergences and one coordinate:

az.plot_pair(trace,
    divergences=True,
    coords={
        'channels':3
    },
    var_names=[
        "retention",
        "saturation_point",
        "costumer_aquisition"]);

Or to plot one coordinate and one chain:

az.plot_pair(trace.posterior.sel(chain=0),
    coords={
        'channels':3
    },
    var_names=[
        "retention",
        "saturation_point",
        "costumer_aquisition"]);

Or even plot one chain using the coord keyword:

az.plot_pair(trace,
    coords={
        'chain':0
    },
    var_names=[
        "retention",
        "saturation_point",
        "costumer_aquisition"]);
2

@OriolAbril

3

Use .sel directly on the inferencedata instead of the coords argument. I haven’t checked the source but I guess nobody expected this usecase and consequently coords are only applied to the posterior group (or whatever group the user choses) not the sample_stats one.

az.plot_pair(trace.sel(chain=0, channels=3),
    var_names=[
        "retention",
        "saturation_point",
        "costumer_aquisition"]);

Note the inferencedata.sel() as opposed to inferencedata.posterior.sel() so that subsetting is applied to all group where it is possible, and all groups are kept. inferencedata.posterior.sel() applies the subsetting to the posterior group only, and returns only that group; a single dataset is a valid input to plot_pair but there is then no sample_stats information so no divergences can be plotted at all.

2 Likes
4

Ah thanks so much Oriol, I didn’t realize you could .sel on the top level of inference data. Super convenient that you can though!