Index CartesionProductUnitRange with CartesianPair (#26)

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.19"
+version = "0.1.20"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/cartesianproduct.jl

@@ -94,6 +94,12 @@ function Base.checkindex(::Type{Bool}, inds::CartesianProductUnitRange, i::Carte
   return checkindex(Bool, arg1(inds), arg1(i)) && checkindex(Bool, arg2(inds), arg2(i))
 end
 
+# Reverse map from CartesianPair to linear index in the range.
+function Base.getindex(inds::CartesianProductUnitRange, i::CartesianPair)
+  i′ = (findfirst(==(arg1(i)), arg1(inds)), findfirst(==(arg2(i)), arg2(inds)))
+  return inds[LinearIndices((length(arg1(inds)), length(arg2(inds))))[i′...]]
+end
+
 using Base.Broadcast: DefaultArrayStyle
 for f in (:+, :-)
   @eval begin

test/test_basics.jl

@@ -36,6 +36,12 @@ elts = (Float32, Float64, ComplexF32, ComplexF64)
   @test length(r) == 6
   @test first(r) == 1
   @test last(r) == 6
+  @test r[1 × 1] == 1
+  @test r[2 × 1] == 2
+  @test r[1 × 2] == 3
+  @test r[2 × 2] == 4
+  @test r[1 × 3] == 5
+  @test r[2 × 3] == 6
 
   r = @constinferred(cartesianrange(2 × 3, 2:7))
   @test r === cartesianrange(Base.OneTo(2) × Base.OneTo(3), 2:7)
@@ -44,6 +50,12 @@ elts = (Float32, Float64, ComplexF32, ComplexF64)
   @test length(r) == 6
   @test first(r) == 2
   @test last(r) == 7
+  @test r[1 × 1] == 2
+  @test r[2 × 1] == 3
+  @test r[1 × 2] == 4
+  @test r[2 × 2] == 5
+  @test r[1 × 3] == 6
+  @test r[2 × 3] == 7
 
   # Test high-dimensional materialization.
   a = randn(elt, 2, 2, 2) ⊗ randn(elt, 2, 2, 2)