TensorBFS / OMEinsumContractionOrders.jl

Tensor network contraction order optimizers for OMEinsum
MIT License
25 stars 3 forks source link

Added a solver based on exact tree width solver #43

Closed ArrogantGao closed 3 months ago

ArrogantGao commented 3 months ago

Added a solver named ExactTreewidth, which is based on the exact tree width solve provided in TreeWidthSolver.jl.

ArrogantGao commented 3 months ago

Added a function tree_reformulate, which can remove a tensor from a given binary contraction order without changing the space complexity of the contraction. Thus, when searching contraction orders of tensor networks with open edges using graph based methods (balanced min cut or tree width), one can add a tensor connecting with all open edges to get a graph with no open edges. Then by remove this manually added tensor, one can get a contraction order of the original tn.

Here is an example:

julia> using OMEinsumContractionOrders, OMEinsum, KaHyPar

julia> using OMEinsumContractionOrders: optimize_kahypar

# given a contraction with open edges
julia> ein_open = OMEinsum.rawcode(ein"ijl, jkm, ik -> lm")
ijl, jkm, ik -> lm

# add a dummy tensor manually
julia> ein_closed = OMEinsum.rawcode(ein"ijl, jkm, ik, lm ->")
ijl, jkm, ik, lm ->

# optimize
julia> opt_ein_closed = optimize_kahypar(ein_closed, uniformsize(ein_closed, 2); max_group_size=10, sc_target=10)
lm, lm ->
├─ jlk, jkm -> lm
│  ├─ ijl, ik -> jlk
│  │  ├─ ijl
│  │  └─ ik
│  └─ jkm
└─ lm

# remove the dummy tensor manually
julia> opt_ein_open = OMEinsumContractionOrders.tree_reformulate(opt_ein_closed, 4)
jlk, jkm -> lm
├─ ijl, ik -> jlk
│  ├─ ijl
│  └─ ik
└─ jkm

julia> OMEinsum.contraction_complexity(opt_ein_closed, uniformsize(ein_closed, 2))
Time complexity: 2^5.169925001442312
Space complexity: 2^3.0
Read-write complexity: 2^5.614709844115208

julia> OMEinsum.contraction_complexity(opt_ein_open, uniformsize(ein_closed, 2))
Time complexity: 2^5.0
Space complexity: 2^3.0
Read-write complexity: 2^5.321928094887363

# the process above has been implemented in optimize_kahypar, one can directly optimize a contraction with open edges
julia> opt_ein_open_direct = optimize_kahypar(ein_open, uniformsize(ein_closed, 2); max_group_size=10, sc_target=10)
jlk, jkm -> lm
├─ ijl, ik -> jlk
│  ├─ ijl
│  └─ ik
└─ jkm

julia> opt_ein_open == opt_ein_open_direct
true
GiggleLiu commented 3 months ago

tree_reformulate is not very intuitive to understand to me, what about using another name, like pivot_tree?

ArrogantGao commented 3 months ago

name of the function has been changed as pivot_tree

GiggleLiu commented 3 months ago

Ready for review?

ArrogantGao commented 3 months ago

Sure! The new package TreeWidthSolver.jl has been registered and the PR is ready for review.

codecov[bot] commented 3 months ago

Codecov Report

Attention: Patch coverage is 92.80000% with 9 lines in your changes missing coverage. Please review.

Project coverage is 94.47%. Comparing base (cc2295f) to head (494aa8c). Report is 1 commits behind head on master.

Files Patch % Lines
src/Core.jl 86.36% 6 Missing :warning:
src/interfaces.jl 0.00% 2 Missing :warning:
src/treewidth.jl 98.59% 1 Missing :warning:
Additional details and impacted files ```diff @@ Coverage Diff @@ ## master #43 +/- ## ========================================== + Coverage 93.16% 94.47% +1.30% ========================================== Files 13 14 +1 Lines 1097 1230 +133 ========================================== + Hits 1022 1162 +140 + Misses 75 68 -7 ```

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

GiggleLiu commented 3 months ago

Good job!

ArrogantGao commented 3 months ago

Thanks a lot! It is much better now.