Jg/switch to extension (#39)

* switch to extensions from Requires

* require julia 1.9

* bump version

* update README and improve complexity printing

CITATION.bib

@@ -0,0 +1,51 @@
+@misc{Liu2022,
+  doi = {10.48550/ARXIV.2205.03718},
+  url = {https://arxiv.org/abs/2205.03718},
+  author = {Liu, Jin-Guo and Gao, Xun and Cain, Madelyn and Lukin, Mikhail D. and Wang, Sheng-Tao},
+  keywords = {Statistical Mechanics (cond-mat.stat-mech), FOS: Physical sciences, FOS: Physical sciences},
+  title = {Computing solution space properties of combinatorial optimization problems via generic tensor networks},
+  publisher = {arXiv},
+  year = {2022},
+  copyright = {Creative Commons Attribution 4.0 International}
+}
+@misc{Pan2021,
+      title={Simulating the Sycamore quantum supremacy circuits}, 
+      author={Feng Pan and Pan Zhang},
+      year={2021},
+      eprint={2103.03074},
+      archivePrefix={arXiv},
+      primaryClass={quant-ph}
+}
+@Article{10.21468/SciPostPhys.7.5.060,
+	title={{Fast counting with tensor networks}},
+	author={Stefanos Kourtis and Claudio Chamon and Eduardo R. Mucciolo and Andrei E. Ruckenstein},
+	journal={SciPost Phys.},
+	volume={7},
+	issue={5},
+	pages={60},
+	year={2019},
+	publisher={SciPost},
+	doi={10.21468/SciPostPhys.7.5.060},
+	url={https://scipost.org/10.21468/SciPostPhys.7.5.060},
+}
+@article{Gray2021,
+   title={Hyper-optimized tensor network contraction},
+   volume={5},
+   ISSN={2521-327X},
+   url={http://dx.doi.org/10.22331/q-2021-03-15-410},
+   DOI={10.22331/q-2021-03-15-410},
+   journal={Quantum},
+   publisher={Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften},
+   author={Gray, Johnnie and Kourtis, Stefanos},
+   year={2021},
+   month={Mar},
+   pages={410}
+}
+@misc{kalachev2021recursive,
+      title={Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits}, 
+      author={Gleb Kalachev and Pavel Panteleev and Man-Hong Yung},
+      year={2021},
+      eprint={2108.05665},
+      archivePrefix={arXiv},
+      primaryClass={quant-ph}
+}

Project.toml

@@ -1,31 +1,35 @@
 name = "OMEinsumContractionOrders"
 uuid = "6f22d1fd-8eed-4bb7-9776-e7d684900715"
 authors = ["Jin-Guo Liu", "Pan Zhang"]
-version = "0.8.2"
+version = "0.8.3"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 BetterExp = "7cffe744-45fd-4178-b173-cf893948b8b7"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 
+[weakdeps]
+KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
+
+[extensions]
+KaHyParExt = ["KaHyPar"]
+
 [compat]
 AbstractTrees = "0.3, 0.4"
 BetterExp = "0.1"
 JSON = "0.21"
-Requires = "1"
+KaHyPar = "0.3"
 Suppressor = "0.2"
-julia = "1"
+julia = "1.9"
 
 [extras]
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
-KaHyPar = "2a6221f6-aa48-11e9-3542-2d9e0ef01880"
 OMEinsum = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TropicalNumbers = "b3a74e9c-7526-4576-a4eb-79c0d4c32334"
 
 [targets]
-test = ["Test", "Random", "KaHyPar", "Graphs", "TropicalNumbers", "OMEinsum"]
+test = ["Test", "Random", "Graphs", "TropicalNumbers", "OMEinsum", "KaHyPar"]

README.md

@@ -1,26 +1,22 @@
 # OMEinsumContractionOrders
-This package provides `optimize_code` function for finding tensor network contraction orders.
+
+OMEinsumContractionOrders is a Julia package that provides an `optimize_code` function for finding optimal contraction orders for tensor networks. It is designed to work with multiple tensor network packages, such as: [OMEinsum.jl](https://github.com/under-Peter/OMEinsum.jl/issues) package and [ITensorNetworks.jl](https://github.com/mtfishman/ITensorNetworks.jl).
 
 [![Build Status](https://github.com/TensorBFS/OMEinsumContractionOrders.jl/workflows/CI/badge.svg)](https://github.com/TensorBFS/OMEinsumContractionOrders.jl/actions)
 [![codecov](https://codecov.io/gh/TensorBFS/OMEinsumContractionOrders.jl/branch/master/graph/badge.svg?token=BwaF0cV6q1)](https://codecov.io/gh/TensorBFS/OMEinsumContractionOrders.jl)
 
 ## Installation
-<p>
-OMEinsumContractionOrders is a Julia Language package. To install OMEinsumContractionOrders,
-    please <a href="https://docs.julialang.org/en/v1/manual/getting-started/">open
-    Julia's interactive session (known as REPL)</a> and press <kbd>]</kbd> key in the REPL to use the package mode, then type the following command
-</p>
 
-For stable release
+To install OMEinsumContractionOrders, please follow these steps:
 
-```julia
-pkg> add OMEinsumContractionOrders
-pkg> add KaHyPar
-```
-The `KaHyPar` package (used in `KaHyParBipartite` optimizer) is optional because some one may have the binary issue, check [this](https://github.com/kahypar/KaHyPar.jl/issues/12) and [this](https://github.com/kahypar/KaHyPar.jl/issues/19).
+1. Open Julia's interactive session (known as [REPL](https://docs.julialang.org/en/v1/manual/getting-started/)) by typing `julia` in your terminal.
+2. Press the <kbd>]</kbd> key in the REPL to enter the package mode.
+3. Type `add OMEinsumContractionOrders` to install the stable release of the package.
+4. (Optional) If you want to use the `KaHyParBipartite` optimizer, which is based on the KaHyPar library, type `add KaHyPar`. Note that this step is optional because some users may have issues with the binary dependencies of KaHyPar (please check issues: [this](https://github.com/kahypar/KaHyPar.jl/issues/12) and [this](https://github.com/kahypar/KaHyPar.jl/issues/19)).
+
+## Example 1: Use it directly
+The contraction order optimizer is implemented in the `optimize_code` function. It takes three arguments: `code`, `size`, and `optimizer`. The `code` argument is the [einsum notation](https://numpy.org/doc/stable/reference/generated/numpy.einsum.html) to be optimized. The `size` argument is the size of the variables in the einsum notation. The `optimizer` argument is the optimizer to be used. The `optimize_code` function returns the optimized contraction order. One can use `contraction_complexity` function to get the time, space and rewrite complexity of returned contraction order.
 
-## Example
-Optimize a contraction order
 ```julia
 julia> using OMEinsumContractionOrders, Graphs, KaHyPar
 
@@ -44,31 +40,43 @@ julia> optcode_kahypar = optimize_code(code, uniformsize(code, 2),
 julia> optcode_sa = optimize_code(code, uniformsize(code, 2),
 	SABipartite(sc_target=30, max_group_size=50));
 
-julia> timespace_complexity(code, uniformsize(code, 2))
-(200.0, 0.0)
-
-julia> timespace_complexity(optcode_kahypar, uniformsize(code, 2))
-(39.00774639886569, 28.0)
-
-julia> timespace_complexity(optcode_sa, uniformsize(code, 2))
-(39.09524927371961, 28.0)
-
-julia> timespace_complexity(optcode_tree, uniformsize(code, 2))
-(31.13883492534964, 27.0)
-
-julia> timespace_complexity(optcode_tree_with_slice, uniformsize(code, 2))
-(31.292828948918775, 22.0)
+julia> contraction_complexity(code, uniformsize(code, 2))
+Time complexity: 2^200.0
+Space complexity: 2^0.0
+Read-write complexity: 2^10.644757592516257
+
+julia> contraction_complexity(optcode_kahypar, uniformsize(code, 2))
+Time complexity: 2^39.5938886486877
+Space complexity: 2^28.0
+Read-write complexity: 2^30.39890775966298
+
+julia> contraction_complexity(optcode_sa, uniformsize(code, 2))
+Time complexity: 2^41.17129641027078
+Space complexity: 2^29.0
+Read-write complexity: 2^31.493976190321106
+
+julia> contraction_complexity(optcode_tree, uniformsize(code, 2))
+Time complexity: 2^35.06468305863757
+Space complexity: 2^28.0
+Read-write complexity: 2^30.351552349259258
+
+julia> contraction_complexity(optcode_tree_with_slice, uniformsize(code, 2))
+Time complexity: 2^33.70760100663681
+Space complexity: 2^24.0
+Read-write complexity: 2^32.17575935629581
 ```
 
-It is already a part of [`OMEinsum`](https://github.com/under-Peter/OMEinsum.jl):
+## Example 2: Use it in `OMEinsum`
+
+`OMEinsumContractionOrders` is shipped with [`OMEinsum`](https://github.com/under-Peter/OMEinsum.jl) package. You can use it to optimize the contraction order of an `OMEinsum` expression.
 
 ```julia
 julia> using OMEinsum
 
 julia> code = ein"ij, jk, kl, il->"
 ij, jk, kl, il -> 
 
-julia> optimize_code(code, uniformsize(code, 2), TreeSA())
+julia> optimized_code = optimize_code(code, uniformsize(code, 2), TreeSA())
 SlicedEinsum{Char, NestedEinsum{DynamicEinCode{Char}}}(Char[], ki, ki -> 
 ├─ jk, ij -> ki
 │  ├─ jk
@@ -79,86 +87,15 @@ SlicedEinsum{Char, NestedEinsum{DynamicEinCode{Char}}}(Char[], ki, ki ->
 )
 ```
 
-## Acknowledge OMEinsum/OMEinsumContractionOrders
-Since we do not have a paper to cite, we wish anyone who finds this package useful in his research can post a comment under this issue:
-https://github.com/TensorBFS/OMEinsumContractionOrders.jl/issues/21
-
 ## References
 
-If you find this package useful in your research, please cite the following papers
-```
-@misc{Liu2022,
-  doi = {10.48550/ARXIV.2205.03718},
-  url = {https://arxiv.org/abs/2205.03718},
-  author = {Liu, Jin-Guo and Gao, Xun and Cain, Madelyn and Lukin, Mikhail D. and Wang, Sheng-Tao},
-  keywords = {Statistical Mechanics (cond-mat.stat-mech), FOS: Physical sciences, FOS: Physical sciences},
-  title = {Computing solution space properties of combinatorial optimization problems via generic tensor networks},
-  publisher = {arXiv},
-  year = {2022},
-  copyright = {Creative Commons Attribution 4.0 International}
-}
-```
-
-To credit the `KaHyParBipartite` and `SABipartite` method,
-```
-@misc{Pan2021,
-      title={Simulating the Sycamore quantum supremacy circuits}, 
-      author={Feng Pan and Pan Zhang},
-      year={2021},
-      eprint={2103.03074},
-      archivePrefix={arXiv},
-      primaryClass={quant-ph}
-}
-```
-
-To credit the `KaHyParBipartite` method,
-```
-@Article{10.21468/SciPostPhys.7.5.060,
-	title={{Fast counting with tensor networks}},
-	author={Stefanos Kourtis and Claudio Chamon and Eduardo R. Mucciolo and Andrei E. Ruckenstein},
-	journal={SciPost Phys.},
-	volume={7},
-	issue={5},
-	pages={60},
-	year={2019},
-	publisher={SciPost},
-	doi={10.21468/SciPostPhys.7.5.060},
-	url={https://scipost.org/10.21468/SciPostPhys.7.5.060},
-}
-```
-
-```
-@article{Gray2021,
-   title={Hyper-optimized tensor network contraction},
-   volume={5},
-   ISSN={2521-327X},
-   url={http://dx.doi.org/10.22331/q-2021-03-15-410},
-   DOI={10.22331/q-2021-03-15-410},
-   journal={Quantum},
-   publisher={Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften},
-   author={Gray, Johnnie and Kourtis, Stefanos},
-   year={2021},
-   month={Mar},
-   pages={410}
-}
-```
-
-To credit the `TreeSA` method,
-```
-@misc{kalachev2021recursive,
-      title={Recursive Multi-Tensor Contraction for XEB Verification of Quantum Circuits}, 
-      author={Gleb Kalachev and Pavel Panteleev and Man-Hong Yung},
-      year={2021},
-      eprint={2108.05665},
-      archivePrefix={arXiv},
-      primaryClass={quant-ph}
-}
-```
+If you find this package useful in your research, please cite the *relevant* papers in [CITATION.bib](CITATION.bib).
 
 ## Multi-GPU computation
-Check this Gist:
+Please check this Gist:
 
 https://gist.github.com/GiggleLiu/d5b66c9883f0c5df41a440589983ab99
 
 ## Authors
-Jin-Guo Liu and Pan Zhang
+
+OMEinsumContractionOrders was developed by Jin-Guo Liu and Pan Zhang.

ext/KaHyParExt.jl

@@ -0,0 +1,30 @@
+module KaHyParExt
+
+using OMEinsumContractionOrders: KaHyParBipartite, SparseMatrixCSC, group_sc, induced_subhypergraph, convert2int
+import KaHyPar
+import OMEinsumContractionOrders: bipartite_sc
+using Suppressor: @suppress
+
+function bipartite_sc(bipartiter::KaHyParBipartite, adj::SparseMatrixCSC, vertices, log2_sizes)
+    n_v = length(vertices)
+    subgraph, remaining_edges = induced_subhypergraph(adj, vertices)
+    hypergraph = KaHyPar.HyperGraph(subgraph, ones(n_v), convert2int(log2_sizes[remaining_edges]))
+    local parts
+    min_sc = 999999
+    for imbalance in bipartiter.imbalances
+        parts = @suppress KaHyPar.partition(hypergraph, 2; imbalance=imbalance, configuration=:edge_cut)
+        part0 = vertices[parts .== 0]
+        part1 = vertices[parts .== 1]
+        sc0, sc1 = group_sc(adj, part0, log2_sizes), group_sc(adj, part1, log2_sizes)
+        sc = max(sc0, sc1)
+        min_sc = min(sc, min_sc)
+        @debug "imbalance $imbalance: sc = $sc, group = ($(length(part0)), $(length(part1)))"
+        if sc <= bipartiter.sc_target
+            return part0, part1
+        end
+    end
+    error("fail to find a valid partition for `sc_target = $(bipartiter.sc_target)`, got minimum value `$min_sc` (imbalances = $(bipartiter.imbalances))")
+end
+
+@info "`OMEinsumContractionOrders` loads `KaHyParExt` extension successfully."
+end

src/OMEinsumContractionOrders.jl

@@ -5,17 +5,8 @@ using SparseArrays
 using Base: RefValue
 using BetterExp
 using Base.Threads
-using Suppressor: @suppress
 using AbstractTrees
 
-using Requires
-function __init__()
-    @require KaHyPar="2a6221f6-aa48-11e9-3542-2d9e0ef01880" begin
-        using .KaHyPar
-        @info "`OMEinsumContractionOrders` loads `KaHyPar` module successfully."
-    end
-end
-
 export CodeOptimizer, CodeSimplifier,
     KaHyParBipartite, GreedyMethod, TreeSA, SABipartite,
     MinSpaceDiff, MinSpaceOut,

src/complexity.jl

@@ -172,9 +172,9 @@ struct ContractionComplexity
 end
 
 function Base.show(io::IO, cc::ContractionComplexity)
-    print(io, "Time complexity (number of element-wise multiplications) = 2^$(cc.tc)
-Space complexity (number of elements in the largest intermediate tensor) = 2^$(cc.sc)
-Read-write complexity (number of element-wise read and write) = 2^$(cc.rwc)")
+    print(io, "Time complexity: 2^$(cc.tc)
+Space complexity: 2^$(cc.sc)
+Read-write complexity: 2^$(cc.rwc)")
 end
 Base.iterate(cc::ContractionComplexity) = Base.iterate((cc.tc, cc.sc, cc.rwc))
 Base.iterate(cc::ContractionComplexity, state) = Base.iterate((cc.tc, cc.sc, cc.rwc), state)

src/kahypar.jl

@@ -35,28 +35,9 @@ function convert2int(sizes::AbstractVector)
     round.(Int, sizes .* 100)
 end
 
-function bipartite_sc(bipartiter::KaHyParBipartite, adj::SparseMatrixCSC, vertices, log2_sizes)
-    n_v = length(vertices)
-    subgraph, remaining_edges = induced_subhypergraph(adj, vertices)
-    if !isdefined(@__MODULE__, :KaHyPar)
-        error("Module `KaHyPar` not found, please type `using KaHyPar` before using the `KaHyParBipartite` optimizer!")
-    end
-    hypergraph = KaHyPar.HyperGraph(subgraph, ones(n_v), convert2int(log2_sizes[remaining_edges]))
-    local parts
-    min_sc = 999999
-    for imbalance in bipartiter.imbalances
-        parts = @suppress KaHyPar.partition(hypergraph, 2; imbalance=imbalance, configuration=:edge_cut)
-        part0 = vertices[parts .== 0]
-        part1 = vertices[parts .== 1]
-        sc0, sc1 = group_sc(adj, part0, log2_sizes), group_sc(adj, part1, log2_sizes)
-        sc = max(sc0, sc1)
-        min_sc = min(sc, min_sc)
-        @debug "imbalance $imbalance: sc = $sc, group = ($(length(part0)), $(length(part1)))"
-        if sc <= bipartiter.sc_target
-            return part0, part1
-        end
-    end
-    error("fail to find a valid partition for `sc_target = $(bipartiter.sc_target)`, got minimum value `$min_sc` (imbalances = $(bipartiter.imbalances))")
+function bipartite_sc(bipartiter, adj::SparseMatrixCSC, vertices, log2_sizes)
+    error("""Guess you are trying to use the `KaHyParBipartite` optimizer.
+Then you need to add `using KaHyPar` first!""")
 end
 
 # the space complexity (external degree of freedoms) if we contract this group

src/treesa.jl

@@ -75,7 +75,7 @@ function Base.replace!(slicer::Slicer, pair::Pair)
     return slicer
 end
 
-function Base.push!(slicer, best)
+function Base.push!(slicer::Slicer, best)
     @assert length(slicer) < slicer.max_size
     @assert !haskey(slicer.legs, best)
     slicer.legs[best] = slicer.log2_sizes[best]  # add best to legs

test/interfaces.jl

@@ -59,8 +59,8 @@ end
     code2 = optimize_code(code, uniformsize(code, 5), TreeSA(ntrials=1, nslices=5))
     pm2 = peak_memory(code2, uniformsize(code, 5))
     tc2, sc2, rw2 = timespacereadwrite_complexity(code2, uniformsize(code, 5))
-    @test 5 * 2^sc1 > pm1 > 2^sc1
-    @test 5 * 2^sc2 > pm2 > 2^sc2
+    @test 10 * 2^sc1 > pm1 > 2^sc1
+    @test 10 * 2^sc2 > pm2 > 2^sc2
 end