Introduce CartesianPair

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "KroneckerArrays"
 uuid = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.18"
+version = "0.1.19"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -23,7 +23,7 @@ KroneckerArraysBlockSparseArraysExt = ["BlockArrays", "BlockSparseArrays"]
 [compat]
 Adapt = "4.3.0"
 BlockArrays = "1.6"
-BlockSparseArrays = "0.7.19"
+BlockSparseArrays = "0.7.20"
 DerivableInterfaces = "0.5.0"
 DiagonalArrays = "0.3.5"
 FillArrays = "1.13.0"

ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl

@@ -1,7 +1,20 @@
 module KroneckerArraysBlockSparseArraysExt
 
+using BlockArrays: Block
+using BlockSparseArrays: BlockIndexVector, GenericBlockIndex
+using KroneckerArrays: CartesianPair, CartesianProduct
+function Base.getindex(b::Block, I1::CartesianPair, Irest::CartesianPair...)
+  return GenericBlockIndex(b, (I1, Irest...))
+end
+function Base.getindex(b::Block, I1::CartesianProduct, Irest::CartesianProduct...)
+  return BlockIndexVector(b, (I1, Irest...))
+end
+
 using BlockSparseArrays: BlockSparseArrays, blockrange
-using KroneckerArrays: CartesianProduct, cartesianrange
+using KroneckerArrays: CartesianPair, CartesianProduct, cartesianrange
+function BlockSparseArrays.blockrange(bs::Vector{<:CartesianPair})
+  return blockrange(map(cartesianrange, bs))
+end
 function BlockSparseArrays.blockrange(bs::Vector{<:CartesianProduct})
   return blockrange(map(cartesianrange, bs))
 end

src/cartesianproduct.jl

@@ -1,4 +1,22 @@
-struct CartesianProduct{A,B}
+struct CartesianPair{A,B}
+  a::A
+  b::B
+end
+arguments(a::CartesianPair) = (a.a, a.b)
+arguments(a::CartesianPair, n::Int) = arguments(a)[n]
+
+arg1(a::CartesianPair) = a.a
+arg2(a::CartesianPair) = a.b
+
+×(a, b) = CartesianPair(a, b)
+
+function Base.show(io::IO, a::CartesianPair)
+  print(io, a.a, " × ", a.b)
+  return nothing
+end
+
+struct CartesianProduct{TA,TB,A<:AbstractVector{TA},B<:AbstractVector{TB}} <:
+       AbstractVector{CartesianPair{TA,TB}}
   a::A
   b::B
 end
@@ -13,15 +31,19 @@ function Base.show(io::IO, a::CartesianProduct)
   return nothing
 end
 
-×(a, b) = CartesianProduct(a, b)
+×(a::AbstractVector, b::AbstractVector) = CartesianProduct(a, b)
 Base.length(a::CartesianProduct) = length(a.a) * length(a.b)
-Base.getindex(a::CartesianProduct, i::CartesianProduct) = a.a[i.a] × a.b[i.b]
+Base.size(a::CartesianProduct) = (length(a),)
 
-function Base.iterate(a::CartesianProduct, state...)
-  x = iterate(Iterators.product(a.a, a.b), state...)
-  isnothing(x) && return x
-  next, new_state = x
-  return ×(next...), new_state
+function Base.getindex(a::CartesianProduct, i::CartesianProduct)
+  return arg1(a)[arg1(i)] × arg2(a)[arg2(i)]
+end
+function Base.getindex(a::CartesianProduct, i::CartesianPair)
+  return arg1(a)[arg1(i)] × arg2(a)[arg2(i)]
+end
+function Base.getindex(a::CartesianProduct, i::Int)
+  I = Tuple(CartesianIndices((length(arg1(a)), length(arg2(a))))[i])
+  return a[I[1] × I[2]]
 end
 
 struct CartesianProductUnitRange{T,P<:CartesianProduct,R<:AbstractUnitRange{T}} <:
@@ -47,10 +69,18 @@ end
 to_range(a::AbstractUnitRange) = a
 to_range(i::Integer) = Base.OneTo(i)
 cartesianrange(a, b) = cartesianrange(to_range(a) × to_range(b))
+function cartesianrange(p::CartesianPair)
+  p′ = to_range(p.a) × to_range(p.b)
+  return cartesianrange(p′)
+end
 function cartesianrange(p::CartesianProduct)
   p′ = to_range(p.a) × to_range(p.b)
   return cartesianrange(p′, Base.OneTo(length(p′)))
 end
+function cartesianrange(p::CartesianPair, range::AbstractUnitRange)
+  p′ = to_range(p.a) × to_range(p.b)
+  return cartesianrange(p′, range)
+end
 function cartesianrange(p::CartesianProduct, range::AbstractUnitRange)
   p′ = to_range(p.a) × to_range(p.b)
   return CartesianProductUnitRange(p′, range)

src/kroneckerarray.jl

@@ -188,6 +188,9 @@ end
 function Base.getindex(a::KroneckerArray{<:Any,N}, I::Vararg{CartesianProduct,N}) where {N}
   return a.a[map(Base.Fix2(getfield, :a), I)...] ⊗ a.b[map(Base.Fix2(getfield, :b), I)...]
 end
+function Base.getindex(a::KroneckerArray{<:Any,N}, I::Vararg{CartesianPair,N}) where {N}
+  return a.a[map(Base.Fix2(getfield, :a), I)...] ⊗ a.b[map(Base.Fix2(getfield, :b), I)...]
+end
 # Fix ambigiuity error.
 Base.getindex(a::KroneckerArray{<:Any,0}) = a.a[] * a.b[]