Skip to content

feat(plot): add exclude_keys to filter trace plots (closes #730)#1074

Open
christianescamilla15-cell wants to merge 1 commit intobilby-dev:mainfrom
christianescamilla15-cell:feat/trace-plot-filter-keys
Open

feat(plot): add exclude_keys to filter trace plots (closes #730)#1074
christianescamilla15-cell wants to merge 1 commit intobilby-dev:mainfrom
christianescamilla15-cell:feat/trace-plot-filter-keys

Conversation

@christianescamilla15-cell
Copy link
Copy Markdown

Summary

Addresses #730 by adding an opt-in way to filter nuisance parameters out of the checkpoint trace plots produced by `bilby_mcmc`. The motivation in the issue is that with ~20 calibration parameters per detector, trace plots can easily hit 60+ axes and blow up memory on long runs.

Closes #730.

Design

Rather than hard-removing calibration parameters (which @ColmTalbot correctly pointed out can still be diagnostically useful), this is opt-in with glob-style patterns, matching the spirit of arviz's `var_names` / `filter_vars` selection.

```python
bilby.run_sampler(
sampler="bilby_mcmc",
plot_exclude_keys=["recalib_*"],
...
)
```

Changes

`bilby/bilby_mcmc/chain.py`

  • Added `exclude_keys` parameter to `Chain.plot()`
  • Uses `fnmatch.fnmatchcase` for glob-style pattern matching
  • Filters `self.keys` before computing the figure layout (so the figure height shrinks with the filtered count)
  • Applied to both trace axes and histogram axes
  • Added comprehensive docstring with the full parameter list

`bilby/bilby_mcmc/sampler.py`

  • Added `plot_exclude_keys` kwarg to `Bilby_MCMC.init`
  • Stored as `self.plot_exclude_keys`
  • Threaded through `print_long_progress` → `plot_progress` → `chain.plot`
  • `plot_progress` signature extended with `exclude_keys=None`

Backwards compatibility

  • Default behaviour is unchanged (`exclude_keys=None` plots everything)
  • Existing users see identical output unless they opt in

Future work (not in this PR)

The dynesty sampler's trace plotting is generated via dynesty's own machinery and would need a separate adapter to accept `exclude_keys`. Happy to follow up in a separate PR once this approach is approved.

Test plan

  • `py_compile` passes for both modified files
  • `fnmatch` patterns work for typical cases (`recalib_*`, `spin_?`)
  • Full test suite in CI (reviewer to confirm)

References

)

Issue bilby-dev#730 reports that bilby_mcmc / dynesty checkpoint trace plots
with ~20 * Ndet axes can balloon in size and cause memory issues on
long runs, particularly for runs with many calibration parameters
(e.g. recalib_* nodes).

Adds an exclude_keys parameter to Chain.plot() accepting glob-style
patterns via fnmatch. Keys matching any pattern are omitted from
both the trace and histogram sub-plots, and the figure height is
computed from the filtered count.

The parameter is threaded through from the Bilby_MCMC sampler as
plot_exclude_keys, so users can configure it at the run level:

    bilby.run_sampler(
        sampler="bilby_mcmc",
        plot_exclude_keys=["recalib_*"],
        ...
    )

As @ColmTalbot noted in the discussion, calibration traces can still
be diagnostically useful. This is opt-in — default behaviour is
unchanged.

Inspired by arviz's var_names / filter_vars selection for plot_trace.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Remove calibration plots from trace (bilby_mcmc, dynesty)

1 participant