[WIP] Define random_unitary constructor

Project.toml

@@ -1,13 +1,14 @@
 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"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 TupleTools = "9d95972d-f1c8-5527-a6e0-b4b365fa01f6"
 TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
@@ -22,7 +23,8 @@ ArrayLayouts = "1.10.4"
 BlockArrays = "1.2.0"
 EllipsisNotation = "1.8.0"
 GradedUnitRanges = "0.1.0"
-LinearAlgebra = "1.10"
+LinearAlgebra = "1.10.0"
+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/TensorAlgebraGradedUnitRangesExt/TensorAlgebraGradedUnitRangesExt.jl

@@ -1,8 +1,29 @@
 module TensorAlgebraGradedUnitRangesExt
-using GradedUnitRanges: AbstractGradedUnitRange, tensor_product
-using TensorAlgebra: TensorAlgebra
+
+using GradedUnitRanges: AbstractGradedUnitRange, dual, tensor_product
+using GradedUnitRanges.BlockArrays: Block, blocklengths, blocksize
+using Random: AbstractRNG
+using TensorAlgebra: TensorAlgebra, random_unitary
 
 function TensorAlgebra.:⊗(a1::AbstractGradedUnitRange, a2::AbstractGradedUnitRange)
   return tensor_product(a1, a2)
 end
+
+function TensorAlgebra.dual(a::AbstractGradedUnitRange)
+  return dual(a)
+end
+
+function TensorAlgebra.random_unitary(
+  rng::AbstractRNG,
+  elt::Type,
+  ax::Tuple{AbstractGradedUnitRange},
+)
+  a = zeros(elt, dual.(ax)..., ax...)
+  # TODO: Define `blockdiagindices`.
+  for i in 1:minimum(blocksize(a))
+    a[Block(i, i)] = random_unitary(rng, elt, Int(blocklengths(only(ax))[i]))
+  end
+  return a
+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,69 @@
+# Version of `sign` that returns one
+# if `x == 0`.
+function nonzero_sign(x)
+  iszero(x) && return one(x)
+  return sign(x)
+end
+
+using LinearAlgebra: LinearAlgebra, Diagonal, diag
+function qr_positive(M::AbstractMatrix)
+  Q, R = LinearAlgebra.qr(M)
+  Q′ = typeof(R)(Q)
+  signs = nonzero_sign.(diag(R))
+  Q′ = Q′ * Diagonal(signs)
+  R = Diagonal(conj.(signs)) * R
+  return Q′, R
+end
+
+using Random: Random, AbstractRNG
+
+dual(x) = x
+
+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)
+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))
+end
+
+function random_unitary(rng::AbstractRNG, elt::Type, dims::Tuple{Integer})
+  Q, _ = qr_positive(randn(rng, elt, (dims..., dims...)))
+  return Q
+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