Cartesian indices, linear indices, and sparse matrices

briochemc · August 20, 2018, 6:56pm

I will go with a minimal working example and ask questions in bold as I go through it.

Before v0.7, I had this working code:

A = sparse([1,3,4],[2,4,3], [1,2,3], 4,5)
idx = find(A)
# make a sparse matrix using subindices
idx1, idx2 = ind2sub(size(A), idx)
B = sparse(idx1, idx2, [4,5,6], 4, 5)
# transform one of the subindices and make another matrix
f(x) = (x + 2) .% 5 + 1
C = sparse(idx1, f(idx2), [7,8,9], 4, 5)

Now a few things have changed with v0.7+. First, I now need to use the SparseArrays package and find is deprecated, so I now start with

using SparseArrays
A = sparse([1,3,4],[2,4,3], [1,2,3], 4,5)
idx = findall(!iszero, A)

Note that idx’s type is 3-element Array{CartesianIndex{2},1} now, which leads me to my first question:

How do I create a sparse matrix B from an array of CartesianIndex and values?

Currently,
```
B = sparse(idx, [4,5,6], 4, 5)
```
is not valid and raises an error.
Same goes for using linear indices:
```
lidx = LinearIndices(A)[idx]
B = sparse(lidx, [4,5,6], 4, 5)
```
also raises an error. So let’s assume I cannot do that directly with cartesian or linear indices, then I need subindices, which leads me to my second question:
How do I access the subindices of a cartesian index array?

Right now, idx contains the following:
```
3-element Array{CartesianIndex{2},1}:
 CartesianIndex(1, 2)
 CartesianIndex(4, 3)
 CartesianIndex(3, 4)
```
and I would like to use the subindices [1,4,3] and [2,3,4]. But the ind2sub function is deprecated, encouraging us to use the CartesianIndex type. I could painfully access them by looping through every CartesianIndex:
```
j = 0
subidx1 = zeros(Int64, length(idx))
subidx2 = copy(subidx1)
for i in idx
    j += 1
    subidx1[j] = i[1]
    subidx2[j] = i[2]
end
B = sparse(subidx1, subidx2, [4,5,6], 4,5)
```
But this feels “dirty” to me… I may probably be writing this the wrong way, but I wished there was a built-in way to get those sub indices. (Maybe there is?) Something like
```
subidx1, subidx2 = ind2sub(size(A), idx)
```
would be nice!
How do I apply a function to a subindex inside a cartesian index?

I may be wrong in just trying to access those subindices, but I have a need to transform one of them sometimes, by applying a function to it. Say I want to shift the second index by 3 to the right (like in the non-v0.7+ example at the start of this post), or I want to pack these indices as close to the left as possible. Again, there might be an easy way to do it, but I have not figured it out yet.

Please help

mbauman · August 20, 2018, 7:59pm

Thanks for writing this up as a comprehensive question!

It’d probably be nice to add a sparse method that takes arrays of CartesianIndexes. Here’s a simple implementation that could serve you in the mean-time:

function SparseArrays.sparse(IJ::Vector{<:CartesianIndex}, v, m, n)
    IJ′ = reinterpret(Int, reshape(IJ, 1, :))
    return sparse(view(IJ′, 1, :), view(IJ′, 2, :), v, m, n)
end

In the above I used reinterpret and lazy views to prevent needlessly allocating new vectors that contain just those “sub” indices you’re after. But you can use comprehensions to great effect to do this more simply than your for loop:
```
julia> [i[1] for i in idx]
3-element Array{Int64,1}:
 1
 4
 3

julia> [i[2] for i in idx]
3-element Array{Int64,1}:
 2
 3
 4
```

To apply a function to transform the cartesian index, I’d just work directly at the level of the cartesian index. E.g.,

julia> f(ind) = CartesianIndex(ind[1]÷2+1, ind[2]÷2+1)
f (generic function with 1 method)

julia> [f(i) for i in idx]
3-element Array{CartesianIndex{2},1}:
 CartesianIndex(1, 2)
 CartesianIndex(3, 2)
 CartesianIndex(2, 3)

julia> f.(idx) # Or just use broadcasting to the same effect
3-element Array{CartesianIndex{2},1}:
 CartesianIndex(1, 2)
 CartesianIndex(3, 2)
 CartesianIndex(2, 3)

briochemc · August 20, 2018, 8:14pm

Thanks so much for your answer! For anybody like me struggling with this, I’ll try to use your response to upgrade my non-v0.7+ code:

briochemc:

A = sparse([1,3,4], [2,4,3], [1,2,3], 4,5) 
idx = find(A) 
# make a sparse matrix using subindices
idx1, idx2 = ind2sub(size(A), idx) 
B = sparse(idx1, idx2, [4,5,6], 4, 5) 
# transform one of the subindices and make another matrix 
f(x) = (x + 2) .% 5 + 1 
C = sparse(idx1, f(idx2), [7,8,9], 4, 5)

to v0.7+ code:

using SparseArrays
A = sparse([1,3,4], [2,4,3], [1,2,3], 4, 5)
idx = findall(!iszero, A)

# 1. make a sparse matrix using the new cartesian indices
function SparseArrays.sparse(IJ::Vector{<:CartesianIndex}, v, m, n)
    IJ′ = reinterpret(Int, reshape(IJ, 1, :))
    return sparse(view(IJ′, 1, :), view(IJ′, 2, :), v, m, n)
end 
B = sparse(idx, [4,5,6], 4, 5) 

# 2. Or use the subindices to create B
subidx1 = [i[1] for i in idx]
subidx2 = [i[2] for i in idx]
B == sparse(subidx1, subidx2, [4,5,6], 4, 5)

# 3. Transform one of the subindices and make another matrix 
f(ind) = CartesianIndex(ind[1], (ind[2] + 2) .% 5 + 1)
C = sparse([f(i) for i in idx], [7,8,9], 4, 5)

Please let me know if I should do things differently of course!

briochemc · August 21, 2018, 1:21am

Follow-up questions. I want to create an adjacency matrix for a 3d grid of the ocean wrapped around the globe. I have working code but it is slow IMHO. Can you help me figure out a faster code?

First I use a function that returns a valid index. That function is valid_index (code below) and will output a valid index given a cartesian index input i = (y, x, z) which represents the (lat, lon, depth) index tuple. valid_index also needs the size of each dimension (nlat, nlon, and ndepth). It replaces i with its closest valid value. In other words, for depth and latitude, it reassigns any out-of-bounds subindex by the corresponding bound. And for longitude, which is periodic, it reassigns the periodic index by using mod1(x, nlon) (which gives modulo values within 1 and nlon). Here is the function:

function valid_index(j, nlat, nlon, ndepth)
    y = max(1, min(nlat, j[1])) 
    x = mod1(j[2], nlon)
    z = max(1, min(ndepth, j[3]))
    return CartesianIndex((y, x, z))
end

This function might be ugly, slow, and improvable. Please let me know how to improve it (after checking how I use it below)!

Now to the adjacency matrix. I inspired myself from this very nice blog post about cartesian indices in Julia. I loop through all the cartesian indices of a 3d matrix of my model, and at each index, I look at its neighbors, and fill in the adjacency matrix with true values.

function adjacency_matrix(A3d)
    n = length(A3d)
    nlat, nlon, ndepth = size(A3d)
    AM = spzeros(Bool, n, n) # Create an empty Adjacency Matrix (AM)
    L = LinearIndices(A3d) # Linear indices to fill in the AM
    R = CartesianIndices((nlat, nlon, ndepth)) # Range of all cartesian indices
    i1 = first(R) # convenient to create R2 for accessing neighbors in the line below
    R2 = CartesianIndices(-i1, i1) # Range of cartesian indices of neighbors 3 directions
    for i in R # Loop through every cartesian index i
        li = L[i] # linear index for rows
        for k in R2 # loop through all neighbor indices j = i + k
            lj = L[valid_index(i + k, nlat, nlon, ndepth)] # "valid" linear index for columns
            AM[li, lj] = true # Filling in the AM
        end
    end
    return AM
end

The code is run as follows (v0.7+):

using SparseArrays
A3d = zeros(10, 10, 10)
AM = adjacency_matrix(A3d)

Feel free to change the sizes to see how slow it becomes for large values. Any ideas on how to improve efficiency?

mbauman · August 21, 2018, 3:15pm

I’d highly recommend profiling so you can see exactly which part is taking up most of the time so you can focus your efforts.

https://docs.julialang.org/en/stable/manual/profile/#Profiling-1

Without even looking myself, though, I can tell you that assigning into sparse matrices with integer indices AM[li,lj] = … is really slow — and that’s gonna be your culprit. It’s much better to create a list of indices and values and then create it all in one go (or otherwise work within the internal CSC structure).

Remember that for each scalar indexing operation CSC, the matrix needs to do a binary search to find the given row with several levels of indirection. This is slow. For assigning into CSCs, you pay that cost as well as the cost of the additional work that needs to be done to shift everything around the inserted element. That is catastrophic in an inner loop.

barche · August 21, 2018, 9:14pm

I couldn’t resist taking a stab at this, especially since I’ve been meaning to take a closer look at the SparseArrays API anyway. Here are 3 versions, including a slightly modified version of your initial code to avoid allocating the dense array, and also fixing an error I got in the line defining R2:

function valid_index(j, nlat, nlon, ndepth)
  y = max(1, min(nlat, j[1])) 
  x = mod1(j[2], nlon)
  z = max(1, min(ndepth, j[3]))
  return CartesianIndex((y, x, z))
end

function adjacency_matrix(nlat, nlon, ndepth)
  L = LinearIndices((nlat, nlon, ndepth)) # Linear indices to fill in the AM
  n = length(L)
  AM = spzeros(Bool, n, n) # Create an empty Adjacency Matrix (AM)
  R = CartesianIndices((nlat, nlon, ndepth)) # Range of all cartesian indices
  i1 = first(R) # convenient to create R2 for accessing neighbors in the line below
  R2 = CartesianIndices(UnitRange.(Tuple(-i1), Tuple(i1))) # Range of cartesian indices of neighbors 3 directions
  for i in R # Loop through every cartesian index i
      li = L[i] # linear index for rows
      for k in R2 # loop through all neighbor indices j = i + k
          lj = L[valid_index(i + k, nlat, nlon, ndepth)] # "valid" linear index for columns
          AM[li, lj] = true # Filling in the AM
      end
  end
  return AM
end

function is_valid_index(j, nlat, nlon, ndepth)
  return 1 <= j[1] <= nlat && 1 <= j[3] <= ndepth
end

function count_nnz(nlat, nlon, ndepth)
  result = 0
  L = LinearIndices((nlat, nlon, ndepth)) # Linear indices to fill in the AM
  n = length(L)
  R = CartesianIndices((nlat, nlon, ndepth)) # Range of all cartesian indices
  i1 = first(R) # convenient to create R2 for accessing neighbors in the line below
  R2 = CartesianIndices(UnitRange.(Tuple(-i1), Tuple(i1))) # Range of cartesian indices of neighbors 3 directions
  for i in R # Loop through every cartesian index i
      li = L[i] # linear index for rows
      for k in R2 # loop through all neighbor indices j = i + k
          if is_valid_index(i + k, nlat, nlon, ndepth)
            result += 1
          end
      end
  end
  return result
end

function set_indices!(target_array, target_start_idx, L, i, R2, nlat, nlon, ndepth)
  nb_valid = 0
  for k in R2
    j = i+k
    if is_valid_index(j, nlat, nlon, ndepth)
      target_array[target_start_idx+nb_valid] = L[valid_index(j, nlat, nlon, ndepth)]
      nb_valid += 1
    end
  end
  sort!(@view(target_array[target_start_idx:target_start_idx+nb_valid-1]))
  return nb_valid
end

function adjacency_matrix2(nlat, nlon, ndepth)
  L = LinearIndices((nlat, nlon, ndepth)) # Linear indices to fill in the AM
  n = length(L)
  R = CartesianIndices((nlat, nlon, ndepth)) # Range of all cartesian indices
  i1 = first(R) # convenient to create R2 for accessing neighbors in the line below
  R2 = CartesianIndices(UnitRange.(Tuple(-i1), Tuple(i1))) # Range of cartesian indices of neighbors 3 directions
  nnz_AM = count_nnz(nlat,nlon,ndepth)
  colptr = Vector{Int}(undef,n+1)
  rowval = Vector{Int}(undef,nnz_AM)
  start_idx = 1
  colptr[1] = start_idx
  for li in 1:n
    i = R[li] # cartesian index for row
    start_idx += set_indices!(rowval, start_idx, L, i, R2, nlat, nlon, ndepth)
    colptr[li+1] = start_idx
  end
  return SparseMatrixCSC(n, n, colptr, rowval, fill(true, nnz_AM))
end

function adjacency_matrix3(nlat, nlon, ndepth)
  L = LinearIndices((nlat, nlon, ndepth)) # Linear indices to fill in the AM
  n = length(L)
  Is = Int[] # linear i-indices
  Js = Int[] # linear j-indices
  R = CartesianIndices((nlat, nlon, ndepth)) # Range of all cartesian indices
  i1 = first(R) # convenient to create R2 for accessing neighbors in the line below
  R2 = CartesianIndices(UnitRange.(Tuple(-i1), Tuple(i1))) # Range of cartesian indices of neighbors 3 directions
  for i in R # Loop through every cartesian index i
      li = L[i] # linear index for rows
      for k in R2 # loop through all neighbor indices j = i + k
          lj = L[valid_index(i + k, nlat, nlon, ndepth)] # "valid" linear index for columns
          push!(Is,li)
          push!(Js,lj)
      end
  end
  return sparse(Is, Js, fill(true,length(Is)))
end

Test code:

using BenchmarkTools
using Test

const N = (15,15,15)

AM = adjacency_matrix(N...)
println("AM time:")
@btime adjacency_matrix(N...)

AM2 = adjacency_matrix2(N...)
@test AM == AM2
println("AM2 time:")
@btime adjacency_matrix2(N...)

AM3 = adjacency_matrix3(N...)
@test AM == AM3
println("AM3 time:")
@btime adjacency_matrix3(N...)

Result on my machine:

AM time:
  133.744 ms (37 allocations: 2.28 MiB)
AM2 time:
  1.673 ms (3382 allocations: 916.23 KiB)
AM3 time:
  3.095 ms (55 allocations: 5.66 MiB)

As @mbauman suggested, directly constructing the sparse array data structures (counting the number of elements before allocation, even) results in a massive performance gain. I haven’t checked how this can be optimized further, 3382 allocations suggests room for improvement, though it could be the @view and I have no quick fix for that.