New broadcasted gaussian integrals

[ziofil]

User description

Cleaned up and refactored Gaussian integrals, now working with batched, non batched and polynomial inputs

---

PR Type

enhancement, tests

---

Description

• Refactored Gaussian integral functions by removing deprecated functions and introducing complex_gaussian_integral_1 and complex_gaussian_integral_2.
• Enhanced join_Abc function to support batch processing and improved input validation.
• Updated various modules to utilize the new Gaussian integral functions.
• Revised tests to cover new functionality and removed tests for deprecated functions.

---

Changes walkthrough 📝

	Relevant files
Enhancement	gaussian_integrals.py Refactor Gaussian integral functions and enhance batch processing mrmustard/physics/gaussian_integrals.py Removed complex_gaussian_integral and contract_two_Abc functions. Introduced complex_gaussian_integral_1 and complex_gaussian_integral_2 functions. Enhanced join_Abc function to support batch processing. Improved input validation and error handling. +253/-266 representations.py Update Gaussian integral usage in representations                mrmustard/physics/representations.py Replaced complex_gaussian_integral with complex_gaussian_integral_1. Updated matrix multiplication to use complex_gaussian_integral_2. +9/-17    triples.py Update triples to use new Gaussian integral function          mrmustard/physics/triples.py Replaced contract_two_Abc with complex_gaussian_integral_2. +5/-7
Tests	test_circuit_components_utils.py Update circuit components utils tests for new integrals    tests/test_lab_dev/test_circuit_components_utils.py Updated tests to use complex_gaussian_integral_1 and complex_gaussian_integral_2. +7/-7      test_gaussian_integrals.py Add tests for new Gaussian integral functions                        tests/test_physics/test_gaussian_integrals.py Added tests for complex_gaussian_integral_1 and complex_gaussian_integral_2. Removed tests for deprecated functions. +133/-72 test_representations.py Update representation tests for new integrals                        tests/test_physics/test_representations.py Updated tests to use complex_gaussian_integral_1 and complex_gaussian_integral_2. +4/-4

---

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

[codiumai-pr-agent-pro]

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

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🏅 Score: 85

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Performance Optimization The new complex_gaussian_integral_1 and complex_gaussian_integral_2 functions might benefit from further optimization, especially for large-scale computations. Code Duplication There's some code duplication in the join_Abc function for different modes. Consider refactoring to reduce redundancy.

[codiumai-pr-agent-pro]

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

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Score
Best practice	✅ Implement error handling for potential exceptions during matrix operations Consider using a context manager or a try-except block to handle potential exceptions that might occur during matrix operations, especially when dealing with potentially singular matrices. mrmustard/physics/gaussian_integrals.py [405-416] -if np.all(math.abs(determinant)) > 1e-12: +try: + inv_M = math.inv(M) c_post = ( c * math.sqrt(math.cast((-1) ** m / determinant, "complex128")) * math.exp(-0.5 * math.sum(bM * math.solve(M, bM), axes=[-1])) ) - A_post = R - math.einsum("bij,bjk,blk->bil", D, math.inv(M), D) + A_post = R - math.einsum("bij,bjk,blk->bil", D, inv_M, D) b_post = bR - math.einsum("bij,bj->bi", D, math.solve(M, bM)) -else: +except np.linalg.LinAlgError: A_post = R - math.einsum("bij,bjk,blk->bil", D, M * np.inf, D) b_post = bR - math.einsum("bij,bjk,bk->bi", D, M * np.inf, bM) c_post = math.real(c) * np.inf [Suggestion has been applied] Suggestion importance[1-10]: 9 Why: Introducing error handling for matrix operations can prevent runtime errors due to singular matrices, enhancing the robustness and reliability of the code. This is an important improvement for handling edge cases.	9
	✅ Use a more robust method for checking if a value is close to zero Consider using a more robust method for checking if the determinant is close to zero, such as np.allclose() with a specified tolerance, instead of directly comparing with a hard-coded value. mrmustard/physics/gaussian_integrals.py [405-416] -if np.all(math.abs(determinant)) > 1e-12: +if np.allclose(determinant, 0, atol=1e-12): + A_post = R - math.einsum("bij,bjk,blk->bil", D, M * np.inf, D) + b_post = bR - math.einsum("bij,bjk,bk->bi", D, M * np.inf, bM) + c_post = math.real(c) * np.inf +else: c_post = ( c * math.sqrt(math.cast((-1) ** m / determinant, "complex128")) * math.exp(-0.5 * math.sum(bM * math.solve(M, bM), axes=[-1])) ) A_post = R - math.einsum("bij,bjk,blk->bil", D, math.inv(M), D) b_post = bR - math.einsum("bij,bj->bi", D, math.solve(M, bM)) -else: - A_post = R - math.einsum("bij,bjk,blk->bil", D, M * np.inf, D) - b_post = bR - math.einsum("bij,bjk,bk->bi", D, M * np.inf, bM) - c_post = math.real(c) * np.inf [Suggestion has been applied] Suggestion importance[1-10]: 8 Why: Using np.allclose() with a specified tolerance is a more robust and numerically stable method for checking if a determinant is close to zero, which can prevent potential issues with floating-point precision.	8
Enhancement	Add error handling tests for complex Gaussian integral functions to ensure they handle invalid inputs correctly Consider adding error handling tests for complex_gaussian_integral_1 and complex_gaussian_integral_2 to ensure they handle invalid inputs correctly. tests/test_physics/test_gaussian_integrals.py [254-273] def test_complex_gaussian_integral_1_not_batched(): """Tests the ``complex_gaussian_integral_1`` method for non-batched inputs.""" A, b, c = triples.thermal_state_Abc(nbar=[0.5, 0.9, 1.0]) Ar, br, cr = triples.vacuum_state_Abc(0) res = complex_gaussian_integral_1((A, b, c), [0, 2, 4], [1, 3, 5]) assert np.allclose(res[0], Ar) assert np.allclose(res[1], br) assert np.allclose(res[2], cr) A1, b1, c1 = triples.vacuum_state_Abc(2) A2, b2, c2 = triples.displacement_gate_Abc(x=[0.1, 0.2], y=0.3) A3, b3, c3 = triples.displaced_squeezed_vacuum_state_Abc(x=[0.1, 0.2], y=0.3) A, b, c = join_Abc((A1, b1, c1), (A2, b2, c2), mode="zip") res = complex_gaussian_integral_1((A, b, c), [0, 1], [4, 5]) assert np.allclose(res[0], A3) assert np.allclose(res[1], b3) assert np.allclose(res[2], c3) + # Test error handling + with pytest.raises(ValueError): + complex_gaussian_integral_1((A, b, c), [0, 1], [4, 5, 6]) # Mismatched idx_z and idx_zconj + + with pytest.raises(ValueError): + complex_gaussian_integral_1((A, b, c), [0, 1, 10], [4, 5, 11]) # Invalid indices + +def test_complex_gaussian_integral_2_error_handling(): + """Tests error handling for the ``complex_gaussian_integral_2`` method.""" + A1, b1, c1 = triples.vacuum_state_Abc(2) + A2, b2, c2 = triples.displacement_gate_Abc(x=[0.1, 0.2], y=0.3) + + with pytest.raises(ValueError): + complex_gaussian_integral_2((A1, b1, c1), (A2, b2, c2), [0, 1], [2, 3, 4]) # Mismatched idx1 and idx2 + + with pytest.raises(ValueError): + complex_gaussian_integral_2((A1, b1, c1), (A2, b2, c2), [0, 1, 10], [2, 3, 11]) # Invalid indices + + with pytest.raises(ValueError): + complex_gaussian_integral_2((A1, b1, c1), (A2, b2, c2), [0, 1], [2, 3], mode="invalid_mode") # Invalid mode + Apply this suggestion Suggestion importance[1-10]: 8 Why: Adding error handling tests is a valuable enhancement that ensures robustness by verifying the functions' behavior with invalid inputs. This increases the reliability of the code.	8
	Refactor the batched join_Abc tests using parameterization to reduce code duplication and improve test coverage Consider using a parameterized test for test_join_Abc_batched_zip and test_join_Abc_batched_kron to reduce code duplication and improve test coverage. tests/test_physics/test_gaussian_integrals.py [145-192] -def test_join_Abc_batched_zip(): - """Tests the ``join_Abc`` method for batched inputs in zip mode (and with polynomial c).""" +@pytest.mark.parametrize("mode, expected_A, expected_b, expected_c", [ + ("zip", np.array([ + [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 8, 9], [0, 0, 10, 11]], + [[5, 6, 0, 0], [7, 8, 0, 0], [0, 0, 12, 13], [0, 0, 14, 15]] + ]), np.array([[5, 6, 12, 13], [7, 8, 14, 15]]), np.array([70, 800])), + ("kron", np.array([ + [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 8, 9], [0, 0, 10, 11]], + [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 12, 13], [0, 0, 14, 15]] + ]), np.array([[5, 6, 12, 13], [5, 6, 14, 15]]), np.array([70, 700])) +]) +def test_join_Abc_batched(mode, expected_A, expected_b, expected_c): + """Tests the ``join_Abc`` method for batched inputs in different modes.""" A1 = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]]) b1 = np.array([[5, 6], [7, 8]]) c1 = np.array([7, 8]) A2 = np.array([[[8, 9], [10, 11]], [[12, 13], [14, 15]]]) b2 = np.array([[12, 13], [14, 15]]) c2 = np.array([10, 100]) - A, b, c = join_Abc((A1, b1, c1), (A2, b2, c2), mode="zip") + if mode == "kron": + A1 = A1[:1] + b1 = b1[:1] + c1 = c1[:1] - assert np.allclose( - A, - np.array( - [ - [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 8, 9], [0, 0, 10, 11]], - [[5, 6, 0, 0], [7, 8, 0, 0], [0, 0, 12, 13], [0, 0, 14, 15]], - ] - ), - ) - assert np.allclose(b, np.array([[5, 6, 12, 13], [7, 8, 14, 15]])) - assert np.allclose(c, np.array([70, 800])) + A, b, c = join_Abc((A1, b1, c1), (A2, b2, c2), mode=mode) + assert np.allclose(A, expected_A) + assert np.allclose(b, expected_b) + assert np.allclose(c, expected_c) -def test_join_Abc_batched_kron(): - """Tests the ``join_Abc`` method for batched inputs in kron mode (and with polynomial c).""" - A1 = np.array([[[1, 2], [3, 4]]]) - b1 = np.array([[5, 6]]) - c1 = np.array([7]) - - A2 = np.array([[[8, 9], [10, 11]], [[12, 13], [14, 15]]]) - b2 = np.array([[12, 13], [14, 15]]) - c2 = np.array([10, 100]) - - A, b, c = join_Abc((A1, b1, c1), (A2, b2, c2), mode="kron") - - assert np.allclose( - A, - np.array( - [ - [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 8, 9], [0, 0, 10, 11]], - [[1, 2, 0, 0], [3, 4, 0, 0], [0, 0, 12, 13], [0, 0, 14, 15]], - ] - ), - ) - assert np.allclose(b, np.array([[5, 6, 12, 13], [5, 6, 14, 15]])) - assert np.allclose(c, np.array([70, 700])) - Apply this suggestion Suggestion importance[1-10]: 7 Why: The suggestion to use parameterized tests is valid as it reduces code duplication and enhances test coverage. This improves maintainability and readability of the test suite.	7
Performance	Utilize numpy's broadcasting to simplify array operations and improve performance Consider using numpy's broadcasting capabilities instead of explicitly repeating and tiling arrays. This can lead to more efficient and readable code. mrmustard/physics/gaussian_integrals.py [270-275] -A1 = np.repeat(A1, batch2, axis=0) -A2 = np.tile(A2, (batch1, 1, 1)) -A1Z = np.concatenate([A1, np.zeros((batch1 * batch2, n1_plus_m1, n2_plus_m2))], axis=-1) -ZA2 = np.concatenate([np.zeros((batch1 * batch2, n2_plus_m2, n1_plus_m1)), A2], axis=-1) -b1 = np.repeat(b1, batch2, axis=0) -b2 = np.tile(b2, (batch1, 1)) +A1 = A1[:, np.newaxis, :, :] +A2 = A2[np.newaxis, :, :, :] +A1Z = np.concatenate([A1, np.zeros((batch1, batch2, n1_plus_m1, n2_plus_m2))], axis=-1) +ZA2 = np.concatenate([np.zeros((batch1, batch2, n2_plus_m2, n1_plus_m1)), A2], axis=-1) +b1 = b1[:, np.newaxis, :] +b2 = b2[np.newaxis, :, :] Apply this suggestion Suggestion importance[1-10]: 7 Why: The suggestion to use numpy's broadcasting can improve both the readability and performance of the code by eliminating the need for explicit repetition and tiling of arrays. This is a valid optimization for the given context.	7

💡 Need additional feedback ? start a PR chat

[codecov]

Codecov Report

Attention: Patch coverage is 99.08257% with 1 line in your changes missing coverage. Please review.

Project coverage is 89.77%. Comparing base (f0077bc) to head (029fcbb).
Report is 1 commits behind head on develop.

Files with missing lines	Patch %	Lines
mrmustard/physics/gaussian_integrals.py	99.02%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop     #502      +/-   ##
===========================================
+ Coverage    89.76%   89.77%   +0.01%     
===========================================
  Files          104      104              
  Lines         7639     7639              
===========================================
+ Hits          6857     6858       +1     
+ Misses         782      781       -1     

Files with missing lines	Coverage Δ
mrmustard/math/backend_tensorflow.py	100.00% <ø> (ø)
mrmustard/physics/representations.py	98.51% <100.00%> (-0.08%)	⬇️
mrmustard/physics/triples.py	100.00% <100.00%> (ø)
mrmustard/physics/gaussian_integrals.py	99.40% <99.02%> (+0.69%)	⬆️

---

Continue to review full report in Codecov by Sentry.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f0077bc...029fcbb. Read the comment docs.