CommutationMatrix type

replace comm_matrix helper functions with a CommutationMatrix and overloaded linalg ops

Author: alyst

src/StructuralEquationModels.jl

@@ -23,6 +23,10 @@ const SEM = StructuralEquationModels
 # type hierarchy
 include("types.jl")
 include("objective_gradient_hessian.jl")
+
+# helper objects and functions
+include("additional_functions/commutation_matrix.jl")
+
 # fitted objects
 include("frontend/fit/SemFit.jl")
 # specification of models

src/additional_functions/commutation_matrix.jl

@@ -0,0 +1,59 @@
+# transpose linear indices of the square n×n matrix
+# i.e.
+# 1 4
+# 2 5 =>  1 2 3
+# 3 6     4 5 6
+transpose_linear_indices(n::Integer, m::Integer = n) =
+    repeat(1:n, inner = m) .+ repeat((0:(m-1)) * n, outer = n)
+
+"""
+    CommutationMatrix(n::Integer) <: AbstractMatrix{Int}
+
+A *commutation matrix* *C* is a n²×n² matrix of 0s and 1s.
+If *vec(A)* is a vectorized form of a n×n matrix *A*,
+then ``C * vec(A) = vec(Aᵀ)``.
+"""
+struct CommutationMatrix <: AbstractMatrix{Int}
+    n::Int
+    n²::Int
+    transpose_inds::Vector{Int} # maps the linear indices of n×n matrix *B* to the indices of matrix *B'*
+
+    CommutationMatrix(n::Integer) = new(n, n^2, transpose_linear_indices(n))
+end
+
+Base.size(A::CommutationMatrix) = (A.n², A.n²)
+Base.size(A::CommutationMatrix, dim::Integer) =
+    1 <= dim <= 2 ? A.n² : throw(ArgumentError("invalid matrix dimension $dim"))
+Base.length(A::CommutationMatrix) = A.n²^2
+Base.getindex(A::CommutationMatrix, i::Int, j::Int) = j == A.transpose_inds[i] ? 1 : 0
+
+function Base.:(*)(A::CommutationMatrix, B::AbstractMatrix)
+    size(A, 2) == size(B, 1) || throw(
+        DimensionMismatch("A has $(size(A, 2)) columns, but B has $(size(B, 1)) rows"),
+    )
+    return B[A.transpose_inds, :]
+end
+
+function Base.:(*)(A::CommutationMatrix, B::SparseMatrixCSC)
+    size(A, 2) == size(B, 1) || throw(
+        DimensionMismatch("A has $(size(A, 2)) columns, but B has $(size(B, 1)) rows"),
+    )
+    return SparseMatrixCSC(
+        size(B, 1),
+        size(B, 2),
+        copy(B.colptr),
+        A.transpose_inds[B.rowval],
+        copy(B.nzval),
+    )
+end
+
+function LinearAlgebra.lmul!(A::CommutationMatrix, B::SparseMatrixCSC)
+    size(A, 2) == size(B, 1) || throw(
+        DimensionMismatch("A has $(size(A, 2)) columns, but B has $(size(B, 1)) rows"),
+    )
+
+    @inbounds for (i, rowind) in enumerate(B.rowval)
+        B.rowval[i] = A.transpose_inds[rowind]
+    end
+    return B
+end

src/additional_functions/helper.jl

@@ -142,109 +142,6 @@ function elimination_matrix(nobs)
     return L
 end
 
-function commutation_matrix(n; tosparse = false)
-    M = zeros(n^2, n^2)
-
-    for i in 1:n
-        for j in 1:n
-            M[i+n*(j-1), j+n*(i-1)] = 1.0
-        end
-    end
-
-    if tosparse
-        M = sparse(M)
-    end
-
-    return M
-end
-
-function commutation_matrix_pre_square(A)
-    n2 = size(A, 1)
-    n = Int(sqrt(n2))
-
-    ind = repeat(1:n, inner = n)
-    indadd = (0:(n-1)) * n
-    for i in 1:n
-        ind[((i-1)*n+1):i*n] .+= indadd
-    end
-
-    A_post = A[ind, :]
-
-    return A_post
-end
-
-function commutation_matrix_pre_square_add!(B, A) # comuptes B + KₙA
-    n2 = size(A, 1)
-    n = Int(sqrt(n2))
-
-    ind = repeat(1:n, inner = n)
-    indadd = (0:(n-1)) * n
-    for i in 1:n
-        ind[((i-1)*n+1):i*n] .+= indadd
-    end
-
-    @views @inbounds B .+= A[ind, :]
-
-    return B
-end
-
-function get_commutation_lookup(n2::Int64)
-    n = Int(sqrt(n2))
-    ind = repeat(1:n, inner = n)
-    indadd = (0:(n-1)) * n
-    for i in 1:n
-        ind[((i-1)*n+1):i*n] .+= indadd
-    end
-
-    lookup = Dict{Int64, Int64}()
-
-    for i in 1:n2
-        j = findall(x -> (x == i), ind)[1]
-        push!(lookup, i => j)
-    end
-
-    return lookup
-end
-
-function commutation_matrix_pre_square!(A::SparseMatrixCSC, lookup) # comuptes B + KₙA
-    for (i, rowind) in enumerate(A.rowval)
-        A.rowval[i] = lookup[rowind]
-    end
-end
-
-function commutation_matrix_pre_square!(A::SparseMatrixCSC) # computes KₙA
-    lookup = get_commutation_lookup(size(A, 2))
-    commutation_matrix_pre_square!(A, lookup)
-end
-
-function commutation_matrix_pre_square(A::SparseMatrixCSC)
-    B = copy(A)
-    commutation_matrix_pre_square!(B)
-    return B
-end
-
-function commutation_matrix_pre_square(A::SparseMatrixCSC, lookup)
-    B = copy(A)
-    commutation_matrix_pre_square!(B, lookup)
-    return B
-end
-
-function commutation_matrix_pre_square_add_mt!(B, A) # comuptes B + KₙA # 0 allocations but slower
-    n2 = size(A, 1)
-    n = Int(sqrt(n2))
-
-    indadd = (0:(n-1)) * n
-
-    Threads.@threads for i in 1:n
-        for j in 1:n
-            row = i + indadd[j]
-            @views @inbounds B[row, :] .+= A[row, :]
-        end
-    end
-
-    return B
-end
-
 # returns the vector of non-unique values in the order of appearance
 # each non-unique values is reported once
 function nonunique(values::AbstractVector)

src/loss/ML/FIML.jl

@@ -91,7 +91,7 @@ struct SemFIML{T, W} <: SemLossFunction{ExactHessian}
 
     imp_inv::Matrix{T}  # implied inverse
 
-    commutation_indices::Dict{Int, Int}
+    commutator::CommutationMatrix
 
     interaction::W
 end
@@ -104,7 +104,7 @@ function SemFIML(; observed::SemObservedMissing, specification, kwargs...)
     return SemFIML(
         [SemFIMLPattern(pat) for pat in observed.patterns],
         zeros(n_man(observed), n_man(observed)),
-        get_commutation_lookup(nvars(specification)^2),
+        CommutationMatrix(nvars(specification)),
         nothing,
     )
 end
@@ -158,13 +158,12 @@ function ∇F_fiml_outer!(G, JΣ, Jμ, fiml::SemFIML, imply::SemImplySymbolic, m
 end
 
 function ∇F_fiml_outer!(G, JΣ, Jμ, fiml::SemFIML, imply, model)
-    Iₙ = sparse(1.0I, size(imply.A)...)
     P = kron(imply.F⨉I_A⁻¹, imply.F⨉I_A⁻¹)
+    Iₙ = sparse(1.0I, size(imply.A)...)
     Q = kron(imply.S * imply.I_A⁻¹', Iₙ)
-    #commutation_matrix_pre_square_add!(Q, Q)
-    Q2 = commutation_matrix_pre_square(Q, fiml.commutation_indices)
+    Q .+= fiml.commutator * Q
 
-    ∇Σ = P * (imply.∇S + (Q + Q2) * imply.∇A)
+    ∇Σ = P * (imply.∇S + Q * imply.∇A)
 
     ∇μ = imply.F⨉I_A⁻¹ * imply.∇M + kron((imply.I_A⁻¹ * imply.M)', imply.F⨉I_A⁻¹) * imply.∇A