Chunking cleanup: remove chunk_iter, auto-derive iteration count from chunk_size

.github/workflows/cd.yml

@@ -51,7 +51,7 @@ jobs:
       - uses: actions/setup-python@v5
         name: Install Python
         with:
-          python-version: "3.10"
+          python-version: "3.13"
       - uses: actions/download-artifact@v4
         with:
           name: artifact

.github/workflows/ci.yml

@@ -22,7 +22,7 @@ jobs:
     strategy:
       fail-fast: true
       matrix:
-        python-version: ["3.10", "3.11", "3.12"]
+        python-version: ["3.11", "3.12", "3.13"]
         os: [ubuntu-latest, windows-latest, macOS-latest]
 
     steps:
@@ -51,15 +51,15 @@ jobs:
           python -m pip install --upgrade pip
           pip install pytest pytest-cov wheel
 
-      - name: Install torch # no torch on 3.11 to test no-torch scenario
-        if: ${{ matrix.python-version != '3.11' }}
+      - name: Install torch # no torch on 3.13 to test no-torch scenario
+        if: ${{ matrix.python-version != '3.13' }}
         run: |
           pip install torch
 
       # We only want to install this on one run, because otherwise we'll have
       # duplicate annotations.
       - name: Install error reporter
-        if: ${{ matrix.python-version == '3.10' }}
+        if: ${{ matrix.python-version == '3.13' }}
         run: |
           python -m pip install pytest-github-actions-annotate-failures
 

README.md

@@ -275,7 +275,6 @@ def pted(
     metric: Union[str, float] = "euclidean",
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     two_tailed: bool = True,
     prog_bar: bool = False,
 ) -> Union[float, tuple[float, np.ndarray, float]]:
@@ -286,8 +285,7 @@ def pted(
 * **permutations** *(int)*: number of permutations to run. This determines how accurately the p-value is computed.
 * **metric** *(Union[str, float])*: distance metric to use. See scipy.spatial.distance.cdist for the list of available metrics with numpy. See torch.cdist when using PyTorch, note that the metric is passed as the "p" for torch.cdist and therefore is a float from 0 to inf. When using JAX arrays, the metric is passed as the "ord" for jnp.linalg.norm and therefore is also a float.
 * **return_all** *(bool)*: if True, return the test statistic and the permuted statistics with the p-value. If False, just return the p-value. bool (default: False)
-* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. This is useful for large datasets. The chunk size is the number of samples to use for each chunk. If None, use the full dataset.
-* **chunk_iter** *(Optional[int])*: The chunk iter is the number of iterations to use with the given chunk size.
+* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. The chunk size is the number of samples per chunk. The number of chunks is determined automatically as `max(len(x), len(y)) // chunk_size`, iterating over the larger dataset once and cycling through the smaller one. If `chunk_size >= len(x)` and `chunk_size >= len(y)`, PTED falls back to the full (non-chunked) computation. If None, use the full dataset.
 * **two_tailed** *(bool)*: if True, compute a two-tailed p-value. This is useful if you want to reject the null hypothesis when x and y are either too similar or too different. If False, only checks for dissimilarity but is more sensitive. Default is True.
 * **prog_bar** *(bool)*: if True, show a progress bar to track the progress of permutation tests. Default is False.
 
@@ -302,7 +300,6 @@ def pted_coverage_test(
     warn_confidence: Optional[float] = 1e-3,
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     sbc_histogram: Optional[str] = None,
     sbc_bins: Optional[int] = None,
     pit_plot: Optional[str] = None,
@@ -316,8 +313,7 @@ def pted_coverage_test(
 * **permutations** *(int)*: number of permutations to run. This determines how accurately the p-value is computed.
 * **metric** *(Union[str, float])*: distance metric to use. See scipy.spatial.distance.cdist for the list of available metrics with numpy. See torch.cdist when using PyTorch, note that the metric is passed as the "p" for torch.cdist and therefore is a float from 0 to inf. When using JAX arrays, the metric is passed as the "ord" for jnp.linalg.norm and therefore is also a float.
 * **return_all** *(bool)*: if True, return the test statistic and the permuted statistics with the p-value. If False, just return the p-value. bool (default: False)
-* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. This is useful for large datasets. The chunk size is the number of samples to use for each chunk. If None, use the full dataset.
-* **chunk_iter** *(Optional[int])*: The chunk iter is the number of iterations to use with the given chunk size.
+* **chunk_size** *(Optional[int])*: if not None, use chunked energy distance estimation. The chunk size is the number of samples per chunk. The number of chunks is determined automatically as `max(len(x), len(y)) // chunk_size`, iterating over the larger dataset once and cycling through the smaller one. If None, use the full dataset.
 * **sbc_histogram** *(Optional[str])*: If given, the path/filename to save a Simulation-Based-Calibration histogram.
 * **sbc_bins** *(Optional[int])*: If given, force the histogram to have the provided number of bins. Otherwise, select an appropriate size: ~sqrt(N).
 * **pit_plot** *(Optional[str])*: If given, the path/filename to save a Probability Integral Transform (PIT) plot of the per-simulation p-values against the expected uniform distribution, with a shaded KS confidence band.
@@ -361,16 +357,20 @@ If a GPU isn't enough to get PTED running fast enough for you, or if you are
 running into memory limitations, there are still options! We can use an
 approximation of the energy distance, in this case the test is still exact but
 less sensitive than it would be otherwise. We can approximate the energy
-distance by taking random subsamples (chunks) of the full dataset, computing the
-energy distance, then averaging. Just set the `chunk_size` parameter for the
-number of samples you can manage at once and set the `chunk_iter` for the number
-of trials you want in the average. The larger these numbers are, the closer the
-estimate will be to the true energy distance, but it will take more compute.
-This lets you decide how to trade off compute for sensitivity.
-
-Note that the computational complexity for standard PTED goes as 
-`O((n_samp_x + n_samp_y)^2)` while the chunked version goes as 
-`O(chunk_iter * (2 * chunk_size)^2)` so plan your chunking accordingly.
+distance by iterating through sequential chunks of the full dataset, computing
+the energy distance on each chunk, then averaging. Just set the `chunk_size`
+parameter for the number of samples you can manage at once. The number of
+iterations is determined automatically as `max(len(x), len(y)) // chunk_size`,
+iterating over the larger dataset once and cycling through the smaller one if
+their sizes differ. The larger the chunk size, the closer the estimate will be
+to the true energy distance, but it will take more compute.
+
+Note that the computational complexity for standard PTED goes as `O((n_samp_x +
+n_samp_y)^2)` while the chunked version goes as `O(n_iter * (2 * chunk_size)^2)`
+where `n_iter = max(n_samp_x, n_samp_y) // chunk_size`, so plan your chunking
+accordingly. For a given chunk size, the computational complexity of PTED grows
+linearly with dataset size, much like other large scale (machine learning
+oriented) two sample tests.
 
 Example:
 ```python
@@ -380,7 +380,7 @@ import numpy as np
 x = np.random.normal(size = (500, 10)) # (n_samples_x, n_dimensions)
 y = np.random.normal(size = (400, 10)) # (n_samples_y, n_dimensions)
 
-p_value = pted(x, y, chunk_size = 50, chunk_iter = 100)
+p_value = pted(x, y, chunk_size = 50)
 print(f"p-value: {p_value:.3f}") # expect uniform random from 0-1
 ```
 

pyproject.toml

@@ -13,7 +13,7 @@ authors = [
 ]
 description = "Implementation of a Permutation Test using the Energy Distance for two sample tests and posterior coverage tests"
 readme = "README.md"
-requires-python = ">=3.9"
+requires-python = ">=3.11"
 license = {file = "LICENSE"}
 keywords = [
         "statistics",
@@ -41,14 +41,14 @@ dev = [
     "pytest-cov>=4.1,<5",
     "pytest-mock>=3.12,<4",
     "torch>=2.0,<3",
-    "jax>=0.4,<1",
+    "jax>=0.7,<1",
     "matplotlib",
 ]
 torch = [
     "torch>=2.0,<3",
 ]
 jax = [
-    "jax>=0.4,<1",
+    "jax>=0.7,<1",
 ]
 
 [tool.hatch.metadata.hooks.requirements_txt]

src/pted/pted.py

@@ -27,7 +27,6 @@ def pted(
     metric: Union[str, float] = "euclidean",
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     two_tailed: bool = True,
     prog_bar: bool = False,
 ) -> Union[float, tuple[float, np.ndarray, float]]:
@@ -92,10 +91,10 @@ def pted(
             bool (default: False)
         chunk_size (Optional[int]): if not None, use chunked energy distance
             estimation. This is useful for large datasets. The chunk size is the
-            number of samples to use for each chunk. If None, use the full
-            dataset.
-        chunk_iter (Optional[int]): The chunk iter is the number of iterations
-            to use with the given chunk size.
+            number of samples per chunk. The number of chunks is determined
+            automatically as ``max(len(x), len(y)) // chunk_size``, iterating
+            over the larger dataset once and cycling through the smaller one.
+            If None, use the full dataset.
         two_tailed (bool): if True, compute a two-tailed p-value. This is useful
             if you want to reject the null hypothesis when x and y are either
             too similar or too different. Default is True.
@@ -109,24 +108,25 @@ def pted(
         samples in x and y, D is the number of dimensions, and P is the number
         of permutations. For large datasets this can get unwieldy, so chunking
         is recommended. For chunking, the energy distance will be estimated at
-        each iteration rather than fully computed. To estimate the energy
-        distance, we take `chunk_size` sub-samples from x and y, and compute the
-        energy distance on those sub-samples. This is repeated `chunk_iter`
-        times, and the average is taken. This is a trade-off between speed and
-        accuracy. The larger the chunk size and larger chunk_iter, the more
-        accurate the estimate, but the slower the computation. PTED remains an
-        exact p-value test even when chunking, it simply becomes less sensitive
-        to the difference between x and y. The chunked pted test has time
+        each iteration rather than fully computed. The dataset is divided into
+        sequential chunks of size `chunk_size`; the number of chunks (iterations)
+        is ``max(len(x), len(y)) // chunk_size``, iterating over the larger
+        dataset once and cycling through the smaller one. The average energy
+        distance over all chunks is the final estimate. PTED remains an exact
+        p-value test even when chunking, it simply becomes less sensitive to
+        the difference between x and y. The chunked pted test has time
         complexity O(c^2 * I * D * P), where c is the chunk size, I is the
         number of iterations, D is the number of dimensions, and P is the number
         of permutations. For chunking to be worth it you should have c^2 * I << n^2.
     """
     assert type(x) == type(y), f"x and y must be of the same type, not {type(x)} and {type(y)}"
     assert len(x.shape) >= 2, f"x must be at least 2D, not {x.shape}"
     assert len(y.shape) >= 2, f"y must be at least 2D, not {y.shape}"
-    assert (chunk_size is not None) is (
-        chunk_iter is not None
-    ), "chunk_size and chunk_iter must both be provided for chunked PTED test"
+    if chunk_size is not None:
+        assert chunk_size > 0, "chunk_size must be > 0"
+        # If chunk_size covers both full datasets, chunking adds no benefit
+        if chunk_size >= len(x) and chunk_size >= len(y):
+            chunk_size = None
     assert (
         x.shape[1:] == y.shape[1:]
     ), f"x and y samples must have the same shape (past first dim), not {x.shape} and {y.shape}"
@@ -142,7 +142,6 @@ def pted(
             permutations=permutations,
             metric=metric,
             chunk_size=int(chunk_size),
-            chunk_iter=int(chunk_iter),
             prog_bar=prog_bar,
         )
     elif is_torch_tensor(x):
@@ -156,7 +155,6 @@ def pted(
             permutations=permutations,
             metric=metric,
             chunk_size=int(chunk_size),
-            chunk_iter=int(chunk_iter),
             prog_bar=prog_bar,
         )
     elif is_jax_array(x):
@@ -168,7 +166,6 @@ def pted(
             permutations=permutations,
             metric=metric,
             chunk_size=int(chunk_size),
-            chunk_iter=int(chunk_iter),
             prog_bar=prog_bar,
         )
     else:
@@ -200,7 +197,6 @@ def pted_coverage_test(
     warn_confidence: Optional[float] = 1e-3,
     return_all: bool = False,
     chunk_size: Optional[int] = None,
-    chunk_iter: Optional[int] = None,
     sbc_histogram: Optional[str] = None,
     sbc_bins: Optional[int] = None,
     pit_plot: Optional[str] = None,
@@ -269,10 +265,10 @@ def pted_coverage_test(
             (default: False)
         chunk_size (Optional[int]): If not None, use chunked energy distance
             estimation. This is useful for large datasets. The chunk size is the
-            number of samples to use for each chunk. If None, use the full
-            dataset.
-        chunk_iter (Optional[int]): The chunk iter is the number of iterations
-            to use with the given chunk size.
+            number of samples per chunk. The number of chunks is determined
+            automatically as ``max(len(x), len(y)) // chunk_size``, iterating
+            over the larger dataset once and cycling through the smaller one.
+            If None, use the full dataset.
         sbc_histogram (Optional[str]): If given, the path/filename to save a
             Simulation-Based-Calibration histogram.
         sbc_bins (Optional[int]): If given, force the histogram to have the provided
@@ -295,14 +291,13 @@ def pted_coverage_test(
         samples in x and y, D is the number of dimensions, and P is the number
         of permutations. For large datasets this can get unwieldy, so chunking
         is recommended. For chunking, the energy distance will be estimated at
-        each iteration rather than fully computed. To estimate the energy
-        distance, we take `chunk_size` sub-samples from x and y, and compute the
-        energy distance on those sub-samples. This is repeated `chunk_iter`
-        times, and the average is taken. This is a trade-off between speed and
-        accuracy. The larger the chunk size and larger chunk_iter, the more
-        accurate the estimate, but the slower the computation. PTED remains an
-        exact p-value test even when chunking, it simply becomes less sensitive
-        to the difference between x and y. The chunked pted test has time
+        each iteration rather than fully computed. The dataset is divided into
+        sequential chunks of size `chunk_size`; the number of chunks (iterations)
+        is ``max(len(x), len(y)) // chunk_size``, iterating over the larger
+        dataset once and cycling through the smaller one. The average energy
+        distance over all chunks is the final estimate. PTED remains an exact
+        p-value test even when chunking, it simply becomes less sensitive to
+        the difference between x and y. The chunked pted test has time
         complexity O(c^2 * I * D * P), where c is the chunk size, I is the
         number of iterations, D is the number of dimensions, and P is the number
         of permutations. For chunking to be worth it you should have c^2 * I << n^2.
@@ -328,7 +323,6 @@ def pted_coverage_test(
             return_all=True,
             two_tailed=False,
             chunk_size=chunk_size,
-            chunk_iter=chunk_iter,
         )
         test_stats.append(test)
         permute_stats.append(permute)