Custom StaticArray matrix type constructors from sub-arrays of a larger matrix type

I have a Matrix2 and Matrix3 type derived from StaticArrays.FieldMatrix:

(Practice) julia> x = Matrix2(1, 2, 3, 4)
2×2 Matrix2{Int64} with indices SOneTo(2)×SOneTo(2):
 1  3
 2  4

(Practice) julia> y = Matrix3(1, 2, 3, 4, 5, 6, 7, 8, 9)
3×3 Matrix3{Int64} with indices SOneTo(3)×SOneTo(3):
 1  4  7
 2  5  8
 3  6  9

I know Julia has some excellent matrix composition functions, and I’m wondering what would be the best way to create a new Matrix2 outer constructor method such that this is produced:

(Practice) julia> Matrix2(y)
2×2 Matrix2{Int64} with indices SOneTo(2)×SOneTo(2):
 1  4
 2  5

That is, it aligns a window of the 2x2 matrix over the the top left of the 3x3 matrix and extracts those elements, discarding the rest. I’m wondering if this can be done in a general fashion to work with other pairs of square matrices, rather than hard-coding the fields to extract for each constructor method.

I’ve been staring at the docs for a while and not seeing anything too useful here. reshape requires the number of elements not to change. resize! is for vectors. The only things that looked possibly relevant were SubArray types and views into arrays, but I don’t know how to use them with immutable StaticArray types, if it’s even possible.

I need some pointers on what a general solution would be for a fixed set of square matrix types. Ideally I’d like one constructor that does the right thing based on the statically-known square matrix size, rather than writing separate methods with hard-coded fields to extract for each.

Any help would be appreciated. Thanks in advance.

Maybe something like this (for it to be efficient, all sizes must be known at compile time):

julia> function sub_matrix(::Val{N}, m::SMatrix{M,M,T}) where {N,M,T}
           SMatrix{N,N,T}(@view(m[CartesianIndices((1:N,1:N))]))
       end
sub_matrix (generic function with 1 method)

julia> m = SMatrix{3,3,Int,9}(1,2,3,4,5,6,7,8,9)
3×3 SMatrix{3, 3, Int64, 9} with indices SOneTo(3)×SOneTo(3):
 1  4  7
 2  5  8
 3  6  9

julia> sub_matrix(Val(2), m)
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 1  4
 2  5

julia> sub_matrix(Val(1), m)
1×1 SMatrix{1, 1, Int64, 1} with indices SOneTo(1)×SOneTo(1):
 1

julia> sub_matrix(Val(3), m)
3×3 SMatrix{3, 3, Int64, 9} with indices SOneTo(3)×SOneTo(3):
 1  4  7
 2  5  8
 3  6  9

julia> @btime sub_matrix(Val(2), $m)
  6.018 ns (0 allocations: 0 bytes)
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 1  4
 2  5


You can also add offsets, to sample some other part of the larger matrix:

julia> function sub_matrix(::Val{N}, m::SMatrix{M,M,T}; offset=(0,0)) where {N,M,T}
           r1 = 1+offset[1]:N+offset[1]
           r2 = 1+offset[2]:N+offset[2]
           SMatrix{N,N,T}(@view(m[CartesianIndices((r1,r2))]))
       end
sub_matrix (generic function with 1 method)

julia> sub_matrix(Val(2), m)
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 1  4
 2  5

julia> sub_matrix(Val(2), m; offset=(1,1))
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 5  8
 6  9

julia> @btime sub_matrix(Val(2), $m; offset=(1,1))
  14.786 ns (0 allocations: 0 bytes)
2×2 SMatrix{2, 2, Int64, 4} with indices SOneTo(2)×SOneTo(2):
 5  8
 6  9


1 Like

Thank you very much! I’m currently thinking how to extend this to support a Val larger than the matrix. Currently Val(n) must be <= m, but it would be nice to be able to say copy only the 3x3 portion of a 4x4 matrix into a new 3x3 matrix.

It would also be nice to be able to extend this so that copying a smaller matrix to a larger one leaves the elements along the main diagonal that it doesn’t touch to 1, with the rest 0 (copy a smaller sub-matrix into a larger identity matrix).

These things are going to take me a while to think about. But thank you for the starting point!

You can have a lot of flexibility if you use a mutable static matrix as an intermediate local to the function:

julia> function newm(::Val{N}, m::SMatrix{M,M,T}) where {N,M,T}
           a = zeros(MMatrix{N,N,T})
           for i in 1:N
               a[i,i] = one(T)
           end
           a[1:M,1:M] .= m
           return SMatrix{N,N,T}(a)
       end
newm (generic function with 1 method)

julia> newm(Val(4), m)
4×4 SMatrix{4, 4, Int64, 16} with indices SOneTo(4)×SOneTo(4):
 1  4  7  0
 2  5  8  0
 3  6  9  0
 0  0  0  1

julia> @btime newm(Val(4), $m)
  11.524 ns (0 allocations: 0 bytes)
4×4 SMatrix{4, 4, Int64, 16} with indices SOneTo(4)×SOneTo(4):
 1  4  7  0
 2  5  8  0
 3  6  9  0
 0  0  0  1


2 Likes

Note that it might be somewhat faster to not allocate the intermediate mutable matrix, and construct this by concatenating block matrices (which all produce immutable intermediate types):

julia> function newm2(::Val{N}, m::SMatrix{M,M,T}) where {N,M,T}
           r = range(0, 0, length=(N-M)*M)
           Z = SMatrix{M,N-M,T,M*(N-M)}(r)
           A = hcat(m, Z)
           B = hcat(Z', Diagonal(SVector{N-M,T}(range(1,1,length=N-M))))
           vcat(A, B)
       end
newm2 (generic function with 1 method)

julia> newm(Val(6), m) == newm2(Val(6), m)
true

julia> @btime newm(Val(6), $m);
  33.430 ns (0 allocations: 0 bytes)

julia> @btime newm2(Val(6), $m);
  19.218 ns (0 allocations: 0 bytes)

Maybe I’m missing some of the requirements (on a phone), but why not index with SVectors?

As in

ind = SA[1, 2]
y[ind, ind]

at least for the simplest case?