Add exponential fitting for ZNE #5972

natestemen · 2024-07-09T04:44:56Z

Context: As part of increasing the utility of ZNE within Catalyst, we'd like to add the option to use exponential fitting with ZNE. The issue

Description of the Change: A function which is written using JAX is added which performs an exponential fit of data, and evaluates said fit at 0.

Benefits: Increased capabilities for users to play with. This is especially important given the benefits of tuning ZNE to the problem and backend at hand.

Possible Drawbacks: The current implementation transforms data using the natural logarithm $\ln$ to linearize the data so that a linear fit can be applied to learn the necessary parameters (see section below on the exponential model). Once they transformed parameters are learned, they can be converted back into those from the exponential model quite simply. The major drawback of this approach is that they expectation values have to be shifted to be in the domain of $\ln$ (i.e. $> 0$), or thrown away if small enough (see $\varepsilon$ eps parameter introduced in expontial_extrapolate).

Related GitHub Issues: PennyLaneAI/catalyst#754

Exponential model: Assuming an exponential model with asymptote $y(x) = A\mathrm{e}^{Bx} + C$ we must first turn this into a linear equation to be able to apply a linear fit.

$$y(x) = A\mathrm{e}^{Bx} + C$$

$$y(x) - C = A\mathrm{e}^{Bx}$$

$$\ln(y(x) - C) = \ln(A) + Bx$$

$$\tilde{y}(x) = \tilde{A} + Bx$$

Once $\tilde{A}$ is learned, the hard part is over as the ZNE value is "simply" $y(0) = A + C$.

pennylane/transforms/mitigate.py

josh146 · 2024-07-11T01:07:43Z

Thanks for this PR @natestemen!

⚠️ draft PR note: I'm currently having trouble getting tests to pass on this change. I am aware and actively working on it. Putting this PR up for visibility and in case anyone has clues where I'm messing things up.

Feel free to tag us directly on the PR, or alternatively let us know if you want anyone to go through the PR.

…xtrap

natestemen · 2024-07-12T06:18:25Z

Guidance would be appreciated on resolving the CI error raised here.

natestemen · 2024-07-12T06:18:45Z

@josh146 ready for the first look!

rmoyard

Thanks @natestemen, it looks great! I suggest you use the math module from PennyLane then the function is compatible with all frameworks and add some tests for the different ML frameworks. Otherwise make sure to parametize the unit test for the fit.

pennylane/transforms/mitigate.py

tests/transforms/test_mitigate.py

pennylane/transforms/mitigate.py

tests/transforms/test_mitigate.py

test both increasing and decreasing exponentials

rmoyard

It looks good to me! The failure seems unrelated. We need two approvals, I will tag someone from the core team.

codecov · 2024-07-17T22:01:06Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.65%. Comparing base (f454f59) to head (155e090).
Report is 294 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5972      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         427      427              
  Lines       41059    40776     -283     
==========================================
- Hits        40920    40636     -284     
- Misses        139      140       +1

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

EmilianoG-byte

Great addition 🤩 ! Just left some comments regarding typos/cosmetics.

pennylane/transforms/mitigate.py

EmilianoG-byte · 2024-07-18T21:31:29Z

Is there a reason why we are not checking differentiability for other backends besides jax? For instance here:

pennylane/tests/transforms/test_mitigate.py

Lines 646 to 647 in 169d53e

    
           @pytest.mark.parametrize("interface", ["auto", "torch"]) 
        
           def test_diffability_torch(self, interface):

why didn't we add

@pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])

like we did with JAX and JAX-JIT?

Co-Authored-By: Cristian Emiliano Godinez Ramirez <[email protected]>

natestemen · 2024-07-19T02:03:39Z

Thanks for the docstring suggestions! Addressed in e0815da.

Is there a reason why we are not checking differentiability for other backends besides jax?

The goal of this PR was to get an exponential fit function in that would be compatible with pennylaneai/catalyst. Originally I thought the exponential_extrapolate function was required to be written in JAX and hence only ensured it was tested as such. I wasn't aware of the testing requirements of the other backends¹ until getting some review on this PR.

Happy to add your suggestion. Let me know.

I'm still unclear if the other backends (torch, tf, autograd) are required in order for this functionality to work with catalyst. Is it? ↩

EmilianoG-byte · 2024-07-19T15:22:41Z

I'm still unclear if the other backends (torch, tf, autograd) are required in order for this functionality to work with catalyst. Is it? ↩

Good question @natestemen 🤔

Strictly speaking, the other interfaces are not needed for this functionality to work with Catalyst. However, since this will be part of the Pennylane features (even without Jitting our workflows), it should be compatible with the interfaces supported by Pennylane - this includes being able to differentiate it.

The good thing is that since you are using qml.math in exponential_extrapolate, adding this support should be rather straightforward by parametrizing the differentiability tests with

@pytest.mark.parametrize("extrapolate", [richardson_extrapolate, exponential_extrapolate])

wherever only richardson_extrapolate was used previously. Let me know if there is any unforeseen difficulty 😃.

pennylane/transforms/mitigate.py

rmoyard · 2024-07-19T15:32:39Z

@natestemen The requirement for it to work with Catalyst is to be Jax-jittable. I have asked for other interfaces because in PL we always make things compatible for the different frameworks.

doc/releases/changelog-dev.md

EmilianoG-byte

Looks good to me! 😄

**Context:** The last PR ([first](#806), [second](PennyLaneAI/pennylane#5972)) to add exponential extrapolation capabilities to catalyst. **Description of the Change:** This PR ensures that `pennylane.transforms.exponential_extrapolate` works inside catalyst's `mitigate_with_zne` function. **Benefits:** ZNE has increased functionality. **Possible Drawbacks:** As part of the testing, I've removed 0.1 and 0.2 from the list of values being used to create different circuits. This is because exponential extrapolation is not necessarily as stable as polynomial fitting. For example, when extrapolating near constant values exponential extrapolation can struggle due to the fact that a linear fit is first performed to understand the "direction" (or `sign`) of the exponential. If the slope is positive, the data is flipped to fit something that looks like $A\mathrm{e}^{-Bx} + C$ where $B\in\mathbb{R}_{> 0}$. An example of this happening can be seen here. ```py >>> exponential_extrapolate([1,2,3], [0.3894, 0.3894183, 0.38941]) -1.0000000000000059e-06 >>> richardson_extrapolate([1,2,3], [0.3894, 0.3894183, 0.38941]) 0.3893551000000072 ``` If we want to ensure to continue testing values 0.1 and 0.2 for polynomial extrapolation, we can create separate tests. Let me know what would be best. **Related GitHub Issues:** fixes #754

natestemen added 2 commits July 8, 2024 21:35

add exponential fitting function

82d0ec2

attempt to test

e808837

cosenal reviewed Jul 10, 2024

View reviewed changes

pennylane/transforms/mitigate.py Outdated Show resolved Hide resolved

natestemen added 5 commits July 11, 2024 08:20

protect against missing jax install

88cde3b

use linear extrapolation on scaled data

dc238e4

add tests/parametrize some

752fed6

Merge remote-tracking branch 'upstream/master' into nts-exponential-e…

b8db629

…xtrap

add description in changelog

b7e68e5

natestemen marked this pull request as ready for review July 12, 2024 06:18

rmoyard self-requested a review July 15, 2024 14:22

rmoyard reviewed Jul 15, 2024

View reviewed changes

tests/transforms/test_mitigate.py Show resolved Hide resolved

cosenal reviewed Jul 15, 2024

View reviewed changes

pennylane/transforms/mitigate.py Outdated Show resolved Hide resolved

pennylane/transforms/mitigate.py Outdated Show resolved Hide resolved

tests/transforms/test_mitigate.py Outdated Show resolved Hide resolved

natestemen and others added 10 commits July 15, 2024 13:03

parametrize multi-return test

31aa0e4

remove print calls

58d1415

make test name accurate

52825e1

use qml.math where possible

f961c23

flush out docstring

113aaa8

Merge branch 'master' into nts-exponential-extrap

c3da4e7

remove jax dependency from expfitting

b60ffc4

add framework specific unit tests

5248c02

output is no longer an Array

a7ec797

parametrize accuracy test

32af13c

test both increasing and decreasing exponentials

natestemen requested a review from rmoyard July 17, 2024 00:55

natestemen added 2 commits July 17, 2024 08:32

uncomment mark

494edcf

add test docstrings

b1c57b6

rmoyard approved these changes Jul 17, 2024

View reviewed changes

Merge branch 'master' into nts-exponential-extrap

169d53e

rmoyard requested a review from a team July 18, 2024 13:45

Alex-Preciado requested review from EmilianoG-byte and removed request for a team July 18, 2024 16:02

EmilianoG-byte reviewed Jul 18, 2024

View reviewed changes

docstring wordsmithing

e0815da

Co-Authored-By: Cristian Emiliano Godinez Ramirez <[email protected]>

Merge branch 'master' into nts-exponential-extrap

470998f

EmilianoG-byte reviewed Jul 19, 2024

View reviewed changes

pennylane/transforms/mitigate.py Outdated Show resolved Hide resolved

natestemen added 2 commits July 19, 2024 08:47

parametrize all tests with extrapolation function

87cbe10

fix unindented docstring argument description

5e78673

natestemen requested a review from EmilianoG-byte July 19, 2024 15:48

rmoyard reviewed Jul 19, 2024

View reviewed changes

doc/releases/changelog-dev.md Show resolved Hide resolved

natestemen and others added 2 commits July 19, 2024 09:35

add myself to contributors

e8c95e7

Merge branch 'master' into nts-exponential-extrap

7c0fff8

EmilianoG-byte approved these changes Jul 19, 2024

View reviewed changes

Merge branch 'master' into nts-exponential-extrap

155e090

natestemen mentioned this pull request Jul 23, 2024

Add tests for exponential extrapolation PennyLaneAI/catalyst#953

Merged

dime10 enabled auto-merge (squash) July 23, 2024 20:09

dime10 disabled auto-merge July 23, 2024 20:11

dime10 merged commit 46fa7ff into PennyLaneAI:master Jul 23, 2024
38 checks passed

natestemen deleted the nts-exponential-extrap branch July 23, 2024 20:20

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add exponential fitting for ZNE #5972

Add exponential fitting for ZNE #5972

natestemen commented Jul 9, 2024 •

edited

Loading

josh146 commented Jul 11, 2024

natestemen commented Jul 12, 2024

natestemen commented Jul 12, 2024

rmoyard left a comment

rmoyard left a comment

codecov bot commented Jul 17, 2024 •

edited

Loading

EmilianoG-byte left a comment •

edited

Loading

EmilianoG-byte commented Jul 18, 2024 •

edited

Loading

natestemen commented Jul 19, 2024

EmilianoG-byte commented Jul 19, 2024

rmoyard commented Jul 19, 2024 •

edited

Loading

EmilianoG-byte left a comment

Add exponential fitting for ZNE #5972

Add exponential fitting for ZNE #5972

Conversation

natestemen commented Jul 9, 2024 • edited Loading

josh146 commented Jul 11, 2024

natestemen commented Jul 12, 2024

natestemen commented Jul 12, 2024

rmoyard left a comment

Choose a reason for hiding this comment

rmoyard left a comment

Choose a reason for hiding this comment

codecov bot commented Jul 17, 2024 • edited Loading

Codecov Report

EmilianoG-byte left a comment • edited Loading

Choose a reason for hiding this comment

EmilianoG-byte commented Jul 18, 2024 • edited Loading

natestemen commented Jul 19, 2024

Footnotes

EmilianoG-byte commented Jul 19, 2024

rmoyard commented Jul 19, 2024 • edited Loading

EmilianoG-byte left a comment

Choose a reason for hiding this comment

natestemen commented Jul 9, 2024 •

edited

Loading

codecov bot commented Jul 17, 2024 •

edited

Loading

EmilianoG-byte left a comment •

edited

Loading

EmilianoG-byte commented Jul 18, 2024 •

edited

Loading

rmoyard commented Jul 19, 2024 •

edited

Loading