I have a matrix of integers, phase_space
of shape (n,n)
, where each entry represents the number of points in that location in space. I also have tw
We can use np.bincount -
M,N = u_x.max()+1,u_y.max()+1
ids = u_x*N+u_y
out = np.bincount(ids.ravel(),phase_space.ravel(),minlength=M*N).reshape(M,N)
Sample run on a more generic setup -
In [14]: u_x
Out[14]:
array([[1, 2, 1],
[0, 1, 4],
[0, 0, 0]])
In [15]: u_y
Out[15]:
array([[2, 1, 2],
[6, 0, 1],
[2, 6, 0]])
In [17]: phase_space
Out[17]:
array([[1, 1, 1],
[5, 1, 1],
[1, 1, 1]])
In [18]: out
Out[18]:
array([[1., 0., 1., 0., 0., 0., 6.],
[1., 0., 2., 0., 0., 0., 0.],
[0., 1., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0.],
[0., 1., 0., 0., 0., 0., 0.]])
We could also make use of sparse matrices, especially if memory is a concern -
from scipy.sparse import csr_matrix,coo_matrix
out = coo_matrix( (phase_space.ravel(), (u_x.ravel(), u_y.ravel())), shape = (M,N))
Output would be a sparse matrix. To convert to a dense one, use out.toarray()
.
You can use pandas.DataFrame.groupby()
to accumulate all moves with same coordinates in phase_space
:
new_phase_space + (pd.DataFrame(phase_space)
.stack()
.groupby([u_x.ravel(), u_y.ravel()])
.sum()
.unstack(fill_value=0)
.values
)
Output:
array([[2., 2., 4.],
[2., 0., 4.],
[0., 4., 0.]])