[WIP] Define random_unitary constructor

Project.toml

@@ -1,28 +1,35 @@
 name = "TensorAlgebra"
 uuid = "68bd88dc-f39d-4e12-b2ca-f046b68fcc6a"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.2.1"
+version = "0.2.2"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 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]
+TensorAlgebraBlockSparseArraysGradedUnitRangesExt = ["BlockSparseArrays", "GradedUnitRanges"]
 TensorAlgebraGradedUnitRangesExt = "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"
+LinearAlgebra = "1.10.0"
+MatrixAlgebraKit = "0.1.1"
+Random = "1.10.0"
 TupleTools = "1.6.0"
-TypeParameterAccessors = "0.2.1, 0.3"
+TypeParameterAccessors = "0.2.1, 0.3.0"
 julia = "1.10"

ext/TensorAlgebraBlockSparseArraysGradedUnitRangesExt/TensorAlgebraBlockSparseArraysGradedUnitRangesExt.jl

@@ -0,0 +1,27 @@
+module TensorAlgebraBlockSparseArraysGradedUnitRangesExt
+
+using BlockArrays: Block, blocksize
+using BlockSparseArrays: BlockSparseMatrix, @view!
+using GradedUnitRanges: AbstractGradedUnitRange, dual
+using Random: AbstractRNG
+using TensorAlgebra: TensorAlgebra, random_unitary!
+
+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,<:NTuple{2,AbstractGradedUnitRange}},
+)
+  # TODO: Define and use `blockdiagindices`
+  # or `blockdiaglength`.
+  for i in 1:blocksize(a, 1)
+    random_unitary!(rng, @view!(a[Block(i, i)]))
+  end
+  return a
+end
+
+end

ext/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl

@@ -1,8 +1,10 @@
 module TensorAlgebraGradedUnitRangesExt
+
 using GradedUnitRanges: AbstractGradedUnitRange, tensor_product
 using TensorAlgebra: TensorAlgebra
 
 function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
   return tensor_product(a1, a2)
 end
+
 end

src/TensorAlgebra.jl

@@ -13,5 +13,6 @@ include("contract/blockedperms.jl")
 include("contract/allocate_output.jl")
 include("contract/contract_matricize/contract.jl")
 include("factorizations.jl")
+include("random_unitary.jl")
 
 end

src/random_unitary.jl

@@ -0,0 +1,55 @@
+using MatrixAlgebraKit: qr_full!
+using Random: Random, AbstractRNG, randn!
+
+function square_zero_map(elt::Type, ax::Tuple{AbstractUnitRange,Vararg{AbstractUnitRange}})
+  return zeros(elt, (ax..., ax...))
+end
+
+# Imports `..` into the namespace.
+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)
+  random_unitary!(rng, a_mat)
+  splitdims!(a, a_mat, biperm)
+  return a
+end
+
+function random_unitary!(rng::AbstractRNG, a::AbstractMatrix)
+  Q, _ = qr_full!(randn!(rng, a); positive=true)
+  a .= Q
+  return a
+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}}}
+)
+  return random_unitary(Random.default_rng(), elt, map(to_axis, dims))
+end
+function random_unitary(elt::Type, dims::Tuple{Vararg{Union{AbstractUnitRange,Integer}}})
+  return random_unitary(Random.default_rng(), elt, dims)
+end
+function random_unitary(
+  rng::AbstractRNG, elt::Type, dims::Union{AbstractUnitRange,Integer}...
+)
+  return random_unitary(rng, elt, dims)
+end
+function random_unitary(elt::Type, dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(Random.default_rng(), elt, dims)
+end
+function random_unitary(rng::AbstractRNG, dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(rng, Float64, dims)
+end
+function random_unitary(dims::Union{AbstractUnitRange,Integer}...)
+  return random_unitary(Random.default_rng(), Float64, dims)
+end

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"