[WIP] Define random_unitary constructor

Project.toml

@@ -14,14 +14,16 @@ TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
+BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
 
 [extensions]
-TensorAlgebraGradedUnitRangesExt = "GradedUnitRanges"
+TensorAlgebraBlockSparseArraysGradedUnitRangesExt = ["BlockSparseArrays", "GradedUnitRanges"]
 
 [compat]
 ArrayLayouts = "1.10.4"
 BlockArrays = "1.2.0"
+BlockSparseArrays = "0.3.6"
 EllipsisNotation = "1.8.0"
 GradedUnitRanges = "0.1.0"
 LinearAlgebra = "1.10.0"

ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl

@@ -0,0 +1,33 @@
+module TensorAlgebraBlockSparseArraysGradedUnitRangesExt
+
+using BlockArrays: Block, blocksize
+using BlockSparseArrays: BlockSparseMatrix, @view!
+using GradedUnitRanges: AbstractGradedUnitRange, dual, tensor_product
+using Random: AbstractRNG
+using TensorAlgebra: TensorAlgebra, random_unitary!
+
+function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
+  return tensor_product(a1, a2)
+end
+
+function TensorAlgebra.square_zero_map(
+  elt::Type, ax::Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}
+)
+  return BlockSparseArray{elt}(undef, (dual.(ax)..., ax...))
+end
+
+function TensorAlgebra.random_unitary!(
+  rng::AbstractRNG,
+  a::BlockSparseMatrix{
+    <:Any,<:Any,<:Any,<:Tuple{AbstractGradedUnitRange,Vararg{AbstractGradedUnitRange}}
+  },
+)
+  # TODO: Define and use `blockdiagindices`
+  # or `blockdiaglength`.
+  for i in 1:blocksize(a, 1)
+    a[Block(i, i)] = random_unitary!(rng, @view!(a[Block(i, i)]))
+  end
+  return a
+end
+
+end

src/fusedims.jl

@@ -21,8 +21,6 @@ function fusedims(::ReshapeFusion, a::AbstractArray, axes::AbstractUnitRange...)
   return reshape(a, axes)
 end
 
-dual(x) = x
-
 ⊗(a::AbstractUnitRange) = a
 function ⊗(a1::AbstractUnitRange, a2::AbstractUnitRange, as::AbstractUnitRange...)
   return ⊗(a1, ⊗(a2, as...))

src/random_unitary.jl

@@ -1,29 +1,32 @@
-using Random: Random, AbstractRNG
+using MatrixAlgebraKit: qr_full!
+using Random: Random, AbstractRNG, randn!
 
-function random_unitary(
-  rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}
-)
-  ax_fused = ⊗(ax...)
-  a_fused = random_unitary(rng, elt, ax_fused)
-  return splitdims(a_fused, dual.(ax), ax)
+function square_zero_map(elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}})
+  return zeros(elt, (ax..., ax...))
 end
 
-# Copy of `Base.to_dim`:
-# https://github.com/JuliaLang/julia/blob/1431bec1bcd205f181ca2a3f1c314247b64076df/base/array.jl#L439-L440
-to_dim(d::Integer) = d
-to_dim(d::Base.OneTo) = last(d)
-
-# Matrix version.
-function random_unitary(rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange})
-  return random_unitary(rng, elt, map(to_dim, ax))
+using EllipsisNotation: : .. function random_unitary!(rng::AbstractRNG, a::AbstractArray)
+  @assert iseven(ndims(a))
+  ndims_codomain = ndims(a) ÷ 2
+  biperm = blockedperm(ntuple(identity, ndims(a)), (ndims_codomain, ndims_codomain))
+  a_mat = fusedims(a, biperm)
+  r_mat = random_unitary!(rng, a_mat)
+  splitdims!(a, r_mat, biperm)
+  return a
 end
 
-using MatrixAlgebraKit: qr_full!
-function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Integer})
-  Q, _ = qr_full!(randn(rng, elt, (dims..., dims...)); positive=true)
+function random_unitary!(rng::AbstractRNG, a::AbstractMatrix)
+  a_r = randn!(rng, a)
+  Q, _ = qr_full!(randn!(rng, a); positive=true)
   return Q
 end
 
+function random_unitary(
+  rng::AbstractRNG, elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}}
+)
+  return random_unitary!(rng, square_zero_map(elt, ax))
+end
+
 # Canonicalizing other kinds of inputs.
 function random_unitary(
   rng::AbstractRNG, elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}}

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+BlockSparseArrays = "2c9a651f-6452-4ace-a6ac-809f4280fbb4"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 GradedUnitRanges = "e2de450a-8a67-46c7-b59c-01d5a3d041c5"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"