A great option! I'll confess I was a bit tired, so after figuring out something that would broadcast over n-dimensional inputs better than using where, I stopped. Your code is clearly better.