Update argument names

Author: mtfishman

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)
 
 function Base.show(io::IO, a::CartesianPair)
-  print(io, a.a, " × ", a.b)
+  print(io, arg1(a), " × ", arg2(a))
   return nothing
 end
 
@@ -20,16 +20,16 @@ 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)
@@ -38,7 +38,7 @@ function Base.show(io::IO, ::MIME"text/plain", a::CartesianProduct)
 end
 
 ×(a::AbstractVector, b::AbstractVector) = CartesianProduct(a, b)
-Base.length(a::CartesianProduct) = length(a.a) * length(a.b)
+Base.length(a::CartesianProduct) = length(arg1(a)) * length(arg2(a))
 Base.size(a::CartesianProduct) = (length(a),)
 
 function Base.getindex(a::CartesianProduct, i::CartesianProduct)

src/kroneckerarray.jl

@@ -25,9 +25,10 @@ function _convert(A::Type{<:Diagonal}, a::AbstractMatrix)
   return isdiag(a) ? _construct(A, a) : throw(InexactError(:convert, A, a))
 end
 
-struct KroneckerArray{T,N,A<:AbstractArray{T,N},B<:AbstractArray{T,N}} <: AbstractArray{T,N}
-  a::A
-  b::B
+struct KroneckerArray{T,N,A1<:AbstractArray{T,N},A2<:AbstractArray{T,N}} <:
+       AbstractArray{T,N}
+  arg1::A1
+  arg2::A2
 end
 function KroneckerArray(a::AbstractArray, b::AbstractArray)
   if ndims(a) != ndims(b)
@@ -38,11 +39,15 @@ function KroneckerArray(a::AbstractArray, b::AbstractArray)
   elt = promote_type(eltype(a), eltype(b))
   return _convert(AbstractArray{elt}, a) ⊗ _convert(AbstractArray{elt}, b)
 end
-const KroneckerMatrix{T,A<:AbstractMatrix{T},B<:AbstractMatrix{T}} = KroneckerArray{T,2,A,B}
-const KroneckerVector{T,A<:AbstractVector{T},B<:AbstractVector{T}} = KroneckerArray{T,1,A,B}
+const KroneckerMatrix{T,A1<:AbstractMatrix{T},A2<:AbstractMatrix{T}} = KroneckerArray{
+  T,2,A1,A2
+}
+const KroneckerVector{T,A1<:AbstractVector{T},A2<:AbstractVector{T}} = KroneckerArray{
+  T,1,A1,A2
+}
 
-arg1(a::KroneckerArray) = a.a
-arg2(a::KroneckerArray) = a.b
+@inline arg1(a::KroneckerArray) = getfield(a, :arg1)
+@inline arg2(a::KroneckerArray) = getfield(a, :arg2)
 
 function mutate_active_args!(f!, f, dest, src)
   (isactive(arg1(dest)) || isactive(arg2(dest))) ||
@@ -81,8 +86,10 @@ function Base.copyto!(dest::KroneckerArray{<:Any,N}, src::KroneckerArray{<:Any,N
   return mutate_active_args!(copyto!, copy, dest, src)
 end
 
-function Base.convert(::Type{KroneckerArray{T,N,A,B}}, a::KroneckerArray) where {T,N,A,B}
-  return _convert(A, arg1(a)) ⊗ _convert(B, arg2(a))
+function Base.convert(
+  ::Type{KroneckerArray{T,N,A1,A2}}, a::KroneckerArray
+) where {T,N,A1,A2}
+  return _convert(A1, arg1(a)) ⊗ _convert(A2, arg2(a))
 end
 
 # Promote the element type if needed.
@@ -140,17 +147,17 @@ function Base.similar(
 end
 
 function Base.similar(
-  arrayt::Type{<:KroneckerArray{<:Any,<:Any,A,B}},
+  arrayt::Type{<:KroneckerArray{<:Any,<:Any,A1,A2}},
   axs::Tuple{
     CartesianProductUnitRange{<:Integer},Vararg{CartesianProductUnitRange{<:Integer}}
   },
-) where {A,B}
-  return similar(A, map(arg1, axs)) ⊗ similar(B, map(arg2, axs))
+) where {A1,A2}
+  return similar(A1, map(arg1, axs)) ⊗ similar(A2, map(arg2, axs))
 end
 function Base.similar(
-  ::Type{<:KroneckerArray{<:Any,<:Any,A,B}}, sz::Tuple{Int,Vararg{Int}}
-) where {A,B}
-  return similar(promote_type(A, B), sz)
+  ::Type{<:KroneckerArray{<:Any,<:Any,A1,A2}}, sz::Tuple{Int,Vararg{Int}}
+) where {A1,A2}
+  return similar(promote_type(A1, A2), sz)
 end
 
 function Base.similar(
@@ -243,7 +250,7 @@ end
 arguments(a::KroneckerArray) = (arg1(a), arg2(a))
 arguments(a::KroneckerArray, n::Int) = arguments(a)[n]
 argument_types(a::KroneckerArray) = argument_types(typeof(a))
-argument_types(::Type{<:KroneckerArray{<:Any,<:Any,A,B}}) where {A,B} = (A, B)
+argument_types(::Type{<:KroneckerArray{<:Any,<:Any,A1,A2}}) where {A1,A2} = (A1, A2)
 
 function Base.print_array(io::IO, a::KroneckerArray)
   Base.print_array(io, arg1(a))
@@ -362,22 +369,22 @@ function Base.reshape(
 end
 
 using Base.Broadcast: Broadcast, AbstractArrayStyle, BroadcastStyle, Broadcasted
-struct KroneckerStyle{N,A,B} <: AbstractArrayStyle{N} end
-arg1(::Type{<:KroneckerStyle{<:Any,A}}) where {A} = A
+struct KroneckerStyle{N,A1,A2} <: AbstractArrayStyle{N} end
+arg1(::Type{<:KroneckerStyle{<:Any,A1}}) where {A1} = A1
 arg1(style::KroneckerStyle) = arg1(typeof(style))
-arg2(::Type{<:KroneckerStyle{<:Any,B}}) where {B} = B
+arg2(::Type{<:KroneckerStyle{<:Any,<:Any,A2}}) where {A2} = A2
 arg2(style::KroneckerStyle) = arg2(typeof(style))
 function KroneckerStyle{N}(a::BroadcastStyle, b::BroadcastStyle) where {N}
   return KroneckerStyle{N,a,b}()
 end
 function KroneckerStyle(a::AbstractArrayStyle{N}, b::AbstractArrayStyle{N}) where {N}
   return KroneckerStyle{N}(a, b)
 end
-function KroneckerStyle{N,A,B}(v::Val{M}) where {N,A,B,M}
-  return KroneckerStyle{M,typeof(A)(v),typeof(B)(v)}()
+function KroneckerStyle{N,A1,A2}(v::Val{M}) where {N,A1,A2,M}
+  return KroneckerStyle{M,typeof(A1)(v),typeof(A2)(v)}()
 end
-function Base.BroadcastStyle(::Type{<:KroneckerArray{<:Any,N,A,B}}) where {N,A,B}
-  return KroneckerStyle{N}(BroadcastStyle(A), BroadcastStyle(B))
+function Base.BroadcastStyle(::Type{<:KroneckerArray{<:Any,N,A1,A2}}) where {N,A1,A2}
+  return KroneckerStyle{N}(BroadcastStyle(A1), BroadcastStyle(A2))
 end
 function Base.BroadcastStyle(style1::KroneckerStyle{N}, style2::KroneckerStyle{N}) where {N}
   style_a = BroadcastStyle(arg1(style1), arg1(style2))
@@ -386,9 +393,11 @@ function Base.BroadcastStyle(style1::KroneckerStyle{N}, style2::KroneckerStyle{N
   (style_b isa Broadcast.Unknown) && return Broadcast.Unknown()
   return KroneckerStyle{N}(style_a, style_b)
 end
-function Base.similar(bc::Broadcasted{<:KroneckerStyle{N,A,B}}, elt::Type, ax) where {N,A,B}
-  bc_a = Broadcasted(A, bc.f, arg1.(bc.args), arg1.(ax))
-  bc_b = Broadcasted(B, bc.f, arg2.(bc.args), arg2.(ax))
+function Base.similar(
+  bc::Broadcasted{<:KroneckerStyle{N,A1,A2}}, elt::Type, ax
+) where {N,A1,A2}
+  bc_a = Broadcasted(A1, bc.f, arg1.(bc.args), arg1.(ax))
+  bc_b = Broadcasted(A2, bc.f, arg2.(bc.args), arg2.(ax))
   a = similar(bc_a, elt)
   b = similar(bc_b, elt)
   return a ⊗ b
@@ -497,12 +506,12 @@ using Base.Broadcast: broadcasted
 # Represents broadcast operations that can be applied Kronecker-wise,
 # i.e. independently to each argument of the Kronecker product.
 # Note that not all broadcast operations can be mapped to this.
-struct KroneckerBroadcasted{A,B}
-  a::A
-  b::B
+struct KroneckerBroadcasted{A1,A2}
+  arg1::A1
+  arg2::A2
 end
-arg1(a::KroneckerBroadcasted) = a.a
-arg2(a::KroneckerBroadcasted) = a.b
+@inline arg1(a::KroneckerBroadcasted) = getfield(a, :arg1)
+@inline arg2(a::KroneckerBroadcasted) = getfield(a, :arg2)
 ⊗(a::Broadcasted, b::Broadcasted) = KroneckerBroadcasted(a, b)
 ⊗(a::Broadcasted, b) = KroneckerBroadcasted(a, b)
 ⊗(a, b::Broadcasted) = KroneckerBroadcasted(a, b)
@@ -525,18 +534,20 @@ function Base.axes(a::KroneckerBroadcasted)
 end
 
 function Base.BroadcastStyle(
-  ::Type{<:KroneckerBroadcasted{A,B}}
-) where {StyleA,StyleB,A<:Broadcasted{StyleA},B<:Broadcasted{StyleB}}
-  @assert ndims(A) == ndims(B)
-  N = ndims(A)
-  return KroneckerStyle{N}(StyleA(), StyleB())
+  ::Type{<:KroneckerBroadcasted{A1,A2}}
+) where {StyleA1,StyleA2,A1<:Broadcasted{StyleA1},A2<:Broadcasted{StyleA2}}
+  @assert ndims(A1) == ndims(A2)
+  N = ndims(A1)
+  return KroneckerStyle{N}(StyleA1(), StyleA2())
 end
 
 # Operations that preserve the Kronecker structure.
 for f in [:identity, :conj]
   @eval begin
-    function Broadcast.broadcasted(::KroneckerStyle{<:Any,A,B}, ::typeof($f), a) where {A,B}
-      return broadcasted(A, $f, arg1(a)) ⊗ broadcasted(B, $f, arg2(a))
+    function Broadcast.broadcasted(
+      ::KroneckerStyle{<:Any,A1,A2}, ::typeof($f), a
+    ) where {A1,A2}
+      return broadcasted(A1, $f, arg1(a)) ⊗ broadcasted(A2, $f, arg2(a))
     end
   end
 end

src/linearalgebra.jl

@@ -1,3 +1,4 @@
+using DiagonalArrays: δ
 using LinearAlgebra:
   LinearAlgebra,
   Diagonal,
@@ -17,7 +18,7 @@ using LinearAlgebra:
 
 using LinearAlgebra: LinearAlgebra
 function KroneckerArray(J::LinearAlgebra.UniformScaling, ax::Tuple)
-  return Eye{eltype(J)}(arg1.(ax)) ⊗ Eye{eltype(J)}(arg2.(ax))
+  return δ(eltype(J), arg1.(ax)) ⊗ δ(eltype(J), arg2.(ax))
 end
 function Base.copyto!(a::KroneckerArray, J::LinearAlgebra.UniformScaling)
   copyto!(a, KroneckerArray(J, axes(a)))
@@ -26,21 +27,15 @@ end
 
 using LinearAlgebra: LinearAlgebra, pinv
 function LinearAlgebra.pinv(a::KroneckerArray; kwargs...)
-  return pinv(a.a; kwargs...) ⊗ pinv(a.b; kwargs...)
+  return pinv(arg1(a); kwargs...) ⊗ pinv(arg2(a); kwargs...)
 end
 
 function LinearAlgebra.diag(a::KroneckerArray)
   return copy(DiagonalArrays.diagview(a))
 end
 
-# Allows customizing multiplication for specific types
-# such as `Eye * Eye`, which doesn't return `Eye`.
-function _mul(a::AbstractArray, b::AbstractArray)
-  return a * b
-end
-
 function Base.:*(a::KroneckerArray, b::KroneckerArray)
-  return _mul(a.a, b.a) ⊗ _mul(a.b, b.b)
+  return (arg1(a) * arg1(b)) ⊗ (arg2(a) * arg2(b))
 end
 
 function LinearAlgebra.mul!(
@@ -53,17 +48,20 @@ function LinearAlgebra.mul!(
         "Can't multiple KroneckerArrays with nonzero β and nonzero destination."
       ),
     )
-  mul!(c.a, a.a, b.a)
-  mul!(c.b, a.b, b.b, α, β)
+  # TODO: Only perform in-place operation on the non-active argument(s).
+  mul!(arg1(c), arg1(a), arg1(b))
+  mul!(arg2(c), arg2(a), arg2(b), α, β)
   return c
 end
 
+using LinearAlgebra: tr
 function LinearAlgebra.tr(a::KroneckerArray)
-  return tr(a.a) ⊗ tr(a.b)
+  return tr(arg1(a)) * tr(arg2(a))
 end
 
+using LinearAlgebra: norm
 function LinearAlgebra.norm(a::KroneckerArray, p::Int=2)
-  return norm(a.a, p) ⊗ norm(a.b, p)
+  return norm(arg1(a), p) * norm(arg2(a), p)
 end
 
 # Matrix functions
@@ -113,45 +111,54 @@ for f in MATRIX_FUNCTIONS
   end
 end
 
-using LinearAlgebra: checksquare
+# `DiagonalArrays.issquare` and `DiagonalArrays.checksquare` are more general
+# than `LinearAlgebra.checksquare`, for example it compares axes and can check
+# that the codomain and domain are dual of each other.
+using DiagonalArrays: DiagonalArrays, checksquare, issquare
+function DiagonalArrays.issquare(a::KroneckerArray)
+  return issquare(arg1(a)) && issquare(arg2(a))
+end
+
+using LinearAlgebra: det
 function LinearAlgebra.det(a::KroneckerArray)
-  checksquare(a.a)
-  checksquare(a.b)
-  return det(a.a) ^ size(a.b, 1) * det(a.b) ^ size(a.a, 1)
+  checksquare(a)
+  return det(arg1(a)) ^ size(arg2(a), 1) * det(arg2(a)) ^ size(arg1(a), 1)
 end
 
 function LinearAlgebra.svd(a::KroneckerArray)
-  Fa = svd(a.a)
-  Fb = svd(a.b)
-  return SVD(Fa.U ⊗ Fb.U, Fa.S ⊗ Fb.S, Fa.Vt ⊗ Fb.Vt)
+  F1 = svd(arg1(a))
+  F2 = svd(arg2(a))
+  return SVD(F1.U ⊗ F2.U, F1.S ⊗ F2.S, F1.Vt ⊗ F2.Vt)
 end
 function LinearAlgebra.svdvals(a::KroneckerArray)
-  return svdvals(a.a) ⊗ svdvals(a.b)
+  return svdvals(arg1(a)) ⊗ svdvals(arg2(a))
 end
 function LinearAlgebra.eigen(a::KroneckerArray)
-  Fa = eigen(a.a)
-  Fb = eigen(a.b)
-  return Eigen(Fa.values ⊗ Fb.values, Fa.vectors ⊗ Fb.vectors)
+  F1 = eigen(arg1(a))
+  F2 = eigen(arg2(a))
+  return Eigen(F1.values ⊗ F2.values, F1.vectors ⊗ F2.vectors)
 end
 function LinearAlgebra.eigvals(a::KroneckerArray)
-  return eigvals(a.a) ⊗ eigvals(a.b)
+  return eigvals(arg1(a)) ⊗ eigvals(arg2(a))
 end
 
-struct KroneckerQ{A,B}
-  a::A
-  b::B
+struct KroneckerQ{A1,A2}
+  arg1::A1
+  arg2::A2
 end
+@inline arg1(a::KroneckerQ) = getfield(a, :arg1)
+@inline arg2(a::KroneckerQ) = getfield(a, :arg2)
 function Base.:*(a::KroneckerQ, b::KroneckerQ)
-  return (a.a * b.a) ⊗ (a.b * b.b)
+  return (arg1(a) * arg1(b)) ⊗ (arg2(a) * arg2(b))
 end
-function Base.:*(a::KroneckerQ, b::KroneckerArray)
-  return (a.a * b.a) ⊗ (a.b * b.b)
+function Base.:*(a1::KroneckerQ, a2::KroneckerArray)
+  return (arg1(a1) * arg1(a2)) ⊗ (arg2(a1) * arg2(a2))
 end
-function Base.:*(a::KroneckerArray, b::KroneckerQ)
-  return (a.a * b.a) ⊗ (a.b * b.b)
+function Base.:*(a1::KroneckerArray, a2::KroneckerQ)
+  return (arg1(a1) * arg1(a2)) ⊗ (arg2(a1) * arg2(a2))
 end
 function Base.adjoint(a::KroneckerQ)
-  return KroneckerQ(a.a', a.b')
+  return KroneckerQ(arg1(a)', arg2(a)')
 end
 
 struct KroneckerQR{QQ,RR}
@@ -165,8 +172,8 @@ function ⊗(a::LinearAlgebra.QRCompactWYQ, b::LinearAlgebra.QRCompactWYQ)
   return KroneckerQ(a, b)
 end
 function LinearAlgebra.qr(a::KroneckerArray)
-  Fa = qr(a.a)
-  Fb = qr(a.b)
+  Fa = qr(arg1(a))
+  Fb = qr(arg2(a))
   return KroneckerQR(Fa.Q ⊗ Fb.Q, Fa.R ⊗ Fb.R)
 end
 
@@ -181,7 +188,7 @@ function ⊗(a::LinearAlgebra.LQPackedQ, b::LinearAlgebra.LQPackedQ)
   return KroneckerQ(a, b)
 end
 function LinearAlgebra.lq(a::KroneckerArray)
-  Fa = lq(a.a)
-  Fb = lq(a.b)
+  Fa = lq(arg1(a))
+  Fb = lq(arg2(a))
   return KroneckerLQ(Fa.L ⊗ Fb.L, Fa.Q ⊗ Fb.Q)
 end