use onehot batch (#16)

* use onehot batch
QuantumBFS · May 17, 2019 · 49df431 · 49df431
1 parent a2462d1
commit 49df431
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "YaoArrayRegister"
 uuid = "e600142f-9330-5003-8abb-0ebd767abc51"
-version = "0.3.7"
+version = "0.3.8"
 
 [deps]
 BitBasis = "50ba71b6-fa0f-514d-ae9a-0916efc90dcf"
@@ -15,7 +15,7 @@ YaoBase = "a8f54c17-34bc-5a9d-b050-f522fe3f755f"
 julia = "^1.0.0"
 LuxurySparse = "~0.4.0"
 YaoBase = "~0.9.0"
-BitBasis = "^0.4.0"
+BitBasis = "^0.5.0"
 
 [extras]
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

diff --git a/src/register.jl b/src/register.jl
@@ -105,7 +105,7 @@ ArrayReg{1,Complex{Float32},Array...}
 ArrayReg(bitstr::BitStr) = ArrayReg(ComplexF64, bitstr)
 ArrayReg(::Type{T}, bitstr::BitStr) where T = ArrayReg{1}(T, bitstr)
 ArrayReg{B}(bitstr::BitStr) where B = ArrayReg{B}(ComplexF64, bitstr)
-ArrayReg{B}(::Type{T}, bitstr::BitStr) where {B, T} = ArrayReg{B}(hcat((onehot(T, bitstr) for k in 1:B)...))
+ArrayReg{B}(::Type{T}, bitstr::BitStr) where {B, T} = ArrayReg{B}(onehot(T, bitstr, B))
 
 
 """
@@ -320,8 +320,7 @@ product_state(total::Int, bit_config::Integer; nbatch::Int=1) = product_state(Co
 product_state(::Type{T}, bit_str::BitStr; nbatch::Int=1) where T = ArrayReg{nbatch}(T, bit_str)
 
 function product_state(::Type{T}, total::Int, bit_config::Integer; nbatch::Int=1) where T
-    raw_state = hcat((onehot(T, total, bit_config) for k in 1:nbatch)...)
-    return ArrayReg{nbatch}(raw_state)
+    return ArrayReg{nbatch}(onehot(T, total, bit_config, nbatch))
 end
 
 """