Feature Request: Reintroduce trainer.mmd_hypothesis_test() equivalent in bayesflow v2

[thegialeo]

Background

In version 1, the library provided a convenience method trainer.mmd_hypothesis_test, making it easy to perform an MMD hypothesis test. However, this method was removed in version 2, requiring users to manually construct and execute the test, leading to additional boilerplate code.

Proposal

I propose reintroducing an equivalent convenience function in version 2 to streamline MMD hypothesis testing. To improve modularity, the new implementation should be a standalone utility function(s) rather than a class method. If necessary for projection into the summary space, an approximator could be provided via dependency injection.

Benefits

• Improves usability: Users migrating from v1 will have an equivalent function.
• Reduces boilerplate: Simplifies the workflow for MMD hypothesis testing.
• Enhances accessibility: Makes it easier for new users to apply MMD tests without diving into low-level implementations.

Possible Implementation

BayesFlow v1 Implementation (for reference)

class Trainer:
   def mmd_hypothesis_test(
        self, observed_data, reference_data=None, num_reference_simulations=1000, num_null_samples=100, bootstrap=False
    ):
        """Performs a sampling-based hypothesis test for detecting Out-Of-Simulation (model misspecification).

        Parameters
        ----------
        observed_data : np.ndarray
           Observed data, shape (num_observed, ...)
        reference_data : np.ndarray
           Reference data representing samples from the well-specified model, shape (num_reference, ...)
        num_reference_simulations : int, default: 1000
           Number of reference simulations (M) simulated from the trainer's generative model
           if no `reference_data` are provided.
        num_null_samples : int, default: 100
           Number of draws from the MMD sampling distribution under the null hypothesis "the trainer's generative
           model is well-specified"
        bootstrap : bool, default: False
           If true, the reference data (see above) are bootstrapped for each sample from the MMD sampling distribution.
           If false, a new data set is simulated for computing each draw from the MMD sampling distribution.

        Returns
        -------
        mmd_null_samples : np.ndarray
           samples from the H0 sampling distribution ("well-specified model")
        mmd_observed : float
           summary MMD estimate for the observed data sets
        """

Proposed Utility Functions in BayesFlow v2

A new module diagnostics.metrics.maximum_mean_discrepancy.py could be introduced:

def maximum_mean_discrepancy(source_samples: NDArray, target_samples: NDArray) -> mmd

Optional additional arguments: kernel type, kwargs based on kernel choice.

Utility functions (not sure where these should go):

# Independent of approximator
def compute_mmd_hypothesis_test_from_summaries(observed_summaries: NDArray, reference_summaries: NDArray) -> mmd_observed, mmd_reference:

# Uses dependency injection for approximator: call approximator.summary_network and compute_mmd_hypothesis_test_from_summaries
def compute_mmd_hypothesis_test_from_data(observed_data: NDArray, reference_data: NDArray, approximator: Approximator) -> mmd_observed, mmd_reference:

Integration with Existing Code

• The existing diagnostics.plots.mmd_hypothesis_test function can be used to visualize the computed MMD values.
• To streamline the user workflow, an optional "workflow function" could be introduced that sequentially calls the proposed functions in a single step.

Would the maintainers be open to this approach? I’d be happy to discuss further or contribute an implementation! 🚀

[vpratz]

Hi, thanks for opening the issue, it would be great if you would supply an implementation.

A new module diagnostics.metrics.maximum_mean_discrepancy.py could be introduced

There is already a maximum_mean_discrepancy function in metrics/functional/maximum_mean_discrepancy.py, which (from the way I read the signature you proposed) does what would be needed here. So we would mainly need the mmd_hypothesis_test part that returns mmd_null_samples and mmd_observed.

I'm not sure which naming/structure would be best here. Either two separate functions as you proposed (maybe we skip the compute_, in line with our other functions), or one function which can handle both combinations.
@stefanradev93 @paul-buerkner What do you think would work best? And where would we put that function? diagnostics, diagnostics.metrics or diagnostics.misspecification (which would be new), ... would be options.

I think an integration into the BasicWorkflow would be nice, do you agree @stefanradev93 ?

[thegialeo]

Thanks for the feedback! Appreciate the heads-up about maximum_mean_discrepancy already being in the codebase. Just to get a better sense of the project structure—what’s the main difference between metrics and diagnostics/metrics?

On Implementation:

• I’m fine with either approach: two separate functions (one for summaries, one for raw data with an approximator) or a single function that handles both cases.
• If we go with two separate functions, we could also introduce a third function that automatically selects the appropriate one based on the inputs. This would improve modularity and usability, though it might add an overhead due to additional function calls compared to the one function approach.
• Totally flexible on naming and structure—just let me know the preferred function signatures, arguments, and return values, and I’ll match that. Also, where would you prefer these functions to go?

On BasicWorkflow Integration:

I don’t have a full view of the codebase, so I’d rather leave that part to someone more familiar with it to avoid introducing regressions. I also don’t want to be responsible for breaking BasicWorkflow for anyone currently using it.

Let me know what you think! 🚀

[vpratz]

Sounds good to me :)

what’s the main difference between metrics and diagnostics/metrics

I have not written this code, but I think the main distinction is that functions in metrics provide the bare-bones implementation, taking in arrays and producing the desired output. Functions in diagnostics/metrics may contain functionality to simplify directly calling them with output from the approximator, without manually selecting the appropriate arrays.

Regarding the final structure, I would wait on feedback from the others, so we can either wait for their comments or, if you already want to start, we adapt the interface later (if needed). In the latter case, I'd suggest you implement two functions, mmd_hypothesis_test (from data), and mmd_hypothesis_test_from_summaries, and put them in a new file diagnostics/metrics/mmd_hypothesis_test.py.

[LarsKue]

I agree with what @vpratz has said so far. Thank you for the contribution @thegialeo!

[thegialeo]

Thanks for the feedback! I’ll go ahead and create a PR with a first draft implementation with the specifications proposed by @vpratz. Further details can be discussed and refined in the PR.

[vpratz]

Closing this with #384, let's discuss the user-facing part of this issue in #428.