Define tensor_product (#2)

Author: ogauthe

Project.toml

@@ -3,5 +3,12 @@ uuid = "decf83d6-1968-43f4-96dc-fdb3fe15fc6d"
 authors = ["ITensor developers <[email protected]> and contributors"]
 version = "0.1.0"
 
+[weakdeps]
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+
+[extensions]
+TensorProductsBlockArraysExt = "BlockArrays"
+
 [compat]
+BlockArrays = "1.2.0"
 julia = "1.10"

docs/make.jl

@@ -1,9 +1,7 @@
 using TensorProducts: TensorProducts
 using Documenter: Documenter, DocMeta, deploydocs, makedocs
 
-DocMeta.setdocmeta!(
-  TensorProducts, :DocTestSetup, :(using TensorProducts); recursive=true
-)
+DocMeta.setdocmeta!(TensorProducts, :DocTestSetup, :(using TensorProducts); recursive=true)
 
 include("make_index.jl")
 

ext/TensorProductsBlockArraysExt/TensorProductsBlockArraysExt.jl

@@ -0,0 +1,23 @@
+module TensorProductsBlockArraysExt
+
+using BlockArrays:
+  AbstractBlockedUnitRange,
+  Block,
+  BlockArrays,
+  blockaxes,
+  blockedrange,
+  blocklengths,
+  blocks
+
+using TensorProducts: OneToOne, TensorProducts
+
+function TensorProducts.tensor_product(
+  a1::AbstractBlockedUnitRange, a2::AbstractBlockedUnitRange
+)
+  new_blocklengths = mapreduce(vcat, Iterators.product(blocks(a1), blocks(a2))) do (x, y)
+    return length(x) * length(y)
+  end
+  return blockedrange(new_blocklengths)
+end
+
+end

src/TensorProducts.jl

@@ -1,5 +1,8 @@
 module TensorProducts
 
-# Write your package code here.
+export ⊗, OneToOne, tensor_product
+
+include("onetoone.jl")
+include("tensor_product.jl")
 
 end

src/onetoone.jl

@@ -0,0 +1,7 @@
+# This files defines the struct OneToOne
+# OneToOne represents the range `1:1` or `Base.OneTo(1)`.
+
+struct OneToOne{T} <: AbstractUnitRange{T} end
+OneToOne() = OneToOne{Int}()
+Base.first(a::OneToOne) = one(eltype(a))
+Base.last(a::OneToOne) = one(eltype(a))

src/tensor_product.jl

@@ -0,0 +1,31 @@
+# This files defines an interface for the tensor product of two axes
+# https://en.wikipedia.org/wiki/Tensor_product
+
+# ==================================  misc  ================================================
+is_offset_axis(a::AbstractUnitRange) = !isone(first(a))
+
+function require_one_based_axis(a::AbstractUnitRange)
+  return is_offset_axis(a) && throw(ArgumentError("Range must be one-based"))
+end
+
+# ==============================  tensor product  ==========================================
+⊗() = tensor_product()
+⊗(a) = tensor_product(a)
+
+# default. No type restriction to allow sectors as input
+⊗(a1, a2) = tensor_product(a1, a2)
+
+# allow to specialize ⊗(a1, a2) to fusion_product
+⊗(a1, a2, as...) = ⊗(⊗(a1, a2), as...)
+
+tensor_product() = OneToOne()
+tensor_product(a) = a
+tensor_product(a1, a2, as...) = tensor_product(tensor_product(a1, a2), as...)
+
+# default
+function tensor_product(a1::AbstractUnitRange, a2::AbstractUnitRange)
+  require_one_based_axis(a1) || require_one_based_axis(a2)
+  return Base.OneTo(length(a1) * length(a2))
+end
+
+tensor_product(::OneToOne, ::OneToOne) = OneToOne()

test/Project.toml

@@ -1,11 +1,13 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
 Aqua = "0.8.9"
+BlockArrays = "1.2.0"
 SafeTestsets = "0.1"
 Suppressor = "0.2"
 Test = "1.10"

test/basics/test_basics.jl

@@ -0,0 +1,15 @@
+using Test: @test, @testset
+
+using BlockArrays: BlockRange, blockaxes
+
+using TensorProducts: OneToOne
+
+@testset "OneToOne" begin
+  a0 = OneToOne()
+  @test a0 isa OneToOne{Int}
+  @test a0 isa AbstractUnitRange{Int}
+  @test eltype(a0) == Int
+  @test length(a0) == 1
+
+  @test blockaxes(OneToOne()) == (BlockRange(OneToOne()),)
+end

test/test_basics.jl

@@ -0,0 +1,8 @@
+using Test: @test, @testset
+
+using TensorProducts: TensorProducts
+
+@testset "Test exports" begin
+  exports = [:⊗, :TensorProducts, :OneToOne, :tensor_product]
+  @test issetequal(names(TensorProducts), exports)
+end

test/test_exports.jl

@@ -0,0 +1,32 @@
+using Test: @test, @test_throws, @testset
+
+using TensorProducts: ⊗, OneToOne, tensor_product
+
+using BlockArrays: blockedrange, blockisequal
+
+r0 = OneToOne()
+b1 = blockedrange([1, 2])
+
+@testset "⊗" begin
+  @test ⊗() isa OneToOne
+  @test ⊗(1:2) == 1:2
+  @test ⊗(1:2, 1:3) == 1:6
+  @test ⊗(1:2, 1:3, 1:4) == 1:24
+
+  @test ⊗(r0, r0) isa OneToOne
+  @test blockisequal(⊗(b1, b1), blockedrange([1, 2, 2, 4]))
+end
+
+@testset "tensor_product" begin
+  @test tensor_product() isa OneToOne
+  @test tensor_product(1:2) == 1:2
+  @test tensor_product(1:2, 1:3) == 1:6
+  @test tensor_product(1:2, 1:3, 1:4) == 1:24
+
+  @test_throws ArgumentError tensor_product(2:3, 1:2)
+  @test_throws ArgumentError tensor_product(1:3, 2:2)
+  @test_throws ArgumentError tensor_product(2:3, 2:2)
+
+  @test tensor_product(r0, r0) isa OneToOne
+  @test blockisequal(tensor_product(b1, b1), blockedrange([1, 2, 2, 4]))
+end