suspendtools/particle.py

26 lines
710 B
Python

def sort(pp,col):
ncol,npart,ntime = pp.shape
assert('id' in col)
for itime in range(0,ntime):
idx = pp[col['id'],:,itime].squeeze().argsort()
pp[:,:,itime] = pp[:,idx,itime]
return pp
def slice_columns(pp,col,keys):
idx_col = []
col_new = {}
ii = 0
for key in keys:
idx_col.append(col[key])
col_new[key] = ii
ii+=1
return pp[idx_col,:,:], col_new
def translate_circular(pp,col,translation,bounds,axis=0):
'''Translates particles while taking into account
the bounding box'''
assert(axis<3)
L = bounds[2*axis+1]
keys = ('x','y','z')
pp[col[keys[axis]],:,:] = (pp[col[keys[axis]],:,:]+translation)%L
return pp