CommutationMatrix type

replace comm_matrix helper functions with a CommutationMatrix and overloaded linalg ops

Author: alyst

src/StructuralEquationModels.jl

@@ -24,6 +24,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,68 @@
+"""
+
+    transpose_linear_indices(n, [m])
+
+Put each linear index of the *n×m* matrix to the position of the
+corresponding element in the transposed matrix.
+
+## Example
+`
+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

@@ -148,106 +148,3 @@ function elimination_matrix(nobs)
     end
     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

src/loss/ML/FIML.jl

@@ -24,7 +24,7 @@ Analytic gradients are available.
 ## Implementation
 Subtype of `SemLossFunction`.
 """
-mutable struct SemFIML{INV, C, L, O, M, IM, I, T, U, W} <: SemLossFunction
+mutable struct SemFIML{INV, C, L, O, M, IM, I, T, W} <: SemLossFunction
     inverses::INV #preallocated inverses of imp_cov
     choleskys::C #preallocated choleskys
     logdets::L #logdets of implied covmats
@@ -37,7 +37,7 @@ mutable struct SemFIML{INV, C, L, O, M, IM, I, T, U, W} <: SemLossFunction
 
     mult::T
 
-    commutation_indices::U
+    commutator::CommutationMatrix
 
     interaction::W
 end
@@ -64,8 +64,6 @@ function SemFIML(; observed, specification, kwargs...)
     ∇ind =
         [findall(x -> !(x[1] ∈ ind || x[2] ∈ ind), ∇ind) for ind in patterns_not(observed)]
 
-    commutation_indices = get_commutation_lookup(get_n_nodes(specification)^2)
-
     return SemFIML(
         inverses,
         choleskys,
@@ -75,7 +73,7 @@ function SemFIML(; observed, specification, kwargs...)
         meandiff,
         imp_inv,
         mult,
-        commutation_indices,
+        CommutationMatrix(get_n_nodes(specification)),
         nothing,
     )
 end
@@ -163,10 +161,9 @@ function ∇F_fiml_outer(JΣ, Jμ, imply, model, semfiml)
     Iₙ = sparse(1.0I, size(A(imply))...)
     P = kron(F⨉I_A⁻¹(imply), F⨉I_A⁻¹(imply))
     Q = kron(S(imply) * I_A⁻¹(imply)', Iₙ)
-    #commutation_matrix_pre_square_add!(Q, Q)
-    Q2 = commutation_matrix_pre_square(Q, semfiml.commutation_indices)
+    Q .+= semfiml.commutator * Q
 
-    ∇Σ = P * (∇S(imply) + (Q + Q2) * ∇A(imply))
+    ∇Σ = P * (∇S(imply) + Q * ∇A(imply))
 
     ∇μ =
         F⨉I_A⁻¹(imply) * ∇M(imply) +