XanaduAI / MrMustard

A differentiable bridge between phase space and Fock space
https://mrmustard.readthedocs.io/
Apache License 2.0
78 stars 27 forks source link

added mm_einsum #515

Open ziofil opened 3 weeks ago

ziofil commented 3 weeks ago

User description

Context: mm_einsum for contracting a TN made of CV and Fock components

Description of the Change: For now just the mm_einsum.py file


PR Type

enhancement, documentation


Description


PRDescriptionHeader.CHANGES_WALKTHROUGH

Relevant files
Enhancement
mm_einsum.py
Implement `mm_einsum` for tensor network contraction         

mrmustard/physics/mm_einsum.py
  • Implemented mm_einsum for contracting tensor networks with CV and Fock
    components.
  • Added functions for calculating contraction costs and optimal paths.
  • Included detailed docstrings for all functions.
  • Utilized Numba for performance optimization.
  • +211/-0 

    πŸ’‘ PR-Agent usage: Comment /help "your question" on any pull request to receive relevant information

    codiumai-pr-agent-pro[bot] commented 3 weeks ago

    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: 92
    πŸ§ͺ No relevant tests
    πŸ”’ No security concerns identified
    ⚑ Recommended focus areas for review

    Performance Optimization
    The `optimal` function uses an exhaustive search algorithm which may become inefficient for large tensor networks. Consider implementing a heuristic approach or using dynamic programming for better scalability. Error Handling
    The code lacks explicit error handling for potential edge cases, such as empty input lists or invalid Fock dimensions. Consider adding appropriate error checks and exception handling.
    codiumai-pr-agent-pro[bot] commented 3 weeks ago

    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:

    CategorySuggestion                                                                                                                                    Score
    Performance
    Use built-in Python functions for simple mathematical operations to reduce dependencies and improve performance ___ **Consider using math.prod() instead of np.prod() for calculating the product of
    integers. It's more efficient for small sequences and doesn't require importing
    NumPy.** [mrmustard/physics/mm_einsum.py [86]](https://github.com/XanaduAI/MrMustard/pull/515/files#diff-4b9ec38e891b9e6fb6cdd434e05a4d92a2580713dc019fcfe6da438d35c535beR86-R86) ```diff -fock_flops = np.prod(fock_contracted_shape) * np.prod(fock_remaining_shape) +from math import prod +fock_flops = prod(fock_contracted_shape) * prod(fock_remaining_shape) ``` - [ ] **Apply this suggestion**
    Suggestion importance[1-10]: 7 Why: The suggestion to use `math.prod()` instead of `np.prod()` is valid as it reduces dependency on NumPy for simple integer operations, which can improve performance and readability for small sequences.
    7
    Best practice
    Use descriptive variable names to enhance code readability and self-documentation ___ **Consider using a more descriptive variable name instead of m in the _CV_flops
    function to improve code readability.** [mrmustard/physics/mm_einsum.py [21-28]](https://github.com/XanaduAI/MrMustard/pull/515/files#diff-4b9ec38e891b9e6fb6cdd434e05a4d92a2580713dc019fcfe6da438d35c535beR21-R28) ```diff @njit -def _CV_flops(nA: int, nB: int, m: int) -> int: +def _CV_flops(nA: int, nB: int, num_contracted_cv: int) -> int: """Calculate the cost of contracting two tensors with CV indices. Args: nA: Number of CV indices in the first tensor nB: Number of CV indices in the second tensor - m: Number of CV indices involved in the contraction + num_contracted_cv: Number of CV indices involved in the contraction """ ``` - [ ] **Apply this suggestion**
    Suggestion importance[1-10]: 6 Why: Renaming the variable `m` to `num_contracted_cv` improves code readability by making the purpose of the variable clearer, which is a good practice for maintainability.
    6
    Use named constants for magic values to improve code clarity and maintainability ___ **Consider using a constant for the infinity value in the optimal function instead of
    float("inf") to improve readability and maintainability.** [mrmustard/physics/mm_einsum.py [162]](https://github.com/XanaduAI/MrMustard/pull/515/files#diff-4b9ec38e891b9e6fb6cdd434e05a4d92a2580713dc019fcfe6da438d35c535beR162-R162) ```diff -best_flops: int = float("inf") +import math +INFINITY = math.inf +best_flops: int = INFINITY + ``` - [ ] **Apply this suggestion**
    Suggestion importance[1-10]: 5 Why: Using a named constant for infinity improves code readability and maintainability by avoiding magic numbers, although the impact is relatively minor.
    5
    Enhancement
    Use defaultdict to simplify dictionary access and avoid explicit key existence checks ___ **Consider using collections.defaultdict for fock_size_dict to simplify the logic in
    the attempt_decomposition function and avoid potential KeyError exceptions.** [mrmustard/physics/mm_einsum.py [100-114]](https://github.com/XanaduAI/MrMustard/pull/515/files#diff-4b9ec38e891b9e6fb6cdd434e05a4d92a2580713dc019fcfe6da438d35c535beR100-R114) ```diff +from collections import defaultdict + def attempt_decomposition( - indices: set[int], fock_size_dict: dict[int, int] + indices: set[int], fock_size_dict: defaultdict[int, int] ) -> tuple[set[int], int]: """Attempt to reduce the number of indices by combining Fock indices when possible. Only possible if there is only one CV index and multiple Fock indices. Args: indices: Set of indices to potentially decompose - fock_size_dict: Dictionary mapping indices to their sizes + fock_size_dict: DefaultDict mapping indices to their sizes, with a default value of 0 Returns: Tuple of (decomposed indices, cost of decomposition) """ - fock_indices_shape = [fock_size_dict[idx] for idx in indices if idx in fock_size_dict] - cv_indices = [idx for idx in indices if idx not in fock_size_dict] + fock_indices_shape = [fock_size_dict[idx] for idx in indices if fock_size_dict[idx] != 0] + cv_indices = [idx for idx in indices if fock_size_dict[idx] == 0] ``` - [ ] **Apply this suggestion**
    Suggestion importance[1-10]: 4 Why: The suggestion to use `defaultdict` can simplify the code by removing the need for explicit key checks, but it introduces a new dependency and changes the logic slightly, which may not be necessary if the current implementation is already clear and correct.
    4

    πŸ’‘ Need additional feedback ? start a PR chat

    codecov[bot] commented 3 weeks ago

    Codecov Report

    Attention: Patch coverage is 0% with 59 lines in your changes missing coverage. Please review.

    Project coverage is 88.87%. Comparing base (e5c02ce) to head (01b56e7).

    Files with missing lines Patch % Lines
    mrmustard/physics/mm_einsum.py 0.00% 59 Missing :warning:
    Additional details and impacted files ```diff @@ Coverage Diff @@ ## develop #515 +/- ## =========================================== - Coverage 89.77% 88.87% -0.90% =========================================== Files 92 93 +1 Lines 6054 6113 +59 =========================================== - Hits 5435 5433 -2 - Misses 619 680 +61 ``` | [Files with missing lines](https://app.codecov.io/gh/XanaduAI/MrMustard/pull/515?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI) | Coverage Ξ” | | |---|---|---| | [mrmustard/physics/mm\_einsum.py](https://app.codecov.io/gh/XanaduAI/MrMustard/pull/515?src=pr&el=tree&filepath=mrmustard%2Fphysics%2Fmm_einsum.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI#diff-bXJtdXN0YXJkL3BoeXNpY3MvbW1fZWluc3VtLnB5) | `0.00% <0.00%> (ΓΈ)` | | ... and [1 file with indirect coverage changes](https://app.codecov.io/gh/XanaduAI/MrMustard/pull/515/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI) ------ [Continue to review full report in Codecov by Sentry](https://app.codecov.io/gh/XanaduAI/MrMustard/pull/515?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI). > **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI) > `Ξ” = absolute (impact)`, `ΓΈ = not affected`, `? = missing data` > Powered by [Codecov](https://app.codecov.io/gh/XanaduAI/MrMustard/pull/515?dropdown=coverage&src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI). Last update [e5c02ce...01b56e7](https://app.codecov.io/gh/XanaduAI/MrMustard/pull/515?dropdown=coverage&src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=XanaduAI).