Merge branch 'main' into white_input

README.md

@@ -1,6 +1,6 @@
-# PQM
+# PQMass: Probabilistic Assessment of the Quality of Generative Models using Probability Mass Estimation
 
-Implementation of the PQMass two sample test from Lemos et al. 2024
+Implementation of the PQMass two sample test from Lemos et al. 2024 [here](https://arxiv.org/abs/2402.04355)
 
 ## Install
 
@@ -21,16 +21,33 @@ import numpy as np
 x_sample = np.random.normal(size = (500, 10))
 y_sample = np.random.normal(size = (400, 10))
 
+# To get pvalues from PQMass
 pvalues = pqm_pvalue(x_sample, y_sample, num_refs = 100, bootstrap = 50)
-
 print(np.mean(pvalues), np.std(pvalues))
+
+# To get chi^2 from PQMass
+chi2_stat, dof = pqm_chi2(x_sample, y_sample, num_refs = 100, bootstrap = 50)
+print(np.mean(chi2_stat), np.std(chi2_stat))
+print(np.unqiue(dof)) # This should be the same as num_refs - 1, if it is not, we suggest you use pqm_pvalue
 ```
 
-If your two samples are drawn from the same distribution then the pvalue should
+If your two samples are drawn from the same distribution, then the p-value should
 be drawn from the random uniform(0,1) distribution. This means that if you get a
-very small value (i.e. 1e-6) then you have failed the null hypothesis test and
+very small value (i.e., 1e-6), then you have failed the null hypothesis test, and
 the two samples are not drawn from the same distribution.
 
+For the chi^2 metric, given your two sets of samples, if they come from the same
+distribution, the histogram of your chi² values should follow the chi² distribution. 
+The peak of this distribution will be at DoF - 2, and the standard deviation will 
+be √(2 * DoF). If your histogram shifts to the right of the expected chi² distribution, 
+it suggests that the samples are out of distribution. Conversely, if the histogram shifts 
+to the left, it indicates potential duplication or memorization (particularly relevant 
+for generative models).
+
+Note that the chi^2 metric faces limitations if you have a few samples. A solution could
+be to use bootstrapping. Another such solution is to pqm_pvalue. We leave it to the user to 
+identify the best solution for their problem.
+
 ## Developing
 
 If you're a developer then:

src/pqm/__init__.py

@@ -1,10 +1,11 @@
-from .pqm import pqm_pvalue
+from .pqm import pqm_pvalue, pqm_chi2
 from .test_gaussian import test
 from ._version import __version__
 
 
 __all__ = (
     "pqm_pvalue",
+    "pqm_chi2",
     "test",
     "__version__",
 )

src/pqm/pqm.py

@@ -4,42 +4,29 @@
 from scipy.stats import chi2_contingency
 from scipy.spatial import KDTree
 
-__all__ = "pqm_pvalue"
+__all__ = ("pqm_chi2", "pqm_pvalue")
 
 
-def pqm_pvalue(
-    x_samples: np.ndarray,
-    y_samples: np.ndarray,
-    num_refs: int = 100,
-    bootstrap: Optional[int] = None,
-    whiten: bool = False,
-):
+def _pqm_test(x_samples: np.ndarray, y_samples: np.ndarray, num_refs: int, whiten: bool):
     """
-    Perform the PQM test of the null hypothesis that `x_samples` and `y_samples` are drawn form the same distribution.
+    Helper function to perform the PQM test and return the results from chi2_contingency.
 
     Parameters
     ----------
     x_samples : np.ndarray
-        Samples from the first distribution. Must have shape (N, *D) N is the number of x samples, and D is the dimensionality of the samples.
+        Samples from the first distribution, test samples.
     y_samples : np.ndarray
-        Samples from the second distribution. Must have shape (M, *D) M is the number of y samples, and D is the dimensionality of the samples.
+        Samples from the second distribution, reference samples.
     num_refs : int
-        Number of reference samples to use. Note that these will be drawn from y_samples, and then removed from the y_samples array.
-    bootstrap : Optional[int]
-        Number of bootstrap iterations to perform. No bootstrap if None (default).
+        Number of reference samples to use.
     whiten : bool
         If True, whiten the samples by subtracting the mean and dividing by the standard deviation.
 
     Returns
     -------
-    float
-        pvalue. Null hypothesis that both samples are drawn from the same distribution.
+    tuple
+        Results from scipy.stats.chi2_contingency function.
     """
-    if bootstrap is not None:
-        return list(
-            pqm_pvalue(x_samples.copy(), y_samples.copy(), num_refs=num_refs, whiten=whiten)
-            for _ in range(bootstrap)
-        )
     if len(y_samples) < num_refs:
         raise ValueError(
             "Number of reference samples must be less than the number of true samples."
@@ -72,5 +59,78 @@ def pqm_pvalue(
     C = (counts_x > 0) | (counts_y > 0)
     counts_x, counts_y = counts_x[C], counts_y[C]
 
-    _, pvalue, _, _ = chi2_contingency(np.array([counts_x, counts_y]))
+    return chi2_contingency(np.array([counts_x, counts_y]))
+
+
+def pqm_pvalue(
+    x_samples: np.ndarray,
+    y_samples: np.ndarray,
+    num_refs: int = 100,
+    bootstrap: Optional[int] = None,
+    whiten: bool = False,
+):
+    """
+    Perform the PQM test of the null hypothesis that `x_samples` and `y_samples` are drawn form the same distribution.
+
+    Parameters
+    ----------
+    x_samples : np.ndarray
+        Samples from the first distribution, test samples. Must have shape (N, *D) N is the number of x samples, and D is the dimensionality of the samples.
+    y_samples : np.ndarray
+        Samples from the second distribution, reference samples. Must have shape (M, *D) M is the number of y samples, and D is the dimensionality of the samples.
+    num_refs : int
+        Number of reference samples to use. Note that these will be drawn from y_samples, and then removed from the y_samples array.
+    bootstrap : Optional[int]
+        Number of bootstrap iterations to perform. No bootstrap if None (default).
+    whiten : bool
+        If True, whiten the samples by subtracting the mean and dividing by the standard deviation.
+
+    Returns
+    -------
+    float or list
+        pvalue(s). Null hypothesis that both samples are drawn from the same distribution.
+    """
+    if bootstrap is not None:
+        return [
+            pqm_pvalue(x_samples, y_samples, num_refs=num_refs, whiten=whiten)
+            for _ in range(bootstrap)
+        ]
+    _, pvalue, _, _ = _pqm_test(x_samples, y_samples, num_refs, whiten)
     return pvalue
+
+
+def pqm_chi2(
+    x_samples: np.ndarray,
+    y_samples: np.ndarray,
+    num_refs: int = 100,
+    bootstrap: Optional[int] = None,
+    whiten: bool = False,
+):
+    """
+    Perform the PQM test of the null hypothesis that `x_samples` and `y_samples` are drawn form the same distribution.
+
+    Parameters
+    ----------
+    x_samples : np.ndarray
+        Samples from the first distribution, test samples. Must have shape (N, *D) N is the number of x samples, and D is the dimensionality of the samples.
+    y_samples : np.ndarray
+        Samples from the second distribution, reference samples. Must have shape (M, *D) M is the number of y samples, and D is the dimensionality of the samples.
+    num_refs : int
+        Number of reference samples to use. Note that these will be drawn from y_samples, and then removed from the y_samples array.
+    bootstrap : Optional[int]
+        Number of bootstrap iterations to perform. No bootstrap if None (default).
+    whiten : bool
+        If True, whiten the samples by subtracting the mean and dividing by the standard deviation.
+
+    Returns
+    -------
+    float or list
+        chi2 statistic(s) and degree(s) of freedom.
+    """
+    if bootstrap is not None:
+        return [
+            pqm_chi2(x_samples, y_samples, num_refs=num_refs, whiten=whiten)
+            for _ in range(bootstrap)
+        ]
+    chi2_stat, _, dof, _ = _pqm_test(x_samples, y_samples, num_refs, whiten)
+    return chi2_stat, dof

src/pqm/test_gaussian.py

@@ -1,5 +1,5 @@
 import numpy as np
-from .pqm import pqm_pvalue
+from .pqm import pqm_pvalue, pqm_chi2
 
 
 def test():