feat: qudit noise model (2/5) (#1854)

[bramathon]

Part of the effort to land the experimental noise model and simulation module (originally PR #1848, branch nonbreaking-noise-model) as a reviewable stack. This is PR 2 of 5, stacked on #1858 — review/merge in order. Closes #1854.

This PR introduces the new quax-backed qudit noise model under pyquil/noise/:

• Channel, MeasurementChannel, ResetChannel, and CycleChannel (_channels.py), and the new NoiseModel container plus DepolarizingNoiseModel / CompositeNoiseModel and fidelity estimators (_noise_model.py).
• The historical pyquil.noise public API is preserved: pyquil/noise.py is moved to pyquil/noise/_legacy_noise.py and re-exported from pyquil/noise/__init__.py, with the legacy NoiseModel/KrausModel now marked deprecated in favor of the quax-based model.

It also adds qudit DefGate support in quilbase.py (matrix dimensions may now be any perfect power, e.g. 3, 9, …, not just powers of two), which the qutrit noise channels and later simulators rely on, and enables float64 (jax_enable_x64) for the test suite.

A handful of ResetChannel tests verify behavior end-to-end through the density-matrix simulator; since that simulator lands in the next PR, those specific assertions are deferred to #1855 (they are restored there). Everything else in the noise module is fully tested here (test_noise_model.py).

Stack

1. Drop Python 3.9 #1852 — Drop Python 3.9
2. Qudit Noise model #1854 — Qudit noise model (this PR)
3. Experimental exact simulators #1855 — Experimental exact simulators
4. Trajectory simulator #1856 — Trajectory simulator
5. Dynamic shape trajectory simulator #1857 — Dynamic-shape trajectory simulator

Tracked in #1863.

[github-actions]

Bencher Report

Project	pyquil
Branch	1854-qudit-noise-model
Testbed	ci-runner-linux

⚠️ WARNING: No Threshold found!

Without a Threshold, no Alerts will ever be generated.

• Latency (nanoseconds (ns))

Click here to create a new Threshold
For more information, see the Threshold documentation.
To only post results if a Threshold exists, set the --ci-only-thresholds flag.

Click to view all benchmark results

Benchmark	Latency	seconds (s)
test/benchmarks/test_program.py::test_copy_everything_except_instructions	📈 view plot ⚠️ NO THRESHOLD	8.21 s
test/benchmarks/test_program.py::test_instructions	📈 view plot ⚠️ NO THRESHOLD	2.59 s
test/benchmarks/test_program.py::test_iteration	📈 view plot ⚠️ NO THRESHOLD	2.60 s

🐰 View full continuous benchmarking report in Bencher

[asaites]

I've made several comments, which I hope will be useful. I might spend a little more time looking through _noise_model.py, but for the most part, I think I've gone through this pretty thoroughly (and am doing my best to keep up with the physics).

Most of my comments are more style than anything, but I have a few concerns. Overall, though, this looks pretty solid!

[asaites]

This might not matter much for the expected values of rows/reasonable values of len(self.matrix), but it is very slow for large numbers. As a much faster alternative, you can calculate it by iteratively checking if $\sqrt[p]{r} \in \mathbb{Z}$ for $p \in \text{PRIME}$, stopping when $\sqrt[p]{r} < 2$. If so, it's a factor of k, and you can reduce the problem to finding the solution for that new base, then multiply its factors by the one you've found.

Concretely:

def _int_n_root(x: int, n: int) -> int | None:
    """Return the largest integer less than or equal to n-th root of x."""
    # This helper quickly checks if `x` is a perfect `n-th` root,
    # and will do so without the numerical instability
    # of the floating point check.
    # There are faster algorithms for very large numbers,
    # but this is already O(log(x)), and simple to implement.

    low, high = 1, x
    while low <= high:
        mid = (low + high) // 2
        if (mid_pow_n := mid ** n) == x:
            return mid
        elif mid_pow_n < x:
            low = mid + 1
        else:
            high = mid - 1
    return low - 1

SMALL_PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61]

def find_k(r: int) -> tuple[int, int]:
    """For positive integer `r`, return integers `(b, k)` such that `r = b**k`
    and `k` is the largest integer for any integer `b >= 1` that satisfies this relation.

    If `r` is not a perfect power, this function returns `(r, 1)`.
    """
    if r < 2:
        return r, 1

    # Check if `r` is a perfect p-th power of a small prime `p`.
    # When it is, reduce `r` to that p-th root and build up the exponent `k`.

    base, k = r, 1
    primes = iter(SMALL_PRIMES)
    p = next(primes, None)
    while p:
        if (b := _int_n_root(base, p)) < 2:
            return base, k
        elif b ** p != base:
            p = next(primes, None)
        else:
            k *= p
            base = b

    raise ValueError("r is too large to efficiently find k")

This approach is $O(\log^{2}{n})$ versus $O(n^2)$.

[bramathon]

On further thought, this entire concept has a huge problem - it doesn't support quarts. Besides qubits and qutrits this the most likely qudit size. I think I'll need to add a option to specify the qudit dims

[bramathon]

No good solution to this right now. I think we will just have to accept only prime qudit dims

[github-actions]

☂️ Python Coverage

current status: ❌

Overall Coverage

Lines	Covered	Coverage	Threshold	Status
8296	7246	87%	87%	🟢

New Files

File	Coverage	Status
pyquil/noise/init.py	100%	🟢
pyquil/noise/_channels.py	85%	🔴
pyquil/noise/_noise_model.py	76%	🔴
TOTAL	87%	🔴

Modified Files

File	Coverage	Status
pyquil/quilbase.py	94%	🟢
TOTAL	94%	🟢

updated for commit: 69edd32 by action🐍

[asaites]

This doesn't work if the exponent is composite.

[bramathon]

I've done a substantial refactor to add support for Lindbladian channels. The way this works is as follows:

ChannelProtocol defines the protocol for a channel which provides a superoperator associated with a gate. There are two concrete classes for this: Channel and SuperopChannel. The Channel has the _LindbladianBacked mixin, which means that it not only has a superoperator, but a lindbladian which generates it. This is now our canonical Channel while SuperopChannel is relegated to special use cases where we don't know the Lindbladian.

ResetChannel receives the same treatment, becoming _LindbladianBacked and with two concrete implementations.

CycleChannel and MeasureChannel are unchanged.

[asaites]

I think there's a problem with the jax or quax dependency version.

[windsurf-bot]

1 file skipped due to size limits:

• pyquil/noise/_channels.py

Looks good to me 🤙

💡 To request another review, post a new comment with "/windsurf-review".

[bramathon]

I think there's a problem with the jax or quax dependency version.

Confirmed, we need to bump the minimum version of quax