Vectorize computation of the steady state distribution

I am trying to speed up the following code that computes a steady state distribution of a dynamic system.

The idea here is that I start with a state (ii,jj,kk) and need to map it to a state (ii',jj',kk'). The transition rule for the first dimension is random so that each index ii can move to any index in (ii1',...,iiN') according to a probability vector z_P(ii,:). The transitions along the 2nd and 3rd dimensions are deterministic according to rules kp_ind(ii,jj,kk) and bp_ind(ii,jj,kk), respectively. The initial state of the system is descirbed by array m_ini. Since I am looking for a steady state distribution, I need to iterate till the transition array converges in the sense specified in the code below.

I have been trying to speed-up my code but without much success. I was wondeing if there is any way to improve upon my code.

------------------------------------------------------------------------------------------------

The code I wote is the following

% kp_ind is array (z_grid_num-by-k_grid_num-by-b_grid_num) 
% bp_ind is array (z_grid_num-by-k_grid_num-by-b_grid_num) 
% z_P is a matrix (z_grid_num-by-z_grid_num); z_P(ii,ii') is the probability that we move from state ii to ii'

while diff_m>eps_stop_m
     
    % updated measure
    m_new = zeros(z_grid_num,k_grid_num,b_grid_num);

    for jj = 1:k_grid_num % old capital
        for kk = 1:b_grid_num % old debt
            for ii = 1:z_grid_num          
                m_new(:,kp_ind(ii,jj,kk),bp_ind(ii,jj,kk)) = m_new(:,kp_ind(ii,jj,kk),bp_ind(ii,jj,kk))...
                                                             + m_ini(ii,jj,kk)*z_P(ii,:)';               
            end
        end
    end

    % checking convergence
    diff_m = sum(abs(m_new -m_ini),'all');
    
    % updating
    m_ini = m_new;
    
end

I am sure that droping one of the loops would help tremendously. The problem is that different states today can be mapped to the same state tomorrow. That's why m_new appears on both sides of the main equations that I loop over. And note that each index (ii,jj,kk) is mapped to a vector of indices (:,kp_ind(ii,jj,kk),bp_ind(ii,jj,kk)) because of the random transition along the first dimension.

------------------------------------

Here is what I tried. I thought I can compute a matrix of indices to which I transition from states (:,jj,kk) (i.e., from a vector of states). The matrix of such indices can be easily computed as

ind_mat = Z+(kp_ind(:,jj,kk)-1)*k_grid_num+(bp_ind(:,jj,kk)-1)*kb_num

where Z is a matrx Z = [(1:z_grid_num)' ... (1:z_grid_num)'] (i.e., possible state along each dimension to which the system can move). The problem is that I will typically obtain ind_mat given as follows (for a simple case where k_grid_num=b_grid_num=z_grid_num=3):

ind_mat = [1 1 1; 1 2 2; 2 3 3]

so that a multiple of states today are mapped to a single state tomorrow. So I cannot make an assignment of the form

 m_new(ind_mat) = m_new(ind_mat) +  z_P.*m_ini(:,jj,kk);