Generic and flexible array-of-struct to struct-of-array conversions

General Usage
,
1

I have a problem that comes up fairly often that I do not really have a satisfying solution to at the moment.
Say I have a container type like the one below, representing many “versions” of a real number

struct Baz{T} <: Real
    a::Vector{T}
end

and that I can have arrays of such containers, for instance the following complex version

R = [Complex(Baz(randn(3)), Baz(randn(3))) for i in 1:2, j in 1:2]

julia> typeof(R)
Matrix{Complex{Baz{Float64}}}

When performing some operations on R, I would essentially like to permute the memory layout of R by unpacking the elements of Baz.a one by one so that I get r illustrated by the pseudo code below, perform some operation on r and then repackage the result back into the structure Baz again.

function eigvals(R::Matrix{Complex{Baz{T}}}) where T
    # Allocate E
    for all k
        r[i,j] = Complex(R[i,j].re.a[k], R[i,j].im.a[k]) # "unpack R into r"
        e = eigvals(r)                                   # operate on r
        E[i].re.a[k] = e[i].re                           # repackage result into Baz
        E[i].im.a[k] = e[i].im
    end
    E
end

So far, this can be done manually without too much of a problem, but there might be variations of the problem, for instance, if instead of eigvals I call svdvals, the result is a vector of real values rather than complex, and I would like to do something like this instead:

function svdvals(R::Matrix{Complex{Baz{T}}}) where T
    for all k
        r[i,j] = Complex(R[i,j].re.a[k], R[i,j].im.a[k])
        s = svdvals(r)
        S[i].a[k] = e[i]
    end
    S
end

The data layout of Baz is chosen to fit well for scalar operations, but when Baz appears in more complicated contexts, I lack a general solution that doesn’t require customizing the implementation for each function I want to handle, i.e., eigvals, svdvals above. I’ve been looking towards Accessors.jl which could potentially serve as a piece of the puzzle, but I can’t really find something that does exactly what I need. For completeness, below is a manual implementation of eigvals that does what I want

using LinearAlgebra
struct Baz{T} <: Real
    a::Vector{T}
end

Base.show(io::IO, baz::Complex{Baz{T}}) where T = print(io, "Complex{Baz{$T}[($(baz.re.a), $(baz.im.a))")

function LinearAlgebra.eigvals(R::Matrix{Complex{Baz{T}}}) where T
    K = length(R[1].re.a)
    # Allocate E
    E = similar(R, size(R,1))
    for i in eachindex(E)
        E[i] = Complex(Baz(zeros(K)), Baz(zeros(K)))
    end
    r = zeros(Complex{T}, size(R)...)
    for k in 1:K
        for j in eachindex(R)
            r[j] = Complex(R[j].re.a[k], R[j].im.a[k]) # "unpack R into r"
        end
        e = eigvals(r)                                 # operate on r
        for i in eachindex(e)
            E[i].re.a[k] = e[i].re                     # repackage result into Baz
            E[i].im.a[k] = e[i].im
        end
    end
    E
end

R = [Complex(Baz(randn(3)), Baz(randn(3))) for i in 1:2, j in 1:2]
E = eigvals(R)

Generally, the array R may have any number of dimensions and produce a result of any number of dimensions, and both R and E may be arrays of either complex or real versions of Baz.
Any idea that would help in implementing functions like the above would be greatly appreciated!

1 Like
2

Have you looked into StructArrays.jl? Even though the package has always mystified me, I believe that some accessor magic can be worked out.

Consider your matrix R in the following form

R = let
    R = [Complex(Baz(randn(3)), Baz(randn(3))) for i in 1:2, j in 1:2]
    StructArray(R; unwrap = t -> t <: Baz)
end;

Then you can easily access the imaginary / real part like so S.re.a or S.im.b. I understand this is still far from solving your problem but may shine some light onto the problem.

1 Like
3

Thank you @FMeirinhos ! I have looked at the docs of StructArrays but got a bit confused by all the options :slight_smile: Thank you for showing me this trick, it can certainly simplify writing these custom functions.

4

It occured to me that if isbitstype(Baz) I could use reinterpret to get an array which I could easily do permutedims on and then operate using mapslices etc., but due to the vector in Baz, thsi is not possible. Perhaps there’s a tool, maybe StructArrays is this tool, that could provide a flatten/reinterpret that is flexible enough to handle R in the example?

5

Undoubtedly the problem is screaming lenses but I am not sure in practical terms (i.e., finite amount of developing time) you can use them… Hopefully I am proven wrong!

I would say that the absolute easiest thing you can do is to have the complex type at the lowest level Baz{Complex{T}} instead of Complex{Baz{T}}, since that way you could leverage the algebraic rules / functions of a complex T instead of complex Baz{T}. But I guess that’s not feasible for your use cases. But it would make all of these custom multiple dispatches essentially trivial.

6

The problem of choosing T{S} vs. S{T] where both T and S are <: Number is something I encounter quite often, and I’ve come to the conclusion that you want both representation in different contexts. This arises when, for instance, you combine any of the following, Unitful.Quantity, Complex, ForwardDiff.Dual, MonteCarloMeasurements.Particles and the list goes on…
Implementing functionality with method overloading only goes so far when it comes to combining special number types like this. It’s primarily the outermost type that determines dispatch, so, to make a number behave like Complex, Complex has to be the outermost type.

If I were to have Baz{Complex{T}}, it would not make sense to define Baz <: Real, and I’m not allowed to define Baz{T} <: T, so then I’m left with either defining Baz <: Number or having a separate ComplexBaz. There are many hacky workarounds to this problem, but they all fell like just that, hacky workarounds :confused:

7

Ohhh, I do feel you. I think I have lost the war against Dual Complexs a long time ago…

Related with your original problem, possibly LazyStack.jl could be used in tandem

using StructArrays
using LazyStack

C = let
    C = [Complex(Baz(randn(3)), Baz(randn(3))) for i in 1:2, j in 1:2]
    C = StructArray(C; unwrap = t -> t <: Baz)
end;

C_re = stack(C.re.a)
C_im = stack(C.im.a)

R = let
    R = [Baz(randn(3)) for i in 1:2, j in 1:2]
    StructArray(R; unwrap = t -> t <: Baz)
end;

R_ = stack(R.a)
1 Like