Completely change how chunked evaluation works

[ConnorStoneAstro]

Chunked evaluation now works on a rectangular distance matrix with a chosen width based on chunk_size

[codecov-commenter]

Codecov Report

❌ Patch coverage is 99.02913% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.84%. Comparing base (ec4a5c8) to head (6914631).

Files with missing lines	Patch %	Lines
src/pted/utils.py	99.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #20      +/-   ##
==========================================
- Coverage   97.65%   96.84%   -0.81%     
==========================================
  Files           4        4              
  Lines         341      317      -24     
==========================================
- Hits          333      307      -26     
- Misses          8       10       +2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Copilot]

🟡 Changes recommended

The new energy-distance computation and chunked permutation logic contain correctness issues that can invalidate test statistics/p-values (and torch chunking may break on GPU).

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR refactors PTED’s backend handling (NumPy / PyTorch / JAX) and replaces the prior chunked-evaluation approach with a landmark-based rectangular distance matrix approach controlled by chunk_size.

Changes:

• Consolidates backend-specific implementations into unified pted / pted_chunk helpers in src/pted/utils.py, and simplifies src/pted/pted.py to delegate to them.
• Reimplements chunked evaluation using a fixed landmark set and row-permutation reindexing of a rectangular distance matrix.
• Updates tests to run key behaviors across NumPy / torch / JAX backends and adjusts documentation to reflect Euclidean-only behavior and the new chunking scheme.

File summaries

File	Description
tests/test_utils.py	Updates two_tailed_p expectations by removing a now-obsolete assertion test.
tests/test_pted.py	Parameterizes tests across backends and updates coverage/chunking tests to match the new implementation.
src/pted/utils.py	Introduces unified backend dispatch, implements landmark-based chunked evaluation, and removes backend-specific public entry points.
src/pted/pted.py	Switches to using unified utils implementations and updates docstrings/complexity notes for the new chunking approach.
README.md	Updates docs for landmark-based chunking and Euclidean-only behavior; adds explanatory notes/footnotes.

Review details

Suppressed comments (2)

src/pted/utils.py:241

• pted_chunk mutates landmark_pos inside the permutation loop, but dmatrix is always indexed from the original (unpermuted) ordering (dmatrix[I]). After the first iteration, landmark_pos = np.argsort(I)[landmark_pos] is no longer applying the inverse permutation to the original landmark positions, so landmark re-splitting will be wrong for subsequent permutations.

            I = np.random.permutation(len(z))
            # Track where each landmark moved to, then re-split columns by their new x/y side.
            landmark_pos = np.argsort(I)[landmark_pos]
            order = np.argsort(landmark_pos >= nx, kind="stable")
            dmatrix_i = dmatrix[I][:, order]
            landmark_pos = landmark_pos[order]
            nxc_i = int(np.sum(landmark_pos < nx))

src/pted/pted.py:166

• pted_coverage_test also removed the metric parameter, which is a breaking change for downstream callers. If the library is now Euclidean-only, consider keeping metric in the signature (defaulting to "euclidean") and rejecting other values with a clear message, similar to pted().

def pted_coverage_test(
    g: Union[np.ndarray, "Tensor", "jax.Array"],
    s: Union[np.ndarray, "Tensor", "jax.Array"],
    permutations: int = 1000,
    warn_confidence: Optional[float] = 1e-3,
    return_all: bool = False,
    chunk_size: Optional[int] = None,
    sbc_histogram: Optional[str] = None,
    sbc_bins: Optional[int] = None,
    pit_plot: Optional[str] = None,
    pit_confidence: float = 0.95,
    prog_bar: bool = False,
) -> Union[float, tuple[np.ndarray, np.ndarray, float]]:
    """

• Files reviewed: 5/5 changed files
• Comments generated: 4
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.