diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 6284ed8..21e2290 100644 --- a/.github/workflows/cd.yml +++ b/.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 diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index fa73e2d..1dc49b6 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/README.md b/README.md index f460e36..56a63db 100644 --- a/README.md +++ b/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 ``` diff --git a/pyproject.toml b/pyproject.toml index 0b3aee1..c96f676 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/pted/pted.py b/src/pted/pted.py index 0daf3b8..590941c 100644 --- a/src/pted/pted.py +++ b/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,14 +108,13 @@ 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. @@ -124,9 +122,11 @@ def pted( 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) diff --git a/src/pted/utils.py b/src/pted/utils.py index 57df5fb..17ce375 100644 --- a/src/pted/utils.py +++ b/src/pted/utils.py @@ -19,9 +19,11 @@ class torch: try: import jax import jax.numpy as jnp + from jax import jit except ImportError: jax = None jnp = None + jit = lambda *a, **k: lambda f: f # type: ignore __all__ = ( @@ -86,56 +88,45 @@ def _energy_distance_torch( return _energy_distance_precompute(D, nx, ny).item() -def _energy_distance_estimate_numpy( - x: np.ndarray, - y: np.ndarray, - chunk_size: int, - chunk_iter: int, - metric: Union[str, float] = "euclidean", -) -> float: +def _chunk_slices(lenx: int, leny: int, chunk_size: int): + """Yield slices for chunking two arrays of lengths lenx and leny. - E_est = [] - for _ in range(chunk_iter): - # Randomly sample a chunk of data - idx = np.random.choice(len(x), size=min(len(x), chunk_size), replace=False) - x_chunk = x[idx] - idy = np.random.choice(len(y), size=min(len(y), chunk_size), replace=False) - y_chunk = y[idy] - - # Compute the energy distance - E_est.append(_energy_distance_numpy(x_chunk, y_chunk, metric=metric)) - return np.mean(E_est) + The smaller of the two is cycled while the larger is iterated through + completely (minus the last incomplete chunk if any). + """ + nx = max(1, lenx // chunk_size) + ny = max(1, leny // chunk_size) + + for i in range(max(nx, ny)): + ix = i % nx + iy = i % ny + yield slice(ix * chunk_size, (ix + 1) * chunk_size), slice( + iy * chunk_size, (iy + 1) * chunk_size + ) -def _energy_distance_estimate_torch( - x: torch.Tensor, - y: torch.Tensor, +def _energy_distance_estimate( + x, + y, chunk_size: int, - chunk_iter: int, - metric: Union[str, float] = "euclidean", + metric: Union[str, float], + energy_distance_fn, ) -> float: + """Estimate energy distance by averaging over sequential sliced chunks. + Iterates ``max(len(x), len(y)) // chunk_size`` times, using plain slicing + on both arrays. The smaller of the two is tiled along axis 0 as needed so + that both arrays are at least ``n_iter * chunk_size`` rows long before the + loop begins. + """ E_est = [] - for _ in range(chunk_iter): - # Randomly sample a chunk of data - idx = np.random.choice(len(x), size=min(len(x), chunk_size), replace=False) - x_chunk = x[torch.tensor(idx)] - idy = np.random.choice(len(y), size=min(len(y), chunk_size), replace=False) - y_chunk = y[torch.tensor(idy)] - - # Compute the energy distance - E_est.append(_energy_distance_torch(x_chunk, y_chunk, metric=metric)) + for cx, cy in _chunk_slices(len(x), len(y), chunk_size): + E_est.append(energy_distance_fn(x[cx], y[cy], metric=metric)) return np.mean(E_est) +@jit(static_argnames=["p"]) def _jax_cdist(x, y, p: float = 2.0): - if p == 2.0: - # Squared-norm identity avoids materializing the (nx, ny, d) diff tensor. - # ||x_i - y_j||^2 = ||x_i||^2 + ||y_j||^2 - 2 * x_i . y_j - x_sq = jnp.sum(x**2, axis=-1) # (nx,) - y_sq = jnp.sum(y**2, axis=-1) # (ny,) - sq_dist = x_sq[:, None] + y_sq[None, :] - 2.0 * (x @ y.T) - return jnp.sqrt(jnp.maximum(sq_dist, 0.0)) # For general p-norms use vmap to avoid the (nx, ny, d) intermediate. return jax.vmap(lambda xi: jnp.linalg.norm(xi - y, ord=p, axis=-1))(x) @@ -150,47 +141,28 @@ def _energy_distance_jax(x, y, metric: Union[str, float] = "euclidean") -> float return float(_energy_distance_precompute(D, nx, ny)) -def _energy_distance_estimate_jax( - x, - y, - chunk_size: int, - chunk_iter: int, - metric: Union[str, float] = "euclidean", -) -> float: - - E_est = [] - for _ in range(chunk_iter): - # Randomly sample a chunk of data - idx = np.random.choice(len(x), size=min(len(x), chunk_size), replace=False) - x_chunk = x[idx] - idy = np.random.choice(len(y), size=min(len(y), chunk_size), replace=False) - y_chunk = y[idy] - - # Compute the energy distance - E_est.append(_energy_distance_jax(x_chunk, y_chunk, metric=metric)) - return np.mean(E_est) - - def pted_chunk_numpy( x: np.ndarray, y: np.ndarray, permutations: int = 100, metric: str = "euclidean", chunk_size: int = 100, - chunk_iter: int = 10, prog_bar: bool = False, ) -> tuple[float, list[float]]: assert np.all(np.isfinite(x)) and np.all(np.isfinite(y)), "Input contains NaN or Inf!" nx = len(x) - test_stat = _energy_distance_estimate_numpy(x, y, chunk_size, chunk_iter, metric=metric) + test_stat = _energy_distance_estimate( + x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy + ) permute_stats = [] + z = np.concatenate((x, y), axis=0) for _ in trange(permutations, disable=not prog_bar): - z = np.concatenate((x, y), axis=0) z = z[np.random.permutation(len(z))] - x, y = z[:nx], z[nx:] permute_stats.append( - _energy_distance_estimate_numpy(x, y, chunk_size, chunk_iter, metric=metric) + _energy_distance_estimate( + z[:nx], z[nx:], chunk_size, metric=metric, energy_distance_fn=_energy_distance_numpy + ) ) return test_stat, permute_stats @@ -201,7 +173,6 @@ def pted_chunk_torch( permutations: int = 100, metric: Union[str, float] = "euclidean", chunk_size: int = 100, - chunk_iter: int = 10, prog_bar: bool = False, ) -> tuple[float, list[float]]: assert torch.__version__ != "null", "PyTorch is not installed! try: `pip install torch`" @@ -210,14 +181,17 @@ def pted_chunk_torch( ), "Input contains NaN or Inf!" nx = len(x) - test_stat = _energy_distance_estimate_torch(x, y, chunk_size, chunk_iter, metric=metric) + test_stat = _energy_distance_estimate( + x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch + ) permute_stats = [] + z = torch.cat((x, y), dim=0) for _ in trange(permutations, disable=not prog_bar): - z = torch.cat((x, y), dim=0) z = z[torch.randperm(len(z))] - x, y = z[:nx], z[nx:] permute_stats.append( - _energy_distance_estimate_torch(x, y, chunk_size, chunk_iter, metric=metric) + _energy_distance_estimate( + z[:nx], z[nx:], chunk_size, metric=metric, energy_distance_fn=_energy_distance_torch + ) ) return test_stat, permute_stats @@ -309,21 +283,23 @@ def pted_chunk_jax( permutations: int = 100, metric: Union[str, float] = "euclidean", chunk_size: int = 100, - chunk_iter: int = 10, prog_bar: bool = False, ) -> tuple[float, list[float]]: assert jax is not None, "JAX is not installed! try: `pip install jax`" assert jnp.all(jnp.isfinite(x)) and jnp.all(jnp.isfinite(y)), "Input contains NaN or Inf!" nx = len(x) - test_stat = _energy_distance_estimate_jax(x, y, chunk_size, chunk_iter, metric=metric) + test_stat = _energy_distance_estimate( + x, y, chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax + ) permute_stats = [] + z = jnp.concatenate([x, y], axis=0) for _ in trange(permutations, disable=not prog_bar): - z = jnp.concatenate([x, y], axis=0) z = z[np.random.permutation(len(z))] - x, y = z[:nx], z[nx:] permute_stats.append( - _energy_distance_estimate_jax(x, y, chunk_size, chunk_iter, metric=metric) + _energy_distance_estimate( + z[:nx], z[nx:], chunk_size, metric=metric, energy_distance_fn=_energy_distance_jax + ) ) return test_stat, permute_stats diff --git a/tests/test_pted.py b/tests/test_pted.py index 35f397a..7d3d166 100644 --- a/tests/test_pted.py +++ b/tests/test_pted.py @@ -31,6 +31,7 @@ def test_inputs_extra_dims(): if torch is None: pytest.skip("torch not installed") # Test with torch tensors + torch.manual_seed(43) g = torch.randn(100, 30, 30) s = torch.randn(50, 100, 30, 30) p = pted.pted_coverage_test(g, s) @@ -84,6 +85,7 @@ def test_pted_torch(): def test_pted_coverage_full(): + np.random.seed(42) g = np.random.normal(size=(100, 10)) # ground truth (n_simulations, n_dimensions) s = np.random.normal( size=(200, 100, 10) @@ -104,11 +106,11 @@ def test_pted_chunk_torch(): D = 10 x = torch.randn(1000, D) y = torch.randn(1000, D) - p = pted.pted(x, y, chunk_size=100, chunk_iter=10) + p = pted.pted(x, y, chunk_size=100) assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))" y = torch.rand(1000, D) - p = pted.pted(x, y, chunk_size=100, chunk_iter=10) + p = pted.pted(x, y, chunk_size=100) assert p < 1e-2, f"p-value {p} is not in the expected range (~0)" @@ -119,16 +121,50 @@ def test_pted_chunk_numpy(): D = 10 x = np.random.normal(size=(1000, D)) y = np.random.normal(size=(1000, D)) - p = pted.pted(x, y, chunk_size=100, chunk_iter=10) + p = pted.pted(x, y, chunk_size=100) assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))" y = np.random.uniform(size=(1000, D)) - p = pted.pted(x, y, chunk_size=100, chunk_iter=10) + p = pted.pted(x, y, chunk_size=100) + assert p < 1e-2, f"p-value {p} is not in the expected range (~0)" + + +def test_pted_chunk_mismatched_sizes(): + """Chunked PTED correctly cycles the smaller dataset for mismatched x/y sizes.""" + np.random.seed(0) + D = 5 + # x has 800 samples, y has 300 samples; chunk_size=100 → 8 iterations, cycling y + x = np.random.normal(size=(800, D)) + y = np.random.normal(size=(300, D)) + p = pted.pted(x, y, chunk_size=100) + assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))" + + # Different distributions should give small p-value even with mismatched sizes + y_diff = np.random.uniform(size=(300, D)) + p = pted.pted(x, y_diff, chunk_size=100) assert p < 1e-2, f"p-value {p} is not in the expected range (~0)" +def test_pted_chunk_size_fallback(): + """chunk_size >= both dataset sizes falls back to regular (non-chunked) PTED.""" + np.random.seed(7) + D = 5 + x = np.random.normal(size=(50, D)) + y = np.random.normal(size=(50, D)) + + # chunk_size equals both lengths → should silently fall back to regular PTED + p_chunk = pted.pted(x, y, chunk_size=50) + # Rerun with explicit non-chunked PTED using same seed for reference + np.random.seed(7) + p_plain = pted.pted(x, y) + # Both should give a valid p-value in U(0,1) + assert p_chunk > 1e-2 and p_chunk < 0.99, f"p-value {p_chunk} outside expected range" + assert p_plain > 1e-2 and p_plain < 0.99, f"p-value {p_plain} outside expected range" + + def test_pted_coverage_edgecase(): # Test with single simulation + np.random.seed(42) g = np.random.normal(size=(1, 10)) s = np.random.normal(size=(100, 1, 10)) p = pted.pted_coverage_test(g, s) @@ -136,6 +172,7 @@ def test_pted_coverage_edgecase(): def test_pted_coverage_progress_bar(capsys): + np.random.seed(42) g = np.random.normal(size=(42, 10)) s = np.random.normal(size=(100, 42, 10)) pted.pted_coverage_test(g, s) @@ -152,6 +189,7 @@ def test_pted_coverage_progress_bar(capsys): def test_pted_coverage_overunder(): if torch is None: pytest.skip("torch not installed") + torch.manual_seed(42) g = torch.randn(100, 3) s = torch.randn(50, 100, 3) with pytest.warns(pted.utils.OverconfidenceWarning): @@ -161,6 +199,7 @@ def test_pted_coverage_overunder(): def test_sbc_histogram(): + np.random.seed(42) g = np.random.normal(size=(100, 10)) # ground truth (nsim, ndim) s = np.random.normal(size=(150, 100, 10)) # posterior samples (nsamp, nsim, ndim) @@ -169,6 +208,7 @@ def test_sbc_histogram(): def test_pit_plot_coverage_test(): + np.random.seed(42) g = np.random.normal(size=(100, 10)) # ground truth (nsim, ndim) s = np.random.normal(size=(150, 100, 10)) # posterior samples (nsamp, nsim, ndim) @@ -179,6 +219,7 @@ def test_pit_plot_coverage_test(): def test_pit_plot_utility_direct(): """pit_plot utility function creates a file and handles edge cases.""" + np.random.seed(42) pvals = np.random.uniform(size=50) pted.utils.pit_plot(pvals, "pit_direct.pdf") assert os.path.exists("pit_direct.pdf"), "PIT plot file was not created" @@ -225,17 +266,17 @@ def test_pted_jax(): def test_pted_chunk_jax(): if jax is None: pytest.skip("jax not installed") - np.random.seed(42) + np.random.seed(0) # example 2 sample test D = 3 - x = jnp.array(np.random.normal(size=(100, D))) - y = jnp.array(np.random.normal(size=(100, D))) - p = pted.pted(x, y, chunk_size=100, chunk_iter=10) + x = jnp.array(np.random.normal(size=(200, D))) + y = jnp.array(np.random.normal(size=(200, D))) + p = pted.pted(x, y, chunk_size=50) assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))" - y = jnp.array(np.random.uniform(size=(110, D))) - p = pted.pted(x, y, chunk_size=100, chunk_iter=10) + y = jnp.array(np.random.uniform(size=(300, D))) + p = pted.pted(x, y, chunk_size=50) assert p < 1e-2, f"p-value {p} is not in the expected range (~0)" @@ -243,17 +284,13 @@ def test_pted_coverage_jax(): if jax is None: pytest.skip("jax not installed") + np.random.seed(42) g = jnp.array(np.random.normal(size=(75, 5))) s = jnp.array(np.random.normal(size=(50, 75, 5))) p = pted.pted_coverage_test(g, s) assert p > 1e-2 and p < 0.99, f"p-value {p} is not in the expected range (U(0,1))" -# --------------------------------------------------------------------------- -# Unit tests for newly-added utils functions -# --------------------------------------------------------------------------- - - def test_is_jax_array_with_jax(): """is_jax_array returns True for a real JAX array and False for other types.""" if jax is None: @@ -298,23 +335,13 @@ def test_energy_distance_jax(): """_energy_distance_jax returns 0 when x and y are identical.""" if jax is None: pytest.skip("jax not installed") + np.random.seed(42) x = jnp.array(np.random.normal(size=(50, 5))) # Identical samples means energy distance should be ~0 ed = pted.utils._energy_distance_jax(x, x) assert abs(ed) < 1e-6 -def test_energy_distance_estimate_jax(): - """_energy_distance_estimate_jax returns a finite scalar.""" - if jax is None: - pytest.skip("jax not installed") - np.random.seed(0) - x = jnp.array(np.random.normal(size=(100, 4))) - y = jnp.array(np.random.normal(size=(100, 4))) - ed = pted.utils._energy_distance_estimate_jax(x, y, chunk_size=20, chunk_iter=5) - assert np.isfinite(ed) - - def test_pted_jax_no_jax(monkeypatch): """pted_jax raises AssertionError when JAX is not installed.""" monkeypatch.setattr("pted.utils.jax", None) @@ -345,11 +372,6 @@ def test_pted_chunk_torch_no_torch(monkeypatch): pted.utils.pted_chunk_torch(np.zeros((5, 2)), np.zeros((5, 2))) -# --------------------------------------------------------------------------- -# Cross-backend consistency tests -# --------------------------------------------------------------------------- - - def test_jax_cdist_matches_scipy(): """_jax_cdist (L2) and scipy cdist produce the same pairwise distances.""" if jax is None: @@ -390,33 +412,21 @@ def test_energy_distance_numpy_torch_jax_agree(): ), f"numpy ({ed_numpy}) and jax ({ed_jax}) energy distances differ" -def test_energy_distance_estimate_numpy_torch_jax_agree(): - """_energy_distance_estimate_{numpy,torch,jax} return close values for the same seed/data.""" - if torch is None: - pytest.skip("torch not installed") - if jax is None: - pytest.skip("jax not installed") - +def test_energy_distance_estimate_matches_energy_distance(): + """_energy_distance_estimate returns a value close to _energy_distance for small inputs.""" np.random.seed(123) - # Use float32 so all backends operate at the same precision - x_np = np.random.normal(size=(200, 5)).astype(np.float32) - y_np = np.random.normal(size=(200, 5)).astype(np.float32) - - # Run with the same seed so the same chunks are sampled - np.random.seed(0) - ed_numpy = pted.utils._energy_distance_estimate_numpy(x_np, y_np, chunk_size=50, chunk_iter=5) - np.random.seed(0) - ed_torch = pted.utils._energy_distance_estimate_torch( - torch.tensor(x_np), torch.tensor(y_np), chunk_size=50, chunk_iter=5 - ) - np.random.seed(0) - ed_jax = pted.utils._energy_distance_estimate_jax( - jnp.array(x_np), jnp.array(y_np), chunk_size=50, chunk_iter=5 + x = np.random.normal(size=(100, 5)) + y = np.random.uniform(size=(100, 5)) + + ed_direct = pted.utils._energy_distance_numpy(x, y) + ed_estimate = pted.utils._energy_distance_estimate( + x, + y, + chunk_size=25, + metric="euclidean", + energy_distance_fn=pted.utils._energy_distance_numpy, ) - assert ed_numpy == pytest.approx( - ed_torch, rel=1e-4 - ), f"numpy ({ed_numpy}) and torch ({ed_torch}) energy distance estimates differ" - assert ed_numpy == pytest.approx( - ed_jax, rel=1e-4 - ), f"numpy ({ed_numpy}) and jax ({ed_jax}) energy distance estimates differ" + assert ed_direct == pytest.approx( + ed_estimate, rel=1.5e-1 + ), f"direct ({ed_direct}) and estimate ({ed_estimate}) energy distances differ"