Partition matrix by rows

README.md

@@ -93,16 +93,19 @@ this can be significant savings.
 The API for computing the color vector is:
 
 ```julia
-matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color())
+matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false)
 ```
 
 The first argument is the abstract matrix which represents the sparsity pattern
 of the Jacobian. The second argument is the optional choice of coloring algorithm.
 It will default to a greedy distance 1 coloring, though if your special matrix
 type has more information, like is a `Tridiagonal` or `BlockBandedMatrix`, the
-color vector will be analytically calculated instead.
+color vector will be analytically calculated instead. The variable argument
+`partition_by_rows` allows you to partition the Jacobian on the basis of rows instead
+of columns and generate a corresponding coloring vector which can be used for
+reverse-mode AD. Default value is false.
 
-The result is a vector which assigns a color to each row of the matrix.
+The result is a vector which assigns a color to each column (or row) of the matrix.
 
 ### Color-Assisted Differentiation
 
@@ -195,14 +198,14 @@ autonum_hesvec(f,x,v)
 numback_hesvec!(du,f,x,v,
                      cache1 = similar(v),
                      cache2 = similar(v))
-                     
+
 numback_hesvec(f,x,v)
 
 # Currently errors! See https://github.com/FluxML/Zygote.jl/issues/241
 autoback_hesvec!(du,f,x,v,
                      cache2 = ForwardDiff.Dual{DeivVecTag}.(x, v),
                      cache3 = ForwardDiff.Dual{DeivVecTag}.(x, v))
-                     
+
 autoback_hesvec(f,x,v)
 ```
 
@@ -211,7 +214,7 @@ the former almost always being more efficient and is thus recommended. `numback`
 `autoback` methods are numerical/ForwardDiff over reverse mode automatic differentiation
 respectively, where the reverse-mode AD is provided by Zygote.jl. Currently these methods
 are not competitive against `numauto`, but as Zygote.jl gets optimized these will likely
-be the fastest. 
+be the fastest.
 
 In addition,
 the following forms allow you to provide a gradient function `g(dx,x)` or `dx=g(x)`

src/coloring/high_level.jl

@@ -11,16 +11,16 @@ struct ContractionColor <: ColoringAlgorithm end
     The coloring defaults to a greedy distance-1 coloring.
 
 """
-function matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color())
+function matrix_colors(A::AbstractMatrix,alg::ColoringAlgorithm = GreedyD1Color(); partition_by_rows::Bool = false)
     _A = A isa SparseMatrixCSC ? A : sparse(A) # Avoid the copy
-    A_graph = matrix2graph(_A)
+    A_graph = matrix2graph(_A, partition_by_rows)
     color_graph(A_graph,alg)
 end
 
 """
     matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular})
 
-    The color vector for dense matrix and triangular matrix is simply 
+    The color vector for dense matrix and triangular matrix is simply
     `[1,2,3,...,size(A,2)]`
 """
 function matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular})
@@ -77,4 +77,4 @@ function matrix_colors(A::BandedBlockBandedMatrix)
     startinds=[endinds[i]-ncolors[i]+1 for i in 1:blockwidth]
     colors=[_cycle(startinds[blockcolors[i]]:endinds[blockcolors[i]],cols[i]) for i in 1:nblock]
     vcat(colors...)
-end
+end

src/coloring/matrix2graph.jl

@@ -6,23 +6,39 @@ sparse matrix, columns are represented with vertices
 and 2 vertices are connected with an edge only if
 the two columns are mutually orthogonal.
 """
-function matrix2graph(SparseMatrix::SparseMatrixCSC{T,Int}) where T<:Number
+function matrix2graph(SparseMatrix::SparseMatrixCSC{T,Int}, partition_by_rows::Bool) where T<:Number
     dropzeros(SparseMatrix)
     (rows_index, cols_index, val) = findnz(SparseMatrix)
 
-    V = cols = size(SparseMatrix, 2)
+    cols = size(SparseMatrix, 2)
     rows = size(SparseMatrix, 1)
 
+    partition_by_rows ? V = rows : V = cols
+
     inner = SimpleGraph(V)
     graph = VSafeGraph(inner)
 
-    for i = 1:length(cols_index)
-        cur_col = cols_index[i]
-        for j = 1:(i-1)
-            next_col = cols_index[j]
-            if cur_col != next_col
-                if rows_index[i] == rows_index[j]
-                    add_edge!(graph, cur_col, next_col)
+    if partition_by_rows
+        for i = 1:length(rows_index)
+            cur_row = rows_index[i]
+            for j = 1:(i-1)
+                next_row = rows_index[j]
+                if cur_row != next_row
+                    if cols_index[i] == cols_index[j]
+                        add_edge!(graph, cur_row, next_row)
+                    end
+                end
+            end
+        end
+    else
+        for i = 1:length(cols_index)
+            cur_col = cols_index[i]
+            for j = 1:(i-1)
+                next_col = cols_index[j]
+                if cur_col != next_col
+                    if rows_index[i] == rows_index[j]
+                        add_edge!(graph, cur_col, next_col)
+                    end
                 end
             end
         end

test/test_matrix2graph.jl

@@ -17,7 +17,7 @@ end
 
 for i in 1:20
     matrix = matrices[i]
-    g = matrix2graph(matrix)
+    g = matrix2graph(matrix, false)
     for e in edges(g)
         src = LG.src(e)
         dst = LG.dst(e)
@@ -27,3 +27,16 @@ for i in 1:20
         @test pr != 0
     end
 end
+
+for i in 1:20
+    matrix = matrices[i]
+    g = matrix2graph(matrix, true)
+    for e in edges(g)
+        src = LG.src(e)
+        dst = LG.dst(e)
+        row1 = abs.(matrix[src, :])
+        row2 = abs.(matrix[dst, :])
+        pr = row1' * row2
+        @test pr != 0
+    end
+end