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

[Copilot]

The chunking API required users to specify both chunk_size and chunk_iter together, with chunk_iter being redundant given the data size. The new interface infers iteration count automatically and handles mismatched dataset sizes by cycling the smaller dataset.

Interface changes

• Removed chunk_iter from pted(), pted_coverage_test(), and all internal pted_chunk_* / _energy_distance_estimate_* functions.
• Iteration count is now max(len(x), len(y)) // chunk_size (remainder samples discarded).
• Mismatched sizes: iterates over the larger dataset's chunks; cycles through the smaller one via modular indexing.
• Added assertion: chunk_size <= max(len(x), len(y)).

# Before
p = pted(x, y, chunk_size=100, chunk_iter=10)

# After — chunk_iter derived automatically (e.g., 1000 samples → 10 iters)
p = pted(x, y, chunk_size=100)

# Mismatched sizes just work — smaller dataset is cycled
p = pted(x_800, y_300, chunk_size=100)  # 8 iterations, y cycles

Implementation

• Sequential chunk traversal replaces random subsampling; data is already shuffled by the permutation step, so random sampling within chunks is unnecessary.
• All three backends (NumPy, PyTorch, JAX) updated consistently.
• Tests updated to remove chunk_iter, with a new test_pted_chunk_mismatched_sizes test covering the cycling behavior.
• README argument tables and the "Memory and Compute limitations" section updated.

• Fixes PTED chunking cleanup #17

[ConnorStoneAstro]

@copilot I think for the energy distance estimate, you can now use slicing instead of fancy indexing. This might mean that there can just be one energy distance estimate function now.

Also, I think if the chunk size is larger than both the datasets, then you can just fallback to regular PTED.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 95.65217% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 97.66%. Comparing base (0bc4398) to head (dea8e0b).

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

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #18      +/-   ##
==========================================
- Coverage   98.05%   97.66%   -0.39%     
==========================================
  Files           4        4              
  Lines         359      342      -17     
==========================================
- Hits          352      334      -18     
- Misses          7        8       +1     

☔ 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]

🟡 Not ready to approve

The new chunk estimator can produce nan for some small-input usages (public pted_chunk_*) and chunk_size validation/documentation mismatches need to be resolved.

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

This review doesn't count toward merge requirements. Sign up for the private preview to control whether Copilot approvals count.

Pull request overview

This PR simplifies PTED’s chunking interface by removing the redundant chunk_iter parameter and deriving the number of chunk iterations from chunk_size, including support for mismatched x/y sizes by cycling the smaller dataset’s chunks.

Changes:

• Removed chunk_iter from pted(), pted_coverage_test(), backend chunk helpers, and tests.
• Implemented sequential chunk traversal with iteration count derived from max(len(x), len(y)) // chunk_size, cycling the smaller dataset’s chunks.
• Updated documentation and tests, including new mismatched-size coverage.

File summaries

File	Description
tests/test_pted.py	Updates chunking tests to the new API, adds mismatched-size and fallback coverage, and revises energy-distance estimate tests.
src/pted/utils.py	Refactors chunked energy-distance estimation to remove chunk_iter and share logic across backends.
src/pted/pted.py	Removes chunk_iter from the public API and updates chunking behavior/parameter documentation.
README.md	Updates argument docs and chunking guidance to match the new chunking interface.

Review details

Suppressed comments (2)

src/pted/utils.py:125

• _energy_distance_estimate will return nan when chunk_size exceeds both dataset lengths (because _chunk_slices yields 0 iterations and np.mean([]) is nan). Since pted_chunk_* are exported in __all__, calling them on small inputs with the default chunk_size=100 can trigger this. Consider falling back to a single full energy-distance computation when no chunks are produced, and update the docstring (it currently claims the smaller input is tiled).

    """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.
    """

src/pted/pted.py:128

• The PR description mentions an added assertion chunk_size <= max(len(x), len(y)), but the implementation instead silently disables chunking when chunk_size covers both datasets (and also accepts chunk_size larger than both). This is a user-visible contract difference; either the PR description/docs should be updated to describe the fallback, or the code should enforce the asserted constraint.

    if chunk_size is not None:
        # If chunk_size covers both full datasets, chunking adds no benefit
        if chunk_size >= len(x) and chunk_size >= len(y):
            chunk_size = None

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

We're testing this review assessment. Please use 👍 or 👎 to tell us if it's correct.