Rewrite MatrixAlgebraKit factorization definitions (#43)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "KroneckerArrays"
 uuid = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
 authors = ["ITensor developers <[email protected]> and contributors"]
-version = "0.1.31"
+version = "0.2.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -12,6 +12,7 @@ GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MapBroadcast = "ebd9b9da-f48d-417c-9660-449667d60261"
 MatrixAlgebraKit = "6c742aac-3347-4629-af66-fc926824e5e4"
+TypeParameterAccessors = "7e5a90cf-f82e-492e-a09b-e3e26432c138"
 
 [weakdeps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
@@ -29,12 +30,13 @@ Adapt = "4.3"
 BlockArrays = "1.6"
 BlockSparseArrays = "0.9, 0.10.3"
 DerivableInterfaces = "0.5.3"
-DiagonalArrays = "0.3.11"
+DiagonalArrays = "0.3.19"
 FillArrays = "1.13"
 GPUArraysCore = "0.2"
 LinearAlgebra = "1.10"
 MapBroadcast = "0.1.10"
 MatrixAlgebraKit = "0.2, 0.3"
 TensorAlgebra = "0.3.10"
 TensorProducts = "0.1.7"
+TypeParameterAccessors = "0.4.2"
 julia = "1.10"

docs/Project.toml

@@ -6,4 +6,4 @@ KroneckerArrays = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
 [compat]
 Documenter = "1"
 Literate = "2"
-KroneckerArrays = "0.1"
+KroneckerArrays = "0.2"

examples/Project.toml

@@ -2,4 +2,4 @@
 KroneckerArrays = "05d0b138-81bc-4ff7-84be-08becefb1ccc"
 
 [compat]
-KroneckerArrays = "0.1"
+KroneckerArrays = "0.2"

ext/KroneckerArraysBlockSparseArraysExt/KroneckerArraysBlockSparseArraysExt.jl

@@ -36,8 +36,7 @@ end
 
 using BlockArrays: AbstractBlockedUnitRange
 using BlockSparseArrays: Block, ZeroBlocks, eachblockaxis, mortar_axis
-using KroneckerArrays: KroneckerArrays, KroneckerArray, ⊗, arg1, arg2, _similar
-using BlockSparseArrays.TypeParameterAccessors: unwrap_array_type
+using KroneckerArrays: KroneckerArrays, KroneckerArray, ⊗, arg1, arg2, isactive
 
 function KroneckerArrays.arg1(r::AbstractBlockedUnitRange)
   return mortar_axis(arg1.(eachblockaxis(r)))
@@ -57,30 +56,25 @@ function block_axes(ax::NTuple{N,AbstractUnitRange{<:Integer}}, I::Block{N}) whe
   return block_axes(ax, Tuple(I)...)
 end
 
+using DiagonalArrays: ShapeInitializer
+
 ## TODO: Is this needed?
 function Base.getindex(
-  a::ZeroBlocks{N,KroneckerArray{T,N,A,B}}, I::Vararg{Int,N}
-) where {T,N,A<:AbstractArray{T,N},B<:AbstractArray{T,N}}
+  a::ZeroBlocks{N,KroneckerArray{T,N,A1,A2}}, I::Vararg{Int,N}
+) where {T,N,A1<:AbstractArray{T,N},A2<:AbstractArray{T,N}}
   ax_a1 = map(arg1, a.parentaxes)
   ax_a2 = map(arg2, a.parentaxes)
-  # TODO: Instead of mutability, maybe have a trait like
-  # `isstructural` or `isdata`.
-  ismut1 = ismutabletype(unwrap_array_type(A))
-  ismut2 = ismutabletype(unwrap_array_type(B))
-  (ismut1 || ismut2) || error("Can't get zero block.")
-  a1 = if ismut1
-    ZeroBlocks{N,A}(ax_a1)[I...]
-  else
-    block_ax_a1 = arg1.(block_axes(a.parentaxes, Block(I)))
-    _similar(A, block_ax_a1)
-  end
-  a2 = if ismut2
-    ZeroBlocks{N,B}(ax_a2)[I...]
-  else
-    block_ax_a2 = arg2.(block_axes(a.parentaxes, Block(I)))
-    a2 = _similar(B, block_ax_a2)
+  block_ax_a1 = arg1.(block_axes(a.parentaxes, Block(I)))
+  block_ax_a2 = arg2.(block_axes(a.parentaxes, Block(I)))
+  # TODO: Is this a good definition? It is similar to
+  # the definition of `similar` and `adapt_structure`.
+  return if isactive(A1) == isactive(A2)
+    ZeroBlocks{N,A1}(ax_a1)[I...] ⊗ ZeroBlocks{N,A2}(ax_a2)[I...]
+  elseif isactive(A1)
+    ZeroBlocks{N,A1}(ax_a1)[I...] ⊗ A2(ShapeInitializer(), block_ax_a2)
+  elseif isactive(A2)
+    A1(ShapeInitializer(), block_ax_a1) ⊗ ZeroBlocks{N,A2}(ax_a2)[I...]
   end
-  return a1 ⊗ a2
 end
 
 using BlockSparseArrays: BlockSparseArrays

src/KroneckerArrays.jl

@@ -6,9 +6,6 @@ include("cartesianproduct.jl")
 include("kroneckerarray.jl")
 include("linearalgebra.jl")
 include("matrixalgebrakit.jl")
-include("fillarrays/kroneckerarray.jl")
-include("fillarrays/linearalgebra.jl")
-include("fillarrays/matrixalgebrakit.jl")
-include("fillarrays/matrixalgebrakit_truncate.jl")
+include("fillarrays.jl")
 
 end

src/cartesianproduct.jl

@@ -1,17 +1,17 @@
-struct CartesianPair{A,B}
-  a::A
-  b::B
+struct CartesianPair{A1,A2}
+  arg1::A1
+  arg2::A2
 end
-arguments(a::CartesianPair) = (a.a, a.b)
+arguments(a::CartesianPair) = (arg1(a), arg2(a))
 arguments(a::CartesianPair, n::Int) = arguments(a)[n]
 
-arg1(a::CartesianPair) = a.a
-arg2(a::CartesianPair) = a.b
+arg1(a::CartesianPair) = getfield(a, :arg1)
+arg2(a::CartesianPair) = getfield(a, :arg2)
 
-×(a, b) = CartesianPair(a, b)
+×(a1, a2) = CartesianPair(a1, a2)
 
 function Base.show(io::IO, a::CartesianPair)
-  print(io, a.a, " × ", a.b)
+  print(io, arg1(a), " × ", arg2(a))
   return nothing
 end
 
@@ -20,25 +20,25 @@ struct CartesianProduct{TA,TB,A<:AbstractVector{TA},B<:AbstractVector{TB}} <:
   a::A
   b::B
 end
-arguments(a::CartesianProduct) = (a.a, a.b)
+arguments(a::CartesianProduct) = (arg1(a), arg2(a))
 arguments(a::CartesianProduct, n::Int) = arguments(a)[n]
 
-arg1(a::CartesianProduct) = a.a
-arg2(a::CartesianProduct) = a.b
+arg1(a::CartesianProduct) = getfield(a, :a)
+arg2(a::CartesianProduct) = getfield(a, :b)
 
 Base.copy(a::CartesianProduct) = copy(arg1(a)) × copy(arg2(a))
 
 function Base.show(io::IO, a::CartesianProduct)
-  print(io, a.a, " × ", a.b)
+  print(io, arg1(a), " × ", arg2(a))
   return nothing
 end
 function Base.show(io::IO, ::MIME"text/plain", a::CartesianProduct)
   show(io, a)
   return nothing
 end
 
-×(a::AbstractVector, b::AbstractVector) = CartesianProduct(a, b)
-Base.length(a::CartesianProduct) = length(a.a) * length(a.b)
+×(a1::AbstractVector, a2::AbstractVector) = CartesianProduct(a1, a2)
+Base.length(a::CartesianProduct) = length(arg1(a)) * length(arg2(a))
 Base.size(a::CartesianProduct) = (length(a),)
 
 function Base.getindex(a::CartesianProduct, i::CartesianProduct)
@@ -118,12 +118,12 @@ end
 function CartesianProductUnitRange(p::CartesianProduct)
   return CartesianProductUnitRange(p, Base.OneTo(length(p)))
 end
-function CartesianProductUnitRange(a, b)
-  return CartesianProductUnitRange(a × b)
+function CartesianProductUnitRange(a1, a2)
+  return CartesianProductUnitRange(a1 × a2)
 end
 to_product_indices(a::AbstractVector) = a
 to_product_indices(i::Integer) = Base.OneTo(i)
-cartesianrange(a, b) = cartesianrange(to_product_indices(a) × to_product_indices(b))
+cartesianrange(a1, a2) = cartesianrange(to_product_indices(a1) × to_product_indices(a2))
 function cartesianrange(p::CartesianPair)
   p′ = to_product_indices(arg1(p)) × to_product_indices(arg2(p))
   return cartesianrange(p′)

src/fillarrays.jl

@@ -0,0 +1,68 @@
+using FillArrays: FillArrays, Ones, Zeros
+function FillArrays.fillsimilar(
+  a::Zeros{T},
+  ax::Tuple{
+    CartesianProductUnitRange{<:Integer},Vararg{CartesianProductUnitRange{<:Integer}}
+  },
+) where {T}
+  return Zeros{T}(arg1.(ax)) ⊗ Zeros{T}(arg2.(ax))
+end
+
+# Simplification rules similar to those for FillArrays.jl:
+# https://github.com/JuliaArrays/FillArrays.jl/blob/v1.13.0/src/fillbroadcast.jl
+using FillArrays: Zeros
+function Base.broadcasted(
+  style::KroneckerStyle,
+  ::typeof(+),
+  a::KroneckerArray,
+  b::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+)
+  # TODO: Promote the element types.
+  return a
+end
+function Base.broadcasted(
+  style::KroneckerStyle,
+  ::typeof(+),
+  a::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+  b::KroneckerArray,
+)
+  # TODO: Promote the element types.
+  return b
+end
+function Base.broadcasted(
+  style::KroneckerStyle,
+  ::typeof(+),
+  a::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+  b::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+)
+  # TODO: Promote the element types and axes.
+  return b
+end
+function Base.broadcasted(
+  style::KroneckerStyle,
+  ::typeof(-),
+  a::KroneckerArray,
+  b::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+)
+  # TODO: Promote the element types.
+  return a
+end
+function Base.broadcasted(
+  style::KroneckerStyle,
+  ::typeof(-),
+  a::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+  b::KroneckerArray,
+)
+  # TODO: Promote the element types.
+  # TODO: Return `broadcasted(-, b)`.
+  return -b
+end
+function Base.broadcasted(
+  style::KroneckerStyle,
+  ::typeof(-),
+  a::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+  b::KroneckerArray{<:Any,<:Any,<:Zeros,<:Zeros},
+)
+  # TODO: Promote the element types and axes.
+  return b
+end