Implement general tensor_pow via Base.power_by_squaring (#54)

Author: akirakyle

src/QuantumInterface.jl

@@ -45,7 +45,7 @@ function tensor end
 const ⊗ = tensor
 tensor() = throw(ArgumentError("Tensor function needs at least one argument."))
 
-function tensor_pow end # TODO should Base.^ be the same as tensor_pow?
+function tensor_pow end
 
 function traceout! end
 

src/bases.jl

@@ -93,12 +93,7 @@ function tensor(b1::Basis, b2::CompositeBasis)
 end
 tensor(bases::Basis...) = reduce(tensor, bases)
 
-function Base.:^(b::Basis, N::Integer)
-    if N < 1
-        throw(ArgumentError("Power of a basis is only defined for positive integers."))
-    end
-    tensor([b for i=1:N]...)
-end
+Base.:^(b::Basis, N::Integer) = tensor_pow(b, N)
 
 """
     equal_shape(a, b)

src/tensor.jl

@@ -8,3 +8,36 @@ tensor(op::AbstractOperator) = op
 tensor(operators::AbstractOperator...) = reduce(tensor, operators)
 tensor(state::StateVector) = state
 tensor(states::Vector{T}) where T<:StateVector = reduce(tensor, states)
+
+"""
+    tensor_pow(a, N)
+
+Gives the tensor product of `a` `N` times.
+"""
+tensor_pow(a, N) = tensor_pow_by_squaring(a, N)
+
+# Copied from Base.power_by_squaring as `mul` keyword dosn't work without implementing `*`
+function tensor_pow_by_squaring(x, p::Integer)
+    if p == 1
+        return x
+    elseif p == 2
+        return tensor(x, x)
+    elseif p < 1
+        throw(DomainError("Cannot take tensor_pow to power less than one"))
+    end
+    t = trailing_zeros(p) + 1
+    p >>= t
+    while (t -= 1) > 0
+        x = tensor(x, x)
+    end
+    y = x
+    while p > 0
+        t = trailing_zeros(p) + 1
+        p >>= t
+        while (t -= 1) >= 0
+            x = tensor(x, x)
+        end
+        y = tensor(y, x)
+    end
+    return y
+end

test/test_bases.jl

@@ -28,7 +28,7 @@ comp_b2 = tensor(b1, b1, b2)
 
 @test b1^3 == CompositeBasis(b1, b1, b1)
 @test (b1⊗b2)^2 == CompositeBasis(b1, b2, b1, b2)
-@test_throws ArgumentError b1^(0)
+@test_throws DomainError b1^(0)
 
 comp_b1_b2 = tensor(comp_b1, comp_b2)
 @test comp_b1_b2.shape == [prod(shape1), prod(shape2), prod(shape1), prod(shape1), prod(shape2)]