Start rewrite

Author: mtfishman

src/KroneckerArrays.jl

@@ -8,7 +8,7 @@ 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/matrixalgebrakit.jl")
+# include("fillarrays/matrixalgebrakit_truncate.jl")
 
 end

src/linearalgebra.jl

@@ -30,7 +30,7 @@ function LinearAlgebra.pinv(a::KroneckerArray; kwargs...)
 end
 
 function LinearAlgebra.diag(a::KroneckerArray)
-  return copy(diagview(a))
+  return copy(DiagonalArrays.diagview(a))
 end
 
 # Allows customizing multiplication for specific types

src/matrixalgebrakit.jl

@@ -9,39 +9,60 @@ using MatrixAlgebraKit:
   default_qr_algorithm,
   default_svd_algorithm,
   eig_full!,
+  eig_full,
   eig_trunc!,
+  eig_trunc,
   eig_vals!,
+  eig_vals,
   eigh_full!,
+  eigh_full,
   eigh_trunc!,
+  eigh_trunc,
   eigh_vals!,
+  eigh_vals,
   initialize_output,
   left_null!,
+  left_null,
   left_orth!,
+  left_orth,
   left_polar!,
+  left_polar,
   lq_compact!,
+  lq_compact,
   lq_full!,
+  lq_full,
   qr_compact!,
+  qr_compact,
   qr_full!,
+  qr_full,
   right_null!,
+  right_null,
   right_orth!,
+  right_orth,
   right_polar!,
+  right_polar,
   svd_compact!,
+  svd_compact,
   svd_full!,
+  svd_full,
   svd_trunc!,
+  svd_trunc,
   svd_vals!,
+  svd_vals,
   truncate!
 
-using MatrixAlgebraKit: MatrixAlgebraKit, diagview
-# Allow customization for `Eye`.
-_diagview(a::AbstractMatrix) = diagview(a)
-function MatrixAlgebraKit.diagview(a::KroneckerMatrix)
-  return _diagview(a.a) ⊗ _diagview(a.b)
+using DiagonalArrays: DiagonalArrays, diagview
+function DiagonalArrays.diagview(a::KroneckerMatrix)
+  return diagview(arg1(a)) ⊗ diagview(arg2(a))
 end
+MatrixAlgebraKit.diagview(a::KroneckerMatrix) = diagview(a)
 
 struct KroneckerAlgorithm{A,B} <: AbstractAlgorithm
-  a::A
-  b::B
+  arg1::A
+  arg2::B
 end
+arg1(alg::KroneckerAlgorithm) = alg.arg1
+arg2(alg::KroneckerAlgorithm) = alg.arg2
 
 using MatrixAlgebraKit:
   copy_input,
@@ -62,10 +83,6 @@ using MatrixAlgebraKit:
   svd_compact,
   svd_full
 
-function _copy_input(f::F, a::AbstractMatrix) where {F}
-  return copy_input(f, a)
-end
-
 for f in [
   :eig_full,
   :eigh_full,
@@ -80,7 +97,7 @@ for f in [
 ]
   @eval begin
     function MatrixAlgebraKit.copy_input(::typeof($f), a::KroneckerMatrix)
-      return _copy_input($f, a.a) ⊗ _copy_input($f, a.b)
+      return copy_input($f, arg1(a)) ⊗ copy_input($f, arg2(a))
     end
   end
 end
@@ -93,105 +110,183 @@ for f in [
   :default_polar_algorithm,
   :default_svd_algorithm,
 ]
-  _f = Symbol(:_, f)
   @eval begin
-    function $_f(A::Type{<:AbstractMatrix}; kwargs...)
-      return $f(A; kwargs...)
-    end
     function MatrixAlgebraKit.$f(
       A::Type{<:KroneckerMatrix}; kwargs1=(;), kwargs2=(;), kwargs...
     )
       A1, A2 = argument_types(A)
       return KroneckerAlgorithm(
-        $_f(A1; kwargs..., kwargs1...), $_f(A2; kwargs..., kwargs2...)
+        $f(A1; kwargs..., kwargs1...), $f(A2; kwargs..., kwargs2...)
       )
     end
   end
 end
 
-# TODO: Delete this once https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/32 is merged.
-function MatrixAlgebraKit.default_algorithm(
-  ::typeof(qr_compact!), A::Type{<:KroneckerMatrix}; kwargs...
-)
-  return default_qr_algorithm(A; kwargs...)
-end
-# TODO: Delete this once https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/pull/32 is merged.
-function MatrixAlgebraKit.default_algorithm(
-  ::typeof(qr_full!), A::Type{<:KroneckerMatrix}; kwargs...
-)
-  return default_qr_algorithm(A; kwargs...)
-end
-
-# Allows overloading while avoiding type piracy.
-function _initialize_output(f::F, a::AbstractMatrix, alg::AbstractAlgorithm) where {F}
-  return initialize_output(f, a, alg)
-end
-_initialize_output(f::F, a::AbstractMatrix) where {F} = initialize_output(f, a)
-
 for f in [
-  :eig_full!,
-  :eigh_full!,
-  :qr_compact!,
-  :qr_full!,
-  :left_polar!,
-  :lq_compact!,
-  :lq_full!,
-  :right_polar!,
-  :svd_compact!,
-  :svd_full!,
+  :eig_full,
+  :eigh_full,
+  :left_polar,
+  :lq_compact,
+  :lq_full,
+  :qr_compact,
+  :qr_full,
+  :right_polar,
+  :svd_compact,
+  :svd_full,
 ]
+  f! = Symbol(f, :!)
   @eval begin
     function MatrixAlgebraKit.initialize_output(
-      ::typeof($f), a::KroneckerMatrix, alg::KroneckerAlgorithm
+      ::typeof($f!), a, alg::KroneckerAlgorithm
     )
-      return _initialize_output($f, a.a, alg.a) .⊗ _initialize_output($f, a.b, alg.b)
+      return nothing
     end
-    function MatrixAlgebraKit.$f(
+    function MatrixAlgebraKit.$f!(
       a::KroneckerMatrix, F, alg::KroneckerAlgorithm; kwargs1=(;), kwargs2=(;), kwargs...
     )
-      $f(a.a, Base.Fix2(getfield, :a).(F), alg.a; kwargs..., kwargs1...)
-      $f(a.b, Base.Fix2(getfield, :b).(F), alg.b; kwargs..., kwargs2...)
-      return F
+      a1 = $f(arg1(a), arg1(alg); kwargs..., kwargs1...)
+      a2 = $f(arg2(a), arg2(alg); kwargs..., kwargs2...)
+      return a1 .⊗ a2
     end
   end
 end
 
-for f in [:eig_vals!, :eigh_vals!, :svd_vals!]
+for f in [
+  :eig_vals,
+  :eigh_vals,
+  :svd_vals,
+]
+  f! = Symbol(f, :!)
   @eval begin
     function MatrixAlgebraKit.initialize_output(
-      ::typeof($f), a::KroneckerMatrix, alg::KroneckerAlgorithm
+      ::typeof($f!), a, alg::KroneckerAlgorithm
     )
-      return _initialize_output($f, a.a, alg.a) ⊗ _initialize_output($f, a.b, alg.b)
+      return nothing
     end
-    function MatrixAlgebraKit.$f(a::KroneckerMatrix, F, alg::KroneckerAlgorithm)
-      $f(a.a, F.a, alg.a)
-      $f(a.b, F.b, alg.b)
-      return F
+    function MatrixAlgebraKit.$f!(
+      a::KroneckerMatrix, F, alg::KroneckerAlgorithm; kwargs1=(;), kwargs2=(;), kwargs...
+    )
+      a1 = $f(arg1(a), arg1(alg); kwargs..., kwargs1...)
+      a2 = $f(arg2(a), arg2(alg); kwargs..., kwargs2...)
+      return a1 ⊗ a2
     end
   end
 end
 
-for f in [:left_orth!, :right_orth!]
+for f in [:left_orth, :right_orth]
+  f! = Symbol(f, :!)
   @eval begin
-    function MatrixAlgebraKit.initialize_output(::typeof($f), a::KroneckerMatrix)
-      return _initialize_output($f, a.a) .⊗ _initialize_output($f, a.b)
+    function MatrixAlgebraKit.initialize_output(::typeof($f!), a::KroneckerMatrix)
+      return nothing
+    end
+    function MatrixAlgebraKit.$f!(a::KroneckerMatrix, F; kwargs1=(;), kwargs2=(;), kwargs...)
+      a1 = $f(arg1(a); kwargs..., kwargs1...)
+      a2 = $f(arg2(a); kwargs..., kwargs2...)
+      return a1 .⊗ a2
     end
   end
 end
 
-for f in [:left_null!, :right_null!]
-  _f = Symbol(:_, f)
+for f in [:left_null, :right_null]
+  f! = Symbol(f, :!)
   @eval begin
     function MatrixAlgebraKit.initialize_output(::typeof($f), a::KroneckerMatrix)
-      return _initialize_output($f, a.a) ⊗ _initialize_output($f, a.b)
+      return nothing
     end
-    function $_f(a::AbstractMatrix, F; kwargs...)
-      return $f(a, F; kwargs...)
+    function MatrixAlgebraKit.$f!(a::KroneckerMatrix, F; kwargs1=(;), kwargs2=(;), kwargs...)
+      a1 = $f(arg1(a); kwargs..., kwargs1...)
+      a2 = $f(arg2(a); kwargs..., kwargs2...)
+      return a1 ⊗ a2
     end
-    function MatrixAlgebraKit.$f(a::KroneckerMatrix, F; kwargs1=(;), kwargs2=(;), kwargs...)
-      $_f(a.a, F.a; kwargs..., kwargs1...)
-      $_f(a.b, F.b; kwargs..., kwargs2...)
-      return F
+  end
+end
+
+# Truncation
+
+using MatrixAlgebraKit: TruncationStrategy, findtruncated, truncate!
+
+struct KroneckerTruncationStrategy{T<:TruncationStrategy} <: TruncationStrategy
+  strategy::T
+end
+
+## # Avoid instantiating the identity.
+## function Base.getindex(a::EyeKronecker, I::Vararg{CartesianProduct{Colon},2})
+##   return a.a ⊗ a.b[I[1].b, I[2].b]
+## end
+## function Base.getindex(a::KroneckerEye, I::Vararg{CartesianProduct{<:Any,Colon},2})
+##   return a.a[I[1].a, I[2].a] ⊗ a.b
+## end
+## function Base.getindex(a::EyeEye, I::Vararg{CartesianProduct{Colon,Colon},2})
+##   return a
+## end
+
+## using FillArrays: OnesVector
+## const OnesKroneckerVector{T,A<:OnesVector{T},B<:AbstractVector{T}} = KroneckerVector{T,A,B}
+## const KroneckerOnesVector{T,A<:AbstractVector{T},B<:OnesVector{T}} = KroneckerVector{T,A,B}
+## const OnesVectorOnesVector{T,A<:OnesVector{T},B<:OnesVector{T}} = KroneckerVector{T,A,B}
+
+axis(a) = only(axes(a))
+
+## # Convert indices determined with a generic call to `findtruncated` to indices
+## # more suited for a KroneckerVector.
+## function to_truncated_indices(values::OnesKroneckerVector, I)
+##   prods = cartesianproduct(axis(values))[I]
+##   I_id = only(to_indices(arg1(values), (:,)))
+##   I_data = unique(arg2.(prods))
+##   # Drop truncations that occur within the identity.
+##   I_data = filter(I_data) do i
+##     return count(x -> arg2(x) == i, prods) == length(arg2(values))
+##   end
+##   return I_id × I_data
+## end
+## function to_truncated_indices(values::KroneckerOnesVector, I)
+##   #I = findtruncated(Vector(values), strategy.strategy)
+##   prods = cartesianproduct(axis(values))[I]
+##   I_data = unique(arg1.(prods))
+##   # Drop truncations that occur within the identity.
+##   I_data = filter(I_data) do i
+##     return count(x -> arg1(x) == i, prods) == length(arg2(values))
+##   end
+##   I_id = only(to_indices(arg2(values), (:,)))
+##   return I_data × I_id
+## end
+function to_truncated_indices(values::KroneckerVector, I)
+  return throw(ArgumentError("Not implemented"))
+end
+
+function MatrixAlgebraKit.findtruncated(
+  values::KroneckerVector, strategy::KroneckerTruncationStrategy
+)
+  I = findtruncated(Vector(values), strategy.strategy)
+  return to_truncated_indices(values, I)
+end
+
+for f in [:eig_trunc!, :eigh_trunc!]
+  @eval begin
+    function MatrixAlgebraKit.truncate!(
+      ::typeof($f), DV::NTuple{2,KroneckerMatrix}, strategy::TruncationStrategy
+    )
+      return truncate!($f, DV, KroneckerTruncationStrategy(strategy))
+    end
+    function MatrixAlgebraKit.truncate!(
+      ::typeof($f), (D, V)::NTuple{2,KroneckerMatrix}, strategy::KroneckerTruncationStrategy
+    )
+      I = findtruncated(diagview(D), strategy)
+      return (D[I, I], V[(:) × (:), I])
     end
   end
 end
+
+function MatrixAlgebraKit.truncate!(
+  f::typeof(svd_trunc!), USVᴴ::NTuple{3,KroneckerMatrix}, strategy::TruncationStrategy
+)
+  return truncate!(f, USVᴴ, KroneckerTruncationStrategy(strategy))
+end
+function MatrixAlgebraKit.truncate!(
+  ::typeof(svd_trunc!),
+  (U, S, Vᴴ)::NTuple{3,KroneckerMatrix},
+  strategy::KroneckerTruncationStrategy,
+)
+  I = findtruncated(diagview(S), strategy)
+  return (U[(:) × (:), I], S[I, I], Vᴴ[I, (:) × (:)])
+end