Update wigner.py

[JakeKitchen]

User description

Optimized Wigner Function Iterative Computation

Context:

The _wigner_discretized_iterative function calculates the Wigner function iteratively,
which is essential for simulating quantum states in continuous variable quantum mechanics.
The original implementation had redundant operations, recalculations, and excessive memory
allocations, leading to performance bottlenecks. The make_grid function was also optimized
with @njit to further enhance performance. This update improves both mathematical efficiency
and memory management.

Description of the Change:

1. Precompute Exponential Terms: The exponential grid is computed only once and reused to
   avoid redundant calculations in each iteration.
2. Precompute Square Roots: All square roots required for the computation are calculated at
   the beginning and stored in an array. This reduces overhead inside nested loops.
3. Efficient Matrix Swapping: Instead of creating new matrices in every loop iteration, the
   function reuses memory by swapping matrices, which reduces memory allocations.
4. Optimized make_grid Function:
   • The make_grid function was updated to use @njit to accelerate its performance.
   • Memory allocations for the Q and P coordinate matrices are optimized to reduce overhead.
   • A precomputed sqrt_factor is used to avoid redundant calculations inside nested loops.
   • By leveraging @njit, the function is now compliant with Numba’s JIT compilation, allowing
     faster and more efficient execution.

Benefits:

1. Improved Mathematical Efficiency:

   • The function minimizes the number of floating-point operations inside loops by precomputing
     values such as exponential terms and square roots.
   • The iterative Wigner calculation avoids recalculating these values multiple times, making
     the function more computationally efficient.
   • Memory-intensive operations are reduced by swapping matrices rather than allocating new ones,
     ensuring cache-friendlier execution.

2. Better Performance:

   • Memory Reuse: By reusing the same memory blocks with matrix swapping, the function avoids
     costly memory allocations, improving speed.
   • Optimized make_grid: Using @njit in make_grid improves both speed and memory efficiency,
     enabling seamless integration with the rest of the optimized functions.

3. Numba Compliance:

   • The optimized function ensures type safety and alignment with Numba’s @njit constraints,
     avoiding common type-related issues.

Possible Drawbacks:

• Increased Memory Usage at Startup: Precomputing the square roots and exponential grid slightly
  increases memory usage at the beginning, but this is offset by performance gains during execution.

Related GitHub Issues:

NA

---

PR Type

Enhancement

---

Description

• Optimized the iterative Wigner function computation for efficiency

  • Precomputed exponential and square root terms to reduce redundant calculations
  • Improved memory usage by reusing and swapping matrices

• Enhanced make_grid function for better performance and Numba compatibility

  • Used explicit loops and precomputed factors for faster grid creation

---

Changes walkthrough 📝

Relevant files

Enhancement

wigner.py Optimize Wigner function and grid computation for speed and memory mrmustard/physics/wigner.py Refactored make_grid to use explicit loops and precomputed factors for Numba JIT Precomputed exponential and square root terms in _wigner_discretized_iterative Reduced memory allocations by swapping matrices instead of creating new ones Improved overall mathematical and computational efficiency of Wigner calculation +34/-20

---

Need help?

Type /help how to ... in the comments thread for any questions about Qodo Merge usage.

Check out the documentation for more information.

[qodo-merge-pro]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

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

🏅 Score: 92

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Missing Initialization The initialization of W for m>0 is missing. The code adds to W but doesn't initialize W[m,m] component before the inner loop. W += rho[m, m].real * Wmat[1, m].real

[qodo-merge-pro]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Possible issue	Prevent numerical instability There's a potential numerical instability when m is very small, as division by sqrt_n[m] could lead to large values. Add a check to handle the case where sqrt_n[m] is close to zero to prevent potential overflow. mrmustard/physics/wigner.py [176-179] # Compute the remaining coefficients and accumulate the Wigner function. for m in range(1, cutoff): - Wmat[1, m] = (2 * np.conj(grid) * Wmat[0, m] - sqrt_n[m] * Wmat[0, m - 1]) / sqrt_n[m] + if sqrt_n[m] > 1e-10: # Avoid division by very small numbers + Wmat[1, m] = (2 * np.conj(grid) * Wmat[0, m] - sqrt_n[m] * Wmat[0, m - 1]) / sqrt_n[m] + else: + Wmat[1, m] = 2 * np.conj(grid) * Wmat[0, m] # Simplified when sqrt(m) ≈ 0 W += rho[m, m].real * Wmat[1, m].real Apply / Chat Suggestion importance[1-10]: 7 __ Why: Adding a check to avoid division by very small numbers in sqrt_n[m] is a reasonable precaution for numerical stability, though in practice m starts at 1 and sqrt(1) is not problematic. The suggestion is valid but its practical impact is limited.	Medium
	Validate input parameter The function should validate that hbar is positive before computing sqrt_factor to prevent potential runtime errors or complex number results when calculating the square root of a negative number. mrmustard/physics/wigner.py [28-38] @njit def make_grid(q_vec: np.ndarray, p_vec: np.ndarray, hbar: float): r"""Returns two coordinate matrices `Q` and `P` from coordinate vectors `q_vec` and `p_vec`, along with the grid over which Wigner functions can be discretized. """ + if hbar <= 0: + raise ValueError("hbar must be positive") n_q = q_vec.size n_p = p_vec.size Q = np.empty((n_q, n_p), dtype=np.float64) P = np.empty((n_q, n_p), dtype=np.float64) sqrt_factor = 1.0 / np.sqrt(2.0 * hbar) Apply / Chat Suggestion importance[1-10]: 7 __ Why: Adding input validation for hbar being positive is a good defensive programming practice, preventing potential runtime errors or unexpected behavior. This improves robustness but is not critical for correctness if the function is always called with valid hbar.	Medium
General	Use precomputed value directly The initialization of W should use exp_grid directly instead of Wmat[0, 0].real since they are identical. This avoids unnecessary memory access and ensures consistency with the precomputed value. mrmustard/physics/wigner.py [161-166] # Precompute the exponential term to avoid recalculating it. exp_grid = np.exp(-2.0 * np.abs(grid) ** 2) / np.pi # Initialize Wmat and W with the |0><0| component. Wmat[0, 0] = exp_grid -W = rho[0, 0].real * Wmat[0, 0].real +W = rho[0, 0].real * exp_grid.real Apply / Chat Suggestion importance[1-10]: 4 __ Why: This suggestion makes a minor optimization by using the already computed exp_grid instead of accessing Wmat[0, 0].real, which is functionally identical. The impact on correctness and performance is minimal.	Low
Update

[codecov]

Codecov Report

Attention: Patch coverage is 30.00000% with 14 lines in your changes missing coverage. Please review.

Project coverage is 90.53%. Comparing base (a095015) to head (baaff4d).

Files with missing lines	Patch %	Lines
mrmustard/physics/wigner.py	30.00%	14 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop     #593      +/-   ##
===========================================
- Coverage    90.73%   90.53%   -0.21%     
===========================================
  Files          101      101              
  Lines         5777     5798      +21     
===========================================
+ Hits          5242     5249       +7     
- Misses         535      549      +14     

Files with missing lines	Coverage Δ
mrmustard/physics/wigner.py	56.25% <30.00%> (-43.75%)	⬇️

---

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 a095015...baaff4d. Read the comment docs.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.