renate.shift.mmd_helpers module

renate.shift.mmd_helpers.mmd_gram(K, z, num_permutations=0)

Maximum mean discrepancy based on a precomputed kernel-Gram matrix.

This computes the test statistic and (optionally) p-value to conduct an MMD two-sample test to decide whether two sets are generated by the same distribution. The inputs are passed implicitly in the form of a kernel Gram matrix, evaluated across the union of both sets, and a binary vector indicating the assignments of data points to the two sets. I.e., a value of z[i] = 0 indicates that the i-th data point belongs to set zero. Optionally, a permutation test is carried out and a p-value is returned alongside the raw test statistic.

MMD tests have been proposed by

[1] Gretton, A., et al. A kernel two-sample test. JMLR (2012).

Parameters:

• K (Tensor) – A tensor containing a kernel-Gram matrix (size (n, n)).

• z (Tensor) – A binary vector of length n indicating the partition.

• num_permutations (int) – If this is a positive number, a permutation test will be carried out and an approximate p-value will be returned.

Return type:

Tuple[Tensor, Optional[Tensor]]

Returns:

A tuple (t, p) of two scalar floats, where t is the value of the MMD test statistic and p is the p-value (or None if num_permutations=0).

renate.shift.mmd_helpers.mmd(X0, X1, kernel, num_permutations=0)

Compute MMD between two samples.

Optionally, return an estimated p-value based on a permutation test.

Parameters:

• X0 (Tensor) – First sample, shape (n0, dx).

• X1 (Tensor) – Second sample, shape (n1, dx).

• kernel (Kernel) – Kernel function to use.

• num_permutations (int) – If this is a positive number, a permutation test will be carried out and a p-value will be returned.

Return type:

Tuple[Tensor, Optional[Tensor]]

Returns:

A tuple (t, p) where t is the value of the MMD test statistic and p is the p-value (or None if num_permutations=0).