generalized find() with multiple occupation a vectorized version

Question

I have the following process where N particles are dealt among M boxes lattice. Each box can be occupied by an arbitrary number of particles. I store the state of my lattice in a vector.

For example, state vector [1 1 1 0] means 1 particle in box 1, 1 particle in box 2, 1 particle in box 3, 0 particles in box 4. While state vector [2 1 3 0] means 2 particles in box 1, 1 particle in box 2, 3 particles in box 3, and 0 particles in box 4.

From this vector, I extract the particles' position in a different vector for further calculations. Thus, for the state [1 1 1 0] (N<=M) the position vector is [1 2 3] which is easily achieved in Matlab using

pos=find([1 1 1 0])

For a more general case vector like [2 1 3 0] (N>M) the positions vector is [1 1 2 3 3 3]. I achieve this via a custom Matlab function

  pos=findGeneral([2 1 3 0])

The code for this function is

function [position]=findGeneral(state,N)

if 1*isempty(state(state>=2))==0  
    r=find(state);
    state=state(state~=0);        
    position=zeros(1,N);
    t=cell(1,length(r));    
    for i=1:length(r)
        t{i}=r(i)*ones(1,state(i));
    end
    position(1:N)=cell2mat(t);    
else
    position=find(state);    
end

end

I am trying to run this function when my state vector is a gpuArray, but apparently, cell2mat() is not compatible with gpuArrays arguments.

Is there a way to vectorise lines 7--11 of my code so I get rid of the for loop and the cellarrays?

Answer 1

You can go from a "position" to a "state" using

position = [1 1 2 3 3 3];
state = sum( position.' == (1:4) );
% state = [2, 1, 3, 0]

You can go from a "state" to a "position" using

state = [2, 1, 3, 0];
position = repelem( 1:numel(state), state );
% position = [1, 1, 2, 3, 3, 3]

Untested on a gpuArray (just on normal arrays), but from the docs, it looks like numel, repelem and sum are all supported by gpuArrays.