Update to BlockSparseArrays v0.9

Author: mtfishman

ext/KroneckerArraysTensorProductsExt/KroneckerArraysTensorProductsExt.jl

@@ -0,0 +1,11 @@
+module KroneckerArraysTensorProductsExt
+
+using KroneckerArrays: CartesianProductOneTo, ×, arg1, arg2, cartesianrange
+using TensorProducts: TensorProducts, tensor_product
+function TensorProducts.tensor_product(a1::CartesianProductOneTo, a2::CartesianProductOneTo)
+  return cartesianrange(
+    tensor_product(arg1(a1), arg1(a2)) × tensor_product(arg2(a1), arg2(a2))
+  )
+end
+
+end

src/fillarrays/kroneckerarray.jl

@@ -41,24 +41,11 @@ function _convert(::Type{AbstractMatrix{T}}, a::RectDiagonal) where {T}
   RectDiagonal(convert(AbstractVector{T}, _diagview(a)), axes(a))
 end
 
-# Like `similar` but preserves `Eye`.
-function _similar(a::AbstractArray, elt::Type, ax::Tuple)
-  return similar(a, elt, ax)
+# Like `similar` but preserves `Eye`, `Ones`, etc.
+using FillArrays: Ones
+function _similar(arrayt::Type{<:Ones}, axs::Tuple)
+  return Ones{eltype(arrayt)}(axs)
 end
-function _similar(A::Type{<:AbstractArray}, ax::Tuple)
-  return similar(A, ax)
-end
-function _similar(a::AbstractArray, ax::Tuple)
-  return _similar(a, eltype(a), ax)
-end
-function _similar(a::AbstractArray, elt::Type)
-  return _similar(a, elt, axes(a))
-end
-function _similar(a::AbstractArray)
-  return _similar(a, eltype(a), axes(a))
-end
-
-# Like `similar` but preserves `Eye`.
 function _similar(a::Eye, elt::Type, axs::NTuple{2,AbstractUnitRange})
   return Eye{elt}(axs)
 end
@@ -77,19 +64,6 @@ end
 # Like `copy` but preserves `Eye`.
 _copy(a::Eye) = a
 
-using DerivableInterfaces: DerivableInterfaces, zero!
-function DerivableInterfaces.zero!(a::EyeKronecker)
-  zero!(a.b)
-  return a
-end
-function DerivableInterfaces.zero!(a::KroneckerEye)
-  zero!(a.a)
-  return a
-end
-function DerivableInterfaces.zero!(a::EyeEye)
-  return throw(ArgumentError("Can't zero out `Eye ⊗ Eye`."))
-end
-
 using Base.Broadcast:
   AbstractArrayStyle, AbstractArrayStyle, BroadcastStyle, Broadcasted, broadcasted
 

src/kroneckerarray.jl

@@ -54,10 +54,19 @@ function Base.convert(::Type{KroneckerArray{T,N,A,B}}, a::KroneckerArray) where
 end
 
 # Like `similar` but allows some custom behavior, such as for `FillArrays.Eye`.
-function _similar(a::AbstractArray, elt::Type, axs::Tuple{Vararg{AbstractUnitRange}})
+function _similar(a::AbstractArray, elt::Type, axs::Tuple)
   return similar(a, elt, axs)
 end
-function _similar(arrayt::Type{<:AbstractArray}, axs::Tuple{Vararg{AbstractUnitRange}})
+function _similar(a::AbstractArray, ax::Tuple)
+  return _similar(a, eltype(a), ax)
+end
+function _similar(a::AbstractArray, elt::Type)
+  return _similar(a, elt, axes(a))
+end
+function _similar(a::AbstractArray)
+  return _similar(a, eltype(a), axes(a))
+end
+function _similar(arrayt::Type{<:AbstractArray}, axs::Tuple)
   return similar(arrayt, axs)
 end
 
@@ -130,6 +139,16 @@ Base.collect(a::KroneckerArray) = kron_nd(collect(arg1(a)), collect(arg2(a)))
 
 Base.zero(a::KroneckerArray) = zero(arg1(a)) ⊗ zero(arg2(a))
 
+using DerivableInterfaces: DerivableInterfaces, zero!
+function DerivableInterfaces.zero!(a::KroneckerArray)
+  ismut1 = ismutable(arg1(a))
+  ismut2 = ismutable(arg2(a))
+  (ismut1 || ismut2) || throw(ArgumentError("Can't zero out immutable KroneckerArray."))
+  ismut1 && zero!(arg1(a))
+  ismut2 && zero!(arg2(a))
+  return a
+end
+
 function Base.Array{T,N}(a::KroneckerArray{S,N}) where {T,S,N}
   return convert(Array{T,N}, collect(a))
 end
@@ -372,13 +391,15 @@ _eltype(x) = eltype(x)
 _eltype(x::Broadcasted) = Base.promote_op(x.f, _eltype.(x.args)...)
 
 using Base.Broadcast: broadcasted
-struct KroneckerBroadcasted{A<:Broadcasted,B<:Broadcasted}
+struct KroneckerBroadcasted{A,B}
   a::A
   b::B
 end
 arg1(a::KroneckerBroadcasted) = a.a
 arg2(a::KroneckerBroadcasted) = a.b
 ⊗(a::Broadcasted, b::Broadcasted) = KroneckerBroadcasted(a, b)
+⊗(a::Broadcasted, b) = KroneckerBroadcasted(a, b)
+⊗(a, b::Broadcasted) = KroneckerBroadcasted(a, b)
 Broadcast.materialize(a::KroneckerBroadcasted) = copy(a)
 Broadcast.materialize!(dest, a::KroneckerBroadcasted) = copyto!(dest, a)
 Broadcast.broadcastable(a::KroneckerBroadcasted) = a

src/linearalgebra.jl

@@ -179,10 +179,3 @@ function LinearAlgebra.lq(a::KroneckerArray)
   Fb = lq(a.b)
   return KroneckerLQ(Fa.L ⊗ Fb.L, Fa.Q ⊗ Fb.Q)
 end
-
-using DerivableInterfaces: DerivableInterfaces, zero!
-function DerivableInterfaces.zero!(a::KroneckerArray)
-  zero!(a.a)
-  zero!(a.b)
-  return a
-end

test/test_blocksparsearrays.jl

@@ -23,7 +23,7 @@ arrayts = (Array, JLArray)
     Block(1, 1) => dev(randn(elt, 2, 2) ⊗ randn(elt, 2, 2)),
     Block(2, 2) => dev(randn(elt, 3, 3) ⊗ randn(elt, 3, 3)),
   )
-  a = dev(blocksparse(d, r, r))
+  a = dev(blocksparse(d, (r, r)))
   @test sprint(show, a) isa String
   @test sprint(show, MIME("text/plain"), a) isa String
   @test blocktype(a) === valtype(d)
@@ -45,7 +45,7 @@ arrayts = (Array, JLArray)
     Block(1, 1) => dev(randn(elt, 2, 2) ⊗ randn(elt, 2, 2)),
     Block(2, 2) => dev(randn(elt, 3, 3) ⊗ randn(elt, 3, 3)),
   )
-  a = dev(blocksparse(d, r, r))
+  a = dev(blocksparse(d, (r, r)))
   @test a[Block(2, 2)[(2:3) × (2:3), (2:3) × (2:3)]] ==
     a[Block(2, 2)][(2:3) × (2:3), (2:3) × (2:3)]
   @test a[Block(2, 2)[(:) × (2:3), (:) × (2:3)]] == a[Block(2, 2)][(:) × (2:3), (:) × (2:3)]
@@ -145,7 +145,7 @@ end
     Block(1, 1) => Eye{elt}(2, 2) ⊗ dev(randn(elt, 2, 2)),
     Block(2, 2) => Eye{elt}(3, 3) ⊗ dev(randn(elt, 3, 3)),
   )
-  a = @constinferred dev(blocksparse(d, r, r))
+  a = @constinferred dev(blocksparse(d, (r, r)))
   @test sprint(show, a) == sprint(show, Array(a))
   @test sprint(show, MIME("text/plain"), a) isa String
   @test @constinferred(blocktype(a)) === valtype(d)
@@ -167,34 +167,34 @@ end
     Block(1, 1) => dev(Eye{elt}(2, 2) ⊗ randn(elt, 2, 2)),
     Block(2, 2) => dev(Eye{elt}(3, 3) ⊗ randn(elt, 3, 3)),
   )
-  a = dev(blocksparse(d, r, r))
+  a = dev(blocksparse(d, (r, r)))
   @test a[Block(2, 2)[(2:3) × (2:3), (2:3) × (2:3)]] ==
     a[Block(2, 2)][(2:3) × (2:3), (2:3) × (2:3)]
   @test a[Block(2, 2)[(:) × (2:3), (:) × (2:3)]] == a[Block(2, 2)][(:) × (2:3), (:) × (2:3)]
   @test a[Block(2, 2)[(1:2) × (2:3), (:) × (2:3)]] ==
     a[Block(2, 2)][(1:2) × (2:3), (:) × (2:3)]
 
-  # Blockwise slicing, shows up in truncated block sparse matrix factorizations.
-  I1 = BlockIndexVector(Block(1), Base.Slice(Base.OneTo(2)) × [1])
-  I2 = BlockIndexVector(Block(2), Base.Slice(Base.OneTo(3)) × [1, 3])
-  I = [I1, I2]
-  b = a[I, I]
-  @test b[Block(1, 1)] == a[Block(1, 1)[(1:2) × [1], (1:2) × [1]]]
-  @test arg1(b[Block(1, 1)]) isa Eye
-  @test iszero(b[Block(2, 1)])
-  @test arg1(b[Block(2, 1)]) isa Eye
-  @test iszero(b[Block(1, 2)])
-  @test arg1(b[Block(1, 2)]) isa Eye
-  @test b[Block(2, 2)] == a[Block(2, 2)[(1:3) × [1, 3], (1:3) × [1, 3]]]
-  @test arg1(b[Block(2, 2)]) isa Eye
+  ## # Blockwise slicing, shows up in truncated block sparse matrix factorizations.
+  ## I1 = BlockIndexVector(Block(1), Base.Slice(Base.OneTo(2)) × [1])
+  ## I2 = BlockIndexVector(Block(2), Base.Slice(Base.OneTo(3)) × [1, 3])
+  ## I = [I1, I2]
+  ## b = a[I, I]
+  ## @test b[Block(1, 1)] == a[Block(1, 1)[(1:2) × [1], (1:2) × [1]]]
+  ## @test arg1(b[Block(1, 1)]) isa Eye
+  ## @test iszero(b[Block(2, 1)])
+  ## @test arg1(b[Block(2, 1)]) isa Eye
+  ## @test iszero(b[Block(1, 2)])
+  ## @test arg1(b[Block(1, 2)]) isa Eye
+  ## @test b[Block(2, 2)] == a[Block(2, 2)[(1:3) × [1, 3], (1:3) × [1, 3]]]
+  ## @test arg1(b[Block(2, 2)]) isa Eye
 
   # Slicing
   r = blockrange([2 × 2, 3 × 3])
   d = Dict(
     Block(1, 1) => dev(Eye{elt}(2, 2) ⊗ randn(elt, 2, 2)),
     Block(2, 2) => dev(Eye{elt}(3, 3) ⊗ randn(elt, 3, 3)),
   )
-  a = dev(blocksparse(d, r, r))
+  a = dev(blocksparse(d, (r, r)))
   i1 = Block(1)[(1:2) × (1:2)]
   i2 = Block(2)[(2:3) × (2:3)]
   I = mortar([i1, i2])
@@ -209,7 +209,7 @@ end
     Block(1, 1) => dev(Eye{elt}(2, 2) ⊗ randn(elt, 2, 2)),
     Block(2, 2) => dev(Eye{elt}(3, 3) ⊗ randn(elt, 3, 3)),
   )
-  a = dev(blocksparse(d, r, r))
+  a = dev(blocksparse(d, (r, r)))
   i1 = Block(1)[(1:2) × (1:2)]
   i2 = Block(2)[(2:3) × (2:3)]
   I = [i1, i2]
@@ -282,7 +282,7 @@ end
     Block(1, 1) => Eye{elt}(2, 2) ⊗ randn(rng, elt, 2, 2),
     Block(2, 2) => Eye{elt}(3, 3) ⊗ randn(rng, elt, 3, 3),
   )
-  a = @constinferred dev(blocksparse(d, r, r))
+  a = @constinferred dev(blocksparse(d, (r, r)))
   if arrayt === Array
     u, s, v = svd_trunc(a; trunc=(; maxrank=6))
     u′, s′, v′ = svd_trunc(Matrix(a); trunc=(; maxrank=5))
@@ -293,10 +293,10 @@ end
 
   @testset "Block deficient" begin
     da = Dict(Block(1, 1) => Eye{elt}(2, 2) ⊗ dev(randn(elt, 2, 2)))
-    a = @constinferred dev(blocksparse(da, r, r))
+    a = @constinferred dev(blocksparse(da, (r, r)))
 
     db = Dict(Block(2, 2) => Eye{elt}(3, 3) ⊗ dev(randn(elt, 3, 3)))
-    b = @constinferred dev(blocksparse(db, r, r))
+    b = @constinferred dev(blocksparse(db, (r, r)))
 
     @test Array(a + b) ≈ Array(a) + Array(b)
     @test Array(2a) ≈ 2Array(a)