Merge branch 'kmp5/feature/FieldTypes' of github.com:kmp5VT/ITensors.…

…jl into kmp5/feature/FieldTypes

NDTensors/Project.toml

@@ -4,6 +4,7 @@ authors = ["Matthew Fishman <[email protected]>"]
 version = "0.2.22"
 
 [deps]
+Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -17,6 +18,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 HDF5 = "f67ccb44-e63f-5c2f-98bd-6dc0ccc4ba2f"
 InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MappedArrays = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
@@ -42,7 +44,9 @@ NDTensorsOctavianExt = "Octavian"
 NDTensorsTBLISExt = "TBLIS"
 
 [compat]
+Accessors = "0.1.33"
 Adapt = "3.7"
+ArrayLayouts = "1.4"
 BlockArrays = "0.16"
 Compat = "4.9"
 Dictionaries = "0.3.5"

NDTensors/src/NDTensors.jl

@@ -27,12 +27,12 @@ for lib in [
   :BroadcastMapConversion,
   :Unwrap,
   :RankFactorization,
+  :GradedAxes,
   :TensorAlgebra,
   :SparseArrayInterface,
   :SparseArrayDOKs,
   :DiagonalArrays,
   :BlockSparseArrays,
-  :GradedAxes,
   :NamedDimsArrays,
   :SmallVectors,
   :SortedSets,

NDTensors/src/lib/BlockSparseArrays/examples/README.jl

@@ -64,6 +64,8 @@ function main()
   @test permuted_b == permutedims(Array(b), (2, 1))
 
   @test b + b ≈ Array(b) + Array(b)
+  @test b + b isa BlockSparseArray
+  @test block_nstored(b + b) == 2
 
   scaled_b = 2b
   @test scaled_b ≈ 2Array(b)

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysExtensions/BlockArraysExtensions.jl

@@ -1,5 +1,6 @@
-using BlockArrays: Block
+using BlockArrays: AbstractBlockArray, Block, blockedrange
 using Dictionaries: Dictionary, Indices
+using ..SparseArrayInterface: stored_indices
 
 # TODO: Use `Tuple` conversion once
 # BlockArrays.jl PR is merged.
@@ -12,3 +13,28 @@ end
 function blocks_to_cartesianindices(d::Dictionary{<:Block})
   return Dictionary(blocks_to_cartesianindices(eachindex(d)), d)
 end
+
+function block_reshape(a::AbstractBlockArray, dims::Tuple{Vararg{Vector{Int}}})
+  return block_reshape(a, blockedrange.(dims))
+end
+
+function block_reshape(a::AbstractBlockArray, dims::Vararg{Vector{Int}})
+  return block_reshape(a, dims)
+end
+
+tuple_oneto(n) = ntuple(identity, n)
+
+function block_reshape(a::AbstractBlockArray, axes::Tuple{Vararg{BlockedUnitRange}})
+  reshaped_blocks_a = reshape(blocks(a), blocklength.(axes))
+  reshaped_a = similar(a, axes)
+  for I in stored_indices(reshaped_blocks_a)
+    block_size_I = map(i -> length(axes[i][Block(I[i])]), tuple_oneto(length(axes)))
+    # TODO: Better converter here.
+    reshaped_a[Block(Tuple(I))] = reshape(reshaped_blocks_a[I], block_size_I)
+  end
+  return reshaped_a
+end
+
+function block_reshape(a::AbstractBlockArray, axes::Vararg{BlockedUnitRange})
+  return block_reshape(a, axes)
+end

NDTensors/src/lib/BlockSparseArrays/src/BlockArraysSparseArrayInterfaceExt/BlockArraysSparseArrayInterfaceExt.jl

@@ -0,0 +1,5 @@
+using ..SparseArrayInterface: SparseArrayInterface, nstored
+
+function SparseArrayInterface.nstored(a::AbstractBlockArray)
+  return sum(b -> nstored(b), blocks(a))
+end

NDTensors/src/lib/BlockSparseArrays/src/BlockSparseArrays.jl

@@ -2,12 +2,16 @@ module BlockSparseArrays
 include("blocksparsearrayinterface/blocksparsearrayinterface.jl")
 include("blocksparsearrayinterface/linearalgebra.jl")
 include("blocksparsearrayinterface/blockzero.jl")
+include("blocksparsearrayinterface/broadcast.jl")
 include("abstractblocksparsearray/abstractblocksparsearray.jl")
+include("abstractblocksparsearray/wrappedabstractblocksparsearray.jl")
 include("abstractblocksparsearray/abstractblocksparsematrix.jl")
 include("abstractblocksparsearray/abstractblocksparsevector.jl")
 include("abstractblocksparsearray/arraylayouts.jl")
 include("abstractblocksparsearray/linearalgebra.jl")
+include("abstractblocksparsearray/broadcast.jl")
 include("blocksparsearray/defaults.jl")
 include("blocksparsearray/blocksparsearray.jl")
 include("BlockArraysExtensions/BlockArraysExtensions.jl")
+include("BlockArraysSparseArrayInterfaceExt/BlockArraysSparseArrayInterfaceExt.jl")
 end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/abstractblocksparsearray.jl

@@ -9,13 +9,7 @@ abstract type AbstractBlockSparseArray{T,N} <: AbstractBlockArray{T,N} end
 # Base `AbstractArray` interface
 Base.axes(::AbstractBlockSparseArray) = error("Not implemented")
 
-# BlockArrays `AbstractBlockArray` interface
-BlockArrays.blocks(::AbstractBlockSparseArray) = error("Not implemented")
-
-blocktype(a::AbstractBlockSparseArray) = eltype(blocks(a))
-
 blockstype(::Type{<:AbstractBlockSparseArray}) = error("Not implemented")
-blocktype(arraytype::Type{<:AbstractBlockSparseArray}) = eltype(blockstype(arraytype))
 
 # Base `AbstractArray` interface
 function Base.getindex(a::AbstractBlockSparseArray{<:Any,N}, I::Vararg{Int,N}) where {N}
@@ -28,30 +22,3 @@ function Base.setindex!(
   blocksparse_setindex!(a, value, I...)
   return a
 end
-
-function Base.setindex!(
-  a::AbstractBlockSparseArray{<:Any,N}, value, I::BlockIndex{N}
-) where {N}
-  blocksparse_setindex!(a, value, I)
-  return a
-end
-
-function Base.setindex!(a::AbstractBlockSparseArray{<:Any,N}, value, I::Block{N}) where {N}
-  blocksparse_setindex!(a, value, I)
-  return a
-end
-
-# `BlockArrays` interface
-function BlockArrays.viewblock(
-  a::AbstractBlockSparseArray{<:Any,N}, I::Block{N,Int}
-) where {N}
-  return blocksparse_viewblock(a, I)
-end
-
-# Needed by `BlockArrays` matrix multiplication interface
-function Base.similar(
-  arraytype::Type{<:AbstractBlockSparseArray{T}}, axes::Tuple{Vararg{BlockedUnitRange}}
-) where {T}
-  # TODO: Make generic for GPU!
-  return BlockSparseArray{T}(undef, axes)
-end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/arraylayouts.jl

@@ -3,6 +3,7 @@ using BlockArrays: BlockLayout
 using ..SparseArrayInterface: SparseLayout
 using LinearAlgebra: mul!
 
+# TODO: Generalize to `BlockSparseArrayLike`.
 function ArrayLayouts.MemoryLayout(arraytype::Type{<:AbstractBlockSparseArray})
   outer_layout = typeof(MemoryLayout(blockstype(arraytype)))
   inner_layout = typeof(MemoryLayout(blocktype(arraytype)))

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/broadcast.jl

@@ -0,0 +1,5 @@
+using Base.Broadcast: Broadcast
+
+function Broadcast.BroadcastStyle(arraytype::Type{<:BlockSparseArrayLike})
+  return BlockSparseArrayStyle{ndims(arraytype)}()
+end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/map.jl

@@ -0,0 +1,42 @@
+using ArrayLayouts: LayoutArray
+
+# Map
+function Base.map!(f, a_dest::AbstractArray, a_srcs::Vararg{BlockSparseArrayLike})
+  blocksparse_map!(f, a_dest, a_srcs...)
+  return a_dest
+end
+
+function Base.copy!(a_dest::AbstractArray, a_src::BlockSparseArrayLike)
+  blocksparse_copy!(a_dest, a_src)
+  return a_dest
+end
+
+function Base.copyto!(a_dest::AbstractArray, a_src::BlockSparseArrayLike)
+  blocksparse_copyto!(a_dest, a_src)
+  return a_dest
+end
+
+# Fix ambiguity error
+function Base.copyto!(a_dest::LayoutArray, a_src::BlockSparseArrayLike)
+  blocksparse_copyto!(a_dest, a_src)
+  return a_dest
+end
+
+function Base.permutedims!(a_dest::AbstractArray, a_src::BlockSparseArrayLike, perm)
+  blocksparse_permutedims!(a_dest, a_src, perm)
+  return a_dest
+end
+
+function Base.mapreduce(f, op, as::Vararg{BlockSparseArrayLike}; kwargs...)
+  return blocksparse_mapreduce(f, op, as...; kwargs...)
+end
+
+# TODO: Why isn't this calling `mapreduce` already?
+function Base.iszero(a::BlockSparseArrayLike)
+  return blocksparse_iszero(a)
+end
+
+# TODO: Why isn't this calling `mapreduce` already?
+function Base.isreal(a::BlockSparseArrayLike)
+  return blocksparse_isreal(a)
+end

NDTensors/src/lib/BlockSparseArrays/src/abstractblocksparsearray/wrappedabstractblocksparsearray.jl

@@ -0,0 +1,63 @@
+using Adapt: WrappedArray
+
+const WrappedAbstractBlockSparseArray{T,N,A} = WrappedArray{
+  T,N,<:AbstractBlockSparseArray,<:AbstractBlockSparseArray{T,N}
+}
+
+const BlockSparseArrayLike{T,N} = Union{
+  <:AbstractBlockSparseArray{T,N},<:WrappedAbstractBlockSparseArray{T,N}
+}
+
+# BlockArrays `AbstractBlockArray` interface
+BlockArrays.blocks(a::BlockSparseArrayLike) = blocksparse_blocks(a)
+
+blocktype(a::BlockSparseArrayLike) = eltype(blocks(a))
+
+# TODO: Use `parenttype` from `Unwrap`.
+blockstype(arraytype::Type{<:WrappedAbstractBlockSparseArray}) = parenttype(arraytype)
+
+blocktype(arraytype::Type{<:BlockSparseArrayLike}) = eltype(blockstype(arraytype))
+
+function Base.setindex!(a::BlockSparseArrayLike{<:Any,N}, value, I::BlockIndex{N}) where {N}
+  blocksparse_setindex!(a, value, I)
+  return a
+end
+
+function Base.setindex!(a::BlockSparseArrayLike{<:Any,N}, value, I::Block{N}) where {N}
+  blocksparse_setindex!(a, value, I)
+  return a
+end
+
+# Fix ambiguity error
+function Base.setindex!(a::BlockSparseArrayLike{<:Any,1}, value, I::Block{1})
+  blocksparse_setindex!(a, value, I)
+  return a
+end
+
+# `BlockArrays` interface
+# TODO: Is this needed if `blocks` is defined?
+function BlockArrays.viewblock(a::BlockSparseArrayLike{<:Any,N}, I::Block{N,Int}) where {N}
+  return blocksparse_viewblock(a, I)
+end
+
+# Needed by `BlockArrays` matrix multiplication interface
+function Base.similar(
+  arraytype::Type{<:BlockSparseArrayLike}, axes::Tuple{Vararg{BlockedUnitRange}}
+)
+  return similar(arraytype, eltype(arraytype), axes)
+end
+
+# Needed by `BlockArrays` matrix multiplication interface
+function Base.similar(
+  arraytype::Type{<:BlockSparseArrayLike}, elt::Type, axes::Tuple{Vararg{BlockedUnitRange}}
+)
+  # TODO: Make generic for GPU! Use `blocktype`.
+  return BlockSparseArray{elt}(undef, axes)
+end
+
+function Base.similar(
+  a::BlockSparseArrayLike, elt::Type, axes::Tuple{Vararg{BlockedUnitRange}}
+)
+  # TODO: Make generic for GPU! Use `blocktype`.
+  return BlockSparseArray{eltype(a)}(undef, axes)
+end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearray/blocksparsearray.jl

@@ -102,3 +102,11 @@ BlockArrays.blocks(a::BlockSparseArray) = a.blocks
 
 # TODO: Use `SetParameters`.
 blockstype(::Type{<:BlockSparseArray{<:Any,<:Any,<:Any,B}}) where {B} = B
+
+# Base interface
+function Base.similar(
+  a::AbstractBlockSparseArray, elt::Type, axes::Tuple{Vararg{BlockedUnitRange}}
+)
+  # TODO: Preserve GPU data!
+  return BlockSparseArray{elt}(undef, axes)
+end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blocksparsearrayinterface.jl

@@ -1,5 +1,10 @@
 using BlockArrays: Block, BlockIndex, block, blocks, blockcheckbounds, findblockindex
-using ..SparseArrayInterface: nstored
+using ..SparseArrayInterface: perm, iperm, nstored
+using MappedArrays: mappedarray
+
+function blocksparse_blocks(a::AbstractArray)
+  return blocks(a)
+end
 
 function blocksparse_getindex(a::AbstractArray{<:Any,N}, I::Vararg{Int,N}) where {N}
   @boundscheck checkbounds(a, I...)
@@ -24,17 +29,35 @@ function blocksparse_setindex!(a::AbstractArray{<:Any,N}, value, I::Block{N}) wh
   # TODO: Create a conversion function, say `CartesianIndex(Int.(Tuple(I)))`.
   i = I.n
   @boundscheck blockcheckbounds(a, i...)
-  blocks(a)[i...] = value
+  blocksparse_blocks(a)[i...] = value
   return a
 end
 
 function blocksparse_viewblock(a::AbstractArray{<:Any,N}, I::Block{N}) where {N}
   # TODO: Create a conversion function, say `CartesianIndex(Int.(Tuple(I)))`.
   i = I.n
   @boundscheck blockcheckbounds(a, i...)
-  return blocks(a)[i...]
+  return blocksparse_blocks(a)[i...]
 end
 
 function block_nstored(a::AbstractArray)
-  return nstored(blocks(a))
+  return nstored(blocksparse_blocks(a))
+end
+
+# Base
+function blocksparse_map!(f, a_dest::AbstractArray, a_srcs::AbstractArray...)
+  map!(f, blocks(a_dest), blocks.(a_srcs)...)
+  return a_dest
+end
+
+# PermutedDimsArray
+function blocksparse_blocks(a::PermutedDimsArray)
+  blocks_parent = blocksparse_blocks(parent(a))
+  # Lazily permute each block
+  blocks_parent_mapped = mappedarray(
+    Base.Fix2(PermutedDimsArray, perm(a)),
+    Base.Fix2(PermutedDimsArray, iperm(a)),
+    blocks_parent,
+  )
+  return PermutedDimsArray(blocks_parent_mapped, perm(a))
 end

NDTensors/src/lib/BlockSparseArrays/src/blocksparsearrayinterface/blockzero.jl

@@ -18,26 +18,22 @@ struct BlockZero{Axes}
   axes::Axes
 end
 
-function (f::BlockZero)(a::AbstractArray, I::CartesianIndex)
+function (f::BlockZero)(a::AbstractArray, I)
   return f(eltype(a), I)
 end
 
-function (f::BlockZero)(
-  arraytype::Type{<:AbstractArray{<:Any,N}}, I::CartesianIndex{N}
-) where {N}
+function (f::BlockZero)(arraytype::Type{<:AbstractArray}, I)
   # TODO: Make sure this works for sparse or block sparse blocks, immutable
   # blocks, diagonal blocks, etc.!
   return fill!(arraytype(undef, block_size(f.axes, Block(Tuple(I)))), false)
 end
 
 # Fallback so that `SparseArray` with scalar elements works.
-function (f::BlockZero)(blocktype::Type{<:Number}, I::CartesianIndex)
+function (f::BlockZero)(blocktype::Type{<:Number}, I)
   return zero(blocktype)
 end
 
 # Fallback to Array if it is abstract
-function (f::BlockZero)(
-  arraytype::Type{AbstractArray{T,N}}, I::CartesianIndex{N}
-) where {T,N}
-  return fill!(Array{T,N}(undef, block_size(f.axes, Block(Tuple(I)))), zero(T))
+function (f::BlockZero)(arraytype::Type{AbstractArray{T,N}}, I) where {T,N}
+  return f(Array{T,N}, I)
 end