Add exponential fitting functionality to ZNE

[natestemen]

Motivation

Zero-Noise Extrapolation (ZNE) as implemented in https://github.com/PennyLaneAI/catalyst/pull/324 and https://github.com/PennyLaneAI/catalyst/pull/414 allow a user to employ ZNE with a polynomial extrapolation of a particular degree. Allowing other extrapolation methods will give catalyst more functionality and flexibility for users. The API should allow for exponential fitting.[^1]

Technical details

The existing ZNE class provides an API that accepts a paramater deg to denote the degree of the polynomial fit. This will need to be modified (design TBD) to accept either the degree, or a flag to mark the fact that exponential fitting will be used instead of polynomial.

Future steps

Once the design is finalized, further inspiration can be taken from Mitiq's PolyExpFactory.extrapolate function which allows for more data and flags to be passed to the extrapolation method.

[^1]: More can be added, such as logarithmic fitting, or the use of adaptive methods, but exponential will be a good first add to get things started.

[josh146]

Requirements

The mitigate_with_zne function is modified to take an extrapolate argument:

mitigate_with_zne(f, *, scale_factors: jax.Array, extrapolate: Callable, deg)

where extrapolate accepts a function that returns an extrapolated result when provided a range of scale factors and corresponding results.

PennyLane already provides two extrapolation functions:

• poly_extrapolate
• richardson_extrapolate

For more details on extrapolation, see the mitiq documentation.

The mitigate_with_zne docstring should be adequately udpated.

Technical notes

• It might make sense to utilize recently added callback support to support Python functions for the extrapolation callable.

[dime10]

It might make sense to utilize recently added callback support to support Python functions for the extrapolation callable.

I would recommend against building Catalyst features that rely on callbacks (except in specific contexts like debugging) because it restricts those features from being available in other environments in the future as well as making them opaque to the compiler.

[josh146]

@dime10 yep good point. Assuming the extrapolation functions are written in pure JAX, they should be lowered down to MLIR without a callback.

[dime10]

Hi @natestemen, the existing Python (user) tests for the ZNE functionality are located here: https://github.com/PennyLaneAI/catalyst/blob/main/frontend/test/pytest/test_mitigation.py