Torch/numpy：groupby 熊猫替代品

Question

Is there some efficient way to rewrite the following code to avoid installing & importing pandas and using torch / numpy instead? I am used to work with pandas , so I wrote it like this, but I am trying to learn numpy and torch , so I am looking for alternative solutions that do not use pandas .

bins = torch.LongTensor(3072).random_(0, 35)
weights = torch.rand((3072))
df = pd.DataFrame({'weights': weights.numpy(), 'bins': bins.numpy()})
bins_sum = df.groupby('bins').weights.sum().values

So, basically: how, without using pandas , get a sum of weights grouped by bins ?

Answer 1

You can compute unique elements of bins via torch.unique (the values to group by) and then use index masks for accessing the corresponding elements in weights :

unique = torch.unique(bins)
result = torch.zeros(unique.size(), dtype=weights.dtype)
for i, val in enumerate(unique):
    result[i] += weights[bins == val].sum()

Answer 2

numpy has isin which is similar to pandas.isin

Pandas groupby , select the data ( row ) and apply function on the group .

def groupby(data, bin_data, grouper, agg):
  '''
    data:     numpy    array
    bin_data: bin's    data
    grouper:  callable, which give returns a list of array values.
    agg:      callable, to be applied on group
  '''

  res = {}
  for key,arr in grouper(data, bin_data):
    res.update({key, agg(arr)})

  return res

# Find the indices where `bins == b` and then use them to select the `arry` values
bin_grouper = lambda arry, bvalue: [(b, arry[np.isin(bvalue, b)]) for b in bvalue]

# Compute result
gdata = groupby(weights.numpy(), bins.numpy(), bin_grouper, np.sum)