Broadcasted gaussian integral #494

ziofil · 2024-09-26T17:24:54Z

User description

Context:
When circuit components have a large batch dimension, it's not efficient to loop over it when computing gaussian integrals.

Description of the Change:
Gaussian integrals now can be computed using broadcasting rules

Benefits:
Faster (about 10x)

Possible Drawbacks:
Implemented the version where two Abc triples are passed, not the one where a single Abc triple is being contracted

Related GitHub Issues:
None

PR Type

Enhancement, Tests

Description

Introduced a new complex_gaussian_integral_2 function in gaussian_integrals.py that supports batched inputs and uses broadcasting rules for efficient computation.
Enhanced error handling to check for batch size mismatches in the new function.
Added comprehensive tests in test_gaussian_integrals.py to validate the functionality of the new batched Gaussian integral computation, including tests for both batched and non-batched inputs and polynomial c parameter handling.

Changes walkthrough 📝

Relevant files

Enhancement

gaussian_integrals.py `Add batched Gaussian integral computation with broadcasting` mrmustard/physics/gaussian_integrals.py Added a new `complex_gaussian_integral_2` function with support for batched inputs. Enhanced the function to compute Gaussian integrals using broadcasting rules. Improved error handling for batch size mismatches.	+118/-80

Tests

test_gaussian_integrals.py `Add tests for batched Gaussian integral computation` tests/test_physics/test_gaussian_integrals.py Added tests for the new `complex_gaussian_integral_2` function. Included tests for both batched and non-batched inputs. Added tests for polynomial `c` parameter handling.	+51/-3

💡 PR-Agent usage: Comment /help "your question" on any pull request to receive relevant information

codiumai-pr-agent-pro · 2024-09-26T17:25:57Z

PR-Agent was enabled for this repository. To continue using it, please link your git user with your CodiumAI identity here.

PR Reviewer Guide 🔍

⏱️ Estimated effort to review: 3 🔵🔵🔵⚪⚪
🏅 Score: 85
🧪 PR contains tests
🔒 No security concerns identified
⚡ Key issues to review Performance Concern The function creates a large zero matrix and then fills it using a loop, which could be inefficient for large batch sizes. Potential Bug The function assumes that the input tensors have at least 3 dimensions when batched, which may not always be true.

mrmustard/physics/gaussian_integrals.py

apchytr · 2024-09-26T18:24:39Z

mrmustard/physics/gaussian_integrals.py

+    idx1: Sequence[int],
+    idx2: Sequence[int],
+    measure: float = -1,
+    batched: bool = False,


we could avoid needing a batched argument if we enforced Abc1, Abc2 to be batched by default

True, but then one needs to go through hoops in order to use this function. Perhaps I can simply detect if the triple isn't batched and batch it if needed, and finally return it as it was given.

Hmm what hoops exactly? If this is a method we only use internally then things should be batched already (assuming we're consistent with that and if we're not we should fix it). If this is something we expose to users then yeah I'm okay with a batched argument

mrmustard/physics/gaussian_integrals.py

apchytr · 2024-09-26T18:36:23Z

mrmustard/physics/gaussian_integrals.py

+    X = math.block([[Z, I], [I, Z]])
+    M = math.gather(math.gather(A, idx, axis=-1), idx, axis=-2) + X * measure
+    bM = math.gather(b, idx, axis=-1)
+    cpart1 = math.sqrt(math.cast((-1) ** m / math.det(M), "complex128"))


I'm remembering having to introduce some logic for handling the case when math.det(M) is zero (see complex_gaussian_integral). E.g. this was an issue when calling (I believe) .normalize on QuadratureEigenstate. Should we have that here?

I'm not sure. What were you doing in those situations?

E.g. see line 136 in complex_gaussian_integral. I check if det is 0 and if it is handle it by using np.inf

### **User description** **Context:** After [to_bargmann() original bargmann data bug](667c516) we always store the original Bargmann data when calling to_fock. This isn't ideal in the case when c is an array where instead keeping the array is more efficient. **Description of the Change:** `to_fock` checks the polynomial shape and only sets `._original_abc_data` if c is a scalar ___ ### **PR Type** Bug fix, Tests ___ ### **Description** - Fixed a bug in `to_fock` method to conditionally store `_original_abc_data` only when `c` is a scalar. - Added tests to ensure `_original_abc_data` is correctly set or not set based on the polynomial shape. - Improved efficiency by avoiding unnecessary storage of data when `c` is an array. ___ ### **Changes walkthrough** 📝 <table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Bug fix</strong></td><td><table> <tr> <td> <details> <summary><strong>circuit_components.py</strong><dd><code>Conditional storage of `_original_abc_data` based on polynomial shape</code></dd></summary> <hr> mrmustard/lab_dev/circuit_components.py <li>Added a condition to check if <code>polynomial_shape[0]</code> is zero before <br>setting <code>_original_abc_data</code>.<br> <li> Ensured <code>_original_abc_data</code> is only set when the condition is met.<br> </details> </td> <td><a href="https://github.com/XanaduAI/MrMustard/pull/493/files#diff-05cc8f1f470b3eab23b8a8b1b1a628a97c87852722ea6b5408e1bb70efd4ab72">+2/-1</a>      </td> </tr> </table></td></tr><tr><td><strong>Tests</strong></td><td><table> <tr> <td> <details> <summary><strong>test_circuit_components.py</strong><dd><code>Add tests for `_original_abc_data` conditional logic</code>          </dd></summary> <hr> tests/test_lab_dev/test_circuit_components.py <li>Added assertions to verify the correct setting of <code>_original_abc_data</code>.<br> <li> Updated tests to check for <code>None</code> when <code>c</code> is not a scalar.<br> </details> </td> <td><a href="https://github.com/XanaduAI/MrMustard/pull/493/files#diff-cf594d34002781ecbd6f00d55be035f5d6821af5cfde94038830e0e3ea5f5bc3">+4/-2</a>      </td> </tr> </table></td></tr></tr></tbody></table> ___ > 💡 **PR-Agent usage**: Comment `/help "your question"` on any pull request to receive relevant information

**Context:** As part of measurement devices in lab_dev we require a sampler class. The sampler can be though of as a black box that takes in a state and outputs either a probability distribution or a collection of samples. Here I implement said class along with subclasses for PNR samples and Homodyne samples where the sampler is designed in way to easily allow for performance shortcuts. Note: some implementation details are subject to change as we continue to develop measurement devices. **Description of the Change:** `Sampler`, `PNRSampler`, `HomodyneSampler`.

### **User description** **Context:** `Ggate.symplectic` should probably be returning the symplectic matrix it's defined with rather than recomputing. **Description of the Change:** Overrride `symplectic` on `Ggate` to return the symplectic matrix stored in the param set. **Benefits:** Avoid having to recompute the symplectic matrix ___ ### **PR Type** enhancement, tests ___ ### **Description** - Enhanced the `Ggate` class by overriding the `symplectic` property to return the symplectic matrix stored in the parameter set, avoiding recomputation. - Updated the initialization method to use `_add_parameter` for adding the symplectic parameter. - Added a test case to verify that the `symplectic` property returns the correct symplectic matrix. ___ ### **Changes walkthrough** 📝 <table><thead><tr><th></th><th align="left">Relevant files</th></tr></thead><tbody><tr><td><strong>Enhancement</strong></td><td><table> <tr> <td> <details> <summary><strong>ggate.py</strong><dd><code>Override `symplectic` property to return stored matrix</code>      </dd></summary> <hr> mrmustard/lab_dev/transformations/ggate.py <li>Added a property <code>symplectic</code> to return the stored symplectic matrix.<br> <li> Replaced <code>parameter_set.add_parameter</code> with <code>_add_parameter</code>.<br> </details> </td> <td><a href="https://github.com/XanaduAI/MrMustard/pull/495/files#diff-bf96755105326dd7077e8fef650082b18d16970f6f49f0be585a4957ff99d023">+7/-3</a>      </td> </tr> </table></td></tr><tr><td><strong>Tests</strong></td><td><table> <tr> <td> <details> <summary><strong>test_ggate.py</strong><dd><code>Add test for `symplectic` property in `Ggate`</code>                        </dd></summary> <hr> tests/test_lab_dev/test_transformations/test_ggate.py <li>Added an assertion to check if <code>symplectic</code> returns the correct matrix.<br> </details> </td> <td><a href="https://github.com/XanaduAI/MrMustard/pull/495/files#diff-e9c4d0c2193b280b6aa6f4c21a02bc9dc4161ce629671f0be05d20a4cb5bfd76">+1/-0</a>      </td> </tr> </table></td></tr></tr></tbody></table> ___ > 💡 **PR-Agent usage**: Comment `/help "your question"` on any pull request to receive relevant information

ziofil · 2024-10-11T17:24:45Z

See #502

added new broadcasted gaussian integral and tests

1ba1d00

ziofil requested a review from apchytr September 26, 2024 17:24

ziofil added the no changelog Pull request does not require a CHANGELOG entry label Sep 26, 2024

codiumai-pr-agent-pro bot added enhancement New feature or request tests Review effort [1-5]: 3 labels Sep 26, 2024

ziofil added 2 commits September 26, 2024 10:26

removed unused variable

7f21c6b

removed unnecessary parenthesis

2d5a33c

XanaduAI deleted a comment from codiumai-pr-agent-pro bot Sep 26, 2024

apchytr reviewed Sep 26, 2024

View reviewed changes

mrmustard/physics/gaussian_integrals.py Outdated Show resolved Hide resolved

apchytr reviewed Sep 26, 2024

View reviewed changes

mrmustard/physics/gaussian_integrals.py Outdated Show resolved Hide resolved

apchytr reviewed Sep 26, 2024

View reviewed changes

mrmustard/physics/gaussian_integrals.py Outdated Show resolved Hide resolved

apchytr reviewed Sep 26, 2024

View reviewed changes

ziofil and others added 10 commits September 27, 2024 10:24

np -> math

a1cbb17

removed c part, avoided loop to fill A

ec38baa

WIP

8fb1bee

works!

44ecf5f

fixed tests

1a57207

tmp

26bca9c

stash

888804a

ziofil closed this Oct 11, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Broadcasted gaussian integral #494

Broadcasted gaussian integral #494

ziofil commented Sep 26, 2024 •

edited by codiumai-pr-agent-pro bot

Loading

codiumai-pr-agent-pro bot commented Sep 26, 2024

apchytr Sep 26, 2024

ziofil Sep 27, 2024

apchytr Oct 3, 2024

apchytr Sep 26, 2024

ziofil Sep 27, 2024

apchytr Oct 3, 2024

ziofil commented Oct 11, 2024

Broadcasted gaussian integral #494

Broadcasted gaussian integral #494

Conversation

ziofil commented Sep 26, 2024 • edited by codiumai-pr-agent-pro bot Loading

User description

PR Type

Description

Changes walkthrough 📝

codiumai-pr-agent-pro bot commented Sep 26, 2024

PR Reviewer Guide 🔍

apchytr Sep 26, 2024

Choose a reason for hiding this comment

ziofil Sep 27, 2024

Choose a reason for hiding this comment

apchytr Oct 3, 2024

Choose a reason for hiding this comment

apchytr Sep 26, 2024

Choose a reason for hiding this comment

ziofil Sep 27, 2024

Choose a reason for hiding this comment

apchytr Oct 3, 2024

Choose a reason for hiding this comment

ziofil commented Oct 11, 2024

ziofil commented Sep 26, 2024 •

edited by codiumai-pr-agent-pro bot

Loading