[NDTensors] BlockSparseArray contract, QR, and Hermitian eigendecomposition

NDTensors/Project.toml

@@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SplitApplyCombine = "03a91e81-4c3e-53e1-a0a4-9c0c8f19dd66"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Strided = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
@@ -52,6 +53,7 @@ LinearAlgebra = "1.6"
 PackageExtensionCompat = "1"
 Random = "1.6"
 SimpleTraits = "0.9.4"
+SparseArrays = "1.6"
 SplitApplyCombine = "1.2.2"
 StaticArrays = "0.12, 1.0"
 Strided = "2"

NDTensors/src/BlockSparseArrays/src/BlockSparseArrays.jl

@@ -1,10 +1,26 @@
 module BlockSparseArrays
-using BlockArrays
-using Compat
-using Dictionaries
-using SplitApplyCombine
-
-using BlockArrays: block
+using BlockArrays:
+  AbstractBlockArray,
+  BlockArrays,
+  BlockVector,
+  Block,
+  BlockIndex,
+  BlockRange,
+  BlockedUnitRange,
+  findblockindex,
+  block,
+  blockaxes,
+  blockcheckbounds,
+  blockfirsts,
+  blocklasts,
+  blocklength,
+  blocklengths,
+  blockedrange,
+  blocks
+using Compat: Returns, allequal
+using Dictionaries: Dictionary, Indices, getindices, set! # TODO: Move to `SparseArraysExtensions`.
+using LinearAlgebra: Hermitian
+using SplitApplyCombine: groupcount
 
 export BlockSparseArray, SparseArray
 
@@ -14,11 +30,14 @@ include("axes.jl")
 include("abstractarray.jl")
 include("permuteddimsarray.jl")
 include("blockarrays.jl")
+# TODO: Split off into `SparseArraysExtensions` module, rename to `SparseArrayDOK`.
 include("sparsearray.jl")
 include("blocksparsearray.jl")
+include("allocate_output.jl")
 include("subarray.jl")
 include("broadcast.jl")
 include("fusedims.jl")
 include("gradedrange.jl")
+include("LinearAlgebraExt/LinearAlgebraExt.jl")
 
 end

NDTensors/src/BlockSparseArrays/src/LinearAlgebraExt/LinearAlgebraExt.jl

@@ -0,0 +1,16 @@
+module LinearAlgebraExt
+using BlockArrays: BlockArrays, blockedrange, blocks
+using ..BlockSparseArrays: SparseArray, nonzero_keys # TODO: Move to `SparseArraysExtensions` module, rename `SparseArrayDOK`.
+using ..BlockSparseArrays: BlockSparseArrays, BlockSparseArray, nonzero_blockkeys
+using LinearAlgebra: LinearAlgebra, Hermitian, Transpose, I, eigen, qr
+using ...NDTensors: Algorithm, @Algorithm_str # TODO: Move to `AlgorithmSelector` module.
+using SparseArrays: SparseArrays, SparseMatrixCSC, spzeros, sparse
+
+# TODO: Move to `SparseArraysExtensions`.
+include("hermitian.jl")
+# TODO: Move to `SparseArraysExtensions`.
+include("transpose.jl")
+include("qr.jl")
+include("eigen.jl")
+include("svd.jl")
+end

NDTensors/src/BlockSparseArrays/src/LinearAlgebraExt/eigen.jl

@@ -0,0 +1,19 @@
+function LinearAlgebra.eigen(a::BlockSparseArray)
+  return error("Not implemented")
+end
+
+# TODO: Maybe make `Hermitian` partially eager for `BlockSparseArray`?
+function LinearAlgebra.eigen(
+  a::Union{Hermitian{<:Real,<:BlockSparseArray},Hermitian{<:Complex,<:BlockSparseArray}}
+)
+  # TODO: Test `a` is block diagonal.
+  # @assert is_block_diagonal(a)
+  d = BlockSparseArray{real(eltype(a))}(axes(a, 1))
+  u = BlockSparseArray{eltype(a)}(axes(a))
+  for b in nonzero_blockkeys(a)
+    d_b, u_b = eigen(@view a[b])
+    d[BlockArrays.Block(b.n[1])] = d_b
+    u[b] = u_b
+  end
+  return d, u
+end

NDTensors/src/BlockSparseArrays/src/LinearAlgebraExt/hermitian.jl

@@ -0,0 +1,4 @@
+# TODO: This needs to be done more carefully by
+# grabbing upper or lower triangles of the parent matrix.
+# TODO: Move to `SparseArraysExtensions`.
+BlockSparseArrays.nonzero_keys(a::Hermitian) = nonzero_keys(parent(a))

NDTensors/src/BlockSparseArrays/src/LinearAlgebraExt/qr.jl

@@ -0,0 +1,137 @@
+# Check if the matrix has 1 or fewer entries
+# per row/column.
+function is_permutation_matrix(a::SparseMatrixCSC)
+  return all(col -> length(nzrange(a, col)) ≤ 1, axes(a, 2))
+end
+
+# Check if the matrix has 1 or fewer entries
+# per row/column.
+function is_permutation_matrix(a::SparseArray{<:Any,2})
+  keys = collect(Iterators.map(Tuple, nonzero_keys(a)))
+  I = first.(keys)
+  J = last.(keys)
+  return allunique(I) && allunique(J)
+end
+
+function findnonzerorows(a::SparseMatrixCSC, col)
+  return view(a.rowval, a.colptr[col]:(a.colptr[col + 1] - 1))
+end
+
+function SparseArrays.SparseMatrixCSC(a::SparseArray{<:Any,2})
+  # Not defined:
+  # a_csc = SparseMatrixCSC{eltype(a)}(size(a))
+  a_csc = spzeros(eltype(a), size(a))
+  for I in nonzero_keys(a)
+    a_csc[I] = a[I]
+  end
+  return a_csc
+end
+
+# Get the sparse structure of a SparseArray as a SparseMatrixCSC.
+function sparse_structure(structure_type::Type{<:SparseMatrixCSC}, a::SparseArray{<:Any,2})
+  # Idealy would work but a bit too complicated for `map` right now:
+  # return SparseMatrixCSC(map(x -> iszero(x) ? false : true, a))
+  # TODO: Change to `spzeros(Bool, size(a))`.
+  a_structure = structure_type(spzeros(Bool, size(a)...))
+  for I in nonzero_keys(a)
+    i, j = Tuple(I)
+    a_structure[i, j] = true
+  end
+  return a_structure
+end
+
+# Get the sparsity structure as a `SparseMatrixCSC` with values
+# of `true` where there are structural nonzero blocks and `false`
+# otherwise.
+function block_sparse_structure(structure_type::Type, a::BlockSparseArray{<:Any,2})
+  return sparse_structure(structure_type, blocks(a))
+end
+
+function is_block_permutation_matrix(a::BlockSparseArray{<:Any,2})
+  return is_permutation_matrix(blocks(a))
+end
+
+qr_rank(alg::Algorithm"thin", a::AbstractArray{<:Any,2}) = minimum(size(a))
+
+# m × n → (m × min(m, n)) ⋅ (min(m, n) × n)
+function qr_block_sparse_structure(alg::Algorithm"thin", a::BlockSparseArray{<:Any,2})
+  axes_row, axes_col = axes(a)
+  a_csc = block_sparse_structure(SparseMatrixCSC, a)
+  F = qr(float(a_csc))
+  # Outputs full Q
+  # q_csc = sparse(F.Q[invperm(F.prow), :])
+  q_csc = (F.Q * sparse(I, size(a_csc, 1), minimum(size(a_csc))))[invperm(F.prow), :]
+  r_csc = F.R[:, invperm(F.pcol)]
+  nblocks = size(q_csc, 2)
+  @assert nblocks == size(r_csc, 1)
+  a_sparse = blocks(a)
+  blocklengths_qr = Vector{Int}(undef, nblocks)
+  for I in nonzero_keys(a_sparse)
+    i, k = Tuple(I)
+    # Get the nonzero columns associated
+    # with the given row.
+    j = only(findnonzerorows(r_csc, k))
+    # @assert is_structural_nonzero(r, j, k)
+    # @assert is_structural_nonzero(q, i, j)
+    blocklengths_qr[j] = qr_rank(alg, @view(a[BlockArrays.Block(i, k)]))
+  end
+  axes_qr = blockedrange(blocklengths_qr)
+  axes_q = (axes(a, 1), axes_qr)
+  axes_r = (axes_qr, axes(a, 2))
+  # TODO: Come up with a better format to ouput.
+  # TODO: Get `axes_qr` as a permutation of the
+  # axes of `axes(a, 2)` to preserve sectors
+  # when using symmetric tensors.
+  return q_csc, axes_q, r_csc, axes_r
+end
+
+# m × n → (m × m) ⋅ (m × n)
+function qr_block_sparse_structure(alg::Algorithm"full", a::BlockSparseArray{<:Any,2})
+  return error("Not implemented")
+end
+
+function qr_blocks(a, structure_r, block_a)
+  i, k = block_a.n
+  j = only(findnonzerorows(structure_r, k))
+  return BlockArrays.Block(i, j), BlockArrays.Block(j, k)
+end
+
+# Block-preserving QR.
+function LinearAlgebra.qr(a::BlockSparseArray{<:Any,2}; alg="thin")
+  return qr(Algorithm(alg), a)
+end
+
+# Block-preserving QR.
+function LinearAlgebra.qr(alg::Algorithm, a::BlockSparseArray{<:Any,2})
+  if !is_block_permutation_matrix(a)
+    # Must have 1 or fewer blocks per row/column.
+    println("Block sparsity structure is:")
+    display(nonzero_blockkeys(a))
+    error("Not a block permutation matrix")
+  end
+  eltype_a = eltype(a)
+  # TODO: `structure_q` isn't needed.
+  structure_q, axes_q, structure_r, axes_r = qr_block_sparse_structure(alg, a)
+  # TODO: Make this generic to GPU, use `similar`.
+  q = BlockSparseArray{eltype_a}(axes_q)
+  r = BlockSparseArray{eltype_a}(axes_r)
+  for block_a in nonzero_blockkeys(a)
+    # TODO: Make thin or full depending on `alg`.
+    q_b, r_b = qr(a[block_a])
+    # Determine the block of Q and R
+    # TODO: Do the block locations change for `alg="full"`?
+    block_q, block_r = qr_blocks(a, structure_r, block_a)
+
+    # TODO Make this generic to GPU.
+    q[block_q] = Matrix(q_b)
+    r[block_r] = r_b
+  end
+  # TODO: If `alg="full"`, fill in blocks of `q`
+  # with random unitaries.
+  # Which blocks should be filled? Seems to be based
+  # on the QNs...
+  # Maybe fill diagonal blocks.
+  # TODO: Also store `structure_r` in some way
+  # since that is needed for permuting the QNs.
+  return q, r
+end

NDTensors/src/BlockSparseArrays/src/LinearAlgebraExt/svd.jl

@@ -0,0 +1,3 @@
+function LinearAlgebra.svd(a::BlockSparseArray; kwargs...)
+  return error("Not implemented")
+end

NDTensors/src/BlockSparseArrays/src/LinearAlgebraExt/transpose.jl

@@ -0,0 +1,6 @@
+# TODO: Move to `SparseArraysExtensions`.
+function BlockSparseArrays.nonzero_keys(a::Transpose)
+  return (
+    CartesianIndex(reverse(Tuple(parent_index))) for parent_index in nonzero_keys(parent(a))
+  )
+end

NDTensors/src/BlockSparseArrays/src/allocate_output.jl

@@ -0,0 +1,123 @@
+#############################################################################
+# Generic
+#############################################################################
+
+function output_type(f, args::Type...)
+  # TODO: Is this good to use here?
+  # Seems best for `Number` subtypes, maybe restrict to that here.
+  return typeof(f(zero.(args)...))
+  # return Base.promote_op(f, args...)
+end
+
+function output_type(f::Function, as::Type{<:AbstractArray}...)
+  @assert allequal(ndims.(as))
+  elt = output_type(f, eltype.(as)...)
+  n = ndims(first(as))
+  # TODO: Generalize this to GPU arrays!
+  return Array{elt,n}
+end
+
+#############################################################################
+# AbstractArray
+#############################################################################
+
+# Related to:
+# https://github.com/JuliaLang/julia/issues/18161
+# https://github.com/JuliaLang/julia/issues/25107
+# https://github.com/JuliaLang/julia/issues/11557
+abstract type AbstractArrayStructure{ElType,Axes} end
+
+# TODO: Make this backwards compatible.
+# TODO: Add a default for `eltype`.
+# TODO: Change `Base.@kwdef` to `@kwdef`.
+Base.@kwdef struct ArrayStructure{ElType,Axes} <: AbstractArrayStructure{ElType,Axes}
+  eltype::ElType
+  axes::Axes
+end
+
+function output_eltype(::typeof(map_nonzeros), fmap, as::Type{<:AbstractArray}...)
+  return output_type(fmap, eltype.(as)...)
+end
+
+function output_eltype(f::typeof(map_nonzeros), fmap, as::AbstractArray...)
+  # TODO: Compute based on runtime information?
+  return output_eltype(f, fmap, typeof.(as)...)
+end
+
+function output_axes(f::typeof(map_nonzeros), fmap, as::AbstractArray...)
+  # TODO: Make this more sophisticated, BlockSparseArrays
+  # may have different block shapes.
+  @assert allequal(axes.(as))
+  return axes(first(as))
+end
+
+# Defaults to `ArrayStructure`.
+# Maybe define a `default_output_structure`?
+function output_structure(f::typeof(map_nonzeros), fmap, as::AbstractArray...)
+  return ArrayStructure(;
+    eltype=output_eltype(f, fmap, as...), axes=output_axes(f, fmap, as...)
+  )
+end
+
+# Defaults to `ArrayStructure`.
+# Maybe define a `default_output_type`?
+function output_type(f::typeof(map_nonzeros), fmap, as::AbstractArray...)
+  return Array
+end
+
+# Allocate an array with uninitialized/undefined memory
+# according the array type and structure (for example the
+# size or axes).
+function allocate(arraytype::Type{<:AbstractArray}, structure)
+  # TODO: Use `set_eltype`.
+  return arraytype{structure.eltype}(undef, structure.axes)
+end
+
+function allocate_zeros(arraytype::Type{<:AbstractArray}, structure)
+  a = allocate(arraytype, structure)
+  # Assumes `arraytype` is mutable.
+  # TODO: Use `zeros!!` or `zerovector!!` from VectorInterface.jl?
+  map!(Returns(false), a, a)
+  return a
+end
+
+function allocate_output(f::typeof(map_nonzeros), fmap, as::AbstractArray...)
+  return allocate_zeros(output_type(f, fmap, as...), output_structure(f, fmap, as...))
+end
+
+#############################################################################
+# SparseArray
+#############################################################################
+
+# TODO: Maybe store nonzero locations?
+# TODO: Make this backwards compatible.
+# TODO: Add a default for `eltype` and `zero`.
+# TODO: Change `Base.@kwdef` to `@kwdef`.
+Base.@kwdef struct SparseArrayStructure{ElType,Axes,Zero} <:
+                   AbstractArrayStructure{ElType,Axes}
+  eltype::ElType
+  axes::Axes
+  zero::Zero
+end
+
+function allocate(arraytype::Type{<:SparseArray}, structure::SparseArrayStructure)
+  # TODO: Use `set_eltype`.
+  return arraytype{structure.eltype}(structure.axes, structure.zero)
+end
+
+function output_structure(f::typeof(map_nonzeros), fmap, as::SparseArrayLike...)
+  return SparseArrayStructure(;
+    eltype=output_eltype(f, fmap, as...),
+    axes=output_axes(f, fmap, as...),
+    zero=output_zero(f, fmap, as...),
+  )
+end
+
+function output_type(f::typeof(map_nonzeros), fmap, as::SparseArrayLike...)
+  return SparseArray
+end
+
+function output_zero(f::typeof(map_nonzeros), fmap, as::SparseArrayLike...)
+  # TODO: Check they are all the same, update for now `axes`?
+  return first(as).zero
+end