Np.select equivalent to deal with multiple conditions

Hi,

I am looking for an equivalent to np.select.

How would you reproduce the following python code in Julia so it would work with n conditions?

x = np.arange(6)
condlist = [x<3, x>3]
choicelist = [x, x**2]
np.select(condlist, choicelist, 42)
array([ 0,  1,  2, 42, 16, 25])

Thanks!

I’m not aware of an equivalent to that function, but it doesn’t seem too hard to implement. (And since this is Julia, our implementations can be as fast as ones from libraries.)

julia> x = 0:5
0:5

julia> condmat = [x .< 3  x .> 3]
6Ă—2 BitMatrix:
 1  0
 1  0
 1  0
 0  0
 0  1
 0  1

julia> choicemat = [x x.^2]
6Ă—2 Matrix{Int64}:
 0   0
 1   1
 2   4
 3   9
 4  16
 5  25

julia> function conditionalselect(conditionmatrix, choicematrix, default = zero(eltype(choicematrix)))
         @assert size(conditionmatrix) == size(choicematrix)
         result = fill(default, size(choicematrix, 1))
         for (rownum, conditionrow) in pairs(eachrow(conditionmatrix))
           choice = findfirst(conditionrow)
           if !isnothing(choice)
             result[rownum] = choicematrix[rownum, choice]
           end
         end
         result
       end
conditionalselect (generic function with 1 method)

julia> conditionalselect(condmat, choicemat, 42)
6-element Vector{Int64}:
  0
  1
  2
 42
 16
 25
1 Like
julia> [ x<3 ? x : x>3 ? x^2 : 42 for x in 0:5]
7-element Vector{Int64}:
  0
  1
  2
 42
 16
 25

is simple and readily extendible.

2 Likes

Since this is Julia discourse: it also only allocates the final Vector, not all those intermediates.

2 Likes

Although it’s possible to implement a select function in Julia, I would re-think your whole approach. For both performance and clarity you really want to do this sort of thing in a single loop or comprehension as demonstrated by @Dan above.

The Matlab or Numpy style of doing a sequence of “vectorized” calls leaves a lot of performance on the table because it constructs a sequence of temporary arrays just to discard them at the end. It’s also, in my opinion, a lot clearer to write a single loop in many cases.

Loops are fast in Julia; you don’t need to fear them like you would in Python. Programming doesn’t have to be an exercise in “mining” the standard library and hoping you can find just the function that you need.

4 Likes

Thanks a lot for the help

I start to really love this language! (few more GIS capabilities and I am sold :slight_smile: )

thanks for the tip - definitely useful to keep in mind when coming from Python