Conditional statistics over 1 dimension of a multidimensional array

General Usage
,
1

I would like to compute conditional statistics over one dimension of a two-dimensional array. I have an array a and I would like to compute a mean over the first dimension, taking into account only the cells where b < 0.5. With numpy this is easy with masked arrays, but I haven’t found an Julia equivalent for this. What is the Julian way of solving this?

using Statistics

nx = 16
nz = 8

a = rand(nx, nz)
b = rand(nx, nz)

# This gives a vertical profile of 8 layers. This is the right shape,
# but without the mask applied.
a_prof = mean(a, dims=1)
println(a_prof)

# This gives a scalar, but I would like a vertical profile of 8 layers,
# with only the indices where b < 0.5 in the computation.
a_prof_b = mean(a[b .< 0.5], dims=1)
println(a_prof_b)
2

One suggestion:

a_prof_b = [mean(a[b[:,i] .< 0.5, i]) for i in 1:size(a,2)]
1 Like
3

Looks like your arrays are related and have corresponding indices, so it makes sense to keep and use them together:

using SplitApplyCombine
using StructArrays

# reproduce your simple mean of a with SplitApplyCombine functions:
map(splitdims(a, 2)) do x
	mean(x)
end

# combine a and b arrays to use them together:
AB = StructArray(; a, b)

# compute the desired conditional mean:
map(splitdims(AB, 2)) do x
	mean(x.a[x.b .< 0.5])
end
1 Like